Linear Algebra for ML · 1.21

Numerical Stability & the Log-Sum-Exp Trick

Why softmax silently produces NaN, and the one-line fix every framework uses.

On this page

In plain English — beginner to advanced

Beginner: computers store numbers with limited precision (floating point). Some mathematically simple expressions — like exp() of a large number, or dividing by a very small number — silently overflow to infinity, underflow to zero, or lose almost all their precision, even though the "true" math is perfectly well-behaved.

Intermediate: the single most common place this bites ML code is the softmax function, exp(x_i) / Σ exp(x_j). If any x_i is even moderately large (say, 1000), exp(1000) overflows to infinity in float32 — even though the final softmax probability would have been a perfectly ordinary number between 0 and 1.

Advanced: the fix, the log-sum-exp trick, subtracts the maximum value before exponentiating: exp(x_i − max(x)) / Σ exp(x_j − max(x)). This is mathematically identical to the original formula (the max cancels out), but now the largest exponent computed is always exp(0) = 1, so nothing ever overflows.

Formula

\text{softmax}(x_i) = \frac{e^{x_i - m}}{\sum_j e^{x_j - m}}, \quad m = \max_j x_j

Subtracting the max is algebraically a no-op (it cancels between numerator and denominator) but is the difference between code that works and code that silently produces NaN.

Derivation: why subtracting the max changes nothing

Start from the plain definition and factor out a constant from every exponential:

e^{x_i - m} = e^{x_i}\,e^{-m}

Substitute this into both the numerator and every term of the denominator's sum:

\frac{e^{x_i-m}}{\sum_j e^{x_j-m}} = \frac{e^{x_i}e^{-m}}{\sum_j e^{x_j}e^{-m}} = \frac{e^{-m}\,e^{x_i}}{e^{-m}\sum_j e^{x_j}}

The e⁻ᵐ factor is identical in every term of the sum, so it can be pulled outside — and then it cancels exactly between numerator and denominator, leaving the original, unshifted softmax formula untouched:

= \frac{e^{x_i}}{\sum_j e^{x_j}} = \text{softmax}(x_i)

The shift by m is therefore a pure identity, true for any choice of m — choosing m = max(x) specifically is what keeps every exponent computed ≤ 0, so the largest term is e⁰ = 1 and nothing can overflow.

Where this is used: this exact algebraic cancellation is why every framework's built-in log_softmax/cross_entropy function is safe to call on raw, unbounded network outputs (logits) without any manual pre-scaling.

Practical example — softmax done wrong, then done right

Run softmax_naive yourself on large inputs — every deep learning framework's built-in softmax and cross-entropy functions use the stable version internally for exactly this reason.

python

import numpy as np

def softmax_naive(x):
    e = np.exp(x)
    return e / e.sum()

def softmax_stable(x):
    e = np.exp(x - x.max())
    return e / e.sum()

x = np.array([1000., 1001., 1002.])
print(softmax_naive(x))    # [nan, nan, nan] -- overflow!
print(softmax_stable(x))    # [0.09, 0.24, 0.67] -- correct

# Same trick for log-likelihoods: log-sum-exp
def logsumexp(x):
    m = x.max()
    return m + np.log(np.sum(np.exp(x - m)))
print(logsumexp(x))   # a normal, finite number

Real-world examples

Cross-entropy loss implementations always combine log and softmax into one numerically stable operation (log_softmax) rather than computing softmax then taking its log separately, which would reintroduce the same overflow risk.
Attention scores in Transformers go through exactly this stabilized softmax — a single unstabilized attention layer could silently produce NaN losses on certain inputs.
Log-determinants (section 1.12) and Gaussian log-likelihoods use the same underlying principle: work in log-space as long as possible, only exponentiate at the very end, if at all.

Common mistakes

Computing probabilities as raw ratios of exponentials instead of using a library's stable softmax/log_softmax function — this is one of the most common sources of mysterious NaN losses in from-scratch model implementations.
Dividing by a norm or standard deviation without checking it isn't (near) zero — add a small eps (e.g. 1e-8) to the denominator as standard practice.

Going deeper

The same shift-before-exponentiate idea generalizes to any "log of a sum of exponentials" expression, which is common throughout probabilistic modeling (mixture models, hidden Markov models, variational inference) — it's universally called the log-sum-exp trick for exactly this reason.

At the master level: float16/bfloat16 mixed-precision training (section 1.5's Tensor Cores) makes numerical stability substantially more fragile than float32 — this is precisely why modern training frameworks use loss scaling (multiplying the loss by a large constant before the backward pass, then dividing gradients by the same constant afterward) to keep small gradient values from underflowing to zero in reduced precision.

Check yourself

Why does subtracting the max before exponentiating not change the softmax result?

Because e^(x-m) = e^x · e^(-m), and the e^(-m) factor appears identically in every term of both the numerator and denominator, so it cancels out completely.

PreviousMatrix Norms: Frobenius, Spectral & Nuclear NextSparse Matrices & Structured Storage

Numerical Stability & the Log-Sum-Exp Trick

Stay in the loop

Newsletter