log-sum-exp trick

Goal: compute $\log \left({\sum }_{i=1}^n{e}^{x_i}\right)$
Problem: numerical stability for large positive or negative $x$
Solution:

$$\begin{array}{rl}\log \left(\sum _{i=1}^n{e}^{x_i}\right)& =\log \left(e^c\cdot e^{-c}\sum _{i=1}^n{e}^{x_i}\right)\\ & =\log \left(e^c\sum _{i=1}^n{e}^{x_i-c}\right)\\ & =c+\log \left(\sum _{i=1}^n{e}^{x_i-c}\right)\end{array}$$

where $c={\max }_i(x_i)$

This way:

• all exponents are $\le 0$
• one term will have exponent $=0$, where $x_i=c$
• → small numbers, but no underflow/overflow

---

normalized_logits = logits - logits.logsumexp(dim=1, keepdim=True)
softmax = normalized_logits.exp() # exp(logits)/sum(exp(logits))

References

logarithm, exponential, softmax

https://gregorygundersen.com/blog/2020/02/09/log-sum-exp/