Numerically Stable logsumexp

Log-sum-exp computes the natural logarithm of a sum of exponentials:

logsumexp(x) = log(sum(exp(x)))

It appears in the normalization term for softmax and in calculations that combine probabilities represented by their logarithms. For a vector input, reducing across all entries produces one scalar.

Why the direct expression can fail

In float32, exp(100) overflows to infinity. So directly evaluating log(sum(exp([100, 100, 100]))) gives infinity, even though the answer is 100 + log(3), about 101.098612.

Very negative inputs cause the opposite problem: exp(-1000) underflows to zero, so taking the logarithm afterwards gives negative infinity.

Shift before exponentiating

For finite inputs, let m = max(x). Then use:

m + log(sum(exp(x - m)))

The shifted exponents are nonpositive, and at least one equals zero. This avoids overflowing those exponentials and keeps their sum nonzero. Adding m back restores the original scale. The identity is exact in real arithmetic; floating-point evaluation still rounds.

import jax
import jax.numpy as jnp

x = jnp.array([100.0, 100.0, 100.0], dtype=jnp.float32)
print(jax.nn.logsumexp(x))  # approximately 101.09861
x = jnp.array([-1000.0, -1000.0], dtype=jnp.float32)
print(jax.nn.logsumexp(x))  # approximately -999.3068

Turn scores into probabilities

Softmax converts a list of scores, also called logits, into probabilities that sum to one. It exponentiates each score and divides by the sum:

$$p_i=\frac{\exp(z_i/T)}{\sum_j\exp(z_j/T)},\qquad T>0.$$

T is the temperature. With the usual choice T = 1, scores [0, 0] give probabilities [0.5, 0.5]. Scores [0, log(3)] give [0.25, 0.75], because their exponentials are [1, 3].

Lower temperature puts more probability on the largest scores. Higher temperature makes the probabilities more even. A positive temperature preserves the ordering of finite scores.

Subtract the maximum before exponentiating

Large positive scores can make exponentials overflow. Subtracting the largest scaled score avoids large positive exponents without changing the probabilities: the common exponential factor cancels between numerator and denominator.

For a finite, nonempty list of scaled scores, every shifted score is at most zero and at least one is exactly zero. Their exponentials are therefore at most one, with at least one equal to one. For example, [1000, 1000] becomes [0, 0] before exponentiation and still gives [0.5, 0.5].

Choose the axis that contains the alternatives

For logits shaped (batch, classes), apply softmax over classes. Each example then gets its own probability distribution. Summing over the entire tensor would incorrectly make examples compete with each other.

When implementing the division yourself, keep the summed axis as a size-one dimension so the denominator broadcasts back across the intended scores.

If you need log probabilities, use log-softmax or a log-sum-exp formulation. Taking the logarithm after softmax can produce negative infinity when a small probability has already rounded to zero.

For masked softmax, each row needs at least one allowed, finite score. A row containing only negative infinity has no valid choices to normalise.