KL Regularisation

A reward model is correct near the responses it was trained on and unconstrained elsewhere. Maximise it freely and the policy leaves that region immediately, because the highest-scoring text under a proxy is usually text the proxy has never seen.

$\pi_{\text{ref}}$ is the supervised fine-tuned model — frozen, and the definition of “still a reasonable language model”. $\beta$ sets how far the policy may travel from it.

The closed form

Eq. 1 has an exact solution, which is worth knowing because two other entries depend on it.

The reference distribution reweighted by exponentiated reward. It cannot be sampled from — $Z(x)$ sums over all sequences — which is why PPO exists. But it can be rearranged: solve for $r$ in terms of $\pi^*$, substitute into the Bradley–Terry likelihood, and the partition function cancels. That is DPO, and it is why DPO has a $\beta$ even though it never computes a KL.

Estimating it

The KL cannot be computed exactly over sequences, so it is estimated from samples, and the choice of estimator matters more than it looks.

$k_1$ is unbiased and high-variance, and is negative about half the time — an estimate of a non-negative quantity that is routinely negative is a poor diagnostic even when its mean is right. $k_2$ is low-variance and biased. $k_3$ is unbiased and non-negative, being a Bregman divergence, and costs one extra exponential. It is what current implementations use, and it is Schulman’s blog post rather than any paper that established it.

Two places to apply it

In the reward. Subtract $\beta \log(\pi/\pi_{\text{ref}})$ per token from the reward before computing advantages. The KL then flows through the value function and the advantage estimator like any other reward.

In the loss. Add the estimated KL as a separate term. GRPO does this, which is cleaner to reason about since the penalty does not interact with the critic.

They are not equivalent — the first shapes credit assignment, the second does not — and papers are frequently unclear about which they mean.

import torch
from torch import Tensor


def kl_k3(logp: Tensor, logp_ref: Tensor) -> Tensor:
    """Schulman's k3: unbiased and never negative. logp: [B, T] per-token."""
    ratio = logp_ref - logp                       # log(π_ref / π)
    return ratio.exp() - ratio - 1.0


def adaptive_beta(beta: float, kl: float, target: float = 6.0, horizon: int = 10_000) -> float:
    """Ziegler et al. §2.2: steer the coefficient towards a KL budget."""
    error = torch.clamp(torch.tensor(kl / target - 1.0), -0.2, 0.2).item()
    return beta * (1.0 + error / horizon)

The adaptive controller is the underused half. A fixed $\beta$ means the realised KL wanders as the policy improves — early in training the same coefficient buys a much smaller step than late. Targeting a KL budget instead makes runs comparable across reward models and across seeds, which a fixed coefficient does not.

The reasoning-model work is the first credible argument against. When the reward is verifiable — a maths answer that is right or wrong, a test suite that passes — it is not a proxy, so there is nothing to overoptimise into. Several R1-style recipes set $\beta = 0$ and report that the long chains of thought the model develops are precisely what a KL leash to the SFT policy would have prevented.

The scope of that argument is narrow and worth keeping straight. It applies exactly where the reward is ground truth. For anything scored by a learned model or an LLM judge, the proxy is back and so is the leash.