SeerAttention-R: Sparse Aanpassing van Aandacht voor Lang Redeneren

SeerAttention-R: Sparse Attention Adaptation for Long Reasoning

June 10, 2025

Auteurs: Yizhao Gao, Shuming Guo, Shijie Cao, Yuqing Xia, Yu Cheng, Lei Wang, Lingxiao Ma, Yutao Sun, Tianzhu Ye, Li Dong, Hayden Kwok-Hay So, Yu Hua, Ting Cao, Fan Yang, Mao Yang

cs.AI

Samenvatting

We introduceren SeerAttention-R, een raamwerk voor sparse aandacht dat specifiek is afgestemd op het lange decoderen van redeneermodellen. Als uitbreiding van SeerAttention behoudt SeerAttention-R het ontwerp van het leren van aandachts-sparsiteit via een zelfgedistilleerd gatingmechanisme, terwijl query-pooling wordt verwijderd om auto-regressief decoderen mogelijk te maken. Met een lichtgewicht plug-in gating is SeerAttention-R flexibel en kan het eenvoudig worden geïntegreerd in bestaande voorgetrainde modellen zonder de oorspronkelijke parameters aan te passen. We tonen aan dat SeerAttention-R, getraind op slechts 0,4B tokens, bijna verliesvrije redeneernauwkeurigheid behoudt met een budget van 4K tokens in de AIME-benchmark onder grote sparse aandachtsblokgroottes (64/128). Met behulp van TileLang ontwikkelen we een sterk geoptimaliseerde sparse decodeerkernel die bijna theoretische snelheidsverbeteringen tot 9x bereikt ten opzichte van FlashAttention-3 op de H100 GPU bij 90% sparsiteit. Code is beschikbaar op: https://github.com/microsoft/SeerAttention.

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.