Softmax, its gradient, and the shift it cannot see

Turning scores into probabilities

A classifier's last layer produces one number per class. They can be negative, they need not sum to anything, and they are called logits. Softmax turns them into probabilities:

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

Compute one. Logits z = (2, 1, 0):

e^2 = 7.389,  e^1 = 2.718,  e^0 = 1.000,  sum = 11.107
p = (7.389, 2.718, 1.000) / 11.107 = (0.665, 0.245, 0.090)

Three properties follow from the exponential. Every output is positive, because e^x is. The outputs sum to 1, by construction. And larger logits give larger probabilities, monotonically.

Logits are log-probabilities up to a constant

Take the log of the softmax:

log p_i = z_i − log Σ_j e^(z_j)

The second term is the same for every class. So the logit is the log-probability, shifted by a shared constant. Differences of logits are therefore logs of probability ratios: z_1 − z_2 = 1 means p_1 / p_2 = e = 2.718, and in the example p_1 / p_2 = 0.665 / 0.245 = 2.72. A logit gap of 4.6 means one class is a hundred times more probable than the other. This is the fastest way to read a row of logits.

The shift it cannot see

Add the same number to every logit and nothing changes. (102, 101, 100) gives exactly the same probabilities as (2, 1, 0), because e^(z+c) = e^c × e^z and the e^c cancels between numerator and denominator.

Two consequences. First, a model's logits have no absolute meaning; only their differences do, so a logit of 12 is not "confident" unless you know the others. Second, you may subtract anything you like before exponentiating, and the next lesson shows why every library subtracts the maximum.

Temperature

Dividing the logits by a temperature T before softmax is the same operation as scaling the gaps. With T = 2 the example becomes (1, 0.5, 0) and the output (0.506, 0.307, 0.186): flatter. With T = 0.5 it becomes (4, 2, 0) and the output (0.867, 0.117, 0.016): sharper. As T goes to zero the largest logit takes everything; as it goes to infinity the output goes uniform. The course how-llms-work covers what that does to generated text. The arithmetic is only this: temperature rescales the logit differences, and a difference of d at temperature T is a probability ratio of e^(d/T).

The gradient, and why it is worth seeing

Pair softmax with cross-entropy, so the loss is −log p_correct, and differentiate with respect to the logits. Several cancellations happen and the result is remarkably clean:

∂L / ∂z_i = p_i − y_i

where y is the one-hot target. For the example, with the correct class first:

gradient = (0.665 − 1, 0.245 − 0, 0.090 − 0) = (−0.335, 0.245, 0.090)

The correct logit is pushed up by 0.335, the others are pushed down by their own probability, and the three numbers sum to zero, as they always do. The size of the push is the size of the mistake. If the model had assigned 0.01 to the correct class the push would be 0.99, nearly the maximum; if it had assigned 0.999 the push would be 0.001. Cross-entropy trains hard on what it gets wrong and barely at all on what it already knows. That is the mechanism behind the shape of most loss curves: fast early, when everything is wrong, and slow late, when the remaining errors are few and the gradient on each is small.

It is also why the derivative has the same form as in linear regression, prediction − target, as the previous lesson promised. The exponential in the softmax and the log in the loss undo each other exactly.

Where it fails

Saturation. When one logit exceeds the rest by 20, its probability is 1 − 2 × 10^-9. If that class is wrong, the gradient on the correct logit is essentially 1 and on the wrong one essentially −1, which sounds healthy, but the input to the softmax is far from any region where small changes matter, and it can take many steps to climb back. Confident mistakes are slow to unlearn. Label smoothing and clipping of logits both exist to keep the model out of that region.

Forced choice. Softmax insists the classes are exclusive and exhaustive. Add an irrelevant fourth class and every other probability shrinks to make room. If an example can belong to several classes at once, or to none, the right tool is one sigmoid per class, not a softmax.

Five lines of NumPy

import numpy as np
def softmax(z):
    z = z - z.max()               # the shift it cannot see, used on purpose
    e = np.exp(z)
    return e / e.sum()
p = softmax(np.array([2.0, 1.0, 0.0]))
grad = p - np.array([1, 0, 0])   # (-0.335, 0.245, 0.090)

Run it once, change a logit, and watch the gradient move. The whole of classification training is that last line, repeated.