Why Softmax Subtracts the Max

A short explainer on one specific detail in the softmax function: why implementations subtract the largest score before exponentiating.

Using logits [1, 2, 3] as a worked example, this video shows two paths side by side — computing softmax directly vs. subtracting the max first — and demonstrates that both land on the exact same final probabilities (0.090 / 0.245 / 0.665). The difference only shows up in the intermediate numbers: with larger logits like [1000, 1001, 1002], the direct path overflows to inf/NaN, while the max-subtracted path stays numerically stable and gives identical results either way.

What this shows:

The mathematical identity: subtracting a constant before exponentiating cancels out during normalization The practical reason: it prevents floating-point overflow (float64's ceiling sits around exp(709)) Real numbers computed from the course's own reference implementation, not invented figures

Built with the Brutalist toolkit (brutalist.art) — (https://www.youtube.com/playlist?list=PLG9A_zHUv8A4)

Transcript

A lightly formatted transcript of the video, with timestamps so you can jump to any section on YouTube.

[0:00] Nick Bear Brown, this is Liam in for Bear. SoftMax has one line in it that reads like a typo and it is the only thing keeping the function alive. Here is the stakes. SoftMax turns scores into probabilities by exponentiating them and exponentials grow faster than almost anything else in your program. A 64-bit float tops out a little past e to the 709. So, a logit of 1,000 does not give you a big number, it gives you infinity. Divide infinity by infinity and every probability in the vector becomes NaN. Real models produce logits that large, so this is not a style question. It is whether the thing returns an answer at

[0:37] All. So, we ask for both paths in one function, not a fixed implementation, a flag. If the subtraction really is free, we should be able to switch it off and watch nothing happen. Read it before you trust it. The third line is the entire lesson. z = x - x.max, but only when shift is true. Everything after that line is identical on both paths. Exponentiate total divide. And because the subtraction happens inside the exponential, it comes out as a constant factor in both the numerator and the denominator, so it cancels exactly. It cannot change the answer. What it changes is whether the numerator is still a number by the time you

[1:15] Divide. Run it on logits 1 2 3 at temperature 1. The direct path exponentiates to 2.718, 7.39, and 20.09, totaling 30.19. The max subtracted path gives 0.135, 0.368, and 1.0, totaling 1.50. Completely different intermediates. Now, normalize both. 0.090, 0.245, 0.665. Identical to every digit on screen. The subtraction bought us nothing here and that is exactly what it is supposed to do. So, if it changes nothing, why is it in every implementation ever written? Change one thing. Add 999 to every logit. Softmax is shift invariant, which means the correct answer is exactly the same as before. Same distribution, same three numbers. The only question is which path can still find it. The direct

[2:08] Path overflows on every element. Infinity, infinity, infinity, total infinity. Then infinity divided by infinity gives NaN, NaN, NaN. Note what did not happen, no crash, no exception. NumPy hands back a vector that is no longer a distribution and carries on. The max subtracted path subtracts 1,002 from everything, so the gaps are untouched and the exponentials are the same three numbers we saw a minute ago. 0.090, 0.245, 0.665. The same answer as the small logits because after the shift, it is the same computation. So, here is the teardown. The max subtraction is not an optimization and it is not defensive coding. It is the reason the function has a usable range

[2:52] At all. Mathematically, it cancels and cannot move the answer. Numerically, it is the difference between a probability and a NaN. And the failure mode is the part worth remembering. It is silent. If you only ever test on small logits, that line looks like superstition inherited from somebody else's code right up until production hands you a big one. Your turn. Take softmax steps and sweep the temperature from 0.1 up to 10 on those same three logits and watch the distribution sharpen and then flatten. Then switch the array to float 32 and find the temperature where the smallest probability underflows to zero. That is worth running because it is the same bug

[3:28] From the other direction. We spent this whole video on the top of the range and the bottom of the range will bite you just as quietly. Why softmax subtracts the max.