KL Divergence Penalty › KL Divergence Penalty

In the previous chapters, we built up the RLHF pipeline: collecting human preferences, training reward models, and applying PPO to optimize language model policies. But there is a necessary problem we have only briefly mentioned: reward hacking. Without constraints, a policy can discover bizarre outputs that score highly on the reward model while being clearly worse by any human measure. The KL divergence penalty is the mechanism that prevents this collapse, keeping the fine-tuned model anchored to the capabilities of its pre-trained foundation.

Think of the KL divergence penalty as the safety tether on a spacewalk. An astronaut tethered to the station can move freely in many directions, exploring the external surface and performing repairs, but cannot float away into open space. Remove the tether, and the same astronaut could drift indefinitely in any direction, including directions from which there is no return. The pre-trained language model is the space station: a large, carefully constructed structure representing enormous accumulated effort. The KL penalty is the tether that keeps the policy from drifting away from it while still letting real exploration of the surrounding space of behaviors.

This analogy also captures an important nuance: the length of the tether matters. A very short tether (large KL coefficient) keeps the astronaut close to the station but severely limits how far they can reach. A very long tether (small KL coefficient) allows wide exploration but risks losing contact entirely. The art of RLHF involves choosing and adapting the tether length appropriately as training progresses, tightening it when the policy shows signs of drifting dangerously and loosening it when progress stalls.

Why does reward hacking happen so readily without constraint? The reward model is a neural network trained on a finite dataset of human preferences. It approximates human judgment but does not perfectly capture it. Given enough freedom, a policy can always find inputs that the reward model scores highly but that humans would rate poorly. This is not a bug in the reward model; it is an inevitable consequence of learning from finite data. The policy is in effect performing adversarial search against the reward model, finding the holes in its coverage. The KL penalty limits how far the policy can search, keeping it in regions where the reward model is well-calibrated because those regions resemble the training data distribution.

The KL divergence penalty also serves a second, subtler purpose beyond preventing reward hacking. The pre-trained model is extraordinarily capable in ways that the reward model does not measure. It can write coherent multi-paragraph arguments, maintain consistent personas, handle edge cases gracefully, and produce text with appropriate register and tone for different audiences. None of these capabilities are directly rewarded during RLHF because they are hard to specify as preferences. By anchoring to the pre-trained model via KL divergence, we preserve these capabilities for free, without needing to explicitly specify them in the reward function. This makes the KL constraint a remarkably efficient mechanism: a single scalar coefficient does the work that would otherwise require carefully engineering dozens of auxiliary reward terms.

This chapter explores why KL divergence is the right tool for this job, how to compute it efficiently for autoregressive models, how to set and adapt the KL coefficient to balance learning against stability, and what happens when the constraint is too tight or too loose. We also examine the worked arithmetic of KL computation, the reward shaping interpretation that connects the penalty to classical RL theory, and the practical considerations that distinguish successful RLHF deployments from unstable ones.