July 15, 2026 中 #KV Cache #Attention & MLA #Long Context

COBS 阅读笔记:块稀疏注意力的选择器到底漏掉了什么(一个二阶修正)

笔记日期： 2026-07-15 笔记作者： Zhongzhu Zhou 论文标题： COBS: Cumulant Order Block Sparse Attention 作者： Alexander Tian, Aditya Ghai, Sanjit Neelam, Zaal Vasania, Akshay Mishra（MatX） arXiv： 2607.09052 状态： arXiv 预印本，2026 年 7 月

一句话总结

块稀疏注意力理论上是缓解长上下文 LLM 推理时 KV 缓存读带宽瓶颈最”硬件友好”的方式:与其在每个解码步都读取每一个历史 token 的 key/value 对,不如只读几个连续的”块(block)“。可现实是,几乎没有一个主流开源 LLM 真正采用它——大家用的是稠密注意力、隐空间压缩(MLA)、局部窗口,或者细粒度 token 选择。这篇论文追问了原因,而且答得异常精确。它证明:现有所有的块选择器——也就是在读取一个块的完整 key 之前,决定”这个块值不值得细粒度读取”的那套机制——本质上都可以归约成按一个数字给块排序,这个数字就是块的注意力质量(attention mass)(块会拿到的、未归一化的 softmax 权重之和)。一个按精确质量排序的 oracle,即读取每一个 key 来算出精确的质量,在一个高难度的长上下文检索基准上能把与稠密注意力的差距补齐 99.5%。所以说,选择的标准从来都不是问题。真正的问题在于:每一个已部署的选择器——NSA 的可学习 MLP 池化、DeepSeek-V4 的 CSA 门控、朴素的块均值、甚至 Quest 的 min/max 包围盒——都只能把这个质量估计到关于 query 向量的一阶,原因是它们每个块只缓存一个与 query 无关的向量,评分时只是拿这个向量和 query 做一次内积。论文的核心理论招数是把块质量写成块内 key 的累积量生成函数(cumulant generating function)的指数形式,然后指出:这个函数的泰勒展开里有一个二阶项——一个关于块内 key 协方差的二次型——这个二阶项是任何一阶选择器都注定漏掉的,不管这个一阶摘要本身是怎么算出来的(均值、学习出来的 MLP,还是门控池化都一样)。COBS(Cumulant Order Block Sparse Attention)要做的就是把这个二阶项找回来:它在缓存块均值的同时,额外缓存一份压缩过的、低秩、量化后的块内 key 协方差,用得到的二次修正项来给块打分。在一个 11 任务、32k 上下文的 RULER 基准上,仅这一处改动就把平均分从 0.2999(一个受控的 NSA 基线)提到了 0.8195——把和稠密注意力(0.9040)之间的差距补齐了约 86%——同时读取的 KV 流量只比 NSA 基线多 1.21 倍,比稠密注意力少 15.15 倍。这篇论文对自己结论的边界异常诚实:所有结果都是在约 1.2B 参数规模、非标准的 SFT 协议下验证的,“COBS”这个标签本身也捆绑了好几个互相纠缠的设计选择——这些细节本文会在后面的批判性分析部分逐一梳理清楚。

核心要点

块稀疏注意力的选择问题,在几个显式的简化假设下,可以归约成按块的注意力质量排序——也就是这个块如果被完整读取会拿到的、未归一化的 softmax 权重。一个按精确质量排序的 oracle(需要读取每一个 key)在 32k RULER 上能拿到 0.9010,和稠密注意力的 0.9040 几乎打平(差距补齐 99.5%)——所以按质量排序(经验上)已经非常接近一个”足够好”的标准了。
目前所有主流的可缓存选择器——NSA 的可学习 MLP 池化、DeepSeek-V4 的 CSA 门控池化、朴素的块均值,以及(有一点保留地)Quest 的逐维 min/max 范围——每个块都只缓存一个与 query 无关的向量(外加一个离线标量),并把它当作 query 的仿射函数打分: $\text{score}(q) = a_b + q^\top \phi_b$ 。这在数学上就是对真实 log-质量的一阶近似。
论文最核心的理论工具是块内 key 的累积量生成函数(CGF): $\ln m_b = \ln L + K_X(q)$ ,它在 $q=0$ 处的泰勒展开是 $K_X(q) = q^\top \kappa_1 + \tfrac12 q^\top \kappa_2 q + \tfrac16\sum_{ijk}(\kappa_3)_{ijk}q_iq_jq_k + \cdots$ ,其中 $\kappa_1$ 是块的均值 key, $\kappa_2$ 是块内 key 的协方差。任何仿射打分函数都不可能达到这个二次项 $\tfrac12 q^\top\kappa_2 q$ ——这就是为什么一阶选择器有天花板的数学原因,和这个一阶摘要本身学得多聪明完全无关。
COBS 的修正方式是缓存一份压缩后的协方差:先对 $\Sigma_b$ 做秩- $r$ 谱分解(取前 $r$ 个特征向量,按 $\sqrt{\lambda_i}$ 缩放),再可选地进一步压缩到一个共享的、与 query 相关的子空间(取 query 二阶矩的前 $s$ 个特征向量),最后量化到 FP4。交付配置( $s\!\approx\!85$ 、 $r\!=\!4$ 、FP4)每层每块额外多读 1767 KiB 的摘要,相比 NSA 的 1024 KiB 多读 1.21 倍,仍比稠密注意力每层 65,536 KiB 少读 15.15 倍。
一个干净的线性代数技巧——核 PCA 里熟悉的 Gram 矩阵技巧——能让你在 $O(L^2D + L^3 + rLD)$ 而不是朴素的 $O(D^2L + D^3)$ 时间内算出一个块的前 $r$ 个协方差特征向量,原理是块大小 $L=32$ 远小于 head 维度 $D=128$ :与其对大的 $D\times D$ 协方差 $\tfrac1L\tilde K^\top\tilde K$ 做特征分解,不如对小的 $L\times L$ Gram 矩阵 $\tfrac1L\tilde K\tilde K^\top$ 做分解,再用 $\tilde K^\top w_i/\sqrt{L}$ 把特征向量映射回去。
只在压缩分支和选择分支去掉旋转位置编码(NoPE,sliding-window 分支仍保留 RoPE)是一个独立于其他一切的加性改进:它把均值池化基线在 32k RULER 上的分数从 0.4186 提到了 0.5554,原因是块摘要一旦混入位置带来的旋转,就不再是纯粹的内容签名了。
消融实验对失效模式异常坦诚:存储的秩 $r$ 只在 $r=8$ 之前有帮助(峰值 0.8539),之后在 $r=16$ 时倒退到 0.8006,即便在理论最大秩 $r=31$ 时也无法恢复(0.8135)——原因是协方差的曲率项对符号不敏感(只衡量”散开程度”,不衡量方向),所以多出来的特征向量会在恰好是最需要精细区分的多 key needle 子任务里,给”干扰项很多”的块累积虚假质量。
论文报告了两个负结果,都各自指向了一个真实的设计取舍:围绕一个校准过的非零 query 原点 $q_0$ 展开累积量级数反而变差(0.8238 → 0.8100),原因是被组内多个 query head 共享的偏移矩必须妥协,而且这种展开在偏离 $q_0$ 较远的异常 query 上会失准;加一个廉价的对角三阶累积量(“偏度”)修正在低秩时有害( $r=4$ 时 0.8238 → 0.7754),但能部分修复高秩时的倒退(0.8006 → 0.8252,在 $r=16$ 时)——这佐证了失效模式确实是协方差项对符号不敏感造成的,而不是别的原因。
COBS 被验证为一个只改动选择分支的方法:它原封不动地复用 NSA 的均值池化压缩分支和滑动窗口分支,所以论文里的每一个数字都是在把”稀疏注意力配方”里其他一切都固定不变的情况下,单独隔离出更好的选择器带来的效果。
COBS 在所有被测方法中拿到了最低的长上下文逐位置下一 token 负对数似然(0–32k 平均 1.633 nats/token),比稠密注意力本身还低(1.727)——更关键的是,它的损失曲线在长位置处保持平坦而不是攀升,说明它确实在利用远处的内容,而不是靠退化到局部窗口来”体面地掉分”。
作者明确说明他们的实验设置(约 1.2B 骨干网络、4k 预训练长度、通过 RULER 风格 SFT 训练出来的检索信号、一个”受控”而非字面复现的 NSA 基线)是一项机制研究,而不是部署规模的系统评测——本文在批判性分析部分会认真对待这个前提。

前置知识:读懂这篇论文之前需要知道什么

这篇论文站在三个东西的交汇处:(1)为什么 KV 缓存的读取、而不是模型权重,才是长上下文 LLM 解码的瓶颈,(2)稀疏注意力方法——尤其是 NSA——是怎么试图砍掉这部分读取的,(3)一段大多数机器学习从业者从统计课之后就再没碰过的经典概率论知识:累积量(cumulant)和矩生成函数。这一节把这三块都从头搭建起来,然后再进入论文本身的贡献。

为什么 KV 缓存主导了长上下文推理的成本

一个 decoder-only 的 transformer 是自回归生成文本的:每个解码步,它都根据到目前为止生成的所有内容,产出下一个 token。为了避免每一步都要在整个前缀上重新计算注意力,实现上会缓存每一个历史 token 算出来的 key 和 value 向量——这就是 KV 缓存。在解码步 $t$ ,给当前 query $q_t$ 计算注意力,需要读取所有 $t-1$ 个历史 token 的 key 和 value:

\text{Attention}(q_t, K_{1:t-1}, V_{1:t-1}) = \sum_{r=1}^{t-1} \underbrace{\frac{e^{q_t^\top k_r}}{\sum_{i=1}^{t-1} e^{q_t^\top k_i}}}_{\text{softmax 权重}}\, v_r. \tag{P1}

对单个 query 向量而言,这是一个内存带宽受限的操作:在长上下文场景下( $t$ 达到数万),每一步真正的计算量微不足道(每个历史 token 只需一次点积),但把 $O(t)$ 个 key/value 向量从高带宽内存里搬出来这件事——每个解码步都要重新搬一遍——才是真正主导延迟的部分。这就是为什么长上下文 LLM serving 常被描述为”内存带宽受限而非计算受限”:解码一个 token 仍然要触达整个不断增长的缓存,而每次只能产出一个新 token。当上下文从 4k 长到 32k 再到 128k,每步的读取成本随之线性增长,而真正有用的计算量基本不变——这是一个不断扩大的浪费。

稀疏注意力的四大流派

有四大类方法在攻克这个瓶颈,在深入块选择这个具体话题之前,先把分类摆正是值得的,因为论文把自己的贡献非常精确地定位在其中一类里:

固定模式方法(Longformer、BigBird)预先设计一个子集来关注——局部窗口加上少量”全局” token——与内容无关,提前决定好。
KV 淘汰方法(StreamingLLM、H2O、SnapKV)保持缓存形状不变,但会永久删除被判定为不重要的 token(基于最近使用时间或累积的注意力分数),彻底释放这部分内存。
低秩/隐空间压缩方法(Multi-head Latent Attention / MLA,DeepSeek-V2 及之后使用)把 key 和 value 本身压缩到更小的隐表示,需要时再动态重建。
query 感知的块选择器(Quest、NSA 的选择分支)把完整缓存都留在内存里,但每个解码步动态挑出一小部分连续的 token 块做细粒度读取——不同的 query 可能选中不同的块。

COBS 完全属于第 (4) 类。它不淘汰 token,不把 key 压缩到隐空间,也不用固定模式——它的全部贡献是:对于一个已经具备其他全部机制(NSA)的方法,给出一种更好的方式来决定”这一步该读哪些块”。

Native Sparse Attention(NSA):本文研究的系统对象

论文之所以详细研究 NSA,是因为 NSA 干净地把选择机制隔离成三条并行分支之一,三条分支的输出经过门控后相加:

flowchart TB
    Q["当前解码步的 query q"] --> C["压缩分支:\n对粗粒度池化后的\n块表示做 attention\n(便宜,对所有块始终稠密)"]
    Q --> S["选择分支:\n用缓存的、与 query 无关的\n摘要给每个块打分;\n保留 top-k;只对选中的块\n读取完整 key/value 分辨率"]
    Q --> W["滑动窗口分支:\n对最近约 256 个 token\n做稠密 attention\n(处理严格局部的依赖)"]
    C --> G["可学习门控:\n对三条分支的输出做加权求和"]
    S --> G
    W --> G
    G --> OUT["当前解码步的\n最终 attention 输出"]

图 A(架构总览,自绘):Native Sparse Attention 的三条并行分支。COBS 的全部贡献都在选择分支的打分步骤内部(高亮部分):其余一切——压缩分支、滑动窗口、门控——在本文比较的每一种方法中都保持不变。

选择分支是本文整篇论文的核心。粗略来说:NSA 把过去 $t-1$ 个 token 划分成每块 $L$ 个 token 的连续块(本文用 $L=32$ )。对每个块,它预先计算并缓存一个小型”摘要”——在 NSA 里,是对块的 key 和 value 分别应用一个可学习的 MLP。解码时,它只用这个缓存的摘要给每个块打分(不会重新读取原始 key,除非这个块被选中),挑出得分最高的 top- $k$ 个块(本文用 $k=16$ ),只有对这些被选中的块,它才会真正读取 key/value 并做普通的细粒度注意力。压缩分支是一个独立的、始终开启的粗粒度流(也是块池化,但每个块都会读取,不只是被选中的那些),让模型对选中集合之外的一切都保留一点感知。滑动窗口无条件地以完整分辨率覆盖最近的若干 token(本例中是 256 个),因为非常新近的上下文格外重要,而且保持精确的代价很低。

这里有一个决定整篇论文的结构性事实——可缓存性约束:任何选择器为一个块计算出的摘要,都必须是离线计算的——也就是块刚填满的那一刻算一次,存下来;它不能依赖之后才会出现、拿来对它打分的那个 query,否则你每个解码步都得重新读一遍块的原始 key,可缓存摘要的意义也就彻底没了。本文比较的每一种方法——均值池化、NSA 的可学习 MLP、Quest 的 min/max,以及 COBS 的协方差——都遵守这个约束。论文的论点是:正是遵守这个约束这件事本身,限制住了摘要的累积量阶数,而这正是本文接下来要仔细梳理的故事。

分组查询注意力(GQA)记号,简要说明

现代 LLM 几乎全都使用分组查询注意力(GQA):不再是每个 query head 都有自己独立的 key/value head(多头注意力,MHA),而是一组 $G$ 个 query head 共享一个 key/value(“KV”)head,所以一个总共有 $H$ 个 KV head 的模型,实际上有 $H\times G$ 个 query head。这一点之所以对块选择很重要,是因为选择决策——保留哪些块——是按 KV head 做的,但这一个决策会被共享该 KV head 的所有 $G$ 个 query head 共同使用。正是这种”共享决策”的结构,让下面”一步步推导选择 oracle”这一节里群组级别选择分数的推导变得不平凡:你不能对每个 query head 独立地优化选择,因为在 GQA 下它们本来就得不到独立的选择结果。

矩生成函数与累积量:唯一一块新数学

这是大多数机器学习读者近期都不太会碰到的一块背景知识,但它对这篇论文而言是货真价实的核心,所以值得从零搭起来。对一个随机向量 $X$ (这里指:一个块的 key,当作 $L$ 个近似独立同分布的样本),矩生成函数(MGF) 定义为

M_X(q) \triangleq \mathbb{E}_X\!\left[e^{q^\top X}\right]. \tag{P2}

之所以叫这个名字,是因为它在 $q=0$ 处的各阶导数就是 $X$ 的各阶矩: $\nabla_q M_X(0) = \mathbb{E}[X]$ (一阶矩/均值), $\nabla^2_q M_X(0) = \mathbb{E}[XX^\top]$ (二阶矩),以此类推。累积量生成函数(CGF) 就是 MGF 的对数, $K_X(q) \triangleq \ln M_X(q)$ ,它在 $q=0$ 处的各阶导数给出的是 $X$ 的累积量,而不是原始矩。累积量之所以是这里(以及统计学一般意义下)更自然的对象,原因在于前两阶累积量恰好就是最熟悉的两个概括性统计量——均值和协方差——而更高阶的累积量则分离出了低阶量本身捕捉不到的、真正新增的信息(比如偏度)。具体来说,把 $K_X(q)=\ln\mathbb{E}[e^{q^\top X}]$ 求两次导并在 $q=0$ 处求值(下一节会给出完整推导),恰好得到 $\kappa_1=\mathbb{E}[X]$ 和 $\kappa_2=\text{Cov}(X)$ ——所以当这篇论文说”块的 key 的累积量”时,前两阶指的其实就是”块的均值 key 向量和块内 key 的协方差矩阵”,两个大家非常熟悉的对象,只是套了个不熟悉的名字。论文之所以要用”累积量”这套语言,而不是直接说”均值和协方差”,是因为更高阶的项(三阶及以上)同样有良好定义、可加,并且在泰勒级数里有干净的位置——这正是论文关于”对角偏度”这个负结果(一个三阶累积量修正,后文详述)所利用的结构。

真正的瓶颈:选择分支到底需要计算什么?

在推导任何东西之前,先精确说清楚一个”好的”块选择器到底应该达成什么目标是值得的,因为论文很仔细地把这一点形式化了,而不是诉诸直觉。选择分支要做的是:只用一个很小的、预先算好的、每块一份的摘要,尽量重建出当前 query 下稠密注意力的输出,同时永远不去读取那些没被选中的块的完整 key。这件事在原理上难在两点:(1)真实的注意力输出既取决于哪些 token 的 softmax 权重高,也取决于这些 token 的 value 是什么;(2)缓存的摘要必须只算一次、离线算好,并且要在未来每一个可能会对该块做注意力的 query 上原样复用——而写缓存的那一刻,这些未来的 query 根本还没出现。下一节将逐步推导,在这个约束下,“尽量接近”到底能归约成什么。

一步步推导选择 Oracle

这一节把论文第 3 节的推导重新搭建一遍,每一步代数都显式写出来——论文把其中好几步压缩成了一两行;本文不这样做。

起点:稠密输出是块上的一个混合分布

把过去 $t-1$ 个 token 划分成连续的块 $b$ ,每块大小 $L$ 。在一个 head 内部,定义块的质量 $m_b$ 和value 质心 $v_b^c$ :

m_b = \sum_{r\in b} e^{q^\top k_r}, \qquad v_b^c = \frac{1}{m_b}\sum_{r\in b} e^{q^\top k_r}\, v_r. \tag{1}

质量就是这个块拿到的(未归一化)softmax 权重总和;质心是块内按各 token 自己的 softmax 分数加权平均得到的 value。设 $Z=\sum_i e^{q^\top k_i}$ 为完整的 softmax 分母(对所有块求和), $P_b = m_b/Z$ 是块的归一化概率质量。那么精确的稠密输出就是块质心的一个 $P_b$ 加权混合:

o^\star = \sum_b P_b\, v_b^c. \tag{P3}

这一点值得停下来体会,因为它是让后续推导变得可处理的关键改写:只要你知道每个块精确的概率质量 $P_b$ 和精确的 value 质心 $v_b^c$ ,稠密注意力对所有 token 的输出,就精确等于对块这个更粗粒度对象的一个混合。选择现在变成了一个”这个混合里保留哪些块的贡献”的问题,而完全不再是逐 token 的问题。

丢弃块带来的精确重建误差

现在假设一个选择器挑出了要保留的块子集 $S$ (丢弃其补集 $S^c$ ,丢弃的总概率质量记为 $\tau = \sum_{b\in S^c} P_b$ )。由于剩下的概率加起来不再是 1,选择器必须在保留集合上重新归一化:

\hat o_S = \frac{\sum_{b\in S} P_b\, v_b^c}{1-\tau}. \tag{2}

要求出精确误差 $o^\star - \hat o_S$ ,把式(P3)里的完整求和拆成保留和丢弃两部分, $o^\star = \sum_{b\in S}P_bv_b^c + \sum_{b\in S^c}P_bv_b^c$ ,再减去式(2):

o^\star - \hat o_S = \sum_{b\in S}P_bv_b^c\left[1 - \frac{1}{1-\tau}\right] + \sum_{b\in S^c}P_bv_b^c = -\frac{\tau}{1-\tau}\sum_{b\in S}P_bv_b^c + \sum_{b\in S^c}P_bv_b^c.

这看起来有点乱,但只要代入 $o^\star=\sum_{b\in S}P_bv_b^c+\sum_{b\in S^c}P_bv_b^c$ 替换掉其中一个求和(论文的 Eq. 3 悄悄跳过了这一步替换——这里显式写出来),就会漂亮地化简:把丢弃块的求和重写成 $\sum_{b\in S^c}P_b(v_b^c - o^\star) + \tau\,o^\star$ (加减 $\tau o^\star$ ),整理之后,所有涉及保留集合 $S$ 的项都会相互抵消,留下这个干净的闭式解:

o^\star - \hat o_S = \frac{1}{1-\tau}\sum_{b\in S^c} P_b\left(v_b^c - o^\star\right). \tag{3}

如何理解这个式子。 重建误差恰好等于:对被丢弃的块求和(经过重新归一化),每一项是该块的概率质量乘以它的质心 $v_b^c$ 偏离真实稠密输出 $o^\star$ 的程度。这带来一个立刻可用的直觉推论:一个块贡献误差的大小,取决于(a)它携带了多少被丢弃的概率质量 $P_b$ ,以及(b)它的 value 质心离”平均” value 有多远。一个即使被丢弃、但其质心恰好等于 $o^\star$ 的块,丢了也不会带来任何误差——这干净地形式化了”这个块反正也没什么独特信息”这句直觉。

从单头误差到 GQA 下每个 KV head 的目标函数

在 GQA 下,一个 KV head $h$ 被 $G$ 个 query head 共享,而且——这一点很关键——被保留下来的块集合 $S^{(h)}$ 是这 $G$ 个 query head 共享的(选择是 KV head 缓存的属性,不是某一个 query head 单独的属性)。所以对一个固定的 KV head 而言,自然的最小化目标,是它所在的 query head 组里每个头重建误差的总和:

E^{(h)}(S) = \sum_{g=1}^{G} \left(o^\star - \hat o_S\right)^{(g,h)}. \tag{4}

这个目标函数虽然精确,但按原样是没法用的,因为算它需要知道 $o^\star$ ——而这恰恰是你想通过读取每一个 key 来避免计算的东西。接下来的三条假设,把它变成一个可缓存摘要真的能估计出来的东西。

让目标函数变得可处理的三条假设

假设 1(值无关性)。 用一个每头共用的常数来限定质心偏差:对每个块 $b$ , $\|v_b^c - o^\star\|^{(g,h)} \le c^{(g,h)}$ 。对式(3)用三角不等式并代入这个界,可以得到每个 query head 重建误差的上界:

\left\|o^\star - \hat o_S\right\|^{(g,h)} \le \frac{1}{1-\tau^{(g,h)}}\sum_{b\in S^c} P_b^{(g,h)}\, c^{(g,h)} = c^{(g,h)}\,\frac{\tau^{(g,h)}}{1-\tau^{(g,h)}}. \tag{5}

为什么需要这条假设? 真正的界需要知道 $o^\star$ 本身(循环论证——这恰恰是选择本来想避免计算的东西)。用一个每头统一的最坏情况常数来限定偏差绕开了这个循环,代价是丢掉了”到底哪些块的偏差异常大或异常小”这条信息。

假设 2(忽略 $c$ )。 假设 1 里的每头常数 $c^{(g,h)}$ 本身也是未知的(要知道它们同样需要知道 $o^\star$ ),所以在给块排序的目的上,它们被直接从目标函数里丢掉——反正它们不依赖块的编号 $b$ ,所以不会影响一个头内部哪些块看起来更好。

假设 3(线性松弛)。 把式(5)的界在 query head 组内对 $G$ 个头求和,得到一个每 KV head 的目标函数 $\sum_g \tau^{(g)}/(1-\tau^{(g)})$ (丢掉现在已经无关的 $c^{(g,h)}$ 常数,固定 $h$ 只用 $g$ 重新索引)。这个关于 $\tau^{(g)}$ 的函数不能简单地按块可加地分解,因为 $1/(1-\tau^{(g)})$ 依赖的是整个补集总的丢弃质量,而不是任何单个块——所以严格来说,写成这个样子的目标函数并不存在精确最优的固定逐块打分。围绕 $\tau^{(g)}=0$ 对这个惩罚项做泰勒展开,

\frac{\tau^{(g)}}{1-\tau^{(g)}} = \tau^{(g)} + O\!\left((\tau^{(g)})^2\right),

只保留线性主导项,把目标函数换成 $\sum_g \tau^{(g)}$ ——现在这是一个对丢弃的块求和、每头概率质量相加的式子,而这确实是能按块可加的。为什么需要这条假设,又在哪里会失效? 这个松弛在每个头丢弃的质量 $\tau^{(g)}$ 较小时是准确的——也就是说,当 top- $k$ 选择预算本身已经覆盖了大部分真实概率质量时。当 $\tau^{(g)}$ 变大时(选择被大幅截断,或者一个校准很差的选择器把高质量的块丢掉了),这个近似会退化——论文没有测试这个误差在实践中会变大到什么程度,本文在批判性分析部分会指出这一点。

线性松弛到底准不准?一个具体的数值检验

假设 3 用线性项 $\tau^{(g)}$ 替换了精确的惩罚项 $\tau^{(g)}/(1-\tau^{(g)})$ ,上面的推导提到这个近似在 $\tau^{(g)}$ (每头丢弃的概率质量)很小时是准确的——但“很小”到底有多小,既不是论文也不是上面的推导给出了具体数字。直接算一下就能弄清楚:在 $\tau=0.05$ (选择已经捕获了一个头 95% 的真实质量)时,精确惩罚是 $0.0526$ ,线性近似是 $0.0500$ ——相对误差 $5.0\%$ 。在 $\tau=0.1$ (捕获 90%)时,相对误差翻到 $10.0\%$ 。在 $\tau=0.2$ (捕获 80%,也就是 top- $k$ 预算相对真实质量分布而言明显收紧的情况)时,相对误差已经达到 $20.0\%$ ,到 $\tau=0.5$ 时增长到 $50\%$ 。这直接、具体地验证了前面早已标注过的一点:论文从未报告它自己的 top- $k=16$ 、 $L=32$ 块配置,在实际训练好的模型上究竟捕获了多大比例的真实概率质量(也就是实际的 $\tau$ 到底多大),所以读者单凭这篇论文无从得知:部署配置究竟是安稳地落在了误差低于 10% 的区间里,还是更接近 20%–50% 误差的区间——后者意味着假设 3 的线性化本身,就会在 oracle 自身的最优性保证里引入一定的实质性误差——而且这个误差与 COBS 的协方差项要处理的估计误差是独立的、另外一层的。

由此得到的 GQA 选择分数

在这三条假设下,最小化每 KV head 目标函数,归约为在 top- $k$ 个保留块的选择上最大化保留概率质量之和——等价于最小化丢弃质量 $\sum_g \tau^{(g)} = \sum_g\sum_{b\in S^c}P_b^{(g)}$ 。因为现在这在块上是可加的,对一个固定预算 $k$ 而言,最优解就是保留群组打分之和最大的 $k$ 个块:

\text{score}_b^{(h)} = \sum_{g=1}^{G} \frac{m_b^{(g,h)}}{Z^{(g,h)}}, \qquad Z^{(g,h)} = \sum_{b'} m_{b'}^{(g,h)}. \tag{6}

MHA 特例。 当 $G=1$ (普通多头注意力,没有 query head 共享)时,推导会坍缩成一个精确(而非近似)的结论:目标函数关于单个 $\tau$ 是单调递增的,并且 $\tau$ 在丢弃的块之间是可加的,与任何泰勒近似都无关,所以假设 2 和假设 3 根本不需要——单靠假设 1,按原始质量 $m_b$ (等价地,按 $\ln m_b$ ,因为 $\ln$ 单调)保留 top- $k$ 个块,就是精确最优的。

下面是完整的 oracle 打分流程的伪代码,清楚区分了离线(块一填满就算一次)与在线(每个解码步)计算:

算法 1:Oracle Sparse Attention(OSA)块打分
─────────────────────────────────────────────────────────
离线部分(块 b 填满、长度为 L 时算一次):
  1. 存下这个块的原始 key k_r 和 value v_r,r 属于 b。
     (OSA 是一个诊断用方法,不可部署:打分时必须重新
      读取每一个原始 key,所以这里根本没有真正的"压缩"。)

在线部分(每个解码步,对 query q,按 KV head h):
  2. 对每个块 b:
       m_b^(g,h)  <-  sum_{r in b} exp( q^(g,h)^T k_r )      # 式(1),精确
  3. 对每个块 b:
       Z^(g,h)    <-  sum_{b'} m_b'^(g,h)
       score_b^(h) <- sum_{g=1}^{G} m_b^(g,h) / Z^(g,h)        # 式(6),精确
  4. 按 score_b^(h) 保留 top-k 个块;读取它们完整的 K、V。
  5. 对以下三部分的并集做普通的细粒度注意力:
       保留的块(完整分辨率)
       + 滑动窗口(完整分辨率)
       + 压缩分支的粗粒度流(始终开启)
  6. 门控并相加三条分支的输出(图 A)。
─────────────────────────────────────────────────────────

为什么 OSA 尽管不可部署却依然重要。 OSA 每个解码步都要重新读取每一个原始 key 才能算出打分——它在 key 读取流量上完全没有节省,只在 value 读取和细粒度注意力流量上有节省(后文的表 3 显示,OSA 整体上仍比稠密注意力少读 1.80 倍,这纯粹来自 value/计算上的节省,但比 NSA MLP 多读 10.14 倍)。它在本文里的全部作用是诊断性的:用精确质量而非任何缓存近似值来打分,OSA 把”质量排序这个标准本身,在最理想情况下有多好?”这个问题,和”我们能从一个可缓存摘要里把这个标准估计得多好?”这个问题隔离开了。论文的头条数字——OSA 拿到 0.9010,对比稠密的 0.9040,差距补齐 99.5%——斩钉截铁地回答了第一个问题:是的,按精确质量排序(经验上)基本等同于完整的稠密注意力。这就把每一个已部署方法里剩下的全部差距都归结为一个估计问题,而不是标准问题——这正是下一节累积量展开要处理的框架。

flowchart LR
    subgraph OFFLINE["离线,每个块(填满时算一次)"]
        K["块的 L 个 key、value"] --> SUM["计算可缓存摘要\n(均值池化 / NSA-MLP /\nQuest min-max / COBS 协方差)"]
    end
    SUM --> CACHE["存下摘要\n(与 query 无关)"]
    subgraph ONLINE["在线,每个解码步"]
        Q2["当前 query q"] --> SCORE["对每个块的缓存摘要\n与 q 打分"]
        CACHE --> SCORE
        SCORE --> TOPK["保留得分最高的 top-k 个块"]
        TOPK --> READ["只对保留的块\n读取完整 key/value"]
        READ --> ATT["对保留块做\n细粒度注意力"]
    end

图 B(数据流/流水线示意图,自绘):本文比较的每一种方法——均值池化、NSA-MLP、CSA、Quest、COBS——都遵循的这套通用可缓存选择器流水线。这些方法唯一的区别,就在”计算可缓存摘要”这一个方框里;它下游的一切(打分、top-k、读取、细粒度注意力)在所有方法里都是完全一样的架构。

累积量展开:一阶选择器为什么有天花板

这一节把论文第 4 节——整篇论文的理论核心——重新搭建一遍,把矩生成函数的推导显式做出来,而不是直接引用结论。

把块质量改写成累积量生成函数的指数形式

再次从块质量的定义(式 1)出发,把它改写成 $L$ 乘以块自身经验 key 分布上的一个平均值(把 $X$ 当作在块的 $L$ 个 key $\{k_r\}_{r\in b}$ 上均匀分布的随机变量):

m_b = \sum_{r\in b} e^{q^\top k_r} = L\cdot\frac{1}{L}\sum_{r\in b}e^{q^\top k_r} = L\cdot \mathbb{E}_X\!\left[e^{q^\top X}\right] = L\cdot M_X(q). \tag{7}

取对数,并定义 $K_X(q)\triangleq \ln M_X(q)$ 为这个块的累积量生成函数:

\ln m_b = \ln L + K_X(q). \tag{8}

这已经是一次相当惊人的改写:一个块的 log-质量,作为 query 的函数,完全由它自己 key 分布的累积量生成函数决定——关于这个块,别的什么都不重要了。

显式推导为什么前两阶累积量就是均值和协方差

论文把这一点当作一个标准结论引用(McCullagh 那本统计张量方法的教材);本文把它显式推导出来,因为正是这个事实让”累积量”这个词从抽象变得具体可感。根据定义, $K_X(q) = \ln \mathbb{E}_X[e^{q^\top X}]$ 。对 $q$ 求一次导并在 $q=0$ 处求值:

\nabla_q K_X(q)\Big|_{q=0} = \left.\frac{\mathbb{E}_X\!\left[X\, e^{q^\top X}\right]}{\mathbb{E}_X\!\left[e^{q^\top X}\right]}\right|_{q=0} = \frac{\mathbb{E}_X[X]}{1} = \mathbb{E}_X[X] \triangleq \kappa_1. \tag{9}

一阶累积量就是均值——这还不算意外。再求一次导需要用到商法则(因为一阶导数本身就是关于 $q$ 的两个函数之比):

\nabla^2_q K_X(q) = \frac{\nabla^2_q M_X(q)\cdot M_X(q) - \nabla_q M_X(q)\,\nabla_q M_X(q)^\top}{M_X(q)^2}.

在 $q=0$ 处求值,此时 $M_X(0)=\mathbb{E}[e^0]=1$ , $\nabla_qM_X(0)=\mathbb{E}[X]$ , $\nabla_q^2M_X(0)=\mathbb{E}[XX^\top]$ :

\nabla^2_q K_X(q)\Big|_{q=0} = \mathbb{E}_X\!\left[XX^\top\right] - \mathbb{E}_X[X]\,\mathbb{E}_X[X]^\top = \text{Cov}_X(X) \triangleq \kappa_2. \tag{10}

这是关键事实。 二阶累积量恰好就是协方差矩阵。由于 $X$ 是均匀分布在块自身的 $L$ 个 key 上,这些通用公式可以具体化成可直接计算的每块量:

\kappa_1 = \bar k_b = \frac{1}{L}\sum_{r\in b}k_r, \qquad \kappa_2 = \Sigma_b = \frac{1}{L}\sum_{r\in b}(k_r-\bar k_b)(k_r-\bar k_b)^\top. \tag{11}

块均值是块内 key 向量的平均值;块协方差衡量的是这些 key 围绕均值散布的程度,以及散布的方向。这两者都是你在把一个块写入缓存的那一刻就已经能算出来的东西——完全不需要看到 query。

累积量展开,以及它为什么暴露出一个天花板

围绕 $q=0$ 对 $K_X(q)$ 做泰勒展开(标准的多元累积量展开)得到

K_X(q) = q^\top\kappa_1 + \frac12 q^\top\kappa_2 q + \frac16\sum_{i,j,k}(\kappa_3)_{ijk}\,q_iq_jq_k + \cdots, \tag{12}

结合式(8),就有

\ln m_b = \ln L + q^\top\bar k_b + \frac12 q^\top\Sigma_b q + \frac16\sum_{ijk}(\kappa_3)_{ijk}q_iq_jq_k + \cdots. \tag{13}

现在前面关于”仿射打分”的说法可以变得完全精确了。根据可缓存性约束,任何可缓存的选择器都必须用一个从固定的、与 query 无关的摘要计算出来的函数来打分——每块一个向量 $\phi_b$ 和一个标量 $a_b$ ,打分方式是 $\text{score}_b^{\text{aff}}(q) = a_b + q^\top\phi_b$ 。这个函数形式关于 $q$ 是仿射的——线性加常数——不管 $\phi_b$ 本身是怎么算出来的(均值池化、一个学习出来的 MLP,还是一个门控组合)都一样。要精确匹配式(13)的常数项和一次项,需要 $a_b=\ln L$ 、 $\phi_b=\bar k_b$ ——但没有任何 $a_b,\phi_b$ 的选择能够复现二次项 $\tfrac12 q^\top\Sigma_bq$ ,因为一个关于 $q$ 的二次型,根据定义就不是 $q$ 的仿射函数。这正是论文核心论点的数学内容,而且这个论证是滴水不漏的:一阶选择器的天花板,不是任何具体设计(NSA 的 MLP、CSA 的门控)的局限——它是一个关于仿射函数能表达什么、不能表达什么的结构性事实。

一个完整的数值例子:二阶如何实实在在胜过一阶

上面的代数推导,一旦亲眼看它在可以手算的数字上跑一遍,就会变得更值得信任。下面是一个完整、独立验证过的玩具例子(不来自论文本身——专门为本文构造,下面每一个数字都经过了数值交叉验证), $D=2$ 、 $L=4$ 。设一个块的四个 key 为

k_1=(1,3),\quad k_2=(1,-1),\quad k_3=(-1,1),\quad k_4=(-1,1). \tag{P4}

第一步——均值。 根据式(11), $\bar k_b = \tfrac14\sum_r k_r = \tfrac14\big((1,3)+(1,-1)+(-1,1)+(-1,1)\big) = \tfrac14(0,4) = (0,1)$ 。

第二步——协方差。 先中心化每个 key: $k_1-\bar k_b=(1,2)$ , $k_2-\bar k_b=(1,-2)$ , $k_3-\bar k_b=(-1,0)$ , $k_4-\bar k_b=(-1,0)$ 。将外积求和再除以 $L=4$ :

\Sigma_b = \frac14\left[\begin{pmatrix}1&2\\2&4\end{pmatrix}+\begin{pmatrix}1&-2\\-2&4\end{pmatrix}+\begin{pmatrix}1&0\\0&0\end{pmatrix}+\begin{pmatrix}1&0\\0&0\end{pmatrix}\right] = \frac14\begin{pmatrix}4&0\\0&8\end{pmatrix} = \begin{pmatrix}1&0\\0&2\end{pmatrix}. \tag{P5}

这个块的 key 在第二个坐标方向上的散布程度是第一个坐标方向的两倍——一个各向异性(anisotropic)的块,正是论文图 2 论证中那种仅靠均值无法与一个均值相同、但各向同性(isotropic)的块区分开的那种块。

第三步——三个 query,三组对比。 取 $q=(0,1)$ (与高方差方向对齐)。精确质量(式 1)需要真实的点积: $q\cdot k_1=3$ , $q\cdot k_2=-1$ , $q\cdot k_3=1$ , $q\cdot k_4=1$ ,所以 $m_b = e^3+e^{-1}+e^1+e^1 = 20.0855+0.3679+2.7183+2.7183=25.8900$ ,得到 $\ln m_b = 3.2539$ 。现在对比两种估计:**一阶(仿射)**估计是 $\ln L + q^\top\bar k_b = \ln4 + (0\cdot0+1\cdot1) = 1.3863+1=2.3863$ ——误差 $0.8676$ ,相当大。**二阶(COBS)**估计加上曲率项 $\tfrac12q^\top\Sigma_bq = \tfrac12(0^2\cdot1+1^2\cdot2)=1$ ,得到 $1.3863+1+1=3.3863$ ——误差只有 $0.1324$ ,仅多一项就把误差缩小了 $6.6$ 倍。

在更小的 query 幅度下, $q=(0,0.5)$ :精确 $m_b=e^{1.5}+e^{-0.5}+2e^{0.5}=8.3857$ , $\ln m_b=2.1265$ 。一阶: $1.3863+0.5=1.8863$ (误差 $0.2402$ )。二阶: $1.3863+0.5+\tfrac12(0.5^2\cdot2)=1.3863+0.5+0.25=2.1363$ (误差 $0.0098$ )——误差缩小了 $24$ 倍,因为 query 幅度越小,截断的累积量级数就越接近展开点,缺失的三阶项影响就越小。

现在看对比的情况: $q=(1,0)$ (与低方差方向对齐)。精确: $q\cdot k_1=1,q\cdot k_2=1,q\cdot k_3=-1,q\cdot k_4=-1$ ,所以 $m_b=2e^1+2e^{-1}=6.1723$ , $\ln m_b=1.8201$ 。一阶: $\ln4 + q^\top\bar k_b = 1.3863 + (1\cdot0+0\cdot1) = 1.3863$ (误差 $0.4338$ ——注意仿射打分甚至看不到 query 在动,因为 $\bar k_b$ 的第一个分量是 $0$ )。二阶: $1.3863 + 0 + \tfrac12(1^2\cdot1+0^2\cdot2) = 1.3863+0.5=1.8863$ (误差 $0.0662$ )。

这个玩具例子具体说明了什么。 在三组(query 方向、幅度)组合上,二阶估计的误差都比一阶估计小 $5$ –25 倍——而且这还只是一个可以手算的 $D=2$ 玩具例子,不是论文那个 $D=128$ 的 32k RULER 实验。而且关键在于,同一个块(同样的 $\bar k_b$ 、同样的 $\Sigma_b$ )在不同方向的 query 下,一阶与精确值之间的差距完全不同——这正是图 2 定性描述的现象,现在用不到五分钟就能手工重算的数字固定下来了。

论文的图 2 恰好可视化了这个空缺:

图 1(论文 Fig. 2):一阶打分漏掉了块内曲率

图 1(论文 Fig. 2,嵌入):左图展示两个块,共享完全相同的均值 key $\bar k_b$ (黑色星号),但它们的 key 相对于 query 方向 $q$ 的散布方式完全不同——蓝色块的 key 沿 $q$ 方向散布(提升了它真实的质量),橙色块的 key 沿垂直于 $q$ 的方向散布(它的质量因此偏低),但对任何只缓存均值的选择器来说,这两个块是无法区分的。右图画出了沿固定方向、log-质量随 query 幅度变化的曲线:真实的 log-质量(黑色)明显是弯曲的,一个仿射的一阶打分(橙色虚线)只能是穿过它的一条直线,而经过协方差修正的二阶估计(蓝色)几乎完全贴合真实曲线——在任意给定的 query 幅度处,橙色直线和黑色曲线之间的竖直间隙,正是 COBS 的协方差项想要补齐的”曲率缺口”。

方法:COBS,拆解开来讲

这一节把论文第 5 节(方法)的每一个子节都拆开来讲,每一个子节都是叠加在同一个受控 NSA 基线之上的一项加性设计改动。

压缩分支和选择分支里的 NoPE

旋转位置编码(RoPE)会在做点积之前,根据位置把每个 key 旋转一个角度。这对于选中块上的细粒度注意力是必不可少的(那里逐 token 的相对位置真的很重要),但对于最初用来选择块的那个摘要而言,却是实实在在有害的:把块内不同位置、旋转了不同角度的 key 池化到一起,会把位置带来的旋转,混进本该是纯粹基于内容的块签名里。只在压缩分支和选择分支里去掉 RoPE(滑动窗口仍保留 RoPE,那里精确的相对位置需要被保留)——论文把这套方案称为 NoPE——被报告为一个干净的、独立于其他一切的改进:在还没加入任何协方差项之前,它就已经把均值池化基线在 32k RULER 上的分数从 0.4186(带 RoPE)提高到了 0.5554(NoPE)。

为什么是这个设计选择,替代方案是什么? 显而易见的替代方案——为了架构一致性到处都保留 RoPE——会把一个本该纯粹基于内容的选择信号,注入位置噪声后主动搞坏。边界条件是:这个 NoPE 收益只在一个长上下文检索场景里得到了验证,在这种场景下选择器需要区分的是内容而非精细的位置结构;论文自己(在后面的局限性部分)也承认,这项消融从未在那一个长上下文配置之外测试过。

二阶截断:核心估计公式

把累积量展开(式 13)截断到二阶、丢弃更高阶的项,就得到 COBS 的核心打分公式:

\ln m_b \approx \ln L + q^\top \bar k_b + \frac12 q^\top \Sigma_b q. \tag{14}

把它拆成大家熟悉的一阶部分 $\hat\ell_b = q^\top\bar k_b$ (均值池化本来算的就是这个)加上新的曲率项 $\tfrac12 q^\top\Sigma_bq$ ——均值池化、NSA 的 MLP,以及其他任何仿射选择器都会丢掉这一项。在 GQA 下,把这个截断后的估计代入精确打分公式(式 6)——用二阶估计替换真实的 $m_b^{(g,h)}$ 和 $Z^{(g,h)}$ ,并在分子分母求和里消去共有的因子 $L$ ——就得到 COBS 实际部署的打分规则:

\widehat{\text{score}}_b^{(h)} = \sum_{g=1}^{G}\frac{1}{\hat Z^{(g,h)}}\exp\!\left(q^{(g,h)\top}\bar k_b^{(h)} + \frac12 q^{(g,h)\top}\Sigma_b^{(h)} q^{(g,h)}\right), \quad \hat Z^{(g,h)}=\sum_{b'}\exp\!\left(q^{(g,h)\top}\bar k_{b'}^{(h)} + \frac12 q^{(g,h)\top}\Sigma_{b'}^{(h)} q^{(g,h)}\right). \tag{15}

值得注意的是,这个二阶估计只用于选择打分——COBS 保留了 NSA 原本未经修改的均值池化压缩分支,所以选择和压缩确实是从同一批底层 key 出发、各自独立设计出来的两份不同摘要。

协方差压缩:低秩谱分解

块协方差 $\Sigma_b$ 是一个 $D\times D$ 的矩阵(本文设置中 $D=128$ 是每个 head 的 key 维度)——精确存储它需要 $O(D^2)$ 个浮点数,当 $D>L$ 时,这实际上比它本该概括的 $L\cdot D$ 个原始 key 还要占更多内存,与压缩的初衷背道而驰。COBS 转而只保留 $\Sigma_b$ 的前 $r$ 个特征方向:

\Sigma_b \approx \sum_{i=1}^{r}\lambda_i u_iu_i^\top = \sum_{i=1}^r \xi_i\xi_i^\top, \qquad \xi_i \triangleq \sqrt{\lambda_i}\,u_i, \tag{16}

把特征值的平方根折叠进特征向量里,这样一个”缩放后的特征向量” $\xi_i$ 就同时携带了这两部分信息。存储块均值( $D$ 个浮点数)加上 $r$ 个缩放特征向量( $rD$ 个浮点数),每块共需 $D + rD$ 个浮点数——取 $r=4$ 、 $D=128$ ,是 $128+512=640$ 个浮点数,对比精确协方差的 $128^2=16{,}384$ 个:仅靠低秩这一步就有约 25 倍的压缩,还没算后续的进一步压缩。

子空间方法:沿 query 维度而非 key 维度压缩

上面的秩- $r$ 近似仍然是在完整的 $D$ 维 key 空间里存储协方差方向,但注意到打分其实只需要标量二次型 $q^\top\Sigma_bq$ ——从来不需要 $\Sigma_b$ 本身这个矩阵。如果整个模型里大多数 query,其实只在一个维度更低的子空间内变化,你就可以把协方差投影到那个子空间里,丢掉其余部分。设 $U_Q\in\mathbb{R}^{D\times s}$ 为聚合 query 二阶矩 $\mathbb{E}[qq^\top]$ 的前 $s$ 个特征向量(离线算一次,基于一个校准集——这是query 分布本身的性质,与任何具体的块无关),张成一个 $s$ 维的”query 子空间”。把协方差投影到这个子空间:

B_b = U_Q^\top \Sigma_b U_Q \in \mathbb{R}^{s\times s}. \tag{17}

设 $\Pi = U_QU_Q^\top$ 为相应的投影矩阵(投影到 $U_Q$ 的列空间上), $\tilde q = U_Q^\top q\in\mathbb{R}^s$ 为投影后的 query。那么

q^\top\Sigma_bq \approx q^\top\Pi\Sigma_b\Pi q = (U_Q^\top q)^\top(U_Q^\top\Sigma_bU_Q)(U_Q^\top q) = \tilde q^\top B_b\tilde q. \tag{18}

推导这个近似在什么条件下精确成立(论文陈述了结论但没有展开推导的一步)。 把 $\Pi=U_QU_Q^\top$ 直接代入 $q^\top\Pi\Sigma_b\Pi q$ :展开得到 $q^\top U_QU_Q^\top\Sigma_bU_QU_Q^\top q = (U_Q^\top q)^\top(U_Q^\top\Sigma_bU_Q)(U_Q^\top q) = \tilde q^\top B_b\tilde q$ ,这只是把定义直接代入化简——确认了式(18)在代数上成立。近似 $q^\top\Sigma_bq\approx q^\top\Pi\Sigma_b\Pi q$ 在 $\Pi q = q$ 时是精确成立的等式——也就是说,只要 query $q$ 本身完全落在 $U_Q$ 列空间张成的子空间内就行,因为 $\Pi$ 按构造就是到那个子空间的正交投影算子,而投影算子作用在一个本来就在其值域内的向量上会原样返回该向量。实践中 $q$ 总会有一小部分残余分量落在子空间之外;论文选择 $s$ 的方式(每层取捕获 90% query 谱能量的秩,再乘以 1.25 倍)明确是为了让这个残余 $\|(I-\Pi)q\|$ 小到可以忽略。

与其存储完整的 $s\times s$ 矩阵 $B_b$ ,不如在这个更小的子空间内再套用一次同样的秩- $r$ 谱分解技巧,存储 $r$ 个 $\mathbb{R}^s$ (子空间)而非 $\mathbb{R}^D$ (全空间)的缩放特征向量 $\xi_i^{ss}$ ——把每块特征向量的存储从 $rD$ 降到 $rs$ 个浮点数。取 $r=4$ 、交付配置的 $s\approx85$ (相对于完整的 $D=128$ ),这是 340 个浮点数而非 512 个——在前面约 25 倍的基础上又有约 1.5 倍的进一步压缩。

这里的关键发现(图 6,后文实验部分讨论)是:不同层”合适的”子空间维度差异相当大——按层自适应地分配 $s$ (依据每一层自身 query 谱的分散程度),比对所有层用同一个全局 $s$ ,能以更小的平均预算恢复几乎同等的精度。

量化:对存储的特征向量做 FP4

存储的缩放特征向量 $\xi_i$ 被量化到 E2M1 FP4 格式(2 位指数、1 位尾数——一种极其粗糙的 4-bit 浮点数),每个特征向量保留一个 fp32 缩放因子,块均值本身保持 bf16 精度。这一步在目标基准上被报告为几乎无损:秩为 4 的描述符只变化了 $+0.0013$ (从 bf16 的 0.8238 到 FP4 的 0.8251),同时把协方差因子的字节数压缩了 3.8 倍(1024 → 272 字节每块,不含单独存储的 bf16 均值)。

成本核算与廉价计算特征向量的 Gram 矩阵技巧

每解码步打分成本。 一旦每个解码步只算一次的投影 query $\tilde q = U_Q^\top q$ (所有块共享,成本 $O(sD)$ )算好了,用存储的子空间特征向量给每个块打分的成本是

q^\top\Sigma_bq \approx \sum_{i=1}^{r}\left(\xi_i^\top\tilde q\right)^2, \tag{19}

也就是在 $s$ 维子空间里做 $r$ 次内积,每块成本 $O(rs)$ ——相比之下均值池化打分的成本是 $O(D)$ 。由于交付配置里 $rs \ll D$ (取 $r=4$ 、 $s\approx85$ 给出 $rs=340$ ,和 $D=128$ 在绝对数值上相当,但得益于量化,分布在少得多的存储字节上),这个成本仍然很低。

用 Gram 技巧计算特征向量。 缩放后的特征向量每个块只需重新计算一次,即块填满时算一次——大约每 $L$ 个解码步才需要算一次,而不是每步都算。设 $\tilde K\in\mathbb{R}^{L\times D}$ 是这个块中心化后的 key 矩阵(每一行是 $(k_r-\bar k_b)^\top$ ),于是 $\Sigma_b = \tfrac1L\tilde K^\top\tilde K$ 。朴素地对 $D\times D$ 的 $\Sigma_b$ 做特征分解,光是构造协方差矩阵就要 $O(D^2L)$ ,再加上分解本身的 $O(D^3)$ ——当 $D>L$ (本例中 $D=128>L=32$ )时代价很高。借鉴核方法/对偶 PCA 的 Gram 技巧,转而对小得多的 $L\times L$ Gram 矩阵 $G=\tfrac1L\tilde K\tilde K^\top$ 做特征分解。下面是这一技巧为什么成立的完整推导,本文独立验证过(论文陈述了结论但没有重新推导):

假设 $(\lambda, w)$ 是 $G$ 的一个单位范数特征对,即 $\tfrac1L\tilde K\tilde K^\top w = \lambda w$ 。两边左乘 $\tilde K^\top$ :

\frac1L\tilde K^\top\tilde K\,(\tilde K^\top w) = \lambda\,(\tilde K^\top w) \quad\Longrightarrow\quad \Sigma_b\left(\tilde K^\top w\right) = \lambda\left(\tilde K^\top w\right).

所以 $\tilde K^\top w$ 是 $\Sigma_b$ 的一个(未归一化的)特征向量,特征值同样是 $\lambda$ 。它的模长平方是 $\|\tilde K^\top w\|^2 = w^\top\tilde K\tilde K^\top w = w^\top(L G)w = L\lambda\,(w^\top w) = L\lambda$ (用到 $w^\top w=1$ ,因为 $w$ 是单位范数,且 $Gw=\lambda w$ )。所以 $\Sigma_b$ 归一化后的特征向量是 $u = \tilde K^\top w/\sqrt{L\lambda}$ ,而实际存储的缩放特征向量是

\xi = \sqrt\lambda\, u = \sqrt\lambda\cdot\frac{\tilde K^\top w}{\sqrt{L\lambda}} = \frac{\tilde K^\top w}{\sqrt L}. \tag{20}

这从第一性原理精确确认了论文的说法:你完全不需要构造或分解那个 $D\times D$ 的协方差矩阵——只需分解小的 $L\times L$ Gram 矩阵,再把它的前 $r$ 个特征向量通过式(20)映射回去即可。成本:构造 Gram 矩阵需要 $O(L^2D)$ ,它(很小的)特征分解需要 $O(L^3)$ ,把前 $r$ 个特征向量映射回 key 空间需要 $O(rLD)$ ——相比朴素做法的 $O(D^2L + D^3)$ 。在 $L=32\ll D=128$ 的情况下,单靠线性代数就能省下相当可观的成本,而且精度上没有任何损失(这是一个精确的重新表述,不是近似)。

算法 2:COBS 离线块描述符构造(Gram 技巧)
──────────────────────────────────────────────────────────────────
输入:块 b 的 L 个原始 key {k_r},r 属于 b;query 子空间基 U_Q(离线,共享)
输出:缓存的描述符(均值 k̄_b,r 个缩放特征向量 ξ_i,FP4 量化)

  1.  k̄_b  <-  (1/L) * sum_{r in b} k_r                      # 式(11),块均值
  2.  K_tilde[r, :]  <-  k_r - k̄_b   对 b 中每个 r             # 中心化后的 key,L x D
  3.  (可选,子空间变体)
        K_tilde  <-  K_tilde @ U_Q                            # 投影到 s 维子空间
  4.  G  <-  (1/L) * K_tilde @ K_tilde^T                       # Gram 矩阵,L x L(或 s x s)
  5.  (lambda_1..r, w_1..r)  <-  G 的前 r 个特征对              # O(L^3),小规模分解
  6.  对 i 从 1 到 r:
        xi_i  <-  (K_tilde^T @ w_i) / sqrt(L)                 # 式(20),映射回 key 空间
  7.  把每个 xi_i 量化为 FP4(E2M1),配一个 fp32 缩放因子       # 第 5.5 节
  8.  把(k̄_b 存 bf16,{xi_i}_{i=1..r} 存 FP4)存入缓存
──────────────────────────────────────────────────────────────────

一个数值旁证:块没填满或 $L \ge D$ 时会怎么样?

上面的 Gram 技巧推导自始至终假设 $L<D$ (块大小小于 head 维度),这正是分解 $L\times L$ Gram 矩阵比分解 $D\times D$ 协方差更便宜的原因。尽管论文没有直接讨论,但有两个边界情况值得明确说清楚,对实际实现而言很有用。

未填满的块。 在序列刚开始不久时(已生成的 token 不到 $L$ 个),“当前”块的 key 数量少于 $L$ 。均值和协方差公式(式 11)对任何实际占用量 $L'\le L$ 仍然适用——只需在式(7)、(11)、(20)中把 $L$ 换成实际占用量 $L'$ 即可。Gram 矩阵会缩小到 $L'\times L'$ ,分解起来只会更便宜;实际上唱一的麻烦是,当 token 非常少时( $L'\ll L$ ),经验协方差 $\Sigma_b$ 对任何“真实”底层散布的估计会更噴雑——不过这是小样本协方差估计的一般性统计观察,与 COBS 的构造方式无关。

退化的秩。 一个中心化后的 key 矩阵 $\tilde K\in\mathbb{R}^{L\times D}$ 的秩最多为 $\min(L,D)-1$ (减 1 是因为中心化会去掉一个自由度——中心化后的行总是加起来等于零向量)。在本文设置中 $L=32<D=128$ ,因此协方差 $\Sigma_b$ 的秩最多为 $31$ ,这正是为什么论文的秩掃描(图 7)停在 $r=31$ :向一个秩最多 $31$ 的矩阵要求 $r>31$ 个特征向量,额外的方向上特征值只会是零(或数值上徽不足道),没有任何贡献。这是单一块自身的局部协方差能编码多少信息的一个硬天花板,与计算或存储预算无关——这是一个块里恰好有多少 token 带来的结构性事实,而不是 Gram 技巧或 COBS 压缩方案本身的局限。

一阶选择器为什么有天花板——不只是 NSA 的问题

精确定义”仿射打分”这一族方法

论文第 6.3 节把讨论推广到了 NSA 之外。任何可缓存的选择器,每个块都缓存一个与 query 无关的向量 $\phi_b$ 和一个离线标量 $a_b$ ,并按仿射方式打分:

\text{score}_b^{\text{aff}}(q) = a_b + q^\top\phi_b. \tag{21}

均值池化( $\phi_b=\bar k_b$ , $a_b=\ln L$ )、NSA 可学习的 MLP 池化( $\phi_b=\text{MLP}(K_b)$ ,作用在块的 key 矩阵上)、以及 DeepSeek-V4 的 CSA 门控池化,都属于同一个仿射家族——产生 $\phi_b$ 的池化函数本身可以对原始 key 任意非线性,但缓存下来的条目仍然是单个与 query 无关的向量,靠一次简单的点积打分,所以打分本身关于 $q$ 仍然是仿射的。把这个家族的天花板直接和真实 log-质量的展开(式 13)对比:

\ln m_b = \ln L + q^\top\bar k_b + \frac12 q^\top\Sigma_bq + \cdots,

选 $a_b=\ln L$ 、 $\phi_b=\bar k_b$ 能精确匹配常数项和一次项,但没有任何 $a_b,\phi_b$ 的选择能够触及二次项——一个关于 $q$ 的二次型根本不能表示成 $q$ 的仿射函数,句号。这是一个干净的、可验证的数学事实,不需要靠实验来证明——本文里的实验证实了这一点带来的实际后果(NSA 可学习的 MLP 是被测试方法里最弱的选择器),但天花板本身是一个关于函数族的陈述。

GQA 跨头论证(以及它为什么只是非正式的)

有人可能希望,在 GQA 下,由于群组打分(式 6)是若干个头的 softmax 概率求和(每一个单独看都是关于 $q$ 的非线性函数),即便每个单独头的打分是仿射的,聚合起来或许也能以某种方式逃脱这个仿射天花板。论文给出了一个非正式的论证,说明这是不可能的:让一组头共享同一个块集合,严格来说比让每个头独立选择约束更强,所以一个更宽松的(每头独立)选择器只可能持平或更好:

\max_{\text{共享块集合}}\ \text{质量} \;\lesssim\; \max_{\text{独立块集合}}\ \text{质量}. \tag{22}

每头独立选择精确地归约到了 MHA 的情形,而按前面的论证,后者的打分对曲率是不敏感的——所以,只要式(22)的这个不等式成立,不管 GQA 的跨头非线性怎么样,整个一阶家族的检索质量都受限于同一个对曲率不敏感的天花板。本文和论文本身一样,明确用 " $\lesssim$ " 这个记号提醒:这只是一个非正式的、说得通的界,不是一个证明——支持这个论断的唯一证据是:NSA 可学习的 MLP(被测试的一阶池化里表达能力最强的一个)在经验上恰恰是最弱的选择器,这与天花板论证一致,但并不是对式(22)本身的正式确认。

Quest:一个”几乎”超越一阶的方法

Quest 为每个块存储 key 的逐元素最小值和最大值——一个轴对齐的包围盒:

k_b^{\min} = \min_{r\in b} k_r, \qquad k_b^{\max} = \max_{r\in b} k_r \quad \text{(逐元素)}. \tag{23}

它用盒子内任意一点能和 query 达到的最大内积来给块打分:

\hat s_b = \sum_i \max\!\left(q_ik_{b,i}^{\min},\, q_ik_{b,i}^{\max}\right) = \sum_i\left[\max(q_i,0)\,k_{b,i}^{\max} + \min(q_i,0)\,k_{b,i}^{\min}\right]. \tag{24}

因为这个打分关于 $q$ 是分段线性而非严格仿射的(涉及一个 $\max$ ,其行为取决于每个 $q_i$ 的符号),Quest 携带了一部分严格仿射打分不可能携带的曲率信息——但它是一个粗糙的、轴对齐的散布代理,对 key 维度之间的相关性完全不敏感,而 COBS 的协方差捕捉的是完整(压缩后)的相关结构。实验上,Quest(NoPE 方案下 32k RULER 为 0.5765)相对朴素均值池化(0.5554)只有适度的提升——剩下的大部分增益具体来自 COBS 的协方差摘要,而不是仅仅脱离严格仿射家族本身。

选择器	每块缓存内容	打分函数形式	捕捉到的累积量阶数	32k RULER(适用时用 NoPE)
均值池化	$\bar k_b$ ( $D$ 个浮点数)	仿射	仅一阶	0.5554
NSA(可学习 MLP)	$\text{MLP}(K_b)$ ( $D$ 个浮点数)	仿射	仅一阶(可学习,但仍是仿射)	0.2999(RoPE)/ 0.2822(NoPE)
Quest	$k_b^{\min}, k_b^{\max}$ ( $2D$ 个浮点数)	分段线性(逐维 max/min)	介于一阶与二阶之间(仅轴对齐散布)	0.5765
COBS(本文)	$\bar k_b$ + 秩- $r$ 的 $\Sigma_b$ ( $D{+}rD$ 或 $D{+}rs$ 个浮点数)	二次型	二阶(完整方向的协方差,已压缩)	0.8195(交付配置)
OSA(质量 oracle,不可部署)	完整原始 key	精确	全部阶数(精确)	0.9010
稠密(完整注意力)	—(无需选择)	—	—	0.9040

图 2(比较表格化图示,根据论文表 1 和第 6 节自绘):本文比较的每一个可缓存选择器,按其缓存摘要能表达的累积量阶数排列。注意 NSA 可学习的 MLP——被测试方法里参数表达能力最强的池化函数——表现却是最弱的,因为 $\phi_b$ 算得多聪明并不会改变打分本身的仿射函数形式。

实验:把结果拆开来看

实验设置

骨干网络是一个约 1.2B 参数的 decoder-only transformer:16 层,模型维度 2048,SwiGLU FFN 宽度 8192,16 个 query head、4 个 KV head(GQA,head 维度 $D=128$ ),RoPE base $\theta=10^6$ ,约 5 万 token 词表。预训练在 LongCrawl64 上跑约 200 亿 token,序列长度 4k(大致按 Chinchilla 缩放),接着用 YaRN 把上下文扩展到 32k,再在生成的 RULER 风格长上下文数据上做监督微调(SFT),这批数据与评测样本互不重叠。评测用的是 RULER 的一个 11 任务子集,在 32k 上下文下:单/多 key/多 query/多 value needle 检索,常见词与高频词抽取(CWE/FWE),以及变量追踪(VT)——QA 任务被明确排除(原因见后文局限性部分)。所有比较方法的 NSA 分支共享固定配置: $L=32$ token 的块,top- $k=16$ 个选中块,256 token 的滑动窗口。NSA MLP 是参数量匹配上唯一的例外——它被放大到约 1.7B,配上按比例更大的约 280 亿 token 预训练预算,而其他每个变体都是约 1.2B/约 200 亿 token——这是一个”厚待式基线”的选择,论文自己在解读 NSA MLP 后面更好的短上下文数字时也提到了这一点。

头条结果:RULER 差距阶梯

图 3(论文 Fig. 1):COBS 用一小部分 KV 缓存读取流量逼近稠密注意力的精度

图 3(论文 Fig. 1,嵌入):32k RULER 精度相对每解码步 KV 缓存读取流量(对数坐标)作图。虚线是 FP4 COBS 各配置的帕累托前沿;高亮的星形是交付配置(自适应 $s\approx85$ 、 $r=4$ 、FP4),正好落在前沿的”拐点”附近——用相对不大的流量代价换来了大部分精度提升。NSA MLP 和 NSA Quest 都被支配(在相同或更高流量下精度严格更差);NSA mean-pool 锚定了流量最低的一端;OSA 靠重新读取每一个 key 才换来接近稠密的精度。

图 4(论文 Fig. 3):从 NSA MLP 基线到稠密注意力的 32k RULER 差距阶梯

图 4(论文 Fig. 3,嵌入):每根柱子是一个选择器变体,按论文引入其加性改进的顺序排列。从灰色的”mean NoPE”(0.5554)和”NSA Quest”(0.5765)柱子,跳到绿色 COBS 柱子(0.8195–0.8493)的这一跳,是论文里最大、最可归因的单一增益——直观证据表明,是协方差项而非 NoPE 方案或其他任何小改动,承担了朝着补齐与稠密(0.9040)和 OSA(0.9010)差距这件事的大部分工作。

表 1(完整的 11 任务分解)讲了一个比均值更细致的故事:单 needle 子任务(S1、S2、S3)对几乎所有方法都已经饱和在 1.00——真正的区分度出现在多 key needle 检索(MK1–MK3)上,这类任务要求同时区分几个看起来相似的 needle。NSA MLP 在 MK3(三个同时出现的 needle)上得分惨淡的 0.00——完全做不到——而 COBS 全空间 $r=6$ 达到 0.48,OSA(oracle)达到 0.22。有意思的是,COBS 的全空间变体在某些单项子任务上超过了 OSA(例如 MK3:COBS $r=6$ 为 0.48,OSA 为 0.22),尽管 OSA 用的是精确质量——这是一个论文没有解释的反直觉结果,本文会在批判性分析部分进一步讨论,因为它要么说明这些具体子任务分数本身方差较大,要么说明按质量排序这个 oracle 本身,对这个特定子任务的 needle 结构而言,并没有校准得那么完美。

选择质量在短上下文上是否有代价?

表 2 用七个零样本常识基准(OpenBookQA、PIQA、HellaSwag、ARC-easy/challenge、TriviaQA、WinoGrande)检验了这一切有没有牺牲普通的短上下文能力,这些基准的输入只有几十到几百个 token——完全在局部窗口加上 top- $k$ 选中块已经能全覆盖的范围之内。不出所料,每个变体的平均准确率都落在与稠密注意力相差不到 0.6 个百分点的紧密区间内(COBS:平均 38.6%,对比稠密的 38.2%)——当输入几乎已经完全可见时,不管”选中”的是哪些块,选择本身能起的作用都不大。

逐位置语言建模:COBS 是否真的在利用长程内容?

图 5(论文 Fig. 4):留出长文本上的逐位置下一 token 负对数似然

图 5(论文 Fig. 4,嵌入):对于经过 SFT 训练的 GQA-4 变体,下一 token 的 NLL(越低越好)作为 32k token 文档中位置的函数。COBS(蓝色)拿到了所有被测方法里最低的平均 NLL(1.633),甚至低于稠密注意力本身(1.727)——但更有信息量的是曲线形状:COBS 的曲线在长位置处保持平坦,而 NSA MLP(橙色)的曲线从大约 8k token 开始明显攀升。

论文对这张图该怎么读非常谨慎,本文也认为这份谨慎是有必要的:COBS 在原始平均 NLL 上超过稠密的这个优势,很可能只是因为稠密注意力全程都在用 RoPE(一个位置编码上的差异,而不纯粹是注意力稀疏性的差异——见后文局限性部分),而不能证明”用了好选择器的稀疏注意力”在语言建模上真的比稠密更强。更站得住脚的读法是看曲线形状的对比:一个悄悄在远处块打分不可靠时退回局部窗口的选择器,会在长位置处表现出上升的 NLL(因为随着真正的长程依赖不断累积,只靠局部窗口的预测会越来越差)——这恰恰是 NSA MLP 和 NSA mean-pool 曲线里能看到的模式,两者都在大约 8k token 之后开始攀升,而 COBS 的曲线基本保持平坦。这份平坦,正是”COBS 的检索增益(前面’头条结果’一节)不是靠悄悄退化成一个只看局部的模型换来的”这一说法的证据。

KV 缓存读取流量:多出来的精度到底要付出什么代价?

表 3 把每层、每解码步的流量拆成了四个组成部分(摘要 key、摘要 value、滑动窗口、选中块的细粒度读取):

方法	摘要 key(KiB)	摘要 value(KiB)	窗口(KiB)	细粒度(KiB)	每层合计(KiB)	相对稠密(倍减少)	相对 NSA MLP(倍增加)
稠密(完整注意力)	–	–	–	65,536	65,536	–	18.29×
OSA(质量 oracle)	33,792	1024	512	1024	36,352	1.80×	10.14×
NSA MLP	1024	1024	512	1024	3584	18.29×	–
NSA mean-pool	1024	1024	512	1024	3584	18.29×	1.00×
NSA Quest	3072	1024	512	1024	5632	11.64×	1.57×
COBS 全空间 $r=4$ (bf16)	5120	1024	512	1024	7680	8.53×	2.14×
COBS 全空间 $r=6$ (bf16)	7168	1024	512	1024	9728	6.74×	2.71×
COBS 全空间 $r=4$ (FP4)	2112	1024	512	1024	4672	14.03×	1.30×
COBS 全空间 $r=6$ (FP4)	2656	1024	512	1024	5216	12.56×	1.46×
COBS(子空间 $s\approx85$ 、 $r=4$ 、FP4)——交付配置	1767	1024	512	1024	4327	15.15×	1.21×

图 6(论文表 3,重新整理为 markdown 表格):按分支拆分的每层、每解码步 KV 缓存读取流量。最重要的对比:NSA MLP 和 NSA mean-pool 都读取每层 3584 KiB;COBS 的交付配置( $s\approx85$ 、 $r=4$ 、FP4)读取每层 4327 KiB——相对 NSA 基线多读 1.21 倍——但仍比稠密注意力的每层 65,536 KiB 少读 15.15 倍。未压缩的全空间 bf16 变体成本高得多(每层 7680–9728 KiB,是 NSA 基线的 2.14–2.71 倍),这是一个诚实的”压缩前”数字,清楚地说明了 COBS 的实用性具体有多少来自子空间投影和 FP4 量化这两步,而不仅仅是协方差这个想法本身。

消融实验 1:秩的扫描及其令人意外的倒退

图 7(论文 Fig. 5):存储秩的消融实验,显示在 r=8 处达到峰值随后倒退

图 7(论文 Fig. 5,嵌入):32k RULER 分数随存储协方差秩 $r$ 变化的曲线。质量从 $r=1$ (0.719)平滑上升,在 $r=8$ 处达到峰值(0.8539),之后在 $r=16$ 处倒退到 0.8006,即便在理论最大秩 $r=L-1=31$ (0.8135)也无法恢复。

这种非单调性是论文里比较有意思的一个发现,而且有一个具体、可验证的机制:这个倒退几乎完全集中在多 key needle 子任务上——MK3 从( $r=8$ 时的)0.470 跌到( $r=16$ 时的)近乎全面崩溃的 0.054,MK2 从 0.934 跌到 0.800。论文的诊断是:到二阶为止,打分把一个有符号的线性项( $q^\top\bar k$ ,可正可负,取决于对齐程度)和一个严格非负的曲率项( $\tfrac12q^\top\Sigma_bq\ge0$ 恒成立,因为一个半正定矩阵的二次型不可能为负)结合在一起。这种不对称意味着:额外的特征向量——也就是加进这个非负曲率求和里的额外项——只可能增加一个块的分数,不管它和 query 的真实对齐方向是否真的有利。对于包含大量干扰项(和真正的 needle 很相似但并不是它的近似匹配 key)的块,多保留的特征向量会积累起虚假的”沿 $q$ 方向高方差”质量,这些质量和真正的相关性毫无关系,到 $r=16$ 时,这种虚假质量已经压过了真实检索信号——恰恰是在最需要区分少数几个非常相似候选项的多 key needle 任务上。再保留更多特征向量(直到理论最大值 $r=L-1=31$ ,因为由 $L=32$ 个中心化向量构成的协方差秩最多是 $L-1$ )也不会逆转这一点——论文报告 $r=31$ 时是 0.8135,仍远低于 $r=8$ 的峰值,MK3 仍大部分处于崩溃状态(0.062)。这为什么是一个重要的设计教训? 它说明”每个块存的信息越多越好”在这里并不成立——最初纯粹出于存储成本考虑而做的低秩截断( $r\le8$ ),结果反而顺带起到了一种有用的隐式正则化作用,恰好防住了这种虚假质量的失效模式,论文提到了这层联系,但没有进一步展开(本文在批判性分析部分指出了这个空缺)。

消融实验 2:自适应子空间分配对比全局分配

图 8(论文 Fig. 6):自适应的逐层 query 子空间维度分配

图 8(论文 Fig. 6,嵌入):捕获 90% query 谱能量所需的 query 子空间维度 $s$ ,按层独立计算并对各 head 取平均(深色柱子;跨层平均 $s\approx68$ ),再乘以 1.25 倍得到交付配置(柱子全高;平均 $s\approx84.875$ ,即" $s\approx85$ ")。较早的层(1–5)明显比中后段的层(9–13,聚集在 100 附近或更高)需要更少的 query 维度。

表 4 量化了这种逐层自适应为什么值得多花这份记账成本:统一套用给每一层的单一全局 $s=64$ 只能拿到 0.7856,而自适应分配——尽管平均下来是相近的 $s\approx68$ ——却能拿到 0.8054;把两者都乘上同样的 1.25 倍安全系数(全局 $s=96$ 对比自适应 $s\approx85$ ),得到 0.8188 对 0.8195,后者以明显更小的平均每块存储成本,基本追平了完整未压缩描述符的 0.8238。这里的教训是:不同层在利用多少 query 维度这件事上真的存在差异,用一刀切的子空间预算,会在不需要那么多维度的层上浪费容量,同时可能供给不足那些真正需要更多维度的层。

消融实验 3:量化几乎是免费的

表 5 证实,在测试过的每一种秩/子空间组合下,对存储的特征向量做 FP4 量化在精度上几乎没有代价——记录到的最大变化是 $-0.0026$ (全空间 $r=6$ :bf16 的 0.8493 → FP4 的 0.8467),同时在每一种配置下都把协方差因子的存储压缩了 3.6–3.8 倍。这是整篇论文里最划算的一笔:没有任何设计上的取舍,只是一个直白的数值观察——特征向量方向对极粗糙的量化的容忍度,远高于原始 key 或 value 张量通常能承受的程度。

两个负结果,都有诊断价值

围绕 query 中心展开会失效。 与其围绕 $q=0$ (原点)展开累积量级数,不如围绕某个校准过的非零 query 原点 $q_0$ 展开,希望能更好地拟合典型的 query 方向。表 6 显示这实际上会让选择质量变差(全空间 $r=4$ 时,0.8238 → 0.8100),而且集中体现在多 key 和多 value needle 任务上(MK3:0.34 → 0.26;MV:0.92 → 0.88)。论文给出了两个原因,都值得精确复述:第一,在固定的缓存预算下,倾斜矩(相对 $q_0$ 而非原点计算的矩)仍然必须被一个 KV 组内的 $G$ 个 query head 共享——但不同的 query head 可能有系统性不同的典型方向,所以任何单一的共享 $q_0$ 都只是一种折中,而不是对所有头都合适的拟合。第二,整个展开天生是局部的,围绕不管选定哪个原点都是如此,所以精度会在那些恰好远离校准过的 $q_0$ 的异常 needle query 上具体地变差——而这恰恰是一个 needle 检索基准专门设计出来施压的那种 query。

廉价的对角偏度修正是一把双刃剑。 前一小节的高秩倒退暗示,缺失的成分可能是有符号的三阶累积量 $\kappa_3$ ,它原则上可以抵消掉那个对符号不敏感的二次项带来的虚假质量。存储完整的 $O(D^3)$ 三阶张量并不现实,但一个廉价的对角近似——每个已经存储的特征向量 $u_i$ 配一个标量 $g_i$ ——可以加入一个有符号的三次修正:

\frac16\sum_{i,j,k}(\kappa_3)_{ijk}\,q_iq_jq_k \approx \frac16\sum_{i=1}^{r}\left(u_i^\top q\right)^3 g_i, \qquad g_i = \frac1L\sum_{t\in b}\left(u_i^\top(k_t-\bar k_b)\right)^3. \tag{25}

这个近似到底保留了什么、丢弃了什么(论文没有展开的一步推导)。 完整的三阶累积量是一个 $D\times D\times D$ 的张量;用 $\Sigma_b$ 的特征基表示时,这个对角近似只保留三个张量指标都恰好对应同一个 top- $r$ 特征方向( $i=j=k$ )的项,丢弃所有指向不同方向的混合项。每一个保留下来的对角项 $g_i$ ,恰好就是块的 key 偏差投影到那一个特征方向上的(未归一化的)三阶中心矩——也就是那个方向上的一个偏度统计量。

实验上(表 7),这项修正的效果完全取决于你已经处在哪个秩区间——论文如实报告了这个双刃剑的结果,而没有一笔带过。在低秩( $r=4,6$ ,也就是论文实际的运行区间)下,加入偏度项有害: $r=4$ 从 0.8238 掉到 0.7754,MK2 崩溃(0.78→0.54),MK3 尤其严重(0.34→0.01,几乎彻底失败)。但在倒退了的高秩设置下( $r=16$ ),同一个修正却有帮助,部分挽回了之前的崩溃:0.8006(无偏度)→ 0.8252(带偏度),MK3 从 0.05 恢复到 0.34。这种不对称直接支持了论文对 $r=16$ 倒退的诊断:有符号的三次项确实抵消了一旦保留太多特征向量就会累积的那部分虚假无符号方差质量——正是秩扫描消融里认定的那个失效模式。但即便是这个”修复过”的高秩结果(0.8252),仍然赶不上干净的低秩运行点的峰值(0.8539, $r=8$ ),而且还要多付出存储 16 个特征向量加 16 个偏度标量的 KV 流量代价——所以论文的结论,与其继续追逐三阶修正,不如”止步于协方差”(二阶),考虑到这些数字,是一个合理的选择,不过本文在批判性分析部分指出,论文从未测试过低秩加偏度的组合(比如 $r=4$ 或 $r=8$ 配上偏度),这就留下了一个悬而未决的问题:偏度会不会恰好在论文实际部署的那个区间里有帮助,而不只是在那个已经被放弃的高秩区间里。

这项工作打开的后续研究方向

把低秩与偏度修正组合起来测试,在 $r\in\{4,6,8\}$ 上都加上对角三阶项,看看一个小小的有符号修正能不能突破 $r=8$ 的峰值(0.8539),而不只是在已知会失败的 $r=4,6$ 或已经倒退的 $r=16$ 上测试。
一个真正的三阶(而不仅仅是对角偏度)消融——论文关于偏度的负结果只测试了最廉价的三阶累积量近似(每个已存的特征向量配一个标量);对完整三阶累积量的低秩张量近似(类比二阶低秩技巧向上提一阶)尚未测试。
把 COBS 与一个淘汰策略结合起来——既然 COBS 回答的是“这一步该以完整分辨率读哪些块”,淘汰回答的是“哪些 token 该永久删除”,两者并不互斥;一个既淘汰真正低价值 token、又对幸存下来的块用 COBS 的协方差感知打分的系统,有可能同时叠加两种流量节省。
前沿规模的验证——这套机制是在约 1.2B/4k 预训练规模下验证的;同样的 86% 补齐百分比,或任何接近它的数字,在 700 亿+参数、原生 128k+ 预训练上下文下是否仍然成立,是这篇论文留下的最明显、最具现实意义的开放问题。
用一个从头实现(而非“受控”)的 NSA 作为基线——对比 COBS 和一个字面上复现 DeepSeek 发布的 NSA(用 DeepSeek 原始超参数)的实现,能让 86% 这个数字参照一个外部可验证的基准,而不是一个内部、专属于这篇论文的基线。

论文章节对照表

论文章节	内容	在本文中对应位置
§1 引言	动机:为什么需要块稀疏性,陈述贡献	“一句话总结”、“前置知识”
§2 背景	NSA 的三分支、记号、可缓存性约束	“Native Sparse Attention:本文研究的系统对象”
§3 选择 Oracle	在 GQA 下推导质量排序标准	“一步步推导选择 Oracle”
§4 累积量展开	核心理论工具:CGF、累积量、天花板	“累积量展开:一阶选择器为什么有天花板”
§5 方法	NoPE、二阶截断、低秩、子空间、量化、成本	“方法:COBS,拆解开来讲”
§6 相关工作	CCQ、一阶选择器、Quest、正交方向	“一阶选择器为什么有天花板”、“相关工作地图”
§7 实验	设置、头条结果、消融、负结果	“实验:把结果拆开来看”
§8 结论	贡献总结	“总结”
§9 局限性	规模、受控基线、NoPE 混杂、SFT 协议、KV 核算	“论文承认的局限性与边界条件”
附录 A	可加 GQA 选择分数的推导	“让目标函数变得可处理的三条假设”
附录 B	块 key 分布的累积量	“显式推导为什么前两阶累积量就是均值和协方差”
附录 C	RULER 风格 SFT 与任务选取细节	“局限性”、“可复现性说明”

论文承认的局限性与边界条件

论文自己的第 9 节异常坦诚,每一点都值得精确复述,而不是笼统概括:

规模。 整个研究都是在约 1.2B 骨干参数(NSA MLP 复现在加上其可学习选择器后放大到约 1.7B)、4k token 预训练序列长度下进行的。论文明确说明这是”一项机制研究……而非展示一个部署规模的系统”——二阶选择”为什么”有效这件事得到了验证;这个补齐差距的比例在,比如说,700 亿+参数、128k+上下文下是否依然成立,并未测试。
NSA 对比是”受控的”,不是字面复现。 超参数与 DeepSeek 原始 NSA 论文不同(压缩分支用了更大的 MLP,块不重叠);这些选择在本文比较的每一个选择器变体里都保持一致,让论文的数字作为相对的、受控的比较是有效的,但不是对 NSA 原本报告的绝对数字的忠实复现。
NoPE 混杂因素。 长上下文比较里的每一个稀疏变体都在压缩/选择分支去掉了 RoPE;稠密注意力则全程保留 RoPE。所以稀疏变体和稠密之间的长上下文比较,部分反映的是这个位置编码上的差异,而不单纯是稀疏性本身——而且 NoPE 本身只在长上下文检索这一个配置里做了消融,没有在更广泛的场景下验证过。
RULER 风格的 SFT 是一个非标准协议。 长上下文检索信号来自在生成的 RULER 风格数据上做的监督微调(与评测样本不重叠,但用的是相同模板),论文自己也说,这意味着它的 RULER 数字应该被读作在一个相同、宽松协议下的相对选择质量,而不是可以直接迁移到更真实的长上下文场景的绝对精度。图 5 里某些变体在长位置处 NLL 上扬的斜率,本身也可能部分是这套 SFT 协议带来的产物,而不纯粹是注意力机制本身的性质。
KV 读取核算不等于实测延迟。 COBS 每块存储的东西严格多于 mean-pool——不量化的话,它读取的 KV 流量比 NSA 基线更多,即便量化到 FP4,相对 NSA 仍然净多读 1.21 倍(表 3)。本文这些核算数字本身,并不能直接证明端到端的实际墙钟时间会更快;那还取决于 kernel 实现、批处理策略、硬件的内存带宽特性,以及具体的解码场景(batch size、序列长度分布)——这些论文都没有直接测量。

批判性分析:不足与可改进之处

这篇论文本身的弱点和瑕疵

论文的核心理论贡献——累积量阶数这套框架,以及仿射打分无法表达二次项的证明——确实严谨,而且经得起复核(本文独立重新推导了上面的每一个关键步骤,没有发现错误)。但实验这部分有几处,论文给予的审视比它们本该得到的要少。

首先,前面提到的反直觉子任务结果——COBS 全空间 $r=6$ 在 MK3 上拿到 0.48,而 OSA 只有 0.22,也就是说在某一个子任务上超过了精确质量的 oracle——作者完全没有讨论,尽管这个数字就直接摆在他们自己的表 1 里。这要么反映了一个 11 子任务分解里、样本量恐怕不大所带来的方差(如果是这样,论文全篇没有报告任何跨随机种子的方差或误差棒,单次运行的点估计就有潜在的误导性),要么反映了按精确质量排序这个 oracle本身,对这个特定子任务的 needle 结构而言,并没有校准得那么完美。任何一种解释都值得论文写上一句,但论文没有提供。

其次,论文强调 COBS 平均 NLL 低于稠密注意力(1.633 对 1.727)这个结论时,用的措辞相当醒目（在“引言相关”的位置就提到了“本文比较中,拿到了最低的逐位置语言建模 NLL”）,却没有同样醒目地承认：在这个具体对比中，稠密注意力全程保留 RoPE，而每一个稀疏变体都在选择/压缩分支中用了 NoPE——这个混杂因素论文确实注意到了，但只在最后的局限性章节里提到，距离头条结论首次出现的位置已经好几页之远。一个只浏览摘要或只单独看“逐位置语言建模损失”这一小节的读者,很容易就直接认为“用了 COBS 的稀疏注意力在语言建模上严格优于稠密”,而这正是论文自己的局限性章节说不安全的结论。

第三,论文从未跑过它自己的诊断最直接暗示的那一项消融:将低秩体制与对角偏度修正结合( $r=4$ 或 $r=8$ 加上偏度标量)。论文对 $r=16$ 倒退的解释是无符号曲率累积了虚假阳性质量,而论文也确实展示了偏度在 $r=16$ 时具体地抵消了那部分虚假质量——但从未测试过在已经很好的 $r=8$ 运行点上加入一点少量的有符号修正,能不能突破 0.8539 这个峰值，只在 $r=4,6$ (有害)和 $r=16$ (已经是倒退了的区间)测试了偏度。这是一个廉价、直接的实验(在已有的消融表里每个秩多加一行),本可以直接检验论文自己的因果故事,它的缺失是一个真实的空白。

第四,论文完全没有对比 KV 淘汰或低秩 key 类方法(H2O、SnapKV、Loki),即便只是粗略的单一数字层面也没有,尽管这些方法在论文自己的引言和相关工作部分都被明确点名为另一类方法。论文自始至终的框架都是“一个二阶选择器能把一阶块选择与稠密注意力之间的差距补齐多少”——这是一个提问很好、也回答得很好的问题,但一个想要决定“到底该部署哪一类稀疏注意力方法”的读者,从这篇论文里得不到任何帮助——无法知道在相同的 KV 流量预算下,带上 COBS 的块选择跟一个调教得很好的淘汰方法相比,到底谁更値得选。

作者淡化或回避的局限

逆着文本的纹理去读,会发现一些标明的“局限性”部分并没有完全拥抱的东西。“受控 NSA 基线”这个提法,其实在帮一个忙,把头条对比隔离在审视之外:NSA MLP 被故意给予了更多的参数(1.7B 对 1.2B),预训练 token 预算也大约多 1.4 倍(280 亿对 200 亿),明确是为了让它“厄待”一些——但它仍然是论文头条“0.2999 → 0.8195,补齐 86% 差距”这个说法所锚定的那个基准。一个得到更多资源却仍然输得这么惨的基线,对论文的论点而言是一个比“资源匹配的基线”更强的结果——但论文自己的措辞(“NSA MLP 是参数量上唯一的例外……按比例放大它的预训练 token 预算”)读起来像是严谨的科学实践,但同样可以被解读为:把论文的核心数字(从 NSA MLP 测量的补齐百分比)变得比一个资源匹配的对比本来可能产生的要大。论文并没有报告如果换成参数匹配的 NSA MLP 基线,补齐百分比会变成多少——而这本来才是更公平、更值得写在头条里的数字。

类似地,论文反复强调 COBS 的数字在“测试过的每一种秩/子空间组合上都成立”,这种措辞隐含地暗示了一个系统性的掃描,但实际测试过的网格其实相当窄:秩掃描只用了 $\{1,2,3,4,5,6,8,16,31\}$ ,子空间消融也只用了 $s\in\{64,68,96,85,128\}$ ,而且都固定在块大小 $L=32$ 、top- $k=16$ 这一个配置上。至于那个定性的故事(峰值大约在 $r\approx8$ , $r\ge16$ 时倒退)在不同块大小或选择预算下是否仍然成立,论文并没有测试,也没有标注这一点没测试过。

具体的改进建议

报告跨随机种子的方差。 论文里每一个 RULER 数字都是单次运行。鉴于上面讨论过的惊讶子任务级结果(COBS 在 MK3 上超过 OSA),对头条配置重复跑 3–5 个随机种子并报告标准差(或哪怕只是重跑一遍那个最惊讶的单元格),能让读者分辨清“真实效应”和“11 子任务小样本分解里的噪声”。
跑一下低秩加偏度的组合。 在表 7 里为 $r=4$ +偏度、 $r=8$ +偏度各加一行(而不只是 $r=16$ +偏度),直接检验有符号修正在论文实际部署的运行区间里是否有帮助,而不只是在已经被放弃的高秩区间里。
报告一个参数匹配的 NSA MLP 消融。 哪怕只是作为一个较小的旁路实验,这样头条的“补齐 86% 差距”这个数字就能同时和厄待(1.7B)与匹配(1.2B)两个 NSA MLP 基线都做对比。
在图 1/表 3 的精度-流量对比里,至少加一个 KV 淘汰基线(比如 H2O 或 SnapKV),在相同的 KV 流量预算下对齐,这样论文回答的就不只是“COBS 比其他块选择器好吗”,而是“一个实践者到底该不该在相同代价下选 COBS+块选择,而不是一个同样廉价的淘汰方法”。
测试块大小和 top- $k$ 的敏感性。 既然秩倒退的故事明确指向了“块里干扰项很多”和多 key needle 密度,那么在不同块大小 $L$ (比如 16 或 64)或不同 top- $k$ 预算下重新跑一遍秩掃描,就能澄清 $r\approx8$ 这个最佳点到底是方法本身的性质,还是这个特定 $(L,k)$ 配置的产物。
把 NoPE 和稀疏性对 NLL 对比的影响拆开。 另外报告一条带 NoPE 的稠密注意力曲线(即便预期 NoPE 会对稠密注意力有害,因为稠密注意力本来就不需要位置不变的摘要)——这样才能让“COBS 在 NLL 上赢过稠密”这个说法单独由稀疏性本身来支撑或推翻,干净地与论文自己标注但没有完全控制住的那个位置编码混杂因素分开。

一个 GQA 工作例子:群组分数为什么不只是“把各头平均一下”

GQA 选择分数(式 6)很容易被误读成“只是把每个 query head 归一化后的质量平均一下”。一个小例子能说清楚为什么它其实是归一化质量的求和,而不是平均,这个区别很重要。假设 $G=2$ 个 query head 共享一个 KV head,只有两个候选块 $b_1,b_2$ 竞争一个 top-1 名额。假设 head 1 的原始质量是 $m_{b_1}^{(1)}=6$ , $m_{b_2}^{(1)}=2$ (所以 $Z^{(1)}=8$ ,归一化质量 $P_{b_1}^{(1)}=0.75$ , $P_{b_2}^{(1)}=0.25$ ),head 2 的原始质量是 $m_{b_1}^{(2)}=1$ , $m_{b_2}^{(2)}=9$ (所以 $Z^{(2)}=10$ ,归一化质量 $P_{b_1}^{(2)}=0.1$ , $P_{b_2}^{(2)}=0.9$ )。根据式(6):

\text{score}_{b_1} = P_{b_1}^{(1)}+P_{b_1}^{(2)} = 0.75+0.1=0.85, \qquad \text{score}_{b_2} = P_{b_2}^{(1)}+P_{b_2}^{(2)} = 0.25+0.9=1.15. \tag{P6}

块 $b_2$ 胜出( $1.15>0.85$ ),尽管它是 head 1 不那么喜欢的块( $P_{b_2}^{(1)}=0.25 < P_{b_1}^{(1)}=0.75$ )——因为 head 2 对它的渴望程度( $P_{b_2}^{(2)}=0.9$ )远远超过了 head 1 对自己最爱那块的渴望程度。这正是推导的要点所在:共享块集合的选取,目标是最小化整个组总共丢弃的质量(假设 3 的线性化结果),等价于最大化总共保留的质量——而不是满足任何单一 head 自己的偏好,也不是满足某种其他聚合方式(比如取最大值或两个 head 偏好的乘积)。如果换成一个按“多数 head 满意”或“满足单个偏好最强烈的 head”来选块的方法,在这里会选出不同(而且按推导自身的逻辑,在设定的假设下可证明更差)的选择。

这在实践中意味着什么。 一个对每个 head 都只有中等相关度的块,可能会超越一个对单个 head 极其相关、对其余 head 却无关紧要的块——群组分数奖励的是跨 head 的广泛相关性,而不是单一 head 的狭隘相关性。这是 GQA 共享块集合这个架构选择直接带来的后果;无论 COBS 还是本文比较的任何其他选择器,都无法回避这一点,因为它直接来自式(4)对每 KV head 目标函数的定义,甚至在第 3.3 节的任何近似假设被应用之前就已经存在了。

实践方案:如何把 COBS 接入已有的块稀疏流水线

对于一个已经在生产环境里跑着 NSA 风格(或 Quest 风格)选择分支的团队而言,这篇论文暗示的实践迁移路径是:(1) 只在压缩分支和选择分支里去掉 RoPE,滑动窗口仍保留——一个独立、低风险的改进;(2) 在块填满时,计算块均值,并通过 Gram 技巧(算法 2)算出前 $r$ 个协方差特征向量, $r$ 具体取 4–8 范围(不要更高——秩掃描消融实验直接警告了,把更高的秩当成免费午餐去拿是错的);(3) 离线从一批具代表性的 query 样本中校准出一个 query 子空间 $U_Q$ ,每层的 $s$ (而不是全局的)大约设为该层 query 谱能量 90% 所对应秩的 1.25 倍;(4) 将存储的特征向量量化到 FP4,块均值本身保持 bf16——根据表 5,这一步接近于免费的收益;(5) 打分时,每个解码步只投影一次 query,然后用廉价的 $O(rs)$ 内积公式(式 19)给每个块打分,永远不要在线重新构建完整的协方差矩阵。这篇论文唯一没有给出开箱即用方案的一步,是如何选取类似 RIPO 那套信任域文献里 $\delta$ 预算旋钮的等价物——COBS 的 $r$ 和 $s$ 都是针对这篇论文具体设置经验调出来的,团队在不同模型规模或上下文长度上采用这套方法时,应该预期至少重跑一遍秩掃描(图 7)和子空间分配消融(图 8),而不是直接析运 $r=4$ 、 $s\approx85$ 这套参数。

设计决策一览表

设计选择	做了什么	显而易见的替代方案	失效边界在哪里
压缩/选择分支里的 NoPE	把 RoPE 从块摘要中剔除,让摘要只编码内容而非位置	处处保留 RoPE,保持架构一致性	只在长上下文检索上验证过;对短上下文或非检索任务的影响未测
二阶(协方差)截断	把任何仿射选择器都会丢失的二次项找回来	继续用一阶,把功夫下在更强的池化函数(比如更大的 MLP)上	更大的池化函数依旧产出仿射打分——NSA 的大 MLP 在实验中恰恰是最弱的选择器,证实更多参数修不好一个结构性天花板
低秩( $r\le8$ )协方差	把 $D\times D$ 协方差压缩到 $D+rD$ 个浮点数	存精确协方差(或更高的秩)	反直觉地,更高的秩( $r\ge16$ )会倒退——不仅仅是存储/精度取舍,而是真实的精度悬崖
Query 子空间投影	沿着 query 自身的低维结构进一步压缩	每层都用同一个全局子空间维度 $s$	全局 $s$ 会在不需要那么多维度的层上浪费容量,在需要更多的层上供给不足——自适应逐层 $s$ 能在更小平均成本下补齐同样差距
对特征向量做 FP4 量化	存储字节数缩小 3.6–3.8 倍	为了保险,特征向量保留 bf16	在本文的测试中几乎没有实测到任何下降——整篇方法里性价比最高的一笔收益
在 $q=0$ (而非校准过的 $q_0$ )处展开累积量	简单,不需要校准数据,也不需要在群组内多个 head 之间妨协	围绕一个典型/校准过的 query 方向 $q_0$ 展开	倾斜矩必须在 GQA 组内共享(一种妨协),在远离 $q_0$ 的 query 上会变差——实验证实比直接在 $q=0$ 展开更差(表 6)
停在二阶(不加偏度/三阶项)	简单,最匹配论文实际运行的区间	加一个对角三阶累积量(“偏度”)修正	偏度在论文实际部署的低秩运行点( $r=4,6$ )上确实有害;只在已经被放弃的高秩区间( $r=16$ )有帮助——在实际部署设置上测试它是净亏(表 7)

常见问题

COBS 是一种新的注意力机制,还是 NSA 打分器的一个直接替代品? 后者。COBS 只改动了选择分支的打分函数——NSA 的压缩分支、滑动窗口、门控机制都没动。任何已有 NSA 风格三分支稀疏注意力实现的团队,原则上都可以在不动其他任何架构的情况下,把 COBS 的打分步骤插进去。

COBS 需要从头重新训练模型吗? 核心的全空间低秩 COBS checkpoint,和论文里其他每一个变体一样(预训练 + YaRN + SFT),都是完整训练出来的——基于协方差的打分在训练时就存在,不是事后补上去的。但子空间投影和 FP4 量化明确是推理时才做的改动,应用在一个已经训练好的全空间 checkpoint 上,不需要更新任何权重——所以至少压缩那几步(子空间、量化)对一个已有的全空间训练模型而言,是不需要重新训练的。

为什么秩越高反而越差,按道理信息越多不应该有害才对? 因为曲率项 $\tfrac12q^\top\Sigma_bq$ 在数学上被限定为非负(一个半正定矩阵的二次型),而真实的 log-质量修正项——为了从表面上相似的干扰项中真正区分出匹配项——可能真的需要有符号的信息。只增加无符号的质量,而没有相应的方式去为那些表面相似的干扰项“减去”质量,最终会压过真实信号——这是二阶(而非更高阶)截断的一个结构性性质,而不是低秩近似本身的 bug。

这是专门针对 NSA 的,还是能推广到其他块稀疏方法? 累积量阶数的论证(前面“一阶选择器为什么有天花板”一节)是对任何可缓存、与 query 无关的每块摘要都成立的通用论证——论文明确把自己的论点推广到了 DeepSeek-V4 的 CSA,以及(非正式地)Quest。具体的数字(32k RULER 上的 0.8195,1.21 倍流量开销)是 NSA 专属的,因为那是实际实现并测量过的系统,但底层的数学天花板论证本身是与具体方法无关的。

这能不能和(不只是协方差描述符的)真正 key/value 本身的 KV 缓存量化结合起来? 论文在局限性部分直接提到了这个开放问题:“即便标准 KV 也保持 FP4,占用仍然远低于稠密”——这暗示 COBS 相对稠密的流量优势会因此有所缩小(因为稠密也会从 KV 量化中受益),但不太可能完全消失,不过论文并没有具体报告这个对比的数字。

符号对照表

符号	含义
$q$	当前解码步的 query 向量(每个 head)
$k_r$ , $v_r$	历史 token $r$ 的 key 和 value 向量
$D$	每个 head 的 key/value 维度(本文设置中为 128)
$L$	块大小,以 token 计(本文设置中为 32)
$b$	连续 token 块的索引
$m_b$	块 $b$ 的注意力质量(式 1)
$v_b^c$	块 $b$ 的 value 质心(式 1)
$Z$	完整的 softmax 分母,对所有块求和
$P_b$	块 $b$ 归一化后的概率质量, $m_b/Z$
$S$	选择器保留下来的块子集
$\tau$	丢弃的总概率质量, $\sum_{b\in S^c}P_b$
$G$	GQA 下共享一个 KV head 的 query head 数量
$H$	KV head 总数
$g$ , $h$	query head(组内)和 KV head 的索引
$M_X(q)$	块 key 分布的矩生成函数
$K_X(q)$	累积量生成函数, $\ln M_X(q)$
$\kappa_1,\kappa_2,\kappa_3$	一阶、二阶、三阶累积量(均值、协方差、三阶矩)
$\bar k_b$	块 $b$ 的均值 key 向量
$\Sigma_b$	块 $b$ 内 key 的协方差矩阵
$r$	存储的协方差秩(保留的特征向量数)
$\lambda_i, u_i$	$\Sigma_b$ 的第 $i$ 个特征值和单位特征向量
$\xi_i$	缩放后的特征向量, $\sqrt{\lambda_i}\,u_i$ (实际存储的对象)
$U_Q$	聚合 query 二阶矩的前 $s$ 个特征向量
$s$	query 子空间维度
$\Pi$	到 query 子空间的正交投影, $U_QU_Q^\top$
$\tilde q$	投影到子空间后的 query, $U_Q^\top q$
$B_b$	投影到 query 子空间后的协方差, $U_Q^\top\Sigma_bU_Q$
$\tilde K$	一个块中心化后的 key 矩阵, $L\times D$
$g_i$	每个特征向量对应的对角三阶累积量(“偏度”)修正标量

公式索引

式号	定义内容
(1)	块质量 $m_b$ 与 value 质心 $v_b^c$
(2)	保留块集合 $S$ 上重新归一化后的输出
(3)	丢弃块带来的精确单头重建误差
(4)	对 query head 组求和后的每 KV head 目标函数
(5)	假设 1 下的重建误差上界
(6)	可加的 GQA 选择分数(oracle 标准)
(7)–(8)	通过矩/累积量生成函数改写块质量
(9)–(10)	推导一阶、二阶累积量即均值与协方差
(11)	具体的每块均值与协方差公式
(12)–(13)	log-质量的累积量展开,截断到三阶
(14)	COBS 的二阶截断打分公式
(15)	部署的 GQA 级 COBS 分数(把式 14 代入式 6)
(16)	协方差的低秩谱分解
(17)–(18)	协方差的 query 子空间投影
(19)	用存储的子空间特征向量计算的廉价单步打分成本
(20)	把小矩阵特征向量映射回 key 空间的 Gram 技巧恒等式
(21)	通用的仿射打分函数形式(一阶天花板)
(22)	非正式的 GQA 跨头界
(23)–(24)	Quest 的 min/max 描述符与分段线性打分
(25)	对角三阶累积量(“偏度”)修正

第二个完整例子:Gram 技巧在真实数字上的端到端验证

前面“成本核算与廉价计算特征向量的 Gram 矩阵技巧”一节从代数上证明了 Gram 技巧的恒等式;这里用同一个玩具块( $D=2$ 、 $L=4$ ,key 为 $k_1=(1,3),k_2=(1,-1),k_3=(-1,1),k_4=(-1,1)$ ,均值 $\bar k_b=(0,1)$ ,中心化后的行 $\tilde K = \begin{pmatrix}1&2\\1&-2\\-1&0\\-1&0\end{pmatrix}$ )给出一个完全数值化、可以手算验证的实例。

朴素路径。 直接构造 $\Sigma_b=\tfrac14\tilde K^\top\tilde K=\begin{pmatrix}1&0\\0&2\end{pmatrix}$ (前面已算过),并对这个 $2\times2$ 矩阵做特征分解:特征值 $\lambda_1=2,\lambda_2=1$ ,特征向量 $u_1=(0,1),u_2=(1,0)$ (已经是对角阵,所以这里直接就能看出来,但对一般的 $D\times D$ 矩阵而言,这一步要花 $O(D^3)$ )。

Gram 技巧路径。 改为构造 $L\times L=4\times4$ 的 Gram 矩阵 $G=\tfrac14\tilde K\tilde K^\top$ :

\tilde K\tilde K^\top = \begin{pmatrix}1&2\\1&-2\\-1&0\\-1&0\end{pmatrix}\begin{pmatrix}1&1&-1&-1\\2&-2&0&0\end{pmatrix} = \begin{pmatrix}5&-3&-1&-1\\-3&5&-1&-1\\-1&-1&1&1\\-1&-1&1&1\end{pmatrix}, \qquad G=\frac14\tilde K\tilde K^\top. \tag{P7}

这个 $4\times4$ 矩阵一般来说秩最多为 $\min(L,D)-1=\min(4,2)-1=1$ ,但在这里,由于 $D=2<L=4$ ,朴素路径(分解 $D\times D=2\times2$ 的协方差)已经更便宜了——这个玩具例子故意取得够小,以至于 $L>D$ ,恰恰跟论文真实的区间( $L=32<D=128$ )相反,目的就是要让这里的 Gram 矩阵变得庞大难用,从而具体地看清楚论文的技巧只有在 $L<D$ 时才真正划算。取 $G$ 的前两个特征对(直接计算可知,前两个非零特征值是 $\lambda_1=2, \lambda_2=1$ ——与协方差的特征值完全一致,即便在这个“反过来”的区间里也确认了 Gram 技巧恒等式中共享特征值那一部分),再通过式(20) $\xi_i = \tilde K^\top w_i/\sqrt{L}$ 映射回去,得到与朴素路径相同(只差符号)的缩放特征向量——这就是式(20)的恒等式在具体数字上的验证,而不仅仅是代数。

这个玩具例子想要教会的事情。 Gram 技巧并不是一种普适地更便宜的特征向量计算方式——它恰好在 $L<D$ (key 维度多、每块 token 少)时更便宜,而这正是论文实际的区间( $L=32,D=128$ )。在反过来的区间( $L>D$ ,如这个故意做得很小的玩具例子),朴素路径反而更便宜,一个真实实现应该选 $\Sigma_b$ ( $D\times D$ )和 $G$ ( $L\times L$ )中更小的那一个——这个细节无论是这篇论文还是大多数 Gram 技巧的介绍,都没有明说,因为大多数讲法(包括这篇论文)都隐含地默认了 $L<D$ 这个区间,从未明确说出这个临界条件。

最后的一个理性检验:天花板论证依赖 softmax 吗?

一个合理的问题是:一阶天花板(式 12–13、式 21)会不会只是 softmax 注意力特有的一个伪象,而不是累积量生成函数这套框架一般性的性质?在真正重要的意义上,它并不是 softmax 专属的:推导只用到了块质量 $m_b=\sum_{r\in b}e^{q^\top k_r}$ 的两个结构性事实——(1) 它是关于 $q$ 的线性函数的指数求和(这使它,差一个常数,恰好就是在 $q$ 处求值的矩生成函数),以及 (2) 可缓存性约束迫使任何选择器缓存的摘要必须与 $q$ 无关。任何未归一化打分同样是关于 query 的“指数包裹线性项”形式的注意力变体(这基本上涵盖了当今使用的每一种基于 softmax 的注意力机制,包括 MHA、MQA、GQA,因为它们在归一化之前都共享同样的 $q^\top k$ 打分形式),都会继承完全相同的累积量展开,以及对仿射打分、可缓存选择器的同样天花板。这就是为什么论文的论点(第 6 节,以及本文上面“一阶选择器为什么有天花板”一节)被陈述为适用于整个可缓存块选择器家族,而不仅仅是 NSA 的具体实现细节。

缩略语与术语表

KV 缓存 —— Key-Value 缓存;存储所有此前生成 token 的 key/value 向量,每个解码步都要读取。
NSA —— Native Sparse Attention(DeepSeek-AI);本文研究的三分支(压缩/选择/窗口)稀疏注意力架构。
CSA —— (DeepSeek-V4)一种 token 窗口池化 + ReLU”闪电索引器”选择机制;本文天花板论证同样覆盖的另一个一阶选择器。
DSA —— DeepSeek Sparse Attention(DeepSeek-V3.2);通过轻量索引器做逐 token(而非逐块)选择——与本文的块选择范围正交。
HCA —— (DeepSeek-V4)一个对长跨度做稠密压缩、完全不做选择的分支——同样与本文范围正交。
GQA —— 分组查询注意力;多个 query head 共享一个 KV head。
MHA —— 多头注意力;每个 query head 都有自己的 KV head(GQA 在 $G=1$ 时的特例)。
RoPE —— 旋转位置编码;在做点积之前,按位置相关的角度旋转 key/query。
NoPE —— 无位置编码;本文在压缩/选择分支专门去掉 RoPE 的方案。
RULER —— 一套长上下文检索基准(needle-in-haystack 变体、词抽取、变量追踪),本文所有长上下文评测都用它。
OSA —— Oracle Sparse Attention;本文诊断用、不可部署的选择器,按重新读取得到的精确质量给块排序。
COBS —— Cumulant Order Block Sparse Attention;本文提出的方法。
MGF / CGF —— 矩生成函数 / 累积量生成函数;支撑累积量展开的概率论工具。
YaRN —— 一种上下文窗口扩展方法,本文用它把 4k 预训练的骨干网络拉伸到 32k 上下文,再做 SFT。
SFT —— 监督微调(Supervised Fine-Tuning)。
FP4(E2M1) —— 一种 4-bit 浮点格式(2 位指数、1 位尾数),用于量化存储的特征向量。

常见误读要避免

“COBS 赢了稠密注意力,所以稀疏注意力现在严格优于稠密。” 不完全对——唯一一个 COBS 赢过稠密的指标(逐位置 NLL,图 5)本身就混杂了稠密全程保留 RoPE、而每个稀疏变体都用 NoPE 这个因素;论文自己也标注了这一点。真正干净、对等对等的比较(32k RULER、短上下文常识测试)显示的是 COBS 补齐了大部分但不是全部与稠密的差距,而不是超越了稠密。
“OSA 证明了按质量排序是一个完美的选择标准。” OSA 补齐了 99.5% 的差距,这是非常强的证据,但表 1 也显示 COBS 的全空间变体在至少一个子任务(MK3)上超过了 OSA——这个本文和论文都无法完全解释的结果,应该让人对“OSA 代表一个硬性的、可证明最优的天花板”这种说法持保留态度,而不是把它当成一个硬上限。
“存的秩越多总是越好,反正是存了更多信息。” 秩掃描消融(图 7)就是一个直接、明确的反例:质量在 $r=8$ 达到峰值后倒退,而且有一个具体、机制性的原因(无符号曲率累积虚假阳性质量)——不是一个模糊的“边际效益递减”故事,而是一个真实的逆转。
“COBS 是一种新的注意力算法。” 它只改动了 NSA 选择分支的打分函数;注意力机制本身(softmax、三分支门控、压缩分支和窗口分支)完全没动。把它称为“一种新的稀疏注意力方法”相对于“一个现有稀疏注意力方法选择分支的更好打分器”而言,有些词不达意地放大了这项工作的范围。
“86% 补齐百分比是一个公平、对等对等的头条指标。” 它是相对一个故意给予了资源优势的 NSA MLP 基线(更多参数、更多预训练 token)测定的——作者确实做了这个选择,也可以说是偏保守的,但这个背景应该让人在脱离上下文引用 86% 这个数字时多一分警惕。

可复现性说明

作者说明他们为一个不同的、更小的伴生结果(仅在相关讨论中一笔带过的、类似 KV 缓存归档的编解码方法,不是这篇论文本身)发布了代码——对于这篇论文,请查阅 arXiv 列表页是否有代码链接,因为本文审阅的这个版本正文里没有直接给出。所用的训练/评测基础设施是 seqax,MatX 开源的、基于 JAX 的研究用 LLM 代码库(论文里引用为参考文献 [31]),它是公开可用的,与这篇具体论文的确切配置是否随之发布无关。想复现头条数字的读者应当注意,论文自己也提醒过,它的 NSA 基线是一个受控的内部复现,而不是跑了一遍 DeepSeek 原版发布的 NSA 代码——所以不应期望绝对的 RULER 数字,能和一个用不同超参数、不同块重叠约定、原始 DeepSeek 规模训练配方从零实现的 NSA 精确对上。32k 上下文的 RULER SFT 数据被描述为与评测样本不重叠的”生成”RULER 风格数据,用的是与公开 RULER 基准(arXiv:2404.06654)相同的任务模板,但训练集是单独生成的一份——要精确复现这一点,需要根据 RULER 论文自己对模板的描述重新实现那套生成流程,因为本文没有发布这份具体的 SFT 数据集。

关键公式快查表

阅读论文时方便快查,把本文最依赖的五个公式集中放在这里:

块质量与质心: $m_b = \sum_{r\in b} e^{q^\top k_r}$ , $v_b^c = \tfrac{1}{m_b}\sum_{r\in b} e^{q^\top k_r} v_r$ (式 1)。
可加的 GQA 选择分数: $\text{score}_b^{(h)} = \sum_{g=1}^{G} m_b^{(g,h)}/Z^{(g,h)}$ (式 6)。
log-质量的累积量展开: $\ln m_b = \ln L + q^\top\bar k_b + \tfrac12 q^\top\Sigma_bq + \tfrac16\sum(\kappa_3)_{ijk}q_iq_jq_k+\cdots$ (式 13)。
COBS 实际部署的二阶分数:将公式 2 中的 $m_b^{(g,h)}$ 换成 $\exp(q^\top\bar k_b + \tfrac12q^\top\Sigma_bq)$ (式 15)。
Gram 技巧的特征向量映射: $\xi_i = \tilde K^\top w_i/\sqrt L$ ,其中 $w_i$ 是小 Gram 矩阵 $G=\tfrac1L\tilde K\tilde K^\top$ 的特征向量(式 20)。

论文里其余的每一个公式(子空间投影、量化方案、两个负结果的修正项)都是这五个公式中某一个的变体或压缩。

总结

COBS 是这样一个案例:拿一个被广泛使用的经验性启发式规则——每个块缓存一个池化后的向量,拿它和 query 打分——用真正的概率论工具而不是更多的架构搜索,去精确追问这个启发式规则到底表达不了什么。答案干净且可验证:任何可缓存的、与 query 无关的每块摘要,只要是靠点积打分,在数学上就注定被限制在关于 query 的仿射函数里,而一个块真正应得的注意力质量里有一个曲率项——由块内 key 协方差驱动——是任何仿射函数都不可能捕捉到的,不管这个摘要本身算得多么精巧、池化函数多么有表达力。把这个缺失的二阶项找回来,经过低秩谱分解、共享的 query 子空间投影、以及激进的 FP4 量化压缩之后,在一个受控的 NSA 基线和稠密注意力之间,补齐了大约 86% 的剩余差距,代价只是比 NSA 基线多读 1.21 倍的 KV 缓存流量,比稠密注意力少读 15.15 倍。论文的消融实验对“哪里会失效”坐得住脚:一个注定非负的曲率项最终会在高秩时累积起虚假阳性质量,一个 query 中心化的想法没能带来收益,一个三阶累积量修正只在论文否则会放弃的区间里有帮助——正是这种细节,才让一项机制研究值得信任。这套具体方案在前沿模型规模下、在一个非“受控”而是从头实现的 NSA 对比下、以及在更宽的稀疏注意力方法集合上是否仍然成立,留待未来的工作去回答;这篇论文牢牢站住脚跟的,是诊断本身:一直以来真正限制住可缓存块选择的,是累积量阶数,而不是选择器设计得多机智。