KL Divergence From Scratch — Measuring the Gap Between Two Distributions

KL divergence measures the gap between two probability distributions. We build it up from a compression metaphor to the definition, its famous asymmetry, and a numpy implementation — then watch it at work as the regularizer in VAEs and the leash in RLHF.

Training language models to follow instructions with human feedback

We need a number for "how far apart are these two distributions?"

Strip away the details and most of machine learning is one job: take the true distribution $P$ and push a model distribution $Q$ toward it. For a language model, $P$ is the distribution of words humans actually write next and $Q$ is the model's prediction. For a weather app, $P$ is what the sky actually does and $Q$ is the forecast.

You cannot push $Q$ toward $P$ until you can measure how far apart they are — one number, computable, differentiable. The standard ruler is KL divergence (Kullback–Leibler divergence). The name is intimidating; the substance is not. It is simply the average cost of acting on a wrong belief. This article builds it up in the usual order — metaphor, intuition, definition, code — and then follows it into two places where it quietly runs the modern AI stack: the loss function of VAEs and the reward objective of RLHF.

The metaphor: overpaying with the wrong codebook

As we saw in the information theory article, the trick to compressing data is to give short codes to frequent symbols and long codes to rare ones. The ideal code length is set by probability: a symbol with probability $p$ deserves about $\log \frac{1}{p}$ bits.

Now imagine a planning mistake. You believe tomorrow's data will follow distribution $Q$, so you build the codebook that is optimal for $Q$. But the data actually arrives from a different distribution $P$. Symbols you bet on rarely show up; symbols you assigned long codes to keep appearing. Your compressed files are consistently fatter than they had to be.

The average number of extra bits you pay because of that wrong belief is exactly the KL divergence $D_{\mathrm{KL}}(P \| Q)$. If $P$ and $Q$ agree, the surcharge is zero. The more they disagree, the more you overpay. That is the KL way of thinking: measure the gap between distributions as real, incurred cost.

The intuition: averaging the excess surprise

Here is the same idea from a second angle. Information theory measures the "surprise" of an event with probability $p$ as $\log \frac{1}{p}$ — the rarer the event, the bigger the surprise.

Suppose symbol $x$ actually arrives. Believing $Q$, your surprise is $\log \frac{1}{q(x)}$. Someone who knows the true distribution feels only $\log \frac{1}{p(x)}$. The difference, $\log \frac{p(x)}{q(x)}$, is how much extra surprise your wrong belief cost you on this one event. Average that over how events actually occur — weight by $p(x)$ — and you have KL. KL divergence is the expected excess surprise.

The definition: one line

Read aloud: for each symbol $x$, take the log of the ratio between the true probability $p(x)$ and the believed probability $q(x)$, then average it weighted by how often $x$ really occurs. $P$ is the true distribution, $Q$ is the approximation (the model), and for continuous variables the sum simply becomes an integral. Use log base 2 and the answer is in bits; use the natural log and it is in nats.

Tie that back to the metaphor and the sum is the surcharge, which says: a wrong belief about a symbol you meet constantly lands hard on the bill, while the same wrong belief about a symbol you almost never see barely moves it. Weighting by $p(x)$ is what puts that on the meter.

Two properties are worth memorizing. First, KL is never negative, and it equals zero only when $P$ and $Q$ match exactly (Gibbs' inequality) — the minimum qualification for anything calling itself a measure of mismatch. Second, it relates to cross-entropy by $H(P,Q) = H(P) + D_{\mathrm{KL}}(P\|Q)$. Since $H(P)$ depends only on the data, it is a constant during training — so minimizing cross-entropy, the loss you use for classification and LLM pretraining, is exactly minimizing KL divergence. The training loop you run every day has been lowering a KL all along.