arXiv:2502.18277cs.CL2025-02EMNLP被引 7

改进Transformer注意力的Softmax,解决梯度消失问题。

Self-Adjust Softmax

  • 将Softmax输出与原输入相乘,动态调整梯度分布。
  • 在27亿参数模型上提升注意力稳定性和训练效率。
  • 无需修改架构,可直接替换现有Transformer注意力模块。

Softmax函数在Transformer注意力中至关重要,通过行归一化使注意力分数和为1,表现优于其他函数。然而当某些注意力分数趋于极端值(如接近0或1)时,传统Softmax会出现梯度消失问题。本文提出Self-Adjust Softmax(SA-Softmax),将softmax(x)改为x·softmax(x),并引入归一化变体:(x - min(x_min,0)) / (max(0,x_max)-min(x_min,0)) · softmax(x)。理论证明,SA-Softmax相比原始Softmax具有更优的梯度特性。此外,该注意力机制可无缝集成至现有Transformer模型,仅需小幅调整。实验在多达27亿参数的模型上进行,覆盖多种数据集、语言任务和位置编码方法,验证了其有效性。

原文摘要 · Abstract (English)

The softmax function is crucial in Transformer attention, which normalizes each row of the attention scores with summation to one, achieving superior performances over other alternative functions. However, the softmax function can face a gradient vanishing issue when some elements of the attention scores approach extreme values, such as probabilities close to one or zero. In this paper, we propose Self-Adjust Softmax (SA-Softmax) to address this issue by modifying $softmax(x)$ to $x \cdot softmax(x)$ and its normalized variant $\frac{(x - min(x_{\min},0))}{max(0,x_{max})-min(x_{min},0)} \cdot softmax(x)$. We theoretically show that SA-Softmax provides enhanced gradient properties compared to the vanilla softmax function. Moreover, SA-Softmax Attention can be seamlessly integrated into existing Transformer models to their attention mechanisms with minor adjustments. We conducted experiments to evaluate the empirical performance of Transformer models using SA-Softmax compared to the vanilla softmax function. These experiments, involving models with up to 2.7 billion parameters, are conducted across diverse datasets, language tasks, and positional encoding methods.

注意力机制Transformer梯度优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。