KL Divergence Penalty › Why KL Divergence?

Recall from the Reward Hacking chapter that reward models are imperfect approximations of human preferences. They are trained on a finite dataset and will inevitably assign high scores to some outputs that humans would reject. When we optimize a policy using PPO, we are searching over a large space of possible behaviors, and the optimization process will find these reward model errors if given enough freedom. This creates a tradeoff in RLHF: we must improve the model using the reward signal without trusting it unconditionally.

The basic insight is that the pre-trained language model is already a strong prior over sensible, grammatical, informative text. It was trained on trillions of tokens of human-written content, absorbing patterns of coherent discourse, factual knowledge, and linguistic conventions. By constraining our policy to stay close to this reference distribution, we inherit its linguistic competence while still allowing targeted improvements in alignment. Think of the reference model as a trusted expert whose judgment we respect: we want to make adjustments based on new feedback, but we do not want to stray so far that we lose the foundational competence that made the model useful in the first place.

What distinguishes KL divergence from other possible measures of distributional similarity? We could imagine using other metrics: total variation distance, Wasserstein distance, or simple L2 distance between parameter vectors. Each of these alternatives has shortcomings in the RLHF setting. Parameter-space constraints like L2 weight regularization penalize changes in weights, not changes in behavior. A model can change behavior sharply while keeping weights close together (via cancellations across layers) or keep behavior almost identical while changing weights materially (via symmetries and redundancies). We care about behavioral similarity, measured in terms of the outputs the model produces. Wasserstein and total variation distances operate in the right space (distributions over outputs), but they are expensive to compute for high-dimensional distributions and lack convenient decomposition properties for sequential models.

KL divergence measures how different one probability distribution is from another. For language models, this means measuring how much the policy's probability assignments over tokens differ from the reference model's. At each position in a sequence, both models assign probabilities to every token in the vocabulary. KL divergence captures the extent to which these probability distributions disagree, weighted by the probability mass the policy assigns. The key properties that make KL divergence ideal for this application are:

• Non-negativity: $D_{KL}(P \| Q) \geq 0$, with equality only when $P = Q$
• Asymmetry: it penalizes the policy for assigning probability to tokens the reference does not expect
• Decomposability: for autoregressive models, we can compute KL at each token position and sum

The non-negativity property ensures that the penalty term always works in one direction: it can only reduce the objective, never artificially inflate it. This mathematical guarantee means the optimization has a clear interpretation where reward pulls the policy toward better outputs while the KL penalty resists moving too far from the reference. The asymmetry property is particularly important for RLHF because it creates an asymmetric treatment of errors. If the policy assigns high probability to something the reference considers unlikely, the penalty is severe. But if the policy assigns low probability to something the reference considers likely, the penalty is more forgiving. This asymmetry encourages the policy to remain conservative, avoiding novel but potentially problematic outputs. Finally, decomposability allows us to compute and analyze the penalty at the level of individual tokens. This gives fine-grained insight into where and how the policy is diverging.

The decomposability property deserves special attention because it is what makes KL divergence tractable at the scale of language model training. The probability of a complete sequence is exponentially small; computing sequence-level distances directly would require summing over an astronomically large space. But because KL divergence over the full sequence decomposes into a sum of token-level KL divergences, we can compute it efficiently by running both models on the same generated sequence and comparing their log probabilities token by token. This makes the per-sequence KL computation as cheap as two forward passes, regardless of vocabulary size.