softmax数值稳定性问题以及CrossEntropyWithLogits的由来

softmax自身导致的数值问题

对于 $x=[x_1,x_2,\cdots,x_n]$，softmax公式：

$softmax(x)=[a_1,a_2,\cdots,a_n], \quad a_i=\frac{e^{x_i}}{\sum_{j=1}^{n}e^{x_j}}$

原始softmax公式的数值问题：

• 若 $x_i$ 过大， $e^{x_i}$ 上溢产生nan
• 若 $x_i$ 过小， $e^{x_i}$ 下溢，结果为0，若所有的$x_i$都过小，则导致 $a_i$ 的分母为0

解决办法：

先减去所有元素中的最大值，记为 $\max x$，再进行标准的softmax。这在数学上是等价的，通过对原始softmax公式分子分母同时除以 $e^{\max x}$ 即可证明。即：

$a_i=\frac{e^{x_i}}{\sum_{j=1}^{n}e^{x_j}}=\frac{e^{x_i}/e^{\max x}}{\sum_{j=1}^{n}e^{x_j}/e^{\max x}}=\frac{e^{x_i-\max x}}{\sum_{j=1}^{n}e^{x_j-\max x}}$

经过减去最大值的步骤后，最大的指数为0，因此不会再遇到第一个问题；而由于该指数0的存在，分母至少为1，因此也不会有第二个问题。

代码实现：

import numpy as np
def softmax(x,axis=-1):
    ex=np.exp(x-np.max(x,axis=axis,keepdims=True))
    return ex/np.sum(ex,axis=axis,keepdims=True)

softmax与交叉熵损失一起使用时导致的数值问题

但是注意，这种处理并不能避免 $x_i$ 过小， $e^{x_i}$ 下溢，结果为0这个现象。而且，假如 $\max x$ 是正数，还会加剧这个现象。

其实这个现象本身没有问题，不会导致 softmax 的结果在数学上产生大的误差，但是我们经常会对 softmax 函数的结果使用交叉熵损失，其中含有对数函数 $\log$，而 $\log(0)=-∞$，因此会导致数值问题。

所以，一般的深度学习框架都会建议使用将softmax和交叉熵一起计算的损失函数，例如Tensorflow中的 tf.nn.softmax_cross_entropy_with_logits，其底层实现可以参考 这个帖子

这样做的理由是：

$\begin{aligned} CrossEntropy&=\sum_i y_i \log(a_i)\\ &=\sum_i y_i \log\left(\frac{e^{x_i-\max x}}{\sum_{j=1}^{n}e^{x_j-\max x}}\right)\\ &=\sum_iy_i\left[(x_i-\max x)-\log \sum_{j=1}^{n}e^{x_j-\max x} \right] \end{aligned}$

可以看到 softamx 中分子的指数运算和交叉熵的对数运算抵消了，所以即使 $x_i-\max x$ 是一个很小的负数，也会在Loss中直接得到这个数，而不会由于经过指数运算后下溢为 0，进而得到 $\log(0)=-∞$