Log-sum-exp: calculate model scores without exponential overflow

Log-sum-exp calculates log(exp(x₁) + ... + exp(xₙ)). Computing the exponentials directly can overflow even when the final logarithm is finite. Subtracting the largest input before exponentiating gives an equivalent expression with much safer intermediate values.

This pattern appears when combining log-probabilities, normalizing classification scores and implementing losses. It illustrates a general numerical lesson: algebraically equivalent formulas can behave differently in floating-point arithmetic. We will derive the shift, reproduce the failure and check a stable implementation using invented model scores.

Derive the maximum shift

Let m be the largest finite input. Factor exp(m) out of the sum, then take the logarithm:

log(sum(exp(xᵢ))) = m + log(sum(exp(xᵢ - m))).

Every shifted input is at most zero, so every shifted exponential is at most one. At least one is exactly one. Large positive exponentials disappear from the intermediate computation, while the final addition restores the original scale.

For scores [1000, 1001, 999], use m = 1001. The shifted scores are [-1, 0, -2]. Their exponentials are entirely ordinary numbers, even though exp(1001) is too large for a standard Python float.

Reproduce and repair the calculation

This script uses only the standard library. It accepts a nonempty sequence of finite scores; other cases need an explicit contract rather than accidental arithmetic.

from math import exp, log, isfinite, isclose

def log_sum_exp(values):
    values = list(values)
    if not values or not all(isfinite(value) for value in values):
        raise ValueError("provide nonempty finite scores")
    maximum = max(values)
    return maximum + log(sum(exp(value - maximum) for value in values))

scores = [1000.0, 1001.0, 999.0]
try:
    log(sum(exp(value) for value in scores))
except OverflowError:
    print("direct: overflow")

normalizer = log_sum_exp(scores)
print(f"log-sum-exp: {normalizer:.6f}")
probabilities = [exp(value - normalizer) for value in scores]
print("softmax:", [round(value, 6) for value in probabilities])
assert isclose(sum(probabilities), 1.0, abs_tol=1e-12)

shifted = [value + 5000 for value in scores]
assert isclose(log_sum_exp(shifted), normalizer + 5000)
shifted_probabilities = [exp(value - log_sum_exp(shifted)) for value in shifted]
assert all(isclose(a, b, abs_tol=1e-12)
           for a, b in zip(probabilities, shifted_probabilities))

Expected output:

direct: overflow
log-sum-exp: 1001.407606
softmax: [0.244728, 0.665241, 0.090031]

The middle score is largest, so it receives the largest normalized probability. Adding the same constant to all logits does not change their softmax probabilities; the final assertions check that invariance.

Connect the result to log-softmax

The log probability for class i is xᵢ - log_sum_exp(x). Keeping that quantity in log space avoids computing a tiny probability only to take its logarithm again. This matters for unlikely classes and for objectives that sum many log probabilities.

Libraries provide operations for this reason. SciPy exposes logsumexp, with axis and weighting options beyond this teaching implementation. Machine-learning frameworks also provide stable log-softmax and loss functions. Prefer the appropriate maintained primitive in production code, especially when gradients, batches or accelerators are involved.

Our example applies ordinary Python arithmetic to three scores. It does not implement differentiation, tensor axes, masking or weighted sums. Replacing a framework's loss with this function would discard those requirements.

Stability has boundaries

The maximum shift removes the obvious exponential overflow. It does not make floating-point arithmetic exact. Very negative shifted values can underflow to zero, and adding very unequal positive values can lose small contributions. Whether that error matters depends on the calculation and precision requirements.

The finite-input restriction is deliberate. If all scores are negative infinity, subtracting the maximum would involve an undefined infinity subtraction. Empty sequences, NaNs, masked classes and positive infinities also need defined behavior. A library primitive may support those cases differently; inspect its documented contract before adapting a result.

Precision matters even in the stable form. To compute a small log probability accurately, a direct log-softmax primitive can be preferable to assembling intermediate values in a lower-precision workflow. Test the actual dtype and framework configuration instead of assuming a successful float64 demonstration proves every implementation is stable.

Check the method on smaller values

Change the inputs to [1, 2, 3] and compare the direct and stable formulas. They should agree within floating-point tolerance. Then decrease all three values by 10,000. The direct expression can fail through underflow, while the shifted expression should remain usable.

Continue through Maths resources and Statistics resources to connect numerical calculations with the probabilistic quantities they represent.