arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2501.13428cs.CLcs.AIcs.LG
收藏

Softplus注意力与重加权提升大语言模型的长度外推能力

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

University of California, Berkeley(加州大学伯克利分校)

Bo Gao, Michael W. Spratling, Letizia Gionfrida

AI总结:

提出一种两阶段注意力机制,用Softplus和l1归一化替代Softmax,并引入基于不变熵的动态缩放因子和重加权机制,以提升数值稳定性、缓解注意力下沉现象,并显著改善长度外推性能。

摘要:

近年来,大语言模型取得了显著成功,这主要归功于自注意力机制。然而,传统的Softmax注意力存在数值不稳定性,并且随着推理令牌数量的增加,性能会下降。本文通过提出一种新的注意力设计原则来解决这些问题,将注意力视为一个两阶段过程。第一阶段(归一化)通过用数值更稳定的Softplus后接$l_{1}$归一化替代Softmax来改进标准注意力。此外,我们引入了一个基于不变熵的动态缩放因子。我们证明,这种新颖的注意力机制优于传统的Softmax注意力和最先进的非Softmax替代方案。我们的第二个提议是引入第二阶段处理(锐化),该阶段由一个重加权机制组成,该机制放大重要的注意力权重,同时削弱较弱的权重。这使得模型能够更有效地聚焦于相关令牌,缓解注意力下沉现象,并从根本上改善长度外推。这种新颖的两阶段自注意力替代方案被证明能确保数值稳定性,并显著提升长度外推能力,在训练长度的16倍时保持几乎恒定的验证损失,同时在具有挑战性的长上下文检索任务和下游基准测试中取得优异结果。此外,符号回归实验表明,我们的方法使模型能够从轨道轨迹序列中恢复牛顿万有引力定律,这为适当的注意力机制对于基础模型发展真正的物理世界模型至关重要提供了证据。我们的代码可在 https://github.com/iminfine/freeattn 获取。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑