🤖 machine learning

COBS: Cumulant Order Block Sparse Attention

本文介绍了 COBS，这是一种块稀疏注意力方法，它通过使用一种基于压缩二阶统计量的创新选择器来更好地逼近注意力质量，从而在保持硬件效率的同时，显著缩小了与密集注意力的质量差距，进而提升了长文本检索性能。

原作者： Alexander Tian, Aditya Ghai, Sanjit Neelam, Zaal Vasania, Akshay Mishra

发布于 2026-07-13

📖 1 分钟阅读☕ 轻松阅读

原作者： Alexander Tian, Aditya Ghai, Sanjit Neelam, Zaal Vasania, Akshay Mishra

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

想象一下，你正试图在一大堆草中寻找一根特定的针。但这堆草实在太大了，以至于如果你试图观察每一根稻草，你的大脑（或计算机）就会因为处理不过来而“断电”。这就是现代 AI 模型在尝试阅读超长文档时面临的问题。它们必须记住目前为止读过的所有内容，而检查每一个记忆标记（token）既慢又昂贵。

为了解决这个问题，研究人员尝试了一种叫做**块稀疏注意力（Block Sparse Attention）**的捷径。与其观察每一根稻草，他们决定观察稻草的小捆（称为“块”），并只挑选出那些看起来最有趣的进行仔细检查。这就像雇佣一名侦察兵来扫描一些稻草捆，并告诉你哪些捆可能包含那根针。

问题所在：侦察兵太简单了

这项研究针对的是一种流行的名为 NSA（原生稀疏注意力）的方法。在这个系统中，侦察兵观察一捆稻草，并快速判断它是否重要。研究发现，这个侦察兵使用了一个非常简单的技巧：它只观察了稻草在捆中的平均位置。

你可以这样理解：想象两捆稻草。

A 捆稻草紧密地挤在中间。
B 捆稻草散布得很开，有的在极左，有的在极右。

如果你只看平均位置，这两捆看起来完全一样！但实际上，B 捆更有可能包含那根针，因为它覆盖的范围更广。旧的侦察兵（一阶方法）对这种“分布”或“曲率”是盲目的。它们就像是一个试图仅通过中心点来猜测云朵形状的人；它们忽略了那些真正重要的、蓬松的边缘。

解决方案：COBS（更聪明的侦察兵）

作者提出了一种新方法，称为 COBS（累积量阶数块稀疏注意力）。COBS 的侦察兵不仅仅提供平均位置，还携带了一张微小的、压缩后的地图，不仅显示稻草平均在哪里，还显示它们是如何分布的。

用数学术语来说，论文称之为“二阶统计量”或“协方差”。在我们的类比中，这就像侦察兵意识到：“嘿，这捆稻草很宽且很乱，所以它有更高的概率拥有那根针！”通过保留这部分额外的信息（但对其进行压缩以节省空间），COBS 可以做出更好的判断。

结果：巨大的飞跃

团队在著名的 32k RULER 基准测试（一项针对 11 种不同长文本检索任务的测试）上测试了该方法。以下是他们的发现：

旧方法 (NSA MLP)： 这个简单的侦察兵得分仅为 0.2999。它很难找到那些针。
完美方法 (OSA)： 如果你能通过某种魔法在不使用任何捷径的情况下获得精确答案（称为“预知法/Oracle”），你会得到 0.9040 的分数。
新方法 (COBS)： 这个带有“分布图”的聪明侦察兵得到了 0.8195 的分数。

这意味着 COBS 填补了旧方法与完美方法之间约 86% 的差距。这是一个巨大的进步！

代价：值得吗？

通常，变得更聪明意味着要做更多的功。但 COBS 非常高效。

旧方法读取了一定量的数据。
完美方法（即阅读所有数据的方法）读取的数据量是 COBS 的 15.15 倍。
而 COBS 只比旧的、挣扎中的方法多读取了 1.21 倍 的数据。

因此，COBS 在只需付出极少额外工作量的情况下，就几乎达到了完美的水平。

论文拒绝了什么

作者非常谨慎地排除了一些看似是好捷径但实际上行不通的想法：

仅仅增加平均值的复杂度： 他们尝试使用一种高级神经网络（MLP）来让“平均值”的猜测变得更聪明，但效果并不理想。问题不在于平均值的复杂度，而在于平均值本身就是错误的工具。你需要的是“分布”信息，而不是一个更好的“平均值”。
在简单的“方框”中观察“分布”： 另一种方法尝试通过观察稻草的最小值和最大值（一个方框）来猜测分布。这确实有一点帮助，但不如 COBS 的分布图那样精确。
加入更复杂的数学（三阶）： 作者测试了加入“偏度”（衡量分布不对称程度的指标）。令人惊讶的是，在低复杂度水平下，这反而让情况变得更糟，导致模型产生混乱。它只有在模型已经非常复杂且表现不佳时才起作用，更像是一种“创可贴”而非根本解决方案。他们决定坚持使用“分布”（二阶）作为最佳平衡点。

他们有多确定？

论文对这些数字非常有信心，因为他们进行了受控实验。他们不仅仅是在猜测，而是通过 32k RULER 测试测量了性能，并发现 COBS 始终优于旧方法。他们还检查了这是否会破坏模型理解短句的能力（并没有），并且证明了它确实能比旧方法更好地帮助模型预测长文本中的下一个词。

然而，作者也诚实地说明了局限性：

他们是在一个拥有约 12 亿参数的模型上测试的。他们不能确定这是否在大型科技公司使用的超大规模模型上表现完全一致，尽管数学逻辑表明应该是如此。
他们使用特定类型的合成数据（RULER 风格）来训练模型进行测试。虽然这是测试长文本能力的标准方法，但现实世界的数据行为可能会略有不同。

总结

论文表明，要在草堆中高效地寻找针，你不能只看那一捆的中心。你需要知道这一捆是如何分布的。通过添加一张微小的、压缩后的分布图，COBS 让 AI 模型能够更准确地阅读长文档，而不会降低速度，以极少的额外努力弥合了“足够好”与“完美”之间的差距。

技术摘要：COBS（累积量阶数块稀疏注意力）

1. 问题陈述

大语言模型（LLM）在长上下文长度下的推理受限于内存带宽而非计算量。具体而言，密集注意力（dense attention）所需的键值（KV）缓存读取操作限制了解码速度，并导致硬件利用率不足。虽然块稀疏注意力（block sparse attention）通过仅读取连续的 token 块提供了一种硬件友好的解决方案，但在领先的开源模型中仍未得到充分利用。现有方法通常依赖于密集注意力、细粒度 token 选择或 KV 驱逐（eviction）。

块稀疏注意力的主要挑战在于块选择（block selection）：即如何在不读取每个块的所有键（keys）和值（values）的情况下，识别出哪些 KV 缓存块包含与给定查询（query）最相关的信息。目前最先进的方法，如 DeepSeek 的原生稀疏注意力（Native Sparse Attention, NSA），采用的是“一阶”选择策略。这些方法将块总结为一个与查询无关的向量（例如通过 MLP 或均值池化），并根据该向量对查询进行线性评分。本文指出，这些一阶近似无法捕捉注意力质量（attention mass）的非线性曲率，从而导致与密集注意力之间存在显著的性能差距。

2. 方法论

2.1 理论基础：选择预言机（Selection Oracle）

作者将块选择形式化为从子集块中重建密集注意力输出的问题。他们推导出，最优的选择标准是根据注意力质量（ $m_b$ ）对块进行排序，其中 $m_b$ 定义为块内注意力得分之和：
$m_b = \sum_{r \in b} e^{q^T k_r}$
一个基于精确注意力质量（需要读取所有键）进行选择的“预言机稀疏注意力”（Oracle Sparse Attention, OSA）几乎能达到密集注意力的性能，这证明了选择标准本身是合理的。瓶颈在于如何从紧凑的、与查询无关的摘要中估计这种质量。

2.2 累积量展开（The Cumulant Expansion）

论文分析了使用累积量展开（等同于对对数矩生成函数的泰勒展开）来估计注意力质量的过程。对数质量展开为：
$\ln m_b = \ln L + q^T \kappa_1 + \frac{1}{2} q^T \kappa_2 q + \dots$
其中：

$\kappa_1$ 是块内的平均键（ $\bar{k}_b$ ）。
$\kappa_2$ 是块内的键协方差（ $\Sigma_b$ ）。

现有的一阶选择器（如 NSA 的 MLP 或均值池化）仅捕捉线性项（ $q^T \bar{k}_b$ ），丢弃了二阶曲率项（ $\frac{1}{2} q^T \Sigma_b q$ ）。这种遗漏导致当块内键的分布相对于查询方向具有显著离散度时，它们会失效，即使其均值是相似的。

2.3 COBS 方法

COBS（累izable 量阶数块稀疏注意力） 通过将二阶项引入选择过程来解决这一问题。该方法基于 NSA 架构（包括压缩、选择和滑动窗口分支）并进行了以下创新：

二阶评分： 选择器不仅存储一个均值向量，还存储块内键的均值（ $\bar{k}_b$ ）以及块内协方差矩阵（ $\Sigma_b$ ）的压缩表示。选择得分变为：
$\text{score}_b \approx q^T \bar{k}_b + \frac{1}{2} q^T \Sigma_b q$
协方差压缩： 由于存储完整的 $D \times D$ $D \times D$ 协方差矩阵在内存上是不可接受的，COBS 采用了三种压缩技术：
- 低秩分解： 仅存储 $\Sigma_b$ 的前 $r$ 个特征向量和特征值。
- 子空间投影： 将协方差投影到较低维度的查询子空间（由查询二阶矩的前几个特征向量所张成的空间），将存储从 $rD $降至$ rs$。
- 量化： 以 FP4 (E2M1) 格式存储缩放后的特征向量，同时保持块均值为 BF16。
NoPE（无位置编码）： 压缩和选择分支移除了旋转位置编码（RoPE），以确保块摘要仅取决于内容，防止位置噪声破坏摘要。

3. 核心贡献

选择过程的形式化： 论文将块选择简化为按注意力质量排序，并证明了使用精确质量的预言机可以恢复几乎所有的密集注意力性能。
累积量阶数分析： 它为现有的一阶选择器的局限性提供了理论解释，指出被丢弃的块内协方差是产生性能差距的根源。
COBS 算法： 提出了一种实用的方法，为每个块存储压缩的二阶统计量（协方差），从而实现对注意力质量更准确的近似，而无需读取完整的键。
实验验证： 证明了二阶选择显著缩小了在长上下文基准测试中稀疏注意力与密集注意力之间的差距。

4. 实验结果

4.1 长上下文检索 (32k RULER)

在 11 项任务的 32k RULER 基准测试中进行评估：

NSA 基线 (MLP)： 平均分 0.2999。
NSA 均值池化 (NoPE)： 平均分 0.5554（显示了 NoPE 的益处）。
COBS (子空间, $r=4$ , FP4)： 平均分 0.8195。
密集注意力： 平均分 0.9040。
差距缩小： COBS 弥补了 NSA 基线与密集注意力之间约 86% 的差距。
流量效率： COBS 使用的 KV 缓存读取流量仅为 NSA 基线的 1.21 $\times$ ，且比密集注意力少 15.15 $\times$ 。

4.2 短上下文与语言建模

短上下文推理： COBS 保留了短上下文常识推理性能，在标准基准测试（OpenBookQA, PIQA 等）上达到了与密集注意力持平的水平。
逐位置 NLL： COBS 实现了最低的平均下一 Token 负对数似然（NLL）为 1.633（相比之下，密集注意力为 1.727），表明检索能力的提升并未以牺牲长上下文语言建模能力为代价。

4.3 消融实验

秩搜索 (Rank Sweep)： 性能在 $r=8$ 时达到顶峰，并在更高秩（ $r=16$ ）时因干扰项中无符号方差质量的累积而下降。作者建议 $r \le 8$ 。
子空间 vs 全局： 自适应的每层子空间维度（ $s \approx 85$ ）表现与全局维度 128 相当，同时降低了存储成本。
量化： 特征向量的 FP4 量化被证明基本上是无损的。

5. 重要性与主张

论文声称，累积量阶数既是现有稀疏注意力方法的诊断视角，也是改进的实际设计轴。通过从一阶（均值）转向二阶（协方差）摘要，COBS 证明了块稀疏注意力可以在保持显著减少内存带宽使用的同时，接近密集注意力的质量。

作者强调，这是一个在 $\approx 1.2$ B 参数规模下进行的机制研究，采用了特定的训练协议（RULER 式 SFT）。他们并不声称这是一个适用于所有 LLM 的部署就绪系统，而是证明了块选择中的理论差距是可以弥合的。这项工作表明，如果选择机制足够具有表达力（即具有更高的累积量阶数），那么块稀疏性（通常在领先模型中缺失）是可行的。

作者注明的局限性：

研究受限于特定的骨干网络规模和预训练序列长度。
对比使用了受控的 NSA 基线，其具有特定的超参数（如非重叠块、大型 MLP），这与原始 NSA 报告有所不同。
在稀疏分支中使用 NoPE 而在密集分支中使用 RoPE，在长上下文比较中引入了一个混淆变量。
结果依赖于 R_RULER 式的微调，这可能无法完美泛化到所有的长上下文场景。