Softmax and the Log-Sum-Exp Trick

A classifier's last layer spits out raw scores — logits — one per class, and they can be any real number: negative, huge, tiny. To train the model or report a confidence level you need probabilities: numbers that are all positive and sum to exactly 1. Softmax is the standard way to get there, but the textbook formula, typed literally into code, crashes on real logits. Fixing that crash without changing the answer is the log-sum-exp trick.

The analogy: converting speed ratings into win odds

Imagine three racehorses with speed ratings 2, 1, and 0.1. You want implied win probabilities, and you want a small rating edge to translate into a disproportionately bigger edge in odds — a horse rated one point faster should be noticeably more likely to win, not just slightly. Exponentiating each rating before comparing does exactly that: it stretches gaps between scores into much bigger gaps between probabilities, while keeping everything positive. Divide each exponentiated rating by the total to get shares that sum to 1, and you have softmax.

The formula

$$\text{softmax}(z_i) = \frac{e^{z_i}}{\sum_{j} e^{z_j}}$$

Here $z_i$ is the raw score for class $i$, $e^{z_i}$ exponentiates it (always positive, and it exaggerates gaps between scores), and dividing by the sum over all classes $j$ rescales the exponentiated scores so they add up to 1. In words: exponentiate every score, then divide each one by the total.

Worked example 1: softmax by hand

Logits $[2, 1, 0.1]$. Exponentiate: $e^2 = 7.389$, $e^1 = 2.718$, $e^{0.1} = 1.105$. Sum $= 11.212$. Divide each by the sum: probabilities $[0.659, 0.242, 0.099]$. The rating gap of 1.9 between the best and worst horse became a probability gap of 56 percentage points — softmax amplified it.

Worked example 2: the overflow, and the fix

Now suppose the logits are $[1000, 999, 998]$ — plausible if a network's outputs are poorly scaled. A computer's floating-point numbers top out around $10^{308}$, but $e^{1000} \approx 10^{434}$ — it overflows to inf. The next class gives another inf. Dividing inf by inf produces NaN: the computation has silently broken.

The fix costs nothing mathematically: subtract the maximum logit, $m = 1000$, from every score before exponentiating.

$$\text{softmax}(z_i) = \frac{e^{z_i - m}}{\sum_j e^{z_j - m}}$$

This is the same function — multiplying top and bottom by $e^{-m}$ doesn't change a ratio — but now the largest exponent is $e^0 = 1$ instead of $e^{1000}$. Redo the arithmetic: $e^0 = 1$, $e^{-1} = 0.368$, $e^{-2} = 0.135$, sum $= 1.503$, giving $[0.665, 0.245, 0.090]$ — nothing overflows, and it is the exact answer you'd get if infinite-precision arithmetic were available. The same shift-then-sum idea, applied to the log of the softmax denominator, is called log-sum-exp: $\text{LSE}(z) = m + \log\sum_j e^{z_j - m}$, and it's what lets cross-entropy loss be computed directly in log-space without ever forming the probabilities at all.

Drag the exponent below and watch how fast $e^x$ leaves any sane numeric range — that runaway growth is exactly what the shift is designed to tame.

What this means in practice

Every deep learning framework's built-in softmax and cross-entropy functions already apply this shift internally — torch.nn.functional.cross_entropy never computes raw softmax probabilities at all, it works in log-space via log-sum-exp for both speed and stability. The practical danger appears when someone reimplements softmax by hand (in a custom loss, in interview code, in a from-scratch project) and skips the shift: it works fine on small toy logits and then silently produces NaN losses the moment the network's outputs drift to large values during training, which is exactly the situation you most need the loss to behave.