KL Divergence as a Regularizer

The "Why": A Leash for Reward Hacking

The core challenge in RLHF is reward hacking. A reward model, trained on a finite set of human preferences, is never perfect. If we optimize our policy solely to maximize the score from this reward model, it can discover strange, unhelpful, or repetitive outputs that "trick" the model into giving a high score.

For example, a policy might learn that generating "I am very helpful and I can assist you" over and over again gets a high reward, even though it's clearly not a good response.

To prevent this, we need to constrain the policy. The key insight is that our initial model (typically the one that has undergone supervised fine-tuning, or SFT) is already very good at producing coherent, grammatical, and sensible text. It has learned this from a massive pre-training dataset. We don't want to lose this foundational competence.

The KL divergence penalty acts as a "mathematical rubber band" or a "leash" that tethers our learning policy (${\pi }_{\theta }$) to this trusted reference policy (${\pi }_{\text{ref}}$). It allows the policy to improve based on the reward signal, but it pulls it back if it strays too far into "weird" territory that the reference model would find unlikely.

The "What": Measuring Divergence with KL

Kullback-Leibler (KL) divergence is a function from information theory that measures how one probability distribution differs from a second, reference probability distribution.

For two discrete probability distributions $P$ and $Q$, the KL divergence from $Q$ to $P$ is defined as:

In our context:

• $P$ is our learning policy, ${\pi }_{\theta }$.
• $Q$ is our frozen reference policy, ${\pi }_{\text{ref}}$.
• The sum is over all possible tokens $x$ in the vocabulary.

This formula has some important properties that make it perfect for our use case.

As the article highlights, the asymmetry of KL divergence is particularly important. It heavily penalizes the policy for assigning high probability to tokens that the reference model considers very unlikely. This discourages the policy from exploring novel but potentially nonsensical outputs.

The "How": The KL-Constrained Objective

While the formal definition of KL divergence involves a sum over the entire vocabulary (which is computationally infeasible), we can leverage a more practical formulation based on expectation. The KL divergence can be rewritten as:

This means we can estimate the KL divergence by sampling outcomes from our policy $P$ and averaging the difference of their log-probabilities.

For an autoregressive language model, this simplifies even further. The KL divergence for a whole sequence is just the sum of the KL divergences at each token. For a single token $y_t$ generated by our policy, its contribution to the KL divergence is simply:

We only need to calculate the log-probabilities from both models for the single token that was actually chosen.

This KL penalty is then subtracted from the reward in the overall RLHF objective function:

Here:

• $r_{\varphi }(x,y)$ is the score from the reward model.
• $\beta$ (beta) is the KL coefficient, a hyperparameter that controls the strength of the penalty. A larger $\beta$ creates a stiffer "rubber band," keeping the policy closer to the reference model.

This process is often called reward shaping, because we are augmenting the final reward with a per-token signal that guides the model's behavior at every step.

Let's watch a video that brings these concepts to life.

This idea of penalizing divergence from a base model is a common theme in modern RL. The following clip discusses its use in the context of the DeepSeek-R1 model, which we'll be studying more later.

Finally, to tie everything together, let's look at a detailed diagram showing where the KL penalty fits into the full PPO pipeline.

Conclusion

In this lesson, we explored the crucial role of the KL divergence penalty in stabilizing RLHF training. While the clipped surrogate objective from our last lesson provides stability within a batch of updates, the KL penalty provides stability across the entire training process by anchoring the policy to a trusted reference.

Key Takeaways:

• Reward Hacking: Unconstrained optimization against a reward model can lead to degenerate outputs.
• The Reference Model: The SFT model serves as a trusted reference, possessing strong foundational language skills that we want to preserve.
• KL Divergence: This measures the "distance" between the learning policy's and the reference policy's probability distributions.
• The Penalty Term: The overall objective is to maximize the reward while minimizing the KL divergence, balanced by the $\beta$ coefficient: $J(\theta )=\text{Reward}-\beta \cdot D_{KL}$.
• Practical Computation: For autoregressive models, the KL penalty is efficiently estimated as a sum of per-token log-probability differences between the policy and reference models.

Preview of the Next Lesson:

Both the clipped objective and the KL penalty are explicit mechanisms added to the PPO algorithm. But what if we could achieve the same regularization effect more directly? In our next lesson, we will introduce Direct Preference Optimization (DPO), an elegant method that bakes the KL constraint directly into its objective function, bypassing the need for an explicit reward model and the complexities of the full PPO pipeline.