KL regularization constrains reward-driven policy optimization by penalizing divergence from a frozen reference policy. In RLHF, the reference is usually the instruction-tuned model or an earlier checkpoint; the penalty preserves its useful behavior while allowing improvement toward higher reward. The usual objective maximizes reward minus β times KL(πRL || πref). Larger β keeps the trained policy nearer the reference, while smaller β permits more reward-seeking drift. section 3.1.4 section 15.1.1
Equation 3.11
\[\label{eq:rlhf_opt_eq} {\max_\pi \; \mathbb{E}_{\tau \sim \pi} \left[r_\theta(s_t, a_t)\right] - \beta \mathcal{D}_{\text{KL}}(\pi_{\text{RL}}(\cdot|s_t) \| \pi_{\text{ref}}(\cdot|s_t)).}\]The objective makes the trade-off explicit: reward promotes preferred outputs, and the KL term charges the policy for moving probability away from the reference. section 3.1.4
For generated tokens, reverse KL compares probabilities under the current policy with those of the reference. It becomes especially costly when the new policy puts substantial probability on tokens the reference considers unlikely. Implementations commonly estimate this quantity from samples of the RL policy rather than computing an exact distribution-level divergence. section 15.1 section 15.4
KL regularization is an operational control rather than a guarantee of alignment:
- Treat KL distance as a budget: excessive drift can signal a broken run or correlate with over-optimization, while a stronger penalty can reduce over-optimization at a given distance but require more training steps. section 14.2 section 14.4
- Choose a static coefficient or a dynamic controller. Earlier PPO-style systems adjusted the penalty to target a measured KL, whereas many modern implementations use a static penalty. section 6.3
- Distinguish KL control from PPO clipping. PPO constrains update size, but with one gradient step per batch its clipping may not activate, leaving the reference-distance penalty as the dominant regularizer. section 6.4.2
GRPO loss placement
A standard RLHF implementation can subtract β·KL from the reward before optimization. GRPO can instead add β·KL directly to the policy-gradient loss. For reasoning batches with multiple completions per prompt, KL is typically computed per token, with batch, sequence-length, and completion dimensions. The placement is an implementation choice, not a different purpose: both discourage departure from the reference policy. section 6.3.6
The penalty also has a probabilistic interpretation. Optimizing reward minus reverse KL is equivalent to moving toward a reward-tilted version of the reference distribution: the reference probability is reweighted by exp(reward/β) and normalized. Thus β sets how sharply reward can reshape the starting policy. This reverse-KL objective differs from supervised fine-tuning's forward-KL interpretation, producing different pressures on prior knowledge. section 15.2.2 section 15.1
