KL Divergence Penalty › Limitations and Practical Considerations

While KL divergence is the standard constraint for RLHF, it has several real limitations worth understanding in depth. These limitations do not invalidate the approach; they are inherent to the problem of aligning language models with human preferences, and understanding them helps practitioners make better decisions about when and how to apply the KL penalty.

The reference model quality bottleneck. The KL penalty anchors the policy to a specific reference, typically the SFT model. This makes the quality of supervised fine-tuning important for RLHF success. If the SFT model has problematic behaviors, the KL penalty resists correcting them because moving away from those behaviors necessarily increases KL. For example, if the SFT model has learned to be verbose (perhaps because the supervised data over-represented detailed responses), the KL penalty will resist training the policy to be more concise, even if the reward model rewards conciseness. Conversely, if the SFT model is highly capable and well-calibrated, the constraint preserves that capability throughout RLHF. This bidirectional constraint is a feature when the reference is good and a bug when it is not, which is why so much effort goes into creating high-quality SFT data.

Distribution distance versus behavioral distance. KL divergence measures similarity between probability distributions, not between actual behaviors. Two policies could have small KL divergence but generate quite different outputs when sampled, especially for long sequences where small per-token differences compound across many conditional distributions. Conversely, policies with large KL might produce similar outputs if the probability differences are concentrated in low-probability regions that rarely get sampled. This mismatch between KL as a mathematical quantity and our actual goal (behavioral similarity) means that the KL penalty sometimes constrains the wrong thing. A policy that has learned to avoid a specific harmful phrase but otherwise generates identically to the reference might show higher KL than a policy that has not changed its harmful behavior at all but has shifted probability mass in obscure parts of the distribution.

Asymmetry creates specific biases. The KL divergence $D_{KL}(\pi_\theta \| \pi_{\text{ref}})$ penalizes the policy more heavily for assigning low probability to tokens the reference model assigns high probability to. This means the policy is discouraged from becoming more confident (more peaked distribution) than the reference, which can be desirable (maintaining diversity) or undesirable (preventing confident correct answers in factual domains). The asymmetry also means that the reverse KL $D_{KL}(\pi_{\text{ref}} \| \pi_\theta)$ would have different properties and might better suit different scenarios. Research on choosing the right divergence direction for alignment remains active.

Computational overhead. Computing KL requires running both the policy and reference model on every training example. For large models, this roughly doubles the forward pass compute. In practice, this means RLHF training is approximately twice as expensive in compute as supervised fine-tuning on the same data, before accounting for the added complexity of running the reward model. Some implementations amortize this by caching reference model log probabilities for a fixed set of prompts, but this requires careful management of the cache and may miss the per-response nature of the KL computation. Memory requirements are also approximately doubled, since both models must reside in memory simultaneously during training.

KL divergence cannot constrain all capabilities simultaneously. The penalty keeps the overall distribution of the policy close to the reference, but it cannot prevent the policy from changing its behavior on specific important subsets of inputs. A policy could decrease KL on common inputs (where it is easy to stay close to the reference) while sharply changing behavior on rare but important inputs (like edge case harmful content). This "gaming" of the aggregate KL constraint is a subtle form of reward hacking that the KL penalty does not fully prevent. Supplementary evaluations on diverse test sets are needed to detect this kind of distributional shift.

Looking ahead, methods like Direct Preference Optimization (DPO), covered in the next chapter, incorporate the KL constraint implicitly in their objective. This eliminates the need for explicit KL computation during training while achieving similar regularization effects. DPO reparameterizes the RLHF objective in terms of policy log ratios, deriving a loss function that can be minimized directly on preference data without the need for a separate reward model or explicit KL penalty term. This is an elegant alternative to the explicit penalty approach we have examined here, trading the flexibility of separate reward modeling for the simplicity and stability of a closed-form training objective.