TL;DR
- KL Divergence measures the extra bits wasted due to model imperfection
- Forward KL ($D_{KL}(P | Q)$) is mean-seeking — covers all modes
- Reverse KL ($D_{KL}(Q | P)$) is mode-seeking — focuses on one peak
- Jensen-Shannon Divergence is a symmetric, bounded alternative
- Mutual Information quantifies how much two variables tell us about each other
In my previous post, I talked about Entropy (the uncertainty of the universe) and Cross-Entropy (how surprised our model is). But as I dug deeper into the math for my graduate studies, I realized there was a missing link.
We usually train models to minimize Cross-Entropy, but in research papers—especially regarding Generative AI or Model Distillation—everyone talks about KL Divergence.
It turns out, they are two sides of the same coin. If Cross-Entropy is the total cost of sending a message, KL Divergence is the unnecessary tax you pay for having a bad model.
Huge shoutout to my good friend Sayuru! We were discussing about the order of the inputs in KL Divergence and how it affects the final output. This blog post is a result of our discussion.
The following is my intuitive breakdown of KL Divergence, why the order of inputs matters (a lot!), and how it connects to other fancy terms like Mutual Information.
1. The Intuition: The “Inefficiency Tax”
Let’s go back to our “Surprise” analogy.
- Entropy $H(P)$: The inevitable surprise. The minimum bits needed to describe the truth. You can’t get lower than this.
- Cross-Entropy $H(P, Q)$: The surprise your model experiences.
Kullback-Leibler (KL) Divergence is simply the difference between them. It measures the “extra” bits you wasted because your model $Q$ wasn’t perfect.
\[D_{KL}(P \| Q) = H(P, Q) - H(P)\]Or, rearranged:
\[H(P, Q) = H(P) + D_{KL}(P \| Q)\]The Insight:
Since $H(P)$ (the truth) is fixed and unchangeable, minimizing Cross-Entropy is mathematically identical to minimizing KL Divergence.
- If your model is perfect ($Q = P$), then $D_{KL} = 0$.
- If your model is wrong, $D_{KL} > 0$.
2. The Trap: “Forward” vs. “Reverse” KL (Why Order Matters)
This is the part that confused me the most in lectures.
In standard distance metrics (like Euclidean distance), the distance from A to B is the same as B to A.
\[\text{Distance}(A, B) = \text{Distance}(B, A)\]KL Divergence is NOT symmetric.
\[D_{KL}(P \| Q) \neq D_{KL}(Q \| P)\]Why? Let’s look at the formula:
\[D_{KL}(P \| Q) = \sum P(x) \log \frac{P(x)}{Q(x)}\]The weighting term is $P(x)$. This leads to two very different behaviors depending on which distribution comes first.
Case A: Forward KL ($D_{KL}(P | Q)$) – “Mean Seeking”
Usage: Standard Supervised Learning (Classification).
The Logic: We weight by the Truth ($P$).
Effect: Wherever $P(x) > 0$ (where data exists), $Q(x)$ MUST cover it. If the real data is there and your model predicts zero probability, the loss explodes to infinity.
Result: The model tries to “cover” everything. It becomes a wide, blurry average of all the data modes.
Case B: Reverse KL ($D_{KL}(Q | P)$) – “Mode Seeking”
Usage: Variational Autoencoders (VAEs), Reinforcement Learning.
The Logic: We weight by the Model ($Q$).
Effect: The model only cares about being right where it thinks the data is.
Result: The model finds one peak (mode) of the data and sits there perfectly, ignoring other valid data points. It prefers being “precise but incomplete” rather than “broad but blurry.”
TL;DR on Order:
- Forward KL: “I must cover every single data point, even if I have to be blurry.” (Mean-Seeking)
- Reverse KL: “I will pick one specific data cluster and fit it perfectly, ignoring the rest.” (Mode-Seeking)
3. The Peacemaker: Jensen-Shannon Divergence (JSD)
Since KL is asymmetric (and can explode to infinity), it’s sometimes annoying to work with. Enter Jensen-Shannon Divergence.
JSD is essentially a symmetric and smoothed version of KL.
It works by averaging the two distributions first ($M = \frac{P+Q}{2}$) and then checking how much $P$ and $Q$ diverge from that average.
\[JSD(P \| Q) = \frac{1}{2} D_{KL}(P \| M) + \frac{1}{2} D_{KL}(Q \| M)\]Why care?
This is the metric used in GANs (Generative Adversarial Networks). It’s stable, bounded between 0 and 1 (if using log base 2), and symmetric.
4. Mutual Information (MI): The Intersection
Finally, how does information relate to correlation?
Mutual Information ($I$) measures how much knowing one variable ($X$) tells you about another ($Y$).
- If $X$ and $Y$ are independent, knowing $X$ tells you nothing about $Y$. Their Mutual Information is zero.
Interestingly, MI is just the KL Divergence between the Joint Distribution and the Product of Marginals:
\[I(X; Y) = D_{KL}( P(X, Y) \| P(X)P(Y) )\]- $P(X, Y)$: The reality of how X and Y occur together.
- $P(X)P(Y)$: The theoretical assumption if X and Y were completely unrelated.
So, Mutual Information (MI) is literally: “How far is the reality of their relationship from the assumption that they are unrelated?”
When I was learning Medical Image Processing module, I learned that Mutual Information is considered the Gold Standard algorithm for multi-modal image registration.
KL Divergence in Vision Transformers (ViTs)
If you are working with Vision Transformers, you are actually using these concepts every time you run a forward pass.
1. Attention is a Probability Distribution
The Softmax operation in the Self-Attention mechanism converts raw scores into a probability distribution.
\[\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V\]When we analyze Model Explainability (like checking if a head is attending to the object or the background), we often look at the Entropy of this attention map.
- Low Entropy: The head is very confident (Spiky distribution). It is looking at one specific patch.
- High Entropy: The head is confused (Flat distribution). It is averaging information from everywhere.
2. Knowledge Distillation (DeiT)
In papers like DeiT (Data-efficient Image Transformers), researchers use a “Teacher-Student” setup. The Student model isn’t just trained on the ground truth labels (Cross-Entropy); it is also trained to match the Teacher’s probability distribution.
The loss function they use to align the Student with the Teacher? KL Divergence.
\[\mathcal{L} = (1-\alpha)\mathcal{L}_{CE} + \alpha \cdot D_{KL}(\text{Teacher} \| \text{Student})\]Here, the Student is penalized if its output distribution diverges from the Teacher’s soft labels. This allows the Student to learn not just “This is a cat,” but “This is a cat, but it looks a little bit like a dog,” capturing the rich, dark knowledge hidden in the Teacher’s probabilities.
Understanding these concepts is crucial for anyone working on generative models, knowledge distillation, or model interpretability.