残差学习的线性注意力增强方法

残差学习的线性注意力增强方法

ENHANCING LINEAR ATTENTION WITH RESIDUALLEARNING

https://arxiv.org/pdf/2509.25223

摘要
线性注意力以线性时间替代自注意力，却常难以捕获长程模式。本文以“预测–校正”视角重新审视线性注意力，证明主流变体均可写成“历史预测 + 单 token 校正”之和，从而造成表达能力瓶颈。为此，我们提出残差线性注意力（RLA），通过显式残差拟合机制增强线性注意力：维护一条辅助递推状态，随时间累积残差误差并修正基础预测。我们进一步实例化出基于 delta 规则的版本——残差 Delta 网络（RDN），引入自适应门控与残差裁剪以增强校正控制与稳定性。实现上，我们复用高度优化的线性注意力核，保持线性时间与内存。在语言建模与强回忆任务中，RLA 与 RDN 均稳定超越对应基线及现代线性注意力方法，缩小与标准 Transformer 的差距，同时维持线性扩展。

1 引言
Transformer 架构已成为大语言模型标配，但其自注意力的二次复杂度仍是处理长序列的关键瓶颈。线性注意力通过将注意力改写为递推过程，实现线性时间训练与推理，成为高效替代方案。RetNet、Mamba 等已展现竞争力；GLA、DeltaNet 等通过数据依赖门控与状态更新规则进一步提升单状态矩阵的信息流。现代线性注意力可被统一为“键到值的直接映射”，等价于测试时训练。例如，delta 更新规则可由一步在线梯度下降导出，这为改进提供多条路径：探索不同在线损失、设计更强映射函数、修改梯度更新机制等。TTT-MLP、Titans 等用 MLP 作深层记忆模块，却牺牲线性递推，增加并行训练难度。

本文提出对注意力输出的新解读：主流线性注意力的输出可分解为“历史状态生成的基础项”与“仅由当前 token 得到的校正项”（见 2.3 节）。依赖单 token 进行系统性校正造成瓶颈，损害模型表达力。为此，我们提出残差线性注意力（RLA），通过辅助状态矩阵显式建模并修正基础线性注意力的系统预测误差，最终输出为“基础预测 + 学习到的误差校正”。该框架可推广至任意线性注意力方法，为构建更强序列模型提供高效策略。

基于现有线性注意力，我们实现两种残差拟合变体：RLA 与 RDN，并在语言建模与强回忆任务上评估。实验表明，二者均优于对应基线及现代线性注意力方法；消融分析验证了框架内各关键设计的重要性。

2 预备知识
2.1 线性注意力作为递推模型

softmax 注意力机制的计算复杂度与序列长度呈二次关系，在处理长序列时构成显著瓶颈。线性注意力（Katharopoulos 等，2020）通过移除 softmax 函数，允许对计算顺序进行重排，从而解决这一问题。

这种循环形式在推理时每步保持恒定的时间和内存复杂度，并通过分块并行算法促进了高效训练 (Yang et al., 2023)。此外，门控机制的应用催生了更多变体，例如 RetNet (Sun et al., 2023)、Lightning Attention (Qin et al., 2024a) 和 Mamba-2 (Dao & Gu, 2024)。

2.2 在线学习视角

该公式使得像Delta Net（Yang et al., 2024b; Schlag et al., 2021）这样的模型能够实现细粒度的记忆控制。Gated Delta Net（Yang et al., 2024a）则通过在学习过程中引入权重衰减，进一步增强了这种方法。

2.3 分解为预测与校正

我们通过预测-校正的视角来解释线性注意力。标准的线性注意力输出 可以被视为来自过去状态的基础预测与基于当前令牌的校正之和：

基于预测-校正的视角，我们引入了一种残差拟合框架来增强线性注意力。该框架通过学习一个更具表达力的校正项，来显式拟合超出当前令牌的上下文信息。

3 方法

本节介绍我们提出的方法，该方法通过残差拟合过程来增强线性注意力。我们首先描述支撑本方法的基础性残差学习框架。接着，引入自适应校正因子以增强建模能力，以及裁剪方法以稳定残差拟合过程。最后，我们给出该方法的两种最终变体。

3.1 显式残差拟合

利用第2节中线性注意力的在线学习视角，我们对辅助状态应用了类似的更新规则。这产生了以下循环过程：

3.2 自适应门控与校正因子

该公式使用衰减因子和校正因子，分别对从基础状态和辅助状态的检索进行动态门控。

3.3 归一化与残差裁剪

这确保了即使基础模型产生暂时性的大预测误差时，误差校正状态 r t
也能保持稳定的学习轨迹。该裁剪方法的详细推导见附录B。

3.4 最终公式

残差拟合原理是一种通用技术，可与多种线性注意力骨干网络相结合。通过将我们的残差机制应用于标准的加法更新规则和 delta 更新规则，我们推导出了两种强大的变体。这构成了我们的最终模型：

4 实验 4.1 设置

实现： 为了最大化效率，我们在 flash-linear-attention 库（Yang & Zhang, 2024）的基础上，使用 Triton（Tillet et al., 2019）实现了自定义的注意力核。我们利用了一个事实，即我们的状态更新规则与线性注意力的规则相同，只需对其内核进行微小修改：我们增强了它，使其同时返回注意力结果和中间残差。这一设计使得同一个高度优化的内核可以在所有残差拟合阶段重复使用，从而确保了高吞吐量。

模型设置：我们通过与多个近期线性注意力架构进行比较来评估我们的模型，包括 Retentive Network (RetNet) (Sun et al., 2023)、Mamba2 (Dao & Gu, 2024) 和 Gated Delta Net (GDN) (Yang et al., 2024a)。此外，我们通过评估标量门控线性注意力（sGLA），一个配备了查询-键归一化和标量门（α 和 β）的线性注意力变体，为 RLA 建立了一个基线。在我们的主要实验中，我们将裁剪阈值设为 c = 1。所有模型都包含大约 15 亿参数，并在相同条件下使用 1000 亿令牌进行训练，以确保公平比较。训练配置的更多细节见附录C。

4.2 主要结果

核效率：我们针对线性注意力基线和 FlashAttention（Dao et al., 2022; Dao, 2023）对我们的核的运行时间进行了基准测试，如图2所示。尽管残差拟合过程增加了计算开销，但我们方法的运行时间随序列长度呈线性增长。这使得其在更长的序列上比呈二次方增长的 FlashAttention 快得多。在吞吐量方面，我们的方法与其他线性注意力机制类似，保持近乎恒定的高吞吐量。相反，受计算限制的 FlashAttention 的吞吐量则随着序列长度增加而迅速下降。

语言建模与常识推理：我们在 WikiText (Merity et al., 2016) 困惑度和一系列评估推理与常识理解的基准测试上评估了 RLA 和 RDN。推理任务包括 ARC-Easy、ARC-Challenge (Clark et al., 2018)、PIQA (Bisk et al., 2020) 和 MMLU (Hendrycks et al., 2020)，而常识理解则通过 HellaSwag (Zellers et al., 2019)、Winogrande (Sakaguchi et al., 2021)、SocialIQA (Sap et al., 2019) 和 LAMBADA (Paperno et al., 2016) 进行评估。我们的主要结果总结于表2，显示我们提出的残差学习变体 RLA 和 RDN 在各自的基线（sGLA 和 GDN）上，困惑度取得了持续改进。此外，我们的模型在多个基准测试中均优于其他领先的线性注意力方法，并能与标准 Transformer 的性能相竞争。

召回密集型任务：为了评估记忆容量，我们在 Arora 等人 (2024) 的召回密集型任务上对我们的模型进行了基准测试。此外，我们还使用"大海捞针"任务 (NIAH) (gkamradt, 2023) 直接评估模型的检索能力，该任务需要从长文档中不同深度插入的键值对中进行检索。如表3所示，这些基准对线性注意力模型具有挑战性，因为其有限的状态空间构成了信息瓶颈。结果表明，我们提出的 RLA 和 RDN 始终优于其对应的基线，尤其在 DROP 和 FDA 基准测试上提升显著。此外，它们在 NIAH 任务上大幅优于其他模型，突显了其增强的信息召回能力。

4.3 消融研究

在本节中，我们进行了一系列消融研究以验证关键组件的贡献。我们首先量化了学习式残差拟合方法相对于预定义校正的优势。接着，我们研究了使用专用校正因子的重要性，然后分析了使用门控机制来结合基础预测与校正的必要性。最后，我们考察了归一化和残差裁剪的效果。

残差拟合
为了验证累积过去误差的重要性，我们测试了一个使用更简单的预定义校正项的变体。在此消融中，我们将我们持久的辅助状态替换为仅从当前残差推导出的无状态校正项：

如表4所示，缺少显式残差拟合的变体性能不如我们的完整方法。尽管这个消融变体在某些基准测试上保持了有竞争力的性能，但其在训练集和评估集上的困惑度均大幅增加。这种性能下降扩展到专业领域，其数学和代码能力严重退化，这由 GSM8k (Cobbe et al., 2021) 和 HumanEval (Chen et al., 2021) 的困惑度衡量指标所证实。这表明辅助状态在累积过去残差以有效精炼模型输出方面起着关键作用。

专用校正因子：通过比较我们的完整模型与校正因子 γ 受更新因子 β 约束的变体，我们分析了使用专用校正因子 γ 的优势。在图3a中，具有独立 γ 的模型始终获得更低的评估损失，其中 RDN 变体显示出更大的改进。这一趋势延续到了下游性能，如图3b的结果所示，该结果也表明专用校正因子能在多个基准测试中带来性能提升。值得注意的是，我们无需额外 γ 的基础架构，相较于基线线性注意力方法，仍标志着显著的改进。

归一化与残差裁剪：最后，我们研究了归一化和残差裁剪的重要性。我们在 RLA 上进行了消融研究，移除了归一化和裁剪。如图4所示，这两个组件对于稳定训练都至关重要；移除它们会导致激活值无界和性能下降。相比之下，RDN 模型对残差裁剪基本不敏感。这种稳健性归因于其 delta 规则更新固有的稳定性，该更新在没有残差裁剪的情况下也能保持一致的损失曲线（图4b）。

5 相关工作

序列建模历史上由循环神经网络（RNNs）主导（Lipton et al., 2015），包括长短期记忆网络（LSTM）（Hochreiter & Schmidhuber, 1997）和门控循环单元（GRU）（Cho et al., 2014）等变体。虽然有效，但其固有的顺序性阻碍了训练的并行化。Transformer 架构（Vaswani et al., 2017）克服了这一限制，成为序列建模的事实标准。然而，其自注意力机制的计算复杂度与序列长度成二次方，对长上下文应用构成了显著瓶颈。

为了应对这些挑战，近期研究重新审视了线性 RNNs，将其作为高效 Transformer 替代方案的基础。通过将序列处理公式化为线性递归，这些模型实现了可并行训练和线性时间推理。该领域的早期探索，如 S4 (Gu et al., 2021)、LRU (Orvieto et al., 2023) 和 RetNet (Sun et al., 2023)，使用了结构化的状态转移矩阵。随后通过引入数据相关的动态实现了性能飞跃。像 Mamba (Gu & Dao, 2023; Dao & Gu, 2024)、HGRN (Qin et al., 2023; 2024b) 和门控线性注意力 (Yang et al., 2023) 等模型利用依赖于输入的动态门控来控制状态转移，从而增强了表达能力。

更先进的方法引入了 delta 学习规则，将状态更新从简单的门控衰减重新定义为细粒度的记忆校正。这种方法以 DeltaNet (Yang et al., 2024b; Schlag et al., 2021) 和 Gated DeltaNet (Yang et al., 2024a) 为代表，实现了更精确的动态记忆修改。该机制可以从在线学习的角度解释，其中状态更新被视为一个优化过程，正如 TTT (Sun et al., 2024) 中所探索的那样。这一观点启发了旨在发现和改进序列模型中内在学习算法的进一步工作（von Oswald et al., 2023; 2025）。

并行的研究集中在增强状态转移的表达能力上。例如，RWKV-7 (Peng et al., 2025) 采用了对角线加低秩结构，而 DeltaProduct (Siems et al., 2025) 通过为每个令牌执行多个更新步骤来泛化 DeltaNet。为了进一步提升容量，近期的架构如 Titans (Behrouz et al., 2024) 和 Miras (Behrouz et al., 2025) 引入了非线性深度记忆，用 MLP 对状态进行参数化。

6 结论

本文介绍了残差线性注意力，这是一个通过显式残差拟合过程来增强线性注意力模型的框架。我们的方法利用一个辅助状态来校正基础模型的预测误差，从而构建更稳健、更准确的上下文表示。该框架具有高度适应性，可应用于多种线性注意力方法。我们的实验证明了这种通用性，表明我们的方法在各项基准上均持续优于其相应的基线。尽管这一改进是以额外计算开销为代价的，但平衡这一权衡为未来研究提供了一个有前景的方向。

原文链接：https://arxiv.org/pdf/2509.25223