KL penalty

A KL penalty is a term in an RL objective that charges the policy for moving away from a fixed reference policy, measured by Kullback-Leibler divergence. In RLHF the reference policy is usually the supervised fine-tuned model, and the penalty stops the policy from drifting into degenerate text that only a reward model could love.

The objective

Here $\pi$ is the policy being trained, ${\pi }_{\mathrm{ref}}$ is the frozen reference policy, $r$ is the reward and $\beta >0$ sets the strength of the penalty. In practice the KL is estimated per sample as $\log \pi (y\mid x)-\log {\pi }_{\mathrm{ref}}(y\mid x)$, summed over tokens, and is either folded into the reward or added to the loss.

What the penalty does

This objective has a closed-form maximiser:

where $Z(x)$ normalises the distribution. The optimal policy reweights the reference rather than replacing it: responses the reference considers impossible stay impossible, and $\beta$ acts as a temperature on the reward. DPO is derived from exactly this formula.

Worked example

Two responses each have reference probability 0.5; one has reward 1, the other 0. With $\beta =0.5$, the optimal policy puts $e^2/(e^2+1)\approx 0.88$ on the better response. With $\beta =2$ it puts only $e^{0.5}/(e^{0.5}+1)\approx 0.62$ there. As $\beta \to 0$ the policy greedily picks the highest reward; as $\beta \to \infty$ it stays equal to the reference.

Choosing beta

• Too small: the policy finds reward-model blind spots, the signature of reward overoptimisation.
• Too large: almost nothing is learned.
• Adaptive schemes, as in Ziegler et al. (2019), tune $\beta$ to hit a target KL per batch. Reporting KL alongside reward makes runs comparable.

Why it matters

KL from the reference is the most useful single dial for how much a model has changed during RL. Gao et al. (2022) plotted true quality against $\sqrt{\mathrm{KL}}$ and found it peaks and then declines, so KL budget is effectively a safety margin against reward hacking. It also preserves capabilities and calibration learned in pretraining that a narrow reward would otherwise erode.

Related terms

• RLHF (reinforcement learning from human feedback)
• DPO (Direct Preference Optimization)
• Reward overoptimization
• Reward hacking
• PPO (Proximal Policy Optimization)
• Best-of-n sampling

Sources

1. Daniel M. Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B. Brown, Alec Radford, Dario Amodei, Paul Christiano, Geoffrey Irving. Fine-Tuning Language Models from Human Preferences, 2019. arxiv.org
2. Leo Gao, John Schulman, Jacob Hilton. Scaling Laws for Reward Model Overoptimization, 2022. arxiv.org