KL Regularization in RLHF

Definition

KL regularization in RLHF is a penalty on the KL divergence between the policy being trained ${\pi }_{\theta }$ and a frozen reference policy ${\pi }_{\text{ref}}$ (usually the SFT model), added to the reward. It is a leash that keeps the policy from drifting into regions where the learned reward model is unreliable — the mechanism that prevents reward hacking / over-optimization and preserves fluency over the whole training run. It is distinct from PPO’s implicit trust-region KL.

Key math

The per-token reward combines a terminal RM score with a dense KL penalty at every position:
$r_t=\underset{\text{RM score, terminal}}{\underbrace{r_{\psi }(x,y_{1:T})1[t=T]}}-\underset{\text{single-sample KL estimate}}{\underbrace{\beta (\log {\pi }_{\theta }(y_t\mid x,y_{<t})-\log {\pi }_{\text{ref}}(y_t\mid x,y_{<t}))}}$

so the total objective PPO maximizes is
${\mathbb{E}}_{x\sim \mathcal{D},y\sim {\pi }_{\theta }}[r_{\psi }(x,y)-\beta \mathrm{KL}({\pi }_{\theta }(\cdot \mid x)\Vert {\pi }_{\text{ref}}(\cdot \mid x)\left)\right].$

$\beta$ is the reward-vs-fidelity dial: too small → the policy chases the RM into its blind spots (high score, degenerate/dishonest text); too large → the policy is leashed so tightly it barely improves. Often adapted to a target KL (a few nats/response), exactly like PPO-Penalty’s adaptive-$\beta$ controller — but against a frozen anchor rather than ${\pi }_{\text{old}}$. This reference-KL objective is also what DPO solves in closed form.

Why it matters

The reward model is a learned, imperfect proxy for human preference, and its errors are largest exactly where the SFT model rarely ventures — so an unconstrained optimizer marches straight into those blind spots. The $\beta \mathrm{KL}({\pi }_{\theta }\Vert {\pi }_{\text{ref}})$ term makes that drift costly, trading a little reward to stay in-distribution where $r_{\psi }$ is trustworthy. Tuning $\beta$ is the central knob of the RLHF pipeline. Two KLs, don’t conflate: the reference KL bounds total drift (long horizon, fixed anchor); the trust-region KL to ${\pi }_{{\theta }_{\text{old}}}$ bounds step size (short horizon, moving anchor, enforced by clip).

Taught in

• 04-ppo — the RM-score-minus-KL reward, the two-KL distinction, and why the reference KL is load-bearing against reward hacking.

See also

• reward-model — the imperfect proxy the KL guards against
• rlhf — the pipeline this regularizer lives in
• ppo — the optimizer applying it
• trust-region — the other (implicit) KL in RLHF-PPO
• rl-for-llms