Can Decoding-Time Realignment Tune Language Model Alignment Without Retraining?

DeRa moves the regularization dial to decoding

Alignment methods such as reinforcement learning from human feedback (RLHF) trade off two things: maximizing a learned reward for preferred responses, and staying close to the supervised-finetuned model that serves as the reference. That closeness is enforced by a KL penalty with strength β. Too little regularization invites reward hacking, where the model exploits the reward and loses capability; too much blocks alignment. Finding the right β normally means retraining a separate model for each candidate value, which is expensive for large models.

DeRa avoids that retraining. The paper shows that aligned models at different KL strengths are geometric mixtures of the reference distribution and one aligned distribution, differing only in mixing weight. DeRa approximates those mixtures autoregressively while decoding: it linearly combines the logits of the reference model and the aligned model and samples tokens from the resulting softmax. A single scalar λ sets the effective strength to β/λ, so λ = 1 reproduces the original aligned model and λ = 0 reproduces the reference.

Win rates compared DeRa against retrained models

The main quantitative test used a summarization task. The authors drew 1000 queries from the test fold, generated responses with each realigned model, and had PaLM 2 Large (text-unicorn-001) judge which response in each pair was better. Win rate is the share of samples preferred over a stated opponent, so each number is relative to a specific baseline rather than an absolute quality score.

They swept λ over {0.5, 2/3, 1.0, 2.0, 5.0, 10.0}, corresponding to KL strengths β/λ of {0.2, 0.15, 0.1, 0.05, 0.02, 0.01}. Values of λ above 1 are extrapolation: the combined logits fall outside the convex hull of the reference and aligned logits. The authors deliberately included them to stress-test whether DeRa still approximates reasonably there. Comparisons were run against the reference model, against the base aligned model, and directly against retrained-from-scratch models at matching strengths.

DeRa matched retraining and found better strengths

Two findings stand out. First, DeRa identified KL strengths whose win rates exceeded the default strength β, shown as points above the red reference lines in the paper's figures. Second, DeRa and retrained models generally agreed well, and that agreement held even for λ > 1 in the extrapolation regime. The authors read this as evidence that DeRa can serve as a cost-effective but accurate way to pick an effective KL strength before committing to a retraining run.

The same comparison also diagnosed the base model. In the summarization setting, DeRa suggested the aligned model was over-regularized: reducing the KL penalty (λ > 1) improved win rate, and retrained models confirmed it. The appendix reports the mirror case, where DeRa flags an under-regularized model and recommends a larger KL strength with λ < 1. The direction of the recommendation, not just the existence of a gap, is what the sweep produces.

Hallucination control shows the same two-sided tradeoff

A second task tested hallucination mitigation for retrieval-augmented generation. The model had to rewrite lists of pro and con arguments in prose while strictly preserving their meaning and adding no new content. A PaLM 2 model served as reward model via LoRA, and a second PaLM 2 was aligned with LoRA-based RLHF at β = 0.1. DeRa then swept λ over {0.011, 0.1, 0.5, 0.67, 1, 2, 5, 10}, giving strengths β/λ from 9 down to 0.01.

The qualitative result is a tradeoff rather than a single optimum. Low λ keeps behavior close to the reference model and raises the tendency to hallucinate. Increasing λ improves the desired rewriting behavior, but pushing λ further degrades naturalness: the model copies the original arguments verbatim and loses coherence, which the paper marks as reward hacking. This is a qualitative illustration, not a win-rate measurement, so it should not be read as a tuned threshold.

A chat-model sweep suggests limited headroom

The appendix applies DeRa to Zephyr-7b, an open-weight chat model fine-tuned from Mistral 7b with DPO, whose release includes both the SFT model and an aligned model at β = 0.1. Sweeping λ from 0 to 2.0 in steps of 0.2 and scoring on MT-Bench's 160 questions with GPT-4 ratings from 1 to 10, the base strength at λ = 1 was already close to the best score. The authors therefore anticipate limited potential improvement from retuning KL strength for this model.

They still identify λ = 0.5 (β/λ = 0.2) as a sensible alternative, because the MT-Bench score stays near its peak while stronger regularization can protect capabilities that alignment tends to erode. Retraining Zephyr-7b at that strength confirmed the reasoning: MT-Bench averaged 7.23 versus 7.37 for the original, and the Open LLM leaderboard average rose to 61.55 from 59.23, with GSM8K recovering to 30.02 from 12.74. The tradeoff is visible in the same table, where TruthfulQA falls to 54.72 from 57.45.

The savings come with doubled decoding cost

DeRa's benefit is at training time: it avoids retraining across a wide range of regularization strengths and lets the strength be chosen per user, prompt, or task at decoding. The paper notes that combining logits from two models doubles decoding time and memory relative to standard decoding. Combining model weights instead of logits is feasible per the appendix but comes with a performance penalty, consistent with prior weight-combining findings.

Scope matters when reading the evidence. The quantitative win-rate comparison comes from one summarization task with PaLM 2 as judge and 1000 test queries; the hallucination case is qualitative; the chat-model evidence is an appendix sweep on Zephyr-7b scored by GPT-4 on MT-Bench. The method assumes access to both a reference checkpoint and an aligned checkpoint, which the paper notes is available for open-weight models such as Llama-2 and Mistral 7B.

Sweep first, then retrain only the winner

The practical priority the paper establishes is a two-stage workflow: use DeRa to sweep λ on a validation dataset and locate promising strengths, then retrain only at the selected strength. The evidence for this ordering is the close agreement between DeRa and retrained models across the swept range, including extrapolation, plus the concrete case where DeRa's recommendation at λ = 0.5 produced a retrained Zephyr-7b that traded a small MT-Bench drop for a higher overall leaderboard average and much stronger GSM8K.

The judgment to carry away is conditional. When the base model is over- or under-regularized, as in the summarization setting, the sweep points to a direction and a magnitude worth retraining. When the base strength already sits near the peak, as in the Zephyr-7b MT-Bench sweep, the honest conclusion is that retuning buys little on that benchmark, and the case for a stronger KL penalty rests on preserving capabilities like math reasoning rather than on raising the alignment score itself.

Sources used in this answer