Numerically stable softmax and log-sum-exp

Every deep learning library hides this. Interviewers ask you to write it because most candidates have never had to.

Where it breaks

A float64 holds numbers up to about 1.8 × 10³⁰⁸. 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 curve does not drift; it vanishes.

The shift

Softmax does not 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 = e^z_i−c / Σ_j e^z_j−c = e^−ce^z_i / e^−cΣ_j e^z_j = softmax(z)_i

Choose c = max(z). Now the largest exponent is e⁰ = 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 is the correct answer to sixteen digits: a probability of e⁻⁸⁰⁰ is zero for every practical purpose.

Softmax of three large logits, step by step.

Step 1 / 5

Three logits from a confident model.
Step 2 / 5

np.exp(1000) overflows float64. Every entry is inf, and inf / inf is nan.
Step 3 / 5

Subtract the largest logit, 1000. Softmax does not change when every logit moves by the same amount.
Step 4 / 5

The largest exponent is now e⁰ = 1, so nothing can overflow, and the sum is at least 1.
Step 5 / 5

Divide by the sum, 1.5032: the probabilities, accurate to the last digit.

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 log 0 = −∞. Work in log space throughout instead:

log softmax(z)_i = z_i − log Σ_j e^z_j = z_i − m − log Σ_j e^z_j−m, m = max(z)

The second term is the log-sum-exp. It is a smooth version of the maximum: it is never less than max(z) and never more than max(z) + log n. Written with the shift, it is 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 is the one PyTorch's BCEWithLogitsLoss uses; deriving it is one of the problems below.

What an interviewer is checking

That you know the formula you were taught is not the one that runs in production.
That you can prove the shift changes nothing, in one line, without hand-waving.
That you reach for log space before you are told to, because a probability of 10⁻⁴⁰⁰ still has a perfectly finite logarithm.
That your gradient is right: for softmax with cross-entropy it is simply p − y, and saying so unprompted is worth more than the code.

Common mistakes

Subtracting the mean instead of the maximum. The mean does not 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.

Practise it

Easy8 min