动态线性注意力

# 动态线性注意力

来源：https://arxiv.org/html/2606.10650  
1\]俄亥俄州立大学 2\]密歇根大学 3\]字节跳动种子团队 \\contribution\*同等贡献。

Hui Shen、Boyuan Zheng、Xueshen Liu、Minkyoung Cho、Zhongwei Wan、Zesen Zhao、Zhuoqing Mao、Shen Yan、Mi Zhang \\[\\[\[[[email protected]](https://arxiv.org/html/2606.10650v1/mailto:[email protected]) | [[email protected]](https://arxiv.org/html/2606.10650v1/mailto:[email protected])

###### 摘要

大型语言模型（LLM）向长上下文的扩展，其根本限制在于标准注意力的二次复杂度，这促使人们采用具有次二次成本的线性注意力机制。为提升长上下文下的表示能力，近期研究以多状态方式组织记忆。然而，现有的多状态线性注意力方法依赖固定的状态合并策略，无法适应动态变化的令牌重要性，导致关键令牌被不可逆地掩盖，并随着序列长度增加而造成严重的误差累积。为解决这一局限，我们提出DLA，一种面向多状态线性注意力的动态记忆建模框架。DLA引入：(i) *信息感知动态状态合并*，根据令牌级别的信息变化自适应确定状态边界，在语义转换周围保持高分辨率表示，同时积极总结稳定区域；以及 (ii) *容量约束记忆建模*，通过选择性合并相邻的低信息状态，维持一个固定大小、按时间顺序排列的状态缓存，从而在信息损失最小化的同时控制记忆增长。我们在两种不同的线性注意力模型上预训练DLA，并在三类16个数据集上进行评估。实验结果表明，DLA优于现有最先进方法。

\\correspondence

Xin Wang 邮箱：，Mi Zhang 邮箱：

## 1 引言

大型语言模型（LLM）在广泛的自然语言理解和生成任务中展现了卓越的能力。然而，由于标准自注意力[wan2023efficient, wang2024iot, DBLP:conf/iclr/WanWZXTZWLXW025]的二次计算和记忆复杂度，将LLM扩展至长上下文场景仍是一个根本性挑战。这一局限促使人们广泛研究高效的注意力机制，以便在不从头重新训练的情况下实现长序列建模。在这些方法中，线性注意力[DBLP:conf/nips/YangWZSK24, DBLP:conf/iclr/YangKH25]成为一个有前景的方向，因为它以次二次复杂度近似全注意力，并具有良好的长上下文可扩展性。

为进一步提升线性注意力在长序列下的表示能力，近期工作以多状态方式组织历史上下文，将长令牌历史划分为块并总结为紧凑的记忆状态。代表性方法如对数线性注意力[DBLP:journals/corr/abs-2506-04761]展示了在长上下文推理中的高效性和实用性。通过基于总结状态而非单个令牌进行操作，这些方法显著减少了记忆占用量和计算成本。

尽管取得了成功，现有的多状态线性注意力方法在上下文长度增加时仍存在显著的性能下降。这一局限源于固定记忆构建策略与长序列非均匀、动态变化的信息结构之间的根本不匹配。具体来说，当前方法通常依赖固定的块大小或基于规则的合并调度，隐式假设序列中信息密度均匀。这种设计无法适应动态出现的语义转换，迫使关键令牌过早地被吸收到粗略总结中。此外，在固定策略下做出的合并决策是不可逆的：一旦异构的令牌被压缩成一个状态，它们的个体贡献就无法恢复，从而导致误差累积。

这些观察表明，有效的长上下文线性注意力需要既信息感知又容量控制的记忆建模机制。一方面，状态构建应适应局部表示变化，为语义波动区域分配更高分辨率，同时积极总结稳定区域。另一方面，记忆状态的总数必须明确受限，以确保推理过程中可预测的计算和记忆成本。

在本工作中，我们提出动态线性注意力（DLA），一种解决这些挑战的多状态线性注意力新框架。DLA与先前方法在两个关键方面不同。首先，DLA引入信息感知动态状态合并，根据令牌级别的信息变化动态确定状态边界。DLA不依赖固定合并策略，而是评估每个新来的令牌相对于当前记忆状态的表示变化，合并低变化令牌，同时在语义转换点开启新状态。其次，DLA融入容量约束记忆建模，维持一个固定大小、按时间顺序排列的状态缓存。当缓存达到容量时，DLA选择性合并相邻的低信息状态，在保留时间顺序的同时最小化信息损失。

我们在两个线性注意力骨干模型（Mamba-2-780M和Gated DeltaNet-1.3B）上预训练DLA，遵循[DBLP:journals/corr/abs-2506-04761]中的设计。我们在16个数据集上进行评估，涵盖三个方面：八个常识推理基准、六个上下文内检索数据集和两个长上下文建模数据集。我们强调三个主要发现：(1) DLA在所有任务上始终优于最先进的多状态方法对数线性注意力。(2) 当应用于Mamba-2时，DLA变体甚至达到与类似参数预算的全注意力Transformer相当的性能。(3) DLA实现了卓越的效率，比对数线性注意力提供更高的吞吐量和更低的运行时内存消耗。

## 2 预备知识

我们考虑一个输入长度为 \(T\)、隐藏维度为 \(d\) 的序列建模任务。令 \(\mathbf{Q},\mathbf{K},\mathbf{V} \in \mathbb{R}^{T \times d}\) 分别表示查询、键和值矩阵。标准自注意力计算输出 \(\mathbf{O} \in \mathbb{R}^{T \times d}\) 为：

\[\mathbf{O} = \mathrm{softmax}(\mathbf{Q}\mathbf{K}^{\top}\odot\mathbf{M})\mathbf{V},\]

其中 \(\mathbf{M}\) 是因果掩码。虽然有效，但此操作在计算和记忆上对 \(T\) 具有二次复杂度，这推动了次二次注意力机制的发展。在本节中，我们回顾线性注意力及其多状态变体，这些构成了我们方法的基础。

### 2.1 线性注意力

线性注意力通过移除 softmax 归一化来缓解 Transformer 的二次成本，从而允许利用结合律重新排列计算顺序。因果线性注意力层可以以并行形式写为：

\[\mathbf{O} = (\mathbf{Q}\mathbf{K}^{\top}\odot\mathbf{M})\mathbf{V}, \quad \mathbf{M}_{ij} = \mathbb{I}(i \geq j).\] (1)

该公式等价于一个递归实现。令 \(\mathbf{q}_t,\mathbf{k}_t,\mathbf{v}_t \in \mathbb{R}^d\) 分别表示时间步 \(t\) 的查询、键和值向量。线性注意力维护一个单一状态矩阵 \(\mathbf{S}_t \in \mathbb{R}^{d \times d}\)，该矩阵总结所有过去令牌：

\[\begin{aligned}
\mathbf{S}_t &= \mathbf{S}_{t-1} + \mathbf{v}_t\mathbf{k}_t^{\top}, &\text{}&&(2)\\
\mathbf{o}_t &= \mathbf{S}_t\mathbf{q}_t. &\text{}&&(3)
\end{aligned}\]

这种递归形式实现了具有常量记忆的线性时间推理，但将整个历史压缩为一个单一状态，这在长上下文下可能限制表示能力。我们用 \(\phi(\cdot)\) 表示线性注意力中使用的特征映射。除非另有说明，\(\phi: \mathbb{R}^d \rightarrow \mathbb{R}^d\) 按照先前工作实现为单位映射或可学习的线性投影。

### 2.2 带Delta规则的线性注意力

为改进状态跟踪并引入受控遗忘，DeltaNet [DBLP:conf/nips/YangWZSK24] 以 delta 风格的更新规则扩展了线性注意力：

\[\mathbf{S}_t = \mathbf{S}_{t-1}(\mathbf{I} - \beta_t\mathbf{k}_t\mathbf{k}_t^{\top}) + \mathbf{v}_t\mathbf{k}_t^{\top},\] (4)

其中 \(\beta_t\) 是数据相关的步长。虽然该公式相比纯累加器提高了自适应性，但它仍依赖单一全局状态，因此无法在长序列上选择性保留细粒度信息。

### 2.3 多状态线性注意力

为在保持次二次复杂度的同时增加建模能力，近期工作通过将历史上下文划分为片段并将每个片段总结为单独状态，以多状态方式组织线性注意力 [DBLP:journals/corr/abs-2506-04761, DBLP:journals/corr/abs-2507-04416]。其中，对数线性注意力 [DBLP:journals/corr/abs-2506-04761] 将单一递归状态替换为通过因果前缀的 Fenwick 树分解构建的对数级多尺度状态。具体来说，在时间步 \(t\)，前缀 \([0,t]\) 被分解为最多 \(L = \lceil \log_2(t+1) \rceil + 1\) 个不相交的桶 \(\{B_t^{(\ell)}\}_{\ell=0}^{L-1}\)，当前位置附近分辨率更细，远处历史分辨率更粗。相应的线性注意力状态和最终聚合输出计算如下：

\[\mathbf{S}_t^{(\ell)} = \sum_{s \in B_t^{(\ell)}} \mathbf{v}_s\mathbf{k}_s^{\top} \in \mathbb{R}^{d \times d}, \quad \mathbf{o}_t = \sum_{\ell=0}^{L-1} \lambda_t^{(\ell)} \mathbf{S}_t^{(\ell)} \mathbf{q}_t.\] (5)

这种设计实现了 \(O(T\log T)\) 的训练复杂度和每个解码步 \(O(\log T)\) 的时间与记忆成本。然而，其记忆状态的粒度由固定的层次化调度决定，与令牌级别的表示变化无关。因此，语义上显著的令牌可能过早地被吸收到粗略总结中，并且在关键位置引入的干扰可能通过固定的多尺度状态传播。这一局限促使了信息感知和自适应记忆构建的需求，我们在下一节中解决这一问题。

## 3 动态线性注意力 (DLA)

参见图注：图1：DLA示意图。图1 (https://arxiv.org/html/2606.10650#S3.F1) 提供了DLA的概览。DLA是一种信息感知的线性注意力框架，它动态构建一组紧凑的记忆状态，用于高效的长上下文建模。与依赖于固定时间调度或预定义块边界的先前方法不同，DLA基于令牌级别的信息变化自适应地确定状态粒度。具体来说，令牌被顺序处理。对于每个新令牌，DLA计算一个轻量级的*状态信息得分*，衡量其相对于最近记忆状态的表示变化。信息变化低的令牌被合并到当前状态，而表现出显著漂移的令牌则开启一个新状态。这允许在语义转换周围进行细粒度建模，同时积极总结稳定的令牌区域。为了控制记忆和计算，DLA维护一个容量受限的状态缓存。当缓存达到其最大大小时，合并两个信息密度最低的相邻状态，在保留时间顺序的同时最小化信息损失。由此产生的记忆由一组固定大小、按时间顺序排列的总结状态组成。在每个解码步，DLA通过线性注意力公式对所维护的记忆状态进行注意力计算来产生输出，其中每个状态以查询相关权重贡献。信息感知的状态构建和容量约束的记忆建模共同使DLA能够实现自适应分辨率、稳定的推理成本和高效的长上下文表示。

算法1 信息感知动态状态合并
1: 输入：令牌状态 \(\{s_t\}_{t=1}^T\)
2: 输出：记忆状态 \(\mathcal{M} = \{S_i\}\)
3:
   \(\mathcal{M} \leftarrow []\)
4: for \(t = 1\) to \(T\) do
5:   if \(\mathcal{M}\) 为空 then
6:       \(\mathcal{M} \leftarrow \{s_t\}\); continue
7:   end if
8:    \(S \leftarrow\) \(\mathcal{M}\) 中的最后一个状态
9:    \(I_t \leftarrow \frac{\|s_t - S\|_F}{\|S\|_F + \varepsilon}\)
10:  用 \(\mathrm{Merge}(S, s_t)\) 替换 \(\mathcal{M}\) 中的最后一个状态
11:   if \(I_t \geq \tau\) then
12:      将 \(s_t\) 追加到 \(\mathcal{M}\);
13:   end if
14: end for
15: return \(\mathcal{M}\)

### 3.1 信息感知动态状态合并

动机：现有的多块线性注意力方法通常依赖固定调度（例如，每 \(K\) 个令牌进行块合并）[DBLP:journals/corr/abs-2506-04761] 或硬性的基于规则边界 [DBLP:journals/corr/abs-2507-04416] 来确定应合并到总结状态的历史令牌块。虽然此类设计提高了记忆和计算效率，但它们很大程度上对序列的语义演变不敏感。实际上，信息密度高度非均匀：关键的语义转换可能突然发生，而连续大段的令牌可能局部冗余。因此，固定或硬性块策略通常遭受*两个关键限制*。首先，它们无法适应动态出现的语义变化，仅仅因为达到预定义边界，就迫使重要的转换过早地被吸收到粗略总结中。其次，未考虑局部语义连续性而做出的合并决策本质上是不可逆的：一旦令牌在固定策略下被合并，即使后续上下文揭示其重要性，它们的个体贡献也无法恢复。合并决策与真实语义结构之间的这些不匹配导致次优的生成，并最终降低表示质量。

下面，我们提供一个理论证明，说明为什么固定合并策略是次优的。

###### 定理 3.1 (状态偏差)。

令 \(\{u_t\}_{t=1}^T \subset \mathbb{R}^d\) 表示对线性注意力状态的每令牌加性贡献。考虑一个将令牌列表 \(\{1,\dots,T\}\) 划分为 \(m\) 个不相交连续块 \(\{C_i\}_{i=1}^m\) 的阻塞策略 \(\pi\)。对于每个块 \(C_i\)，令 \(\bar{u}_i \in \mathbb{R}^d\) 为一个代表性总结向量。那么，对于任何查询向量 \(q \in \mathbb{R}^d\)，精确输出 \(y(q)\) 和总结输出 \(\tilde{y}_\pi(q)\) 为：

\[y(q) \triangleq \sum_{t=1}^T \langle q, u_t \rangle, \quad \tilde{y}_\pi(q) \triangleq \sum_{i=1}^m \sum_{t \in C_i} \langle q, \bar{u}_i \rangle\] (6)

由总结引起的偏差为 \(\operatorname{Err}(\pi; q) \triangleq |y(q) - \tilde{y}_\pi(q)|\)。