From Exact KL Divergence to Schulman's $k_3$ Estimator

This document derives Schulman’s $k_3$ estimator for the reverse KL divergence between a student distribution and a teacher distribution, as used, for example, in knowledge distillation or RLHF-style fine-tuning, where a student model is trained against a fixed teacher (or reference) model. The unbiasedness results derived here hold for any student and teacher. The setting of interest, however, is one where the two are not too far apart, as is typical when a student is being fine-tuned against a teacher it already approximates reasonably well; it is precisely in that regime that $k_3$’s reduced-variance property, over the simpler $k_1$ estimator, becomes useful.

Setup

Let $h$ be a fixed prefix, and let $w$ denote the next token. The student and teacher define categorical distributions over the vocabulary $\mathcal{V}$, $S(w\mid h)$ and $T(w\mid h)$. Since $h$ is fixed throughout this derivation, it is suppressed from the notation below and both distributions are written simply as $S(w)$ and $T(w)$. Both distributions are normalized:

\[\begin{equation} \sum_{w\in\mathcal{V}}S(w)=1, \qquad \sum_{w\in\mathcal{V}}T(w)=1 \tag{1} \end{equation}\]

We want the reverse KL divergence, in which the student is the first distribution:

\[\begin{equation} D_{\mathrm{KL}}(S\parallel T) \tag{2} \end{equation}\]

The word “reverse” is a naming convention. What matters mathematically is the order of the arguments: tokens are sampled from the first distribution, the student.

Start from the exact KL divergence

For this fixed prefix $h$, the exact reverse KL is

\[\begin{equation} D_{\mathrm{KL}}(S\parallel T) = \sum_{w\in\mathcal{V}} S(w) \log\frac{S(w)}{T(w)} \tag{3} \end{equation}\]

Computing Equation (3) directly requires the student and teacher probabilities for every vocabulary token (a language-model vocabulary can have tens or hundreds of thousands of tokens, so storing or computing the full distribution from both models, for every generated token, is expensive).

For any function $f(w)$, the expectation under the student is defined by

\[\begin{equation} \mathbb{E}_{w\sim S}[f(w)] = \sum_{w\in\mathcal{V}}S(w)f(w) \tag{4} \end{equation}\]

Choose

\[\begin{equation} f(w)=\log\frac{S(w)}{T(w)} \tag{5} \end{equation}\]

Substituting Equation (5) into Equation (4) gives an equivalent expression for the exact KL:

\[\begin{equation} D_{\mathrm{KL}}(S\parallel T) = \mathbb{E}_{w\sim S} \left[ \log\frac{S(w)}{T(w)} \right] \tag{6} \end{equation}\]

Equations (3) and (6) are exactly equal. No approximation has been made yet.

Replace the expectation by student samples

Suppose that the student samples one word $w\sim S$. For this sampled token, only retain the student’s scalar log-probability for the sampled word:

\[\begin{equation} \ell_{S} := \log S(w) \tag{7} \end{equation}\]

The teacher does not independently sample a token. Instead, it scores the same student-sampled token, with the same prefix:

\[\begin{equation} \ell_{T} := \log T(w) \tag{8} \end{equation}\]

Their difference is

\[\begin{align} \ell_{S}-\ell_{T} &=\log S(w)-\log T(w) =\log\frac{S(w)}{T(w)} \tag{9} \end{align}\]

This single-sample quantity is Schulman’s $k_1$ estimator:

\[\begin{equation} k_1 := \ell_{S}-\ell_{T} \tag{10} \end{equation}\]

Because the sampled token comes from the student, its expectation is the exact reverse KL:

\[\begin{equation} \mathbb{E}_{w\sim S}[k_1] = D_{\mathrm{KL}}(S\parallel T) \tag{11} \end{equation}\]

In practice, $k_1$ is not evaluated once and reported directly. It is applied at every position of a sequence the student decodes, and the prefix is different at each position (it is everything generated so far). So each application of $k_1$ targets a different reverse KL, one per prefix. If the student decodes $N$ tokens, this produces $N$ single-sample estimates $k_{1,1},\ldots,k_{1,N}$, one per prefix, each unbiased for the reverse KL at its own prefix by Equation (11). Their average

\[\begin{equation} \widehat{D_{\mathrm{KL}}}_{k_1} = \frac{1}{N}\sum_{i=1}^{N} \left(\ell_{S,i}-\ell_{T,i}\right) \tag{12} \end{equation}\]

is, by linearity of expectation, an unbiased estimate of the average reverse KL across the $N$ prefixes visited. This is not the same as reducing the noise on one fixed KL value by resampling a single prefix: each term in the sum targets its own, generally different, prefix, and averaging remains valid regardless of whether those prefixes coincide.

Why sampled-token log-probabilities are sufficient

Here, a sampled-token log-probability means the log-probability, under a given model, of only the one token $w$ that the student actually sampled, namely $\log S(w)$ or $\log T(w)$, as opposed to the full log-probability vector over the whole vocabulary $\mathcal{V}$.

The exact KL (Equation (3)) is a weighted average over the vocabulary, where each token’s contribution is scaled by how likely the student is to produce it, $S(w)$. Sampling a token from the student already performs this weighting: a token with student probability $0.5$ appears in about half of the samples, one with probability $0.01$ in about one in a hundred. So whichever token happens to be sampled, its own log-ratio $\log S(w)-\log T(w)$ already carries the right amount of weight simply by how likely it was to be the one sampled, and nothing further needs to be multiplied in.

This is exactly what $k_1$ computes (Equation (10)), and it is why $k_1$ depends only on the two scalars $\ell_{S}=\log S(w)$ and $\ell_{T}=\log T(w)$ for the sampled token, and on no other vocabulary entry. Thus, for every student-sampled token, whether these tokens come from one fixed prefix or from a different prefix at every position of a generated sequence, the stored data need only be the two scalars $\ell_{S}$ and $\ell_{T}$ (the student’s log-probability of the token it sampled, and the teacher’s log-probability assigned to that same student-sampled token). There is no need to retain two vectors of size $|\mathcal{V}|$ (the full distribution).

There is one subtlety. To calculate a normalized log-probability from logits, the model internally uses all vocabulary logits. For each vocabulary token $v\in\mathcal{V}$, let $z_v$ denote the model’s raw output logit for $v$, before softmax normalization. Then

\[\begin{equation} \log S(w) = z_w-\log\sum_{v\in\mathcal{V}}e^{z_v} \tag{13} \end{equation}\]

Therefore, “store only the log-probability” does not mean that the other logits never exist during the forward pass. It means that the full probability distribution does not need to be retained for later KL computation. Only the selected token’s normalized log-probability must be saved.

Obtaining a simple estimator: $k_1$

Define the teacher-to-student probability ratio

\[\begin{equation} r:= \frac{T(w)}{S(w)} \tag{14} \end{equation}\]

This direction may initially look surprising because the target is reverse KL. However,

\[\begin{equation} -\log r = \log\frac{S(w)}{T(w)} \tag{15} \end{equation}\]

which is exactly the reverse-KL integrand. From the definition of $r$ and the two scalar log-probabilities,

\[\begin{equation} \log r = \log\frac{T(w)}{S(w)} = \log T(w)-\log S(w) =\ell_{T}-\ell_{S} \tag{16} \end{equation}\]

Exponentiating both sides recovers $r$ itself, since $\exp(\log r)=r$ for $r>0$:

\[\begin{equation} r =\exp\!\left(\ell_{T}-\ell_{S}\right) \tag{17} \end{equation}\]

Consequently,

\[\begin{equation} k_1=-\log r \tag{18} \end{equation}\]

This form is useful because it re-expresses $k_1$ using only the ratio $r$, which is exactly the quantity used in the rest of this derivation. In practice, estimating the reverse KL from this form means computing $r$ for each sampled token and averaging $-\log r$ over the sequence of $N$ tokens:

\[\begin{equation} \widehat{D_{\mathrm{KL}}}_{k_1} = \frac{1}{N}\sum_{i=1}^{N}(-\log r_i) \tag{19} \end{equation}\]

equivalent to Equation (12), now written in terms of $r$ instead of the two log-probabilities directly.

Obtain an equivalent estimator with better properties: $k_3$

Find a term with zero expectation

We now calculate the expectation of $r$ under student sampling. This is a theoretical derivation; it is not a vocabulary-wide sum that must be performed in the implementation:

\[\begin{align} \mathbb{E}_{w\sim S}[r] &= \sum_{w\in\mathcal{V}} S(w) \frac{T(w)}{S(w)} \tag{20}\\ &= \sum_{w\in\mathcal{V}}T(w) \tag{21}\\ &=1 \tag{22} \end{align}\]

The move from Equation (21) to Equation (22) follows only from the fact that the teacher is a normalized probability distribution over the vocabulary, as stated in Equation (1). It follows that

\[\begin{equation} \mathbb{E}_{w\sim S}[r-1] = \mathbb{E}_{w\sim S}[r]-1 =0 \tag{23} \end{equation}\]

This argument assumes that the ratio is well-defined wherever needed. For ordinary softmax language models, all vocabulary probabilities are positive, so this is normally satisfied.

Starting from the expectation form of reverse KL,

\[\begin{equation} D_{\mathrm{KL}}(S\parallel T) = \mathbb{E}_{w\sim S}[-\log r] \tag{24} \end{equation}\]

Equation (23) shows that $r-1$ has expectation zero. We may therefore add it inside the expectation without changing the expected value:

\[\begin{align} D_{\mathrm{KL}}(S\parallel T) &= \mathbb{E}_{w\sim S} \left[-\log r+(r-1)\right] \tag{25}\\ &= \mathbb{E}_{w\sim S} \left[r-1-\log r\right] \tag{26} \end{align}\]

Schulman’s $k_3$ single-sample estimator is therefore defined as

\[\begin{equation} \boxed{k_3:=r-1-\log r} \tag{27} \end{equation}\]

Substituting Equations (16) and (17) gives a form that uses only the sampled token’s two log-probabilities:

\[\begin{equation} \boxed{ k_3 = \exp\!\left(\ell_{T}-\ell_{S}\right) -1 -\left(\ell_{T}-\ell_{S}\right) } \tag{28} \end{equation}\]

The sample estimator $k_3$ and the simpler estimator $k_1$ have the same expectation:

\[\begin{equation} \mathbb{E}_{w\sim S}[k_3] = \mathbb{E}_{w\sim S}[k_1] = D_{\mathrm{KL}}(S\parallel T) \tag{29} \end{equation}\]

Because both estimators share this expectation, either $k_1$ or $k_3$ can be used, in the same single-sample-and-average manner shown above, to estimate the reverse KL divergence; the choice between them is a question of estimator variance, addressed next.

Why $k_3$ is useful

The unbiasedness of $k_1$ and $k_3$ shown above holds for any $S$ and $T$; no closeness assumption was used anywhere in that derivation. The argument in this section is different in kind: it explains why $k_3$ has lower sampling variance than $k_1$, and that explanation does assume $S\approx T$. This assumption is local to this section only.

An individual $k_1=-\log r$ sample can be negative, even though the true KL is nonnegative. This is not a correctness problem: only the expectation of $k_1$ must equal the KL. The more important advantage of $k_3$ is reduced sampling noise when the two distributions are close: specifically, when $S(w)\approx T(w)$ for essentially every token $w\in\mathcal{V}$, so that $r=T(w)/S(w)$ is close to $1$ for whichever token happens to be sampled, not merely for one particular token in isolation.

Define

\[\begin{equation} \epsilon:=r-1, \qquad |\epsilon|\ll 1 \tag{30} \end{equation}\]

where $|\epsilon|\ll 1$ is precisely the closeness condition on $S$ and $T$ just described. Equivalently,

\[\begin{equation} r=1+\epsilon \tag{31} \end{equation}\]

Using

\[\begin{equation} \log(1+\epsilon) = \epsilon-\frac{\epsilon^2}{2}+O(\epsilon^3) \tag{32} \end{equation}\]

the $k_1$ estimator is, by Taylor approximation,

\[\begin{equation} k_1 = -\log(1+\epsilon) = -\epsilon+\frac{\epsilon^2}{2}+O(\epsilon^3) \tag{33} \end{equation}\]

The first-order term $-\epsilon$ has zero expectation (since $\epsilon=r-1$ and $\mathbb{E}_{w\sim S}[r-1]=0$ was already shown in Equation (23); this is not an assumption about random noise, but a direct consequence of that earlier result) but can contribute substantial sample-to-sample noise.

For $k_3$,

\[\begin{align} k_3 &=(1+\epsilon)-1-\log(1+\epsilon) \tag{34}\\ &=\frac{\epsilon^2}{2}+O(\epsilon^3) \tag{35} \end{align}\]

The first-order noise cancels. Also, the inequality

\[\begin{equation} \log r\leq r-1, \qquad r>0 \tag{36} \end{equation}\]

implies that $k_3$ is nonnegative:

\[\begin{equation} k_3=r-1-\log r\geq 0 \tag{37} \end{equation}\]

Nonnegativity is convenient for monitoring a quantity that is theoretically nonnegative, but the cancellation of first-order noise is the main statistical motivation: both approximate the KL, but $k_3$ has reduced noise in comparison to $k_1$.

Reference

The estimators discussed here follow John Schulman’s note, Approximating KL Divergence (2020).