KL Divergence Penalty › Summary

The KL divergence penalty is needed for stable RLHF training. It prevents reward hacking by keeping the policy close to a trusted reference model, preserving the language capabilities acquired during pre-training while letting targeted improvements in alignment. Without it, RLHF optimization rapidly degenerates into adversarial exploitation of reward model errors, creating outputs that are high-scoring but meaningless or harmful.

Key takeaways from this chapter:

• KL divergence measures distribution difference: For autoregressive models, it decomposes into a sum of per-token log probability ratios, allowing efficient computation along sampled trajectories. The expectation-based form means we only need to evaluate the ratio for sampled tokens, not for all tokens in the vocabulary.

• The coefficient $\beta$ controls the reward-constraint tradeoff: Low $\beta$ allows rapid learning but risks exploitation of reward model errors. High $\beta$ maintains stability but slows improvement. Typical values range from 0.01 to 0.2, with the right value depending on reward model quality, desired behavior change magnitude, and training stability.

• Adaptive KL control maintains consistent constraints: By adjusting $\beta$ to target a specific KL budget, adaptive methods give stable training dynamics throughout optimization, preventing both early-training instability and late-training stagnation. The InstructGPT-style multiplicative dead-zone controller is a practical and reliable choice.

• Per-token KL reveals what the policy learns: Analyzing where KL accumulates in generated sequences helps diagnose whether the policy is making targeted modifications or developing a fundamentally different generation style. High KL at the beginning of responses suggests changed discourse opening strategies; high KL at specific content words suggests vocabulary shifts; uniformly distributed KL suggests broad stylistic adaptation.

• The reference model quality determines the ceiling: The KL constraint anchors the policy to the SFT model, making SFT quality a necessary determinant of RLHF outcomes. Strong SFT models let more effective alignment; weak SFT models transfer their limitations to the aligned model.

The KL penalty is one approach to constraining policy optimization. As we will see in the following chapters on DPO, there are alternative formulations that reach similar goals through different mechanisms. It offers tradeoffs in implementation complexity, computational efficiency, and training stability. Understanding the KL penalty deeply, as we have done here, gives the basis for understanding why these alternatives work and when to prefer them.