Part I · Maths · Chapter 9

Softmax, cross-entropy and KL divergence

Softmax turns a list of raw scores into probabilities. Cross-entropy measures how surprised those probabilities were by the right answer, and it is the loss every language model is trained to reduce. Entropy and KL divergence complete the set.

Why it matters

The last step of a language model is a softmax over its whole vocabulary; its training loss is cross-entropy; the "temperature" setting you see in every model API is a knob inside the softmax; and distillation (chapter 40) trains a small model by minimising a KL divergence. One chapter of maths covers all four.

Softmax: scores to probabilities

A model's final layer produces one raw score, called a logit, per possible answer. Scores can be any number, negative or positive, and they do not add up to anything in particular. Say the scores for cat, dog and bird are 2, 1 and 0.1. To turn them into probabilities:

Make every score positive with exp (chapter 3): e² = 7.389, e¹ = 2.718, e^0.1 = 1.105.
Add them: 7.389 + 2.718 + 1.105 = 11.212.
Divide each by the total: 7.389 / 11.212 = 0.659, 2.718 / 11.212 = 0.242, 1.105 / 11.212 = 0.099.

The results are positive, add up to 1 and keep the order of the scores. That is softmax:

softmax(z)_i = e^z_iΣ_j e^z_j

Because exp grows so fast, softmax exaggerates differences: a score only 1 higher gets e ≈ 2.7 times the weight. The name says what it does: a soft version of "pick the maximum".

A practical trick: subtract the maximum

Scores of 1,000, 999 and 998 would make exp overflow, since e¹⁰⁰⁰ is far too big for a computer. Subtracting the largest score from all of them first (giving 0, −1, −2) leaves the answer unchanged, because the common factor cancels between the top and bottom of the fraction. Every real implementation does this.

Temperature

Divide every score by a number T before the softmax. That number is the temperature.

T = 1: the plain softmax, 0.659 / 0.242 / 0.099.
T = 0.5: the scores double to 4, 2, 0.2, and the distribution sharpens to 0.864 / 0.117 / 0.019. As T approaches 0, all the probability goes to the top score.
T = 2: the scores halve to 1, 0.5, 0.05, and the distribution flattens to 0.502 / 0.304 / 0.194. As T grows, it approaches equal odds.

When a language model samples its next token, low temperature makes it pick the likeliest token nearly every time (predictable, repetitive); high temperature lets unlikely tokens through (varied, sometimes nonsensical).

Cross-entropy: the training loss

The model said cat 0.659, dog 0.242, bird 0.099. The picture really was a cat. How wrong was it? Use the "surprise" from chapter 3: minus the log of the probability given to the right answer.

Right answer cat: −ln 0.659 = 0.417. A modest loss: the model favoured the right answer but was not sure.
Right answer bird: −ln 0.099 = 2.317. A big loss: the model thought bird unlikely.
A perfect model gives the right answer probability 1, and −ln 1 = 0.

loss = −ln p_correctThis is cross-entropy against a one-hot target (1 for the right answer, 0 for the rest). It is also called log loss or negative log-likelihood.

Why the log, rather than, say, 1 − p? Three reasons. Being confidently wrong is punished without limit: −ln 0.001 = 6.9 and −ln 0.000001 = 13.8. Losses of independent examples add up the way their log probabilities do, so the loss of a whole sentence is the sum of the losses of its tokens. And the gradient comes out beautifully simple.

The gradient of softmax plus cross-entropy

Chain the softmax and the loss together and take the slope with respect to each score. After the algebra, which is a good chain-rule exercise, the answer is just

∂ loss∂ z_i = p_i − y_iy_i is 1 for the correct answer and 0 otherwise.

For the cat example: (0.659 − 1, 0.242 − 0, 0.099 − 0) = (−0.341, 0.242, 0.099). Gradient descent subtracts these, so the cat score goes up and the others go down, each in proportion to how wrong it was. Every language model's training starts its backward pass with this vector.

Try it

Scores (0, 0, 0, 0) for four answers. What are the probabilities, and what is the loss whichever answer is correct?

Answer

Each gets e⁰ = 1, total 4, so 0.25 each. The loss is −ln 0.25 = ln 4 = 1.386. An untrained language model with a vocabulary of 50,000 starts near ln 50,000 ≈ 10.8, which is a useful sanity check on a first training step.

Entropy: how uncertain is a distribution?

Entropy is the average surprise you expect from a distribution, if you draw from it and are surprised by −ln of what came up:

H(p) = − Σ_i p_i ln p_i

Worked for (0.659, 0.242, 0.099): 0.659 × 0.417 + 0.242 × 1.418 + 0.099 × 2.317 = 0.275 + 0.343 + 0.229 = 0.847. Measured with ln the unit is the "nat"; with log₂ it is the "bit" (here 1.22 bits).

A fair coin: 1 bit, the most uncertain a two-way choice can be.
A coin that lands heads 90% of the time: 0.47 bits. At 99%: 0.08 bits.
A certain outcome: 0. The most uncertain a k-way choice can be is ln k, when all are equally likely.

Entropy of a model's next-token distribution is a rough signal of how sure it is. It is one ingredient, not a reliable hallucination detector on its own (chapter 36).

Cross-entropy and KL divergence between two distributions

The target need not be one-hot. Suppose a large "teacher" model gives (0.7, 0.2, 0.1) and a small "student" gives (0.5, 0.3, 0.2). Cross-entropy of the student q measured against the teacher p is

H(p, q) = − Σ_i p_i ln q_i = −(0.7 ln 0.5 + 0.2 ln 0.3 + 0.1 ln 0.2) = 0.887

The teacher's own entropy is H(p) = 0.802. The difference, 0.887 − 0.802 = 0.085, is the KL divergence: the extra surprise caused by using q when the truth is p.

KL(p ‖ q) = Σ_i p_i ln p_iq_i = 0.7 ln 1.4 + 0.2 ln 0.667 + 0.1 ln 0.5 = 0.236 − 0.081 − 0.069 = 0.085

KL is zero only when the two distributions are identical, and never negative.
It is not symmetric: KL(q ‖ p) here is 0.092, not 0.085. It is a "divergence", not a distance.
With a one-hot target, the entropy term is 0, so cross-entropy and KL are the same number. Minimising cross-entropy on real text is the same as pulling the model's distribution towards the data's.

KL turns up in two places later: distillation, where the student minimises KL to the teacher's softened distribution, and RLHF and DPO (chapter 19), where a KL penalty stops the tuned model drifting too far from the original.

Softmax behind an API

Every hosted language model ends in a softmax, but most APIs never show you the logits or probabilities: you receive text and structured output. Some APIs also let you set a temperature, and some current models do not accept one at all. Where a system needs a fixed, repeatable answer, it should not lean on low temperature. It can remove the choice from the model entirely, with fixed replies for certain cases and plain calculators for every number. That is a design choice worth noticing: temperature 0 reduces variation but does not make a model correct.

Build it

In chapter 20 you train a tiny language model whose loss is the cross-entropy from this chapter, and sample from its softmax at different temperatures.

Self-check

Scores (3, 3) for two answers. What does softmax give, and what if you add 100 to both?

(0.5, 0.5) both times. Softmax depends only on the differences between scores, which is why subtracting the maximum is safe.

A model gives the right token probability 0.05. What is its cross-entropy loss on that token?

−ln 0.05 = 3.0.

What happens to the next-token distribution as temperature goes towards 0?

It becomes one-hot on the highest-scoring token, so sampling always picks that token ("greedy" decoding).

Why is KL divergence not a distance?

It is not symmetric: KL(p ‖ q) and KL(q ‖ p) usually differ. A distance must be the same in both directions.