Softmax与交叉熵的两大数值Bug及修复方案

在深度学习模型中，softmax、log-softmax和交叉熵是几乎每个分类网络都会用到的组件。然而，标准实现中隐藏着两个数值bug，它们可能导致梯度计算错误、训练不稳定，甚至影响模型泛化。最近一篇技术文章详细分析了这两个问题并提出了简洁的修复方案。

两个隐蔽的数值Bug

文章指出，标准实现存在两个缺陷：Softmax Collapse 和 Softmax Cancellation。

Softmax Collapse：梯度消失与不归零

当目标logit远大于其他logit时，标准softmax会将该logit的梯度舍入为零。这导致梯度之和不再为零，破坏了优化过程中的平衡。这种现象在模型高度自信时尤为明显，可能引发训练后期的异常行为。

Softmax 的梯度性质：对于 softmax 输出 (p_i = e^{z_i}/\sum_j e^{z_j})，损失对 logit (z_k) 的梯度可写为 (\partial L/\partial z_k = p_k (\partial L/\partial p_k) - p_k \sum_j p_j (\partial L/\partial p_j))。当使用交叉熵且目标为 (t) 时，梯度简化为 (p_k - \delta_{kt})。理论上所有分量之和恒为零，这是概率单纯形约束的必然结果。然而在浮点运算中，当 (p_t) 极度接近 1 时，(p_t - 1) 会因舍入误差变为 0，而其他 (p_k) 也因下溢为 0，导致梯度向量全零。这种“梯度消失”并非来自激活函数饱和，而是数值表示精度不足造成的。它破坏了梯度平衡，使优化器无法感知目标类与其他类的相对关系，从而在模型高度自信时引发训练异常。修复方案通常采用 log-softmax 并利用 log-sum-exp 技巧稳定计算，确保梯度在极端 logit 差异下仍保持正确符号与量级。

Softmax Cancellation：蒸馏中的虚假梯度

在知识蒸馏中，当学生logits接近教师logits时，标准实现会产生非零梯度——即使两者完全相等。这意味着学生网络会被推向错误的方向，导致漂移。文章将这种现象称为“Softmax Cancellation”，它解释了蒸馏中学生网络性能不稳定的部分原因。

知识蒸馏的梯度构成：蒸馏损失通常由硬标签交叉熵和软标签 KL 散度组成。软标签部分为 (L_{soft} = \tau^2 \mathrm{KL}(p_T^{\tau} | p_S^{\tau}))，其对学生的 logit 梯度为 (\tau (p_S^{\tau} - p_T^{\tau}))。当学生与教师的 logit 完全相等时，该梯度理论上应为零。但标准实现中，softmax 的分子分母分别计算，两个相等的大数相减会产生舍入误差，使得 (p_S - p_T) 并不严格为零，而是残留一个微小但非零的“虚假梯度”。这个梯度方向随机，会持续推动学生参数漂移，尤其在训练后期学生已充分拟合教师时，漂移会累积成性能退化。修复方案是联合计算 log-softmax 并直接对 logit 差值做减法，避免分别归一化再相减，从而保证相等输入产生严格零梯度。

修复方案与实证效果

作者提出了简单的log-softmax、softmax和交叉熵实现，同时修复了上述两个bug。这些修改不增加计算复杂度，却带来了多项训练动态的改善。

消除“Slingshot”现象

此前观察到的“Slingshot”现象——训练损失在某个阶段突然反弹——在修复后消失。这表明梯度计算的数值稳定性对训练轨迹有直接影响。

加速Grokking与防止精度崩溃

使用权重衰减时，修复后的实现加快了grokking（模型从记忆到泛化的突变）的出现；不使用权重衰减时，则防止了精度崩溃。这两点说明数值bug会干扰正则化效果和泛化行为。

Grokking 与权重衰减：Grokking 指模型在训练集上早已过拟合（训练损失接近零）后，经过长时间训练，验证集性能突然从随机水平跃升至泛化良好的现象。它通常出现在小数据集、强权重衰减的设定下，本质是优化器在权重衰减的驱动下逐渐找到更简单的解。数值 bug 会干扰这一过程：Softmax Collapse 导致的梯度失衡可能让模型过早锁定在记忆解，延迟或阻止向泛化解的过渡；而修复后梯度平衡恢复，权重衰减能更有效地引导参数向低范数区域移动，从而加速 grokking。另一方面，在不使用权重衰减时，logit 可能无界增长，标准 softmax 的数值下溢会使梯度完全消失，模型无法继续调整，表现为精度突然崩溃；修复后的实现通过稳定 log-softmax 保持梯度有效，防止了这种崩溃。

稳定知识蒸馏

修复后，学生网络不再从教师网络漂移，蒸馏过程更加稳定。这对于依赖软标签的任务尤为重要。

揭示logit增长的平衡机制

文章进一步发现，在模型记忆训练数据后，logit的增长受到权重衰减与Adam优化器中ε超参数之间平衡的约束。这一发现为理解训练后期的动力学提供了新视角。

Adam 的 ε 与 logit 范数：Adam 更新中，参数步长近似为 (\eta \cdot \hat{m}_t / (\sqrt{\hat{v}_t} + \epsilon))。当梯度持续较小时，(\sqrt{\hat{v}_t}) 也很小，分母中的 (\epsilon) 开始主导，实际步长被压缩为 (\eta \cdot \hat{m}_t / \epsilon)，与梯度大小近似成正比而非自适应缩放。在训练后期，权重衰减将参数拉向零，而分类损失推动 logit 增大以维持正确分类，两者形成拉锯。若没有权重衰减，logit 会持续增长直到数值极限；有权重衰减时，参数范数被约束，logit 增长到某一水平后，梯度中权重衰减项与分类损失项平衡，增长停止。(\epsilon) 决定了这个平衡点附近的有效步长，从而影响 logit 的最终量级。这一机制解释了为什么调整 (\epsilon) 会改变训练后期的动态，也为理解数值实现与优化器超参数的交互提供了新视角。

对实践者的启示

对于从事模型训练和调优的工程师来说，这些发现提醒我们：基础组件的数值实现细节可能产生深远影响。尤其是在知识蒸馏、大规模分类任务以及使用权重衰减的场景中，采用修复后的softmax和交叉熵实现可能带来更稳定的训练和更好的最终性能。

文章提供的实现简单直接，建议在自定义训练循环中替换标准实现，以规避这些隐蔽的数值问题。