The log-sum-exp trick, and why exp(1000) is not a number
The overflow that produces NaN
A float32 number cannot exceed about 3.4 × 10^38. The exponential passes that at x = 88.7, so exp(89) in single precision is not a large number: it is inf. Logits of 89 are not exotic; an untrained network with a bad initialisation or a training run whose learning rate is too high produces them routinely.
Apply the naive softmax formula to z = (1000, 999). Both exponentials are inf. The sum is inf. Each ratio is inf / inf, which the floating-point standard defines as NaN. The answer should have been (0.731, 0.269), since the logits differ by exactly 1 and e / (e + 1) = 0.731. The arithmetic was easy; the intermediate values were not representable.
The fix is the shift invariance
The previous lesson showed that subtracting any constant from every logit leaves the softmax unchanged. Subtract the maximum:
z − max(z) = (0, −1)
e^0 = 1, e^(−1) = 0.368, sum = 1.368
p = (0.731, 0.269)After the shift the largest exponent is exactly 0, so the largest exponential is exactly 1, and nothing can overflow. The smaller logits become negative, and very negative ones underflow to zero, which is harmless: a class with probability 10^-50 contributes nothing to the sum, and the answer is unchanged to every digit float32 can hold.
This is not an optimisation. It is the difference between an answer and NaN, and every deep-learning library does it silently inside softmax.
Log-sum-exp
The same idea gives a stable way to compute log Σ e^(z_i), which appears in every log-probability, every normalising constant and every loss:
LSE(z) = m + log Σ_i e^(z_i − m), where m = max(z)For z = (1000, 999): m = 1000, the shifted sum is 1 + 0.368 = 1.368, its log is 0.313, and LSE = 1000.313. No overflow, and the result is exact to float32 precision. The unstabilised version would have returned inf.
Log-softmax follows immediately:
log p_i = z_i − LSE(z)Compute log-probabilities this way, never as log(softmax(z)). In the second form, a class whose exponential underflowed to zero produces log(0) = −inf, and a single −inf in a loss makes the gradient NaN for the whole batch. The first form gives 1000 − 1000.313 = −0.313 for the top class and a large finite negative number for the others, which is what you wanted.
The bug this explains
torch.nn.CrossEntropyLoss takes raw logits. Internally it computes log-softmax with the trick above and picks out the correct class, in one fused, stable operation. A common mistake is to apply softmax first and pass the probabilities in:
probs = torch.softmax(logits, dim=-1)
loss = F.cross_entropy(probs, targets) # wrong, and it runsNothing errors. But the loss function now treats the probabilities, all between 0 and 1, as logits, and applies a second softmax to them. Logits that differ by at most 1 give probabilities that are nearly uniform: three classes with input (0.9, 0.05, 0.05) come out as (0.53, 0.23, 0.23). The model trains, slowly, toward a ceiling it cannot pass, and the loss never goes far below ln(number of classes). If a classifier learns but its loss stalls near that value, look for a softmax that should not be there.
Adding probabilities in log space
Sequence models, beam search and hidden Markov models all need log(p_1 + p_2) when they hold only log p_1 and log p_2, and the probabilities themselves are too small to form. The answer is log-sum-exp of the two logs: LSE(log p_1, log p_2). For log p_1 = −700 and log p_2 = −701, the direct route underflows both to zero and returns log 0; the trick returns −700 + log(1 + e^(−1)) = −699.69. This is the single identity that makes probabilistic computation over long sequences possible at all.
What the trick does not fix
Subtracting the maximum prevents overflow. It does not add precision. Float32 carries about seven significant digits, so a term whose shifted exponential is below 10^-7 relative to the largest is lost from the sum entirely. That happens once a logit trails the maximum by more than about 16. The effect on the top class is nil, but the gradient with respect to a logit that far behind is computed as exactly zero rather than a tiny positive number. Usually harmless; occasionally the reason a rare class never moves. Module 8 explains where the seven digits come from.
In code
import numpy as np
from scipy.special import logsumexp # does the shift for you
z = np.array([1000.0, 999.0])
log_p = z - logsumexp(z) # (-0.313, -1.313)
p = np.exp(log_p) # (0.731, 0.269)Write the naive version once, feed it (1000, 999), and see the NaN with your own eyes. After that you will never wonder why the library code has that extra line.
The one thing to keep
Because float32 overflows at exp(88.7), softmax is computed by subtracting the maximum logit first, and log-probabilities are computed as z minus log-sum-exp rather than as the log of a softmax, which is why a loss function takes raw logits and why passing it probabilities trains toward a ceiling.
Before you move on
A learner writes `probs = softmax(logits)` and then `loss = F.cross_entropy(probs, y)`. Training runs without error but the loss will not go far below ln(number of classes). Why?
Pick the one you would defend. Nobody sees your answer.