SeerAttention-R: Sparse Attention Adaptation for Long Reasoning

要約

Seerattention-Rを紹介します。これは、推論モデルの長いデコードに合わせて特別に調整されたまばらな注意フレームワークです。
Seerattentionから拡張されたSeerattention-Rは、自己誘導ゲーティングメカニズムを介して注意のスパースを学習することの設計を保持しながら、自動回帰デコードに対応するためにクエリプーリングを削除します。
軽量のプラグインゲーティングを使用すると、Seerattention-Rは柔軟性があり、元のパラメーターを変更せずに既存の前提型モデルに簡単に統合できます。
わずか0.4bトークンで訓練されたSeerattention-Rは、大規模な注意ブロックサイズ（64/128）の下で、AIMEベンチマークで4Kトークン予算でほぼ紛れもない推論の精度を維持していることを実証します。
Tilelangを使用して、H100 GPUで90％スパースでFlashattention-3を超える最大9倍のほぼ理論的スピードアップを達成する高度に最適化されたスパースデコードカーネルを開発します。
コードは、https：//github.com/microsoft/seerattentionで入手できます。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google