Numerical stability of softmax

NOTE

This is part of the 7th homework, titled Deep Learning I, for the course Machine Learning (IN2064) in the Winter Semester 2024/25 at TUM.

Problem 3:

In machine learning you often come across problems which contain the following quantity:

$$y=\log \sum _{i=1}^N{e}^{x_i}$$

For example, if we want to calculate the log-likelihood of a neural network with a softmax output, we get this quantity due to the normalisation constant. If you try to calculate it naively, you will quickly encounter underflows or overflows, depending on the scale of $x_i$. Despite working in log-space, the limited precision of computers is not enough, and the result will be $\infty$ or $-\infty$. To combat this issue, we typically use the following identity:

$$y=\log \sum _{i=1}^N{e}^{x_i}=a+\log \sum _{i=1}^N{e}^{x_i-a}$$

for an arbitrary $a$. This means you can shift the center of the exponential sum. A typical value is setting $a$ to the maximum ($a={\max }_i{x}_i$), which forces the greatest value to be zero, and even if the other values would underflow, you get a reasonable result.

Your task is to show that the identity holds.

Solution:

$$\begin{array}{rl}a+\log \sum _{i=1}^N{e}^{x_i-a}& =a+\log \sum _{i=1}^N\frac{e^{x_i}}{e^a}\\ & =a+\log \left(\frac{1}{e^a}\sum _{i=1}^N{e}^{x_i}\right)\\ & =a+\log (e^{-a}\cdot \sum _{i=1}^N{e}^{x_i})\\ & \stackrel{\ast }{=}a+\log (e^{-a})+\log \sum _{i=1}^N{e}^{x_i}\\ & =a-a+\log \sum _{i=1}^N{e}^{x_i}\\ & =\log \sum _{i=1}^N{e}^{x_i}.\end{array}$$

In ${}^{\ast }$ we used the fact that $\log (M\cdot N)=\log (M)+\log (N)$. In this manner we have proved that LHS = RHS. Q.E.D.

Problem 4:

Similar to the previous exercise, we can compute the output of the softmax function

$${\pi }_i=\frac{e^{x_i}}{{\sum }_{i=1}^N{e}^{x_i}}$$

in a numerically stable way by shifting by an arbitrary constant (a):

$$\frac{e^{x_i}}{{\sum }_{i=1}^N{e}^{x_i}}=\frac{e^{x_i-a}}{{\sum }_{i=1}^N{e}^{x_i-a}}.$$

Often, $a={\max }_i{x}_i$. Show that the above identity holds.

Solution:

The proof is rather trivial:

$$\begin{array}{r}\frac{e^{x_i}}{{\sum }_{i=1}^N{e}^{x_i}}=\frac{e^{-a}{e}^{x_i}}{e^{-a}{\sum }_{i=1}^N{e}^{x_i}}=\frac{e^{x_i-a}}{{\sum }_{i=1}^N{e}^{x_i-a}}\end{array}$$

Q.E.D