arXiv:2607.09052cs.LG2026-07被引 1

提出COBS方法,让大模型高效实现接近稠密注意力的长文本处理能力。

COBS: Cumulant Order Block Sparse Attention

  • 通过二阶统计量近似注意力质量,提升块稀疏注意力选择精度
  • 在32k上下文任务中得分从0.2999提升至0.8195,接近稠密注意力表现
  • 仅需1.21倍缓存读取量,适合部署在资源受限的推理环境

块稀疏注意力是缓解大语言模型中键值缓存读取瓶颈的硬件友好方式。然而,主流开源大模型仍采用稠密注意力或细粒度选择策略,这促使我们展开分析。以DeepSeek的原生稀疏注意力(NSA)为代表,其三分支结构使块选择阶段得以独立研究。我们形式化了选择过程,将其简化为仅用一个量——注意力质量(即块内注意力分数之和)进行排序。若选择注意力质量最高的块,块稀疏注意力可达到稠密注意力的性能水平。但精确计算注意力质量需读取所有键,因此问题转化为如何从紧凑摘要中近似该质量。通过累积量展开分析,我们发现现有方法失败原因在于仅使用一阶近似。为此,我们提出COBS(Cumulant Order Block Sparse Attention),基于NSA设计,引入新选择器,为每块存储压缩的二阶统计量。在32k RULER长文本检索基准上,COBS将NSA基线的平均得分从0.2999提升至0.8195,接近稠密注意力的0.9040,填补约86%差距,同时仅需1.21倍于NSA的缓存读取流量,比稠密注意力少15.15倍。相同模型在短文本任务中保持良好行为,并在位置级负对数似然(NLL)上优于稠密注意力。

原文摘要 · Abstract (English)

Block sparse attention is a hardware friendly way to alleviate the key-value (KV) cache read bottleneck in large language models (LLMs). However, it is not prevalent among leading open-weight LLMs, which rely instead on dense attention or fine-grained selection, thereby motivating our analysis. We study DeepSeek's Native Sparse Attention (NSA) as a representative method, whose three-branch design lets us isolate block selection, the most challenging and consequential stage. We formalize selection and reduce it to ranking blocks by a single quantity, the attention mass: the sum of a block's attention scores. We show that if selection retrieves the blocks with the largest attention mass, block sparse attention can match the quality of dense attention. However, computing the exact attention mass requires reading every key, so the problem of block selection ultimately reduces to approximating this mass from a compact summary instead of the full keys. Via a cumulant expansion, we show why existing methods falter: their selection strategies attempt to estimate the attention mass, but are confined to a first-order approximation. Therefore, we propose COBS (Cumulant Order Block Sparse Attention), an attention method that builds on NSA, incorporating a novel selector that stores a compressed second-order statistic per block. On the 32k RULER long-context retrieval benchmark, COBS raises the NSA baseline's mean score from 0.2999 to 0.8195, approaching dense attention at 0.9040 and closing about 86% of the gap, while using only 1.21x the KV cache read traffic of the NSA baseline and 15.15x less read traffic than dense. The same model preserves short-context behavior and attains lower position-wise negative log-likelihood (NLL) than dense attention in our comparison.

稀疏注意力长文本生成缓存优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。