Learn Your Reference Model for Real Good Alignment

The anchor problem in offline alignment

Alignment starts from a bargain: make the model produce what people prefer, but keep it close to a trusted starting point. The classic objective maximizes reward under a KL penalty, $\max_{\pi_\theta}\,\mathbb{E}\left[r_\phi(x,y)\right]-\beta\,\mathbb{D}_{\mathrm{KL}}\left[\pi_\theta\,\|\,\pi_{\mathrm{ref}}\right]$, where the reference $\pi_{\mathrm{ref}}$ is usually the supervised fine-tuned model. Offline methods fold the reward away and train directly on preference pairs.

DPO is the workhorse of that family. For a prompt $x$ with a preferred completion $y_w$ and a rejected one $y_l$, it pushes the policy's preference margin past the reference's:

IPO and KTO reshape the same ingredients, and all three share one convention: $\pi_{\mathrm{ref}}$ stays frozen for the entire run. All three also share a failure mode. Model quality rises, then falls, as the KL divergence from the reference grows, even though no reward model exists to hack. Rafailov and colleagues traced this offline overoptimization to probability mass escaping onto out-of-domain text.

Why a frozen reference stalls training

The DPO gradient is scaled by $\sigma(s)$, where $s$ measures how far the policy's log-preference margin has moved past the reference's. At the start $s\approx 0$ and updates are strong. As training widens the margin, $s$ drifts negative and $\sigma(s)$ decays toward zero, and we show the Hessian decays with it. The loss landscape flattens exactly where the model has moved furthest.

Here is the trap. In practice the probabilities of chosen and rejected completions both fall during DPO training, an effect Pal and colleagues documented. A probability distribution has to put that mass somewhere, so it flows onto sequences the training data never mentions. Once the curvature has vanished, there is no force left to pull it back.

We can watch this happen in a system small enough to enumerate: a depth-three MDP with three actions per state and an RNN policy. We fine-tune on three trajectories, then align with the single pair $(a_1,a_1,a_0)\succ(a_0,a_0,a_0)$, averaging 100 seeds. DPO steadily inflates the seven trajectories neither dataset contains. The same run with an updated reference inflates the chosen trajectory instead.

Let the anchor walk

Our fix is to move the reference during training, the way trust-region optimizers re-center their constraint on the current iterate. The soft variant merges the policy into the reference at every step, $\pi_{\mathrm{ref}}\leftarrow\alpha\,\pi_\theta+(1-\alpha)\,\pi_{\mathrm{ref}}$, under a stop-gradient. The hard variant simply copies, $\pi_{\mathrm{ref}}\leftarrow\pi_\theta$, every $\tau$ steps. Setting $\alpha=0$ recovers the vanilla method exactly.

The recipe applies to any offline method that keeps an implicit closeness constraint, so we get TR-DPO, TR-IPO, and TR-KTO for free. The price is small: keeping and refreshing the reference adds between 2.9 and 7.4 percent to training time across the three methods, with hard updates the cheaper option.

The update rate has a cliff on each side

How fast should the anchor walk? We swept $\alpha$ and $\tau$ with Pythia 2.8B on Anthropic-HH, judging side-by-side against the frozen-reference baseline with GPT-4. We report the net preference margin, wins minus losses, so zero means no better than the vanilla method.

Both extremes fail for a reason. Update too often and the leash chases the policy so closely that nothing constrains it, which is where TR-KTO posts its worst number, a margin of $-85.4$ at $\tau=64$. Update too rarely and you are back to the frozen anchor. We found $\alpha=0.9$ and $\tau<16$ outright unstable, and the gains at $\alpha$ of 0.5 to 0.6 and $\tau=512$ are the statistically significant ones under probability-of-improvement analysis.

It transfers to real assistants

The sweep used a 2.8B model and one dialogue dataset, so the test that matters is a modern assistant setup. On Llama3-Base 8B, trained on UltraChat and aligned on UltraFeedback, the Trust Region versions beat every vanilla method on AlpacaEval 2 length-controlled win rate, with $\alpha=0.8$ and $\tau=32$ chosen for this setting.

Task-specific results agree. On Reddit TL;DR with Pythia 6.9B, updating the reference improves side-by-side win rates by 8.4 points for DPO, 14.3 for IPO, and 15 for KTO. And the overoptimization analysis closes the loop: at matched KL divergence from the SFT policy, our variants score higher on coherence, correctness, detail, helpfulness, and harmlessness, and hold that advantage at matched generation diversity as well.

What it costs and where it bends

The honest reading of the sweeps is that $\alpha$ and $\tau$ are real hyperparameters. Good defaults exist, near $\alpha=0.6$ and $\tau=512$ for the task-specific runs, but the optimum shifted to $\alpha=0.8$ and $\tau=32$ on Llama3, and KTO punishes a bad setting harshly. Tuning bought us further gains; skipping the tuning can cost more than the method gives.

Two caveats frame the result. Very large $\tau$ means few updates on small datasets, so the hard variant may not scale down gracefully, and adaptive schedules look like the natural next step. And all our judgments flow through GPT-4 as a proxy evaluator, on AutoSxS, AlpacaEval 2, and Arena-Hard alike, so the field's open question about judge reliability applies to us too.