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)=-∞$