Stable softmax

A classifier, a program that assigns a label to an input, gives two scores, both equal to 710, and the softmax of the pair, the rule that turns scores into probabilities, should be one half and one half. Run the textbook formula in binary64 and the exponential of 710, the value $e^{710}$, overflows, exceeding the largest value the format holds: that value is $1.7976931348623157\times 10^{308}$, whose logarithm is 709.78, since $308\ln 10+\ln 1.797=709.78$ [S:66:72]. The exponential returns infinity, the sum is infinity, and each ratio is infinity divided by infinity, which prints as nan; no exception is raised by default, so the program continues and returns two "probabilities" that are not numbers.

How do you compute a probability without overflowing? Subtract the largest score from every score before exponentiating. The exponential is the function that raises $e$ to its argument. The subtraction moves the largest exponent, the input of the exponential rather than the exponent field of a float, to zero and leaves every ratio unchanged, because the same factor appears in the numerator and the denominator [S:66:73]. The answer is the same probability, and only the intermediate numbers move. !The exponential is representable only while the logit stays inside a band that the format sets. The inner band is 32-bit, about $-90$ to $90$; the outer band is 64-bit, about $-710$ to $710$, and its upper edge is the logarithm of the largest finite number, 709.78. Both scores are 710, so the unshifted form asks for $e^{710}$ and the format has nothing there. Subtracting the maximum puts both logits at zero, where the exponential is exactly one and the answer is one half and one half. The two marks for 710 and 709.78 coincide at this scale, so the label carries the difference. Status: schematic.

A logit is a raw score, the number a model produces before any probability is formed. The softmax turns a list of logits into positive numbers that sum to one, by exponentiating each one and dividing by the total [S:66:72]. For a list of scores x, the stable form subtracts the largest score first.

The failure is a property of the format rather than of the mathematics. In the 32-bit format of the first section the same wall sits at 88.72, the logarithm of $3.40\times 10^{38}$. The interval from −90 to 90 below is that binary32 threshold rounded outward: any list of scores whose largest term lies outside it makes the unshifted result unstable, returning nan or zero instead of the probability [S:69:185]. The mathematics assigns a finite probability to every finite list, and the format runs out of numbers on the way there.

$$\text{softmax}(x{)}_i=\frac{\exp \left(x_i-\max (x)\right)}{{\sum }_j\exp \left(x_j-\max (x)\right)}$$

That line says: shift every score down by the largest one, exponentiate, and divide each result by the sum of them all. The index $j$ runs over every entry, so it is a dummy distinct from the free index $i$. [F01]

Subtracting the maximum is the choice that keeps every shifted logit at or below zero. Every exponential then lies between zero and one, because $e^0=1$ and the exponential falls as its argument falls below zero, and the largest is exactly one [S:66:73]. The denominator then lies between one and the number of logits, and neither the sum nor the ratio can overflow [S:69:185].

The shift is not an approximation: it multiplies the numerator and the denominator by the same number, and a common factor cancels in a ratio. The cost is one pass to find the maximum and one subtraction per logit, which is small next to the exponentiations the formula already performs. This is the form a numerical library implements for the log-sum-exp, the logarithm of the shifted sum [S:72:107].

Worked example. Take the two scores 710 and 710, the pair that killed the naive formula. The largest score is 710, so both shifted scores are zero. Both exponentials are one, the sum is 2, and the two probabilities are 0.5 and 0.5. The subtraction is exact here because a number minus itself is zero, and $\exp (0)=1$ exactly; no rounding occurs anywhere in the shifted form for this input.

The logarithm of the sum is $\log 2=0.6931471805599453$, so the log-sum-exp, the maximum plus the logarithm of the shifted sum, is 710.6931471805599, a number the format holds without strain. The naive form asks for the exponential of 710 before it divides anything, and that is where it dies. The stable form asks for the exponential of zero, which is one, and the large number never enters an exponent. Read the two results together: one fails before the division, and the other never forms a large number at all.

The repair handles overflow and leaves underflow in place, where underflow means a value below the smallest the format holds. A logit far below the maximum has a probability too small for the format, so the numerator rounds to zero and the probability is lost even though the log-sum-exp is finite [S:66:73]. The fix is to keep the logarithm of the probability rather than the probability itself, which is why a loss function, the number training tries to shrink, takes logits and never the softmax output [S:69:186]. That is a second idea, and the cost of skipping it is a screen of nan values, the not-a-number value, in the backward pass, the pass that computes the gradients.

The formal statement. For a vector $x$ of $n$ real numbers, the softmax of entry $i$ is the exponential of $x_i$ divided by the sum of the exponentials of all entries. For any constant added to every entry, the result is unchanged. Choose the constant to be the largest entry. Every shifted entry is then at most zero, every exponential lies between zero and one, the largest is exactly one, and the denominator lies between one and $n$ [S:69:185]. In a format whose largest finite value is $m$, the unshifted formula overflows once some entry exceeds the logarithm of $m$, and the shifted formula cannot overflow at all.

The derivation. The proof is one line. The exponential of a difference is the product of the two exponentials. So the factor from the constant appears once in every numerator and once in the denominator, and cancels [S:66:73]. The choice of the maximum makes the largest exponent zero. The cost is one pass to find the maximum and one subtraction per entry.

The edge case. The shift repairs overflow and leaves underflow in place. An entry far below the maximum gives an exponential smaller than the format's smallest positive number, so the numerator rounds to zero and the probability is lost [S:66:73]. The logarithm of that probability is then minus infinity, and the backward pass produces NaN. The repair is to keep the logarithm and to pass logits into the loss instead of probabilities [S:69:186], which is the log-sum-exp form [S:72:107].

What the core hid. The core now names both widths. One threshold is the logarithm of the largest finite value of binary64, and the other is the interval from −90 to 90 that binary32 sets [S:69:185]. The lesson is that the threshold is a property of the format, not of the softmax: move the same logits from one width to another and the failure point moves with the width.

Where the analogy breaks. "Move the largest exponent to zero" suggests the numbers change. They do not, because the probabilities are identical in exact arithmetic and only the intermediate values move. The phrase also suggests that the largest entry becomes exactly one after the exponential. It does, but only because subtracting a number from itself is exact. A shift by any other value would round, so the exactness of that one entry is a property of the chosen shift and not of the method.

Original sources. The stable form and the argument for it are section 4.1 of Goodfellow [S:66:72][S:66:73]. The bound on the denominator and the merge with the cross-entropy loss are section 4.5.2 of Dive into Deep Learning [S:69:185][S:69:186]. The log-sum-exp identity and its stable implementation are section 4.4.2 of Blondel and Roulet [S:72:107].

An open question. Whether a model should ever materialise probabilities is unsettled. The shifted form keeps the probabilities, and the log form keeps their logarithms. Both are exact, and they differ in what the backward pass can still recover [S:69:186]. At eight bits the question sharpens, because a probability below the smallest value of the format is lost whichever form the forward pass uses. Evidence that would settle it: a training run in a narrow format with the loss computed both ways from the same logits, comparing the gradients at every step. The comparison has to run long enough for a lost gradient to change the parameters, since a single step hides it. It also has to report the fraction of entries that underflow, not only the final loss.

A frontier result. FlashAttention computes the softmax in blocks and never writes the full matrix. It keeps a running maximum and rescales the partial result as each block arrives [S:245:4]. That is the same subtraction applied to a row that is read in pieces, and it is why attention fits in fast memory [S:105:65]. The identity that makes it exact is the one in the formal statement.

What to read next. Read section 4.1 of Goodfellow for the failure and the fix [S:66:72]. Read section 4.5.2 of Dive into Deep Learning for the denominator bound [S:69:185], and section 4.4.2 of Blondel and Roulet for the log-domain identity [S:72:107]. Then read the online-softmax section of FlashAttention-2 to see the shift applied block by block [S:245:4].