Markdown

KL regularization in LLM RL¶

Scope: the KL divergence term that anchors an RL policy to a reference model, the sampled estimators used to compute it without materializing full vocabulary distributions, and the two structurally different places it can be applied. This page covers the k1, k2, and k3 estimators and their bias and variance, why the choice of estimator interacts with where the penalty is applied, the reference policy versus the old policy, and when to drop the penalty entirely. It is the shared component under PPO, GRPO, REINFORCE and RLOO, and DPO.

The numpy block is executed and asserted here by exact enumeration over the vocabulary, so the bias, variance, and gradient results are proven rather than estimated from samples.

What it is¶

A KL penalty measures how far the policy being trained has drifted from a frozen reference policy, usually the checkpoint that existed before RL started, and subtracts that drift from what the run is optimizing. The canonical RLHF objective:

max_theta  E[ r(x, y) ]  -  beta * KL( pi_theta(. | x) || pi_ref(. | x) )

Computing that KL exactly means comparing full next-token distributions over a vocabulary of 100,000-plus entries at every position of every completion, which is memory the run cannot spare. Instead RL implementations use sampled Monte Carlo estimators that need only the log-probability of the token actually emitted. Three appear in practice, all from Schulman's note on approximating KL:

Estimator	Formula (per sampled token, `r = pi_ref / pi_theta`)	Unbiased	Non-negative	Notes
`k1`	`-log r`	Yes	No	High variance, can report a negative divergence
`k2`	`0.5 * (log r)^2`	No	Yes	Low variance, biased
`k3`	`r - log r - 1`	Yes	Yes	Unbiased and non-negative; the GRPO default

Schulman's own comparison, on Gaussians with modest mean shifts, concludes that k3 "has even lower standard deviation than k2 while being unbiased, so it appears to be a strictly better estimator." That holds in the regime it was measured in. Check (4) below constructs the regime where it does not: a discrete distribution in which the reference assigns far more mass than the policy to a token the policy still emits. Both statements are true of their own settings, and the second is the one a drifting asynchronous run walks into.

Two policies are easy to confuse and are not interchangeable. The reference policy pi_ref is frozen for the whole run and exists to stop the model wandering away from its pre-RL behaviour. The old policy pi_old is the policy that generated the current rollout batch and is refreshed every time new rollouts are sampled; it appears in the importance ratio, not in the KL penalty. TRPO constrained the KL against pi_old as a hard trust region; PPO replaced that constraint with clipping and moved the KL against pi_ref into the reward. Both terms can be present at once and they do different jobs.

Two places to apply it¶

In the reward. Compute the per-token penalty once when the rollout is sampled, subtract it from the reward, and treat the result as a constant for the rest of the update. PPO does this, as does TRL's RLOOTrainer. The KL is detached by construction.
In the loss. Add the penalty to the differentiable objective so gradients flow through it and it is recomputed at every inner update. The original GRPO formulation does this.

The distinction is not cosmetic, and the executed block below shows why: the k1 estimator has exactly zero gradient in expectation, so a run that adds a k1 term to its loss and differentiates it has added no regularization at all.

Why use it¶

It is the only thing holding a policy to its pre-RL distribution. Reward maximization has no incentive to preserve general capability, formatting, or safety behaviour that the reward does not measure.
The estimator choice is a real stability lever. k3 is the better estimator when the policy and reference are close and the worse one when they are far apart, by orders of magnitude in variance. Getting this backwards is a plausible cause of a run that destabilizes late rather than early.
KL drift is the best available early-warning signal. It moves before the reward curve does, and it is the quantity that correlates with how much prior capability a run is destroying (continual learning with RL).
It costs a frozen forward pass, not a trained model. Unlike a critic, the reference is inference-only, so at roughly 2 GB per billion parameters in half precision it is the cheapest model in the loop.

When to use it (and when not)¶

Keep a KL penalty for RLHF against a learned reward model. A reward model is a proxy that can be gamed, and the penalty bounds how far the policy can travel while gaming it.
Keep it whenever the run must preserve behaviour the reward does not measure: instruction formats, refusal behaviour, or performance on tasks outside the RL data.
Consider dropping it for large-scale verifiable-reward reasoning runs. The argument is that anchoring to the pre-RL checkpoint is too restrictive when a large behavioural change is the point, and that verifiable rewards are much harder to hack than a reward model. This is now the library default rather than an exotic choice: TRL's GRPOConfig ships beta: float = 0.0, so a GRPO run has no KL penalty unless one is asked for.
Do not use it as a substitute for clipping. A KL penalty toward pi_ref does not bound the size of a single update; that is what PPO-style clipping against pi_old is for. Pages that conflate the two produce runs that are regularized and still unstable.
Do not use k1 in a differentiable loss. Use k3 there, or move the k1 penalty into the reward where it is detached.

Architecture¶

flowchart TB
  ROLL["Rollout from pi_old"] --> LP["Per-token log-probs: pi_theta, pi_ref"]
  LP --> EST{"Estimator"}
  EST -->|"k1 = -log r"| K1["Unbiased, high variance,<br/>zero gradient if differentiated"]
  EST -->|"k3 = r - log r - 1"| K3["Unbiased, non-negative,<br/>usable in a loss"]
  K1 --> INR["Into the REWARD<br/>(detached, computed once)"]
  K3 --> INL["Into the LOSS<br/>(differentiable, recomputed)"]
  INR --> ADV["Advantage / return"]
  INL --> OBJ["Surrogate objective"]
  ADV --> OBJ
  OBJ --> UPD["Policy update"]
  CLIP["Clip vs pi_old"] -.->|"bounds the step, not the drift"| UPD

How to use it¶

Every property in the table above is provable by enumerating the vocabulary rather than sampling from it. The block below does that, including the two results that contradict the usual "k3 is simply better" summary.

# kl_estimators.py -- the k1 and k3 sampled KL estimators, by EXACT enumeration over the
# vocabulary, so bias, variance and gradient claims are proven rather than sampled.
# numpy only.
import numpy as np

def softmax(x):
    e = np.exp(x - x.max())
    return e / e.sum()

def kl(p, q):
    return float((p * np.log(p / q)).sum())

def estimators(p, q):
    """Per-token k1 and k3 values, as a policy sampling from p would see them."""
    k1 = np.log(p / q)                       # log(pi/pi_ref) for the sampled token
    log_r = np.log(q / p)                    # log(pi_ref/pi)
    k3 = np.exp(log_r) - log_r - 1
    return k1, k3

def moments(p, q):
    k1, k3 = estimators(p, q)
    true = kl(p, q)
    return true, float(p @ k1), float(p @ k3), float(p @ (k1 - true) ** 2), float(p @ (k3 - true) ** 2)

theta = np.array([0.4, -0.2, 1.1, 0.0])      # current policy logits
ref   = np.array([0.1,  0.0, 0.9, 0.2])      # reference policy logits
p, q = softmax(theta), softmax(ref)

# (1) BOTH ARE UNBIASED. Each matches the true KL in expectation to machine precision.
true, e1, e3, v1, v3 = moments(p, q)
assert np.isclose(e1, true, atol=1e-14) and np.isclose(e3, true, atol=1e-14)

# (2) k3 IS NON-NEGATIVE PER SAMPLE; k1 is not. A single k1 sample can report a negative
#     divergence, which the true KL never is.
k1, k3 = estimators(p, q)
assert (k3 >= 0).all()
assert (k1 < 0).any()

# (3) CLOSE POLICIES: k3 has far lower variance, which is why GRPO's objective uses it.
assert v3 < v1 / 100

# (4) ADVERSARIAL: k3's variance is not universally lower. When the reference assigns much
#     more mass than the policy to a token the policy still emits, the exp() in k3 dominates
#     and its variance explodes past k1's. This is the regime a drifting rollout policy
#     enters, and the reason a large asynchronous run may deliberately pick k1 over k3.
V = 8
rows = []
for eps in (1e-1, 1e-3, 1e-5):
    pp = np.ones(V); pp[0] = eps * V; pp /= pp.sum()
    qq = np.ones(V); qq[0] = 4.0 * V; qq /= qq.sum()
    t, _, _, a1, a3 = moments(pp, qq)
    rows.append((eps, round(t, 4), round(a1, 4), round(a3, 1)))
    assert a3 > a1                          # k3 is the WORSE estimator here
assert rows[-1][3] > 1e4 * rows[-1][2]

# (5) THE GRADIENT TRAP. Putting a KL term in a differentiable loss only works for k3.
#     Differentiating the k1 term through the sampled log-probs gives EXACTLY zero in
#     expectation, because E_pi[grad log pi] = 0. A run that "added a k1 KL to the loss"
#     added nothing at all; k1 belongs in the reward, where it is detached.
def score(a):                                # grad_theta log pi(a) for a softmax policy
    g = -p.copy(); g[a] += 1.0
    return g

grad_k1_in_loss = sum(p[a] * score(a) for a in range(4))                    # d/dtheta of k1
grad_k3_in_loss = sum(p[a] * (1 - q[a] / p[a]) * score(a) for a in range(4))  # d/dtheta of k3
assert np.allclose(grad_k1_in_loss, 0.0, atol=1e-15)
assert np.linalg.norm(grad_k3_in_loss) > 1e-3

# ...and k3's loss gradient does point where the true KL gradient points. Finite differences
# on the exact KL are the independent reference.
def kl_of(th):
    return kl(softmax(th), q)
true_grad = np.array([(kl_of(theta + 1e-6 * np.eye(4)[i]) - kl_of(theta - 1e-6 * np.eye(4)[i])) / 2e-6
                      for i in range(4)])
cos = float(grad_k3_in_loss @ true_grad
            / (np.linalg.norm(grad_k3_in_loss) * np.linalg.norm(true_grad)))
assert cos > 0.99

# (6) BOUNDARY: k3 exponentiates a log-ratio, so it overflows long before k1 does. In fp32
#     a log-ratio of 100 is already infinite; k1 handles it as an ordinary number.
lr = np.float32(100.0)
with np.errstate(over="ignore"):
    assert not np.isfinite(np.exp(lr) - lr - 1)
assert np.isfinite(np.float32(-100.0))

print("true KL:", round(true, 6), "| E[k1]:", round(e1, 6), "| E[k3]:", round(e3, 6))
print("variance close policies  -> k1:", round(v1, 6), " k3:", round(v3, 6))
print("variance drifting policy (eps, KL, var_k1, var_k3):")
for r in rows:
    print("   ", r)
print("grad of k1-in-loss:", np.round(grad_k1_in_loss, 18))
print("grad of k3-in-loss:", np.round(grad_k3_in_loss, 4), "| cos to true grad KL:", round(cos, 4))

Executed output:

true KL: 0.021007 | E[k1]: 0.021007 | E[k3]: 0.021007
variance close policies  -> k1: 0.040058  k3: 0.000279
variance drifting policy (eps, KL, var_k1, var_k3):
    (0.1, 1.2311, 1.2525, 1.6)
    (0.001, 1.707, 0.0784, 575.3)
    (1e-05, 1.7175, 0.0019, 58887.2)
grad of k1-in-loss: [-7.0e-18 -1.7e-17 -1.4e-17  0.0e+00]
grad of k3-in-loss: [ 0.0452 -0.0432  0.0507 -0.0527] | cos to true grad KL: 0.9986

Three results are load-bearing. Check (3) is the textbook case: with close policies, k3's variance is 0.000279 against k1's 0.040058, a factor of 143. Check (4) reverses it: once the reference assigns much more probability than the policy to a token the policy still emits, k3's variance reaches 58,887 while k1's has fallen to 0.0019, a gap of seven orders of magnitude in the other direction. This is the regime a long asynchronous run drifts into, and it is why a production recipe can rationally prefer k1 despite k3 being the better estimator on paper. Check (5) is the implementation trap: differentiating k1 inside a loss gives a gradient of order 1e-17, which is machine zero, because the expected score function is zero. k3 differentiated in a loss gives a real gradient that agrees with the true KL gradient to a cosine of 0.9986.

How to develop with it¶

The two placements are a few lines each and are not interchangeable.

# Reference template (needs torch). Both KL placements, side by side. The maths each
# implements is the k1 / k3 arithmetic validated above.
import torch

# per-token log-probs of the SAMPLED tokens, shapes (B, L)
# per_token_logps        : current policy pi_theta
# old_per_token_logps    : rollout policy pi_old   (fixed for the batch)
# reference_per_token_logps : frozen reference pi_ref

# --- Placement A: into the REWARD (PPO, REINFORCE, TRL RLOO). Detached by construction,
#     computed ONCE from the rollout log-probs and held fixed across inner updates.
per_token_kl_k1 = (old_per_token_logps - reference_per_token_logps) * completion_mask
per_token_rewards = -kl_beta * per_token_kl_k1                       # (B, L)
per_token_rewards[torch.arange(B), last_token_idx] += outcome_reward  # outcome at the end

# --- Placement B: into the LOSS (original GRPO). Differentiable, recomputed every inner
#     update. MUST be k3: k1 here contributes no gradient (validated above).
log_ratio = reference_per_token_logps - per_token_logps
per_token_kl_k3 = torch.exp(log_ratio) - log_ratio - 1               # (B, L)
loss = policy_loss + kl_beta * per_token_kl_k3
loss = ((loss * completion_mask).sum(dim=-1) / completion_mask.sum(dim=-1)).mean()

Implementation rules that follow from the validation:

Work in log space and exponentiate late. torch.exp(log_ratio) on a log-ratio the policy has driven to 100 is infinite in fp32. Clamp the log-ratio before the exponential, and treat a clamp that fires often as a signal that the run has drifted, not as a fix.
Do not detach inside placement B. Detaching the policy log-probs in the k3 expression turns it back into a constant and silently removes the regularization, which is the same defect as using k1 there.
Do detach in placement A. The reward is a constant in the policy-gradient derivation. A gradient path from the reward back into the policy optimizes a different objective.
Aggregate the per-token estimates the same way as the policy loss. Summing KL per sequence while averaging the policy loss per token makes the effective beta a function of response length.

How to maintain it¶

Log the KL separately from the reward, always. When the penalty is folded into the reward, a rising raw reward and a rising KL cancel into a flat curve that hides both. This is the single most common way a KL-regularized run goes wrong unobserved.
Log the estimator's variance, not just its mean. A k3 mean that looks stable while its variance climbs is check (4) starting to happen. That is the point to consider switching to k1 or reducing drift, not after the run diverges.
Watch the clamp rate on the log-ratio. A rising fraction of clamped tokens says the policy and reference have separated on tokens the policy still samples.
Distinguish drift from instability. KL against pi_ref growing steadily is drift and may be intended. The importance ratio against pi_old moving outside the clip band is instability and is not.

How to run it in production¶

Pick the estimator against the expected drift, not by default. k3 for runs that stay near the reference, which is most RLHF. Consider k1 for long asynchronous runs where the rollout and training policies separate; the Composer 2 report states it used the k1 estimator, -log r, because k3's variance increased as the rollout and training policies diverged, which is exactly check (4).
Set beta against the reward's hackability. A learned reward model needs a penalty; a deterministic verifier mostly does not. Many large-scale reasoning recipes run beta = 0 and rely on the verifier plus clipping instead.
Budget the reference forward pass. It is inference-only, so it adds roughly 2 GB per billion parameters plus a forward pass per rollout, not a second trained model. With beta = 0 the reference can be dropped entirely, which removes that pass from the step.
Refresh pi_old on schedule and never confuse it with pi_ref. pi_old moves every time rollouts are sampled; pi_ref does not move at all. A codebase that reuses one variable for both is one refactor away from a silent bug.

Failure modes¶

k1 in a differentiable loss. Contributes exactly zero gradient in expectation, per check (5). The run is unregularized and the KL number in the logs still looks reasonable.
k3 with a drifted policy. Variance explodes, gradients become spiky, and the run destabilizes at a point unrelated to any change in the reward.
Overflow in exp(log_ratio). Produces NaN losses that propagate to every parameter in one step.
KL folded into the reward and never logged separately. Masks both reward progress and drift.
Using the KL penalty to control step size. It bounds drift from the reference, not the size of an update. Reach for clipping instead.
Aggregation mismatch. Per-sequence KL against per-token policy loss makes the regularization strength length-dependent.
Anchoring too hard on a reasoning run. A large beta against the pre-RL checkpoint prevents the behavioural change the run exists to produce.

References¶

Schulman, Approximating KL Divergence (the k1, k2, and k3 estimators): http://joschu.net/blog/kl-approx.html
Schulman et al., Trust Region Policy Optimization (the hard KL constraint): https://arxiv.org/abs/1502.05477
Schulman et al., Proximal Policy Optimization Algorithms: https://arxiv.org/abs/1707.06347
Shao et al., DeepSeekMath (GRPO, k3 in the objective): https://arxiv.org/abs/2402.03300
Cursor Research, Composer 2 Technical Report (k1 chosen over k3 under policy divergence): https://cursor.com/resources/Composer2.pdf
TRL GRPO Trainer documentation (the beta coefficient): https://huggingface.co/docs/trl/grpo_trainer
TRL RLOO Trainer documentation (KL folded into the reward): https://huggingface.co/docs/trl/rloo_trainer
Cameron R. Wolfe, Reinforcement Learning for LLMs: The Complete Guide: https://cameronrwolfe.substack.com/p/llm-rl