Numerical stability: log-sum-exp and overflow: study notes

 

1.29 Numerical stability: log-sum-exp and overflow: study notes

Learn AI by Story

1. The idea in one sentence

Computers store numbers with a limited range, so a formula that is perfectly correct on paper (like softmax) can blow up into infinity or NaN unless you rewrite it in an algebraically identical but safer form, and the safest rewrite is to subtract the biggest value first (the log-sum-exp trick).

2. The story from the video

Arjun's model produces three scores for a photo: 1000, 999 and 998. He feeds them into the textbook softmax formula, e to the power of each score divided by the sum, and the program prints NaN. Nothing is wrong with the maths: exp(1000) is simply too large for the computer, so it becomes infinity, and infinity divided by infinity is not a number. Then his friend shows the fix: subtract the largest score (1000) from every score first. The scores become 0, -1, -2, everything fits, and the answer comes out as a sensible 0.665, 0.245, 0.090. Same answer on paper, but now the computer can actually compute it.

3. Key definitions (glossary)

  • Floating-point number: the way computers store real numbers, as a sign, a significand (the digits) and an exponent (the scale), like scientific notation in binary. Plain English: a number with a sliding decimal point, but a limited number of digits and a limited size range.
  • float16, float32, float64: floating-point formats using 16, 32 and 64 bits. More bits means more range and more precision. float32 is the usual default in deep learning; float64 is the usual default in NumPy and plain Python.
  • Maximum representable value: the largest finite number a format can hold. float16: 65,504. float32: about 3.4028235e38. float64: about 1.7977e308.
  • Overflow: a result is bigger than the format can hold, so it becomes plus or minus infinity (inf). Plain English: the number does not fit in the box.
  • Underflow: a result is so close to zero that it is rounded to exactly zero. Plain English: the number is too tiny to see, so the computer calls it 0.
  • inf: the special value meaning "larger than anything representable".
  • NaN (not a number): the special value produced by undefined operations such as inf / inf, 0 / 0 or inf - inf. Once a NaN enters a calculation it spreads to everything it touches, including gradients.
  • Softmax: turns a list of real scores (logits) into probabilities: softmax_i(z) = exp(z_i) divided by the sum over j of exp(z_j). Outputs are positive and add to 1.
  • Logit: a raw, unnormalised score before softmax or sigmoid.
  • Log-sum-exp (LSE): LSE(z) = ln( sum over i of exp(z_i) ). Plain English: a smooth version of "take the maximum".
  • Shift invariance of softmax: softmax(z + c) = softmax(z) for any constant c, because exp(z_i + c) = exp(c) times exp(z_i) and the exp(c) cancels top and bottom.
  • Log-softmax: ln of softmax: log softmax_i(z) = z_i - LSE(z). Used inside cross-entropy.
  • Cross-entropy loss: for the true class k, the loss is -ln(softmax_k(z)) = LSE(z) - z_k. Plain English: how surprised the model is by the right answer.
  • Numerical stability: an algorithm is numerically stable if rounding and range limits do not wreck its answer. Plain English: it still works on awkward inputs.
  • Mixed precision: training with a mix of float16 (or bfloat16) and float32 to save memory and time; it makes range problems more likely, so stable formulas matter more.
  • Machine epsilon: the smallest gap that matters near 1.0 (about 1.19e-7 for float32); it describes precision, whereas overflow concerns range.

4. Step-by-step explanation, from basics to the full picture

Step 1: numbers on a computer have limits. In school, numbers can be as big as you like. A computer stores each number in a fixed number of bits. With 32 bits (float32) the biggest finite value is about 3.4 x 10^38. The Deep Learning book (Goodfellow, Bengio and Courville, chapter 4, section 4.1 "Overflow and Underflow") explains the two failure modes: underflow when a number near zero is rounded to zero, and overflow when a large number becomes infinity and then often NaN.

Step 2: exp grows explosively. The exponential function doubles its growth rate every time: exp(1) is about 2.7, exp(10) is about 22,026, exp(88) is about 1.65 x 10^38 (just under the float32 limit) and exp(89) is already infinity in float32. The cut-off is ln(3.4028 x 10^38), which is about 88.72. For float64 the cut-off is about 709.78. For float16 it is only about 11 (exp(11) = 59,874 fits under 65,504, but exp(12) = 162,755 does not). Logits of 100 or 1000 are not unusual in practice, so overflow is a real danger.

Step 3: the naive softmax fails. With scores 1000, 999, 998 the naive formula computes exp(1000), which is inf even in float64. The sum is inf, and inf / inf is NaN. All three outputs are NaN.

Step 4: very negative scores fail the other way. With scores -1000, -1001, -1002, every exp(...) underflows to exactly 0. The sum is 0, and 0 / 0 is NaN. Even if you avoid the NaN, ln(0) is minus infinity, which would make a loss infinite.

Step 5: the shift trick. Because softmax(z + c) = softmax(z), choose c = minus the maximum of z. After subtracting, the largest entry is exactly 0, so the largest exp is exactly 1 and nothing can overflow. At least one term in the sum equals 1, so the sum is at least 1 and never zero, so there is no 0 / 0. Smaller terms may underflow to 0, but that is harmless because they were negligible anyway. This is the method described in Goodfellow et al. section 4.1 and in the d2l.ai section "Softmax Revisited".

Step 6: the same idea gives log-sum-exp. Write LSE(z) = m + ln( sum of exp(z_i - m) ) where m is the maximum of z. This is exact algebra: ln( sum exp(z_i) ) = ln( exp(m) times sum exp(z_i - m) ) = m + ln( sum exp(z_i - m) ). Now the big number m is added outside the exponential, where size does not hurt.

Step 7: log-softmax directly. Log-softmax_i = z_i - LSE(z) = (z_i - m) - ln( sum exp(z_j - m) ). Compute this directly instead of computing softmax and then taking the log. If you do softmax first and a probability underflows to 0, then ln(0) = -inf. The Deep Learning book says log-softmax must be implemented with its own stable function for exactly this reason. The d2l.ai text gives the combined formula log y-hat_j = o_j - o-bar - ln( sum exp(o_k - o-bar) ), where o-bar is the maximum logit.

Step 8: fuse softmax and cross-entropy. Since cross-entropy is minus the log of the true-class probability, you never need the probabilities themselves. Libraries therefore take raw logits and compute loss = LSE(z) - z_k in one stable operation. This is why, for example, PyTorch's cross-entropy loss expects unnormalised logits, and why feeding it softmax outputs is a mistake.

5. Everything the video skipped

Why "smooth maximum". LSE always sits close to the maximum: max(z) <= LSE(z) <= max(z) + ln(n), where n is the number of entries. Reason for the left side: the sum contains exp(max) plus other positive terms, so it is at least exp(max). Reason for the right side: each term is at most exp(max), so the sum is at most n times exp(max). Boyd and Vandenberghe (Convex Optimization, section 3.1) state this bound and note that LSE is a convex function. When one score dominates, LSE is almost exactly the maximum; when all scores are equal, it is the maximum plus ln n.

Temperature. Divide the logits by a temperature T before softmax. Small T pushes toward the hard maximum (argmax); large T pushes toward uniform probabilities. With LSE you can write T times LSE(z / T), which approaches max(z) as T goes to 0. The shift trick works unchanged.

Softplus and the sigmoid. The softplus function ln(1 + exp(x)) is log-sum-exp of the two values 0 and x. Computing it naively overflows for large x. The stable form is max(x, 0) + ln(1 + exp(-abs(x))). The sigmoid 1 / (1 + exp(-x)) has the same issue for very negative x; stable implementations use different formulas for positive and negative inputs. See roadmap 1.28 for sigmoid and softmax basics, and the Deep Learning book section 3.10 for softplus.

Underflow in probabilities (long products). Multiplying many small probabilities, as in a language model or a hidden Markov model, shrinks toward 0 quickly: 300 probabilities of 0.001 would be 10^-900, which is below even float64's smallest positive range (roughly 10^-308, or about 10^-324 for denormal numbers). The fix is to work with log-probabilities, turning products into sums. When you must add probabilities (marginalise), you add them in log space with log-sum-exp.

Other places the same trick appears. Log-likelihood of mixtures (Gaussian mixture models), the forward algorithm in sequence models, attention softmax in transformers (implementations subtract the row maximum; fast variants such as flash attention compute a running maximum and rescale), and partition functions in energy-based models.

Precision versus range. Overflow is about range (the largest and smallest sizes). A different problem, rounding error, is about precision: float32 keeps roughly 7 decimal digits. Adding a tiny number to a huge one can lose the tiny one entirely. Subtracting two nearly equal numbers (catastrophic cancellation) loses significant digits. The shift trick helps range; other tricks help precision.

Other stability tools. Gradient clipping limits huge gradients; adding a small epsilon (for example 1e-8) inside a log or a division avoids ln(0) or division by zero; clamping probabilities to [epsilon, 1 - epsilon]; loss scaling in float16 training to avoid gradient underflow; using bfloat16, which has the same range as float32 but fewer precision digits. These are patches; the fused stable formula is better when available.

Background (history). Floating-point arithmetic was standardised as IEEE 754 in the 1980s, which is why inf and NaN behave the same on almost every machine. The log-sum-exp trick is old folklore in statistics and speech recognition and is now built into every major library, for example torch.logsumexp, whose documentation says the computation is numerically stabilised.

Limits. The shift trick cannot rescue a model whose logits genuinely diverge to huge values; that points to a training problem (learning rate too high, bad initialisation). It also does not fix the fact that a softmax output of exactly 0 or 1 gives zero gradient. And in float16, even logits of 20 or so can lose the small probabilities; many systems keep the loss computation in float32.

Roadmap links. Before: 1.28 Logs, exponentials, sigmoid and softmax (this video fixes the problem hidden there). After: 1.30 Constrained optimisation and Lagrange multipliers, 1.31 Convex sets and functions (LSE is a convex function), 1.32 Gradient descent convergence, and later 2.05 Loss functions where cross-entropy is used in earnest.

6. Worked examples with real numbers

All numbers below were checked with Python.

Example A: the story, scores (1000, 999, 998).

  1. Naive: exp(1000) = inf (float64 overflows above about 709.78), so every term is inf and the result is inf / inf = NaN.
  2. Shift by the maximum 1000: new scores are (0, -1, -2).
  3. Exponentials: exp(0) = 1, exp(-1) = 0.3679, exp(-2) = 0.1353.
  4. Sum = 1 + 0.3679 + 0.1353 = 1.5032.
  5. Probabilities: 1 / 1.5032 = 0.665, 0.3679 / 1.5032 = 0.245, 0.1353 / 1.5032 = 0.090. They add to 1.
  6. LSE = 1000 + ln(1.5032) = 1000 + 0.408 = 1000.408.
  7. Log-softmax: (0 - 0.408, -1 - 0.408, -2 - 0.408) = (-0.408, -1.408, -2.408). If the true class is the first one, the cross-entropy loss is 0.408. Check: LSE - z_1 = 1000.408 - 1000 = 0.408.

Example B: the underflow side, scores (-1000, -1001, -1002).

  1. Naive: exp(-1000) is about 0 in float64 (underflow), so all three terms are 0, the sum is 0 and the result is 0 / 0 = NaN.
  2. Shift by the maximum -1000: scores become (0, -1, -2), the very same as Example A.
  3. So the probabilities are again 0.665, 0.245, 0.090. This shows shift invariance: only the differences between the scores matter.

Example C: the smooth-maximum bound, scores (3, 1).

  1. LSE(3, 1) = 3 + ln(1 + exp(-2)) = 3 + ln(1.1353) = 3 + 0.127 = 3.127.
  2. The maximum is 3 and ln 2 = 0.693, so the bound says 3 <= LSE <= 3.693. And 3.127 is inside it, close to the maximum because 3 is much bigger than 1.
  3. For equal scores (2, 2), LSE = 2 + ln 2 = 2.693, the top of the allowed range.

Example D: where float32 stops. exp(88) = 1.65 x 10^38 fits in float32 (limit 3.40 x 10^38), but exp(89) overflows to inf. So a float32 model with a logit of 89 breaks the naive softmax, while a logit of 80 does not. After subtracting the max, the largest exponent is always 0.

7. Comparison tables

Table 1: number formats

Format

Bits

Largest finite value

Naive exp overflows above

Typical use

float16

16

65,504

about 11

mixed-precision training

float32

32

about 3.4e38

about 88.72

default deep learning

float64

64

about 1.8e308

about 709.78

scientific computing, NumPy default

Table 2: naive versus stable computation

Task

Naive formula

Problem

Stable formula

Softmax

exp(z_i) / sum exp(z_j)

overflow (inf / inf) or underflow (0 / 0)

exp(z_i - m) / sum exp(z_j - m)

Log-sum-exp

ln( sum exp(z_i) )

overflow, or ln(0)

m + ln( sum exp(z_i - m) )

Log-softmax

ln( softmax(z) )

softmax underflows to 0, ln(0) = minus infinity

z_i - LSE(z)

Cross-entropy

minus ln( softmax_k )

the same two failures

LSE(z) - z_k from raw logits

Softplus

ln(1 + exp(x))

overflow for large x

max(x, 0) + ln(1 + exp(minus abs(x)))

Table 3: overflow versus underflow

Property

Overflow

Underflow

What happens

number too big becomes inf

number too close to 0 becomes 0

Typical trigger

exp of a large positive number

exp of a large negative number, long products of probabilities

Follow-on damage

inf / inf = NaN

0 / 0 = NaN, ln(0) = minus infinity

Main cure

subtract the maximum

subtract the maximum, or work in log space

8. Common mistakes and myths

  • Myth: "NaN means my formula is wrong." Often the formula is right but the computer cannot hold the intermediate numbers.
  • Mistake: applying softmax in the model and then using a cross-entropy loss that expects logits. You apply softmax twice, which both hurts training and wastes the stability built into the loss.
  • Mistake: computing log(softmax(z)) as two steps. Use a log-softmax function or log-sum-exp.
  • Mistake: subtracting the minimum or the mean instead of the maximum. Subtracting the minimum leaves the largest term as exp(max - min), which can still overflow. Subtracting the maximum guarantees all exponents are 0 or negative.
  • Mistake: thinking the shift changes the answer. It does not; softmax(z + c) = softmax(z) exactly.
  • Myth: "float64 solves it." It only moves the cliff from about 88 to about 710. The shift trick is free, so use it anyway.
  • Mistake: adding a tiny epsilon everywhere without understanding. An epsilon inside ln helps, but it slightly biases the result; the fused stable form is cleaner.
  • Mistake: ignoring NaN in the loss curve. One NaN gradient ruins all the weights. Check for it early (for example with anomaly detection or by monitoring the loss).
  • Myth: "underflow to zero is always harmless." It is harmless for terms that are truly negligible but fatal when you then take a log or divide by it.

9. Real-life and AI uses

  • Classification networks: every image or text classifier that ends in softmax uses fused log-softmax and cross-entropy on raw logits.
  • Transformers and attention: the softmax over attention scores subtracts the row maximum; efficient kernels track a running maximum across blocks.
  • Language models: probabilities of whole sentences are sums of log-probabilities; sampling with temperature uses logits divided by T.
  • Mixture models and EM: responsibilities of Gaussian mixtures are computed with log-sum-exp.
  • Sequence models: the forward algorithm in HMMs and CTC loss for speech recognition run in log space.
  • Mixed-precision training: float16 overflows at 65,504, so stable formulas and loss scaling are essential.
  • Finance and statistics: log-likelihoods of large datasets, Bayesian model evidence, and partition functions all use log-sum-exp.

10. Self-check questions with answers

  1. Why does the naive softmax of (1000, 999, 998) give NaN? exp(1000) overflows to infinity, the sum is infinity, and infinity divided by infinity is NaN.
  2. What is the largest argument of exp that float32 can handle? About 88.72; exp(88) fits (1.65e38), exp(89) is inf.
  3. Why can we subtract any constant from all logits? Because exp(z_i - c) = exp(z_i) times exp(-c), and the common factor cancels in the ratio, so softmax(z - c) = softmax(z).
  4. Why subtract the maximum specifically? It makes the largest exponent 0, so no exp can overflow, and the sum is at least 1, so we never divide by zero.
  5. Compute softmax(1000, 999, 998). About 0.665, 0.245, 0.090.
  6. Write the stable formula for log-sum-exp. LSE(z) = m + ln( sum exp(z_i - m) ), with m the maximum of z.
  7. What is LSE(3, 1)? 3 + ln(1 + e^-2), about 3.127, between 3 and 3 + ln 2 = 3.693.
  8. Why feed raw logits to a cross-entropy loss? The loss LSE(z) - z_k can be computed stably in one step; going through probabilities risks underflow and ln(0).
  9. What goes wrong if all logits are very negative and you use the naive formula? Every exp underflows to 0, giving 0 / 0 = NaN (and ln(0) = minus infinity).
  10. Is float64 a full fix? No. Its limit is about 709.78, and exp(1000) still overflows. Use the shift trick regardless.

11. Key takeaways, go deeper, next

Key takeaways

  • Computers have a finite range: overflow gives inf, underflow gives 0, undefined operations give NaN.
  • exp grows so fast that float32 overflows at an input of about 88.7.
  • Softmax is unchanged when you add or subtract a constant, so subtract the maximum.
  • LSE(z) = max + ln( sum exp(z_i - max) ) and satisfies max <= LSE <= max + ln n.
  • Use fused log-softmax or cross-entropy functions that take raw logits.
  • Work in log space when multiplying many probabilities.

Go deeper

  • Read Goodfellow et al. section 4.1 and try the examples in NumPy with np.float32 and np.float16.
  • Implement logsumexp by hand and compare with torch.logsumexp on inputs like 1000 and -1000.
  • Look at how attention kernels keep a running maximum.
  • Prove the bound max <= LSE <= max + ln n yourself.

Next on the roadmap: 1.30 Constrained optimisation and Lagrange multipliers.

12. Sources and further reading

  • Goodfellow, Bengio and Courville, Deep Learning (2016), chapter 4 "Numerical Computation", section 4.1 "Overflow and Underflow"; and chapter 3, section 3.10 "Useful Properties of Common Functions" (sigmoid, softplus). https://www.deeplearningbook.org
  • Zhang et al., Dive into Deep Learning, section 4.1 "Softmax Regression". https://d2l.ai/chapter_linear-classification/softmax-regression.html
  • Zhang et al., Dive into Deep Learning, section 4.5 "Concise Implementation of Softmax Regression", subsection "Softmax Revisited" (overflow, underflow, subtracting the maximum, combined log-softmax). https://d2l.ai/chapter_linear-classification/softmax-regression-concise.html
  • Boyd and Vandenberghe, Convex Optimization (2004), section 3.1 "Basic properties and examples" (log-sum-exp is convex; max <= LSE <= max + log n). https://web.stanford.edu/~boyd/cvxbook/bv_cvxbook.pdf
  • PyTorch documentation, torch.logsumexp (the computation is numerically stabilised). https://docs.pytorch.org/docs/2.3/generated/torch.logsumexp.html

Accuracy note: all numbers (float limits, softmax and LSE values) were checked with Python on 2026-11-03; the claims about PyTorch's cross-entropy taking logits, flash-attention running maxima and IEEE 754 history are standard background rather than taken from the pages listed.

Comments

Popular posts from this blog

What is AI? Narrow vs General Intelligence (Study Notes)

The five kinds of learning: study notes

Rules vs learning: study notes