全注意力回归：在百步训练内将全注意力转化为稀疏注意力

# 全面注意力卷土重来：在数百训练步内将全面注意力迁移为稀疏注意力
来源：https://arxiv.org/html/2605.16928

###### 摘要

大语言模型中的长上下文推理受到完全注意力二次成本的瓶颈。现有的高效替代方案通常依赖于原生稀疏训练或启发式令牌驱逐，在效率、训练成本和准确率之间造成了不理想的权衡。在这项工作中，我们表明完全注意力大语言模型本质上是内在稀疏的，并且可以通过极小的适应性调整转化为高度稀疏的模型。我们的方法基于三个观察：(1) 只有一小部分注意力头真正需要完整的远距离上下文处理；(2) 远距离检索主要由一个低维子空间控制，使得可以使用一个16维的索引器高效地检索相关令牌；(3) 有用的令牌预算强烈依赖于查询，这使得动态 top-p 选择比固定的 top-k 稀疏化更合适。基于这些见解，我们提出了 RTPurbo，该方法仅对检索头保留完整的 KV 缓存，并引入一个轻量级令牌索引器用于稀疏注意力。通过利用模型的内在稀疏性，RTPurbo 仅需几百个训练步骤即可实现稀疏化。在长上下文基准测试和推理任务上的实验表明，RTPurbo 在保持近乎无损的准确率的同时，提供了显著的效率提升，包括在 1M 上下文下高达 9.36 倍的预填充加速和约 2.01 倍的解码加速。这些结果表明，无需昂贵的原生稀疏预训练，从标准的完全注意力训练中也能获得强大的稀疏推理能力。

††脚注：† 项目负责人 § 通讯作者 ♯ 在阿里巴巴实习期间完成的工作参见图 1：RTPurbo 的效率和准确率增益概览

## 1 引言

长上下文能力已成为现代大语言模型（LLM）的核心需求，尤其适用于多轮对话、长程推理和文档理解等应用 [deepseekR1, kimi2, qwen1M, gemini25]。然而，完全注意力的成本随着上下文长度快速增长，使得长上下文推理成为一个主要的效率瓶颈。因此，稀疏注意力成为降低推理成本的自然方向 [streamLLM, spargeattn, zucchet2026the]。

尽管该领域最近的许多进展用更高效的替代方案替换了标准的完全注意力，例如 Kimi Delta Attention [kimiteam2025kimilinearexpressiveefficient] 和 DeepSeek Sparse Attention [dsa]，但我们的研究表明，经过完全注意力训练的模型已经展现出显著的内在稀疏性。先前的工作已部分揭示了这一现象。具体来说，稀疏性在头部层和令牌层都存在：大多数头部主要依赖局部信息 [streamingLLM, razorattn, duoattn]，而每个查询只对应一小部分令牌获得显著注意力 [fasa, Quest, snapkv]。这一观察自然引出一个关键问题：将完全注意力模型转化为高度稀疏模型的同时保持其能力，所需的最小手术是什么？

我们识别出三个挑战：

- **头部选择**：需要一个稳健的度量来识别真正需要完整上下文访问的头部。
- **高效令牌索引**：需要一个轻量级的选择器来高效地识别必要的令牌。
- **自适应稀疏性**：由于不同查询需要不同数量的被关注令牌，静态稀疏预算可能导致信息丢失。

我们的方法 RTPurbo 旨在以最小的适应性调整应对这些挑战。RTPurbo 的设计基于 LLM 可解释性和理论分析。先前关于归纳头的研究表明，一些头部通过关注先前相似的令牌来实现检索机制 [olsson2022incontextlearninginductionheads]。后续工作进一步表明，在长上下文设置中，这些头部主要负责远程检索，而其余头部专注于局部上下文 [razorattn]。这一观察激发了我们的头部级设计：我们仅对检索头保留完整的 KV 缓存，而对局部头丢弃远程令牌。

对于检索头，关键挑战是高效地识别相关令牌。我们的分析表明，高频成分对远距离检索贡献甚微，甚至会干扰检索，这表明检索过程主要由一个低维子空间控制。这一假设得到了实验的有力支持：使用我们训练的低维投影器，仅用16个维度就能达到超过90%的召回率。此外，我们的分析表明，静态的 Top-k 选择器在某些情况下可能失败，而 Top-p 选择器能更好地适应注意力分布，并在推理和长上下文任务中带来显著更高的准确率。

最后，我们发现自蒸馏对恢复稀疏化模型的性能特别有效。将稀疏模型的输出与原始模型对齐，大大降低了过拟合的风险，并且此对齐阶段仅需几百个训练步骤（约 100 万标签令牌）。这一结果进一步支持了我们的论断：RTPurbo 仅对原始模型进行了最小手术。

据我们所知，RTPurbo 是第一个通过轻量级持续训练实现这种近乎无损压缩的方法。结合我们自定义的稀疏内核，RTPurbo 在预填充阶段实现了高达 9.36 倍的加速，在解码阶段实现了 2.01 倍的加速（图 1 (https://arxiv.org/html/2605.16928#S0.F1)）。重要的是，RTPurbo 的稀疏化范式仍然高度可解释。更广泛地说，我们的结果强调了一个被忽视的观点：对于完全注意力模型，即使没有原生稀疏训练，一个完全训练好的模型也可以以非常小的额外成本进行稀疏化，同时保持强大的性能。这一发现表明，完全注意力训练仍然是一个极具竞争力且实用的选择。

## 2 RTPurbo 背后的见解

参见图 2：不同于大多数主要关注局部信息的注意力头，检索头会关注与当前查询令牌在语义上相关的区域（例如相似模式），即使这些区域在上下文中相距甚远。

### 2.1 头部特化作为稀疏注意力的自然先验

最近的研究表明，预训练 LLM 中的注意力头并非同质的，而是特化为不同的功能角色。特别是，先前的工作表明，只有一小部分头部负责检索远距离的相关内容，而许多其他头部主要处理局部信息 [duoattn, razorattn]。我们将这一子集称为*检索头*。其特征行为是将强烈的注意力放在包含语义相关内容的早期上下文上，从而展现出信息检索模式，如图 2 (https://arxiv.org/html/2605.16928#S2.F2) 所示。

这一观察为我们方法提供了一个重要的设计动机：*我们可以自然地利用模型已经形成的稀疏结构。*具体来说，我们仅对检索头保留完整的 KV 缓存，而对于其余已经本质上稀疏的头部，我们可以安全地丢弃远程令牌。

### 2.2 RoPE 为检索头诱导出可压缩的几何结构

检索头应将高注意力分配给语义相关的令牌，即使它们相距甚远。

然而，检索头的这一特性乍看之下似乎与 RoPE [rope] 存在张力。对于一个在位置 m 的查询令牌和在位置 n 的键令牌，维度 d=2D，RoPE 通过旋转矩阵注入位置信息：

R_i(m) = ( cos(mθ_i)  -sin(mθ_i) ; sin(mθ_i)  cos(mθ_i) ), q_m = R(m)q, k_n = R(n)k, (1)

其中 R(m) = diag(R_1(m), ..., R_D(m))，且 θ_i 随着通道索引增加而减小。产生的查询-键分数仅依赖于相对偏移 Δ=m-n：

s(m,n) = q_m^⊤ k_n = ∑_{i=1}^D [ a_i(q,k) cos(θ_i Δ) + b_i(q,k) sin(θ_i Δ) ], (2)

其中 a_i 和 b_i 是由第 i 个旋转对产生的双线性系数。方程 (2 (https://arxiv.org/html/2605.16928#S2.E2)) 直接揭示了关键区别：高频成分随 Δ 快速变化，并在长距离上变得距离敏感，而低频成分变化平缓，更好地保留了检索信号。这引出了我们的第二个核心见解：*我们可以在一个低得多的维度空间中重建检索头的注意力。*

因此，我们将这种低频结构用作紧凑的检索子空间，从而实现低成本的令牌选择，无需全维度评分。

### 2.3 检索头需要动态阈值化

参见 (a) 由查询令牌“Galápagos”在长段落中触发的扩散检索。
参见 (b) NIAH 查询中的集中检索。

图 3：检索头的行为强烈依赖于查询。(a) 查询令牌“Galápagos”引发了对许多语义相关早期令牌的扩散检索：需要约 8k 个令牌才能恢复 90%+ 的注意力质量，而 top-4k 仅恢复了约 75%。(b) 对于“大海捞针”查询，检索高度集中：两个令牌恢复了 96.6% 的注意力质量，而 top-4k 保留了许多不必要的令牌。

剩下的问题是，一旦能够高效估计相关性，检索头应该保留多少令牌。我们的发现表明，这个数量根本上依赖于查询。即使在同一检索头内，不同的输入也可能诱导出非常不同的模式：某些查询会触发对许多远距离令牌的广泛检索，而其他查询则仅锁定少数关键令牌。因此，所需的稀疏程度并非头部的固定属性；它会随查询而变化。

图 3 (https://arxiv.org/html/2605.16928#S2.F3) 说明了这一点。在一种情况下，查询激活了一个广泛的语义场，因此检索头必须保留一个宽广的支撑集以恢复大部分注意力质量。在另一种情况下，查询仅需恢复一个关键事实，因此头部自然高度集中。

表 1：固定的 top-k 牺牲召回率来换取稀疏性：top-16k 比 top-p 多计算约 8k 个额外令牌，但仅多恢复了 3.8% 的注意力质量。

这正是固定预算规则（如 top-k 采样）变得有问题的地方。当 k 太小时，扩散查询恢复的注意力质量太少，近似变得不准确。当 k 太大时，保留的集合不再足够稀疏，额外的计算大多被浪费。表 1 (https://arxiv.org/html/2605.16928#S2.T1) 使这种权衡具体化：top-16k 比动态 top-p 仅多恢复了 3.8% 的注意力质量，但需要计算约 8k 个额外令牌。因此问题不在于选择一个更好的全局 k；任何固定的 k 都与检索头依赖于查询的特性不匹配。

## 3 方法

我们介绍 RTPurbo，一个具有精确令牌级稀疏计算的头部级注意力框架。本节组织如下。我们首先在 3.1 节 (https://arxiv.org/html/2605.16928#S3.SS1) 描述如何通过离线校准识别检索头。接着在 3.2 节 (https://arxiv.org/html/2605.16928#S3.SS2) 介绍我们的稀疏计算模式。然后在 3.3 节 (https://arxiv.org/html/2605.16928#S3.SS3) 描述 RTPurbo 所需的两阶段训练管道。最后在 3.4 节 (https://arxiv.org/html/2605.16928#S3.SS4) 描述硬件感知的解码内核。

### 3.1 离线头部级校准

为了识别检索头，我们构建一个轻量级的校准序列，方法是在从 FineWeb [fineweb] 采样的长文档的开头和结尾插入一个相同的“针”片段。我们通过测量从后面的针向前面的针分配的注意力质量来量化头部的检索能力。设 N_pre 和 N_post 分别表示较早和较晚针片段的令牌索引。头部 h 的检索分数紧凑地定义为：

R_h = (1/|N_post|) ∑_{t∈N_post} ∑_{j∈N_pre} A_h(t,j), (3)

其中 A_h(t,j) 表示从令牌 t 到令牌 j 的归一化注意力分数（即 softmax 后）。

头部的检索行为高度稳定，且在很大程度上与输入无关。因此，在实践中，仅对单个长文本序列运行此校准，就足以稳健地对所有查询头部进行评分并划分为检索集 H_ret（得分最高的头部）和局部集 H_loc。此划分过程仅离线执行一次。

参见图 4：RTPurbo 的整体架构。

### 3.2 自适应稀疏注意力机制

在推理过程中，局部头 h∈H_loc 在预填充和解码阶段均一致地应用带有注意力汇聚点 [streamingLLM] 的滑动窗口。相比之下，检索头 h∈H_ret 在预填充期间执行全密集注意力以构建完整的 KV 缓存，但在解码期间切换到查询感知的动态稀疏选择。正如在第 2.2 节 (https://arxiv.org/html/2605.16928#S2.SS2) 中分析的，高频 RoPE 成分会降低远距离亲和性。为了避免这一点，我们使用低秩投影 W^Q_h, W^K_h ∈ R^{r×d_h} (r ≪ d_h) 来估计查询-键相关性，这些投影应用于 RoPE 注入之前的特征：

s_h(m,n) = (W^Q_h q_{m,h}^{pre})^⊤ (W^K_h k_{n,h}^{pre}), (4)

其中 q_{m,h}^{pre} 和 k_{n,h}^{pre} 是 RoPE 前的表示。然后我们根据投影得分构建一个动态激活集，并计算稀疏注意力：

O_h(m) = ∑_{n∈S_h(m)} (exp(q_{m,h}^⊤ k_{n,h} / √d_h) / ∑_{j∈S_h(m)} exp(q_{m,h}^⊤ k_{j,h} / √d_h)) v_{n,h}, S_h(m) = Top-P(s_h(m,·), p). (5)

通过这种方式，低秩 RoPE 前投影严格作为高效路由机制，而最终令牌生成保留了完整的特征空间和精确的相对位置几何。对于 MQA 和 GQA 模型，产生的稀疏性应从两个角度解释，因为我们的头部划分是针对查询头定义的。*计算稀疏性*在查询头级别度量，可以视为头部平均关注的令牌数。*内存稀疏性*在 KV 头级别度量：对于每个 KV 头，