Compare Model Distributions with KL Divergence

AI systems often produce probability distributions rather than single answers. A classifier assigns probabilities to classes, a language model assigns probabilities to possible next tokens, and a teacher model can provide a soft target distribution for a smaller student. In all of these cases, developers need a way to ask: how different is one probability distribution from another?

Kullback-Leibler divergence, usually shortened to KL divergence, is one answer. It measures how much a comparison distribution Q differs from a reference distribution P, with the differences weighted by what P considers important.

That last detail is the key. KL divergence is directional, can become infinite, and is not a distance metric. This article builds the calculation from a small example, explains those properties, and shows how to reason about KL divergence when it appears in model training and evaluation.

Start with two probability distributions

Suppose a three-class image classifier predicts cat, dog, or rabbit. A reference distribution says:

P:
cat     0.70
dog     0.20
rabbit  0.10

A second model produces:

Q:
cat     0.60
dog     0.30
rabbit  0.10

Both distributions sum to 1, and they are fairly similar. We want one number that summarizes their mismatch while respecting the fact that errors on high-probability outcomes in P matter more to the calculation.

For discrete distributions, KL divergence from Q to P is conventionally written:

D_KL(P || Q) = sum_i P(i) * ln(P(i) / Q(i))

The notation P || Q is worth reading carefully:

P = reference distribution used for weighting
Q = distribution being compared with P

Using natural logarithms gives the result in nats. Base-2 logarithms give bits. The choice of logarithm changes the unit and numerical scale, not the underlying ordering for a fixed base.

For the example:

D_KL(P || Q)
= 0.70 * ln(0.70 / 0.60)
+ 0.20 * ln(0.20 / 0.30)
+ 0.10 * ln(0.10 / 0.10)

~= 0.70 * 0.1542
 + 0.20 * -0.4055
 + 0

~= 0.0268 nats

Individual terms can be negative, as the dog term is here. The complete KL divergence is nevertheless non-negative when both inputs are valid probability distributions and the quantity is well defined. It is zero exactly when the two distributions agree on all outcomes with probability mass.

Use the expected log-ratio mental model

The formula becomes easier to reason about when written as an expectation:

D_KL(P || Q) = E_{x ~ P}[ln P(x) - ln Q(x)]

Imagine outcomes are drawn according to P. For each outcome, compare how much log probability P gives it with how much Q gives it. Then average those differences according to P.

This explains why P controls what matters. If P(i) is large, outcome i receives a large weight. If P(i) is zero, that outcome contributes nothing to D_KL(P || Q), regardless of the probability that Q assigns to it.

It also connects KL divergence to log loss. Expanding the expression gives:

D_KL(P || Q)
= sum_i P(i) * ln P(i) - sum_i P(i) * ln Q(i)

Define the entropy of P as:

H(P) = -sum_i P(i) * ln P(i)

and the cross-entropy between P and Q as:

H(P, Q) = -sum_i P(i) * ln Q(i)

Then:

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

If P is fixed during optimization, H(P) is constant. Minimizing cross-entropy with respect to Q therefore also minimizes D_KL(P || Q). This is one reason KL divergence and cross-entropy appear so closely related in machine-learning objectives.

Direction changes the question

KL divergence is not symmetric:

D_KL(P || Q) != D_KL(Q || P)   in general

Using the same classifier distributions, reverse the direction:

D_KL(Q || P)
= 0.60 * ln(0.60 / 0.70)
+ 0.30 * ln(0.30 / 0.20)
+ 0.10 * ln(0.10 / 0.10)

~= 0.60 * -0.1542
 + 0.30 * 0.4055
 + 0

~= 0.0291 nats

The values are close in this example because the distributions are close, but they are not identical.

The difference is conceptual, not merely numerical. D_KL(P || Q) asks how poorly Q represents outcomes weighted according to P. Reversing the arguments weights outcomes according to Q instead.

This matters when one distribution assigns probability to regions that the other nearly ignores. A developer should therefore choose the direction from the role of each distribution, not from whichever ordering happens to produce a smaller number.

Zero probabilities create an important boundary condition

Consider:

P:
cat  0.9
dog  0.1

Q:
cat  1.0
dog  0.0

The dog contribution to D_KL(P || Q) contains:

0.1 * ln(0.1 / 0)

Because Q assigns zero probability to an event that has positive probability under P, the divergence is infinite.

This is not a numerical accident. Under the log-loss interpretation, Q says the dog event is impossible, yet P says it can occur. No finite log penalty represents assigning exact zero probability to an event that actually receives positive reference mass.

In floating-point model outputs, probabilities may be extremely small rather than mathematically zero. Implementations also commonly work with logits or log probabilities instead of first materializing probabilities, because direct exponentiation and logarithms can lose numerical precision.

Do not silently add an arbitrary epsilon merely to make the formula finite without understanding what that changes. Smoothing a distribution can be a valid modeling choice, but it changes the distribution and therefore the quantity being measured.

KL divergence is not a distance metric

It is tempting to read a divergence as an ordinary geometric distance. KL divergence does not satisfy the requirements for a metric.

Two properties already show why:

D_KL(P || Q) != D_KL(Q || P)   # not symmetric
D_KL(P || Q) can be infinite    # not necessarily finite

So a statement such as “model A is 0.2 nats away from model B” should not be interpreted like a Euclidean distance.

The absolute magnitude also needs context. A KL value depends on the distributions being compared, their support, and the logarithm base. Comparing KL values is most informative when the evaluation setup is held fixed: the same outcome space, direction, preprocessing, and aggregation procedure.

If an application specifically requires a symmetric, finite comparison, a different divergence may be more appropriate. That is a requirement decision, not a reason to transform KL values informally and continue calling the result KL divergence.

Understand how KL appears in AI training

KL divergence is useful whenever training tries to make one predictive distribution resemble another. The exact roles of P and Q depend on the objective, so inspect the definition rather than assuming every library’s kl_div function uses the same argument convention.

Matching a student to a teacher

Suppose a teacher classifier produces a soft target:

teacher:
cat     0.75
dog     0.20
rabbit  0.05

and a student produces:

student:
cat     0.55
dog     0.35
rabbit  0.10

An objective can penalize divergence between the teacher and student distributions. Unlike a one-hot label that only says cat, the teacher distribution also communicates that dog is considered more plausible than rabbit for this input.

Knowledge-distillation methods often use temperature-scaled distributions and combine a soft-target term with other training terms. Those details affect the objective. The reusable point is simpler: KL divergence can quantify mismatch between two complete predictive distributions.

Keeping an updated policy near a reference policy

In some model-training objectives, a KL-related penalty discourages an updated policy from moving too far from a reference policy. The purpose is not to claim that the reference is perfect. The penalty controls how much distributional change the optimizer is encouraged to make.

The practical trade-off is explicit: a stronger penalty can constrain behavioral change, while a weaker penalty permits larger updates. The effect depends on the objective, data, parameterization, and how the KL term is estimated.

Compute KL carefully in code

For teaching purposes, a direct implementation for small discrete distributions is short:

import math


def kl_divergence(p, q):
    if len(p) != len(q):
        raise ValueError("p and q must have the same length")

    total = 0.0
    for pi, qi in zip(p, q):
        if pi < 0 or qi < 0:
            raise ValueError("probabilities must be non-negative")
        if pi == 0:
            continue
        if qi == 0:
            return math.inf
        total += pi * math.log(pi / qi)

    return total


p = [0.70, 0.20, 0.10]
q = [0.60, 0.30, 0.10]
print(kl_divergence(p, q))  # about 0.0268

This example intentionally exposes the formula. Production numerical code should also validate or deliberately normalize inputs, define tolerances, and prefer stable framework primitives when distributions come from logits. Framework APIs differ in whether they expect probabilities, log probabilities, logits, or targets, and they can use different reduction conventions. Check the API contract before mapping mathematical P and Q onto function arguments.

A common implementation mistake is to pass raw logits into a function that expects log probabilities. Another is to average over every tensor element when the intended metric is a per-example divergence summed over classes and then averaged across examples. Both mistakes can produce plausible-looking numbers with the wrong meaning.

Aggregate divergences at the right level

Suppose a validation set contains N examples and each example has its own reference and comparison distributions. A common evaluation summary is:

mean_KL = (1 / N) * sum_n D_KL(P_n || Q_n)

This gives every example equal weight after summing over its outcome space.

But the mean can hide important variation. A model may match the reference closely on most examples and diverge sharply on a small subset. Depending on the application, inspect the distribution of per-example KL values, useful percentiles, and examples with the largest divergence rather than relying on one aggregate alone.

For variable-length sequence models, aggregation needs even more care. Summing token-level quantities makes longer sequences contribute more simply because they contain more positions. Averaging per token answers a different question from summing per sequence. Neither is universally correct; choose the reduction that matches the decision you are trying to make and document it.

Common mistakes to avoid

The most frequent errors come from losing track of what the formula means.

Treating KL as symmetric. Swapping P and Q changes the weighting and can change the result substantially.

Calling KL a probability of disagreement. A value such as 0.2 nats is not a 20% disagreement rate and is not bounded by 1.

Ignoring support mismatch. If P(i) > 0 where Q(i) = 0, D_KL(P || Q) is infinite.

Comparing values from incompatible setups. Different directions, logarithm bases, vocabularies, class spaces, temperatures, or reductions can make numbers incomparable.

Assuming a lower training KL guarantees better task quality. A model can closely imitate a flawed reference distribution. Distribution matching and downstream correctness are separate evaluation questions.

Forgetting numerical contracts. Probabilities, log probabilities, and logits are different representations. Use the representation expected by the implementation.

When KL divergence is the right tool

KL divergence is a strong fit when you have two probability distributions over the same outcome space and the direction has a meaningful interpretation. Examples include matching a model to soft targets, measuring how a predictive distribution changes relative to a reference, and analyzing distributional differences under a controlled evaluation setup.

A simpler metric is often better when the real question is simpler. If you only care whether a classifier chose the correct class, accuracy or a task-appropriate classification metric answers that question more directly. If you need calibrated decision probabilities, evaluate calibration explicitly. If you need a symmetric notion of distributional difference, choose a measure designed for that requirement.

KL divergence should therefore be selected because its weighting and log-ratio interpretation match the problem, not merely because two arrays of probabilities are available.

Conclusion

KL divergence measures an expected log-probability mismatch: D_KL(P || Q) evaluates Q while weighting outcomes according to P. That mental model explains its most important properties. Direction matters, exact support mismatches can produce infinity, and the result is a divergence rather than an ordinary distance.

When using KL in an AI system, define which distribution is the reference, verify the numerical representation expected by the implementation, and choose an aggregation that matches the decision you care about. With those choices explicit, KL divergence becomes a precise tool for comparing predictive distributions instead of just another opaque training loss.