What is KL Divergence? The Math Explained Step by Step

KL Divergence is a number that tells us how different one probability distribution is from another. The bigger the number, the more different they are, and it is zero only when both are exactly the same.

In this blog, we will learn about KL Divergence. We will also see what a probability distribution is, how to build the KL Divergence formula piece by piece, how to compute it step by step with real numbers, its important properties, how it is related to cross-entropy, and where it is used in AI.

We will cover the following:

• What is a Probability Distribution?
• How different are two distributions?
• Building the KL Divergence formula
• Worked example step by step
• Properties of KL Divergence
• KL Divergence and Cross-Entropy
• Where is KL Divergence used?

I am Amit Shekhar, Founder @ Outcome School, I have taught and mentored many developers, and their efforts landed them high-paying tech jobs, helped many tech companies in solving their unique problems, and created many open-source libraries being used by top companies. I am passionate about sharing knowledge through open-source, blogs, and videos.

I teach AI and Machine Learning at Outcome School.

Let's get started.

What is a Probability Distribution?

A probability distribution is a list of all the possible outcomes along with how likely each outcome is. All the chances add up to 1.

The best way to learn this is by taking an example.

Let's say we live in a city where the weather tomorrow can be only one of three things: Sunny, Cloudy, or Rainy. After watching the weather for many years, we know the real chances:

• Sunny: 0.6 (60 percent)
• Cloudy: 0.3 (30 percent)
• Rainy: 0.1 (10 percent)

Here, 0.6 + 0.3 + 0.1 = 1. This list is a probability distribution. We will call it P. P is the truth.

Now, suppose a weather app makes its own guess for the same city:

• Sunny: 0.3
• Cloudy: 0.3
• Rainy: 0.4

Here also, 0.3 + 0.3 + 0.4 = 1. This is another probability distribution. We will call it Q. Q is the guess.

An LLM (Large Language Model) works in the same way. For the next token, it gives a chance to every token it knows, and all those chances add up to 1. That is also a probability distribution. We have a detailed blog on the Softmax Function that explains how a model turns its raw scores into these chances.

How different are two distributions?

Now, the question is: how far is the weather app's guess Q from the truth P?

Just by looking, we can say that the app is wrong about Sunny and Rainy, and right about Cloudy. But "wrong" is not a number. We need a single number that tells us how different Q is from P.

Here comes the KL Divergence into the picture. KL stands for Kullback-Leibler, the names of the two mathematicians who introduced it. Divergence means "how much two things move apart".

We write it as KL(P || Q) and read it as "the KL Divergence of P from Q". Here, P is the truth and Q is the guess. In AI also, P is usually the truth from the data, and Q is the model's guess. The order matters, and we will see why later in this blog.

Building the KL Divergence formula

Let's build the formula piece by piece, one outcome at a time.

Step 1: We compare the two chances for one outcome using a ratio.

For Sunny, the truth says 0.6 and the guess says 0.3. So, P(x) / Q(x) = 0.6 / 0.3 = 2. Means, the truth gives Sunny two times more chance than the guess does.

• If the ratio is 1, both agree.
• If the ratio is bigger than 1, the guess gives too little chance.
• If the ratio is smaller than 1, the guess gives too much chance.

Why a ratio and not a simple difference? Let's say the truth gives an outcome 0.02 and the guess gives it 0.01. The difference is only 0.01, which looks tiny. But the guess is two times off. The ratio 0.02 / 0.01 = 2 captures this, and the difference does not.

Step 2: We take the log of the ratio.

Why the log? Because the log turns "agree" into zero. The log of 1 is 0. So, when P and Q agree on an outcome, that outcome adds nothing. When the ratio is bigger than 1, the log is positive. When the ratio is smaller than 1, the log is negative.

The log also treats "two times more" and "two times less" as equal and opposite. The log of 2 is 0.6931, and the log of 0.5 is -0.6931.

In this blog, we use the natural log, which is written as ln in math and as np.log in NumPy. With the natural log, the unit of KL Divergence is called nats. If we use log base 2, the unit is called bits. The idea stays the same.

Step 3: We weight the log ratio by P(x).

Not every outcome is equally important. Sunny happens 60 percent of the time in reality, and Rainy happens only 10 percent of the time. So, a mistake on Sunny matters more than a mistake on Rainy. That is why we multiply the log ratio by P(x), the true chance of that outcome.

Step 4: We add up the weighted values for all outcomes.

Putting all four steps together, we get the formula:

KL(P || Q) = sum over all x of P(x) * log(P(x) / Q(x))

Here:

• x is one outcome (Sunny, Cloudy, or Rainy).
• P(x) is the true chance of that outcome.
• Q(x) is the guessed chance of that outcome.
• log is the natural log.

In simple words, KL Divergence is the weighted average of the log ratio between the truth and the guess, where the weights are the true chances P(x).

Note: If P(x) is 0 for an outcome, that outcome adds 0, because it is multiplied by 0. But if Q(x) is 0 for an outcome that P says can happen, we divide by 0, and KL Divergence becomes infinite. Means, a guess that says "this can never happen" for something that does happen gets the biggest possible penalty.

Planning your move into AI?

No matter which tech domain you work in today, with your effort, you can transition to these roles:

• AI Engineer
• Forward Deployed Engineer
• AI Agent Engineer
• Harness Engineer
• Voice AI Engineer
• AI Architect
• LLM Engineer
• LLM Inference Engineer
• Gen AI Engineer
• MLOps Engineer
• ML Engineer
• Deep Learning Engineer

Our AI and Machine Learning Program at Outcome School takes you there.

Worked example step by step

Now, let's compute KL(P || Q) for our weather example.

Let's go row by row.

Sunny: The ratio is 0.6 / 0.3 = 2. The log of 2 is 0.6931. We multiply by P(Sunny) = 0.6 and get 0.4159.

Cloudy: The ratio is 0.3 / 0.3 = 1. The log of 1 is 0. So, Cloudy adds 0. Both distributions agree here.

Rainy: The ratio is 0.1 / 0.4 = 0.25. The log of 0.25 is -1.3863. We multiply by P(Rainy) = 0.1 and get -0.1386.

Finally, we add up the last column.

KL(P || Q) = 0.4159 + 0 + (-0.1386) = 0.2773 nats

Let's see the four steps of the formula as a diagram:

┌──────────────────────────────────┐
│ Sunny: P(x) = 0.6, Q(x) = 0.3    │
└──────────────────────────────────┘
                  ↓
┌──────────────────────────────────┐
│ Step 1: ratio = 0.6 / 0.3 = 2    │
└──────────────────────────────────┘
                  ↓
┌──────────────────────────────────┐
│ Step 2: log(2) = 0.6931          │
└──────────────────────────────────┘
                  ↓
┌──────────────────────────────────┐
│ Step 3: 0.6 * 0.6931 = 0.4159    │
└──────────────────────────────────┘
                  ↓
┌──────────────────────────────────┐
│ Step 4: add all three outcomes   │
│ 0.4159 + 0 + (-0.1386)           │
│ = 0.2773 nats                    │
└──────────────────────────────────┘

Here, we can see Steps 1 to 3 for Sunny. Cloudy and Rainy go through the same three steps, and Step 4 adds the results of all three outcomes.

One more thing to notice: one term is negative. That is fine. The total is still positive, and we will see in the next section that the total is never negative.

We can write the code for this as below:

import numpy as np

P = np.array([0.6, 0.3, 0.1])  # real weather
Q = np.array([0.3, 0.3, 0.4])  # weather app

def kl(p, q):
    return np.sum(p * np.log(p / q))

print(round(kl(P, Q), 4))
print(round(kl(Q, P), 4))
print(round(kl(P, P), 4))

This will print the following:

0.2773
0.3466
0.0

Here, we have:

• p / q divides the chances outcome by outcome.
• np.log takes the natural log of each ratio.
• p * ... weights each log ratio by the true chance.
• np.sum adds everything up.

The first line matches our hand calculation. We will use the other two lines in the next section.

Properties of KL Divergence

Property 1: KL Divergence is always greater than or equal to 0.

Some terms can be negative, like Rainy in our example. But the total never goes below 0.

Let's see this with our example. Both P and Q add up to 1. So, when Q gives too much chance to Rainy, it must take that chance away from some other outcome. Here, Q took it away from Sunny. The log ratio for Rainy is -1.3863, but it is weighted by a small P(Rainy) = 0.1, so it becomes only -0.1386. The log ratio for Sunny is 0.6931, and it is weighted by a big P(Sunny) = 0.6, so it becomes 0.4159. The positive term wins.

The total is never negative for any P and Q. This is a proven result in math known as Gibbs' inequality.

Property 2: KL Divergence is 0 only when the two distributions are exactly the same.

In the code above, kl(P, P) printed 0.0. When P and Q are the same, every ratio is 1, every log is 0, and the total is 0. If the chances differ for even one outcome, the total becomes greater than 0.

Property 3: KL Divergence is not symmetric.

Means, KL(P || Q) is not the same as KL(Q || P). Let's compute KL(Q || P). Now, Q takes the place of P in the formula:

KL(Q || P) = -0.2079 + 0 + 0.5545 = 0.3466 nats

So, KL(P || Q) = 0.2773 and KL(Q || P) = 0.3466. They are different.

Why? Because the weights change. In KL(Q || P), the weights come from Q. Q gives 0.4 to Rainy, while P gives Rainy only 0.1. So, the big mistake on Rainy gets a bigger weight, and the total becomes bigger.

This is why KL Divergence is called a "divergence" and not a "distance". A distance from A to B is the same as from B to A. KL Divergence does not follow this rule. So, we must always be careful about the order.

A quick note for you

No matter which tech domain you work in, get familiar with these topics:

• LLM
• RAG
• MCP
• Agent
• Fine-tuning
• Quantization

We put it all together in one video:

AI Engineering Explained: LLM, RAG, MCP, Agent, Fine-Tuning, and Quantization

No need to stop reading - bookmark it and watch later when you get time. Future you will thank you.

Now, let's get back to the topic.

KL Divergence and Cross-Entropy

Before jumping into the relation, we must know two terms. And for both, we need the idea of surprise.

The surprise of an outcome is -log(chance of that outcome). A rare outcome gives a big surprise. For example, Rainy with chance 0.1 gives -log(0.1) = 2.3026, and Sunny with chance 0.6 gives only -log(0.6) = 0.5108.

Entropy of P is the average surprise when outcomes come from P and we also use P to measure the surprise.

Entropy H(P) = - sum over all x of P(x) * log(P(x))

Cross-entropy of P and Q is the average surprise when outcomes come from P, but we use Q to measure the surprise.

Cross-entropy H(P, Q) = - sum over all x of P(x) * log(Q(x))

Now, here is the relation:

KL(P || Q) = H(P, Q) - H(P)

This comes directly from the KL formula. The log of a ratio is the difference of two logs: log(P(x) / Q(x)) = log(P(x)) - log(Q(x)). So, the KL formula splits into two parts. The part with log(Q(x)) gives the cross-entropy, and the part with log(P(x)) gives the entropy.

Means, KL Divergence is the extra surprise we get because we used the wrong distribution Q instead of the true distribution P.

Let's verify this with our numbers.

• Entropy H(P) = 0.3065 + 0.3612 + 0.2303 = about 0.898
• Cross-entropy H(P, Q) = 0.7224 + 0.3612 + 0.0916 = about 1.175
• H(P, Q) - H(P) = 1.175 - 0.898 = 0.277

This matches KL(P || Q) = 0.2773 that we computed earlier.

This relation is very useful in training. When we train a classifier or an LLM, the true distribution P comes from the data, so H(P) is fixed. It does not change when the model changes. So, making the cross-entropy loss smaller is the same as making the KL Divergence smaller. This is why minimizing cross-entropy loss pushes the model's guess Q closer to the truth P.

For example, when the correct answer is one fixed label, P gives 1 to that label and 0 to all others. In that case, H(P) = 0, and the cross-entropy loss is exactly equal to the KL Divergence.

To master Loss Functions, Logits, and Cross-Entropy, check out our AI and Machine Learning Program at Outcome School.

Where is KL Divergence used?

Now, let's see where KL Divergence is used in AI.

RLHF and PPO: the KL penalty

In RLHF (Reinforcement Learning from Human Feedback), we improve an LLM using a reward model. The reward model gives a score to each answer. The LLM learns to get higher scores.

But, the LLM can find tricks to get a high score while writing strange or broken text. This is called reward hacking. To stop this, we keep a frozen copy of the model from before this training started, called the reference model. Then, for each answer, we compare the chances that the new model gives to the next tokens with the chances that the reference model gives to the same tokens. KL Divergence tells us how far the new model has moved away from the reference model.

Final reward = reward model score - beta * KL(new model || reference model)

Here, beta is a small number, like 0.1, that controls how strong the penalty is. The new model comes first in the KL Divergence because the answers are written by the new model, so its chances act as the weights.

If the new model moves too far from the reference model, the KL Divergence becomes big, and the final reward goes down. So, the model learns to get better scores while staying close to the reference model.

PPO (Proximal Policy Optimization) is the algorithm that is commonly used for this training.

Let's see this as a diagram with beta = 0.1:

┌──────────────────────────────────┐
│ New model writes an answer       │
└──────────────────────────────────┘
                  ↓
┌──────────────────────────────────┐
│ Reward model gives score = 8.0   │
└──────────────────────────────────┘
                  ↓
┌──────────────────────────────────┐
│ Compare token chances:           │
│ new model vs reference model     │
│ KL = 2.0                         │
└──────────────────────────────────┘
                  ↓
┌──────────────────────────────────┐
│ Final reward = 8.0 - 0.1 * 2.0   │
│ = 7.8                            │
└──────────────────────────────────┘
                  ↓
┌──────────────────────────────────┐
│ PPO updates the new model        │
│ using the final reward 7.8       │
└──────────────────────────────────┘

Here, we can see that the reward model gave a score of 8.0, but the new model had moved away from the reference model by a KL Divergence of 2.0. So, the final reward dropped from 8.0 to 7.8.

Stay updated: Subscribe to our newsletter to get our latest AI and Machine Learning blogs straight to your inbox.

Knowledge Distillation

In knowledge distillation, a big model (the teacher) teaches a small model (the student). For each input, the teacher gives a probability distribution over the answers. For example, for a photo of a cat, the teacher can say: cat 0.7, tiger 0.2, dog 0.1. This tells the student much more than just "cat". It also tells the student that a cat looks a bit like a tiger. The student is trained to make its distribution close to the teacher's distribution. The loss we use is KL(teacher || student). The smaller this number, the more the student behaves like the teacher.

VAEs (Variational Autoencoders)

A VAE is a model that learns to compress data, like images, into a small set of numbers and then rebuild the data from them. During training, the VAE adds a KL Divergence term to its loss. This term keeps the distribution of the compressed numbers close to a simple bell-shaped distribution called the standard normal distribution. Because of this, we can later pick random numbers from that simple distribution and generate new images.

Now we must have understood what KL Divergence is, how its formula is built step by step, how to compute it, and where it is used in AI.

Frequently Asked Questions

What is the unit of KL Divergence?

It depends on the log we use. With the natural log, the unit of KL Divergence is called nats. With log base 2, the unit is called bits. The idea stays the same in both cases, only the scale of the number changes.

Can KL Divergence be infinite?

Yes. If the guess Q gives a chance of 0 to an outcome that the truth P says can happen, the formula divides by 0, and KL Divergence becomes infinite. Means, a guess that says "this can never happen" for something that does happen gets the biggest possible penalty.

Why is KL Divergence not a distance?

KL Divergence is not a distance because it is not symmetric. A distance from A to B is the same as from B to A, but KL(P || Q) is usually not equal to KL(Q || P). In our weather example, one is 0.2773 nats and the other is 0.3466 nats, so we must always be careful about the order.

Is minimizing cross-entropy loss the same as minimizing KL Divergence?

Yes, when the true distribution P comes from the data. KL Divergence equals cross-entropy minus the entropy of P, and the entropy of P does not change when the model changes. So, making the cross-entropy loss smaller is the same as making the KL Divergence smaller, which pushes the model's guess Q closer to the truth P.

That's it for now.

Thanks

Amit Shekhar
Founder @ Outcome School

You can connect with me on:

• X
• LinkedIn
• GitHub

Follow Outcome School on:

• X
• YouTube
• LinkedIn
• GitHub

Read all of our high-quality blogs here.