Scalable-Softmax Is Superior for Attention

Softmax函数输出的向量中的最大元素在输入向量尺寸增大时趋向于零。基于Transformer的语言模型依赖Softmax来计算注意力分数，这导致随着上下文大小的增长，注意力分布变得平坦。这降低了模型有效优先处理关键信息的能力，并可能限制其长度泛化能力。为了解决这个问题，我们提出了可扩展Softmax（Scalable-Softmax，简称SSMax），它在输入向量尺寸变化的情况下替代Softmax。SSMax可以无缝集成到现有的基于Transformer的架构中。语言建模的实验结果表明，使用SSMax的模型不仅在预训练期间实现更快的损失减少，而且在长上下文和关键信息检索方面显著提高了性能。此外，对注意力分数的分析显示，SSMax使模型即使在长上下文中也能集中注意力于关键信息。另外，尽管从预训练开始就使用SSMax的模型在长度泛化方面表现更好，但那些已经开始预训练的模型也可以通过在预训练过程中或之后将注意力层中的Softmax替换为SSMax，获得部分这种能力。

该论文旨在解决Softmax函数在处理大输入向量时，最大元素接近零的问题，特别是在Transformer模型中用于计算注意力分数时，随着上下文长度增加，注意力分布趋于平坦化，导致模型难以有效优先处理关键信息，并可能限制其对长文本的泛化能力。这是一个需要解决的重要问题，因为当前的Transformer模型在处理长文本时表现不佳。

论文提出了一种名为Scalable-Softmax (SSMax)的新方法来替代传统的Softmax函数。SSMax特别适用于输入向量大小变化的情况，并且可以无缝集成到现有的Transformer架构中。相比传统Softmax，SSMax能够更好地保持注意力分布的集中性，即使在长上下文中也能使模型聚焦于关键信息。这为改进Transformer模型处理长文本的能力提供了新的思路。

实验结果表明，使用SSMax的模型不仅在预训练期间损失减少得更快，而且在长上下文和关键信息检索任务中的性能也显著提升。此外，研究还发现，即使是已经开始预训练的模型，在注意力层中替换为SSMax后，仍然可以在预训练过程中或之后获得更好的长度泛化能力。论文中使用的实验设计合理，涵盖了语言建模等多个方面，并且展示了详细的注意力分数分析。虽然论文未提及是否开源代码，但这些实验结果为进一步研究提供了坚实的基础。

最近在这个领域，有几项相关研究值得关注。例如，《Reformer: The Efficient Transformer》提出了一种更高效的Transformer变体；《Longformer: The Long-Document Transformer》专注于处理超长文档；以及《Performer: Generalized Attention with RFF Kernels for Long Sequences》引入了随机特征核方法来优化注意力机制。这些研究都试图从不同角度解决Transformer在处理长序列时遇到的问题。