Vol. III · Learning Machinery Plate 2 of 10logits, softmax, cross-entropy, loss, gradient, probabilitiesAL

Softmax and Cross-Entropy.

Scores in, probabilities out, and one number that says how wrong the model was — with a gradient so simple it fits on a line.

A language model’s final layer produces one real number per vocabulary entry: a logit. Logits can be any size and any sign. The softmax exponentiates them and divides by their sum, so that the results are positive and add up to one: a probability distribution. Adding the same constant to every logit changes nothing, which is why only differences between logits matter.

Training needs a single number that says how bad a prediction was: the cross-entropy loss, minus the logarithm of the probability the model gave to the token that actually came next. It is zero if that probability was one and grows without bound as it approaches zero. Its gradient with respect to the logits is strikingly simple — the predicted probabilities minus one at the right token — and that is the signal that flows back through the whole network on every step.

Below: four tokens with logits set by the sliders. The bars show the logits, the softmax probabilities and the gradient of the loss with respect to each logit; the curve shows the loss as the logit of the right token moves.

What you will learn

Compute a softmax by hand.
Compute the cross-entropy loss for a given target.
Derive and read the gradient p − y.
Explain why adding a constant to all logits changes nothing.

Workshop

Left: the logits, the softmax probabilities, and the gradient p − y for each of four tokens (the target is highlighted). Right: the loss as the target’s logit varies with the others held fixed; the dot is the current value.

The signal that trains everything

With logits 2, 1, 0 and −1, the softmax gives the first token a probability of 0.644, and if that token was indeed next the loss is 0.440 nats. The gradient on its logit is −0.356: raise it. Every other logit gets a positive gradient equal to its probability: lower it, the more so the more probability it was taking. If instead the last token had been right, the model gave it only 0.032 and the loss is 3.440 nats — the price of confident mistakes is steep, which is exactly what the curve on the right shows: the loss falls only slowly as the right logit rises past the others but climbs steeply as it falls below them. With all four logits equal, every token gets 0.25 and the loss is ln 4 = 1.386.

In a real model this happens over 8 192 or 128 000 tokens at every position of every sequence at once, and the gradient p − y is the only thing that ever reaches the network from the data. Two numerical details matter in practice. The softmax is computed after subtracting the largest logit, which leaves the result unchanged but keeps the exponentials from overflowing; and the loss is computed as a log-sum-exp minus the right logit rather than by taking the log of a probability, which may have rounded to zero. Vol. XI returns to what can still go wrong in low precision.

The softmax is older than neural networks: it is the Boltzmann distribution of statistical mechanics, with logits playing the role of negative energies and the temperature of plate I.5 playing the role of temperature.

Terms

Logit: A raw score for one token; logits are turned into probabilities by the softmax.
Softmax: Exponentiate each score and divide by the sum, giving positive numbers that add to one.
Cross-entropy loss: Minus the log-probability the model gave to the correct token.
Log-sum-exp: The stable way to compute log Σ e^z, by first subtracting the largest z.

Softmax turns logits into probabilities; cross-entropy is −log p of the right token; and its gradient on the logits is simply p − y.