可扩展 Softmax 对注意力机制至关重要

Softmax 函数输出向量的最大元素随输入向量大小的增加而趋近于零。基于 Transformer 的语言模型依赖 Softmax 计算注意力得分，导致注意力分布随上下文规模增大而趋于平坦。这削弱了模型有效优先处理关键信息的能力，并可能限制其长度泛化。为此，我们提出可扩展 Softmax（SSMax），在输入向量大小可变的场景下取代 Softmax。SSMax 可无缝集成到现有 Transformer 架构中。实验表明，使用 SSMax 的模型在预训练期间不仅实现更快的损失下降，还在长序列上下文和关键信息检索方面显著提升性能。此外，对注意力得分的分析表明，SSMax 使模型即使在长序列上下文中也能聚焦于关键信息。尽管从预训练初期即使用 SSMax 的模型在长度泛化方面取得更好效果，但已开始预训练的模型仍可通过在注意力层中替换 Softmax 为 SSMax（无论在预训练期间或之后），从中获益。