COBS: Cumulant Order Block Sparse Attention

📄 arXiv: 2607.09052v1 📥 PDF

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

分类: cs.LG

发布日期: 2026-07-10

💡 一句话要点

提出COBS以解决大语言模型中的KV缓存读取瓶颈问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 块稀疏注意力 大语言模型 注意力机制 二阶统计量 性能提升 自然语言处理 计算效率

📋 核心要点

  1. 现有方法在块选择阶段面临挑战，无法有效估计注意力质量，导致性能不足。
  2. COBS通过引入压缩的二阶统计量，改进了块选择策略，从而提高了块稀疏注意力的效果。
  3. 实验结果显示，COBS在长上下文检索任务中显著提升了性能，且读取流量低于密集注意力。

📝 摘要（中文）

块稀疏注意力是一种硬件友好的方法，用于缓解大语言模型中的键值（KV）缓存读取瓶颈。然而，目前在主流的开放权重大语言模型中，块稀疏注意力并不普遍，主要依赖于密集注意力或细粒度选择。本文分析了DeepSeek的原生稀疏注意力（NSA），并提出COBS（累积阶块稀疏注意力），通过压缩的二阶统计量来改进块选择，显著提高了注意力质量。实验结果表明，COBS在32k RULER长上下文检索基准上将NSA的平均分数从0.2999提升至0.8195，接近密集注意力的0.9040，同时在KV缓存读取流量上也表现出色。

🔬 方法详解

问题定义：本文旨在解决大语言模型中块稀疏注意力的块选择问题。现有方法在选择阶段依赖于一阶近似，导致无法准确估计注意力质量，影响模型性能。

核心思路：COBS的核心思路是通过累积展开引入二阶统计量，改进块选择策略，从而更准确地估计注意力质量。这样设计的目的是为了在保持计算效率的同时，提升块稀疏注意力的效果。

技术框架：COBS建立在DeepSeek的NSA基础上，采用三分支设计来隔离块选择阶段。整体流程包括块的选择、注意力质量的估计和最终的注意力计算。

关键创新：COBS的主要创新在于引入了压缩的二阶统计量作为块选择的依据，这与现有方法的单一一阶近似形成了本质区别，显著提升了选择的准确性。

关键设计：在参数设置上，COBS使用了压缩的二阶统计量来替代传统的全键读取，减少了计算开销。同时，损失函数和网络结构经过优化，以确保在长上下文下的性能提升。通过这些设计，COBS在KV缓存读取流量上表现出色。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，COBS在32k RULER长上下文检索基准上将NSA的平均分数从0.2999提升至0.8195，接近密集注意力的0.9040，缩小了约86%的性能差距。同时，COBS的KV缓存读取流量仅为NSA基线的1.21倍，远低于密集注意力的15.15倍。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、机器翻译和对话系统等。通过提高大语言模型的效率和性能，COBS能够在实际应用中提供更快速和准确的响应，具有重要的实际价值和未来影响。

📄 摘要（原文）

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.