Numerical stability in exponentials and softmax

Softmax converts a vector of logits into nonnegative values that sum to one. Directly exponentiating large logits can overflow, even when the final normalized result is perfectly representable. Subtracting the maximum before exponentiation avoids this problem for ordinary finite input ranges.

Use the original teaching logits (1000,1001,1002). Their absolute level is large, but their differences are only one and two. Those differences determine the softmax result.

Cancel the common scale before computing it

For each logit z_i, softmax is exp(z_i)/sum(exp(z_j)). Multiplying every numerator and the denominator by exp(-m), where m is the largest logit, leaves the ratio unchanged.

The shifted logits become (-2,-1,0). Their exponentials are small finite values, including one for the largest entry. The normalized probabilities are approximately (0.090031,0.244728,0.665241).

The SciPy softmax documentation describes its shifted computation. Use a maintained numerical implementation in application code; the small helper here makes the operation visible for learning.

Verify the invariant and the result

import numpy as np
from scipy.special import softmax, logsumexp
from inference_core import stable_probabilities

z = np.array([1000.,1001.,1002.])
p = stable_probabilities(z)
assert np.allclose(p,[.09003057317038046,.24472847105479764,.6652409557748218])
assert np.allclose(p,softmax(z))
assert np.allclose(p,stable_probabilities(z-1000))
assert np.isfinite(p).all() and np.all(p>=0) and np.isclose(p.sum(),1)
assert np.isclose(logsumexp(z),1002.4076059644444)
assert np.allclose(z-logsumexp(z),np.log(p))
print({'softmax':p.tolist(),'log_normalizer':float(logsumexp(z))})

Run from the mathematics lab. The comparisons check both a known reference vector and agreement with SciPy, rather than accepting a result merely because its entries sum to one.

Use log-space calculations when that is the actual target

The related stable log-sum-exp calculation is m + log(sum(exp(z-m))). If a loss needs log probabilities, computing them directly as z-logsumexp(z) can avoid losing a tiny probability to underflow and then taking its logarithm.

Very negative shifted logits can underflow to zero in floating-point arithmetic. Whether that is acceptable depends on the downstream calculation. Taking log(0) after softmax is not equivalent numerically to a stable log-softmax calculation, even if the formulas are equivalent in exact arithmetic.

The teaching helper accepts a finite, nonempty one-dimensional vector and rejects nonfinite inputs. Its demonstrated range is the supplied examples; subtracting opposite-sign values near the largest representable float can itself overflow. This is a reason to test actual numerical ranges and use established routines, not to claim that one algebraic trick eliminates every floating-point issue.

Numerical correctness does not establish probability quality

A vector can be finite, nonnegative and normalized while being badly calibrated for the real task. The largest softmax value is not automatically the observed correctness rate among predictions receiving that score.

Calibration needs suitable labelled evaluation data. Distribution shift, class definitions and training choices can all affect the interpretation. Changing a temperature or clipping outputs also changes behavior and should be evaluated rather than presented as a purely cosmetic numerical fix.

Exercise: test logits (0,0,0), then (0,-1000,-2000). Compare softmax followed by a logarithm with direct log-space results. Explain the difference between an underflowed probability and an impossible outcome in the statistical model.

NeuraPath's Data Science course connects numerical methods with trustworthy evaluation. Check that a computation is stable, then separately check whether its predictions deserve the interpretation you assign them.

Continue learning

This article is part of the Mathematics and statistical foundations sequence. Use the neighbouring tasks when you need the prerequisite or the next application.

• Review the prerequisite or neighbouring task in Log transformations: interpret predictions on the original scale.
• Continue with Build a mathematics revision plan from model failure cases.