让大模型长推理更高效，用稀疏注意力节省计算资源
SeerAttention-R: Sparse Attention Adaptation for Long Reasoning
- 通过自蒸馏门控机制动态学习注意力稀疏模式
- 仅用0.4B词元训练，4K上下文仍保持接近无损的推理准确率
- 可无缝接入现有模型，适配长序列推理场景
我们提出SeerAttention-R，一种专为长推理模型设计的稀疏注意力框架。基于SeerAttention改进，保留自蒸馏门控机制以学习注意力稀疏性，同时移除查询池化以支持自回归解码。采用轻量级插件式门控结构，无需修改原有参数即可灵活集成至预训练模型。实验表明，仅在0.4B词元上训练的SeerAttention-R，在AIME基准下使用4K令牌预算与64/128的大稀疏注意力块时，仍能维持近似无损的推理准确率。借助TileLang开发的高度优化稀疏解码内核，在H100 GPU上90%稀疏度下，速度比FlashAttention-3提升接近理论极限的9倍。代码已公开：https://github.com/microsoft/SeerAttention。
原文摘要 · Abstract (English)
We introduce SeerAttention-R, a sparse attention framework specifically tailored for the long decoding of reasoning models. Extended from SeerAttention, SeerAttention-R retains the design of learning attention sparsity through a self-distilled gating mechanism, while removing query pooling to accommodate auto-regressive decoding. With a lightweight plug-in gating, SeerAttention-R is flexible and can be easily integrated into existing pretrained model without modifying the original parameters. We demonstrate that SeerAttention-R, trained on just 0.4B tokens, maintains near-lossless reasoning accuracy with 4K token budget in AIME benchmark under large sparse attention block sizes (64/128). Using TileLang, we develop a highly optimized sparse decoding kernel that achieves near-theoretical speedups of up to 9x over FlashAttention-3 on H100 GPU at 90% sparsity. Code is available at: https://github.com/microsoft/SeerAttention.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。