Re-understanding KL Approximation from an RL-for-LLM Lens: Notes on “Approximating KL Divergence”

What's the difference between KL-divergence estimation methods used in PPO and GRPO?

John Schulman’s blog post “Approximating KL Divergence” talks about how to approximate KL divergence via sampling (Monte Carlo), and it introduces three estimators (\(k_1\), $k_{2}k\_2$, $k_{3}k\_3$) along with their bias–variance behaviors. But the original post is framed in the context of general probability distributions; it doesn’t touch the reinforcement-learning-for-LLM training setting. This write-up records the questions I had while reading, the thoughts I formed after mapping things to RL for LLMs, and a few places where I felt the original explanations could be pushed on a bit.

What “Approximating KL Divergence” Says (in my own words)

In this section, I’m imagining readers who haven’t read the original post yet, so let’s quickly pass through the most important bits. Put simply, the post is about how we can build reasonable Monte Carlo–style estimators for it, when we can’t directly compute KL divergence.

$$\mathrm{K}\mathrm{L}(q,p)=\sum _{x}q(x)\log \frac{q(x)}{p(x)}={\mathbb{E}}_{x\sim q}\text{\hspace{0.17em}⁣}[\log \frac{q(x)}{p(x)}]\mathrm{.}\backslash mathrm\{KL\}(q,\; p)\; =\; \backslash sum\_x\; q(x)\backslash ,\backslash log\backslash frac\{q(x)\}\{p(x)\}\; =\; \backslash mathbb\{E\}\_\{x\backslash sim\; q\}\backslash !\backslash left[\backslash log\backslash frac\{q(x)\}\{p(x)\}\backslash right].$$

As the formula shows: when estimating the KL between two (complicated) distributions, there’s a coding trick people often use: just approximate KL by the sample mean of $\log \text{\hspace{0.17em}⁣}\left(\frac{q(x)}{p(x)}\right)\backslash log\backslash !\backslash big(\backslash frac\{q(x)\}\{p(x)\}\backslash big)$ with samples drawn from $qq$ (as opposed to trying to evaluate the full expectation exactly). The post then points out another move: use the sample average of $\frac{1}{2}(\log r{)}^2\backslash tfrac\{1\}\{2\}(\backslash log\; r)^2$ to replace the more “standard” $\log r\backslash log\; r$ form, where $r=\frac{q(x)}{p(x)}r=\backslash frac\{q(x)\}\{p(x)\}$. The write-up explains why this expression can be a good (albeit biased) estimator for KL, and how to make it unbiased while keeping its low variance.

How we compute KL depends on how we can access $pp$ and $qq$. Here we assume we can evaluate $p(x)p(x)$ and $q(x)q(x)$ (probabilities or densities) for any $xx$, but we can’t analytically sum/integrate over $xx$. Why might we fail to do the analytic sum/integral? Maybe the exact computation is too expensive in compute or memory, maybe there’s no closed form, or maybe we only store log-probs instead of full distributions to keep code simpler, especially fine when KL is just a diagnostic (as is often the case in RL). The most common strategy to approximate sums or integrals is Monte Carlo. Given samples $x_1,x_2,\dots ,x_n\sim qx\_1,\; x\_2,\; \backslash dots,\; x\_n\; \backslash sim\; q$, how do we build a good estimator?

A good estimator should be unbiased (right mean) and low-variance. We know one unbiased estimator:

$$k_1=\log \frac{q(x)}{p(x)}\mathrm{.}k\_1\; =\; \backslash log\backslash frac\{q(x)\}\{p(x)\}.$$

But it has high variance: by definition KL is a nonnegative quantity, yet for the estimator above, roughly “half” the sample values can be negative (if we assume no prior structure on $pp$ and $qq$), and this swings the average around a lot, hence high variance. For notational convenience, set $r=\frac{q(x)}{p(x)}r\; =\; \backslash frac\{q(x)\}\{p(x)\}$. Then the original KL can be written as

$$\mathrm{K}\mathrm{L}[q,p]={\mathbb{E}}_{x\sim q}[\log r]\mathrm{.}\backslash mathrm\{KL\}[q,\; p]\; \backslash ;=\backslash ;\; \backslash mathbb\{E\}\_\{x\backslash sim\; q\}\backslash ,[\backslash log\; r].$$

To reduce variance, we can design an alternative estimator: $$k_2=\frac{1}{2}(\log r{)}^2\mathrm{.}k\_2\; =\; \backslash frac\{1\}\{2\}(\backslash log\; r)^2.$$ It has lower variance, but it’s biased. Intuitively, $k_{2}k\_2$ feels nicer because each sample gives a nonnegative “distance” between $pp$ and $qq$, so it stays positive. Empirically, $k_{2}k\_2$ really does have much lower variance than $k_{1}k\_1$, and the bias can be quite small. As for why $k_{2}k\_2$ enjoys such a variance drop compared with $k_{1}k\_1$, the original post uses an $ff$-divergence view to give an analytic explanation, and I won’t repeat that here.

Now, can we get an estimator that is both unbiased and low-variance? A general trick is to use a control variate: start from the unbiased $k_{1}k\_1$ and add something whose expectation is zero and is negatively correlated with it to reduce variance. A very convenient zero-mean quantity here is $r-1r-1$. Thus, for any $\lambda \backslash lambda$, $$k=-\log r+\lambda (r-1)k\; \backslash ;=\backslash ;\; -\backslash log\; r\; +\; \backslash lambda\backslash ,(r-1)$$ is still an unbiased KL estimator. In theory we could minimize the variance over $\lambda \backslash lambda$, but the closed-form depends on $pp$ and $qq$ and isn’t easy to get. Notice, though, that since $\log (x)\backslash log(x)$ is concave, $$\log (x)\le x-1,\backslash log(x)\; \backslash ;\backslash le\backslash ;\; x-1,$$ so if we pick $\lambda =1\backslash lambda=1$, the expression is guaranteed nonnegative. Here, $r-1r-1$ is the tangent of $\log r\backslash log\; r$ at $r=1r=1$. So with $\lambda =1\backslash lambda=1$ we’re really measuring the vertical gap between $\log (x)\backslash log(x)$ and its tangent. This leads to the estimator $$k_3=(r-1)-\log r,k\_3\; \backslash ;=\backslash ;\; (r\; -\; 1)\; \backslash ;-\backslash ;\; \backslash log\; r,$$ which is always nonnegative. And $k_{3}k\_3$ is exactly the piece where, in practice, GRPO differs from PPO in how KL is estimated (PPO uses $k_{1}k\_1$).

Discussing KL Estimation from an RL-for-LLM Perspective

In RL (think PPO, GRPO, etc.), we often tack a KL divergence term onto the loss to keep the new policy from drifting too far from the old one. Here, $qq$ is the old policy distribution ( ${\pi }_{\text{old}}\backslash pi\_\{\backslash text\{old\}\}$ ), $pp$ is the new policy distribution ( ${\pi }_{\text{new}}\backslash pi\_\{\backslash text\{new\}\}$ ), and $xx$ is a complete action sample (in an LLM this means a token or a token sequence). We usually use $ss$ to denote the state (in an LLM, that’s the prompt or context), and $xx$ is a specific token generated in that context. When we compute KL, what we’re really doing is taking the KL over the action distribution given a state, and then averaging over states:

$$\mathrm{K}\mathrm{L}[p,q]={\mathbb{E}}_s[\sum _{x}p(x\mathrm{\mid }s)\log \frac{p(x\mathrm{\mid }s)}{q(x\mathrm{\mid }s)}]\mathrm{.}\backslash mathrm\{KL\}[p,\; q]\; =\; \backslash mathbb\{E\}\_\{s\}\; \backslash left[\; \backslash sum\_x\; p(x|s)\; \backslash log\; \backslash frac\{p(x|s)\}\{q(x|s)\}\; \backslash right].$$

At sampling time, we typically fix a prompt (state) and then estimate this KL for that prompt.

So why can’t we just compute KL exactly instead of estimating it? The reasons are exactly those listed in the original blog post; in RL for LLMs, the main culprit is Reason #1: the action space (token space) is too large to sum/integrate over all possible $xx$. For example, if a tokenizer has 50,000 vocabulary entries, even computing the KL for a single token means summing over 50,000 actions; and in RL we’re usually doing multi-step (sequence) generation, so the space blows up exponentially, which is completely impractical. There’s also a pragmatic reason: during training we generally don’t store the full distribution (all token probabilities); we only keep the log-probs of the tokens actually generated along the trajectory, to save GPU memory and I/O. So we have to use Monte Carlo sampling: draw $xx$ from some distribution (usually $qq$, the old policy), and use those samples to approximate KL. And that drops us squarely into the territory the blog post is about.

In that post, the estimator we keep talking about is really just a function of a sample: it takes $p(x)p(x)$ and $q(x)q(x)$ for some sampled $xx$ (or their ratio $r=\frac{q(x)}{p(x)}r\; =\; \backslash frac\{q(x)\}\{p(x)\}$) and spits out a number. We then take the average of those numbers over our samples to approximate KL. For example:

• $k_1(x)=-\log rk\_1(x)\; =\; -\backslash log\; r$
• $k_2(x)=\frac{1}{2}(\log r{)}^{2}k\_2(x)\; =\; \backslash frac12\; (\backslash log\; r)^2$
• $k_3(x)=(r-1)-\log rk\_3(x)\; =\; (r\; -\; 1)\; -\; \backslash log\; r$

These $k_{i}k\_i$ are just different KL-estimator formulas. They all approximate KL by averaging over samples, but differ in bias and variance. Once we pick an estimator, we’re really just committing to a specific formula for approximating KL. The process looks like this:

1. Sampling
   Sample a batch of tokens (or sequences) $x_1,x_2,\dots ,x_{N}x\_1,\; x\_2,\; \backslash dots,\; x\_N$ from the old policy $qq$.
2. Compute log-probs
   For each sample, compute the log-probabilities under both new and old policies:

$$\log p(x_i),\log q(x_i)\backslash log\; p(x\_i),\backslash \; \backslash log\; q(x\_i)$$

and get $r_i=\frac{q(x_i)}{p(x_i)}r\_i\; =\; \backslash frac\{q(x\_i)\}\{p(x\_i)\}$ or $\log r_i\backslash log\; r\_i$. 3. Plug into the estimator formula
For example, if we choose $k_{3}k\_3$:

$$k_3(x_i)=(r_i-1)-\log r_{i}k\_3(x\_i)\; =\; (r\_i\; -\; 1)\; -\; \backslash log\; r\_i$$

4. Average

$$\widehat{\mathrm{K}\mathrm{L}}\approx \frac{1}{N}\sum _{i=1}^N{k}_3(x_i)\backslash widehat\{\backslash mathrm\{KL\}\}\; \backslash approx\; \backslash frac1N\; \backslash sum\_\{i=1\}^N\; k\_3(x\_i)$$

That’s the approximate KL value, standing in for the true KL.

If we compare this to computing the true KL (no estimation) for a discrete probability distribution (LLM single-token step): we’d need to iterate over every possible token $xx$: $$\mathrm{K}\mathrm{L}(p\mathrm{\parallel }q)=\sum _{x}p(x)\log \frac{p(x)}{q(x)}\backslash mathrm\{KL\}(p\backslash |q)\; =\; \backslash sum\_x\; p(x)\; \backslash log\; \backslash frac\{p(x)\}\{q(x)\}$$ You can see immediately that with an estimator, the computational load is much smaller than doing the full sum, especially in high-dimensional action spaces.

Talking About Variance in Different KL Estimators

Important to note: the “variance” we’re talking about here is the variance of the values the estimator outputs over samples: $${\mathrm{V}\mathrm{a}\mathrm{r}}_{x\sim q}[k(x)]\backslash mathrm\{Var\}\_\{x\; \backslash sim\; q\}[k(x)]$$ That is, how much $k(x)k(x)$ fluctuates across the sample space. An unbiased estimator means that with infinitely many samples, its mean equals the true KL. But a high-variance estimator means that even if the mean is right (unbiased), with a small number of samples the average can be way off. In RL for LLMs, the KL term is often a regularization factor in the loss (e.g., $\beta \cdot \mathrm{K}\mathrm{L}\backslash beta\; \backslash cdot\; \backslash mathrm\{KL\}$). If the KL estimator’s variance is large, it makes the loss noisy, which in turn makes gradients noisy and training unstable.

In the original post, to give readers an intuition for why $k_{1}k\_1$ is not low-variance, the author writes:

However, it ( $k_{1}k\_1$ ) has high-variance, as it’s negative for half of the samples, whereas KL is always positive.

The author points out that although $k_{1}k\_1$ is unbiased, without prior constraints on $pp$ and $qq$, half the samples will have one bigger than the other, so half of the $k_{1}k\_1$ values are positive and half negative. Up to here I’m fine. But then the author says: because KL is always greater than 0 (a basic inequality), $k_{1}k\_1$ must therefore have high variance. And here I think the causal link doesn’t actually hold: you can’t use the sign of the expectation to dictate the sign of individual samples. A quick counterexample: in computing the expectation, $p(x)\log \frac{p(x)}{q(x)}p(x)\; \backslash log\; \backslash frac\{p(x)\}\{q(x)\}$ is also sometimes positive and sometimes negative; that fact by itself tells you nothing about variance. In reality, a single-sample log ratio (whether $\log \frac{q(x)}{p(x)}\backslash log\; \backslash frac\{q(x)\}\{p(x)\}$ or $\log \frac{p(x)}{q(x)}\backslash log\; \backslash frac\{p(x)\}\{q(x)\}$) can be positive or negative, just like $k_{1}k\_1$, so sign-flipping alone is not the sole reason for high variance.

From the KL definition: $$\mathrm{K}\mathrm{L}(q\mathrm{\parallel }p)={\mathbb{E}}_{x\sim q}[\log \frac{q(x)}{p(x)}]\backslash mathrm\{KL\}(q\; \backslash |\; p)\; =\; \backslash mathbb\{E\}\_\{x\backslash sim\; q\}\backslash left[\; \backslash log\; \backslash frac\{q(x)\}\{p(x)\}\; \backslash right]$$ The expectation is guaranteed nonnegative, but the integrand $\log \frac{q(x)}{p(x)}\backslash log\backslash frac\{q(x)\}\{p(x)\}$ can be positive or negative for individual samples. And $k_{1}k\_1$ is exactly that integrand: $$k_1(x)=\log \frac{q(x)}{p(x)}k\_1(x)\; =\; \backslash log\; \backslash frac\{q(x)\}\{p(x)\}$$ So each sample value can indeed be positive or negative, same as the integrand in the KL definition.

So why does $k_{1}k\_1$ have high variance?

It’s not the mere “sign flipping.” The real reason is that $k_{1}k\_1$’s value distribution is often wide (heavy-tailed). For example, if $p(x)p(x)$ is tiny for some sample, then $\log \frac{q}{p}\backslash log\backslash frac\{q\}\{p\}$ can be huge (positive or negative). These extreme values dominate the finite-sample average, pushing variance up. In other words, it’s the combination of extreme values + positive/negative cancellation: cancellation means you need more samples to converge to the true mean, and extreme values make the sample variance itself larger. So the “half negative” comment in the blog is more of an intuition hook than a complete explanation.

From this perspective, if we look at the other estimators $k_{2}k\_2$ and $k_{3}k\_3$, we see: $k_2=\frac{1}{2}(\log r{)}^{2}k\_2\; =\; \backslash frac12\; (\backslash log\; r)^2$ is always positive, so there’s no cancellation, but this introduces bias; squaring also smooths the magnitude, reducing variance. $k_{3}k\_3$ uses a control variate to knock out part of the fluctuation source, lowering variance while keeping unbiasedness (details next).

In PPO/GRPO, if you use $k_{1}k\_1$ and the batch is small or the distributions are far apart, the KL estimate will jump around (because a few extreme samples can swing the mean hard). That makes the KL-penalty coefficient unstable: it might suddenly be way too strong or too weak. Switching to a lower-variance estimator ( $k_{2}k\_2$ or $k_{3}k\_3$ ) makes each sample’s KL contribution steadier, less likely to be dominated by a handful of extreme samples.

Why can $k_{3}k\_3$ be unbiased and low-variance?

At first glance, $k_{3}k\_3$ is always positive, so you might think its mean must be larger than $k_{1}k\_1$’s.
But remember: $k_{3}k\_3$ is derived from $k_{1}k\_1$ via a control variate. The blog’s reasoning goes like this: $$\tilde{k}(x)=k_1(x)+\lambda \cdot h(x)\backslash tilde\{k\}(x)\; =\; k\_1(x)\; +\; \backslash lambda\; \backslash cdot\; h(x)$$ where $h(x)=r-1h(x)\; =\; r\; -\; 1$, and under $x\sim qx\backslash sim\; q$ its expectation is: $${\mathbb{E}}_{x\sim q}[h(x)]={\mathbb{E}}_q[\frac{p(x)}{q(x)}-1]=\sum _{x}p(x)-1=1-1=0.\backslash mathbb\{E\}\_\{x\backslash sim\; q\}[h(x)]\; =\; \backslash mathbb\{E\}\_q\backslash left[\backslash frac\{p(x)\}\{q(x)\}\; -\; 1\backslash right]\; =\; \backslash sum\_x\; p(x)\; -\; 1\; =\; 1\; -\; 1\; =\; 0.$$ So adding any multiple of $h(x)h(x)$ doesn’t change the expectation. When $\lambda =1\backslash lambda\; =\; 1$: $$\tilde{k}(x)=-\log r+(r-1)=(r-1)-\log r=k_3(x)\mathrm{.}\backslash tilde\{k\}(x)\; =\; -\backslash log\; r\; +\; (r\; -\; 1)\; =\; (r\; -\; 1)\; -\; \backslash log\; r\; =\; k\_3(x).$$ This explains why $k_{3}k\_3$’s expectation equals $k_{1}k\_1$’s expectation, and equals the KL, making it an unbiased estimator.

The reason $k_{3}k\_3$ has lower variance than $k_{1}k\_1$ is: $k_{1}k\_1$ only has $-\log r-\backslash log\; r$, which can swing wildly (both positive and negative, with occasional huge values). But $r-1r\; -\; 1$ and $-\log r-\backslash log\; r$ are numerically highly correlated (when one grows, the other grows/shrinks), and that correlation is negative. Adding $(r-1)(r\; -\; 1)$ is like injecting a negatively correlated term to cancel fluctuations. After cancellation, what’s left in $k_{3}k\_3$ is tighter in range, always positive, and therefore lower in sample variance.