arXiv:2501.19399v1 Announce Type: cross 摘要：Softmax函数输出向量的最大元素在输入向量大小增加时趋近于零。基于Transformers的语言模型依赖Softmax来计算注意力分数，导致随着上下文大小的增加，注意力分布趋于扁平化。这降低了模型有效优先处理关键信息的能力，并且可能限制其长度泛化能力。为解决这个问题，我们提出了Scalable-Softmax（SSMax），它可以在输入向量大小变化的情况下替代Softmax。SSMax可以无缝集成到现有的基于Transformers的架构中。在语言建模中的实验结果显示，使用SSMax的模型不仅在预训练过程中更快地降低损失，还在长上下文中以及关键信息检索方面显著提高了性能。此外，对注意力分数的分析表明，SSMax使模型即使在长上下文中也能够集中注意力于关键信息。另外，虽然从预训练开始就使用SSMax的模型在长度泛化方面表现更好，但已经开始了预训练的模型也可以通过在注意力层中用SSMax替换Softmax，要么在预训练期间，要么在预训练之后，获得部分这种能力。