Table of Contents:

• Overview
• SGD
• Adam

KL Divergence

KL divergence is a key tool in RL (PPO), distillation, and more today. In this post, we’ll cover self-information $\to$ entropy $\to$ cross-entropy $\to$ KL divergence, forward versus reverse KL, and practical Monte Carlo estimators.

Jia Bin Markdown

Fantastic KL Divergence and How to (Actually) Compute It

Contents

1. Why Compare Probability Distributions?
2. Self-Information: How Surprising Is an Event?
3. Entropy: Average Surprise
4. Cross-Entropy: Coding $P$ Using $Q$
5. KL Divergence
6. Forward KL vs. Reverse KL
7. Exact Computation Can Be Expensive
8. Monte Carlo Estimator $k_1$: Unbiased but High Variance
9. Estimator $k_2$: Low Variance but Biased
10. Control Variates: Unbiased and Low Variance
11. Numerically Stable Implementation
12. Comparison of the Three Estimators
13. Machine-Learning Applications
14. Takeaways
15. References

Why Compare Probability Distributions?

Suppose \(P\) is the true or reference distribution and \(Q\) is an approximation, model, or alternative distribution. We want a quantity that answers:

KL divergence provides exactly this information-theoretic penalty.

Self-Information: How Surprising Is an Event?

For an event \(x\) with probability \(P(x)\), its self-information is

Rare events carry more information because they are more surprising.

• If the logarithm is base $2$, information is measured in bits.
• If the logarithm is natural, information is measured in nats.

For example, an event with probability $1/8$ carries

Entropy: Average Surprise

The entropy of a discrete distribution $P$ is the expected self-information:

For a continuous density $p$, the analogous differential entropy is

Entropy measures uncertainty, but it can also be interpreted as the optimal average code length for samples from $P$.

A coding example

Consider three weather outcomes with

An optimal prefix code can assign one bit to the most likely event and two bits to each of the others:

The average code length is

and indeed

Cross-Entropy: Coding $P$ Using $Q$

If the data is generated from $P$, but we use a coding or predictive distribution $Q$, the average surprise is the cross-entropy

Using

we obtain

The wrong model costs an extra $0.25$ bits per outcome on average.

KL Divergence

The KL divergence from $Q$ to $P$ is the excess cross-entropy:

For discrete distributions,

and for continuous densities,

For the weather example,

Important properties

Assuming \(P\) is absolutely continuous with respect to \(Q\):

1. Non-negativity: \(D_{\mathrm{KL}}(P\|Q)\ge 0\).
2. Identity of indiscernibles: \(D_{\mathrm{KL}}(P\|Q)=0\) exactly when \(P=Q\) almost everywhere.
3. Asymmetry: generally, \(D_{\mathrm{KL}}(P\|Q)\neq D_{\mathrm{KL}}(Q\|P)\).
4. Not a metric: it is not symmetric and does not obey the triangle inequality.
5. Support sensitivity: if \(P(x)>0\) where \(Q(x)=0\), then \(D_{\mathrm{KL}}(P\|Q)=+\infty\).

Forward KL vs. Reverse KL

With \(P\) as the target and \(Q\) as the approximation:

Because the expectation is taken under different distributions, minimizing these objectives can lead to different behavior.

Forward KL: mode-covering behavior

In \(D_{\mathrm{KL}}(P\|Q)\), regions where \(P(x)\) is substantial but \(Q(x)\) is tiny are heavily penalized. Thus, an approximation often spreads out to cover all important modes of \(P\).

Reverse KL: mode-seeking behavior

In \(D_{\mathrm{KL}}(Q\|P)\), regions where \(Q(x)\) places mass but \(P(x)\) is tiny are heavily penalized. Thus, an approximation often concentrates on one high-density mode rather than covering low-density gaps.

Qualitative illustrations of mode-covering and mode-seeking behavior.

Exact Computation Can Be Expensive

For a large discrete space,

requires summing over every possible \(x\). In a language model, the relevant space can include a very large vocabulary at every token position, or even the much larger space of complete sequences.

When we can sample from \(P\) and evaluate both \(P(x)\) and \(Q(x)\), Monte Carlo estimation avoids summing over the whole space.

Monte Carlo Estimator \(k_1\): Unbiased but High Variance

Draw independent samples

Then

This estimator is unbiased:

However, an individual log-ratio can be negative even though the true KL divergence is nonnegative, and the estimator can have high variance.

Define the likelihood ratio

Then the per-sample estimator is

Estimator \(k_2\): Low Variance but Biased

When \(P\) and \(Q\) are close, \(\log r\) is small. A second-order local approximation motivates

and therefore

This estimator is nonnegative and often has lower variance, but it is generally biased and need not converge to the exact KL value as \(N\to\infty\).

Control Variates: Unbiased and Low Variance

The key identity is

Thus, for any scalar \(\lambda\),

is still unbiased.

The variance-minimizing coefficient is

When $P$ and $Q$ are close, $\log r\approx r-1$, so $\lambda^*\approx 1$.

Setting $\lambda=1$ gives the practical estimator

and

This estimator has three useful properties:

• Unbiased: the added control variate has expectation zero.
• Nonnegative samplewise: $\log r\le r-1$ for every $r>0$.
• Typically lower variance: the first-order fluctuation cancels when $P$ and $Q$ are close.

Why the first-order fluctuation cancels

Write $r=1+\varepsilon$. Around $\varepsilon=0$,

The $O(\varepsilon)$ term vanishes in $k_3$, explaining its reduced variance near $P=Q$.

Numerically Stable Implementation

Let

Then

For numerical stability, compute $e^\Delta-1$ with an expm1 operation:

# x is sampled from P
log_p = log_prob_P(x)
log_q = log_prob_Q(x)
delta = log_q - log_p
kl_per_sample = torch.expm1(delta) - delta
kl_estimate = kl_per_sample.mean()

For direct summation over a finite support, use log probabilities but compute

Comparison of the Three Estimators

Let $r=Q/P$.

Machine-Learning Applications

Classification

For a one-hot target distribution concentrated on class $y$,

Then $H(P)=0$ and

which is the usual cross-entropy or negative log-likelihood loss.

Knowledge distillation

A student model can match a teacher’s soft predictive distribution by minimizing

Unlike a one-hot label, the teacher distribution communicates similarities among classes.

Policy optimization and RLHF

A learned policy $\pi_\theta$ is often regularized toward a reference policy $\pi_{\mathrm{ref}}$:

If samples are drawn from $P=\pi_\theta$ and $Q=\pi_{\mathrm{ref}}$, the $k_3$ estimator uses

Takeaways

The core chain

\[\boxed{ \text{self-information} \;\longrightarrow\; \text{entropy} \;\longrightarrow\; \text{cross-entropy} \;\longrightarrow\; \text{KL divergence} }\]

1. KL divergence is the expected extra information cost of using $Q$ when data comes from $P$.
2. Forward and reverse KL are not interchangeable; their asymmetry produces different approximation behavior.
3. Exact computation may require a large sum or integral.
4. The naive Monte Carlo log-ratio estimator is unbiased but can be noisy.

5. The control-variate estimator

   \[\boxed{ \frac{Q(x)}{P(x)}-1-\log\frac{Q(x)}{P(x)} }\]

   is unbiased, samplewise nonnegative, and often substantially lower variance when the distributions are close.

References

1. Jia-Bin Huang, Fantastic KL Divergence and How to (Actually) Compute It. YouTube video.
2. John Schulman, Approximating KL Divergence (2020). joschu.net/blog/kl-approx.html.

Hinton Markdown

Here is the Markdown with MathJax:

KL Direction in Distilling the Knowledge in a Neural Network

The paper uses the KL divergence direction

This is commonly called the forward KL, where the teacher distribution is the target and the student distribution is optimized to match it.

Soft-Target Distillation Loss

The paper defines the soft-target loss as cross-entropy between the teacher’s softened probability distribution and the student’s softened probability distribution:

where:

• (\(p_i^{(T)}\)) is the teacher probability for class (i) at temperature (T),
• (\(q_i^{(T)}\)) is the student probability for class (i) at temperature (T).

Cross-entropy can be decomposed as

Because the teacher distribution (p) is fixed, its entropy (H(p)) does not depend on the student parameters. Therefore, minimizing

is equivalent to minimizing

[ D_{\mathrm{KL}}\left( p_{\text{teacher}}^{(T)} ;\middle|; q_{\text{student}}^{(T)} \right). ]

Thus, the direction is

[ \boxed{ D_{\mathrm{KL}}\left( p_{\text{teacher}} ;\middle|; q_{\text{student}} \right) } ]

and not

[ D_{\mathrm{KL}}\left( q_{\text{student}} ;\middle|; p_{\text{teacher}} \right). ]

Gradient of the Distillation Loss

The paper gives the gradient with respect to a student logit (z_i) as

This is the gradient produced by cross-entropy with the teacher distribution as the target.

It therefore corresponds to

[ D_{\mathrm{KL}}(p_{\text{teacher}}|q_{\text{student}}), ]

rather than reverse KL.

Complete Training Objective

The paper recommends combining the soft-target loss with ordinary supervised cross-entropy on the hard labels:

[ \mathcal{L} ===========

\alpha T^2 H\left( p_{\text{teacher}}^{(T)}, q_{\text{student}}^{(T)} \right) + (1-\alpha) H\left( y, q_{\text{student}}^{(1)} \right). ]

Here:

• (y) is the ground-truth label distribution,
• (\alpha) controls the balance between soft and hard targets,
• (T) is the distillation temperature,
• (T^2) compensates for the reduction in gradient magnitude caused by increasing the temperature.

The objective therefore combines:

1. A forward-KL-equivalent soft-target loss:

   [ D_{\mathrm{KL}}\left( p_{\text{teacher}}^{(T)} ;\middle|; q_{\text{student}}^{(T)} \right). ]

2. Standard hard-label cross-entropy:

   [ H\left( y, q_{\text{student}}^{(1)} \right). ]

Specialist-Model Aggregation

The paper also uses the same KL direction when combining predictions from a generalist model and specialist models.

The combined distribution (q) is selected by minimizing an objective of the form

[ D_{\mathrm{KL}}(p^g|q) + \sum_m D_{\mathrm{KL}}(p^m|q), ]

where:

• (p^g) is the generalist distribution,
• (p^m) is the distribution from specialist (m),
• (q) is the combined distribution.

Again, the source distributions appear in the first argument and the optimized distribution appears in the second argument.

High-Temperature Interpretation

At sufficiently high temperature, the soft-target cross-entropy becomes approximately equivalent to matching the teacher and student logits with a squared-error objective.

Informally,

[ D_{\mathrm{KL}}\left( p_{\text{teacher}}^{(T)} ;\middle|; q_{\text{student}}^{(T)} \right) \approx \frac{1}{2T^2} \sum_i \left( z_i^{\text{student}} ——————–

z_i^{\text{teacher}} \right)^2, ]

subject to assumptions about centering the logits.

This is a limiting interpretation of the same distillation objective, not a recommendation to replace the loss with reverse KL.

Conclusion

The paper uses

[ \boxed{ \text{teacher};|;\text{student} } ]

or, explicitly,

[ \boxed{ D_{\mathrm{KL}}\left( p_{\text{teacher}} ;\middle|; q_{\text{student}} \right) }. ]

It does not use

[ D_{\mathrm{KL}}\left( q_{\text{student}} ;\middle|; p_{\text{teacher}} \right). ]

Schulman MARKDOWN

Approximating KL Divergence

A Markdown/MathJax adaptation of John Schulman’s March 7, 2020 blog post

Abstract

This note describes Monte Carlo estimators for the Kullback–Leibler (KL) divergence when probabilities or densities can be evaluated pointwise, but the full sum or integral cannot be computed efficiently. It compares three estimators, conventionally called (k_1), (k_2), and (k_3), and explains why (k_3) combines unbiasedness, nonnegativity, and low variance.

Problem setup

For distributions (q) and (p), define

[ D_{\mathrm{KL}}(q|p) = \sum_x q(x)\log\frac{q(x)}{p(x)} = \mathbb{E}_{x\sim q}!\left[\log\frac{q(x)}{p(x)}\right]. ]

The same formulas extend to continuous distributions by replacing the sum with an integral.

Assume that (p(x)) and (q(x)) can be evaluated for any sampled (x), but that an exact sum or integral is impractical. This can happen because the state space is large, no closed form is available, or an implementation stores only sampled log-probabilities. Given samples

[ x_1,x_2,\ldots,x_n \overset{\mathrm{iid}}{\sim} q, ]

we seek a Monte Carlo estimator with small bias and variance.

Throughout, define the likelihood ratio

[ r(x)=\frac{p(x)}{q(x)}. ]

The naive estimator: (k_1)

The direct estimator is

[ k_1(x) = \log\frac{q(x)}{p(x)} = -\log r(x). ]

It is unbiased because

[ \mathbb{E}q[k_1] = D{\mathrm{KL}}(q|p). ]

Its weakness is variance: individual samples may be negative even though KL divergence is nonnegative. When (p) and (q) are close, the true KL can be tiny relative to the fluctuations of (-\log r).

A low-variance biased estimator: (k_2)

A commonly useful alternative is

[ k_2(x) = \frac{1}{2}\left(\log\frac{p(x)}{q(x)}\right)^2 = \frac{1}{2}(\log r(x))^2. ]

Every sample is nonnegative, and its variance is often dramatically smaller than that of (k_1). However, (k_2) is generally biased.

The bias is often small when (p) and (q) are nearby. To see why, recall that an (f)-divergence has the form

[ D_f(p,q) = \mathbb{E}_{x\sim q}!\left[f!\left(\frac{p(x)}{q(x)}\right)\right], ]

where (f) is convex. For a smooth parametric family (p_\theta) near (p_0),

[ D_f(p_0,p_\theta) = \frac{f’‘(1)}{2}\,\theta^\top F\theta

• O(|\theta|^3), ]

where (F) is the Fisher information matrix evaluated at (p_\theta=p_0).

For (k_2),

[ f_2(r)=\frac{1}{2}(\log r)^2, ]

whereas (D_{\mathrm{KL}}(q|p)) corresponds to

[ f_{\mathrm{KL}}(r)=-\log r. ]

Both satisfy (f’‘(1)=1), so their local second-order behavior agrees when (p\approx q).

An unbiased low-variance estimator: (k_3)

A control variate can reduce the variance of (k_1) without changing its mean. Because

[ \begin{aligned} \mathbb{E}_{x\sim q}[r(x)-1] &= \int q(x)\left(\frac{p(x)}{q(x)}-1\right)\,dx
&= \int p(x)\,dx - 1
&= 0, \end{aligned} ]

for any scalar (\lambda), the estimator

[ -\log r+\lambda(r-1) ]

is unbiased for (D_{\mathrm{KL}}(q|p)).

The variance-minimizing (\lambda) depends on unknown moments of (p) and (q). A simple robust choice is (\lambda=1). Since (\log r\le r-1) for all (r>0), this choice is always nonnegative:

[ \boxed{k_3(x)=(r(x)-1)-\log r(x)}. ]

Thus, (k_3) is unbiased, nonnegative sample by sample, and typically has low variance.

Geometrically, (k_3) is the vertical gap between the concave function (\log r) and its tangent at (r=1). Equivalently, it can be viewed through the geometry of Bregman divergences.

General positive estimators for (f)-divergences

The same tangent-line idea works for a convex generator (f). Since (\mathbb{E}_q[r-1]=0), the quantity

[ f(r)-f’(1)(r-1) ]

has the same expectation as (f(r)). Convexity ensures it is nonnegative because a convex function lies above its tangent line.

For the reverse direction,

[ D_{\mathrm{KL}}(p|q) = \mathbb{E}_{x\sim q}[r\log r], ]

so (f(r)=r\log r) and (f’(1)=1). A nonnegative unbiased estimator under (x\sim q) is therefore

[ r\log r-(r-1). ]

In summary, for (x\sim q) and (r=p(x)/q(x)):

[ \begin{aligned} D_{\mathrm{KL}}(q|p)&:\quad (r-1)-\log r,
D_{\mathrm{KL}}(p|q)&:\quad r\log r-(r-1). \end{aligned} ]

Gaussian experiments

Consider unit-variance normal distributions. The following values report the relative bias and relative standard deviation of a single-sample estimator:

[ \text{relative bias} = \frac{\mathbb{E}[k]-D_{\mathrm{KL}}}{D_{\mathrm{KL}}}, \qquad \text{relative standard deviation} = \frac{\operatorname{Std}(k)}{D_{\mathrm{KL}}}. ]

(q=\mathcal{N}(0,1)) and (p=\mathcal{N}(0.1,1))

The true KL divergence is (0.005).

Here (k_2) has only about (0.2\%) relative bias, while both (k_2) and (k_3) have far lower variance than (k_1).

(q=\mathcal{N}(0,1)) and (p=\mathcal{N}(1,1))

The true KL divergence is (0.5).

At the larger divergence, the bias of (k_2) is more visible. The (k_3) estimator remains unbiased and has slightly smaller variance.

PyTorch reproduction

The following independently written script reproduces the comparison. To test the first experiment, set shift = 0.1; for the second, use shift = 1.0.

import torch
import torch.distributions as dist

shift = 0.1
num_samples = 10_000_000

q = dist.Normal(loc=torch.tensor(0.0), scale=torch.tensor(1.0))
p = dist.Normal(loc=torch.tensor(shift), scale=torch.tensor(1.0))

x = q.sample((num_samples,))
true_kl = dist.kl_divergence(q, p)

log_r = p.log_prob(x) - q.log_prob(x)
r = log_r.exp()

k1 = -log_r
k2 = 0.5 * log_r.square()
k3 = (r - 1.0) - log_r

print(f"true KL: {true_kl.item():.6f}")
for name, estimate in [("k1", k1), ("k2", k2), ("k3", k3)]:
    relative_bias = (estimate.mean() - true_kl) / true_kl
    relative_std = estimate.std() / true_kl
    print(
        f"{name}: relative bias={relative_bias.item():.4f}, "
        f"relative std={relative_std.item():.4f}"
    )

Takeaway

The estimators have the following practical tradeoffs:

• (k_1=-\log r) is unbiased but may have very high variance.
• (k_2=\tfrac12(\log r)^2) is nonnegative and low variance, but biased.
• (k_3=(r-1)-\log r) is unbiased, nonnegative, and usually low variance.

Consequently, (k_3) is often the strongest default when estimating (D_{\mathrm{KL}}(q|p)) from samples drawn from (q).

Aidan Luisie

Intuitively Understanding the KL Divergence

Structured mathematical notes based on a video by Aidan Liusie

Source video: youtube.com/watch?v=SxGYPqCgJWM

These notes are a faithful mathematical reconstruction and paraphrase of the video’s central argument, not a word-for-word transcript.

The question

Suppose we have two probability distributions, $P$ and $Q$, with probability mass functions $p(x)$ and $q(x)$. We want a principled way to quantify how different they are.

A useful starting point is this idea:

If two distributions are similar, then they should assign similar probabilities to the same typical sequences of observations.

The Kullback–Leibler divergence arises naturally when we formalize that statement.

Compare the likelihood of one sequence

Draw a sequence of independent observations from $P$:

For a realized sequence $x_{1:n}=(x_1,\ldots,x_n)$, its likelihood under $P$ is

The probability that $Q$ assigns to that same sequence is

If $P$ and $Q$ describe the data similarly, these two likelihoods should be comparable. Their ratio is

A ratio near $1$ means the two distributions assign approximately equal likelihood to the sequence. A large ratio means the sequence looks much more plausible under $P$ than under $Q$.

Take logarithms and normalize

Products of many probabilities quickly become extremely small. Taking a logarithm converts the product into a sum:

Because this quantity grows with the sequence length, divide by $n$:

Thus, $R_n$ is the average log-likelihood advantage that $P$ has over $Q$ on samples generated by $P$.

The large-sample limit

Since the observations are drawn independently from $P$, the law of large numbers gives

This expectation is precisely the KL divergence from $P$ to $Q$:

For a discrete sample space,

For continuous densities,

whenever the integral is well defined.

Core interpretation

The KL divergence measures the asymptotic, per-observation log-likelihood gap between the true distribution $P$ and an alternative model $Q$:

Equivalently, for a long sequence generated from $P$,

So if $D_{\mathrm{KL}}(P|Q)$ is large, typical data from $P$ become exponentially more likely under $P$ than under $Q$ as the sample size grows.

Why the direction matters

KL divergence is generally asymmetric:

This is not an accident. In $D_{\mathrm{KL}}(P|Q)$, the observations are imagined to come from $P$, so the mismatch is averaged using $p(x)$ as the weight:

Reversing the arguments asks a different question: how well does $P$ explain sequences that are typical under $Q$?

Therefore, KL divergence is a divergence, not a true metric distance.

Immediate consequences

Identical distributions

If $P=Q$, then $p(x)/q(x)=1$ wherever $p(x)>0$, so

Non-negativity

KL divergence satisfies Gibbs’ inequality:

with equality exactly when $P$ and $Q$ agree almost everywhere.

Missing support

If some event can occur under $P$ but is assigned zero probability under $Q$, then

This matches the likelihood interpretation: observing such an event rules out $Q$ entirely.

Units

Using the natural logarithm measures KL divergence in nats. Using $\log_2$ measures it in bits.

A small Bernoulli example

Let

and compare it with a fair Bernoulli model

Then

For a typical sequence of $n$ observations from $P$, the likelihood ratio therefore grows approximately like

Even a modest per-sample divergence compounds strongly over a long sequence.

Connection to cross-entropy and maximum likelihood

Expanding the logarithm gives

where

is the entropy of $P$, and

is the cross-entropy of $P$ relative to $Q$.

When $P$ is fixed and $Q=Q_\theta$ is a model, minimizing

is equivalent to minimizing cross-entropy, or equivalently maximizing the expected log-likelihood

This is why KL divergence appears naturally throughout statistics and machine learning.

One-line summary

KL divergence is the average per-sample log-likelihood advantage of $P$ over $Q$ on data generated from $P$.

References

1. Aidan Liusie, Intuitively Understanding the KL Divergence, YouTube video: youtube.com/watch?v=SxGYPqCgJWM.
2. Solomon Kullback and Richard A. Leibler, “On Information and Sufficiency,” The Annals of Mathematical Statistics, 22(1):79–86, 1951.

Schulman LATEX

This note describes Monte Carlo estimators for the Kullback–Leibler (KL) divergence when probabilities or densities can be evaluated pointwise, but the full sum or integral cannot be computed efficiently. It compares three estimators, conventionally called $k_1$, $k_2$, and $k_3$, and explains why $k_3$ combines unbiasedness, nonnegativity, and low variance.

Problem setup

For distributions $q$ and $p$, define \begin{equation} \KL(q|p) = \sum_x q(x)\log\frac{q(x)}{p(x)} = \E_{x\sim q}!\left[\log\frac{q(x)}{p(x)}\right]. \end{equation} The same formulas extend to continuous distributions by replacing the sum with an integral.

Assume that $p(x)$ and $q(x)$ can be evaluated for any sampled $x$, but that an exact sum or integral is impractical. This can happen because the state space is large, no closed form is available, or an implementation stores only sampled log-probabilities. Given samples [ x_1,x_2,\ldots,x_n \overset{\mathrm{iid}}{\sim} q, ] we seek a Monte Carlo estimator with small bias and variance.

Throughout, define the likelihood ratio \begin{equation} r(x) = \frac{p(x)}{q(x)}. \end{equation}

The naive estimator: $k_1$

The direct estimator is \begin{equation} k_1(x) = \log\frac{q(x)}{p(x)} = -\log r(x). \end{equation} It is unbiased because \begin{equation} \E_q[k_1] = \KL(q|p). \end{equation} Its weakness is variance: individual samples may be negative even though KL divergence is nonnegative. When $p$ and $q$ are close, the true KL can be tiny relative to the fluctuations of $-\log r$.

\section{A low-variance biased estimator: $k_2$}

A commonly useful alternative is \begin{equation} k_2(x) = \frac{1}{2}\left(\log\frac{p(x)}{q(x)}\right)^2 = \frac{1}{2}(\log r(x))^2. \end{equation} Every sample is nonnegative, and its variance is often dramatically smaller than that of $k_1$. However, $k_2$ is generally biased.

The bias is often small when $p$ and $q$ are nearby. To see why, recall that an $f$-divergence has the form \begin{equation} D_f(p,q) = \E_{x\sim q}!\left[f!\left(\frac{p(x)}{q(x)}\right)\right], \end{equation} where $f$ is convex. For a smooth parametric family $p_\theta$ near $p_0$, \begin{equation} D_f(p_0,p_\theta) = \frac{f’‘(1)}{2}\,\theta^\top F\theta + O(|\theta|^3), \end{equation} where $F$ is the Fisher information matrix evaluated at $p_\theta=p_0$.

For $k_2$, \begin{equation} f_2(r)=\frac{1}{2}(\log r)^2, \end{equation} whereas $\KL(q|p)$ corresponds to \begin{equation} f_{\mathrm{KL}}(r)=-\log r. \end{equation} Both satisfy $f’‘(1)=1$, so their local second-order behavior agrees when $p\approx q$.

An unbiased low-variance estimator: $k_3$

A control variate can reduce the variance of $k_1$ without changing its mean. Because \begin{equation} \E_{x\sim q}[r(x)-1] = \int q(x)\left(\frac{p(x)}{q(x)}-1\right)\,dx = \int p(x)\,dx - 1 = 0, \end{equation} for any scalar $\lambda$ the estimator \begin{equation} -\log r + \lambda(r-1) \end{equation} is unbiased for $\KL(q|p)$.

The variance-minimizing $\lambda$ depends on unknown moments of $p$ and $q$. A simple robust choice is $\lambda=1$. Since $\log r\le r-1$ for all $r>0$, this choice is always nonnegative: \begin{equation} \boxed{ k_3(x)=(r(x)-1)-\log r(x) }. \end{equation} Thus $k_3$ is unbiased, nonnegative sample by sample, and typically has low variance.

Geometrically, $k_3$ is the vertical gap between the concave function $\log r$ and its tangent at $r=1$. Equivalently, it can be viewed through the geometry of Bregman divergences.

\section{General positive estimators for $f$-divergences}

The same tangent-line idea works for a convex generator $f$. Since $\E_q[r-1]=0$, the quantity \begin{equation} f(r)-f’(1)(r-1) \end{equation} has the same expectation as $f(r)$. Convexity ensures it is nonnegative because a convex function lies above its tangent line.

For the reverse direction, \begin{equation} \KL(p|q) = \E_{x\sim q}[r\log r], \end{equation} so $f(r)=r\log r$ and $f’(1)=1$. A nonnegative unbiased estimator under $x\sim q$ is therefore \begin{equation} r\log r-(r-1). \end{equation}

In summary, for $x\sim q$ and $r=p(x)/q(x)$: \begin{align} \KL(q|p)&:\quad (r-1)-\log r,
\KL(p|q)&:\quad r\log r-(r-1). \end{align}

\section{Gaussian experiments}

Consider unit-variance normal distributions. The following values report the relative bias and relative standard deviation of a single-sample estimator: [ \text{relative bias} = \frac{\E[k]-\KL}{\KL}, \qquad \text{relative standard deviation} = \frac{\operatorname{Std}(k)}{\KL}. ]

$q=\mathcal{N}(0,1)$ and $p=\mathcal{N}(0.1,1)$

The true KL divergence is $0.005$. \begin{center} \begin{tabular}{lcc} \toprule Estimator & Relative bias & Relative standard deviation
\midrule $k_1$ & $0$ & $20$
$k_2$ & $0.002$ & $1.42$
$k_3$ & $0$ & $1.42$
\bottomrule \end{tabular} \end{center} Here $k_2$ has only about $0.2\%$ relative bias, while both $k_2$ and $k_3$ have far lower variance than $k_1$.

\subsection{$q=\mathcal{N}(0,1)$ and $p=\mathcal{N}(1,1)$}

The true KL divergence is $0.5$. \begin{center} \begin{tabular}{lcc} \toprule Estimator & Relative bias & Relative standard deviation
\midrule $k_1$ & $0$ & $2.00$
$k_2$ & $0.25$ & $1.73$
$k_3$ & $0$ & $1.70$
\bottomrule \end{tabular} \end{center} At the larger divergence, the bias of $k_2$ is more visible. The $k_3$ estimator remains unbiased and has slightly smaller variance.

\section{PyTorch reproduction}

The following independently written script reproduces the comparison. To test the first experiment, set \texttt{shift = 0.1}; for the second, use \texttt{shift = 1.0}.

import torch
import torch.distributions as dist

shift = 0.1
num_samples = 10_000_000

q = dist.Normal(loc=torch.tensor(0.0), scale=torch.tensor(1.0))
p = dist.Normal(loc=torch.tensor(shift), scale=torch.tensor(1.0))

x = q.sample((num_samples,))
true_kl = dist.kl_divergence(q, p)

log_r = p.log_prob(x) - q.log_prob(x)
r = log_r.exp()

k1 = -log_r
k2 = 0.5 * log_r.square()
k3 = (r - 1.0) - log_r

print(f"true KL: {true_kl.item():.6f}")
for name, estimate in [("k1", k1), ("k2", k2), ("k3", k3)]:
    relative_bias = (estimate.mean() - true_kl) / true_kl
    relative_std = estimate.std() / true_kl
    print(
        f"{name}: relative bias={relative_bias.item():.4f}, "
        f"relative std={relative_std.item():.4f}"
    )

Applications in Distillation

Hinton et al, 2014 uses forward KL from the teacher to the student:

[ D_{\mathrm{KL}}!\left(p_{\text{teacher}} ,|, q_{\text{student}}\right). ]

The paper describes the main distillation loss as cross-entropy with the teacher’s softened probabilities:

[ \mathcal L_{\text{soft}} ========================

-\sum_i p_i^{(T)}\log q_i^{(T)}. ]

Since

[ H(p,q)=H(p)+D_{\mathrm{KL}}(p|q), ]

and the teacher distribution (p) is fixed, minimizing this cross-entropy is equivalent to minimizing

[ D_{\mathrm{KL}}!\left(p_{\text{teacher}}^{(T)} \middle| q_{\text{student}}^{(T)}\right). ]

The paper’s gradient,

[ \frac{\partial C}{\partial z_i} ===============================

\frac{1}{T}(q_i-p_i), ]

also corresponds to teacher-target cross-entropy, not reverse KL.

The complete recommended objective is actually a weighted combination:

[ \mathcal L ==========

\alpha T^2 H!\left(p_{\text{teacher}}^{(T)},q_{\text{student}}^{(T)}\right) + (1-\alpha) H!\left(y,q_{\text{student}}^{(1)}\right), ]

so it combines:

1. Forward-KL-equivalent soft-target cross-entropy.
2. Ordinary hard-label cross-entropy.

The (T^2) factor compensates for the soft-target gradients shrinking approximately as (1/T^2).

The paper also explicitly uses the same direction later for specialist-model aggregation:

[ D_{\mathrm{KL}}(p^{g}|q) + \sum_m D_{\mathrm{KL}}(p^{m}|q), ]

where the generalist/specialist distributions are first and the combined distribution (q) is second.

So, using equation direction rather than potentially ambiguous terminology:

[ \boxed{\text{teacher} ,|, \text{student}} ]

—not

[ D_{\mathrm{KL}}(\text{student}|\text{teacher}). ]

At very high temperature, the paper additionally shows that this becomes approximately squared logit matching, but that is a limiting interpretation rather than a different recommended loss.

Takeaway

The estimators have the following practical tradeoffs: \begin{itemize} \item $k_1=-\log r$ is unbiased but may have very high variance. \item $k_2=\tfrac12(\log r)^2$ is nonnegative and low variance, but biased. \item $k_3=(r-1)-\log r$ is unbiased, nonnegative, and usually low variance. \end{itemize} Consequently, $k_3$ is often the strongest default when estimating $\KL(q|p)$ from samples drawn from $q$.

References

1. Geoffrey Hinton, Oriol Vinyals, Jeff Dean. Distilling the Knowledge in a Neural Network. NIPS 2014 Deep Learning Workshop.

2. John Schulman. Approximating KL Divergence. Blog.

3. Jia-Bin Huang. Fantastic KL Divergence and How to (Actually) Compute It. YouTube.

4. Rethinking Kullback-Leibler Divergence in Knowledge Distillation for Large Language Models [PDF].