Why softmax returns NaN at a logit of 710, and the one-line fix

If a training run has ever turned its loss into nan with no warning worth reading, there's a fair chance softmax did it. The formula on the whiteboard is fine. The floating-point version of it isn't, and the fix is one line.

Where it breaks

A float64 holds numbers up to about $1.8 \times 1 0^{308}$ . Since $e^{709.8}$ is already that large, np.exp(710.0) returns inf with a warning. At the other end, np.exp(-746.0) underflows to exactly zero. Logits from a trained network rarely reach those values, but logits during training, in a badly scaled layer, or from a model that has become very confident, do. When they do, the textbook formula

p = np.exp(z) / np.exp(z).sum()

returns inf / inf = nan, and a single nan spreads through every weight it touches on the next update. The loss doesn't creep up first. It just disappears.

The shift

Softmax doesn't change when you add the same constant to every logit, because the constant factors out of the numerator and the denominator alike:

softmax (z - c)_{i} = \frac{e ^{z_{i} - c}}{\sum _{j} e ^{z_{j} - c}} = \frac{e ^{- c} e ^{z_{i}}}{e ^{- c} \sum _{j} e ^{z_{j}}} = softmax (z)_{i}

Choose $c = max (z)$ . Now the largest exponent is $e^{0} = 1$ , so nothing overflows, and the denominator is at least 1, so nothing divides by zero. Small entries may still underflow to zero, but that's the correct answer to sixteen digits: a probability of $e^{- 800}$ is zero for every practical purpose.

Log-sum-exp

Cross-entropy needs log softmax, not softmax. Computing softmax first and then taking the log throws away exactly the information you need: a probability that underflowed to 0 becomes $lo g 0 = - \infty$ . Work in log space throughout instead:

lo g softmax (z)_{i} = z_{i} - lo g j \sum e^{z_{j}} = z_{i} - m - lo g j \sum e^{z_{j} - m}, m = max (z)

The second term is the log-sum-exp. It's a smooth version of the maximum: it's never less than $max (z)$ and never more than $max (z) + lo g n$ . Written with the shift, it's exact and finite for any finite input.

def logsumexp(z, axis=-1): m = np.max(z, axis=axis, keepdims=True) return m + np.log(np.sum(np.exp(z - m), axis=axis, keepdims=True))

Sigmoid and binary cross-entropy

The sigmoid has the same problem in a different shape. 1 / (1 + np.exp(-z)) is fine for large positive z and overflows inside the exponential for large negative z. The usual repair is to branch on the sign, so that the exponential only ever sees a non-positive argument. For the loss itself there is a closed form that needs no branch at all, and it's the one PyTorch's BCEWithLogitsLoss uses; deriving it's one of the problems below.

If this comes up in an interview

Say early that the textbook formula isn't what runs in production, and prove the shift changes nothing in one line. Move to log space before anyone asks, since a probability of $1 0^{- 400}$ still has a perfectly finite logarithm. And if you can say without being prompted that the gradient of softmax with cross-entropy is just $p - y$ , I'd rate that above the code.

Common mistakes

Subtracting the mean instead of the maximum. The mean doesn't bound the exponent.
Shifting a batch by one global maximum rather than row by row. One very large logit in row 3 then underflows every entry of row 7.
Adding a small epsilon inside the log, np.log(p + 1e-12). It hides the problem, biases the loss, and caps the gradient at the wrong value.

First published here: https://geodesicpartners.co.uk/ml-numerical-stability