Computing the Gradient of the KL Penalty
RL post-training methods like GRPO and PPO often add a KL penalty against a
reference policy to keep the trained policy from drifting too far. The penalty
is usually estimated with one of Schulman's k1 or k3
estimators, but differentiating those estimators directly turns out to be
subtler than it looks. These are my notes working through why, and through what
the correct gradient of a sequence-level KL penalty actually is.
The k1 gradient does not depend on the reference policy
Say the loss term is the naive Monte Carlo average of \(\log(p(x)/q(x))\),
\(x \sim p\) — the k1 estimator of the reverse KL
\(\mathbb{E}_{x \sim p}[\log(p(x)/q(x))]\). Differentiating it directly gives:
The gradient does not depend on \(\log q(x)\) at all, so the KL penalty fails to pull \(p\) toward \(q\). As a result, the k1 estimator cannot be used as a direct loss term — only as an extra term folded into the GRPO/PPO advantage.
The k1 gradient estimates a zero-valued expectation
Worse, the expectation of \(\nabla_{\theta} \log p\) is:
$$ \mathbb{E}_{x \sim p}[\nabla_{\theta} \log p] \\ = \sum_{x} p \cdot \nabla_{\theta} \log p \\ = \sum_{x} p \cdot \frac{1}{p} \nabla_{\theta} p \\ = \sum_{x} \nabla_{\theta} p \\ = \nabla_{\theta} \sum_{x} p \\ = \nabla_{\theta} 1 \\ = 0 $$So the gradient carries no real signal — it just boosts the logprob of whichever token was sampled, weighted by the current probability distribution. The update can still nudge the distribution around, but there is no external signal telling it which way to go.
The k3 gradient depends on the inverse policy ratio
The k3 estimator, \(q/p - 1 + \log(p/q)\), is often preferred
because its value is always non-negative. Its gradient is:
So the coefficient of the gradient, \(1 - q/p\), depends on the inverse ratio \(q/p\).
The k3 gradient estimates the forward KL, not the reverse KL
Taking the expectation of that gradient — i.e.\ the value it is a Monte-Carlo estimator of — gives:
$$ \mathbb{E}_{x \sim p}\left[\left(1 - \frac{q}{p}\right) \nabla_{\theta} \log p\right] \\ = \mathbb{E}_{x \sim p}[\nabla_{\theta} \log p] - \mathbb{E}_{x \sim p}\left[\frac{q}{p} \nabla_{\theta} \log p\right] \\ = \left(\sum_{x} p \cdot \frac{1}{p} \nabla_{\theta} p\right) - \mathbb{E}_{x \sim p}\left[\frac{q}{p} \nabla_{\theta} \log p\right] \\ = \left(\sum_{x} \nabla_{\theta} p\right) - \mathbb{E}_{x \sim p}\left[\frac{q}{p} \nabla_{\theta} \log p\right] \\ = \nabla_{\theta} \sum_{x} p - \mathbb{E}_{x \sim p}\left[\frac{q}{p} \nabla_{\theta} \log p\right] \\ = \nabla_{\theta} 1 - \mathbb{E}_{x \sim p}\left[\frac{q}{p} \nabla_{\theta} \log p\right] \\ = 0 - \mathbb{E}_{x \sim p}\left[\frac{q}{p} \nabla_{\theta} \log p\right] \\ = -\mathbb{E}_{x \sim p}\left[\frac{q}{p} \nabla_{\theta} \log p\right] \\ = -\mathbb{E}_{x \sim q}[\nabla_{\theta} \log p] $$But \(-\mathbb{E}_{x \sim q}[\nabla_{\theta} \log p]\) is exactly what the forward KL wants to optimize:
$$ \nabla_{\theta} D_{\text{KL}}(q \Vert p) \\ = \nabla_{\theta} \sum_{x} q \cdot (\log q - \log p) \\ = \nabla_{\theta} \sum_{x} q \cdot \log q - \nabla_{\theta} \sum_{x} q \cdot \log p \\ = 0 - \nabla_{\theta} \sum_{x} q \cdot \log p \\ = -\sum_{x} q \, \nabla_{\theta} \log p \\ = -\mathbb{E}_{x \sim q}[\nabla_{\theta} \log p] $$So even though the value of the k3 estimator is an unbiased estimator of the reverse KL, its gradient is not an unbiased estimator of the reverse-KL gradient — it is an unbiased estimator of the forward-KL gradient instead.
Sequence-level KL is not the sum of token-level KL gradients
The value of a sequence \(y_{1:T}\)'s KL does equal the sum of its tokens' KLs:
$$ D_{\text{KL}}(p(y_{1:T}) \Vert q(y_{1:T})) = \mathbb{E}_{y_{1:T} \sim p}\left[\log \frac{p(y_{1:T})}{q(y_{1:T})}\right] = \mathbb{E}_{y_{1:T} \sim p}\left[\sum_{t=0}^{T} \log \frac{p(y_t \mid x, y_{1:t-1})}{q(y_t \mid x, y_{1:t-1})}\right] $$Training frameworks such as VeRL implement the sequence KL as the right-hand side above, which is correct — but then differentiate each token-level log-ratio independently, which is not, because the implicit \(p(y_{1:T})\) inside the expectation requires the product rule:
$$ \nabla_{\theta} \ \mathbb{E}_{y_{1:T} \sim p}\left[\log \frac{p(y_{1:T})}{q(y_{1:T})}\right] \neq \mathbb{E}_{y_{1:T} \sim p}\left[\sum_{t=0}^{T} \nabla_{\theta} \log \frac{p(y_t \mid x, y_{1:t-1})}{q(y_t \mid x, y_{1:t-1})}\right] $$The correct token-KL estimator
Working the product rule through properly, the correct token-KL gradient is:
$$ \nabla_{\theta} \ \mathbb{E}_{x \sim p}\left[\log \frac{p(x)}{q(x)}\right] \\ = \nabla_{\theta} \sum_{x} p(x) \log \frac{p(x)}{q(x)} \\ = \sum_{x} (\nabla_{\theta} p(x)) \log \frac{p(x)}{q(x)} + \sum_{x} p(x) \left(\nabla_{\theta} \log \frac{p(x)}{q(x)}\right) \\ = \sum_{x} (\nabla_{\theta} p(x)) \log \frac{p(x)}{q(x)} + \sum_{x} p(x) (\nabla_{\theta} \log p(x)) \\ = \sum_{x} (\nabla_{\theta} p(x)) \log \frac{p(x)}{q(x)} + \sum_{x} p(x) \frac{1}{p(x)} \nabla_{\theta} p(x) \\ = \sum_{x} (\nabla_{\theta} p(x)) \log \frac{p(x)}{q(x)} + \sum_{x} \nabla_{\theta} p(x) \\ = \sum_{x} (\nabla_{\theta} p(x)) \log \frac{p(x)}{q(x)} + \nabla_{\theta} \sum_{x} p(x) \\ = \sum_{x} (\nabla_{\theta} p(x)) \log \frac{p(x)}{q(x)} + \nabla_{\theta} 1 \\ = \sum_{x} (\nabla_{\theta} p(x)) \log \frac{p(x)}{q(x)} + 0 \\ = \sum_{x} \left(p(x) \nabla_{\theta} \log p(x)\right) \log \frac{p(x)}{q(x)} \\ = \mathbb{E}_{x \sim p}\left[\log \frac{p(x)}{q(x)} \, \nabla_{\theta} \log p(x)\right] $$So the correct estimator is simply \(\log(p(x)/q(x)) \, \nabla_{\theta} \log p(x)\), which checks out as a policy-gradient loss with \(\log(p(x)/q(x))\) playing the role of the advantage.
The correct sequence-KL estimator
The same product-rule argument at the sequence level gives:
$$ \nabla_{\theta} \ \mathbb{E}_{y_{1:T} \sim p}\left[\log \frac{p(y_{1:T})}{q(y_{1:T})}\right] \\ = \nabla_{\theta} \sum_{y_{1:T}} p(y_{1:T}) \log \frac{p(y_{1:T})}{q(y_{1:T})} \\ = \sum_{y_{1:T}} (\nabla_{\theta} p(y_{1:T})) \log \frac{p(y_{1:T})}{q(y_{1:T})} + \sum_{y_{1:T}} p(y_{1:T}) \, \nabla_{\theta} \left(\log \frac{p(y_{1:T})}{q(y_{1:T})}\right) \\ = \sum_{y_{1:T}} (\nabla_{\theta} p(y_{1:T})) \log \frac{p(y_{1:T})}{q(y_{1:T})} + \sum_{y_{1:T}} p(y_{1:T}) \sum_{t=1}^{T} \nabla_{\theta} \left(\log \frac{p(y_t \mid y_{1:t-1})}{q(y_t \mid y_{1:t-1})}\right) \\ = \sum_{y_{1:T}} (\nabla_{\theta} p(y_{1:T})) \log \frac{p(y_{1:T})}{q(y_{1:T})} + \sum_{y_{1:T}} p(y_{1:T}) \sum_{t=1}^{T} \nabla_{\theta} \log p(y_t \mid y_{1:t-1}) \\ = \sum_{y_{1:T}} (\nabla_{\theta} p(y_{1:T})) \log \frac{p(y_{1:T})}{q(y_{1:T})} + \sum_{y_{1:T}} p(y_{1:T}) \, \nabla_{\theta} \sum_{t=1}^{T} \log p(y_t \mid y_{1:t-1}) \\ = \sum_{y_{1:T}} (\nabla_{\theta} p(y_{1:T})) \log \frac{p(y_{1:T})}{q(y_{1:T})} + \sum_{y_{1:T}} p(y_{1:T}) \, \nabla_{\theta} \log p(y_{1:T}) \\ = \sum_{y_{1:T}} (\nabla_{\theta} p(y_{1:T})) \log \frac{p(y_{1:T})}{q(y_{1:T})} + \sum_{y_{1:T}} p(y_{1:T}) \frac{1}{p(y_{1:T})} \nabla_{\theta} p(y_{1:T}) \\ = \sum_{y_{1:T}} (\nabla_{\theta} p(y_{1:T})) \log \frac{p(y_{1:T})}{q(y_{1:T})} + \sum_{y_{1:T}} \nabla_{\theta} p(y_{1:T}) \\ = \sum_{y_{1:T}} (\nabla_{\theta} p(y_{1:T})) \log \frac{p(y_{1:T})}{q(y_{1:T})} + \nabla_{\theta} \sum_{y_{1:T}} p(y_{1:T}) \\ = \sum_{y_{1:T}} (\nabla_{\theta} p(y_{1:T})) \log \frac{p(y_{1:T})}{q(y_{1:T})} + \nabla_{\theta} 1 \\ = \sum_{y_{1:T}} (\nabla_{\theta} p(y_{1:T})) \log \frac{p(y_{1:T})}{q(y_{1:T})} + 0 \\ = \sum_{y_{1:T}} p(y_{1:T}) \left(\nabla_{\theta} \log p(y_{1:T})\right) \log \frac{p(y_{1:T})}{q(y_{1:T})} \\ = \mathbb{E}_{y_{1:T} \sim p}\left[\log \frac{p(y_{1:T})}{q(y_{1:T})} \left(\nabla_{\theta} \log p(y_{1:T})\right)\right] \\ = \mathbb{E}_{y_{1:T} \sim p}\left[\sum_{s=1}^{T} \log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})} \left(\nabla_{\theta} \sum_{t=1}^{T} \log p(y_t \mid y_{1:t-1})\right)\right] \\ = \mathbb{E}_{y_{1:T} \sim p}\left[\sum_{s=1}^{T} \log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})} \left(\sum_{t=1}^{T} \nabla_{\theta} \log p(y_t \mid y_{1:t-1})\right)\right] $$So the correct estimator still only differentiates each token's logprob — it just weights and sums those logprobs the correct way, coupling every token's log-ratio to every other token's gradient.
Simplifying the sequence-KL estimator
That coupling means the correct estimator is a double sum over token pairs \((s, t)\). Expanding it:
$$ \mathbb{E}_{y_{1:T} \sim p}\left[\sum_{s=1}^{T} \log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})} \left(\sum_{t=1}^{T} \nabla_{\theta} \log p(y_t \mid y_{1:t-1})\right)\right] \\ = \mathbb{E}_{y_{1:T} \sim p}\left[\sum_{s=1}^{T} \sum_{t=1}^{T} \log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})} \, \nabla_{\theta} \log p(y_t \mid y_{1:t-1})\right] \\ = \mathbb{E}_{y_{1:T} \sim p}\left[\sum_{t=1}^{T} \sum_{s=1}^{t-1} \log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})} \, \nabla_{\theta} \log p(y_t \mid y_{1:t-1})\right] + \mathbb{E}_{y_{1:T} \sim p}\left[\sum_{t=1}^{T} \sum_{s=t}^{T} \log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})} \, \nabla_{\theta} \log p(y_t \mid y_{1:t-1})\right] $$For a single term in either double sum:
$$ \mathbb{E}_{y_{1:T} \sim p}\left[\log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})} \, \nabla_{\theta} \log p(y_t \mid y_{1:t-1})\right] \\ = \sum_{y_{1:T}} p(y_{1:T}) \log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})} \, \nabla_{\theta} \log p(y_t \mid y_{1:t-1}) \\ = \sum_{y_{1:t-1}} p(y_{1:t-1}) \sum_{y_t} p(y_t \mid y_{1:t-1}) \sum_{y_{t+1:T}} p(y_{t+1:T} \mid y_{1:t}) \, \log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})} \, \nabla_{\theta} \log p(y_t \mid y_{1:t-1}) $$For the first term, where \(s \lt t\):
$$ \sum_{y_{1:t-1}} p(y_{1:t-1}) \sum_{y_t} p(y_t \mid y_{1:t-1}) \sum_{y_{t+1:T}} p(y_{t+1:T} \mid y_{1:t}) \, \log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})} \, \nabla_{\theta} \log p(y_t \mid y_{1:t-1}) \\ = \sum_{y_{1:t-1}} p(y_{1:t-1}) \underbrace{\log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})}}_{\text{same for every } y_{t:T} \text{ if } s \lt t} \sum_{y_t} p(y_t \mid y_{1:t-1}) \sum_{y_{t+1:T}} p(y_{t+1:T} \mid y_{1:t}) \, \nabla_{\theta} \log p(y_t \mid y_{1:t-1}) \\ = \sum_{y_{1:t-1}} p(y_{1:t-1}) \log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})} \sum_{y_t} p(y_t \mid y_{1:t-1}) \underbrace{\nabla_{\theta} \log p(y_t \mid y_{1:t-1})}_{\text{same for every } y_{t+1:T}} \sum_{y_{t+1:T}} p(y_{t+1:T} \mid y_{1:t}) \\ = \sum_{y_{1:t-1}} p(y_{1:t-1}) \log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})} \sum_{y_t} p(y_t \mid y_{1:t-1}) \, \nabla_{\theta} \log p(y_t \mid y_{1:t-1}) \cdot 1 \\ = \sum_{y_{1:t-1}} p(y_{1:t-1}) \log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})} \, \mathbb{E}_{y_t \sim p}[\nabla_{\theta} \log p(y_t \mid y_{1:t-1})] \\ = \sum_{y_{1:t-1}} p(y_{1:t-1}) \log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})} \cdot 0 \\ = 0 $$The \(s \lt t\) term vanishes, using the same zero-expectation identity from the k1 estimator above. We can't do the same for the \(s \geq t\) term, because there's no way to pull the log-ratio out of the inner expectation to get a 0-valued term. So the estimator reduces to just the \(s \geq t\) half:
$$ \mathbb{E}_{y_{1:T} \sim p}\left[\sum_{t=1}^{T} \sum_{s=t}^{T} \log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})} \underbrace{\nabla_{\theta} \log p(y_t \mid y_{1:t-1})}_{\text{same for every } s \geq t}\right] = \mathbb{E}_{y_{1:T} \sim p}\left[\sum_{t=1}^{T} \nabla_{\theta} \log p(y_t \mid y_{1:t-1}) \underbrace{\left(\sum_{s=t}^{T} \log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})}\right)}_{R_t}\right] $$Implementing the sequence-KL estimator
Given a complete sequence \(y_{1:T}\) with logprobs under both \(p\) and \(q\), this comes down to two steps:
- Compute \(R_t\) for every token. Each position \(t\) needs
\(\log(p_t/q_t) + R_{t+1}\), i.e.\ a reversed cumulative sum:
log_ratio.flip(-1).cumsum(-1).flip(-1). - Compute the loss for each token in parallel, the regular policy-gradient way, using \(R_t\) in place of the usual advantage.
Sequence-KL estimator with importance sampling
Using the importance-sampling ratio, the same expectation over \(p\) can be rewritten as an expectation over the rollout policy \(\mu\) — not the reference policy \(q\), which stays exactly where it is, inside the log-ratio:
$$ \mathbb{E}_{y_{1:T} \sim p}\left[\sum_{t=1}^{T} \nabla_{\theta} \log p(y_t \mid y_{1:t-1}) \left(\sum_{s=t}^{T} \log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})}\right)\right] \\ = \mathbb{E}_{y_{1:T} \sim \mu}\left[\frac{p(y_{1:T})}{\mu(y_{1:T})} \sum_{t=1}^{T} \nabla_{\theta} \log p(y_t \mid y_{1:t-1}) \left(\sum_{s=t}^{T} \log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})}\right)\right] $$Take a single \((s, t)\) term out of that double sum, with \(s \geq t\) as established above. The trick is to split the rollout \(y_{1:T}\) at position \(s\): everything the term depends on — the importance weight up to \(s\), the gradient at \(t \leq s\), and the log-ratio at \(s\) — is fixed once \(y_{1:s}\) is fixed, so the sum over the remaining tokens \(y_{s+1:T}\) integrates away to 1:
$$ \mathbb{E}_{y_{1:T} \sim \mu}\left[\frac{p(y_{1:T})}{\mu(y_{1:T})} \left(\nabla_{\theta} \log p(y_t \mid y_{1:t-1})\right) \left(\log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})}\right)\right] \\ = \sum_{y_{1:s}} \sum_{y_{s+1:T}} \underbrace{\mu(y_{1:s})}_{\text{fixed by } y_{1:s}} \, \mu(y_{s+1:T} \mid y_{1:s}) \left(\underbrace{\frac{p(y_{1:s})}{\mu(y_{1:s})}}_{\text{fixed by } y_{1:s}} \frac{p(y_{s+1:T} \mid y_{1:s})}{\mu(y_{s+1:T} \mid y_{1:s})}\right) \underbrace{\left(\nabla_{\theta} \log p(y_t \mid y_{1:t-1})\right) \left(\log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})}\right)}_{\text{fixed by } y_{1:s}} \\ = \sum_{y_{1:s}} \mu(y_{1:s}) \frac{p(y_{1:s})}{\mu(y_{1:s})} \left(\nabla_{\theta} \log p(y_t \mid y_{1:t-1})\right) \left(\log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})}\right) \sum_{y_{s+1:T}} \mu(y_{s+1:T} \mid y_{1:s}) \frac{p(y_{s+1:T} \mid y_{1:s})}{\mu(y_{s+1:T} \mid y_{1:s})} \\ = \sum_{y_{1:s}} \mu(y_{1:s}) \frac{p(y_{1:s})}{\mu(y_{1:s})} \left(\nabla_{\theta} \log p(y_t \mid y_{1:t-1})\right) \left(\log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})}\right) \sum_{y_{s+1:T}} p(y_{s+1:T} \mid y_{1:s}) \\ = \sum_{y_{1:s}} \mu(y_{1:s}) \frac{p(y_{1:s})}{\mu(y_{1:s})} \left(\nabla_{\theta} \log p(y_t \mid y_{1:t-1})\right) \left(\log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})}\right) \cdot 1 \\ = \mathbb{E}_{y_{1:T} \sim \mu}\left[\frac{p(y_{1:s})}{\mu(y_{1:s})} \left(\nabla_{\theta} \log p(y_t \mid y_{1:t-1})\right) \left(\log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})}\right)\right] $$So the \((s,t)\) term's importance weight collapses from a ratio over the full sequence, \(p(y_{1:T})/\mu(y_{1:T})\), down to a ratio over just the prefix, \(p(y_{1:s})/\mu(y_{1:s})\). Putting every term back into the sum:
$$ \mathbb{E}_{y_{1:T} \sim \mu}\left[\frac{p(y_{1:T})}{\mu(y_{1:T})} \sum_{t=1}^{T} \nabla_{\theta} \log p(y_t \mid y_{1:t-1}) \left(\sum_{s=t}^{T} \log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})}\right)\right] \\ = \mathbb{E}_{y_{1:T} \sim \mu}\left[\sum_{t=1}^{T} \nabla_{\theta} \log p(y_t \mid y_{1:t-1}) \sum_{s=t}^{T} \frac{p(y_{1:T})}{\mu(y_{1:T})} \log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})}\right] \\ = \mathbb{E}_{y_{1:T} \sim \mu}\left[\sum_{t=1}^{T} \nabla_{\theta} \log p(y_t \mid y_{1:t-1}) \underbrace{\left(\sum_{s=t}^{T} \frac{p(y_{1:s})}{\mu(y_{1:s})} \log \frac{p(y_s \mid y_{1:s-1})}{q(y_s \mid y_{1:s-1})}\right)}_{\widetilde{R}_t}\right] $$The derivation hinges on each term in the outer sum over \(t\) not depending on \(s\), so it can be pulled inside the inner sum over \(s\). The upshot: importance sampling doesn't change the shape of the estimator at all — it only replaces \(R_t\) with a version that reweights each \(s\)-term by \(p(y_{1:s})/\mu(y_{1:s})\), the prefix importance ratio.
The gradient of on-policy distillation
On-policy distillation, as implemented in Thinking Machines' blog post, follows the correct token-level KL estimator above, using the teacher as the reference policy. It's presented purely as a token-level estimator — it doesn't claim to estimate the sequence-level KL, and, per the derivation above, it wouldn't be correct if it did.
← Back to all posts