Softmax without overflow — the max-subtraction trick, proved

Softmax is four characters of maths and one of the most common sources of nan in a training run. The fix is one line, and it is worth understanding rather than copying.

The function

Softmax turns a vector of scores into a distribution:

softmax(z)ᵢ = e^zᵢ / Σⱼ e^zⱼ

Every output is positive and they sum to 1. Larger scores get exponentially more mass, which is where the “soft” maximum comes from — as the gap between the largest score and the rest grows, the distribution approaches a one-hot vector.

Where it breaks

float64 tops out at about 1.8 × 10³⁰⁸, and e^709 ≈ 8.2 × 10³⁰⁷. One step further and it is inf:

z = np.array([800.0, 801.0])
np.exp(z)                  # [inf, inf]
np.exp(z) / np.exp(z).sum()  # [nan, nan]

inf / inf is nan, the nan propagates through the backward pass, every weight becomes nan, and the loss prints as nan from that step onward. In float32 the ceiling is far lower — e^88 — so this is not an exotic case. Logits of a few hundred appear routinely in an unnormalised network or after a bad initialisation.

The trick, and why it is exact

Subtract the maximum score before exponentiating:

m = max(z)
softmax(z)ᵢ = e^(zᵢ - m) / Σⱼ e^(zⱼ - m)

This is not an approximation. Multiply numerator and denominator by e^(-m):

e^zᵢ / Σⱼ e^zⱼ  =  (e^zᵢ · e^-m) / (Σⱼ e^zⱼ · e^-m)
                =  e^(zᵢ - m) / Σⱼ e^(zⱼ - m)

The constant cancels, so the output is identical. What changes is the range: the largest exponent is now exactly e⁰ = 1, and everything else is in (0, 1]. Overflow is impossible. Underflow can still occur for very negative shifted scores, but underflow to 0 is a harmless rounding — that class had approximately zero probability anyway.

def softmax(z):
    z = z - z.max()          # the entire fix
    e = np.exp(z)
    return e / e.sum()

softmax(np.array([800.0, 801.0]))   # [0.269, 0.731]

Check the answer by hand: the gap is 1, so the ratio is e¹ ≈ 2.718, and 2.718 / 3.718 = 0.731. Only the differences between logits matter — adding a constant to every score leaves softmax unchanged. That shift-invariance is the same property the trick exploits.

Log-sum-exp

Training needs log(softmax(z)), and computing the softmax first then taking its log throws away precision twice. Expand instead:

log softmax(z)ᵢ = zᵢ - log Σⱼ e^zⱼ
                = zᵢ - m - log Σⱼ e^(zⱼ - m)

The right-hand form is stable for the same reason. NumPy and PyTorch both ship it — scipy.special.logsumexp, torch.logsumexp — and every fused loss uses it internally.

# Two rounding steps, one of them near zero
loss = -np.log(softmax(z)[label])

# One stable expression
loss = logsumexp(z) - z[label]

The rule this generalises to

Any time an expression exponentiates something unbounded, ask what its largest possible value is:

• Softmax → subtract the max.
• Attention → the /√d in softmax(QKᵀ/√d) exists for this reason. Without it, dot products of d-dimensional vectors grow like √d, the logits spread, softmax saturates, and the gradient vanishes.
• Sigmoid + log → use logaddexp, or a fused BCEWithLogitsLoss.
• Products of probabilities → sum logs instead of multiplying, or a long sequence underflows to zero.

The practical rule I now follow: keep logits raw until the last possible moment, and hand them to a loss function that expects logits. Every framework provides one, and every one of them is doing this trick inside.

What to take forward

• exp overflows around 709 in float64 and 88 in float32.
• Subtracting the max is exact, not an approximation — the constant cancels.
• Softmax depends only on differences between logits.
• Prefer logsumexp and logits-based losses over composing log with softmax.

---

— Ishaan SandhwarNumerical stabilitySoftmaxFundamentals