为什么大模型损失函数采用交叉熵，而不是 MSE？
大语言模型在训练时，本质是在做分类任务，即给定前文，预测下一个 token 的概率分布。
例如：模型需要从 30,000 个词中选出最可能的“下一个词”，这就是一个30k 分类问题。
所以，问 “为什么大模型损失函数采用交叉熵，而不是 MSE？” 实质上就等价于问“为什么分类模型采用交叉熵，而不是 MSE？”
例如在三分类问题中，假定某个样本的标签为 [0,0,1]，按道理说我们使用MSE作为损失，使得模型预测的概率分布逼近 [0,0,1] 不就行了？
那为什么分类任务又不用MSE呢？
对此，最主流的解释就是，因为分类模型输出的是一个概率分布，而衡量两个概率分布相似度最好的方法就是交叉熵，因此应该使用交叉熵。
不过今天，我们打算从另外一个角度——反向传播更新梯度——来看一下为什么分类模型要选择交叉熵。
最后你会惊奇的发型， 邪修版 Softmax + MSE 的组合居然能后 Softmax + 交叉熵的组合达到一样的效果。
摘要：交叉熵、大模型、损失函数、均方差
1. 除了交叉熵还有什么？#
既然是说为什么要选择交叉熵，那一定还有其它选项。
为了进行更好的对比，这里列了几种常见的组合，我们一起来进行分析并实验。
① Softmax + 交叉熵
② Sigmoid+交叉熵
③ Softmax + MSE
④ Sigmoid+MSE
其中第①种就是目前的分类任务中优选答案。
2. 网络优化的到底是什么？#
根据反向传播梯度的计算公式可知，此时定义：$L$ 表示神经网络总共包含的层数，$S_l$ 表示第$l$层的神经元数目，$K$ 表示输出层的神经元数目，$w_{ij}^l$ 表示第$l$层第$j$个神经元与第$l+1$层第$i$个神经元之间的权重值。
则前向传播过程为
$$ \begin{aligned} z^{l+1}_i&=a^l_1w^l_{i1}+a^l_2w^l_{i2}+\cdots+a^l_{S_l}w^l_{iS_l}+b^l\\[1ex] z^{l+1}&=a^lw^l+b^l\\[1ex] a^{l+1}&=f(z^{l+1}) \end{aligned}\tag{1} $$反向传播求解梯度过程为
$$ \frac{\partial J}{\partial w^l}=(a^l)^T\otimes\delta^{l+1} \tag {2} $$$$ \frac{\partial J}{\partial b^l}=\delta^{l+1} \tag {3} $$$$ \delta^l=\delta^{l+1}\otimes(w^l)^T\odot f^{\prime}(z^l),\;\;(0 < l \leq L-1)\tag {4} $$其中$\otimes$表示矩阵乘法，$\odot$表示按位乘操作，$f$ 表示每一层对应的激活函数。
此时可以发现，在反向传播过程中，上述式(2)到式(4)的过程对于使用任何形式的目标函数来说没有任何区别，唯独产生差异的就是
$$ \delta^L_i=\frac{\partial J}{\partial z_j}\tag{5} $$3. 使用不同损失函数的梯度分析#
采用不同方式的目标函数，那么计算得到的 $\delta^L_i$ 就会有很大区别，进而就会影响后续整个梯度计算的结果。
此时，如果使用交叉熵作为损失函数，那么有
① Softmax + 交叉熵时，目标函数 $J$ 对最后一层 $z^L$ 的梯度为
$$ \delta^L={\color{red}{(a^L-y)}},a^L=\text{softmax}(z^L)\tag{6} $$② Sigmoid + 交叉熵时，目标函数 $J$ 对最后一层 $z^L$ 的梯度为
$$ \delta^L=y\odot{\color{red}{a^L-y}},\;a^L=\text{sigmoid}(z^L)\tag{7} $$$$ \begin{aligned} \delta^{L}_i&=\frac{\partial J}{\partial z^L_i}=\frac{\partial }{\partial z^L_i}\left[-\sum_{k=1}^{S_L}y_k\cdot\log{a^L_k}\right]\text{sigmoid}^{\prime}(z^L_i)\\[3ex] &=-y_i\frac{1}{a^L_i}a^L_i(1-a^L_i)=y_ia^L_i-y_i \end{aligned} $$
如果使用MSE作为损失函数，那么有
③ Softmax + MSE时，目标函数 $J$ 对最后一层 $z^L$ 的梯度为
$$ \delta^L=a^L\odot{\color{red}{(a^L-y)}}-\phi\cdot a^L,a^L=\text{softmax}(z^L),\phi=\sum a^L_i(a^L_i-y_i)\tag{8} $$$$ \begin{aligned} \delta^{L}_j&=\frac{\partial J}{\partial z^L_j}=\sum_{i=1}^{S_L}\frac{\partial J}{\partial a^L_i}\frac{\partial a^L_i}{\partial z^L_j}=\sum_{i\neq j}\frac{\partial J}{\partial a^L_i}\frac{\partial a^L_i}{\partial z^L_j}+\frac{\partial J}{\partial a^L_i}\frac{\partial a^L_i}{\partial z^L_j}\\[2ex] &=\left[\sum_{i\neq j}-(a^L_i-y_i)a^L_ia^L_{j}\right]+\left[(a^L_j-y_j)a^L_{j}(1-a^L_{j})\right]\\[2ex] &=-a^L_{j}\left[\sum_{i\neq j}(a^L_i-y_i)a^L_i\right]+\left[(a^L_j-y_j)a^L_{j}-(a^L_j-y_j)a^L_{j}a^L_{j})\right]\\[2ex] &=-a^L_{j}\left[\sum_{i\neq j}(a^L_i-y_i)a^L_i+(a^L_j-y_j)a^L_{j}\right]+\left[(a^L_j-y_j)a^L_{j}\right]\\[2ex] &=-a^L_{j}\left[\sum_{i=1}^{S_L}(a^L_i-y_i)a^L_i\right]+\left[(a^L_j-y_j)a^L_{j}\right]\\[2ex] &=a^L_{j}(a^L_j-y_j)-a^L_{j}\left[\sum_{i=1}^{S_L}a^L_i(a^L_i-y_i)\right] \end{aligned} $$
④ Sigmoid + MSE时，目标函数 $J$ 对最后一层 $z^L$ 的梯度为
$$ \delta^{L}={\color{red}{(a^L-y)}}\odot \text{sigmoid}^{\prime}(z^L)={\color{red}{(a^L-y)}}\odot a^L\odot(1-a^L),\;a^L=\text{sigmoid}(z^L)\tag {9} $$在这4种情况中我们可以看出，第①种 Softmax + 交叉熵时，$\delta^L$ 的幅度是最大的；而第②种 Sigmoid + 交叉熵时，$\delta^L$ 的幅度是最小的。
4. 理论分析为什么Softmax + 交叉熵更好#
下面分别做一个简单的分析。
① Softmax + 交叉熵时，$\delta^L={\color{red}{(a^L-y)}}$,$a^L=\text{softmax}(z^L)$ ，假定 $a^L=[0.2894, 0.3199, 0.3907]$ 且 $y=[0,0,1]$，那么此时第0、1两个类别对应的残差明显会小于第2个类别的残差，进一步第2个类别对应的梯度就会增大，使得模型学到更高的置信度。相反，如果此时 $y=[1,0,0]$，那么此时第0个类别的残差是最大的，模型则会更大幅度调整第0个类别对应的参数梯度。同时，如果 $a^L=[0.0894, 0.0199, 0.8907]$ 且 $y=[0,0,1]$ 那么每个类别对应的残差都非常小，对应梯度也会相对更小，那么此时该样本对模型的学习过程影响便会减小，而这也是我们想要的情况，因为模型已经能够正确分类。
② Sigmoid + 交叉熵时 $\delta^L=y\odot{\color{red}{a^L-y}}$,$\;a^L=\text{sigmoid}(z^L)$，当 $\text{sigmoid}$ 接近于1时，则 $\delta^L=0$，这意味着后续所有梯度均为0，无法进行学习；当$\text{sigmoid}$ 接近于0时，$\delta^L=-y$ ，此时虽然对应类别上有残差但是不管置信度如何都不会有缓和情况，因为残差一直不变，因此同样无法有效学习。
③ Softmax + MSE时，$\delta^L=a^L\odot{\color{red}{(a^L-y)}}-\phi\cdot a^L$,$a^L=\text{softmax}(z^L)$,$\phi=\sum a^L_i(a^L_i-y_i)$，首先红色部分依旧满足 Softmax + 交叉熵时的情况，其次尽管乘以了 $a^L$ 以及减去了 $\phi\cdot a^L$ 部分，但是整体上依旧满足第①种时的特性，只是学习速度会变慢，因为新增加的部分会减弱残差的强度，但性质没变。
④ Sigmoid + MSE时，$\delta^{L}={\color{red}{(a^L-y)}}\odot a^L\odot(1-a^L)$,$\;a^L=\text{sigmoid}(z^L)$ ，首先红色部分依旧是有梯度的，且也满足类似 Softmax + 交叉熵时的性质，但是后面两部分其对应的是 $\text{sigmoid}$ 的梯度，根据图1可知其值域为 $[0,0.25]$，会限制残差，因此也会减缓模型的学习速度。
总结就是，“① Softmax + 交叉熵” 是最优选择；其次 “③ Softmax + MSE” 和 “④ Sigmoid + MSE” 也可以，但是收敛速度会变慢；最后，“ ② Sigmoid + 交叉熵” 对应的情况是最糟糕的，很有可能得不到学习。
5. 实验论证Softmax + 交叉熵就是好#
上面是理论分析，下面我们进行一个简单的实验。
这里以一个简单的2层神经网络对MNIST分类为了，分别就这4种情况进行实验，然后看看模型在准确率上的变化情况。
从图2中可以看出，① Softmax + 交叉熵 对应的情况是最好的；②Sigmoid + 交叉熵 对应的情况是最差的；③ Softmax + MSE 和 ④ Sigmoid + MSE 也能进行学习，只是在速度以及最终的效果上不如Softmax + 交叉熵。
到这里，我们总算从理论+实践把为什么要用 Softmax + 交叉熵 的原因给梳理一遍，但是内容并没有结束。
6 邪修版 Softmax + MSE 一样强#
我很好奇！
好奇是不是我只要构造出了和使用 Softmax + 交叉熵时的 $\delta^L$ 对应的残差，那么它就能达到和Softmax + 交叉熵一样的效果？
对于上面③ Softmax + MSE 和 ④ Sigmoid + MSE这两种情况，其对应的 $\delta^L=a^L\odot{\color{red}{(a^L-y)}}-\phi\cdot a^L$ 和 $\delta^{L}={\color{red}{(a^L-y)}}\odot a^L\odot(1-a^L)$ 。其中红色部分是因为使用了MSE而来，剩下部分则是由 $\text{softmax}$ 和 $\text{sigmoid}$ 对应的梯度产生而来。
想一想，如果在前向传播计算 $a^L$ 时保留使用 $\text{softmax}$ 和 $\text{sigmoid}$ ，但是在反向传播计算梯度时去掉，那么是不是 $\delta^L$ 就都只剩下了 ${\color{red}{(a^L-y)}}$ 这部分了？
根据式(2)到式(5)的过程可知，不管使用什么损失函数，只要 $\delta^L$ 一样，那么后续过程计算得到的结果都是一样的。
换句话说，计算出来的损失值可能有效衡量模型的优劣，不同的损失函数精度肯定不同。例如相较于其它损失函数，交叉损失是一个更好的选择。
但是，站在优化的角度，反向传播关心的是每一层对应的残差 $\delta^l$ 。
基于这样的猜想，我们分别实现③和④对应的无梯度版本，⑤ Nograde_Sofmax + MSE 和 ⑥ Nograde_Sigmoid + MSE 这两种情况，然后看看其准确率的变化。
在这之前，从理论上看 ①和⑤效果应该一样，⑥在最后。
如图3所示便是①⑤⑥这3种情况的实验对比结果，可以发现三者已经非常接近了，几乎重合。这就说明不管前向传播使用 $\text{softmax}$ 还是 $\text{sigmoid}$ ，只要 $\delta^L=a^L-y$ 都能有不错的效果。
我们再继续放大看一看！
如图4所示，经过放大后我们发现，至少在一开始的时候①和⑤是在上面的，也就是说使用 $\text{softmax}$ 比 $\text{sigmoid}$ 速度更快。
但是，意外的是①和⑤居然没有重合，这不科学，和理论分析不一样！
经过细致排查后发现，原来是两个模型运行时候初始化的权重值不一样导致的。于是我们又再次实验，终于得到了重合的曲线。
如图5所示便是重新运行后的对比结果，同时还去掉了平滑处理。可以看到此时两条曲线几乎已经重合，略微差异之处可能在于浮点数的处理上。
最后，根据上述理论分析及实验论证我们可以得出的结论就是，在反向传播及使用最原始的梯度下降优化模型时，只要能够构造得到相同的 $\delta^L$ ，那么最后优化得到的结果就是一样的，而所谓的损失函数仅仅只是用来衡量拟合的程度，本质上决定模型应该往何处优化的还是梯度。
尽管如此，但是你知道为什么现在更多的还是使用 $\text{softmax}$ + 交叉熵这一组合吗？