The Role of the KL Divergence Penalty

In the PPO phase of RLHF, our primary objective is to adjust the language model's policy, $\pi_{\theta}$, to generate responses that maximize the expected reward given by the learned reward model (RM). However, optimizing solely for the RM score presents significant risks. The policy might rapidly shift into regions of the policy space that yield high rewards according to the RM but produce outputs that are nonsensical, repetitive, or stylistically inconsistent with the desired behavior learned during Supervised Fine-Tuning (SFT). This phenomenon can be viewed as the policy "overfitting" to the reward model, potentially exploiting its inaccuracies or limitations (a form of reward hacking), or simply forgetting the fundamental language generation capabilities it previously possessed.

To mitigate this, PPO incorporates a penalty term based on the Kullback-Leibler (KL) divergence. The KL divergence, denoted as $D_{KL}(\pi_{\theta} || \pi_{ref})$, measures the difference between two probability distributions. In the context of RLHF, it quantifies how much the current policy $\pi_{\theta}$ has deviated from a reference policy, $\pi_{ref}$. Typically, this reference policy is the model obtained after the SFT phase, let's call it $\pi_{SFT}$.

The KL divergence between the current policy and the SFT policy for a given state (prompt) $s$ and action (token) $a$ is calculated as:

$$D_{KL}(\pi_{\theta}(\cdot|s) || \pi_{SFT}(\cdot|s)) = \sum_{a} \pi_{\theta}(a|s) \log \frac{\pi_{\theta}(a|s)}{\pi_{SFT}(a|s)}$$

A low KL divergence indicates that the current policy's output distribution is similar to the SFT policy's distribution, while a high KL divergence signifies a substantial deviation.

The core idea is to augment the PPO objective function. Instead of purely maximizing the expected advantage (which relates to the reward), we maximize a modified objective that includes a penalty proportional to the KL divergence:

$$\text{Objective} \approx E_{t} [ \text{Reward}_t ] - \beta D_{KL}(\pi_{\theta}(\cdot|s_t) || \pi_{SFT}(\cdot|s_t))$$

Where:

• $E_{t} [ \text{Reward}_t ]$ represents the expected reward obtained by the policy $\pi_{\theta}$, often approximated using advantage estimates $A_t$ in the actual PPO loss function. This term drives the policy towards generations preferred by the reward model.
• $D_{KL}(\pi_{\theta}(\cdot|s_t) || \pi_{SFT}(\cdot|s_t))$ is the KL divergence between the current policy and the frozen SFT policy for the tokens generated at timestep $t$.
• $\beta$ is a hyperparameter that controls the strength of the KL penalty.

This KL penalty acts as a regularization term. It discourages the policy $\pi_{\theta}$ from deviating too drastically from the SFT policy $\pi_{SFT}$ during optimization. By penalizing large changes in the output probability distribution for each token, it helps ensure that the model retains the general language fluency, knowledge, and stylistic characteristics acquired during the SFT phase, even as it adapts to maximize the reward signal.

Implementation Details

In practice, during the PPO training loop:

1. A batch of prompts is sampled.
2. The current policy $\pi_{\theta}$ generates responses to these prompts.
3. For each generated token, both the current policy $\pi_{\theta}$ and the frozen reference policy $\pi_{SFT}$ compute their respective probability distributions (or logits) over the vocabulary.
4. The KL divergence between these two distributions is calculated for each token.
5. The per-token KL divergences are typically averaged over the generated sequence and the batch.
6. This average KL divergence, scaled by $\beta$, is subtracted from the reward signal (or incorporated into the PPO loss calculation alongside the policy and value loss terms).

The reference policy $\pi_{SFT}$ remains fixed throughout the PPO training process; its weights are not updated. It serves as a stable anchor point representing the behavior learned from the initial supervised dataset.

The KL Coefficient ($\beta$)

The choice of the KL coefficient $\beta$ is important for balancing exploration and stability.

• Low $\beta$: The penalty is weak. The policy $\pi_{\theta}$ has more freedom to move away from $\pi_{SFT}$ to maximize the reward. This can lead to faster optimization towards the reward signal but increases the risk of policy collapse, instability, or generating undesirable text if the policy diverges too far.
• High $\beta$: The penalty is strong. The policy $\pi_{\theta}$ is heavily constrained to stay close to $\pi_{SFT}$. This promotes stability and preserves the SFT model's characteristics but may hinder the policy's ability to sufficiently adapt to the reward signal, potentially resulting in only minor improvements over the SFT model.

Policy updates during PPO. $\pi_{SFT}$ is the starting policy. Low $\beta$ allows larger steps towards high reward regions, risking instability. High $\beta$ restricts steps, maintaining closeness to $\pi_{SFT}$.

Finding an appropriate value for $\beta$ often requires experimentation. Furthermore, adaptive KL controllers are commonly used. These controllers dynamically adjust $\beta$ during training based on the observed KL divergence values in each batch. The goal is to keep the actual $D_{KL}(\pi_{\theta} || \pi_{SFT})$ within a predefined target range (e.g., maintain an average KL of 6 nats). If the observed KL exceeds the target, $\beta$ is increased to strengthen the penalty; if it falls below the target, $\beta$ is decreased to allow more optimization. Libraries like Hugging Face's TRL provide implementations of such adaptive KL controllers.

In summary, the KL divergence penalty is a mechanism in PPO for RLHF. It prevents the language model policy from diverging too far from the initial SFT model while optimizing for human preferences encoded in the reward model. This promotes training stability, preserves desirable characteristics of the base model, and provides a tunable balance between reward maximization and policy constraint.