Log-Sum-Exp & Log-Softmax

The problem

Language models are trained on the logarithm of probabilities. If the model gave probability $p$ to the right token, the training loss is $-\log p$ : close to 0 when $p$ is close to 1, and large when $p$ is small.

The obvious way, np.log(softmax(x)), breaks on very small probabilities. A probability such as $e^{-800}$ is too small for float64 and becomes exactly 0, and then log(0) is -inf, even though the true answer, $-800$ , is a perfectly ordinary number.

Example. For scores $x = (0, -800)$ , the second probability is about $e^{-800}$ . np.log(softmax(x)) returns [0, -inf]. The correct answer is [0, -800].

The fix is to never compute the tiny probability at all, and to work with logarithms the whole way. This needs one function, log-sum-exp.

Log-sum-exp

The LSE (log-sum-exp) of a list of scores is the log of the sum of their exponentials:

\mathrm{LSE}(x) = \log \sum_i e^{x_i}

where $x = (x_1, \ldots, x_n)$ is the list of scores (the logits).

Example. For $x = (2, 1, 0)$ : $e^2 + e^1 + e^0 \approx 7.389 + 2.718 + 1 = 11.107$ , and $\log 11.107 \approx 2.408$ .

Computed this way it overflows, like softmax: for $x = (1000, 1000)$ , $e^{1000}$ is inf. The fix is to pull the largest score $m$ out of the sum first:

\mathrm{LSE}(x) = m + \log \sum_i e^{x_i - m}

where $m = \max_i x_i$ is the largest score. After the subtraction the largest exponent is $e^0 = 1$ , so nothing overflows; and the sum contains that 1, so it is at least 1, and its log can never be $\log 0$ .

Example. For $x = (1000, 1000)$ : $m = 1000$ , the shifted scores are $(0, 0)$ , their exponentials add up to $1 + 1 = 2$ , so $\mathrm{LSE}(x) = 1000 + \log 2 \approx 1000.693$ .

For $x = (-1000, -1000)$ the naive formula has the opposite problem: $e^{-1000}$ is 0, so it computes $\log 0 = -\infty$ . With the shift, $m = -1000$ , the sum is again 2, and the result is $-1000 + \log 2 \approx -999.307$ .

Log-softmax

The log of the softmax can be written with LSE, with no division and no tiny probability in between:

\begin{aligned} \log \mathrm{softmax}(x)_i &= \log \frac{e^{x_i}}{\sum_j e^{x_j}} \\ &= x_i - \mathrm{LSE}(x) \end{aligned}

where $x_i$ is score $i$ . The log of a fraction is the log of the top minus the log of the bottom, and $\log e^{x_i} = x_i$ .

Example. For $x = (2, 1, 0)$ , $\mathrm{LSE}(x) \approx 2.408$ , so the log-softmax is $(2 - 2.408,\ 1 - 2.408,\ 0 - 2.408) = (-0.408, -1.408, -2.408)$ . These are the logs of the probabilities $(0.665, 0.245, 0.090)$ .

For $x = (0, -800)$ , $\mathrm{LSE}(x)$ is essentially $0$ , and the log-softmax is [0, -800], the right answer.

What to implement

logsumexp(x, axis=-1, keepdims=False): the LSE along axis. Like np.sum, it removes that axis from the shape, unless keepdims=True, in which case the axis stays with size 1. For a table of shape (4, 6) and axis=-1, the result has shape (4,), or (4, 1) with keepdims=True: one value per row.
log_softmax(x, axis=-1): the log-softmax along axis, with the same shape as x.

Requirements

Stable for very large and very negative scores: logsumexp([1000, 1000]) is $1000 + \log 2$ and logsumexp([-1000, -1000]) is $-1000 + \log 2$ .
Works when some scores are -inf. This is how attention hides positions it must not look at: $e^{-\infty} = 0$ , so they add nothing to the sum. In log_softmax, a -inf score stays -inf.
No Python loops.

Hint

Compute the maximum with keepdims=True, so it lines up with x when you subtract it, then remove the axis at the end if keepdims is False (np.squeeze).
log_softmax can call your logsumexp: with keepdims=True its result lines up with x.