KL Divergence Between Gaussians

1. What is KL Divergence?

Imagine you have two probability distributions, $P$ and $Q$. How would you quantify how different they are with a single number? This is the core question that KL divergence answers. It's not a true distance metric in the formal mathematical sense (like Euclidean distance), primarily because it's not symmetric, but it serves a similar purpose.

To build an intuition for how the KL divergence formula comes about, let's watch a segment from the video "The KL Divergence : Data Science Basics". The presenter uses a simple analogy of a kindergarten teacher tracking students' favorite fruits over two years to constructively derive the formula from first principles.

As you saw, the KL divergence from a "reference" or "true" distribution $Q$ to a "model" or "approximating" distribution $P$ is defined as the expectation of the logarithmic difference between the probabilities. In the video, year 1 was the reference distribution $Q$ and year 2 was the current distribution $P$. The final form for discrete distributions is:

For continuous distributions, the sum is replaced by an integral:

An alternative and powerful perspective comes from information theory. KL divergence can be interpreted as the "extra information" required to encode samples from $P$ when using a code optimized for $Q$. It's the difference between the cross-entropy of $P$ and $Q$, and the entropy of $P$ itself.

2. Deriving the KL Divergence for Gaussian Distributions

In machine learning, we frequently approximate complex, unknown distributions with simpler, tractable ones like the Gaussian distribution. This is a core component of variational inference. Therefore, being able to calculate the KL divergence between two Gaussians analytically is extremely useful.

We will now work through the derivation for the KL divergence between two multivariate normal distributions, $P=\mathcal{N}({\mu }_1,{\Sigma }_1)$ and $Q=\mathcal{N}({\mu }_2,{\Sigma }_2)$. The derivation relies on properties of logarithms, expectations, and the trace operator, which should be familiar from your background in applied mathematics.

The following resource, "Kullback-Leibler divergence for the multivariate normal distribution", provides a concise and rigorous proof. I've broken it down into a few steps.

The expression is now an expectation taken with respect to $p(x)$. The next step is to manipulate the terms inside the expectation so that we can apply standard formulas for expectations of quadratic forms. This involves the clever use of the trace operator, $\mathrm{tr}(\cdot )$. A key property used here is that for a scalar $a$, $a=\mathrm{tr}(a)$, which allows us to convert the quadratic forms into a trace, and then use the cyclic property of the trace, $\mathrm{tr}(ABC)=\mathrm{tr}(BCA)$.

Finally, we can resolve the expectations using known results for multivariate normal distributions and simplify the expression to get our closed-form solution.

After these steps, we arrive at the final expression for the KL divergence between two multivariate normal distributions $P\sim \mathcal{N}({\mu }_1,{\Sigma }_1)$ and $Q\sim \mathcal{N}({\mu }_2,{\Sigma }_2)$:

where $n$ is the dimensionality of the vector $x$.

For the simpler univariate case, where $P\sim \mathcal{N}({\mu }_1,{\sigma }_1^2)$ and $Q\sim \mathcal{N}({\mu }_2,{\sigma }_2^2)$, this formula simplifies to:

Take a moment to inspect these formulas. It should be clear that if you swap the roles of $P$ and $Q$ (i.e., swap ${\mu }_1,{\Sigma }_1$ with ${\mu }_2,{\Sigma }_2$), you will get a different result. This brings us to the crucial property of asymmetry.

3. The Asymmetry of KL Divergence and Its Consequences

The fact that $D_{KL}(P||Q)\ne D_{KL}(Q||P)$ is not just a mathematical curiosity; it has profound practical consequences in machine learning. The choice of which distribution comes first determines the behavior of optimization algorithms that rely on minimizing KL divergence.

Let's use $P$ to denote the true, often complex, data distribution and $Q$ to be our simpler, approximate model distribution.

1. Forward KL Divergence: $D_{KL}(P||Q)$
   The expectation is taken over the true data distribution $P$. The expression is ${\mathbb{E}}_{x\sim P}\left[\log \frac{P(x)}{Q(x)}\right]$. For this value to be small, our model $Q$ must assign a non-trivial probability anywhere that $P$ does. If $P(x)>0$ but $Q(x)\to 0$ for some $x$, the ratio $P(x)/Q(x)$ becomes very large, and the log term blows up to infinity. This forces $Q$ to "cover" all the modes of $P$. This is often called mode-covering or mean-seeking behavior.

2. Reverse KL Divergence: $D_{KL}(Q||P)$
   The expectation is taken over our model distribution $Q$. The expression is ${\mathbb{E}}_{x\sim Q}\left[\log \frac{Q(x)}{P(x)}\right]$. For this to be small, $Q$ must avoid assigning probability mass where $P$ has none. If $Q(x)>0$ but $P(x)\to 0$, the term blows up. This encourages $Q$ to find and focus on one of the modes of $P$ where the probability is high, ignoring other modes. This is known as mode-seeking or zero-avoiding behavior.

The following blog post and video provide excellent visual intuition for this.

The video "Fantastic KL Divergence and How to (Actually) Compute It" also has a great segment that animates this exact concept.

The image below, from a blog post by Jessica Stringham, provides a static summary of this effect. The "real" distribution $P$ is a bimodal mixture of two Gaussians (black curve), and the "approx" distribution $Q$ is a single Gaussian (blue curve).

In variational inference, which we will study next, the objective function we derive (the ELBO) is equivalent to minimizing the reverse KL divergence, $D_{KL}(Q||P)$. This is why VAEs are sometimes criticized for producing "average" or blurry samples—their approximate posterior is encouraged to be mode-seeking and may not capture the full variation of the true posterior distribution.

Conclusion

In this introductory lesson, we have established a foundational understanding of KL divergence. The key takeaways are:

• Definition: KL divergence measures the inefficiency or "surprise" of using an approximate distribution $Q$ to represent a true distribution $P$.
• Asymmetry: It is not a true distance metric because $D_{KL}(P||Q)\ne D_{KL}(Q||P)$. This is because the expectation is taken with respect to the first distribution in the argument.
• Gaussian Case: We worked through the derivation for the closed-form expression of KL divergence between two Gaussian distributions, a result that will be crucial in our study of VAEs.
• Optimization Behavior: The asymmetry leads to distinct outcomes in optimization. Minimizing the forward KL, $D_{KL}(P||Q)$, is mode-covering. Minimizing the reverse KL, $D_{KL}(Q||P)$, is mode-seeking.

In our next lesson, we will use what we've learned about KL divergence to derive the Evidence Lower Bound (ELBO). You will see how maximizing the ELBO is equivalent to minimizing the reverse KL divergence between an approximate and true posterior distribution, forming the core optimization objective for all variational methods, including VAEs.