Posts

A note on log-sum-exp

Subtracting the maximum turns an overflow into an exact identity. The same shift makes softmax safe.

October 2, 2024 2 min read By Diogo Ribeiro

Log-likelihoods, mixture models, and softmax layers all need the same quantity,

\operatorname{LSE}(x_1, \dots, x_n) = \log \sum_{i=1}^{n} e^{x_i},

and the direct formula fails at the first large argument. In double precision, $e^{x}$ overflows to infinity once $x$ exceeds about 709.78, and underflows to zero once $x$ falls below about $-745$ . Log-likelihoods of a few thousand observations easily reach either range.

The fix is one line of algebra. For any constant $m$ ,

\log \sum_{i=1}^{n} e^{x_i} = m + \log \sum_{i=1}^{n} e^{x_i - m},

and choosing $m = \max_i x_i$ makes the largest exponent exactly zero. Every term is then at most one, so nothing overflows, and the largest term is exactly one, so the sum is at least one and its logarithm is finite.

python

import math

def logsumexp(xs):
    m = max(xs)
    if m == -math.inf:
        return -math.inf
    return m + math.log(sum(math.exp(x - m) for x in xs))

For $x = (1000, 1001, 1002)$ , the direct formula returns inf, while the shifted one returns 1002.4076059644444. For $x = (-1000, -1001, -1002)$ , every exponential underflows to zero and the direct formula returns -inf; the shifted one returns -999.5923940355556. Both answers differ by exactly 2002, as the symmetry of the inputs requires.

The guard for $m = -\infty$ matters when every term has zero probability: without it, $x - m$ is $-\infty - (-\infty)$ , which is not a number.

Libraries already do this: SciPy provides scipy.special.logsumexp, and deep-learning frameworks implement softmax and cross-entropy with the shift built in. The reason to know the trick is the code they do not cover, such as a hand-written likelihood, a log-space dynamic program, or a reduction over a custom axis.