Stratified Sparse Attention: Reducing the Sparse-to-Dense Gap at Extreme Sparsity

📄 arXiv: 2610.28441v1 📥 PDF

作者: Sahil Joshi, Prithvi Dixit, Aryan Katiyar, Agniva Chowdhury, Joseph E. Gonzalez, Anshumali Shrivastava, Aditya Desai, Kumar Krishna Agrawal

分类: cs.AI

发布日期: 2026-10-20

备注: 26 pages, 5 figures, 13 tables

🔗 代码/项目: GITHUB

💡 一句话要点

提出分层稀疏注意力以解决极端稀疏下的性能瓶颈问题

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

关键词: 稀疏注意力 长上下文推理 大型语言模型 性能优化 自然语言处理 深度学习 模型加速

📋 核心要点

  1. 现有稀疏注意力方法在确定性检索与采样之间存在性能权衡，导致在极端稀疏下与密集注意力的性能差距较大。
  2. 本文提出的分层稀疏注意力（SSA）框架，结合了高得分令牌的确定性检索与基于得分的采样，从而有效恢复注意力质量。
  3. 在多个基准测试中，SSA在相同稀疏度下显著提高了准确率，并在长上下文推理中实现了高达2.96倍的解码速度提升。

📝 摘要（中文）

大型语言模型在长上下文推理中，因处理日益增长的键值缓存而面临成本瓶颈。稀疏注意力通过仅访问部分缓存来降低成本，但现有方法在确定性检索与采样之间存在权衡。本文提出分层稀疏注意力（SSA），结合了确定性检索与基于得分的采样，能够在高得分令牌上进行确定性关注，同时通过代理注意力得分对剩余部分进行采样，恢复超出检索集的注意力质量。实验结果表明，SSA在多个基准测试中均优于现有基线，且在长上下文推理中实现了显著的加速。

🔬 方法详解

问题定义：本文旨在解决大型语言模型在长上下文推理中，因稀疏注意力方法导致的性能损失与计算成本问题。现有方法在确定性检索与采样之间存在显著的性能差距，尤其在极端稀疏情况下。

核心思路：分层稀疏注意力（SSA）通过结合确定性检索与基于得分的采样，确保在高得分令牌上进行有效关注，同时对剩余部分进行采样，从而恢复注意力质量，减少信息损失。

技术框架：SSA的整体架构包括两个主要模块：首先，进行高得分令牌的确定性检索；其次，使用代理注意力得分对剩余令牌进行采样，并根据其包含概率进行校正。SSA还涵盖了top-k、top-p和混合top-k+采样等特殊情况。

关键创新：SSA的核心创新在于其将确定性检索与采样相结合的统一框架，克服了现有方法在极端稀疏下的性能瓶颈，能够有效恢复注意力质量。

关键设计：在实现中，SSA采用了优化的CUDA内核，以提高计算效率，并在多个模型上进行了验证，确保其在不同参数规模下均能实现显著的性能提升。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在RULER-HARD基准测试中，SSA在Ministral-3-8B模型下的准确率从81.8%提升至84.2%，在32K上下文时，恢复了94%的密集注意力性能；在128K上下文时，准确率从63.8%提升至69.3%，恢复了85%的性能。此外，SSA在1M上下文下实现了高达2.96倍的解码速度提升。

🎯 应用场景

分层稀疏注意力（SSA）在长上下文推理中具有广泛的应用潜力，特别是在大型语言模型的推理任务中。其高效的注意力机制能够显著降低计算成本，提高模型在实际应用中的响应速度，适用于对实时性要求较高的自然语言处理任务，如对话系统和文本生成等。

📄 摘要（原文）

Large Language Model inference over long contexts is increasingly bottlenecked by the cost of attending to a growing Key-Value (KV) cache. Sparse attention reduces this cost by accessing only a fraction of the cache, but existing methods face a trade-off between deterministic retrieval, which captures high-scoring tokens but discards residual attention mass, and sampling, which improves coverage at the cost of higher variance. Hybrid methods combine top-$k$ retrieval with uniform sampling, yet still exhibit a substantial gap from dense attention at aggressive sparsity. We introduce \emph{Stratified Sparse Attention (SSA)}, a unified framework combining deterministic retrieval with score-guided sampling. SSA deterministically attends to high-scoring tokens while sampling the residual using proxy attention scores and correcting for their inclusion probabilities, thereby recovering attention mass beyond the retrieval set. SSA encompasses top-$k$, top-$p$, and hybrid top-$k$+sampling as special cases. Across RULER, LongBench, LOFT, and AIME, and five models ranging from 3B to 27B parameters, SSA consistently improves over the baselines at matched sparsity. On RULER-HARD with Ministral-3-8B, SSA improves accuracy over the strongest practical baseline from 81.8\% to 84.2\% at 32K and from 63.8\% to 69.3\% at 128K under $50\times$ attention sparsity, recovering 94\% and 85\% of dense-attention performance, respectively. With optimized CUDA kernels, SSA achieves up to $2.96\times$ decode speedup over FlashInfer at 1M context, demonstrating that its accuracy gains translate into practical long-context inference acceleration. We open-source our code at https://github.com/skylight-org/SSA-Paper-Code.