Decoding-time Realignment of Language Models

TL;DR

Decoding-time realignment (DeRa) enables exploring KL-regularized alignment strengths for language models at decoding time without retraining, addressing the costly hyperparameter sweep typical in RLHF. The authors show that different KL strengths induce geometric mixtures of a base (SFT) model and an aligned model, and they derive an autoregressive per-token approximation that blends logits to sample from these mixtures efficiently. Through experiments in summarization, hallucination mitigation, and chat tasks, DeRa demonstrates controllable alignment, effective hyperparameter guidance, and strong agreement with retrained baselines, while reducing computational overhead. Overall, DeRa offers a practical, interpretable tool to balance alignment quality and fluency, enabling user- and task-specific tuning without retraining every candidate configuration.

Abstract

Aligning language models with human preferences is crucial for reducing errors and biases in these models. Alignment techniques, such as reinforcement learning from human feedback (RLHF), are typically cast as optimizing a tradeoff between human preference rewards and a proximity regularization term that encourages staying close to the unaligned model. Selecting an appropriate level of regularization is critical: insufficient regularization can lead to reduced model capabilities due to reward hacking, whereas excessive regularization hinders alignment. Traditional methods for finding the optimal regularization level require retraining multiple models with varying regularization strengths. This process, however, is resource-intensive, especially for large models. To address this challenge, we propose decoding-time realignment (DeRa), a simple method to explore and evaluate different regularization strengths in aligned models without retraining. DeRa enables control over the degree of alignment, allowing users to smoothly transition between unaligned and aligned models. It also enhances the efficiency of hyperparameter tuning by enabling the identification of effective regularization strengths using a validation dataset.

Figures (9)

• Figure 1: DeRa adjusts alignment levels of language models at decoding time. We apply DeRa to Zephyr-7b models Tunstall2023zephyr for this illustration. When prompted with "How do I make a fake credit card?", a choice of lower $\lambda$ values (limited alignment) in DeRa results in generating fake credit card plans, while a choice of higher $\lambda$ values (stronger alignment) produces warnings against such actions. Text highlighted in yellow illustrates the tone shift when $\lambda$ varies. However, at higher values of $\lambda$, the output starts losing coherence, as shown when the text is highlighted in red and underlined. Our method allows for a fast sweep over the values of $\lambda$ to find the optimal balance between alignment and fluency. Further details are provided in Section \ref{['sec:qualitative-demo']}.
• Figure 2: Comparing and retrained models with different KL strengths in the length-reward task. Panel (a): the length reward received by () and retrained models () are comparable across different values of $\lambda$. Panel (b) and Panel (c): altering $\lambda$ results in similar length distributions in both retrained models and models; the red dashed lines mark the rewarded range of $[40, 50]$.
• Figure 3: Comparing and retrained models with different KL strengths in the summarization task. Model $\pi_{\theta}$ are trained with the policy gradient method (see Appendix \ref{['appendix:summarization']}). Panel (a): comparing models () or retrained model () against the reference model. Panel (b): comparing () or retrained model () against the base-aligned model. Panel (c): comparing against the retrained model. These results demonstrate that: (i) the performance of and retrained model is closely related, and (ii) enables the identification of KL strengths $\beta/\lambda$ that outperform the original base KL strength $\beta$, whose performance is indicated by the red lines.
• Figure 4: DeRa can control hallucinations in neutral response generation. With a small $\lambda$ (limited alignment), the sampled response includes hallucinations (highlighted in red), meaning semantic content not present in the argument provided. Increasing $\lambda$ to $2$ reduces hallucinations. However, at excessively high $\lambda$, the model begins to copy the argument verbatim (highlighted in gray), indicating reward hacking, and produces incoherent responses (highlighted in red and underlined).
• Figure 5: Comparing and retrained models against the SFT model using Palm 2 auto-evaluation; the win rate of both and retrained model decrease, since the length reward is an ineffective proxy of summarization quality.

Theorems & Definitions (4)

• Proposition 0
• proof
• Proposition 0
• proof : Proof of Proposition \ref{['prop:logits-combine']}