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

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

来源：https://arxiv.org/html/2607.07953

Tommaso Cerruti¹,\* & Tim Rieder¹,\* & George Rowlands¹,\* & Lingfeng Jin¹,\* & Imanol Schlag¹,²

¹苏黎世联邦理工学院，D-INFK  
²苏黎世联邦理工学院AI中心，苏黎世联邦理工学院  
\*同等贡献

###### 摘要

自注意力机制让每个词元能从整个上下文中检索信息，但其在序列长度上的二次方成本限制了长上下文下的训练和推理。本文对softmax注意力与四种近期提出的循环线性注意力架构进行了比较研究：DeltaNet、门控DeltaNet、Kimi Delta注意力以及门控DeltaNet-2。我们使用一种统一的循环记忆符号表示这些机制，明确了它们在表达能力、记忆衰减、擦除和写入控制、训练吞吐量以及实现复杂度上的差异。我们的实验以150亿tokens上训练的3.5亿参数模型为核心，并包括优化器与学习率比较、混合与纯堆栈比较、序列长度运行时间测量、更大规模（13亿和30亿参数）的DeltaNet实验，以及少量下游评估任务。报告中给出的速度结果衡量的是训练吞吐量和迭代时间；我们未提供推理速度的实证基准。在报告的3.5亿参数、150亿tokens的扫描实验中，使用Muon优化器的Kimi Delta注意力取得了最低的最终验证损失；使用AdamW训练的纯门控DeltaNet堆栈具有最高的归一化训练吞吐量；混合堆栈通常以吞吐量为代价改善损失；而在我们评估的匹配架构设置中，Muon相对于AdamW持续降低了最终验证损失。我们引入并评估了针对DeltaNet风格记忆的轻量级跨层路由机制。最自然的DeltaNet启发式公式——将较低层的delta规则写入误差转发到下一层的值目标中——并未优于匹配的基线。将路由引入对齐的隐藏流并转发写入值而非写入误差，在我们报告的匹配运行中带来适度的改善：跨层值路由（CLVR）同时降低了DeltaNet和门控DeltaNet的最终验证损失。

**关键词**：线性注意力 · 循环联想记忆 · DeltaNet · 跨层路由

## 1 引言

Transformer语言模型依赖自注意力机制[1]（https://arxiv.org/html/2607.07953#bib.bib1）让每个词元能从上下文中的其他词元检索信息。这种词元间的信息交换，通常称为序列混合，是其表达能力的关键。然而，使softmax注意力具有表达能力的同一机制也使其代价高昂：计算注意力矩阵需要对输入序列中每对词元进行显式的两两比较，导致成本随序列长度呈二次方增长。随着模型部署在更大的上下文窗口中，这一成本在训练和推理中都成为主导因素，并促使人们寻找具有更优扩展性的序列混合器。

*线性注意力*是对这一挑战最突出的回应之一。通过用特征映射分解替换softmax核，线性注意力可以重新表述为对固定大小记忆矩阵的循环更新，从而实现线性时间训练和每词元恒定时间推理。早期的线性注意力变体为此效率牺牲了相当程度的准确性，但近期工作已显著缩小了这一差距。特别是，DeltaNet[2]（https://arxiv.org/html/2607.07953#bib.bib2）将线性注意力重新解释为一种快速权重编程器，并用纠错*delta规则*取代了朴素加法更新：模型不是每一步都写入完整值，而是仅写入当前值与固定大小记忆已预测值之间的残差，从而减少了来自重叠键值关联的干扰。后续变体在此基础上增加了越来越精细的控制记忆机制。门控DeltaNet[3]（https://arxiv.org/html/2607.07953#bib.bib3）引入了记忆状态上的学习标量衰减；Kimi Delta注意力[4]（https://arxiv.org/html/2607.07953#bib.bib4）通过通道级衰减门进一步细化；门控DeltaNet-2[5]（https://arxiv.org/html/2607.07953#bib.bib5）进一步将主动delta规则编辑解耦为独立的通道级擦除门和写入门。这些架构共同定义了一个规模虽小但快速演化的循环记忆家族，它们共享基本骨架，但在选择性、遗忘机制和控制粒度上有所不同。

与这些进展并行的，还有另一条独立的工作线认为深层语言模型遭受*信息稀释*问题：随着表示在多个层中传播，在较低层提取的有用信号可能逐渐难以恢复。诸如Attention Residuals[6]（https://arxiv.org/html/2607.07953#bib.bib6）和Mixture-of-Depths Attention[7]（https://arxiv.org/html/2607.07953#bib.bib7）的提案通过引入显式跨层路径或深度级注意力来解决这一问题。这些机制是有效的，但将其直接应用于线性循环架构会部分抵消其效率优势。这自然引出一个问题：是否存在一种轻量级的方式在深度之间共享信息，同时尊重DeltaNet风格记忆的线性时间结构？

本文做出四项贡献：
1. 我们用统一的循环记忆符号表示softmax注意力、DeltaNet、门控DeltaNet、Kimi Delta注意力和门控DeltaNet-2。该框架隔离了循环记忆\(W\)、delta规则残差\(r\)以及不同衰减、擦除和写入机制的作用。
2. 我们提供了研究中使用的较新线性注意力变体的可训练Megatron实现和集成，包括Kimi Delta注意力、门控DeltaNet-2以及本文引入的跨层路由变体。该实现可在[![[无标题图片]](https://arxiv.org/html/2607.07953v1/assets/github-mark.png)tommasocerruti/linear-attention-architectures](https://github.com/tommasocerruti/linear-attention-architectures)获取。
3. 我们利用此共享设置来绘制架构间的经验权衡，比较验证损失、吞吐量、优化器和学习率敏感性、混合与纯堆栈结构、序列长度时序、更大规模的DeltaNet实验以及下游行为。
4. 我们为DeltaNet风格记忆引入了轻量级跨层路由。从*跨层误差残差（CLER）*开始——将较低层的delta规则写入误差转发到下一层的值目标中——我们发现这种DeltaNet启发的公式并不优于匹配基线。然后我们将路由引入对齐的隐藏流，并发现层的写入*值*而非其写入误差是有用的信号。由此产生的方法，*跨层值路由（CLVR）*，在我们报告的匹配DeltaNet和门控DeltaNet运行中略微降低了最终验证损失，同时保持了宿主架构的线性时间结构。

我们沿三个轴组织比较：
- • **机制**：我们比较每种循环更新的代数结构。
- • **权衡**：我们讨论每种设计如何平衡选择性、记忆衰减、擦除和写入控制以及记忆粒度。
- • **经验行为**：我们报告在所研究变体中的验证损失、吞吐量、序列长度时序和下游结果。

我们的目标不是声明一个单一的最佳架构，而是使设计空间变得清晰可读。将这些机制并列观察，更容易看出哪些选择负责哪些属性，在哪里扩展性最强，以及哪些跨层路由变体对未来的评估仍是开放问题。

## 2 背景：从Softmax注意力到线性注意力

本节回顾从标准softmax注意力到线性注意力的路径，并介绍作为本文研究的所有DeltaNet风格架构基础的循环记忆视角。我们专注于代数结构而非实现细节，并使用将在本文其余部分使用的同一单头符号。

#### Softmax注意力。

给定查询和键\(q^{(i)}, k^{(i)} \in \mathbb{R}^{d_k}\)以及值\(v^{(i)} \in \mathbb{R}^{d_v}\)，其中\(i=1,\dots,T\)，因果softmax注意力计算位置\(i\)的输出为：
\[
y^{(i)} = \sum_{j \leq i} \frac{\exp\left(q^{(i)\top} k^{(j)} / \sqrt{d_k}\right)}{\sum_{\ell \leq i} \exp\left(q^{(i)\top} k^{(\ell)} / \sqrt{d_k}\right)} v^{(j)}.
\tag{1}
\]
当前查询会与每个先前键进行显式比较，并通过归一化加权平均组合值。这种公式表达能力很强：每个词元原则上可以从任意先前位置检索信息，并且softmax归一化赋予了机制一种内建的软选择形式。然而，相同的显式比较也是其二次方成本的来源。在长度为\(T\)的序列上训练需要计算\(\mathcal{O}(T^2)\)的查询-键交互矩阵，且存储完整的注意力矩阵也需要\(\mathcal{O}(T^2)\)的内存。对于长上下文，时间和内存成本都成为限制因素。

#### 核视角。

线性注意力源于一个简单观察：如果未归一化的注意力权重可以写成查询和键的特征映射之间的内积，那么对先前位置的求和可以重新排列以避免显式的成对比较。具体来说，假设存在一个特征映射\(\phi: \mathbb{R}^{d_k} \to \mathbb{R}^{d_\phi}\)，使得（未归一化的）相似度可以近似为：
\[
\exp\left(q^{(i)\top} k^{(j)} / \sqrt{d_k}\right) \approx \phi\left(q^{(i)}\right)^\top \phi\left(k^{(j)}\right).
\tag{2}
\]
将此代入注意力输出未归一化的分子并交换求和顺序得到：
\[
\sum_{j \leq i} \phi\left(q^{(i)}\right)^\top \phi\left(k^{(j)}\right) v^{(j)} = \phi\left(q^{(i)}\right)^\top \underbrace{\sum_{j \leq i} v^{(j)} \otimes \phi\left(k^{(j)}\right)}_{\text{累积记忆}}.
\tag{3}
\]
括号内的求和不再依赖于查询，并且可以在\(i\)前进时增量式维护。然后查询与单个矩阵交互，而不是与所有先前词元单独交互。

#### 循环记忆。

将运行总和定义为记忆矩阵：
\[
W^{(i)} = \sum_{j \leq i} v^{(j)} \otimes \phi\left(k^{(j)}\right) \in \mathbb{R}^{d_v \times d_\phi},
\tag{4}
\]
线性注意力的（未归一化）输出可以写成一个循环：
\[
W^{(i)} = W^{(i-1)} + v^{(i)} \otimes \phi\left(k^{(i)}\right), \qquad y^{(i)} = W^{(i)} \phi\left(q^{(i)}\right).
\tag{5}
\]
每个词元向记忆贡献一次加法写入，每个输出是单个矩阵-向量乘积。每词元成本在序列长度上是常数，总训练成本在\(T\)上是线性的。记忆矩阵可以用两种等价方式解释：作为所有先前（键，值）对的压缩表示，或者遵循快速权重编程器观点[2]（https://arxiv.org/html/2607.07953#bib.bib2），作为一组由周围网络写入和读取的*快速权重*。

#### 线性注意力放弃的东西。

效率提升是有代价的。Softmax注意力执行逐查询归一化，有效重新加权所有先前词元的贡献，其指数核使注意力分布变得尖锐。线性注意力以其朴素的加法形式，两者都不具备。记忆\(W^{(i)}\)累积先前的写入，没有学习的衰减或擦除机制，因此旧的关联保持活跃，除非被间接覆盖。也没有内建机制来锐化检索分布或遗忘过时信息。随着更多的（键，值）对被写入固定大小的记忆，存储关联之间的干扰增加：旨在检索某个值的查询也可能从与其在\(\phi\)下恰好相关的无关键中拾取虚假贡献。这种干扰是较新线性注意力变体旨在解决的核心失败模式。

#### 从加法写入到delta规则写入。

DeltaNet及其后继者可以理解为对干扰问题的原则性回答。这些架构不是总是将原始值\(v^{(i)}\)添加到记忆，而是首先询问当前记忆对当前键的预测：
\[
\bar{v}^{(i)} = W^{(i-1)} \phi\left(k^{(i)}\right),
\tag{6}
\]
然后仅写入残差\(r^{(i)} = v^{(i)} - \bar{v}^{(i)}\)。这将记忆更新从纯累加器转换为纠错写入器，并赋予架构一种关于*记忆尚不知道什么*的概念。后续变体通过遗忘和编辑机制增强此delta规则更新，范围从单一标量衰减门到通道级衰减，最后到独立的通道级擦除和写入门。这些添加使模型能够显式控制旧信息被衰减的激进程度、哪些关联被擦除以及哪些值通道被提交到记忆。

下一节以统一符号形式化这些架构，并使其差异精确化。

## 3 架构

我们使用单头符号描述所有注意力机制，并为了清晰省略输出投影、归一化层和前馈块。在词元位置\(i\)，输入表示\(x^{(i)}\)被映射到查询、键和值向量：
\[
q^{(i)}, k^{(i)} \in \mathbb{R}^{d_k}, \qquad v^{(i)} \in \mathbb{R}^{d_v}.
\]
对于线性注意力变体，键和查询通过特征映射\(\phi(\cdot)\)，循环状态由矩阵\(W^{(i)} \in \mathbb{R}^{d_v \times d_\phi}\)表示。我们将\(W^{(i)}\)解释为存储直到词元\(i\)的键值信息的主要联想记忆。在这种取向下，记忆将变换后的键映射到值，因此在键\(k^{(i)}\)上的读取返回值空间中的向量。为简洁起见，我们记：
\[
\kappa^{(i)} = \phi\left(k^{(i)}\right) \in \mathbb{R}^{d_\phi}
\tag{7}
\]
为特征映射后的键。在循环记忆变体更新\(W^{(i)}\)之后，其词元输出通过用变换后的查询查询记忆得到：
\[
y^{(i)} = W^{(i)} \phi\left(q^{(i)}\right).
\tag{8}
\]

几个量在DeltaNet风格变体中是共享的。使用式(7)中的特征映射键，我们定义当前键处的记忆预测为：
\[
\bar{v}^{(i)} = W^{(i-1)} \kappa^{(i)},
\]