KL-Divergence Penalty for Policy Regularization in RL Fine-tuning

Welcome to the final lesson in our first module. In our previous lesson, we constructed the full Reinforcement Learning from Human Feedback (RLHF) pipeline, from pre-training a base model to fine-tuning it with PPO using a reward model. This provided the essential context for our study of policy optimization algorithms.

In the standard PPO-based RL fine-tuning loop, the objective is to generate responses that maximize the score from the reward model. However, if we pursue this goal naively, a significant problem arises. This lesson delves into that problem and its elegant solution: the Kullback-Leibler (KL) divergence penalty. You will learn why this penalty is a crucial component for regularizing the language model's policy, preventing it from straying into undesirable behaviors while it learns.

The Problem: Reward Hacking and Distributional Collapse

When we use reinforcement learning to optimize a policy (${\pi }_{\theta }$) to maximize an expected reward (${\mathbb{E}}_{x\sim {\pi }_{\theta }}[r(x)]$), the optimization process is relentlessly focused on that single goal. If the reward function is imperfect—and in RLHF, the reward model is always an imperfect proxy for true human preference—the policy will inevitably find and exploit its flaws. This is known as reward hacking.

In the context of language models, this often manifests as the model generating repetitive, nonsensical, or stylistically bizarre text that happens to receive a high score from the reward model. From a generative modeling perspective, this is a form of distributional collapse. The model, which started as a rich and diverse generator of text, degenerates into a low-entropy distribution that puts all its probability mass on a few high-reward (but low-quality) sequences.

To understand this problem more formally, please read the following section of the blog post "RL with KL penalties is better viewed as Bayesian inference" by Tomek Korbak.

The Solution: The KL-Divergence Penalty

To prevent the policy from "forgetting" how to generate fluent, coherent language, we need to constrain it. We want it to improve based on the reward signal, but not at the expense of its fundamental linguistic competence. The solution is to add a regularization term to the objective that penalizes the policy for straying too far from its original distribution.

The reference point for this is the supervised fine-tuned (SFT) model, which we'll call the reference policy, ${\pi }_{\text{ref}}$. This model is already well-behaved and instruction-following. The KL-regularized RL objective is then:

$$J(\theta )={\mathbb{E}}_{x\sim {\pi }_{\theta }}[r(x)]-\beta \cdot \text{KL}({\pi }_{\theta }\Vert {\pi }_{\text{ref}})$$

Here:

• ${\pi }_{\theta }$ is the policy we are fine-tuning.
• ${\pi }_{\text{ref}}$ is the fixed reference policy (the SFT model).
• $\text{KL}({\pi }_{\theta }\Vert {\pi }_{\text{ref}})$ is the Kullback-Leibler (KL) divergence, which measures the "distance" between the two policies' output distributions.
• $\beta$ is a hyperparameter that controls the strength of this penalty.

This penalty acts like what one researcher aptly calls a "mathematical rubber band," tethering the learning policy ${\pi }_{\theta }$ to the trusted reference policy ${\pi }_{\text{ref}}$.

To get a more intuitive feel for this, let's watch a short video segment.

The following article provides a clear, structured explanation of why KL divergence is well-suited for this task and how it's defined.

As you've just read, a key challenge is that computing the exact KL divergence requires summing over the entire vocabulary, which is computationally infeasible. Instead, we use a simple and effective estimate based on the tokens that were actually sampled during generation.

For a generated sequence $y=(y_1,...,y_T)$, the KL divergence is estimated by summing the per-token log-probability differences:

$$\widehat{\text{KL}}({\pi }_{\theta }\Vert {\pi }_{\text{ref}})\approx \sum _{t=1}^T\left(\log {\pi }_{\theta }(y_t|y_{<t},x)-\log {\pi }_{\text{ref}}(y_t|y_{<t},x)\right)$$

This is simply the difference between the log-probability of the generated token under the current policy and its log-probability under the reference policy. This makes the computation highly efficient.

A Principled View: KL Regularization as Bayesian Inference

Your background in statistical science provides a powerful lens through which to view this process. The KL penalty isn't just an ad-hoc trick; it arises naturally from a Bayesian inference framework.

Consider the alignment problem this way:

• Prior belief ($p(h)$): Our prior belief over what constitutes good text is captured by the SFT model, ${\pi }_{\text{ref}}(x)$.
• Evidence ($p(e|h)$): The reward model provides new evidence. We can model this as a likelihood, for instance $\exp (r(x)/\beta )$, which assigns higher probability to higher-reward sequences.
• Posterior belief ($p(h|e)$): Using Bayes' theorem, our updated belief about the desired text distribution is proportional to the prior times the likelihood: ${\pi }^{\ast }(x)\propto {\pi }_{\text{ref}}(x)\exp (r(x)/\beta )$.

It turns out that maximizing the KL-regularized objective is equivalent to performing variational inference, where we try to find a policy ${\pi }_{\theta }$ that is the best possible approximation of this true posterior ${\pi }^{\ast }(x)$.

This perspective offers a principled justification for the KL penalty. It's the term that emerges when we frame RL fine-tuning as the problem of updating our prior (the SFT model) with new evidence (the reward model) to arrive at a posterior (the aligned model).

The blog post from earlier provides a clear, non-mathematical summary of this viewpoint, with the full derivation available in its appendix for your interest.

KL Penalty in Practice: Reward Shaping

Now that we understand the 'what' and 'why', let's look at the 'how'. How is this KL penalty actually integrated into an algorithm like PPO?

Instead of being a separate term in the final PPO loss function, the KL penalty is typically used to "shape" the reward signal itself. The reward that the PPO algorithm sees is an augmented reward, which is a combination of the score from the reward model and the per-token KL penalty.

The total reward for a sequence is:

$$R_{\text{total}}=r_{\text{RM}}(x,y)-\beta \sum _{t=1}^T\left(\log {\pi }_{\theta }(y_t|...)-\log {\pi }_{\text{ref}}(y_t|...)\right)$$

This total reward is then distributed across the timesteps. The reward model score $r_{\text{RM}}$ is applied only at the final token, while the KL penalty is applied at each token generation step. This provides a dense, per-token feedback signal to the PPO algorithm, encouraging the policy to stay close to the reference at every step.

The following video provides a crystal-clear, step-by-step example of this process.

The effect of the KL penalty, controlled by the $\beta$ hyperparameter, is a trade-off between maximizing the reward and minimizing the divergence from the reference policy.

The KL Penalty in PPO and GRPO

It's important to distinguish between the two forms of regularization we've seen.

1. PPO's Clipped Objective: Constrains how much the policy can change from one update to the next (${\pi }_{\text{new}}$ vs. ${\pi }_{\text{old}}$). This ensures training stability.
2. The KL Penalty: Constrains how far the policy drifts from a fixed reference model (${\pi }_{\theta }$ vs. ${\pi }_{\text{ref}}$). This preserves the model's core capabilities.

This principle of regularizing against a reference model is fundamental to modern RLHF and is carried over directly into GRPO. The GRPO loss function, which we will derive soon, also includes this exact KL penalty term.

Let's watch a final clip that explicitly shows the KL penalty as part of the GRPO loss function, setting the stage for our next module.

Conclusion

In this lesson, we explored the critical role of the KL-divergence penalty in RL fine-tuning. You now have a comprehensive understanding of this essential regularization mechanism.

Here are the key takeaways:

• Purpose: The KL penalty prevents reward hacking and distributional collapse by constraining the learning policy to stay close to a trusted reference model (the SFT model).
• Mechanism: It works by adding a penalty term, $-\beta \cdot \text{KL}({\pi }_{\theta }\Vert {\pi }_{\text{ref}})$, to the RL objective, acting as a "tether" that preserves the model's linguistic fluency and world knowledge.
• Practical Calculation: The KL divergence is efficiently estimated as the sum of per-token log-probability differences between the policy and reference models for a sampled sequence.
• Theoretical Grounding: From a Bayesian perspective, KL-regularized RL is equivalent to performing variational inference to approximate a target posterior distribution, providing a principled foundation for the method.
• Implementation: The penalty is typically implemented as a form of reward shaping, where a per-token negative reward is added to the signal from the reward model.

With this, we conclude our foundational module on policy optimization for LLMs. You have revisited the core mechanics of PPO, placed it within the broader RLHF context, and now understand the crucial role of KL regularization.

In our next module, "The Theory of Group Relative Policy Optimization (GRPO)," we will begin our deep dive into GRPO itself. The first lesson will explain the motivation for GRPO's critic-free approach, directly contrasting it with the PPO architecture we have thoroughly explored.