SeerAttention-R：面向长推理的稀疏注意力适配
机器学习
2025-06-11 v1 人工智能
摘要
我们提出了SeerAttention-R，一个专门为推理模型的长解码设计的稀疏注意力框架。SeerAttention-R从SeerAttention扩展而来，保留了通过自蒸馏门控机制学习注意力稀疏性的设计，同时移除了查询池化以适应自回归解码。通过轻量级插件门控，SeerAttention-R具有灵活性，可以轻松集成到现有预训练模型中而无需修改原始参数。我们证明SeerAttention-R在仅使用0.4B token训练的情况下，在AIME基准测试中在使用大稀疏注意力块大小（64/128）时以4K token预算保持了接近无损的推理精度。使用TileLang，我们开发了一个高度优化的稀疏解码内核，在H100 GPU上以90%稀疏度实现了相对于FlashAttention-3高达9倍的理论加速。代码可在https://github.com/microsoft/SeerAttention获取。
引用
@article{arxiv.2506.08889,
title = {SeerAttention-R: Sparse Attention Adaptation for Long Reasoning},
author = {Yizhao Gao and Shuming Guo and Shijie Cao and Yuqing Xia and Yu Cheng and Lei Wang and Lingxiao Ma and Yutao Sun and Tianzhu Ye and Li Dong and Hayden Kwok-Hay So and Yu Hua and Ting Cao and Fan Yang and Mao Yang},
journal= {arXiv preprint arXiv:2506.08889},
year = {2025}
}