What is a numerically stable softmax, and how do you implement it?

In this article

1. Before you start
2. Why interviewers ask this
3. Numerically stable softmax from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

This page is part of the PyTorch and machine learning (ML) coding round guide, where softmax shows up inside attention, losses and decoding.

Why interviewers ask this

Softmax is the smallest piece of code in a transformer that breaks on real data, so candidates report being asked for it in machine learning coding rounds, often next to scaled dot-product attention, causal masks and normalization layers. One candidate also reported a streaming-entropy prompt that needs the same idea. The question separates people who can write the formula from people who know what a computer does with it.

A weak answer writes exp(x) / sum(exp(x)) and stops. A strong answer says that large inputs overflow to infinity and the division returns not a number (NaN), subtracts the maximum, and says why the result does not change. A Staff-level answer adds the log-softmax for losses, the fully masked row that still returns NaN, and how a fused kernel does the max and the sum in one pass.

Numerically stable softmax from first principles

Numerically stable softmax

A numerically stable softmax is a softmax computed after subtracting the largest input from every input, which leaves the output unchanged but keeps every exponent at zero or below, so the exponentials cannot overflow and the denominator is at least one.

Overflow

Overflow is a floating-point failure in which a calculation produces a number larger than the biggest value the format can store, so the result becomes infinity. Dividing one infinity by another then gives NaN, so a single overflow can spoil a whole row.

Log-sum-exp

Log-sum-exp is the function that returns the natural log of the sum of exponentials of a list, computed by pulling out the maximum so it does not overflow; it is the stable way to get the log of a softmax denominator.

In plain words

Think of a scoreboard where only the gaps between scores matter. Teams with 1000, 1001 and 1002 points stand the same as teams with 0, 1 and 2 points. The small numbers are easier to write and cannot overflow the board. The analogy breaks here. Subtracting the same number keeps softmax exactly the same, because exp turns the gap into a fixed ratio. Multiplying every score is different. The scores 2000, 2002 and 2004 keep the order of 1000, 1001 and 1002, but the softmax changes to 0.0159, 0.1173 and 0.8668.¹

In one precise sentence, softmax is unchanged when the same constant is subtracted from every input, so choosing that constant to be the maximum gives exponents that are at most zero.

A logit is a raw, unnormalized score that a model outputs before softmax turns it into a probability. Softmax of the logits is a list of probabilities, and the shift is applied to the logits.

Step 1: The naive formula overflows

Softmax of an input list x gives, for each position i, the exponential of x_i divided by the sum of the exponentials of all inputs. Computing it directly fails because exp grows very fast. In a 32-bit float, exp overflows to infinity once its input passes about 88.7228, and in a 64-bit float it does so past about 709.7827.¹

Take the inputs 1000, 1001 and 1002. Each exponential is infinity, the sum is infinity, and infinity divided by infinity is NaN, so the naive result is NaN for every position in both 32-bit and 64-bit floats.¹ The failure is not a rare corner: even inputs of 100, 101 and 102 give NaN in 32 bits.¹

Input: a list of logits. Output: NaN. You now know the problem; the next step removes it without changing the mathematics.

Step 2: Subtract the maximum before exponentiating

The stable form subtracts the maximum m of the inputs from every input first, then applies the same formula. The answer is unchanged because each exponential of (x_i minus m) equals exp(x_i) times exp(minus m), and that common factor appears in the numerator and in the sum, so it cancels.

Two things follow, and the second is the one interviewers listen for. Every exponent is now zero or less, so every exponential lies between 0 and 1 and overflow is impossible. And the largest input becomes exactly zero, whose exponential is exactly one, so the sum is at least one and the division can never be by zero.

For the inputs 1000, 1001 and 1002, the maximum is 1002, the shifted inputs are minus 2, minus 1 and 0, and their exponentials are 0.1353, 0.3679 and 1.0.¹

Input: logits. Output: exponentials in (0, 1], with at least one equal to 1.

Step 3: Divide by the sum

The stable softmax is each shifted exponential divided by their sum. For the worked inputs, the sum is 1.5032, so the outputs are 0.09, 0.2447 and 0.6652, and they add to one.¹

Notice what the shift did not fix. The smallest outputs can still be so small that a float rounds them to exactly zero. That is harmless when the outputs are used as attention weights, because a weight that small contributes nothing visible. It matters when you take a log, which is the next step.

For an input of minus 1000, minus 1001 and minus 1002 the naive formula also fails (every exponential is zero, so it computes zero over zero), while the stable version returns 0.6652, 0.2447 and 0.09.¹ The shift handles both ends of the number line.

Input: shifted exponentials. Output: probabilities that sum to one.

Step 4: Take the log with log-sum-exp

The log of a softmax output is x_i minus m minus the log of the sum of the shifted exponentials, and this is called log-softmax. It is the stable way to get log-probabilities, which losses and beam-search scores need.

Taking the log of a stable softmax output is not the same. With the two inputs 0 and minus 200 in a 32-bit float, the second output underflows to exactly 0.0, and its log is minus infinity.¹ Log-softmax returns 0 and minus 200 instead, because it never forms the tiny probability.¹ For the worked inputs the log of the sum is 0.4076, so the log-probabilities are minus 2.4076, minus 1.4076 and minus 0.4076.¹

The same idea is why a cross-entropy loss is best fed logits and not probabilities: PyTorch's CrossEntropyLoss expects unnormalized logits and applies log-softmax internally (PyTorch documentation, read October 2026). Blanchard, Higham and Higham, 2021 study the accuracy of computing log-sum-exp and softmax in floating point.

1 · Naive formula2 · Subtract the max3 · Divide by the sum4 · Take the logReplay

Step-by-step description

1. Exponentials overflow and the division gives NaN.
2. Every exponent is zero or less.
3. The weights add up to one.
4. Log-softmax never forms a tiny probability.

• Subtracting the maximum 1002 gives exponentials 0.1353, 0.3679 and 1.0.
• The sum is 1.5032 and the softmax is 0.09, 0.2447 and 0.6652.
• The log-softmax is minus 2.4076, minus 1.4076 and minus 0.4076.

Worked example

The worked example runs the whole method on the inputs 1000, 1001 and 1002, which are illustrative values chosen so the naive formula fails.¹ The two functions below are the code that was run, in assets/data/numerically-stable-softmax.py; the first takes the maximum along the last axis, subtracts it, exponentiates and divides.

import numpy as np

def stable_softmax(x, axis=-1):
    m = np.max(x, axis=axis, keepdims=True)
    e = np.exp(x - m)
    return e / np.sum(e, axis=axis, keepdims=True)

def log_softmax(x, axis=-1):
    s = x - np.max(x, axis=axis, keepdims=True)
    return s - np.log(np.sum(np.exp(s), axis=axis, keepdims=True))

The code does one reduction for the maximum, one for the sum, and each reduction is linear in the length of the row, so the cost is linear in the number of inputs. The code keeps keepdims=True so the subtraction and the division broadcast across a batch of rows. Run on a batch of two rows, 1000, 1001, 1002 and 0, 1, 2, stable_softmax returns 0.09, 0.2447 and 0.6652 for both rows, and log_softmax returns minus 2.4076, minus 1.4076 and minus 0.4076 for both.¹

All values come from one run of the saved script.¹ The softmax values are the same whether the inputs are 1000 to 1002 or 0 to 2, which is the shift invariance in action.

Text description of the diagram

• The stable softmax of 1000, 1001 and 1002 is 0.09, 0.2447 and 0.6652.
• The outputs add up to one.

Trade-offs and edge cases

Stable softmax costs an extra pass but no accuracy: the maximum is found first, then the sum, then the division. The cost matters inside attention, where rows are long. Milakov and Gimelshein (2018) describe an online softmax normalizer that computes the maximum and the sum in a single pass by rescaling the running sum when a new maximum appears, and report that softmax runs up to 1.3 times faster and a fused softmax with top-k selection up to 5 times faster. On the worked inputs, the running sum goes 1.0, 1.3679 and 1.5032, the same final sum as the three-pass version.¹

The shift also makes the choice of number format matter less. Shifted inputs are at most 0, which is below every overflow threshold of exp: 11.0899 for fp16, 88.7228 for fp32 and 709.7827 for fp64.¹

Text description of the diagram

• Shifted inputs are at most 0, below every overflow threshold of exp.
• The thresholds are 11.0899 for fp16, 88.7228 for fp32 and 709.7827 for fp64.

Edge cases to know:

• Masked rows. Masks set blocked scores to minus infinity before softmax, which works while at least one score is finite. If every score in a row is minus infinity, the maximum is minus infinity, the shift computes minus infinity minus minus infinity, and the row is NaN.¹ Using a large finite negative number such as minus one billion avoids NaN but returns a uniform row of 0.3333 for three entries, which silently attends to masked positions.¹ Handle fully masked rows explicitly; the causal mask page shows where masks enter attention.
• Low-precision formats. In IEEE half precision (fp16), exp overflows at an input of only 11.0899, against 88.7228 in 32 bits.¹ The shift still removes the overflow, but the sum and division lose more precision, so many implementations accumulate in a wider format. Bfloat16 keeps the 32-bit exponent range, so its overflow threshold is about 88.7189, but it has fewer mantissa bits, so the sum still loses precision.¹
• Beam search. Decoding sums log-probabilities along a path, which is why beam search works with log-softmax and not with products of small probabilities.

How to answer it in an interview

The spoken structure: say what breaks, say the shift and why the output is unchanged, say why the denominator is at least one, then say log-softmax for anything that takes a log.

Common mistakes

• "Subtracting the max changes the answer." (wrong) It does not, because the factor exp(minus m) cancels between numerator and sum. (right) "Softmax is invariant to adding or subtracting a constant."
• "The shift prevents underflow too." (wrong) It prevents overflow and a zero denominator, but small outputs can still round to zero. (right) "The shift handles overflow; for log-probabilities I use log-softmax."
• "I apply softmax, then take the log for the loss." (wrong) A probability that rounds to zero gives minus infinity. (right) "I compute log-softmax directly, or feed logits to a loss that does it internally."
• "Masked positions are safe because I used minus infinity." (wrong) A row where everything is masked returns NaN. (right) "I handle fully masked rows explicitly."

Follow-up questions

How do you implement a numerically stable softmax in Python?

Subtract the row maximum with keepdims=True, exponentiate, and divide by the sum along the same axis, as in the worked example. Do the reduction along the last axis so a batch of rows is handled together.

Why not just take the log of softmax?

A stable softmax can still return exactly zero for a small probability, and the log of zero is minus infinity. In the worked check, the input 0 and minus 200 gave a log of minus infinity but a log-softmax of minus 200.¹ Log-softmax computes x minus the maximum minus the log of the sum and never forms the tiny probability.

What happens if every score in a row is masked?

The row returns NaN with minus infinity masks, because the maximum is minus infinity and the shift subtracts infinity from infinity.¹ The fix is a check for fully masked rows, not a different formula. The from-scratch generative pre-trained transformer (GPT) build is a good place to see masks used in practice.

How do you implement streaming entropy with numerical stability?

For counts c_i with total S, the entropy is the formula below, where a count of zero contributes zero.

$$H=\ln S-\frac{1}{S}\sum _i{c}_i\ln c_i$$

For the counts 5, 3 and 2 this gives 1.0297 nats, equal to the direct formula, and for 5, 0 and 5 it gives 0.6931, which is ln 2.¹ Keep S and the running sum of c ln c while reading the stream. When one count dominates, the two terms are nearly equal and the subtraction loses digits, so for extreme counts compute the probabilities first.

How does online softmax avoid an extra pass?

It keeps a running maximum and a running sum, and when a new maximum arrives it multiplies the old sum by exp(old maximum minus new maximum) before adding the new term.¹ Milakov and Gimelshein use this to merge the maximum and sum passes. Softmax sits inside attention, so the same trick appears in multi-head attention written in PyTorch without undefined.

Do PyTorch and NumPy already compute softmax stably?

Library softmax and log-softmax functions such as torch.softmax, torch.nn.functional.log_softmax, torch.logsumexp and scipy.special.logsumexp are built to be stable, so in real code you call them. PyTorch's CrossEntropyLoss takes unnormalized logits and applies log-softmax internally (PyTorch documentation, read October 2026). In an interview, write the shift by hand when asked to implement softmax, then say you would call the library function in production. In IEEE half precision the overflow threshold of exp falls from 88.7228 to 11.0899, so unshifted code fails at much smaller inputs, and accumulating the sum in a wider format is a common choice.¹

Sources

1. Online normalizer calculation for softmax (arXiv 1805.02867) — arXiv (Milakov and Gimelshein), 28 Jul 2018
2. Accurately computing the log-sum-exp and softmax functions — IMA Journal of Numerical Analysis (Blanchard, Higham and Higham), 19 Aug 2020
3. torch.nn.CrossEntropyLoss — PyTorch documentation, 26 Oct 2026