softmax+cross-entropy的前向计算、反向传播的公式推导

本文主要是讲述Softmax和CrossEntropy的公式推导，并用代码进一步佐证。

1. Softmax前向计算

我们把$Softmax$输出的概率定义为 $p_i$：
$$Softmax(a_i)=p_i=\frac{e^{a_i}}{{\sum }_j^N{e}^{a_j}}$$
模型输出$[a_1,a_2,...,a_N]$，共N个值。
其中$a_i$代表第$i$个输出值，$p_i$代表第$i$个输出值经过$Softmax$计算过后的概率。
且 $p_1+p_2+...+p_N=1$

1.1 数值稳定

因为Softmax涉及到指数函数，且底数 $e$ 大于1，在计算机中是可能会有溢出风险的。结合指数、对数函数的转换规则，我们可以制定一些数值稳定的优化策略。（当然这些是在框架中实现的，学习更多是为了扩展视野）

数值稳定的主要思路在于$a_i$减去$A=[a_1,a_2,...,a_N]$中的最大值$max(A)$
$$\begin{array}{rl}{p}_i& =\frac{e^{a_i}}{{\sum }_j^N{e}^{a_j}}\\ & =\frac{C\cdot e^{a_i}}{C\cdot {\sum }_j^N{e}^{a_j}}\\ & =\frac{e^{\log (C)}\cdot e^{a_i}}{e^{\log (C)}\cdot {\sum }_j^N{e}^{a_j}}\\ & =\frac{e^{a_i+\log (C)}}{{\sum }_j^N{e}^{a_j+\log (C)}}\\ & =\frac{e^{a_i-max(A)}}{{\sum }_j^N{e}^{a_j-max(A)}}\end{array}$$
因为C是常数，$log(C)$也是常数，所以我们可以把$log(C)$定义为-max(A)，且并不会改变$p_i$的计算结果。
减去A中的最大值，就能确保A中所有的值都不会上溢出。

2. Cross-Entropy前向计算

我们把交叉熵公式定义为 $H$：
$$CrossEntropy(y_i,p_i)=H(y_i,p_i)=-\sum _i^N{y}_i\cdot \log (p_i)$$
在多分类问题中，我们的 label 通常以独热码（one-hot）的形式展现和学习，因此在 $Y=[y_1,y_2,...,y_N]$ 中，只有一项为 $1$，其余项为 $0$，即 $[0,0,...,1,...,0,0]$。
所以 $H(y_i,p_i)$ 也等于 $-y_i\cdot \log (p_i)$，$y_i=1$ 对应label的类别。

3. Softmax反向传播求导

据 $Softmax$ 公式可知，每个 $p_i$ 均是所有 $a$ 都有参与（分母）运算的，因此梯度的形式为：
$$\frac{\partial p_i}{\partial a_j}=\frac{\partial (\frac{e^{a_i}}{{\sum }_j^N{e}^{a_j}})}{\partial a_j}$$
而且$i$ 和 $j$ 的关系要分类讨论。

这里要先复习下含分母的求导公式：
$$(\frac{h(x)}{g(x)}{)}^{\prime }=\frac{h^{\prime }(x)\cdot g(x)-h(x)\cdot g^{\prime }(x)}{g(x{)}^2}$$
并且简化一下符号：
$$\sum _j^N{e}^{a_j}=\sum$$
当 $i=j$：
$$\begin{array}{rl}\frac{\partial p_i}{\partial a_j}& =\frac{e^{a_i}\cdot \sum -e^{a_i}\cdot e^{a_j}}{\sum \cdot \sum }\\ & =\frac{e^{a_i}\cdot (\sum -e^{a_j})}{\sum \cdot \sum }\\ & =p_i\cdot (1-p_j)\end{array}$$
当$i\ne j$（$e^{a_i}$相当于常数，导数为 0）：
$$\begin{array}{rl}\frac{\partial p_i}{\partial a_j}& =\frac{0\cdot \sum -e^{a_i}\cdot e^{a_j}}{\sum \cdot \sum }\\ & =-p_i\cdot p_j\end{array}$$

4. Cross-Entropy + Softmax反向传播求导

$$\frac{\partial H}{\partial a_j}=\frac{\partial H}{\partial p_i}\cdot \frac{\partial p_i}{\partial a_j}=-\sum _i{y}_i\frac{1}{p_i}\cdot \frac{\partial p_i}{\partial a_j}---①$$
当 $i=j$：
$$\begin{array}{rl}①& =-\sum _{i=j}{y}_i\frac{1}{p_i}\cdot p_i\cdot (1-p_j)\\ & =-\sum _{i=j}{y}_i\cdot (1-p_j)\\ & =-y_i+y_i{p}_j（因为只有i，可以把\sum 去掉）---②\end{array}$$
当$i\ne j$：
$$\begin{array}{rl}①& =-\sum _{i\ne j}{y}_i\frac{1}{p_i}\cdot (-p_i\cdot p_j)\\ & =\sum _{i\ne j}{y}_i{p}_j---③\end{array}$$
因为②和③其实是①的互斥情况，所以可以合并：
$$\begin{array}{rl}①& =②+③（记住在②中，i=j）\\ & =-y_i+y_i{p}_j+\sum _{i\ne j}{y}_i{p}_j\\ & =-y_i+(\sum _{i=j}{y}_i{p}_j+\sum _{i\ne j}{y}_i{p}_j)\\ & =-y_i+\sum _i^N{y}_i{p}_j（因为y_i是one-hot，\sum _i^N{y}_i=1）\\ & =p_j-y_j（y_i=y_j）\end{array}$$
整个Softmax+CrossEntropy的求导推导下来发现，$H$对于$a_j$的偏导值，就是让他的$p_j$去减对应的label值($y_j$)。
举例P = [0.5, 0.3, 0.2]，Y=[1, 0, 0]，对应的导数就是 [-0.5, 0.3, 0.2]。

5. 代码验证

先看下x在softmax+cross entorpy前向计算并且BP后，所产生的梯度是多少，即$\frac{\partial H}{\partial a_j}$，在这个例子中分别对a1，a2，a3求偏导：

x = torch.randn((1, 3), requires_grad=True)
# tensor([[-0.3876,  0.2697, -1.6527]], requires_grad=True)
y = torch.randint(3, (1,), dtype=torch.int64)
# tensor([1])

loss = F.cross_entropy(x, y) 
# F.cross_entropy含了softmax+cross_entropy
# 因此直接调用即可，无需先使用F.softmax
print(loss)
# tensor(0.5095, grad_fn=)

loss.backward()
print(x.grad)
# tensor([[ 0.3113, -0.3992,  0.0879]])

下面再看下 $p_i$ 的值：

F.softmax(x, dim=1)
# tensor([[0.3113, 0.6008, 0.0879]], grad_fn=)

发现没有！发现没有！除了 $a_1$ 比 $p_1$ 减了1之外，其他都没变！正正验证了上面的公式推导！

6. 总结

总结一下，在多分类问题中，softmax+cross entropy是比较普遍，且计算速度较快的损失函数（loss function），因为它的梯度仅仅只用把概率值（pi）减去标签（yi）即可！
这在训练的初期，可以提供较快的训练速度，以提供后续优化的方向。当然，后续也包括对损失函数的优化！