面试常问系列(一)-神经网络参数初始化-之-softmax

背景

本文内容还是对之前关于面试题transformer的一个延伸，详细讲解一下softmax
面试常问系列(二)-神经网络参数初始化之自注意力机制-CSDN博客

Softmax函数的梯度特性与输入值的幅度密切相关，这是Transformer中自注意力机制需要缩放点积结果的关键原因。以下从数学角度展开分析：

1. Softmax 函数回顾

给定输入向量 z = z₁, z₂, ..., zₖ，Softmax 输出概率为：

其中 S 是归一化因子。

2. 梯度计算目标

计算 Softmax 对输入 z 的梯度，即 对所有 i,j∈{1,...,k}。

3. 梯度推导

根据链式法则，对 σi​ 关于 zj​ 求导：

具体推到过程就不展示了，感兴趣的有需要的可以评论下。因为本次重点不是通用的softmax分析，而是偏实战分析。

4. 与交叉熵损失结合的梯度

在实际应用中，Softmax 通常与交叉熵损失 结合使用。此时梯度计算更简单：

其中 是真实标签的 one-hot 编码。

5. 推导

1. 交叉熵损失对 的梯度：

2. 通过链式法则：

3. 代入在上面求解出的：

• 当时，
• 当时，

4.合并上述结果

6. 梯度消失问题

• 极端输入值 ：若远大于其他，则，其他。此时：
  • 对的梯度：（若yk=1，梯度接近0）。
  • 对其他zi的梯度：，梯度趋近于0。
• 后果：梯度消失导致参数更新困难，模型难以训练。

7. 缩放的作用

在Transformer中，点积结果除以dk​​后：

• 输入值范围受限 ：缩放后的方差为1，避免极端值。
• 梯度稳定性提升 ：分布更均匀，和不会趋近于0，梯度保持有效。

5. 直观示例

• 未缩放：若dk=512，点积标准差结果可能达±22，Softmax输出接近0或1，梯度消失。
• 缩放后：点积结果范围约±5，σ(zi)分布平缓，梯度稳定。
• 这个示例在最开始的跳转链接有详细解释，可以参考。

总结

Softmax的梯度对输入值敏感，过大输入会导致梯度消失。Transformer通过除以dk​​控制点积方差，确保Softmax输入值合理，从而保持梯度稳定，提升训练效率。这一设计是深度学习中处理高维数据时的重要技巧。