自适应 Softmax

Softmax 函数在 Transformer 注意力机制中至关重要，它通过求和归一化注意力分数的每一行为 1，从而优于其他替代函数。然而，当注意力分数的某些元素趋近极端值（如接近 1 或 0 的概率）时，softmax 函数会面临梯度消失问题。本文提出 Self-Adjust Softmax（SA-Softmax），通过将 $softmax(x)$ 修改为 $x \cdot softmax(x)$ 及其归一化变体 $\frac{(x - \min(x_{\min},0))}{\max(0,x_{max})-\min(x_{\min},0)} \cdot softmax(x)$ 来解决该问题。我们从理论上证明 SA-Softmax 相比标准 softmax 函数具有增强的梯度特性。此外，SA-Softmax Attention 可以无缝集成到现有 Transformer 模型的注意力机制中，仅需少量调整。我们进行了实验，评估使用 SA-Softmax 的 Transformer 模型相对于标准 softmax 函数的实证性能。这些实验涉及最多 27 亿参数的模型，涵盖了多样化的数据集、语言任务和位置编码方法。