💻 computer science

Uncertainty-gated selection for block-sparse attention

本文引入了一种用于块稀疏注意力的不确定性门控路由机制，该机制通过为具有模糊 top-k 分数的查询动态扩展所选键块，在保持多种模型架构下接近稠密效率的同时，显著提升了长上下文检索的准确率和召回率。

原作者： Thomas Rossi

发布于 2026-07-10

📖 1 分钟阅读☕ 轻松阅读

原作者： Thomas Rossi

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

想象一下，你正试图在一大堆干草中寻找一根特定的针，但你手里只有一个微型手电筒，而且有一个非常严格的规则：你一次只能照亮几小块干草。这正是现代人工智能在尝试一次性阅读超长文档（比如一整部小说）时所面临的问题。

问题所在：“近视”的手电筒

大多数 AI 模型使用一种叫做**块稀疏注意力（block-sparse attention）**的技巧来节省时间。与其阅读整个 10 万字的完整故事，它们会将故事切分成若干块（blocks），并使用一个“选择器”来挑选出前 k 个最有趣的块。

但问题在于，这个选择器是**近视（myopic）**的。想象一下，选择器正在观察两块干草。块 A 的得分是 9.9，而块 B 的得分是 9.8。规则规定：“只选前 1 名。”选择器立即选择了块 A，并将块 B 丢弃了。

论文指出，这是一个错误的举动。如果块 B 实际上包含了问题的答案，而那微小的分差仅仅是一个偶然的误差呢？一旦块 B 被丢弃，AI 就再也找不回它了。这就像一名侦探因为一个线索“差点就没那么好”就把线索扔掉，结果后来发现自己竟然需要那个线索来破案一样。

解决方案：“不确定性门控”智能开关

作者们（由 Thomas Rossi 领导）提出了一个聪明的修复方案，叫做不确定性门控选择（Uncertainty-Gated Selection）。你可以把它想象成给手电筒增加了一个“置信度计”。

在 AI 做出最终决定之前，它会问自己：“我对挑选出正确的块有多大的把握？”

置信度检查： AI 会观察排名靠前的块的得分。如果排名第一的块比第二名好得多（差距很大），AI 就很有信心。它会遵循规则，仅挑选前 k 个块。
“等等，也许呢？”时刻： 如果排名第一的块和排名第二的块得分非常接近（差距极小），AI 会意识到：“喔！我不确定！我可能会丢掉正确答案。”
安全网： 当 AI 感到不确定时，它会触发一条特殊规则：“预算翻倍！” 它不再仅仅挑选 k 个块，而是针对故事中的这个特定部分抓取 2k 个块。它会多花一点能量来确保万无一失。

这并不是一个改变整个 AI 的魔法咒语。它是一个巧妙的、位于现有 AI 选择方法之上的微型智能层。它就像一个只有在飞行员看起来困惑时才会接管方向盘的副驾驶。

论文实际发现了什么（证据）

作者们不仅仅是在猜测；他们在四个不同的 AI 模型（包括 Qwen 和 Mistral）以及两个主要的测试集上进行了测试。以下是数据说明：

巨大的胜利： 在一个名为 LongBench-v2 的困难测试中，标准方法（仅仅挑选前 k 个）获得的“配对召回率（paired recall）”得分为 0.47。这意味着它找到正确线索的概率不到一半。而新的“不确定性门控”方法将该得分提升到了 0.75。这是一个高达 28 个百分点 的巨大飞跃。
速度： 你可能会认为检查不确定性会降低速度。令人惊讶的是，事实并非如此。在极长长度（128K token）下，新方法的运行时间为原始、缓慢的“稠密（dense）”方法（即阅读全部内容的方法）的 0.62×。它实际上比那些标准的快捷方法更快，同时更加聪明。
“大海捞针”测试： 在一个名为 RULER NIAH 的合成测试中（AI 必须找到隐藏的具体事实），新方法帮助 AI 找到了完美（但缓慢）方法所能找到答案的 0.81 到 0.89，同时运行速度依然很快。

论文排除了什么（“禁区”）

了解这种方法不做什么非常重要，作者对此也非常明确：

它不是短篇故事的魔力修复方案： 作者在 LongBench-v1 上测试了该方法，其中的故事足够短，以至于 AI 可以轻松看到所有内容。在这些情况下，新方法没有起到作用。这种“提升”只发生在故事如此之长，以至于 AI 必须变得挑剔的时候。如果你有足够的空间，这种额外的检查是不必要的。
它不是“评分”系统的替代品： 论文测试了两种不同的块评分方式（一种称为“K-mean”，另一种称为“Quest”）。新方法对两者都有效。无论你使用哪种评分系统，这种“不确定性检查”都能让你的系统变得更好。
它并非解决所有问题的完美方案： 作者承认，在某些非常具体的、高难度的推理任务（如带有 3 跳逻辑的“变量追踪”）中，即使是最好的模型也表现挣扎，且新方法无法完全解决问题。他们认为这是因为模型本身需要变得更聪明，而不仅仅是选择器的问题。

底线

这篇论文表明，通过在 AI 的决策过程中加入一个简单的“置信度检查”，我们可以防止它仅仅因为分数接近就丢弃重要的线索。

结果显示，这种方法显著提高了 AI 阅读长文本的能力，且不会降低速度。它将一种“盲目的猜测”转变为一种“在关键时刻进行的谨慎复查”。作者发现，这种方法适用于不同类型的 AI 模型和不同长度的文本，证明了有时，想要跑得快，最好的办法就是学会何时聪明地慢下来。

技术摘要：用于块稀疏注意力的不确定性门控选择

问题陈述

块稀疏注意力机制（如 SSA、Quest、H2O、SnapKV）通过将 $O(N^2)$ 的 softmax 替换为针对每个查询（query）的每 $k$ 个键块（key blocks）的 top- $k$ 选择，来扩展长上下文语言模型。虽然这种方式很高效，但其选择过程是短视的（myopic）。当第 $k$ 个和第 $k+1$ 个块的分数非常接近时，选择器会在没有额外预算的情况下做出决策。如果被丢弃的块包含关键证据（例如在多跳推理或“大海捞针”检索中），该信息将不可挽回地丢失，因为下游层无法恢复它。这是一种结构性失效，而非仅仅是噪声，并且在相关性依赖于潜在上下文而非表面查询-键匹配的场景中，性能退化尤为严重。

现有方法将截断视为一种固定的决策规则。它们没有考虑到该决策的不确定性，也没有考虑通过在截断模糊时扩大注意力集合所能获得的信息价值（VoI）。

方法论

本文提出了一种不确定性门控路由器（uncertainty-gated router），它可以叠加在现有的块评分骨干网络（如 SSA 式的均值池化或 Quest 的 K-min/K-max 上界）之上，且无需重新训练或引入额外参数。该方法分为五个步骤：

块评分（与骨干网络无关）： 将键（Keys）分组为块。使用现有的骨干网络（例如均值池化内积或坐标向上的上界）计算分数。
逐分块（Per-Tile）选择： 为了降低计算开销，选择是基于每个查询分块（Query Tile，即连续的若干行查询）而非逐行进行的。分块分数是该分块内所有行中的最大分数。通过 top- $k$ 选择确定初始集合（ $kv\_idx$ ）。
不确定性测量（VoI 信号）： 该方法为每个分块和每个头计算一个归一化截断裕度（ $\sigma$ $σ$ ）：
$\sigma = \frac{s_{(k-1)} - s_{(k)}}{s_{(0)} - s_{(k)}}$
其中 $s_{(i)}$ $s_{(i)}$ 是排序后的块分数。
- $\sigma \to 1$ ：截断是决定性的（保留集远高于尾部）。
- $\sigma \to 0$ ：截断是模糊的（第 $k$ 个和第 $k+1$ 个块得分相等）。
  $\sigma$ 作为 top- $k$ 决策预期损失的代理指标；低 $\sigma$ 值表示获取更多信息的价值较高。
聚合与触发： 每个头的 $\sigma$ 值被平均为每个分块的信号（ $\bar{\sigma}$ ）。如果 $\bar{\sigma}$ 低于特定层的经验分位数阈值（ $q$ ），则激活触发器。这识别出了决策存在歧义的“风险”分块。
受控扩展： 被触发的分块将其注意力集合扩大 $\rho$ 倍（例如 $\rho=2$ ），即关注前 $\rho k$ 个块，而不是前 $k$ 个。未被触发的分块保留基准预算。实现过程使用了一个融合算子（fused kernel），通过对未触发的分块填充哨兵值（sentinels）来保持统一的调度形状，确保成本随扩展呈亚线性增长。

核心贡献

VoI 公式化： 提出了一种新颖的、与骨干网络无关的公式，将选择器截断视为一种不确定性下的决策。它每个分块仅增加一个标量（ $\sigma$ ）和一个每层一个的分位数阈值，无需重新训练或额外参数。
可组合性： 通过实证证明了该路由器可以与不同的评分骨干网络（如 SSA 式的 K-mean 和 Quest 的 K-max）组合。该路由器始终能提升特定模型所使用的最优骨干网络的性能，证明了“更好的评分”与“更好的预算分配”是两个正交的改进维度。
融合实现： 开发了一个自定义的 Triton 算子，能够在一个代码路径中处理选择与扩展，并保持形状统一的调度。这使得系统在 128K 上下文下的运行时间仅为密集注意力的 $0.62\times$ 到 $0.80\times$ 。
诊断性基准测试： 引入了**指针追踪大海捞针（Pointer-Chase Haystack, PCH）**基准测试，以隔离选择器质量与模型能力，并使用 LongBench-v1 作为负对照实验，以证明路由器的收益取决于选择器预算相对于证据密度的紧缩程度。

结果

该方法在四种模型（Qwen2.5-14B, Qwen2.5-7B-1M, Mistral-Nemo, Qwen3.6-35B）上进行了评估，上下文范围从 32K 到 128K。

性能增益：
- 在 LongBench-v2 medium（32K 上下文）上，采用“路由器+Quest”配置的配对召回率（paired recall）为 0.75，而 top- $k$ 基准仅为 0.47（提升了 28 个百分点），McNemar $p < 0.01$ 。
- 在 RULER NIAH 多键（multi-key） 测试中，Qwen2.5-14B 上的“路由器+Quest”配置达到了 0.98 的准确率（对比 top- $k$ 的 0.93），而在 Qwen2.5-7B-1M 上达到了 1.00。
- 路由器始终能提升每个模型中最优骨干网络的表现。值得注意的是，最优骨干网络会根据架构发生变化：Quest 的 K-max 在 Qwen2.5/Mistral 上表现更好，而 SSA 式的 K-mean 在 Qwen3.6 上表现更好（由于 QK-Norm 平衡了幅度方差）。路由器提升了这两者。
效率：
- 在 128K 上下文下，路由器在 Qwen2.5-7B-1M 上保留了 0.81 的密集注意力准确率，在 Qwen3.6 上保留了 0.89，同时运行时间分别为密集注意力的 0.62 $\times$ 和 0.80 $\times$ 。
- 随着上下文变长（64K–128K），由于注意力占总计算量的比例增加，呈现出“帕累托改进”（在相当质量下获得更快的速度）的现象，这符合阿姆达尔定律（Amdahl's law）。
预算 vs. 选择性： 消融研究表明，路由器的提升既来自于增加的预算，也来自于选择性的分配。在接近性能天花板的模型（如 Qwen-14B）上，提升主要由预算驱动；而在具有更多提升空间的模型（如 Mistral-Nemo）上，选择性（仅在歧义分块上投入预算）提供了大部分收益。

意义与主张

论文声称，不确定性门控路由器为块稀释注意力提供了一个通用的、即插即用层，缓解了短视 top- $k$ 选择的结构性失效模式。其意义在于：

鲁棒性： 它能适应不同模型架构的具体评分特性（例如处理 Qwen3.6 中的 QK-Norm 效应），且无需人工调优。
效率： 它在 128K 上下文中实现了接近密集注意力的准确度，同时显著降低了运行时间，使长上下文推理更具实用性。
理论基础： 它通过信息价值（VoI）视角对选择问题进行了建模，并在高斯噪声假设下给出了有限样本的遗憾界限（尽管论文指出，在目前该方法运作的噪声主导区间内，该界限较弱）。

作者明确指出，路由器的提升是有条件的：它仅在每查询预算相对于答案相关证据的位置较小时才会激活。在预算充足的场景下（例如具有原生短提示词的 LongBench-v1），路由器几乎没有收益，这作为负对照验证了该方法对紧缩预算、长上下文场景的特异性。