🤖 machine learning

SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers

SpotAttention 是一种轻量级的、插件式的块稀疏路由机制，它附加在冻结的预训练 Transformer 上，通过 KL 蒸馏来学习注意力分布，从而实现高达 128K token 的准确长文本推理，且与现有基准相比，其解码速度显著提升并降低了内存使用量。

原作者： Huzama Ahmad, Se-Young Yun

发布于 2026-06-23

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Huzama Ahmad, Se-Young Yun

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你有一位非常聪明的图书管理员（AI 模型），她读过数百万本书。当你向她提问时，这位图书管理员通常会试图记住她读过的每一页内容，以此来寻找答案。

问题在于，随着图书馆规模的扩大，这种“记住一切”的方法变得极其缓慢且昂贵。这就像是在一个每秒钟都在不断变大的草堆中寻找一根特定的针。

SpotAttention 是一个全新的、轻量级的“助手”，它能帮助图书管理员在无需重新阅读整个图书馆的情况下，找到正确的页面。以下是它的工作原理，通过简单的概念进行拆解：

1. 问题所在：“沉重的背包”

当 AI 尝试根据一份非常长的文档（例如一整部小说或庞大的代码库）来回答问题时，它必须保留一个包含目前所见所有重要词汇的“背包”。

旧方法： 每当 AI 思考下一个词时，它都会把整个背包倾倒在桌子上，以寻找线索。随着文档变长，背包变得越来越重，桌面也变得过于拥挤。这使得 AI 运行起来既慢又贵。
“稀疏化”（Sparse）的想法： 其他研究人员尝试提出：“嘿，AI 其实只需要看其中的几页特定页面，对吧？我们只看那些就行了。”但确定“看哪些页面”本身就是一个缓慢且昂贵的任务。这就像是仅仅为了决定看哪几页，就专门雇佣了一个新人，而这个人的效率几乎和读整本书一样慢。

2. 解决方案：“观察员”（SpotAttention）

作者创建了 SpotAttention，它就像是一个附着在图书管理员身上的微型、超快速的“观察员”。

它是如何学习的： 观察员不需要从零开始学习。它观察“专家级”图书管理员（预训练好的 AI）的工作过程。它通过模仿专家的眼神，学会识别专家自然会看向哪些页面。
神奇的技巧： 观察员并不需要逐字阅读来决定保留哪些内容，它通过观察文本块（如段落）并快速进行评分。这就像是一位图书管理员，只需扫一眼书架，就能瞬间知道：“我只需要取出这三本书，其他的可以留在书架上。”

3. “双 Top-p”规则：一个聪明的预算

论文引入了一个巧妙的规则，称为 Dual Top-p。想象一下，图书管理员有一个可以查看多少页面的“预算”。

旧方法： 有些系统会规定：“你只能看前 50% 的页面。”这是僵化的。有时你需要看 80%，有时只需要 20%。
SpotAttention 的方式： 观察员会观察页面的“重要性”并表示：“好吧，针对这个问题，我们需要保留前几页（开头）、最后几页（刚刚提到的内容）以及中间最关键的几页。”
它会动态调整预算。如果问题很简单，它抓取的页面就少；如果问题很复杂，它抓取的页面就多。它会自动完成这一切，而不需要第二个缓慢的步骤来做决策。

4. 结果：速度与记忆

论文在多个大型 AI 模型（特别是 Qwen 系列）上测试了其处理超长上下文（高达 128,000 个词）的能力。

速度： 在 128,000 词的长度下，SpotAttention 比标准方法（FlashAttention）快 3.9 倍，比目前最好的替代方案（Twilight）快 1.8 倍。
准确度： 尽管跳过了大部分文本，AI 得到答案的正确程度与阅读全部内容时一样。它并没有损失任何“智能”。
内存： 观察员的“笔记本”（用于做决策的缓存）可以通过特殊的压缩技术缩减到极小的体积，且不会损失准确性。这节省了大量的计算机内存。

5. 为什么它与众不同

无需重新训练： 你不需要重新教整个 AI。你只需要将这个微小的观察员插件到已经完成且正在运行的 AI 上即可。
是学习而不仅仅是硬编码： 之前的方法使用固定的规则（例如“始终跳过中间部分”）。SpotAttention 则通过学习什么是重要的模式，使其变得更聪明、更灵活。
适用于各种场景： 他们在不同规模的 AI 模型（从 40 亿参数的小模型到 320 亿参数的大模型）上进行了测试，它在所有模型上都表现出色。

总结： SpotAttention 是一个轻量级的、经过学习的助手，它通过快速识别最重要的部分来进行关注，从而帮助 AI 模型阅读长文档。它使 AI 在处理长文本时显著变快且成本更低，同时不会降低其智能水平。

技术摘要：SpotAttention

问题陈述

长文本能力已成为预训练大语言模型（LLM）的标准配置，使其能够处理诸如整个代码库或基于大量证据进行推理的任务。然而，自注意力机制的计算成本随序列长度呈平方级增长。在预填充（prefill）阶段，计算量随序列长度呈二次方增长；而在解码（decode）阶段，每个新 Token 都需要重新读取随上下文规模线性增长的键值（KV）缓存。虽然 FlashAttention 通过避免实例化完整的注意力矩阵提高了 I/O 效率，但它并未改变根本性的二次方预填充成本以及每步内存读取随长度线性增长的问题。

现有的稀疏注意力方法试图通过仅关注过去 Token 的相关子集来缓解这一问题。然而，选择该子集的过程通常非常昂贵。例如 Quest 对 Token 块进行评分，而 Twilight 则增加了 top- $p$ 剪枝，但选择器本身成为了一个新的瓶颈。此外，无训练（training-free）的选择器通常依赖于手工设计的启发式算法，这些算法可能难以泛化，或者需要多个选择阶段，从而产生累积开销。

方法论：SpotAttention

SpotAttention 是一种插件式的块稀疏路由机制，旨在为已训练好的冻结 Transformer 模型进行功能增强。它由一个轻量级的“选择器”（或索引器）组成，并附加在每一个全注意力层上。

1. 选择器架构

选择器是一个极小的多头 Q-K 评分器，具有缩减后的宽度（ $d_{idx}$ ）和 $H_{idx}$ 个头。它不是对单个 Token 进行评分，而是对查询（Query）和键（Key）块（例如每 16 个 Token 一个块）进行操作。

评分机制： 对于位置 $q$ 处的查询和位置 $t$ 处的键块，选择器计算加权 ReLU 评分：
$s_s[q, t] = \sum_{h=1}^{H_{idx}} w_h \text{ReLU}(\mathbf{q}^I_h \cdot \mathbf{k}^I_{h,t})$
其中 $\mathbf{q}^I_h$ 和 $\mathbf{k}^I_{h,t}$ 是逐头投影， $w_h$ 是由隐藏状态导出的混合权重。
效率： 通过进行块级（block-wise）评分，选择过程可以作为单次张量核心（tensor-core）矩阵乘法运行，避免了逐 Token 操作带来的开销。

2. 通过 KL 蒸馏进行训练

选择器针对冻结骨干网络（即“教师”）的稠密注意力分布进行独立训练。

目标： 训练过程最小化教师的注意力分布（ $p_\tau$ ）与选择器的估计分布（ $p_\sigma$ ）之间的 Kullback-Leibler（KL）散度。
变体： 论文探讨了两种 KL 变体：
- DenseKL： 在所有 $T$ 个键上计算损失。
- SparseKL： 仅在选择器自身的 top- $K$ 被选集合上计算损失，并对分布进行重新归一化。
收敛性： 在各种骨干网络规模（4B 到 32B）和架构下，选择器都能快速收敛（约 200 步），学习匹配教师的分布形状。

3. 双重 Top- $p$ 选择规则

在推理时，选择器会对所有过去的块进行评分。为了确定应关注哪些块，SpotAttention 引入了一种双重 top- $p$ 规则，直接从经过校准的选择器分布中读取预算：

保留块： 固定前缀的“汇点”（sink）块（例如前 128 个 Token）和固定的后缀“近期”（recency）块（例如最后 256 个 Token）始终被保留，以防止它们被核采样（nucleus selection）所主导。
核采样（Nucleus Selection）： 对剩余分布（排除汇点和近期块后的部分）应用 top- $p$ 采样。当累积概率质量达到阈值 $p$ 时停止选择。
动态预算： 这种方法为每个查询和每层生成一个可变的预算，能够适应内容变化，而无需单独的剪枝阶段。设置最小底限（ $K_{min}$ ）以防止出现核采样过小的退化情况。

4. 量化

选择器的 K-cache 可以量化为 INT4 或 FP4 微尺度。由于选择器的缓存与骨干网络的 KV 缓存是独立的，这可以在不影响骨干网络精度或准确性的情况下，将选择器的内存占用减少约 3.5 倍。

核心贡献

插件式块稀疏注意力： 一种通过 KL 蒸馏训练的轻量级选择器，可附加到冻结的 Transformer 上。它作为单次张量核心矩阵乘法运行，避免了以往工作中沉重的选择器开销。
双重 Top- $p$ 规则： 一种内容自适应的选择机制，直接从学习到的分布中导出每查询、每层的预算。它保留了汇点和近期块以确保稳定性，在保持较小平均预算的同时，达到了静态 top- $K$ 的准确度。
训练与推理行为： 论文证明了选择器能够学习到骨干网络中每层预算的不对称性（即早期层比后期层关注更多的上下文），这是静态 $K$ 无法复制的模式。该方法从 16K 的训练长度泛化到 128K 的评估长度，且未出现准确度下降。

实验结果

该方法在五个骨干网络（Qwen3: 4B–32B; Qwen3.5: 4B–9B）上进行了评估，测试了四个基准测试（RULER, BABILong, InfiniteBench, LongBench-v2），上下文长度最高达 128K。

准确度： 在上下文长度高达 128K（训练长度的 8 倍）时，SpotAttention 能匹配稠密骨干网络的准确度。在预算为 $K=0.5L$ 时，在所有模型上均保持在与稠密基准线一个自助法标准误差之内。
延迟与吞吐量：
- 在 $L=128K$ 的 Qwen3-8B 上，SpotAttention 的运行速度比 FlashAttention 快 3.9 倍，比 Twilight（最强的无训练基准）快 1.8 倍。
- 解码吞吐量在 128K 时维持在 70 tokens/s，在 512K 时维持在 36 tokens/s；而 FlashAttention 和 Twilight 分别在 256K 和 512K 时发生显存溢出（OOM）。
- 选择器本身的运行速度显著快于 Twilight 的选择器（在 128K 时为 1.55 ms vs 4.59 ms）。
量化： 将选择器的 K-cache 量化为 INT4 或 FP4 可将其内存占用缩小 3.5 倍，且没有明显的准确度损失。虽然这不会提高单流解码延迟（受限于骨干网络 KV 带宽），但能显著降低内存使用量和预填充成本。

重要性与声明

论文将 SpotAttention 定位为解决稀疏注意力中“选择器瓶颈”问题的方案。通过利用 KL 蒸馏学习校准分布，该选择器实现了一种基于核采样的动态预算，能够根据特定的查询和层进行调整。

作者声称：

效率： 在长上下文（高达 512K）下，它实现了对比方法中的最高吞吐量，克服了导致稠密和其他稀疏基准线停滞的内存限制。
泛化性： 在 16K 上下文长度下训练的选择器可以有效地泛化到 128K 及更高长度，无需重新训练。
模块化： 它是一种“插件式”解决方案，不需要重新训练骨干网络，因此适用于现有的预训练模型。
局限性： 作者指出，目前评估的最大骨干网络为 32B，且结果目前仅限于 Qwen 系列及其英文基准测试。他们还承认，尚未对稀疏预填充（sparse-prefill）延迟进行基准测试，尽管其架构支持该功能，且双重 top- $p$ 规则依赖于选择器的良好校准（在 SparseKL 训练下成立，但在 DenseKL 下会失效）。