Predicted-KL Control for Policy Updates

Why People Wanted It

The problem is the price of a policy update: a step that looks small in parameter space can still move next-token probabilities substantially. So the real question was never just how to produce an update, but how large a fraction of that proposed update can be applied while keeping the predicted change in the policy within a chosen budget.

The Core Move

Estimate the local average KL caused by the proposed update, then scale the update so its predicted divergence fits the budget. For proposal $D$ and multiplier $\alpha$, the local estimate is $\operatorname{KL}_{\mathrm{avg}}\approx\alpha^2q(D)$, so the constraint is $\alpha^2q(D)\leq\delta$.

Why Token-Score Changes Become a KL Estimate

Here is the part that gets technical: the estimate comes from the first-order change in every token's score at each sampled prediction context. Let $\dot z_v(s)$ be the score's rate of change along the proposed parameter update, and let $p_v(s)=\operatorname{softmax}(z(s)/T)_v$ be the current probability at temperature $T$. The per-context contribution is the probability-weighted variance of those score changes:

$q(D)=\frac{1}{2|\mathcal P|T^2}\sum_{s\in\mathcal P}\operatorname{Var}_{v\sim p(s)}[\dot z_v(s)].$

The variance matters because adding the same shift to every token score changes no softmax probabilities; only differences among token-score changes affect the distribution. The factor $T^{-2}$ captures how the same score movement translates into probability movement at the evaluation temperature. Averaging across the examined contexts gives the unscaled proposal's predicted KL coefficient.

Why Predicted KL Is Quadratic in Update Size

Now for the local derivation: write the KL at a context as a function of a scalar multiplier $a$ on the fixed update. At $a=0$, the two policies match, so KL is zero. Its first derivative is also zero: the probabilities sum to one, leaving no first-order KL change at the point where the distributions coincide. The leading local term is therefore curvature, giving $\operatorname{KL}_{\mathrm{avg}}(\alpha)=\alpha^2q(D)+O(\alpha^3)$ for small $\alpha$, under the stated smoothness conditions on the logits.

In this approximation, doubling the update scale quadruples predicted KL; halving it cuts the prediction by four. That is a local approximation, not a guarantee about the realized divergence for a finite step. When does this quadratic rule provide a useful control knob rather than a false sense of certainty?

How the Budget Sets the Multiplier

The controller compares the predicted divergence with a target budget $\delta>0$. For finite positive $q(D)$, it chooses

$\alpha=\min\left(\alpha_{\max},\sqrt{\frac{\delta}{q(D)}}\right).$

The square root appears because divergence scales with the square of the multiplier: if the proposal predicts four times the budget, the multiplier is at most one half. The cap $\alpha_{\max}$ prevents the rule from enlarging a proposal without limit when its predicted KL is small. A fixed budget holds the target steady; an annealed budget lowers it over training to permit less policy change later.

Minimal Causal Code

This loop shows the rescaling decision after a predicted-KL estimate has been computed. The estimate is treated as a scalar; production code must obtain it from token-score changes on sampled contexts.

import math

# D is the proposed parameter update; q_D is its predicted-KL coefficient.
# The loop makes the scale decision visible rather than applying an optimizer.
delta = 1e-4
alpha_max = 1.0
q_D = estimate_predicted_kl_coefficient(D, sampled_contexts)

if q_D > 0.0:
    alpha = min(alpha_max, math.sqrt(delta / q_D))
else:
    alpha = alpha_max

for parameter, update in zip(parameters, D):
    parameter.data.add_(update, alpha=alpha)

What You Gain, and What You Lose

You gain a direct, interpretable control over the locally predicted average change in next-token probabilities, using a scale that preserves the proposed update's direction. You also get a natural way to tighten the target over training by lowering the budget.

You lose certainty beyond the local approximation: realized KL can differ from predicted KL, especially when the update is not small. The estimate is based on selected response contexts, so it summarizes those examined predictions rather than every possible context. And a strict budget can shrink updates, trading policy movement for tighter predicted control.

Quiz

If the proposed update has predicted-KL coefficient q(D), what happens to the local predicted KL when the update multiplier is halved?

It becomes zero for any nonzero proposal.

It is unchanged.

It is halved.

It is reduced to one quarter.

$ check