Softmax and Logits

The last thing a language model computes before it predicts a token is a vector of raw scores, one per vocabulary entry. These are logits: unbounded real numbers, positive or negative, with no probabilistic meaning on their own. Softmax is the function that turns them into a probability distribution, and it shows up twice in every transformer, inside attention and at the output head, so understanding it pays off everywhere.

The function

For a logit vector z, softmax is:

softmax(z)_i = exp(z_i) / sum_j exp(z_j)

Exponentiate every logit (making them all positive), then divide by the total (making them sum to one). The result is a valid distribution. Two design choices are doing the work. The exponential turns additive differences in logits into multiplicative ratios of probability, so a logit that is 2 larger than another is exp(2) ≈ 7.4 times more likely, regardless of their absolute values. And the normalisation couples every output together: raising one logit lowers every other token's probability, because the denominator is shared.

Why the exponential, specifically

You could normalise logits many ways; softmax uses exp for a reason beyond positivity. It is the distribution that, given the logits as constraints, assumes the least beyond them (the maximum-entropy distribution). It also pairs with cross-entropy loss to give that clean predicted - true gradient (see next-token-prediction-cross-entropy). And it connects to attention: the same softmax normalises the alignment scores in softmax(QK^T / sqrt(d_k))V, which is why attention weights are forced to sum to one and why "spare" attention mass has to be dumped somewhere, the mechanism behind attention sinks (see context-windows-long-context).

The numerical stability trick

Naive softmax overflows. Real logits can reach 30 or 40, and exp(40) is astronomically large in float32. The universal fix is to subtract the maximum logit before exponentiating:

softmax(z)_i = exp(z_i - max(z)) / sum_j exp(z_j - max(z))

Subtracting a constant from every logit leaves the result mathematically identical (the constant cancels top and bottom) but caps the largest exponent at exp(0) = 1, so nothing overflows. Every production softmax does this, and the same max-subtraction idea, applied blockwise, is what lets FlashAttention compute exact softmax attention without ever materialising the full score matrix.

When it falls down

• Saturation kills gradients. When one logit dominates, softmax outputs near 1 and 0, and the gradient through it vanishes. This is why attention scores are scaled by sqrt(d_k) before softmax: without it, large dot products saturate the softmax and training stalls (see attention-mechanism).
• Temperature reshapes it. Dividing logits by a temperature before softmax sharpens or flattens the distribution; this is a decoding knob, not a training one (see sampling-decoding).
• It never outputs exactly zero. Every token keeps some probability, however tiny, which is why decoders add an explicit truncation step (top-k, top-p, min-p) to cut the long tail before sampling.

Further reading

• Goodfellow, Bengio, Courville, 2016, Deep Learning, Ch. 6 (softmax and cross-entropy) - the maximum-likelihood and numerical-stability treatment.
• Milakov and Gimelshein, 2018, Online normalizer calculation for softmax, arXiv:1805.02867 - the streaming softmax that FlashAttention later built on.