线性注意力架构：机制、权衡与跨层路由

系统比较四种线性注意力变体并提出跨层值路由

前置知识

Softmax Attention

标准的自注意力机制，通过计算查询与键的指数归一化内积来生成注意力权重，然后对值进行加权平均。因果 softmax attention 的输出公式为 $y(i) = \sum_{j \leq i} \frac{\exp(q(i)^\top k(j)/\sqrt{d_k})}{\sum_{\ell \leq i} \exp(q(i)^\top k(\ell)/\sqrt{d_k})} v(j)$，其中 $q(i), k(i) \in \mathbb{R}^{d_k}$ 是查询和键，$v(i) \in \mathbb{R}^{d_v}$ 是值。这种机制让每个 token 能够选择性地从完整上下文中检索信息，但需要计算 $O(T^2)$ 的注意力矩阵，计算和存储成本随序列长度 $T$ 二次增长。

本文比较线性注意力与标准 softmax attention 的性能权衡，理解 softmax attention 的计算局限性是理解线性注意力优势的基础。

Linear Attention

通过特征映射 $\phi: \mathbb{R}^{d_k} \to \mathbb{R}^{d_\phi}$ 将查询和键映射到特征空间，使得未归一化的相似度可以表示为特征映射的内积。通过交换求和顺序，可以维护一个累积记忆矩阵 $W(i) = \sum_{j \leq i} v(j) \otimes \phi(k(j))$，每个 token 的输出变为 $y(i) = W(i) \phi(q(i))$。这种递归更新使得每 token 的计算成本恒定，总训练成本线性于序列长度。线性注意力的关键优势在于将注意力计算从 $O(T^2)$ 降低到 $O(T)$。

本文研究的四种线性注意力架构都是基于线性注意力的递归记忆表示，理解其核心原理是后续内容的前提。

Delta Rule

DeltaNet 引入的错误校正更新规则。与线性注意力简单地将值 $v(i)$ 写入记忆不同，DeltaNet 首先计算当前记忆对当前键的预测 $\bar{v}(i) = W(i-1)\kappa(i)$，然后只写入残差 $r(i) = v(i) - \bar{v}(i)$。记忆更新公式为 $W(i) = W(i-1) + \beta(i) r(i) \otimes \kappa(i)$，其中 $\beta(i)$ 是学习到的写入门控。这种 delta 规则将记忆从纯累加器转换为错误校正写入器，通过只写入记忆未能预测的部分，减少了重叠键值关联之间的干扰。

Delta rule 是本文所有比较的线性注意力变体的核心机制，Gated DeltaNet、Kimi Delta Attention 和 Gated DeltaNet-2 都在此基础上增加了更细粒度的门控机制。

Recurrent Associative Memory

一种递归更新的记忆矩阵 $W(i) \in \mathbb{R}^{d_v \times d_\phi}$，存储从键到值的关联。在每一步，记忆被读取以预测当前键对应的值，然后根据 delta rule 进行更新。记忆矩阵将映射键映射到值，所以键处的读取返回值空间向量。这个记忆可以解释为所有先前键值对的压缩表示，或者作为快速权重编程器。DeltaNet 家族架构共享这个递归记忆骨架，但在选择性、遗忘和编辑控制上有所不同。

本文将所有注意力机制统一在递归记忆符号下进行比较，理解这个框架有助于看到不同架构如何平衡选择性、记忆衰减、擦除和写入控制。

研究动机

Transformer 语言模型依赖自注意力机制让每个 token 从上下文中的其他 token 检索信息，这种 token 间信息交换是模型表达能力的核心。然而，同样的机制使 softmax attention 具有表达力，也使其昂贵：计算注意力矩阵需要对输入序列中每对 token 进行显式成对比较，导致成本随序列长度二次增长 $O(T^2)$。在长序列上，训练和推理的内存限制都变得显著。例如，在 $32k$ token 序列长度下，softmax attention 的迭代时间增长因子约为 $2.9\times$。早期线性注意力变体为这种效率牺牲了相当大的准确性，但最近的工作显著缩小了差距。然而，现有线性注意力架构之间的系统性比较仍然缺乏，缺乏统一框架来理解它们在表达力、记忆衰减、擦除和写入控制、训练吞吐量和实现复杂度方面的差异。同时，深层语言模型中信息稀释问题未得到充分解决，随着表示通过许多层传播，在较低深度提取的有用信号可能变得越来越难以恢复。

本文的目标是本文的目标是提供一个系统性的比较研究，将 softmax attention 与四种最新的循环线性注意力架构放在统一的递归记忆表示下，明确它们在表达力、记忆衰减、擦除和写入控制、训练吞吐量和实现复杂度方面的差异。作者使用共享的 $350M$ 参数、$15B$ token 训练设置来映射架构间的经验权衡，包括验证损失、吞吐量、优化器和学习率敏感性、混合与纯堆栈结构、序列长度计时、更大 DeltaNet 运行和下游行为。同时，作者探索在保持线性时间结构的前提下，通过轻量级跨层路由机制在 DeltaNet 风格记忆之间共享信息，以解决深层语言模型中信息稀释的问题。

与已有工作不同的是，本文的独特切入点在于提供了一个统一的递归记忆表示框架，将不同的线性注意力机制放在相同符号下进行比较，使得设计选择变得清晰。与现有的单一架构论文不同，本文不试图声明单一最佳架构，而是展示多目标权衡。同时，本文探索的跨层路由方向与现有的深度注意力机制不同，后者部分抵消了线性注意力的效率优势。本文询问是否存在一个轻量级的方式在深度间共享信息，同时尊重 DeltaNet 风格记忆的线性时间结构。这种做法是新的，因为不是引入新的深度注意力算子，而是重用递归记忆内部产生的信号。

核心方法

本文方法分为两个主要部分：架构比较和跨层路由。在架构比较部分，作者将 softmax attention、DeltaNet、Gated DeltaNet、Kimi Delta Attention 和 Gated DeltaNet-2 表达在统一的递归记忆符号下。所有机制都维护一个记忆矩阵来存储键值信息，输出通过查询记忆获得。不同之处在于如何更新记忆。DeltaNet 使用错误校正的 delta 规则，只写入记忆未能预测的部分。Gated DeltaNet 增加了学习到的标量衰减门控。Kimi Delta Attention 将标量衰减替换为通道衰减门控。Gated DeltaNet-2 进一步解耦主动 delta 规则编辑为独立的通道擦除和写入门控。在跨层路由部分，作者探索了 CLER、CLER-H 和 CLVR 三种机制。

核心创新点有两个。首先，统一框架下的系统性比较使得不同架构的设计选择变得清晰，包括递归记忆 $W$、delta 规则残差 $r$ 以及不同的衰减、擦除和写入机制。这个框架隔离了记忆如何读取、旧内容如何遗忘或擦除、新值信息如何写入的设计选择。其次，提出的跨层值路由机制通过零初始化投影 $P_l \in \mathbb{R}^{d_{model} \times d_v}$ 将层内的写值 $v_{l,t}$ 投影到模型维度并添加到共享残差流中：$\varepsilon_{l,t} = P_l v_{l,t}$，$h_{l,t} \leftarrow h_{l,t} + \varepsilon_{l,t}$。这种机制与直接将写误差注入到下一层的值目标空间的 CLER 方法不同，CLVR 将信号注入到对齐的共享空间中，使得路由的贡献从训练开始时的零逐渐增长。

方法步骤详情

方法步骤分为架构比较和跨层路由两部分。架构比较的步骤：对于每个 token 位置 $i$，首先将输入表示 $x(i)$ 映射到查询、键和值向量；对于线性注意力变体，通过特征映射得到映射键 $\kappa(i) = \phi(k(i))$；计算记忆预测 $\bar{v}(i) = W(i-1)\kappa(i)$；计算 delta 规则残差 $r(i) = v(i) - \bar{v}(i)$；根据不同架构更新记忆；通过查询记忆得到输出 $y(i) = W(i)\phi(q(i))$。跨层路由的步骤：对于每个路由层 $l$ 和序列位置 $t$，计算写值 $v_{l,t}$ 和写误差 $r_{l,t} = v_{l,t} - \bar{v}_{l,t}$；选择要路由的信号；通过零初始化投影将信号投影到模型维度 $\varepsilon_{l,t} = P_l s_{l,t}$；将投影添加到共享残差流 $h_{l,t} \leftarrow h_{l,t} + \varepsilon_{l,t}$；后续层和输出头通过标准残差路径读取注入的信号。

技术新颖性

本文的技术新颖性体现在三个方面。首先，统一的递归记忆表示框架是新颖的，它将不同的线性注意力机制放在相同符号下进行比较，使得设计选择变得清晰。其次，系统性比较包括最新的线性注意力变体的 Megatron 实现和集成是新的，这些实现公开发布，促进了社区的研究和复现。第三，提出的跨层值路由机制是新的，它路由线性注意力算子内部的写值 $v_{l,t}$ 而不是层输出，通过零初始化投影 $P_l$ 添加到对齐的残差流中。这与现有的深度注意力机制本质不同，因为它们路由层输出并通过学习到的深度 softmax 组合子层输出。据作者所知，将线性注意力或 delta 规则记忆的内部写值路由到共享残差流尚未被研究。

**Figure 1**: Cross-Layer Error Residuals (CLER)

**Figure 2**: Cross-Layer Value Routing (CLVR)

实验结果

实验结果呈现了验证损失、吞吐量和序列长度缩放之间的多目标权衡。在 350M 参数、15B token 的主要扫描中，Kimi Delta Attention 与 Muon 在混合堆栈中达到最佳最终验证损失 2.273，DeltaNet 与 Muon 在混合堆栈中最强的非 KDA 条目达到 2.299。最快的条目是纯 Gated DeltaNet 与 AdamW，归一化相对速度为 100%，但最终损失显著更高为 2.433。序列长度迭代时间缩放显示，从 4k 到 32k token，softmax attention 的增长因子约为 2.9 倍，Gated DeltaNet 混合堆栈为 1.7 倍，纯 Gated DeltaNet 堆栈为 1.1 倍。CLVR 在所有报告行中提供了一致的小幅改进。

**Table 1**: Core configuration information for the main experiments

**Table 2**: 350M-parameter, 15B-token architecture comparison

**Table 3**: Larger DeltaNet-only runs at 1.3B-parameter/40B-token and 3B-parameter/60B-token scale

**Table 4**: Downstream checks on HellaSwag, PIQA, and WinoGrande

**Table 5**: Matched 350M-parameter, 15B-token Muon downstream evaluations for Gated DeltaNet and DeltaNet routing checkpoints

**Table 6**: Matched CLER comparisons at 350M parameters and approximately 1B training tokens

**Table 7**: Matched hidden-stream routing comparisons under Muon

**Table 8**: Supplementary 1B-Token Baselines

**Figure 3**: Validation-loss scaling for the 350M-parameter runs trained to 15B tokens

**Figure 4**: Learning-rate ablation for 350M-parameter hybrid models trained for 2000 steps

**Figure 5**: Learning-rate ablation iso plots for 350M-parameter hybrid models trained for 2000 steps

**Figure 6**: Iteration Time vs. Sequence Length

查看结构化数据

任务	指标	本文	基线	提升
语言建模	最终验证损失	2.273 (Kimi Delta Attention + Muon hybrid)	2.413 (Softmax + AdamW dense)	0.140 绝对改进
训练吞吐量	归一化相对速度	100% (Gated DeltaNet + AdamW pure)	81.7% (Softmax + AdamW dense)	18.3% 相对加速
序列长度缩放	迭代时间增长因子 (4k to 32k)	1.1× (Gated DeltaNet pure)	2.9× (Softmax)	2.6× 更好的缩放
跨层路由效果	验证损失降低	-0.0119 (DeltaNet 350M/1B)	无路由	0.0119 绝对改进
下游任务 HellaSwag	归一化准确率	0.4305 (DeltaNet + Muon hybrid)	0.4158 (Softmax + AdamW dense)	0.0147 提升

局限与改进

本文存在几个局限性。首先，经验比较是审核运行集而不是完整的统计研究。报告的架构和路由行是单次运行，可用记录不支持种子平均排名或标准差。小验证损失差异（低于大约 $10^{-3}$ 到 $10^{-2}$）被解释为建议性的。其次，超参数覆盖在不同变体间不均匀，所以优化器和学习率效应应该一起阅读。第三，下游评估仅限于 HellaSwag、PIQA 和 WinoGrande，可能错过长上下文或内存密集用例中重要的行为。第四，跨层路由效果随规模扩大而减弱，更大的 $1.3B$ 参数运行中增益缩小到 $-0.0019$。第五，报告的速度结果测量训练吞吐量和迭代时间，没有提供经验推理速度基准。

独立分析的弱点

本文的独立分析弱点包括：首先，优化器和学习率敏感性表明单个默认学习率可能会扭曲架构比较，变体可能表现不佳是因为其优化器和学习率对不匹配。其次，混合堆栈虽然改善损失但减少了部分激励线性注意力的速度优势，在 $32k$ token 时混合堆栈迭代时间增加约 $65\%$。第三，跨层路由效果虽然在不同设置中一致，但增益很小且随训练时间延长和模型增大而缩小。第四，缺乏对推理吞吐量、解码内存占用和长上下文质量的系统评估。第五，单次运行设计使得小差异的统计显著性不明确。改进方向包括进行多次种子的统计研究，扩展下游评估到长上下文和内存密集任务，系统测量推理速度和内存占用。

未来方向

未来研究方向包括作者提出的和基于成果可延伸的几个方向。首先，最清楚的问题是增益是否持续或增长在作者未评估路由的主机上：Kimi Delta Attention 和 Gated DeltaNet-2。Gated DeltaNet-2 特别解耦了写值与独立的擦除和写入残差，所以不清楚哪个内部信号应该被路由。其次，观察到的随着更长训练和更大 Gated DeltaNet 行而收益缩小的模式是否持续直到 CLVR 变为中性，这可以通过在更大 token 预算下的额外重复运行解决。第三，未来的路由运行应该记录直接测量影响所需的量：投影后的路由信号、注入前后接收器的残差流以及学习到的投影范数。第四，CLVR 应该在跨层恢复有更清晰角色的任务上测试：合成键值检索、passkey 风格任务、长上下文 QA、多干扰上下文学习等。

复现评估

本文的复现评估显示开源情况和实现细节良好。实现和实验记录在 tommasocerruti/linear-attention-architectures 公开可用。发布的仓库是专注的 Megatron-LM 分支，包含线性注意力实现、最终训练启动器、FineWeb-Edu 和 LLaMA2 数据准备路径。硬件配置使用 CSCS Alps 系统上的一个 GH200 节点和四个 GPU，bf16 精度。训练数据是 FineWeb-Edu，使用 LLaMA2 分词器，序列长度 $4096$，全局批量大小 $128$。主要 $350M$ 比较运行 $15B$ token，受控路由消融运行 $1B$ token。代码、配置和训练脚本的可用性使得复现主要结果是可行的。

线性注意力架构：机制、权衡与跨层路由 Linear Attention Architectures: Mechanisms, Trade-offs, and Cross-Layer Routing