The best estimate of KL ≢ the best estimate of its gradient

KL divergence \mathrm{KL}[p\,\|\,q] = \mathbb{E}_{x\sim p}\!\left[\log\frac{p(x)}{q(x)}\right] is an important quantity in AI.

In reinforcement learning for language models, \mathrm{KL}[\pi_\theta\,\|\,\pi_\mathrm{ref}] is used as a regularizer to keep a learned policy \pi_\theta near a reference policy \pi_\mathrm{ref}, and on-policy distillation (OPD) lets a student policy \pi_\mathrm{student} learn from a teacher policy \pi_\mathrm{teacher} by minimizing \mathrm{KL}[\pi_\mathrm{student}\,\|\,\pi_\mathrm{teacher}].

Exactly computing KL between two language models is intractable, because the expectation \mathbb{E}_p is a sum over every possible sequence of tokens. Instead we estimate it from samples x \sim p.

Three estimators of KL are common:

You might recall these estimators from an excellent post by John Schulman. It shows that k_3 in particular is unbiased like k_1 and never negative like k_2, with lower variance.

We can use these functions to estimate KL, but how should we use them to minimize it? Let’s explore how to minimize reverse KL (\min_p\mathrm{KL}[p\,\|\,q]), since it is relevant for reinforcement learning.

Estimating the value of KL #

Before trying to minimize KL, let’s compare the estimators on two Gaussians, p = \mathcal{N}(\mu, 1) and q = \mathcal{N}(0, 1). Here, \mathrm{KL}[p\,\|\,q] is exactly \mu^2/2why?\displaystyle\log\frac{p}{q} = -\frac{(x-\mu)^2}{2} + \frac{x^2}{2} = \mu x - \frac{\mu^2}{2}. Taking the expectation under p, where \mathbb{E}_p[x] = \mu, gives \displaystyle\mathrm{KL}[p\,\|\,q] = \mu\mathbb{E}_p[x] - \frac{\mu^2}{2} = \frac{\mu^2}{2}..

Let’s first examine bias. An estimator is unbiased when its expectation equals the true KL.

k_1 is unbiased by the definition of KL:

k_2 is biased. For these Gaussians,

k_3 adds a zero-mean term to k_1, so it stays unbiased:

As the graph shows, k_1 and k_3 are unbiased, while k_2 overestimates KL. The key to k_3 staying unbiased is that its added term, q/p - 1, is zero in expectation. It can be nonzero for an individual sample, but it does not change the estimator’s mean.

This zero-mean term is a control variate: we add it to an estimator to adjust its variance without changing its expectation. It can cancel noise and decrease variance, or add noise and increase variance. For these Gaussians, k_3 has lower variance near q and higher variance farther away:

Because the control variate is zero-mean, any multiple of it can be added without biasing the estimator, yielding the family of estimators:

Here, \lambda = 0 gives k_1, and \lambda = 1 gives k_3. We’ll revisit how to choose \lambda.

So choosing the best estimate depends on how far p is from q. If all you want is to estimate KL, k_3 is the best choice near q and k_1 far from it.

The gradient of KL #

In RL, the KL estimate is not the end product. We differentiate it. Using the log derivative trick \nabla p = p\,\nabla\log pwhy?By the chain rule, \displaystyle\nabla\log p = \frac{\nabla p}{p}. Multiplying by p gives \nabla p = p\,\nabla\log p. This lets us rewrite \sum_x \nabla p\,f(x) as \mathbb{E}_p[f(x)\,\nabla\log p]., the gradient of KL is as follows:

We can drop the second term, \mathbb{E}_p[\nabla\log p], as its expectation is zerowhy?Probabilities always sum to one, so their changes sum to zero: \displaystyle\mathbb{E}_p[\nabla\log p] = \sum_x p\,\frac{\nabla p}{p} = \sum_x \nabla p = \nabla \sum_x p = \nabla 1 = 0.:

This tidies up the math. However, notice that \nabla\log p plays the same role for the gradient that \frac{q}{p} - 1 plays for the estimate: a control variate with mean zero.

\nabla\,\mathrm{KL} \equiv k_1 in the reward \equiv k_2 as a loss #

The most common implementations of minimizing KL use one of two mechanisms: adding k_1 to the reward, or k_2 as a surrogate loss.

k_1 in the reward yields the correct gradient. In policy gradient, a sample x with a fixed reward R(x) contributes R(x)\,\nabla\log p. Notice that directly setting R(x) as k_1 yields the correct gradient:

Similarly, using k_2 directly as a loss function also yields the correct gradient:

\nabla\,\mathrm{KL} \equiv \nabla\mathbb{E}_p[k_\lambda] \equiv scaling the control variate #

Now return to the family k_\lambda = \log\frac{p}{q} + \lambda\left(\frac{q}{p} - 1\right) from the estimation section, where \lambda sets the magnitude of the estimate’s control variate. Since k_\lambda is unbiased, we should be able to show that the gradient of its expectation is the same as the gradient of KL:

The second term can be canceled, as \mathbb{E}_p[ (1 - \lambda)\,\nabla\log p] = 0why?Probabilities always sum to one, so their changes sum to zero: \displaystyle\mathbb{E}_p[\nabla\log p] = \sum_x p\,\frac{\nabla p}{p} = \sum_x \nabla p = \nabla \sum_x p = \nabla 1 = 0.:

Notice that again, we’ve found a zero-mean term (1 - \lambda)\,\nabla\log p that does not bias the expectation when added, i.e. a control variate. In k_\lambda it is \lambda\left(\frac{q}{p} - 1\right), but in the \nabla\mathbb{E}_p[k_\lambda] it is (1 - \lambda)\,\nabla\log p. This suggests that choosing \lambda to minimize the variance of KL might not minimize the variance of its gradient.

\nabla\mathbb{E}_p[k_1] (\lambda=0) keeps all of the \nabla\log p term in the gradient, and is the exact same as \nabla\,\mathrm{KL}. Similarly, \nabla\mathbb{E}_p[k_3] (\lambda=1) keeps none of it, and is the exact same as setting k_1 as the reward or k_2 as a loss. All of these are unbiased gradients of KL, but with different variances as a result of adding different magnitudes of the control variate.

In RL language, we are shifting the reward (\log\frac{p}{q}) by a baseline (1-\lambda), the standard tool for reducing a policy gradient’s variance.

The best estimate does not yield the best gradient #

Schulman’s note says that choosing \lambda to minimize the variance of k_\lambda depends on p and q and is hard to calculate analytically. Instead, it uses \lambda = 1, because then the estimate can never be negative.

What if we wanted to choose \lambda optimally? The variance of k_\lambda is quadratic in \lambda, smallest at

However, in training we usually don’t use the value of KL. We use its gradient. The corresponding per-sample gradient is \big(\log\frac{p}{q} + 1 - \lambda\big)\nabla\log p, where \lambda sets the size of a control variate (1 - \lambda)\,\nabla\log p. The variance of this gradient is smallest at

Let’s consider the Gaussians again, p = \mathcal{N}(\mu, 1) and q = \mathcal{N}(0, 1).

Taking the gradient taken with respect to \mu, \nabla\log p = x - \mu. Let g_\lambda = \big(\log\frac{p}{q} + 1 - \lambda\big)\nabla\log p for the gradient. Then

These variances are smallest at

The two optima move in opposite directions. As p moves away from q, the best estimate slides from k_3 toward k_1, while the best gradient wants a larger and larger baseline. Move the slider to see it.

A simple bandit experiment: optimizing KL with each \lambda #

Does the difference matter in training?

Let’s take a softmax policy over 1,000 actions (i.e. a bandit) and a fixed reference q, and initializing the policy q far from q (\mathrm{KL} = 2.43). We’ll minimize \mathrm{KL}[p\,\|\,q] alone by plain SGD on the logits using policy gradient. Each step samples 16 actions and averages \big(\log\frac{p}{q} + 1 - \lambda\big)\nabla\log p, with \lambda set four ways: \lambda = 0 (k_1 differentiated exactly); \lambda = 1 (k_3); \lambda_\text{value}; and \lambda_\text{gradient}, the last two computed exactly from the current policy. Every choice gives an unbiased gradient, so any difference comes from variance.

Epilogue #

A KL estimator can be used to either estimate KL, or estimate the gradient of KL. Usually, we want the latter, which can be unstable if we mistakenly optimize for lower KL estimate variance instead of lower KL gradient variance.

Practically, this doesn’t change much for the LLM status-quo of using k_1 in the reward or k_2 as a loss, but it is useful to understand the distinction.

Appendix #

Mixing up “the estimator” with “how the estimator enters the gradient” is prevalent despite others pointing out these pitfalls before (Tang and Munos; Liu et al.). This shows up in popular code and reports: open-source RL frameworks, the Composer 2 report, and DeepSeek’s GRPO papers.