Exp-normalize和LogSumExp技巧

softmax模型中，会计算出模型输出然后将此输出送入交叉熵损失。
对较大值进行指数运算时，可能会出现上溢 (overflow) 或下溢 (underflow) 问题。

Exp-normalize技巧

使用softmax预测概率分布时，
$${\hat{y}}_j=\frac{\exp \left(x_j\right)}{{\sum }_{k=1}^n\exp \left(x_k\right)}$$
其中${\hat{y}}_j$是预测的概率分布，由向量$\mathbf{x}\in {\mathbb{R}}^n$决定。如果$x_k$中的一些数值非常大， $exp(x_k)$可能会大于数据类型容许的最大数字，造成上溢 (overflow)，这将使这将使分母或分子变为$inf$ (无穷大)，最后得到的是0、$inf$或$nan$ (不是数字) 的 ${\hat{y}}_j$。 在这些情况下，我们无法得到一个明确定义的交叉熵值。这种情况下可以使用exp-normalize技巧避免数值上溢。
$${\hat{y}}_j=\frac{\exp \left(x_i-b\right)\exp (b)}{{\sum }_{k=1}^n\exp \left(x_k-b\right)\exp (b)}=\frac{\exp \left(x_i-b\right)}{{\sum }_{k=1}^n\exp \left(x_k-b\right)}$$
可以选择$b={\max }_{i=1}^n{x}_i$，这种情况下，${\hat{y}}_j$移位不变，不可能出现上溢，因为此时最大的指数为${\max }_{i=1}^n{x}_i-b=0$。

例子：

x = np.array([1, -10, 1000])
np.exp(x) / np.exp(x).sum()

RuntimeWarning: overflow encountered in exp
RuntimeWarning: invalid value encountered in true_divide
array([ 0.,  0., nan])

使用exp-normalize后：

def exp_normalize(x):
    b = x.max()
    y = np.exp(x - b)
    return y / y.sum()

exp_normalize(x)

array([0., 0., 1.])

上溢问题得到解决。

使用exp-normalize技巧的sigmoid函数

sigmoid 函数可以使用exp-normalize技巧进行计算，避免数值溢出。
$$sigmoid(x)=\frac{1}{1+\exp (-x)}=\frac{\exp (x)}{1+\exp (x)}$$
如果x很大，$exp(x)$会出现上溢；如果x很小，$exp(-x)$可能会出现上溢，为了避免这种情况，在$x>=0$时使用第一种表示，$x<0$时使用第二种。

def sigmoid(x):
    "Numerically stable sigmoid function."
    if x >= 0:
        z = exp(-x)
        return 1 / (1 + z)
    else:
        # if x is less than zero then z will be small, denom can't be
        # zero because it's 1+z.
        z = exp(x)
        return z / (1 + z)

LogSumExp技巧

经过减法和规范化步骤之后，可能有些$x_k-b$具有较大的负值，由于精度受限，$exp(x_k-b)$将有接近0的值，即下溢 (underflow) ，这些值可能会四舍五入为零，使${\hat{y}}_j$为零，并且使得$\log ({\hat{y}}_j)$的值为-inf。反向传播几步后，我们可能会发现自己面对一屏幕可怕的nan结果。
尽管我们要计算指数函数，但在后续计算交叉熵损失函数的过程中需要取他们的对数， 通过将softmax和交叉熵结合在一起，可以避免反向传播过程中可能会困扰我们的数值稳定性问题。

LogSumExp (LSE) 操作可以被表示为：
$$\mathrm{LSE}\left(x_1,\dots ,x_N\right)=\log \left(\sum _{n=1}^N\exp \left(x_n\right)\right)$$
且 max ⁡ { x 1 , … , x n } < LSE ⁡ ( x 1 , … , x n ) ≤ max ⁡ { x 1 , … , x n } + log ⁡ ( n ) \max \left\{x_{1}, \ldots, x_{n}\right\}<\operatorname{LSE}\left(x_{1}, \ldots, x_{n}\right) \leq \max \left\{x_{1}, \ldots, x_{n}\right\}+\log (n) max{x1​,…,xn​}<LSE(x1​,…,xn​)≤max{x1​,…,xn​}+log(n)

在softmax中使用这种技巧，将交叉熵和softmax结合在一起，对原式取对数得：
$$\begin{array}{rl}\log \left({\hat{y}}_j\right)& =\log \frac{\exp \left(x_j\right)}{{\sum }_{k=1}^n\exp \left(x_k\right)}\\ & =x_j-\log \sum _{k=1}^n\exp \left(x_k\right)\end{array}$$
此时$\log {\sum }_{k=1}^n\exp \left(x_k\right)$仍有可能出现上溢，故令
$$\begin{array}{rl}LogSumMax=\log \sum _{k=1}^n\exp \left(x_k\right)& =\log \sum _{k=1}^n\exp \left(x_k-b\right)\exp (b)\\ & =b+\log \sum _{k=1}^n\exp \left(x_k-b\right)\end{array}$$
此时$\log (softmax)$函数可表示为：
$$\begin{array}{rl}\log \left({\hat{y}}_j\right)& =x_j-LogSumMax\\ & =x_j-b-\log \sum _{k=1}^n\exp \left(x_k-b\right)\end{array}$$
$${\hat{y}}_j=\exp \left(x_j-LogSumMax\right)$$
对LogSumMax求导即可得到exp-normalize：
$$\frac{\partial \left(b+\log {\sum }_{k=1}^n\exp \left(x_k-b\right)\right)}{\partial x_j}=\frac{\exp \left(x_j-b\right)}{{\sum }_{k=1}^n\exp \left(x_k-b\right)}$$

如果需要对概率分布求log（如使用交叉熵损失的softmax）则使用LogSumMax，如果只求概率分布，使用exp-normalize即可。由于exp-normalize使用exp更少，在数值上更稳定。

例子：

def logsumexp(x):
    c = x.max()
    return c + np.log(np.sum(np.exp(x - c)))  

>>> x = np.array([1000, 1000, 1000])
>>> np.exp(x)
array([inf, inf, inf])

>>> logsumexp(x)
1001.0986122886682
>>> np.exp(x - logsumexp(x))
array([0.33333333, 0.33333333, 0.33333333])

>>> x = np.array([-1000, -1000, -1000])
>>> np.exp(x)
array([0., 0., 0.])
>>> np.exp(x - logsumexp(x))
array([0.33333333, 0.33333333, 0.33333333])

>>> x = np.array([-1000, -1000, 1000])
>>> np.exp(x - logsumexp(x))
array([0., 0., 1.])

exp-normalize和log-sum-max有时都被叫做"softmax"，但其实exp-normalize更接近"soft argmax"，LSE函数才是真正意义上的softmax函数。也就是说我们常常在神经网络中用的softmax其实是soft argmax。

Reference

1. Exp-normalize trick
2. The Log-Sum-Exp Trick