arXiv:2501.13428cs.CLcs.AI2025-01中稿 · ICML被引 1

用软加法+重加权改进注意力,让大模型推理更长文本时更稳定、更准确。

Softplus Attention with Re-weighting Boosts Length Extrapolation in Large Language Models

  • 将Softmax替换为软加法再归一化,提升数值稳定性。
  • 引入动态缩放因子和权重重加权,显著改善长文本推理表现。
  • 适合需要处理超长文本的模型优化与物理规律建模任务。

近年来大语言模型取得显著进展,主要依赖于自注意力机制。然而,传统Softmax注意力在推理令牌数增加时易出现数值不稳定性且性能下降。本文提出一种新的注意力设计原则,将其视为两阶段过程:第一阶段(归一化)用更稳定的软加法替代Softmax,并进行l₁归一化;同时引入基于不变性熵的动态缩放因子。第二阶段(锐化)通过重加权机制放大重要注意力权重,抑制弱权重,缓解注意力汇聚现象,从根本上提升长度外推能力。该新型两阶段注意力机制确保了数值稳定性,在训练长度16倍的推理长度下仍保持近似恒定的验证损失,并在挑战性长上下文检索任务和下游基准测试中表现优异。符号回归实验表明,该方法使模型能从轨道轨迹序列中恢复牛顿引力定律,证明合适的注意力机制对基础模型构建真实物理世界认知至关重要。代码已开源:https://github.com/iminfine/freeattn。

原文摘要 · Abstract (English)

Large language models have achieved remarkable success in recent years, primarily due to self-attention. However, traditional Softmax attention suffers from numerical instability and reduced performance as the number of inference tokens increases. This work addresses these issues by proposing a new design principle for attention, viewing it as a two-stage process. The first stage (normalisation) refines standard attention by replacing Softmax with the more numerically stable Softplus followed by $l_{1}$-normalisation. Furthermore, we introduce a dynamic scale factor based on invariance entropy. We show that this novel attention mechanism outperforms conventional Softmax attention, and state-of-the-art Softmax-free alternatives. Our second proposal is to introduce a second processing stage (sharpening) which consists of a re-weighting mechanism that amplifies significant attentional weights while diminishing weaker ones. This enables the model to concentrate more effectively on relevant tokens, mitigating the attention sink phenomenon, and fundamentally improving length extrapolation. This novel, two-stage, replacement for self-attention is shown to ensure numerical stability and dramatically improve length extrapolation, maintaining a nearly constant validation loss at 16$\times$ the training length while achieving superior results on challenging long-context retrieval tasks and downstream benchmarks. Furthermore, symbolic regression experiments demonstrate that our method enables models to recover Newton's gravitational law from orbital trajectory sequences, providing evidence that appropriate attention mechanisms are crucial for foundation models to develop genuine physical world models. Our code is available at https://github.com/iminfine/freeattn.

注意力机制长文本生成模型优化

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