Softmax: Why Subtract the Maximum
Evaluating exp(x) directly in softmax can overflow. Subtracting the row’s maximum logit preserves the mathematical probabilities while making every exponential input nonpositive. For log probabilities, log-softmax avoids another loss of information: producing a zero first and then taking its logarithm.
Contents
A valid formula can have unsafe intermediate values
Softmax assigns pᵢ=exp(xᵢ)/Σⱼexp(xⱼ). Here x is one row of finite real numbers, and the sum covers that same candidate set. The formula defines relative weights; it does not guarantee that floating-point arithmetic can represent the intermediate exponentials.
For the artificial logits [1000,1001,1002], the differences are modest and the probabilities should vary smoothly. Direct exponentiation can nevertheless exceed the available range. Normalizing afterward cannot recover an intermediate value that was already lost.
Why subtracting the maximum preserves probabilities
Let m=max(x). Multiplying both numerator and denominator by exp(-m) gives pᵢ=exp(xᵢ-m)/Σⱼexp(xⱼ-m). The common factor cancels. With finite inputs, each shifted value is nonpositive, and at least one exponential weight is exactly 1.
The example becomes [-2,-1,0], giving approximately [0.090031,0.244728,0.665241]. These are teaching values, not model predictions. A common additive shift preserves the distribution. Multiplying the logits by a coefficient changes their relative gaps and does not have that property.
The numerical analysis paper by Blanchard and the two Highams examines floating-point behavior of equivalent formulas. The shifted form is an established method, not a new 2026 result. This article does not reproduce the paper’s complete error analysis.
Keep tiny probabilities in the log domain
Shifting avoids large positive exponential inputs, but very small weights can still underflow. For [0,-1000], the second probability becomes 0 in this Python floating-point example, even though it is mathematically positive.
Taking log(0) afterward loses an otherwise finite log probability. Instead compute log pᵢ=(xᵢ-m)-log Σⱼexp(xⱼ-m). This avoids the premature conversion to a tiny probability. Floating-point rounding still exists; it is not a guarantee of perfect accuracy for every range and precision.
This is also why perplexity evaluation usually benefits from working with model log probabilities rather than repeatedly converting between probabilities and logarithms.
A runnable minimal implementation
import math
def softmax(xs):
if not xs or not all(math.isfinite(x) for x in xs):
raise ValueError("Expected nonempty finite logits")
m = max(xs)
weights = [math.exp(x - m) for x in xs]
total = math.fsum(weights)
return [w / total for w in weights]
def log_softmax(xs):
if not xs or not all(math.isfinite(x) for x in xs):
raise ValueError("Expected nonempty finite logits")
m = max(xs)
shifted = [x - m for x in xs]
log_total = math.log(math.fsum(math.exp(x) for x in shifted))
return [x - log_total for x in shifted]
p = softmax([1000.0, 1001.0, 1002.0])
assert math.isclose(math.fsum(p), 1.0)
assert all(math.isclose(a, b) for a, b in zip(p, softmax([0., 1., 2.])))
assert softmax([7.0]) == [1.0]
assert math.isfinite(log_softmax([0., -1000.])[1])
assert softmax([0., -1000.])[1] == 0.0
for bad in ([], [math.nan], [math.inf]):
try:
softmax(bad)
except ValueError:
pass
else:
raise AssertionError("Invalid input accepted")
print([round(x, 6) for x in p])
print(log_softmax([0., -1000.]))
This original Python example was executed. Checks cover normalization, shift invariance, a single candidate, underflow versus a finite log result, and empty or nonfinite inputs. It prints [0.090031,0.244728,0.665241] and the log probabilities [0,-1000] for the second example.
No model was trained and no GPU kernel was benchmarked. The code takes O(n) time and O(n) extra space for n candidates, treating each basic floating-point operation as constant cost.
Attention masks add another boundary
Production code often uses negative infinity for masked positions. This implementation accepts only finite logits and does not support that convention. A fully masked row would subtract negative infinity from negative infinity; the ordinary shift does not define a useful result automatically.
Specify mask and empty-row handling in the calling code, and use a tested implementation from the chosen framework. When investigating NaNs, check inputs, normalization axis and fully masked rows before precision. Subtracting the maximum cannot repair an upstream NaN. Unlike attention’s √d scaling, shifting preserves probabilities; scaling changes the distribution.
Sources
- eprints.maths.manchester.ac.uk — original paper / versioned documentation
- github.com — original paper / versioned documentation
Backfilled on October 13, 2026. The article date retains its planned Beijing slot. References predate that slot; the examples are original teaching calculations.
Support
If this article helped you, you can support this site.
Click a QR code to enlarge. More options: support page。


