The KL penalty and the reference model

The constrained objective and its closed-form solution#

RLHF optimises

$$ \max_{\pi}\;\mathbb{E}{x\sim\mathcal{D},\,y\sim\pi}[r(x,y)] - \beta\,D{\mathrm{KL}}(\pi \,|\, \pi_{\text{ref}}) $$

This is a KL-regularised bandit problem, and it has an exact solution:

$$ \pi^*(y \mid x) = \frac{1}{Z(x)}\,\pi_{\text{ref}}(y \mid x)\exp!\left(\tfrac{1}{\beta} r(x,y)\right), \qquad Z(x) = \sum_{y}\pi_{\text{ref}}(y\mid x)\exp!\left(\tfrac{1}{\beta}r(x,y)\right) $$

The optimal policy is the reference reweighted exponentially by reward. Two consequences follow.

$\beta$ is a temperature on the reward. As $\beta \to \infty$, $\pi^* \to \pi_{\text{ref}}$. As $\beta \to 0$, $\pi^*$ collapses onto $\arg\max_y r(x,y)$ — precisely the 99%-on-the-exploit behaviour in the code above.

$Z(x)$ is intractable, since it sums over all sequences. That intractability is the reason RLHF used RL at all — and rearranging this same equation to eliminate $Z$ is exactly how DPO is derived.

Mode-seeking, not mode-covering#

The objective uses reverse KL, $D_{\mathrm{KL}}(\pi_\theta | \pi_{\text{ref}})$, which is zero-avoiding in $\pi_{\text{ref}}$ but mode-seeking in $\pi_\theta$: the trained policy is free to abandon regions the reference covers. Forward KL would instead force $\pi_\theta$ to cover everything $\pi_{\text{ref}}$ does.

This choice is why RLHF reliably reduces output diversity — mode collapse toward a narrower, safer distribution — and why base models remain more diverse generators than their aligned counterparts. Kirk et al., 2024 (Understanding the Effects of RLHF on LLM Generalisation and Diversity) measured this directly: RLHF improves out-of-distribution generalisation and measurably reduces output diversity relative to SFT, and the trade-off is not incidental to the method.

Wang et al., 2024 (Beyond Reverse KL) generalise the regulariser to any $f$-divergence, giving an explicit knob on the diversity–alignment trade-off. TRL exposes this as f_divergence_type on DPOConfig.

KL estimators#

Only samples from $\pi_\theta$ are available, so KL must be estimated. Writing $\rho = \log\frac{\pi_\theta(y)}{\pi_{\text{ref}}(y)}$ for a sampled $y$:

k3 is the standard choice (Schulman's blog on approximating KL divergence is the canonical reference). It is a control-variate correction of k1: since $\mathbb{E}{\pi\theta}[e^{\rho}] = 1$ under the change of measure, subtracting $e^\rho - 1$ removes variance without introducing bias.

Over-optimisation, quantified#

Gao et al., 2023 (Scaling Laws for Reward Model Overoptimization) fitted the true reward as a function of $d = \sqrt{D_{\mathrm{KL}}(\pi | \pi_{\text{ref}})}$:

$$ R_{\text{bo}n}(d) = d\,(\alpha_{\text{bo}n} - \beta_{\text{bo}n} d), \qquad R_{\text{RL}}(d) = d\,(\alpha_{\text{RL}} - \beta_{\text{RL}}\log d) $$

Both are inverted-U shaped in $d$ — which is what the first table in the code reproduces, at a scale you can run in three seconds. Their empirical findings:

• The peak's location moves further out with larger reward models and more preference data. Policy size barely affects it.
• Best-of-$n$ sampling and RL trace different curves at the same KL, with best-of-$n$ more robust at small KL.
• The KL penalty coefficient does not change the shape; it changes how fast you travel along it. Early stopping and a KL penalty are substitutes rather than complements.

That last point is the most practically useful and the least widely known.

Adaptive KL control#

Fixing $\beta$ fixes the price of divergence, not the divergence. Since the reward scale changes with every reward model, the achieved KL varies wildly across runs at the same $\beta$.

The standard controller (Ziegler et al., 2019) targets a KL value $\mathrm{KL}_{\text{target}}$ and updates

$$ \beta \leftarrow \beta \left(1 + K_\beta \cdot \mathrm{clip}!\left(\frac{\mathrm{KL}{\text{actual}} - \mathrm{KL}{\text{target}}}{\mathrm{KL}_{\text{target}}}, -0.2, 0.2\right)\right) $$

This makes runs comparable across reward models, which is why it is worth the extra hyperparameter.

When the reference can be dropped#

Two situations remove the need for a reference model entirely.

Verifiable rewards. When the reward is a program — a unit test, an exact-match check — there is no proxy to over-optimise, so the Goodhart argument does not apply in the same form. DeepSeek-R1's GRPO recipe and TRL's GRPOConfig default (beta=0.0) both reflect this; see RL with verifiable rewards. The model can still degrade in other ways — readability, language mixing — which is why R1 added separate format rewards rather than a KL term.

Reference-free preference objectives. SimPO (Meng et al., 2024) replaces the reference log-ratio with a length-normalised average log-probability and adds a target margin, removing the reference model's memory and compute cost entirely. It performs competitively; it also gives up the theoretical grounding above.

Papers#

• Ziegler et al., Fine-Tuning Language Models from Human Preferences, 2019 — arxiv.org/abs/1909.08593
• Schulman, Approximating KL Divergence, 2020 — joschu.net/blog/kl-approx.html
• Gao et al., Scaling Laws for Reward Model Overoptimization, ICML 2023 — arxiv.org/abs/2210.10760
• Rafailov et al., Direct Preference Optimization, NeurIPS 2023 — arxiv.org/abs/2305.18290
• Wang et al., Beyond Reverse KL, ICLR 2024 — arxiv.org/abs/2309.16240
• Kirk et al., Understanding the Effects of RLHF on LLM Generalisation and Diversity, ICLR 2024 — arxiv.org/abs/2310.06452
• Meng et al., SimPO: Simple Preference Optimization with a Reference-Free Reward, NeurIPS 2024 — arxiv.org/abs/2405.14734

What to learn next#

• Direct preference optimisation — what happens when you solve this objective in closed form.
• Reward hacking and sycophancy — the failure this penalty is defending against.
• PPO — the algorithm that carries out the optimisation.