SeerAttention-R: Адаптация разреженного внимания для длинных рассуждений

SeerAttention-R: Sparse Attention Adaptation for Long Reasoning

June 10, 2025

Авторы: 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

Аннотация

Мы представляем SeerAttention-R — разреженный механизм внимания, специально разработанный для длительного декодирования в моделях логического вывода. Развивая идеи SeerAttention, SeerAttention-R сохраняет подход к обучению разреженности внимания через механизм самообученного гейтинга, при этом исключая пулинг запросов для поддержки авторегрессивного декодирования. Благодаря легковесному подключаемому гейтингу, SeerAttention-R является гибким и может быть легко интегрирован в существующие предобученные модели без изменения их исходных параметров. Мы демонстрируем, что SeerAttention-R, обученный всего на 0,4 миллиарда токенов, сохраняет почти безошибочную точность логического вывода при бюджете в 4K токенов в бенчмарке AIME при использовании крупных блоков разреженного внимания (64/128). С использованием TileLang мы разработали высокооптимизированное ядро разреженного декодирования, которое достигает почти теоретического ускорения до 9 раз по сравнению с FlashAttention-3 на GPU H100 при 90% разреженности. Код доступен по адресу: 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.