Softmax计算技巧

初始做法

在softmax回归中，定义
$$\begin{gathered} \hat{\mathbf{y}}=\mathrm{softmax}(\mathbf{o}){\hat{y}}_j=\frac{\exp (o_j)}{{\sum }_k\exp (o_k)}(1) \\ (i=\mathrm{1...}n,k=\mathrm{1...}q) \end{gathered}$$
对于任何标签 $y$ 和模型预测 $\hat{y}$ ，损失函数为:
$$l(\mathbf{y},\hat{\mathbf{y}})=-\sum _{j=1}^q{y}_j\log {\hat{y}}_j(2)$$
将 $(1)$ 代入 $(2)$ 中：
$$\begin{array}{r}\begin{array}{rl}l(\mathbf{y},\hat{\mathbf{y}})& =-\sum _{j=1}^q{y}_j\log \frac{\exp (o_j)}{{\sum }_{k=1}^q\exp (o_k)}\\ & =\sum _{j=1}^q{y}_j\log \sum _{k=1}^q\exp (o_k)-\sum _{j=1}^q{y}_j{o}_j\\ & =\log \sum _{k=1}^q\exp (o_k)-\sum _{j=1}^q{y}_j{o}_j.\end{array}\end{array}(3)$$
考虑相对于任何未规范化的预测 $o_j$ 的导数，我们得到：
$${\partial }_{o_j}l(\mathbf{y},\hat{\mathbf{y}})=\frac{\exp (o_j)}{{\sum }_{k=1}^q\exp (o_k)}-y_j=\mathrm{softmax}(\mathbf{o}{)}_j-y_j(4)$$

问题1： $\exp (o_k)$ 可能特别大或特别小

softmax函数 ${\hat{y}}_j=\frac{\exp (o_j)}{{\sum }_k\exp (o_k)}$， 其中 $o_j$ 是预测 $\mathbf{o}$ 的概率分布。$o_j$ 是未规范化的预测的第 $j$ 个元素。
如果 $o_j$ 中的一些数值非常大， 那么 $\exp (o_k)$ 可能大于数据类型容许的最大数字，即上溢（overflow）。 这将使分母或分子变为 $\inf$（无穷大）， $\frac{+\infty }{+\infty }$最后得到的是 $0$、$\inf$ 或 nan（不是数字）的 $\widehat{y_j}$。
另一方面 $\exp (o_k)$ 都特别小，${\sum }_k\exp (o_k)$ 在实际计算中为 $0$，这样就出现了 $0/0$ 的错误。
在这些情况下，我们无法得到一个明确定义的交叉熵值。

解决这个问题的一个技巧是： 在继续softmax计算之前，先从所有 $o_k$ 中减去 $\max (\mathbf{o})$。 这里可以看到每个 $o_k$ 按常数进行的移动不会改变softmax的返回值：
$$\begin{array}{r}\begin{array}{rl}{\hat{y}}_j& =\frac{\exp (o_j-\max (\mathbf{o}))\exp (\max (\mathbf{o}))}{{\sum }_k\exp (o_k-\max (\mathbf{o}))\exp (\max (\mathbf{o}))}\\ & =\frac{\exp (o_j-\max (\mathbf{o}))}{{\sum }_k\exp (o_k-\max (\mathbf{o}))}.\end{array}\end{array}(5)$$
这样分子$\max (\exp (o_j-\max (\mathbf{o})))=1$，分母 ${\sum }_k\exp (o_k-\max (\mathbf{o}))\ge 1$

问题2： $\log ({\hat{y}}_j)=-\infty$

在正向、反向传播中都要计算公式 $(2)$ 中的 $\log ({\hat{y}}_j)$，按照上面做归一化后，由于精度受限， $\exp (o_j-\max (o_k))$ 将有接近零的值，即下溢（underflow），此时 $\log ({\hat{y}}_j)=-\infty$ 。
但是实际上，这个问题在实际数据运算时可以避免掉，在数学上，有下面的运算（永远可行！）
$$\begin{array}{r}\begin{array}{rl}\log ({\hat{y}}_j)& =\log \left(\frac{\exp (o_j-\max (\mathbf{o}))}{{\sum }_k\exp (o_k-\max (\mathbf{o}))}\right)\\ & =\log (\exp (o_j-\max (\mathbf{o})))-\log \left(\sum _k\exp (o_k-\max (\mathbf{o}))\right)\\ & =o_j-\max (\mathbf{o})-\log \left(\sum _k\exp (o_k-\max (\mathbf{o}))\right).\end{array}\end{array}(6)$$
这样在计算 $\log ({\hat{y}}_j)$ 时不用先计算 ${\hat{y}}_j$ 然后计算 $\log ({\hat{y}}_j)$，而是计算公式 $(6)$ 的最后部分：
$$o_j-\max (\mathbf{o})-\log \left(\sum _k\exp (o_k-\max (\mathbf{o}))\right)$$
这一部分已经在机器运算时不会出现问题。

启发

由于数据结构的限制，数值计算在实际运算中要避免 $\log (0)$、$\frac{0}{0}$、$\frac{\infty }{\infty }$ 的情况，考虑在数学能不能进一步变换，找到等价的表达式是一个很好的思路，本文中的做法在其它问题中也可借鉴。

参考

3.7. softmax回归的简洁实现