I was working on a variational autoencoder, which uses KL Divergence as one of the loss functions to optimize the model. The formula for KL divergence seemed intimidating at first, so I wanted to understand it in detail. I found a very interesting YouTube video by Ritvik Math, and decided to replicate the explanation here for my future reference and for anyone who prefers reading over watching an 18-minute video.
KL divergence, or Kullback-Leibler divergence, measures how one probability distribution differs from another reference distribution. Sometimes, we want to compare two different probability distributions to understand how similar or different they are. This kind of comparison is useful in many areas, such as statistics, machine learning, and data science. For example, we might want to know how the population height distribution in one country compares to that in another—this could reveal differences in nutrition, healthcare, or genetics.
Now, how do we define this metric? Let's start by creating our own metric to quantify the difference, then discuss its limitations, and finally move towards the actual KL divergence metric by addressing those limitations.
The graphs above illustrate the percentage distribution of students’ fruit preferences for the 2024 and 2025 batches. We’ll refer to the 2024 distribution as P distribution and the 2025 distribution as Q distribution. Now, you might be thinking—a student could like more than one fruit, or none at all—so shouldn’t the percentages sometimes exceed or fall short of 100%? To keep things simple, let's assume we live in a world where each student likes exactly one type of fruit—no more, no less. Please don’t overthink it.
Figure 1
Our goal is to define a metric that tells us how similar or different the two distributions shown in Figure 1 are. Before we get into the details, let’s outline the key attributes this metric should have:
Quantification: It should provide a clear and meaningful measure of the similarity or difference between the two distributions.
Asymmetry: The metric should reflect that the difference from the second distribution to the first may not be the same as the reverse. Since we are using the first distribution as a benchmark, the metric should be able to capture directionality—indicating not just how much the distributions differ, but also in which direction. This could be reflected in the sign or scale of the value.
Let’s examine how the values scale for each fruit relative to one another, and then compute the average of those scaling values. In other words, we want to understand how distribution P changes with respect to distribution Q. $$ \text{Divergence} = \frac{1}{3} \left( \frac{P(\text{Orange})}{Q(\text{Orange})} + \frac{P(\text{Mango})}{Q(\text{Mango})} + \frac{P(\text{Apple})}{Q(\text{Apple})} \right) \tag{1} $$ $$ \text{Divergence} = \frac{1}{3} \left( \frac{{50}}{{50}} + \frac{10}{40} + \frac{40}{10} \right) $$ $$ \text{Divergence} = 1.75 $$
As shown in Equation 1, the divergence value quantifies the difference between the two distributions and also captures the asymmetry in the comparison. For example, if we are measuring the divergence of P with respect to Q, the ratio is \( \frac{P}{Q} \); whereas if we compare Q with respect to P, the ratio becomes \( \frac{Q}{P} \).
As we take the average over each discrete category (e.g., Mango, Orange, and Apple), we observe that the value \( \frac{P(\text{Apple})}{Q(\text{Apple})} = 4 \) is significantly higher compared to \( \frac{P(\text{Mango})}{Q(\text{Mango})} = \frac{1}{4} \). This large value of 4 can disproportionately skew the importance of the smaller value \( \frac{1}{4} \).
How can we address this imbalance? One solution is to define a transformation function \( f(x) = y \) such that \( f(4) = k \) and \( f\left(\frac{1}{4}\right) = -k \). Fortunately, we already have such a function: \( f(x) = \log(x) \).
Therefore, let’s update Equation 1 to reflect this transformation and define it as Equation 2.
$$ \text{Divergence} = \frac{1}{3} \sum_{x \in \{\text{Orange},\,\text{Mango},\,\text{Apple}\}} \log \frac{P(x)}{Q(x)} \tag{2} $$
Equation 2 helps address the issue of skewness by applying the logarithm, but can we go further to better handle edge cases related to asymmetry? For instance, although asymmetry is introduced at the category level, the overall divergence remains the same when summed across all categories— regardless of whether we compute it with respect to P or Q.
To tackle this, instead of assigning equal weight (\( \frac{1}{n} \)) to each term like \( \log \frac{P(\text{Mango})}{Q(\text{Mango})} \), we should weigh each term according to the importance of that category. This importance is determined by the value of the reference distribution for the given category.
With this adjustment, Equation 2 can be modified to form Equation 3.
$$ \text{Divergence} = \sum_{x in \{\text{Orange},\,\text{Mango},\,\text{Apple}\}} {P(x)} \log \frac{P(x)}{Q(x)} \tag{3} $$
Equation 3 is, in fact, the KL divergence—though it doesn't appear intimidating, right? That's because we derived it using a discrete distribution, which gives us a simpler summation form. However, if we replace the summation with an integral, we arrive at the exact KL divergence formula, as shown in Equation 4. $$ D_{KL}(P \parallel Q) = \int_{-\infty}^{\infty} P(x) \log \frac{P(x)}{Q(x)} \, dx \tag{4} $$