论文

COBS:累积订单块稀疏注意力

COBS: Cumulant Order Block Sparse Attention

模型推理KV Cache

摘要

块稀疏注意力是一种硬件友好的方法,可以缓解 大语言模型 (LLM) 中的键值 (KV) 缓存读取瓶颈。然而,它在领先的开放权重 LLM 中并不普遍,而是依赖于密集的注意力或细粒度的选择,从而激发了我们的分析。我们研究 DeepSeek 的原生稀疏注意力(NSA)作为代表性方法,其三分支设计让我们能够隔离块选择,这是最具挑战性和最重要的阶段。我们将选择形式化,并将其简化为通过单个数量(注意力质量:块注意力分数的总和)对块进行排名。我们证明,如果选择检索具有最大注意力质量的块,则块稀疏注意力可以与密集注意力的质量相匹配。然而,计算精确的注意力质量需要读取每个键,因此块选择的问题最终减少为从紧凑的摘要而不是完整的键中近似该质量。通过累积扩展,我们展示了现有方法为何会失败:它们的选择策略试图估计注意力质量,但仅限于一阶近似。因此,我们提出了 COBS(累积阶块稀疏注意力),这是一种基于 NSA 的注意力方法,结合了一个新颖的选择器,该选择器存储每个块的压缩二阶统计量。在 32k RULER 长上下文检索基准上,COBS 将 NSA 基线的平均得分从 0.2999 提高到 0.8195,接近密集注意力 0.9040,缩小了约 86% 的差距,同时仅使用 NSA 基线 1.21 倍的 KV 缓存读取流量,比密集注意力少 15.15 倍的读取流量。在我们的比较中,相同的模型保留了短上下文行为,并且获得了比密集注意力更低的位置负对数似然(NLL)。