Source-linked AI summary
Decoding-time Realignment of Language Models
Tianlin Liu, Shangmin Guo, Leonardo Bianco, Daniele Calandriello, Quentin Berthet, Felipe Llinares, Jessica Hoffmann, Lucas Dixon, Michal Valko, Mathieu Blondel
TL;DR
Choosing regularization strength is costly because conventional exploration trains multiple aligned models, and the wrong strength can limit alignment or risk reward hacking. DeRa derives alternative strengths during decoding from an aligned model; it identifies strengths that outperform the default and generally agrees with retrained models.
Problem
Finding an effective KL regularization strength conventionally requires repeated alignment training, which is computationally demanding, especially for large models and many strengths.
Source: §3 · p. 3Method
Results
DeRa identifies KL strengths that outperform the default and generally agrees well with retrained models, including in the λ > 1 extrapolation regime.
Source: §5.4 · p. 7Takeaways & Limitations
DeRa can serve as a cost-effective and accurate method for determining effective KL strength hyperparameters.
Source: §5.4 · p. 7Combining logits from two models doubles decoding time and memory compared with standard decoding; combining weights is feasible but incurs a performance penalty.
Source: §4 · p. 4
Abstract
from arXiv · showhide
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.
1. Introduction
The introduction frames alignment as balancing desirable behavior against preserving the original model’s fluency and expressive power. It presents DeRa as a decoding-time way to explore regularization strengths without retraining, supported by a geometric-mixture result and experiments on alignment control and tuning.
1. Introduction
- Language models can exhibit factual errors, biases, and other undesirable behaviors, while alignment aims to guide responses toward human standards.Source: §1 · p. 1
- Too much regularization limits alignment, while too little can cause reward hacking and regressions in other performance characteristics.Finding a balance by sweeping strengths typically requires computationally demanding repeated training, especially for large models.Source: §1 · p. 1
- Under the KL-regularized alignment objective, models at varying regularization strengths are geometric mixtures of one reference model and one aligned model, differing only in mixing weights.Source: §1 · p. 2
- DeRa approximates those mixtures autoregressively, enabling evaluation of different regularization strengths at decoding time without retraining.It allows regularization strength to be adjusted, for example, by user or task.Source: §1 · pp. 1–2
- DeRa experiments show control of alignment strength, faster hyperparameter tuning, and improved downstream-task performance tradeoffs.The illustration shows lower λ producing fake-card plans and higher λ producing warnings, while very high λ reduces coherence.Source: §1 · p. 2, Figure 1 · p. 2
2. Background
The background describes language models as autoregressive response distributions and reviews the finetuning pipeline from demonstrations to human preferences. It frames alignment as a reward–KL tradeoff whose strength affects both alignment and preservation of model capabilities.
2. Background
- Autoregressive factorization makes response log-probabilities easy to compute for maximum-likelihood training and enables sampling at decoding time.The paper identifies transformers as the state-of-the-art language models for modeling these distributions.Source: §2 · p. 2
- Supervised finetuning trains on curated expert demonstrations, while preference finetuning uses human judgments of preferred and unpreferred responses.Preference data can be collected more affordably and abundantly than expert-generated demonstrations, and can support RLHF reward learning.Source: §2 · p. 2
- Alignment typically trades off maximizing expected reward against staying close, in KL divergence, to the SFT model.The query distribution is pX, and β controls deviation from πsft.Source: §2 · p. 2
- A large β limits alignment by keeping the model near SFT, whereas a small β can invite reward hacking and loss of coherence or topicality.Prior work described reward gains alongside less natural samples and summarization quality that eventually deteriorates as models overfit.Source: §2 · p. 3
3. Decoding-time realignment
Decoding-time realignment addresses the cost of training separate aligned models for each regularization strength. It derives a realigned model from an existing aligned model and expresses the change as multiplicative response reweighting controlled by λ.
3. Decoding-time realignment
- Evaluating multiple aligned models trained at different KL strengths is computationally demanding, motivating investigation of the reward–KL tradeoff without retraining each model.Source: §3 · p. 3
- DeRa computes a realigned model π*(β/λ) from an aligned model π*(β) without retraining, with β denoting training-time strength and λ decoding-time strength.Source: §3 · p. 3
- The realigned model multiplicatively reweights response probabilities using an importance ratio, whose importance λ modulates.The approach can also extend to a linear combination of multiple rewards.Source: §3 · pp. 3–4
4. Approximation and implementation
DeRa approximates realigned sequence distributions with per-token sampling, using a configurable mixture of reference and aligned-model logits. It avoids retraining to explore regularization strengths, but combining two models increases decoding cost.
4. Approximation and implementation
- The sequence-level realigned distribution is intractable to compute, so DeRa approximates it with per-token conditional distributions for one-token-at-a-time generation.The approximation supports sampling responses without computing a normalization constant over all possible sequences.Source: §4 · p. 4
- The per-token distribution is simplified using the equivalence between a geometric mean and an arithmetic mean in log space.This yields a formulation based on linearly combining the reference and aligned-model logits.Source: §4 · p. 4
- The realignment parameter λ controls regularization: λ = 0 recovers the reference model, λ = 1 the aligned model, and λ > 1 uses weaker regularization than the base KL strength.The method permits λ values above 1, including 2, 5, and 10 in the experiments.Source: §4 · p. 4
- DeRa samples from softmax probabilities of linearly combined reference and aligned-model logits, generating a response token by token.The SFT model is used as the reference in the illustration, but other reference models, including pretrained models, may be used.Source: §4 · p. 4
- DeRa avoids retraining to test regularization strengths, but combining two models doubles decoding time and memory compared with standard decoding.Combining model weights can reduce inference cost, but incurs a performance penalty; alternatively, DeRa can guide a narrower retraining sweep without decoding overhead.Source: §4 · pp. 4–5
5. Experiments
Across summarization, alignment control, and hallucination mitigation experiments, DeRa is used to vary regularization at decoding time and compared with retrained models. The results indicate that it can help identify effective KL strengths and reproduce important behaviors of retrained models.
5.1. Experimental setup
- The experiments compare decoding-time DeRa against models retrained at corresponding KL strengths and cover both policy-gradient and DPO alignment approaches.The paper describes DeRa as independent of the alignment approach, with policy gradient using online reward annotations and DPO using offline preference data.Source: §5.1 · p. 5, §5 · p. 5
5.2. Toy problem: summarization with a length reward
- DeRa and retrained models show comparable length rewards and similar response-length distributions as λ varies in the toy summarization task.The toy reward targets responses in the range [40, 50]; greater length reward can also encourage copying and reduce overall summarization quality.Source: §5.2 · p. 5, Figure 2 · p. 6
5.3. Controlling alignment: a qualitative demonstration
- Changing λ meaningfully controls the degree of alignment in Zephyr-7b models, with quantitative MT-Bench results reported for λ values from 0 to 2.The qualitative demonstration uses λ values of 0, 1/6, 1, 5, and 100.Source: §5.3 · pp. 5–6
5.4. Learning to summarize from human feedback
- DeRa identifies KL strengths that outperform the default, and its performance generally agrees with retrained models, including in the extrapolation regime λ > 1.This agreement supports using DeRa to determine effective KL strengths before retraining.Source: §5.4 · p. 7, Figure 3 · p. 7
- For summarization, DeRa indicates that the base model may be over-regularized: reducing KL strength improves win rate, while under-regularized models benefit from greater KL strength.The under-regularization result is reported in Appendix Figure 7.Source: §5.4 · p. 7
5.5. Hallucination mitigation
- For rewriting pro/con arguments into natural prose without hallucinations, increasing λ improves faithful rewriting; at still higher λ, naturalness declines as the model copies arguments verbatim and loses coherence.The experiment uses PaLM 2 models with KL strength β = 0.1 and varies λ over 0.011, 0.1, 0.5, 0.67, 1, 2, 5, and 10, corresponding to β/λ strengths from 9 to 0.01.Source: §5.5 · p. 7
6. Related work
Related work covers methods for combining multiple reward-trained models, proxy approaches that merge models at the output level, and other decoding methods that merge logits. DeRa instead targets the reward–regularization balance using a reference model and an aligned model.
Multi-reward RLHF
- Multi-reward methods combine independently trained models for different rewards, whereas DeRa combines a reference model and an aligned model.Prior approaches use parameter interpolation or model ensembles to weight rewards such as helpfulness and harmlessness.Source: §6 · pp. 7–8
Proxy approaches for finetuning
- Emulated fine-tuning and proxy-tuning merge trained models at the output level, while DeRa targets the reward–regularization balance and can proxy retraining at different alignment strengths.The cited proxy-tuning work reports performance competitive with standard fine-tuning across benchmarks.Source: §6 · p. 8
Other decoding approaches that merge logits
- Other logit-merging decoding methods target machine translation, text diversity, prompt adherence, or faster generation rather than DeRa’s alignment objective.The cited approaches include fusion, contrastive decoding, classifier-free guidance, and speculative decoding.Source: §6 · p. 8
7. Conclusion
DeRa adjusts KL regularization during decoding by approximating geometric mixtures of a reference and aligned model. It can tailor alignment to users, prompts, or tasks and help identify promising strengths before retraining.
7. Conclusion
- Aligned models at different KL regularization strengths are geometric mixtures of a reference model and an aligned model, differing only in mixing weights.DeRa uses this result to approximate the mixtures autoregressively during decoding.Source: §7 · p. 8
- DeRa adjusts KL regularization during decoding, approximating those mixtures autoregressively with a simple implementation and clear interpretation.Source: §7 · p. 8
- DeRa can tailor regularization to individual users, prompts, or tasks and help identify promising strengths to retrain, avoiding unnecessary retraining across many strengths.The authors report this streamlines hyperparameter tuning and reduces computational costs.Source: §7 · p. 8
Impact statement
The paper frames DeRa as a decoding-time sweep over regularization strengths in language-model alignment. It presents this as a way to streamline hyperparameter selection and reduce retraining runs.
Impact statement
- DeRa sweeps over regularization strengths at decoding time, streamlining hyperparameter selection and reducing retraining runs to find an optimal strength.Source: Impact statement · p. 9
- The work is situated among alignment techniques intended to promote friendly, safe, and harmless language-model responses.Source: Impact statement · p. 9
A. Proof of Proposition 1
The appendix states an equivalent form for the approximate realigned model and extends the realignment argument to linear combinations of multiple rewards.
A. Proof of Proposition 1
- Proposition 1 states that the approximate realigned model defined in (6) can be written equivalently in another form.The supplied proof text identifies entrywise product and division, but does not state the equivalent expression.Source: A. Proof of Proposition 1 · p. 12
B. Linear combination of multiple rewards
- For K reward functions, the appendix defines a linear combination using weights λ and says the realigned-distribution derivation carries through to multiple combinations.It also states that the autoregressive approximation carries through similarly.Source: B. Linear combination of multiple… · p. 12
C. Alignment methods
The paper describes policy-gradient and direct-preference alignment methods, with PPO or vanilla policy gradients used for different tasks. DPO is used for summarization and chat-model alignment.
C. Alignment methods
- For the length-reward problem the authors use PPO, while summarization uses vanilla policy-gradient optimization.PPO replaces reward with general advantage estimation and introduces clipped probability ratios.Source: C. Alignment methods · p. 13
- DPO directly optimizes the policy through a Bradley–Terry-based loss without using a reward function.Its preference data contain a query and favored and unfavored human responses.Source: C. Alignment methods · p. 13
- The authors use DPO for summarization and chat-model alignment.Source: C. Alignment methods · p. 13
D.1. Length-reward experiments
The length-reward experiment illustrates a tradeoff: weaker KL regularization gains reward while stronger regularization retains summarization quality. A separate Zephyr-7b demonstration uses existing SFT and aligned checkpoints to explore sampling settings.
D.1. Length-reward experiments
- T5-small models are trained on XSum and aligned with PPO to maximize a length reward.The SFT model uses 15,000 training steps; policy and value learning rates are 5e-6 and 1e-5.Source: §D.1 · p. 13
- Greater KL strength retains higher summarization quality, while lower strength gains more length reward.Summarization quality is evaluated by Palm 2 comparisons against the SFT model.Source: §D.1 · p. 13
D.2. Controlling alignment: a qualitative demonstration
- The Zephyr-7b qualitative demonstration samples from existing SFT and aligned checkpoints using 4-bit quantization and specified decoding settings.Sampling uses 256 maximum new tokens, temperature 0.7, top-k 50, and top-p 0.95.Source: §D.2 · p. 13
D.3. Summarization experiments
These experiments compare DeRa with retrained models across summarization, hallucination mitigation, and general-purpose chat settings. Results indicate close agreement in some settings and show that changing KL strength can affect task performance.
D.3. Summarization experiments
- In the length-reward summarization experiment, both DeRa and retrained models lose win rate against SFT because length reward is an ineffective proxy for summarization quality.Source: Figure 5 · p. 14
- For small λ, DeRa and retrained models have overall close pairwise accuracy, but their gap increases as λ grows beyond 1.The passage attributes this to DeRa’s approximation being poor in the extrapolation region.Source: §D.3 · p. 14
D.4. Hallucination mitigation
- In hallucination mitigation, DeRa and retrained-model performances are correlated in comparisons against both the reference and base-aligned models.Source: §D.4 · p. 15
- The hallucination reward model’s classifier achieves ROC-AUC 0.986 on the evaluation set.Source: §D.4 · p. 15
D.5. Aligning general-purpose chat models
- DeRa sweeps λ from 0 to 2.0 in increments of 0.2 to evaluate Zephyr-7b variants on MT-Bench, including a parameter-combining variant.GPT-4 scores responses to 160 multi-turn questions spanning eight knowledge areas.Source: §D.5 · p. 16
- 7.23 versus 7.37: retrained Zephyr-7b at λ = 0.5 closely matches the original model’s MT-Bench score while preserving math performance and raising the Open LLM leaderboard score.The original model’s average MT-Bench score is 7.37; the retrained model’s is 7.23.Source: §D.5 · p. 16, Figure 9 · p. 17