arXiv cs.AI· Sahil Joshi, Prithvi Dixit, Aryan Katiyar, Agniva Chowdhury, Joseph E. Gonzalez, Anshumali Shrivastava, Aditya Desai, Kumar Krishna Agrawal·· 11 小时前AI 评分40
Stratified Sparse Attention（SSA）：在极端稀疏度下缩小稀疏注意力与稠密注意力的差距
Stratified Sparse Attention: Reducing the Sparse-to-Dense Gap at Extreme Sparsity
AI 导读
研究者提出 Stratified Sparse Attention（SSA），将确定性检索与分数引导采样结合，用代理注意力分数对残余部分采样并按纳入概率校正，从而找回检索集之外的注意力质量，top-k、top-p 与 top-k+采样均为其特例。
正文 · AI 翻译
摘要：长上下文下的大语言模型推理正日益受到不断增长的键值（KV）缓存注意力计算成本的制约。稀疏注意力通过仅访问缓存的一小部分来降低这一成本，但现有方法面临两难权衡：确定性检索能捕获高分 token，却会丢弃残余注意力质量；采样则能提升覆盖范围，但代价是方差更高。混合方法将 top-$k$ 检索与均匀采样相结合，但在激进稀疏度下与稠密注意力仍存在显著差距。我们提出 \emph{分层稀疏注意力（Stratified Sparse Attention, SSA）}，这是一个将确定性检索与分数引导采样相统一的框架。SSA 确定性地关注高分 token，同时利用代理注意力分数对残余部分进行采样，并校正其纳入概率，从而恢复检索集之外的注意力质量。SSA 将 top-$k$、top-$p$ 以及 top-$k$+采样混合方法作为特例涵盖其中。在 RULER、LongBench、LOFT 和 AIME 上，以及从 3B 到 27B 参数的五种模型上，SSA 在相同稀疏度下始终优于基线。在 RULER-HARD 上使用 Ministral-3-8B 时，在 $50\times$ 注意力稀疏度下，SSA 将准确率从最强实用基线的 81.8\% 提升至 32K 下的 84.2\%，从 63.8\% 提升至 128K 下的 69.3\%，分别恢复了稠密注意力性能的 94\% 和 85\%。借助优化的 CUDA 内核，SSA 在 1M 上下文下相比 FlashInfer 实现最高 $2.96\times$ 的解码加速，表明其准确率提升可转化为实际的长上下文推理加速。我们在 此 https URL 开源了代码。
| 评论： | 26 页，5 幅图，13 张表 |
| 主题： | 人工智能（cs.AI） |
| 引用为： | arXiv:2610.28441 [cs.AI] |
| （或 arXiv:2610.28441v1 [cs.AI] 用于此版本） | |
| https://doi.org/10.48550/arXiv.2610.28441 由 DataCite 颁发的 arXiv DOI（待注册） |
提交历史
来自：Sahil Joshi [查看邮箱]
[v1]
2026 年 10 月 20 日星期二 16:02:12 UTC（2,105 KB）
来源：arXiv cs.AI · arxiv.org