arXiv:2501.19399cs.CLcs.AI2025-01被引 43

改进注意力机制,让模型在长文本中更聚焦关键信息。

Scalable-Softmax Is Superior for Attention

  • 用可扩展的Softmax替代传统Softmax,适应不同长度输入。
  • 在长文本任务中显著提升信息检索准确率和训练收敛速度。
  • 无需重训练,替换注意力层即可增强模型长程泛化能力。

随着输入向量规模增大,Softmax输出向量的最大元素趋近于零,导致Transformer语言模型在上下文变长时注意力分布趋于平坦,削弱了对关键信息的聚焦能力,影响模型的长度泛化性能。为此,我们提出可扩展Softmax(SSMax),在输入维度变化的场景下替代传统Softmax,可无缝集成至现有Transformer架构中。实验表明,采用SSMax的模型在预训练过程中损失下降更快,并在长上下文任务与关键信息检索中表现显著提升。注意力分析显示,即使在长序列中,SSMax仍能有效聚焦关键内容。此外,即便预训练已开始,通过在注意力层替换为SSMax,模型仍可获得一定的长度泛化能力,无论是在训练中还是训练后进行替换。

原文摘要 · Abstract (English)

The maximum element of the vector output by the Softmax function approaches zero as the input vector size increases. Transformer-based language models rely on Softmax to compute attention scores, causing the attention distribution to flatten as the context size grows. This reduces the model's ability to prioritize key information effectively and potentially limits its length generalization. To address this problem, we propose Scalable-Softmax (SSMax), which replaces Softmax in scenarios where the input vector size varies. SSMax can be seamlessly integrated into existing Transformer-based architectures. Experimental results in language modeling show that models using SSMax not only achieve faster loss reduction during pretraining but also significantly improve performance in long contexts and key information retrieval. Furthermore, an analysis of attention scores reveals that SSMax enables the model to focus attention on key information even in long contexts. Additionally, although models that use SSMax from the beginning of pretraining achieve better length generalization, those that have already started pretraining can still gain some of this ability by replacing Softmax in the attention layers with SSMax, either during or after pretraining.

注意力机制Transformer长文本

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