KL Divergence Is Not a Distance

The Question

One characteristic of a distance metric is that distance is symmetric, i.e., the distance between \(a\) and \(b\) must be the same as the distance between \(b\) and \(a\). For KL divergence, this is not the case, so \(D_{\mathrm{KL}}(p\parallel q)=D_{\mathrm{KL}}(q\parallel p)\) is not guaranteed.

Why is \(D_{\mathrm{KL}}\) not symmetric? And then, how could it help us determine how close two distributions are?

Short answer: the order matters because \(D_{\mathrm{KL}}(c\parallel t)\) averages over samples from \(c\), while \(D_{\mathrm{KL}}(t\parallel c)\) averages over samples from \(t\). Each direction therefore asks a different question.

Why I Care

Today, I was reviewing “Step-by-Step Diffusion: An Elementary Tutorial” and found this side note:

“Formally, \(p_T\) is close in KL divergence to \(\mathcal{N}(0,T\sigma^2)\), assuming \(p_0\) has bounded moments.”

It reminded me about this long-overdue post about KL divergence. When working with generative models, we find plenty of situations where we want to compare distributions. For example, the side note in the tutorial refers to the situation where we want to compare a Gaussian distribution to the final distribution of a forward process in diffusion models.

Setup and Definitions

Imagine that you visit your parents at their home every Sunday. To go there, you could drive your own car or take the train. On one hand, traveling by train takes longer, but the duration is very predictable. On the other hand, using the car is usually faster, with the disadvantage of having less certainty about the total trip duration due to unexpected changes in traffic conditions.

The goal is to use KL divergence to check how similar these two distributions are.

We'll assume the commute time is normally distributed as \(\mathcal{N}(\mu,\sigma^2)\), where \(\mu\) is the mean trip duration and \(\sigma\) is the standard deviation, in minutes:

p(\text{car}) \Rightarrow p(c)=\mathcal{N}(\mu_c,\sigma_c^2),

p(\text{train}) \Rightarrow p(t)=\mathcal{N}(\mu_t,\sigma_t^2),

where

\mu_c=35, \qquad \sigma_c^2=3^2,

and

\mu_t=55, \qquad \sigma_t^2=1^2.

For our derivation, we'll also use the log of a normal density:

\mathcal{N}(x;\mu,\sigma^2) = \frac{1}{\sigma\sqrt{2\pi}} e^{-\frac{(x-\mu)^2}{2\sigma^2}}.

\[ \log\mathcal{N}(x;\mu,\sigma^2) = -\log\left(\sigma\sqrt{2\pi}\right) - \frac{(x-\mu)^2}{2\sigma^2}. \] \((1)\)

Minimal Derivation

KL divergence is defined as:

\[ \begin{aligned} D_{\mathrm{KL}}(c\parallel t) &= \mathbb{E}_{x\sim c} \left[\log c(x)-\log t(x)\right] \\ &= \mathbb{E}_{x\sim c} \left[\log\left(\frac{c(x)}{t(x)}\right)\right]. \end{aligned} \] \((2)\)

In words, we draw samples from \(c\) and ask how much more surprising they are under \(t\) than under \(c\). The expectation averages that extra surprise over samples from \(c\).

Luckily for us, we can compute this analytically:

D_{\mathrm{KL}}(c\parallel t) = \mathbb{E}_{x\sim c} \left[\log c(x)-\log t(x)\right].

By equation (1) in the Setup and Definitions section:

\log c(x) = -\log\left(\sigma_c\sqrt{2\pi}\right) - \frac{(x-\mu_c)^2}{2\sigma_c^2},

\log t(x) = -\log\left(\sigma_t\sqrt{2\pi}\right) - \frac{(x-\mu_t)^2}{2\sigma_t^2}.

Replacing these in (2):

\begin{aligned} D_{\mathrm{KL}}(c\parallel t) = \mathbb{E}_{x\sim c}\Bigg[ &-\log\left(\sigma_c\sqrt{2\pi}\right) -\frac{(x-\mu_c)^2}{2\sigma_c^2} \\ &+\log\left(\sigma_t\sqrt{2\pi}\right) +\frac{(x-\mu_t)^2}{2\sigma_t^2} \Bigg]. \end{aligned}

Taking terms out from the expectation:

\begin{aligned} D_{\mathrm{KL}}(c\parallel t) &= \log\left(\frac{\sigma_t}{\sigma_c}\right) -\frac{1}{2\sigma_c^2} \mathbb{E}_{x\sim c}\left[(x-\mu_c)^2\right] \\ &\quad +\frac{1}{2\sigma_t^2} \mathbb{E}_{x\sim c}\left[(x-\mu_t)^2\right]. \end{aligned}

Since \(x\sim c\),

\mathbb{E}_{x\sim c}[x-\mu_c]=0, \qquad \mathbb{E}_{x\sim c}\left[(x-\mu_c)^2\right]=\sigma_c^2.

And

x-\mu_t=(x-\mu_c)+(\mu_c-\mu_t),

such that

\begin{aligned} \mathbb{E}_{x\sim c}\left[(x-\mu_t)^2\right] &= \mathbb{E}_{x\sim c} \left[ \left((x-\mu_c)+(\mu_c-\mu_t)\right)^2 \right] \\ &= \mathbb{E}_{x\sim c} \left[ (x-\mu_c)^2 +2(x-\mu_c)(\mu_c-\mu_t) +(\mu_c-\mu_t)^2 \right] \\ &= \mathbb{E}_{x\sim c}\left[(x-\mu_c)^2\right] +2\mathbb{E}_{x\sim c}[x-\mu_c](\mu_c-\mu_t) +(\mu_c-\mu_t)^2 \\ &= \sigma_c^2+(\mu_c-\mu_t)^2. \end{aligned}

So

\begin{aligned} D_{\mathrm{KL}}(c\parallel t) &= \log\left(\frac{\sigma_t}{\sigma_c}\right) + \frac{\sigma_c^2+(\mu_c-\mu_t)^2}{2\sigma_t^2} - \frac{\sigma_c^2}{2\sigma_c^2} \\ &= \log\left(\frac{\sigma_t}{\sigma_c}\right) + \frac{ \sigma_c^2+(\mu_c-\mu_t)^2-\sigma_t^2 }{2\sigma_t^2}. \end{aligned}

Similarly,

D_{\mathrm{KL}}(t\parallel c) = \log\left(\frac{\sigma_c}{\sigma_t}\right) + \frac{ \sigma_t^2+(\mu_t-\mu_c)^2-\sigma_c^2 }{2\sigma_c^2}.

Substituting the values from our example, we get:

D_{\mathrm{KL}}(c\parallel t)\approx 202.901 \text{ nats}, \qquad D_{\mathrm{KL}}(t\parallel c)\approx 22.876 \text{ nats}.

Since we used the natural logarithm, these values are measured in nats.

Visualizing the Asymmetry

Deriving the two KL values gives us the result, but seeing where each distribution puts its mass makes the asymmetry easier to internalize. In the playground below, move the commute time and switch the KL direction to compare the two densities, the pointwise extra surprise, and each location's contribution to the final KL value.

What I Think I Learned

These results show a huge asymmetry, almost 9 to 1, and we can intuitively see why. Imagine that you said to your mother it took you 35 minutes to arrive there. Your mom would immediately assume you used the car. It seems almost impossible to arrive in 35 minutes using the train. Indeed:

z=\frac{35-55}{1}=-20.

You are 20 train standard deviations away from the train mean.

In the other direction, if you said it took you 55 minutes, you most certainly took the train, but the possibility of having used the car is not as extreme as the reverse case:

z=\frac{55-35}{3}\approx 6.67.

Both outcomes are extremely unlikely under the wrong model, but 20 standard deviations is vastly more extreme than 6.67.

These z-scores do not calculate the KL; they only help us see where the asymmetry comes from. Looking back at the formula, the same 20-minute gap is being judged using the uncertainty of the model on the right.

Under our assumptions, using the train model to explain commute times that come from the car distribution would cause a lot of surprises, much more than using the car model to explain train commute times:

D_{\mathrm{KL}}(c\parallel t) \gg D_{\mathrm{KL}}(t\parallel c).

Which makes sense. We started the story by saying that train travels were very predictable.

I hope that this example from daily life helps us understand why KL divergence is not symmetric (hence not a distance) and how it is used to compare distributions. Notice that if \(c\) and \(t\) were the same, we would get:

D_{\mathrm{KL}}(c\parallel t) = \mathbb{E}_{x\sim c} \left[ \log\left(\frac{c(x)}{c(x)}\right) \right] = \mathbb{E}_{x\sim c}[\log(1)] = 0.

So we should be careful with the word close. KL does not give us a universal boundary between close and far. It tells us how much worse the model on the right is at explaining samples from the distribution on the left.

Open Questions

3Blue1Brown released a video about cross-entropy. At the end, he asks about using KL divergence instead of cross-entropy when training a specific model. What are the consequences?
What other measurements could we use to compare two distributions, and what are their advantages/disadvantages? For example, a metric closer to the distance between the surfaces of the distributions.