Cross-entropy and KL, from surprise
A language model reads the cat sat on the and gives the word that actually came next a probability of 0.02. How bad is that? We want an answer that adds up over a dataset, has a known best-possible value, and can be differentiated.
That one requirement gets you all of this lesson: surprise, cross-entropy, entropy, KL, the softmax gradient, and why SFT and RL use KL in opposite directions.
How surprised should the model be?
Score a single prediction. The model gave probability \(q\) to the word that actually occurred, and we want a penalty \(s(q)\). Two requirements are obvious: \(s(1) = 0\), because certain and right costs nothing, and \(s\) grows as \(q\) shrinks.
The third requirement decides everything. Two independent tokens in a row have probability \(q_1 q_2\). Scoring the pair should give the same total as scoring each token and adding the two scores.
Test it with two tokens at \(q = \tfrac13\) each, so the pair has probability \(\tfrac19\).
- \(1-q\): the tokens cost \(\tfrac23 + \tfrac23 = \tfrac43\), but the pair costs \(\tfrac89\).
- \(1/q\): the tokens cost \(3 + 3 = 6\), but the pair costs 9.
- \((1-q)^2\): fails the same way.
Only a logarithm turns products into sums, \(\log(q_1q_2) = \log q_1 + \log q_2\), and the minus sign makes the penalty positive.
With base 2 the unit is bits, and a bit has a literal meaning. Probability \(1/8\) is exactly as unlikely as three fair coins all landing heads, so it scores 3 bits. Drag the point, or use the slider. The hook's 0.02 is where it starts.
Exactly 1 bit each: \(\log_2(0.90/0.45) = \log_2(0.02/0.01) = 1\). Surprise depends on the ratio of probabilities, not the difference, so every halving costs one bit wherever it happens. Check it on the slider: each step of 1 on the log scale is one halving.
That is why log loss punishes confident mistakes. Going from 0.01 to 0.001 costs 3.3 bits, the same as going from 1.0 to 0.1.
From one prediction to a dataset
Now score a model on data. The context the cat sat on the appears 20 times in a corpus. What came next:
The model gives the same distribution \(q\) every time it sees this context: mat \(\tfrac12\), floor \(\tfrac14\), sofa \(\tfrac18\), moon \(\tfrac18\). So each mat occurrence costs 1 bit, each floor 2 bits, and each sofa or moon 3 bits.
\((10\cdot 1 + 5\cdot 2 + 4\cdot 3 + 1\cdot 3)/20 = 35/20 = 1.75\) bits. The average runs over occurrences, not over words, so each word's surprise counts as many times as the word occurs. Averaging the four surprises (2.25) would treat moon as if it were as common as mat.
Here is that computation drawn out. Each column is one occurrence, and its height is that occurrence's surprise. Now change the model and watch which columns move.
Grouping identical words turns counts into frequencies, \(10/20 = 0.50\) and so on:
Name the frequencies \(p(x)\) and you have written down cross-entropy, the model's surprise averaged over how often each outcome really happens:
In training code it is the mean of \(-\log q(y_n)\) over a batch: the same number, computed one example at a time.
What is the best possible model?
You can move the model freely. How low can the average go?
| Model | mat ×10 | floor ×5 | sofa ×4 | moon ×1 | Average |
|---|---|---|---|---|---|
| A | 0.04 | 6.64 | 6.64 | 6.64 | 3.34 |
| B | 1.00 | 2.00 | 2.32 | 4.32 | 1.68 |
| C | 0.51 | 2.47 | 3.18 | 6.64 | 1.84 |
| D | 2.00 | 2.00 | 2.00 | 2.00 | 2.00 |
Surprise per occurrence, in bits.
B. Betting on the most frequent word maximizes accuracy but gives the worst model here: the 10 non-mat occurrences each cost 6.6 bits. Sharpening (C) makes the mat columns cheaper and every other column more expensive, and the other columns lose more than mat gains.
No \(q\) beats \(q = p\). Try it in the chart: you won't get below 1.68 bits. Cross-entropy is minimized by the true frequencies, and the minimum has its own name:
Entropy is the average surprise of a model that knows the true frequencies exactly. It isn't zero, because the data itself is uncertain: even a perfect model can't know whether this particular occurrence will be mat or floor. Entropy is the floor under the training loss.
Why q = p is the minimum (two short derivations)
Lagrange. Minimize \(-\sum_x p(x)\log q(x)\) subject to \(\sum_x q(x) = 1\). Setting the derivative of the Lagrangian to zero gives \(-p(x)/q(x) + \lambda = 0\), so \(q(x) = p(x)/\lambda\). Summing over \(x\) gives \(\lambda = 1\), so \(q = p\).
Jensen (this also proves the gap in section 4 is never negative). Log is concave, so
Equality needs \(q(x)/p(x)\) to be constant, which means \(q = p\).
Where does a worse model pay?
The uniform model averages 2.00 bits and the perfect model 1.68. The uniform model is worse on average. Is it worse on every word?
The uniform model: \(-\log_2 0.25 = 2\) bits against \(-\log_2 0.05 = 4.32\). On moon, and on sofa too, the uniform model beats the perfect one, because it gives them more probability than they deserve. That probability had to come from somewhere: it pays on mat, 2 bits instead of 1, ten times over.
The charts now show the uniform model. Each word has a dashed tick at \(-\log_2 p(x)\), the perfect model's surprise on that word. Hatched means the model pays more than the perfect model; a dashed outline means it pays less.
Average the difference between the model's surprise and the perfect model's over all occurrences. That average has a name:
Individual terms can be negative, like moon's. Their \(p\)-weighted sum never is: \(\mathrm{KL}(p\|q) \ge 0\), and it is 0 only when \(q = p\). Probability is conserved, so every word you over-predict forces you to under-predict another, and the words you under-predict are the ones that occur more.
Two consequences you'll use constantly:
- Minimizing cross-entropy over \(q\) is minimizing \(\mathrm{KL}(p\|q)\). \(H(p)\) doesn't depend on the model, so the two have the same gradients.
- A loss value alone doesn't tell you how good a model is. Loss is \(H(p) + \mathrm{KL}\), and for real data you don't know \(H(p)\).
Perplexity is the same quantity on another scale: \(2^{H(p,q)}\) in bits (or \(e^{\text{loss}}\) in nats), the number of equally likely options that would be as surprising. Uniform over 4 words gives 4. The perfect model here gives \(2^{1.68} = 3.2\).
Bet on the next word
Same question, with money on it. Each round you split 20 chips across the four words, a word is drawn, and you are paid 4 × the share you put on it. Ten chips on the winner doubles your bankroll. Two chips cut it to 40%. Zero chips mean you lose everything.
10 / 5 / 4 / 1. Money multiplies, so take logs: each round adds \(\log_2(4q)\) bits to your log-bankroll.
Your growth per round is 2 bits minus your surprise, on average \(2 - H(p,q)\). You already know which \(q\) minimizes that: \(q = p\). Anyone betting differently falls behind by exactly \(\mathrm{KL}(p\|q)\) bits per round. All-on-mat goes bankrupt the first time another word comes up. This is Kelly betting (Cover & Thomas, ch. 6), and its score is the loss from section 2.
Now play it. Level 1 checks the answer you just gave. Levels 2 and 3 take away what you know.
How does training find q = p?
Real models don't set \(q\) directly. They output logits \(z\), and \(q = \mathrm{softmax}(z)\). From here on we use the natural log (nats), like PyTorch's cross_entropy: 1 bit = 0.693 nats, and the floor here is \(H(p) = 1.165\) nats.
Differentiate the second form of \(L_n\). The first term touches only the observed word, and the log-sum-exp gives back softmax:
For mat the gradient is \(q_{\text{mat}} - 1 < 0\), so descent raises it. Every other logit has gradient \(+q_k\), so each goes down, in proportion to how much probability it currently holds. The words compete for a fixed budget.
The one-hot vectors \(e_{y_n}\) average to the frequencies \(p\):
The gradient is zero exactly when \(q = p\): the minimum from section 3, now reached from the training side.
Step it, starting from a model that knows nothing (all logits 0, so \(q\) is uniform):
Moon. It starts 5× over-predicted (0.25 against 0.05), yet its gradient \(q - p\) is small because both numbers are small, and it carries only 1/20 of the loss. Press "10 steps" twice and compare: mat is within 1% while moon is still about 30% off. The long tail is learned last, and the average loss hides it. That's why you look at loss per slice.
Forward vs reverse KL: which mistakes are expensive
KL isn't symmetric. The difference matters when the model can't represent \(p\) exactly, which in practice is always.
Target \(p\): the lengths of responses a reference model gives to one prompt. It answers in two styles, short (about 20 tokens) and long (about 80). Model \(q\): a single bell curve, which can't have two peaks. Which bump is "closest" to \(p\) depends on which way round you measure.
Forward KL averages \(\ln(p/q)\) over samples from \(p\). If \(q\) is tiny anywhere \(p\) has mass, that log ratio explodes, so \(q\) must cover both styles even if that means putting mass in the valley. For a Gaussian \(q\), the optimum matches \(p\)'s mean and variance: center 50, width 31. Check both directions with the buttons below.
The lower two panels show where each KL collects its cost: the integrand at each length. Minimize reverse KL once starting from μ = 30 and once from μ = 70.
What the buttons show:
| Averaged over | Expensive mistake | Behavior | Where you meet it | |
|---|---|---|---|---|
| Forward KL(p‖q) | data \(p\) | \(q \approx 0\) where data exists | covers every mode, hedges | MLE, pretraining, SFT, distillation on the teacher's distribution |
| Reverse KL(q‖p) | the model's own samples | \(q > 0\) where \(p \approx 0\) | locks onto one mode, can drop others cheaply | variational inference; the KL(πθ‖πref) penalty in RLHF-PPO and GRPO |
Reverse KL on one mode costs about \(\ln 2 \approx 0.69\) here: \(q\) matches one component of \(p\), which holds half of \(p\)'s mass, so \(\ln(q/p) \approx \ln 2\) everywhere \(q\) lives. Forward KL for the same fit is 8 to 24 nats, because the other style's samples are nearly impossible under \(q\).
The 1-D picture only isolates the asymmetry. An LLM's distribution over sequences is far richer than one bell curve, but the pull is the same. SFT on data containing two answer styles pushes the model to keep both. An RL policy on a reverse-KL leash can settle on one style the reference model used and pay little for dropping the other.
Apply it somewhere else
B (0.1020 nats vs 0.1044; the floor \(H(p)\) is 0.0980). The ratio rule holds for one outcome, but cross-entropy averages over both outcomes. On clicks, A and B are off by the same factor in opposite directions: \(\pm 0.02 \ln 2\). The non-clicks decide it. A gives them 0.96 instead of 0.98, which costs \(0.98\ln(0.98/0.96) = +0.020\). B gives them 0.99, which is cheaper than the truth: \(-0.010\). KL: A 0.0064, B 0.0039.
Cross-entropy is minimized at \(q = \) target, and the minimum is the target's entropy: \(-(0.925\ln 0.925 + 3\cdot 0.025\ln 0.025) \approx 0.35\). Predicting 1.0 now costs \(0.025 \cdot \ln(1/0) \cdot 3 = \infty\). The smoothed target is the floor, so the model is actively pushed to stay at 92.5% confidence. That's the whole regularization effect.
\(-\log 0 = \infty\), and one infinite term makes the average infinite. This is forward KL's "cover everything" pressure in its purest form. Softmax never outputs an exact 0, but top-k or nucleus truncation does, so perplexity must be computed on the full untruncated distribution.
Per-token cross-entropy averages over different outcome spaces and different numbers of tokens. Convert to total surprise for the same text: multiply by the token count, then divide by bytes or characters. A tokenizer with coarser tokens has fewer, individually harder predictions.
Write it without a framework
Plain editor, no completion, real Python with numpy, hidden tests. Treat it like a screen: state shapes and edge cases out loud before running. Tab indents, Shift+Tab dedents, ⌘/Ctrl+Enter runs. Your code is saved in this browser.
Exercise 1: stable cross-entropy and its gradient
Exercise 2: KL divergence with the zero conventions
Answer out loud, then check
Give yourself about a minute per question. Short answer first, then the reason. Reveal only after you've answered.
- For i.i.d. data the likelihood is \(\prod_n q(y_n)\); its log is \(\sum_n \log q(y_n)\).
- Mean cross-entropy is \(-\frac1N \sum_n \log q(y_n)\): negative log-likelihood per example.
- Equivalently it is \(H(\hat p) + \mathrm{KL}(\hat p\|q)\) for the empirical distribution \(\hat p\), so MLE minimizes forward KL from data to model.
- Loss \(= H(p) + \mathrm{KL}\). You don't know \(H(p)\), so the plateau could be the entropy floor, or capacity, optimization, or data limits.
- Evidence: does loss still fall with more parameters or data (scaling trend)? Train-val gap? Loss by slice (rare tokens, domains, long contexts)? Does the plateau move with LR schedule changes?
- Decide with downstream evals too: equal loss can hide very different behavior on the slices you care about.
- SFT/MLE: forward \(\mathrm{KL}(p_{\text{data}}\|\pi_\theta)\), an expectation over data. Mode-covering: \(\pi_\theta\) must put mass on everything in the data.
- RL penalty: reverse \(\mathrm{KL}(\pi_\theta\|\pi_{\text{ref}})\), estimated on the policy's own samples. It punishes going where \(\pi_{\text{ref}}\) has little mass and barely punishes dropping some of \(\pi_{\text{ref}}\)'s modes.
- The KL-regularized objective \(\mathbb{E}_{\pi}[r] - \beta\,\mathrm{KL}(\pi\|\pi_{\text{ref}})\) has optimum \(\pi^* \propto \pi_{\text{ref}}\, e^{r/\beta}\). DPO starts from this closed form.
- In DeepSeekMath's GRPO the KL term is subtracted in the objective as \(\beta\,\mathrm{KL}\), not folded into the reward as in InstructGPT-style PPO, where a per-token KL penalty is part of the reward.
- It's estimated per token on sampled completions with \(\frac{\pi_{\text{ref}}}{\pi_\theta} - \ln\frac{\pi_{\text{ref}}}{\pi_\theta} - 1\) (Schulman's "k3"). It's unbiased for \(\mathrm{KL}(\pi_\theta\|\pi_{\text{ref}})\) under samples from \(\pi_\theta\), and every sample is non-negative.
- Implementations differ: some later variants (e.g. DAPO) drop the KL term entirely. Name the variant before arguing about behavior.
- Log loss is a strictly proper scoring rule: in expectation it is minimized only by the true probability (section 3). It measures calibration and ranking together.
- AUC is unchanged by any monotone transform of the scores, so it measures ranking only.
- Insist on calibration when the value is used downstream: auctions (pCTR × bid), blending multiple objectives, thresholds, budget pacing. Check with reliability plots or ECE by slice.
- The loss is cross-entropy to soft targets \(p_T\), which is \(\mathrm{KL}(p_T\|q) + H(p_T)\). The gradient on the logits is \(q - p_T\), so every class gets a signal about how plausible it is, not just the top one.
- Temperature \(T\) raises the teacher's entropy so small probabilities carry signal. The loss is typically scaled by \(T^2\) to keep gradient magnitudes comparable.
- Forward KL makes the student cover the teacher's modes; on-policy or reverse-KL distillation makes it mode-seeking, which can suit a small student that can't cover everything.
Sources: C. E. Shannon, A Mathematical Theory of Communication (1948); Cover & Thomas, Elements of Information Theory, ch. 2; Shao et al., DeepSeekMath (2024), for GRPO; J. Schulman, Approximating KL divergence (2020); Hinton, Vinyals & Dean, Distilling the Knowledge in a Neural Network (2015).