Softplus Attention with Re-weighting Boosts Length Extrapolation in Large Language Models
作者: Bo Gao, Michael W. Spratling
分类: cs.CL, cs.AI, cs.LG
发布日期: 2025-01-23 (更新: 2025-08-11)
备注: 10 pages and 3 figures
💡 一句话要点
提出Softplus注意力机制与重加权策略，显著提升大语言模型长度外推能力
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: Softplus注意力 长度外推 大语言模型 自注意力机制 重加权策略
📋 核心要点
- 传统Softmax注意力机制在长序列推理时面临数值不稳定和性能下降的挑战。
- 论文提出Softplus注意力机制，结合动态尺度因子和重加权策略，增强模型对关键token的关注。
- 实验表明，该方法显著提升了模型在长文本上的外推能力，并在长上下文检索任务中表现出色。
📝 摘要（中文）
大型语言模型近年来取得了显著的成功，这主要归功于自注意力机制的实现。然而，传统的Softmax注意力机制存在数值不稳定性和随着推理token长度增加而性能下降的问题。本文通过将注意力视为一个两阶段过程，提出了一种新的注意力设计原则来解决这些问题。我们首先将Softmax操作分解为一个非线性正性变换和一个$l_1$-归一化步骤，并将后者确定为维持模型性能的关键。在第一阶段，我们用数值上更稳定的Softplus激活函数代替标准的指数函数，并引入一个基于不变熵的动态尺度因子，从而创建了一种优于传统Softmax注意力机制的新型注意力机制。在第二阶段，我们引入了一种重加权机制，以锐化注意力分布，放大显著权重，同时减小较弱的权重。这使得模型能够更有效地集中于相关的token，并从根本上改善长度外推。当结合使用时，这种两阶段方法可确保数值稳定性并显着改善长度外推，在训练长度的16倍时保持几乎恒定的验证损失，同时在具有挑战性的长上下文检索任务和标准下游基准测试中获得卓越的结果。
🔬 方法详解
问题定义：现有的大型语言模型依赖于自注意力机制，但标准的Softmax注意力在处理长序列时会遇到数值不稳定的问题，导致性能下降。具体来说，当序列长度增加时，Softmax函数计算出的注意力权重可能过于集中或分散，难以有效捕捉长距离依赖关系，从而限制了模型的外推能力。
核心思路：论文的核心思路是将注意力机制分解为两个阶段：非线性正性变换和归一化。通过替换Softmax中的指数函数为更稳定的Softplus函数，并引入动态尺度因子，来缓解数值不稳定性。此外，通过重加权机制锐化注意力分布，使模型更加关注重要的token，从而提升长度外推能力。
技术框架：该方法包含两个主要阶段。第一阶段是Softplus注意力计算，使用Softplus函数替代指数函数，并引入基于不变熵的动态尺度因子。第二阶段是注意力重加权，通过一个可学习的函数对注意力权重进行调整，放大重要权重，抑制不重要权重。整体流程是：输入Query、Key、Value，首先计算Softplus注意力权重，然后进行重加权，最后加权求和得到最终的上下文向量。
关键创新：最重要的创新点在于将Softmax注意力分解为两个阶段，并分别进行改进。Softplus激活函数的使用提高了数值稳定性，动态尺度因子能够自适应地调整注意力分布，重加权机制则增强了模型对关键token的关注。与传统Softmax注意力相比，该方法在长序列上表现出更好的鲁棒性和外推能力。
关键设计：动态尺度因子基于不变熵计算，旨在保持注意力分布的信息量。重加权机制使用一个可学习的函数，例如一个小型神经网络，对注意力权重进行调整。损失函数方面，论文可能使用了标准的交叉熵损失或类似的损失函数来训练模型。具体的网络结构细节（如Softplus注意力中的动态尺度因子的计算方式，重加权函数的具体形式）在论文中可能有所描述，但摘要中未明确提及。
🖼️ 关键图片
📊 实验亮点
该方法在长度外推任务中表现出色，在训练长度的16倍时仍能保持几乎恒定的验证损失。此外，在具有挑战性的长上下文检索任务和标准下游基准测试中也取得了优于传统Softmax注意力的结果。这些实验结果表明，该方法能够有效提升模型在长序列上的性能和泛化能力。
🎯 应用场景
该研究成果可广泛应用于需要处理长文本的自然语言处理任务中，例如长文档摘要、机器翻译、问答系统和信息检索等。通过提升模型对长序列的理解能力，可以改善这些应用在处理复杂和长篇内容时的性能，并为开发更强大的语言模型奠定基础。
📄 摘要（原文）
Large language models have achieved remarkable success in recent years, primarily due to the implementation of self-attention mechanisms. However, traditional Softmax attention suffers from numerical instability and reduced performance as the length of inference tokens increases. This paper addresses these issues by proposing a new design principle for attention, viewing it as a two-stage process. We first decompose the Softmax operation into a non-linear positivity transformation and an $l_1$-normalisation step, identifying the latter as essential for maintaining model performance. In the first stage, we replace the standard exponential function with the more numerically stable Softplus activation and introduce a dynamic scale factor based on invariance entropy, creating a novel attention mechanism that outperforms conventional Softmax attention. In the second stage, we introduce a re-weighting mechanism that sharpens the attention distribution, amplifying significant weights while diminishing weaker ones. This enables the model to concentrate more effectively on relevant tokens and fundamentally improves length extrapolation. When combined, this two-stage approach ensures numerical stability and dramatically improves 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 standard downstream benchmarks.