The naive formula, why it breaks, and the identity that fixes it without changing the answer.
import math scores = [1, 2, 3] exponentials = [math.exp(s) for s in scores] total = sum(exponentials) [e / total for e in exponentials] # -> [0.09, 0.245, 0.665] math.exp(1000) # OverflowError: math range error # The identity: exp(a) / exp(b) == exp(a - c) / exp(b - c) # because the exp(c) factors cancel top and bottom. # # So subtract the largest score from all of them first. The biggest # exponent becomes exp(0) == 1, and nothing can overflow. shifted = [s - max(scores) for s in scores] # -> [-2, -1, 0]
- The shift cancels, so it costs nothing
Every term in the numerator and the denominator picks up the same factor of exp(−c), and a common factor top and bottom cancels exactly. So the probabilities are identical — this is not an approximation or a trade, it is the same number computed through smaller intermediates. Softmax is shift invariant, which is also why a model’s raw scores have no absolute meaning.
- Subtracting the maximum is the specific choice that is always safe
After shifting, the largest exponent is exactly 0 and every other is negative — so every exp() lands in (0, 1] and overflow becomes impossible, whatever the inputs were. Underflow can still happen for very distant scores, and that is harmless: those terms round to 0, which is very nearly their true value, and the denominator always contains at least the 1 from the maximum.
- It fails on confidence, which is why testing misses it
Scores of 1, 2, 3 work perfectly. The failure needs a genuinely confident model — a well-trained network producing large logits — so the naive version passes every small test and then raises on the clearest image in the batch. Bugs that only appear once the model is good are extraordinarily annoying to trace.
Three steps: shift the scores so the largest is zero, exponentiate, then divide each by the total. The first step is the one the naive version skips.
Handle the empty list first, because max() raises on nothing. Then largest = max(scores), and exponentiate (score - largest) for each — so the biggest becomes exp(0), which is 1.
Compute the maximum and the total once each, above the final comprehension. The total can never be zero: the shifted maximum contributes exactly 1, so no guard is needed on the division — which is worth knowing rather than defending against.
import math def softmax(scores): if not scores: return ____ largest = ____(scores) exponentials = [math.exp(score - ____) for score in scores] total = ____(exponentials) return [value / ____ for value in exponentials]
import math
def softmax(scores):
# max() raises on an empty sequence, so this comes first.
if not scores:
return []
# Subtracting the largest score cancels out of the ratio, and makes
# every exponent zero or negative — so exp() cannot overflow.
largest = max(scores)
exponentials = [math.exp(score - largest) for score in scores]
# Always at least 1, because the largest score contributes exp(0).
total = sum(exponentials)
return [value / total for value in exponentials]
Every term is exp of something at most 0, so every term is in (0, 1] — and the maximum’s term is exactly 1. The total is therefore at least 1, always, so the division needs no guard. That is a stronger statement than 'it probably will not happen', and it comes directly from the shift: without it, a list of very negative scores could underflow every term to 0 and the naive version would divide by zero.
exponentials is built once and read twice — for the sum and for the division. Writing it as a single expression would call exp() on every score twice, doubling the most expensive operation in the function. It is the same reasoning as decorate-sort-undecorate: compute the expensive thing once and keep it.
Softmax produces a valid probability distribution, and that does not make it a calibrated one — a network reporting 0.99 is wrong far more than one per cent of the time, because nothing in training asked those numbers to be honest about uncertainty. Treat them as ranked confidence, not as odds you would bet at, unless the model has been calibrated on purpose.
In real code use scipy.special.softmax or your framework’s — torch.softmax, jax.nn.softmax — all of which do this shift internally and work on whole batches along an axis. More importantly: if the next thing you do is take a logarithm, do not compute softmax at all. log_softmax is a separate function for a good reason — log(exp(x)/sum) can be computed without the intermediate probabilities ever existing, which avoids underflow entirely, and every framework’s cross-entropy loss takes raw scores rather than probabilities for exactly that reason. Handing it softmax output is one of the commonest silent mistakes in a first training loop.
- Add a temperature argument dividing the scores before the shift, and see what temperatures above and below 1 do to how confident the output looks.
- Write log_softmax without computing any exponentials larger than 1, and check it against math.log of your softmax on ordinary inputs.
- Feed the naive version [1000, 1001] and then feed it [-1000, -1001], and note that the two failures are different — one raises and the other divides by zero.