XAttention：対角線スコアリングを備えたブロックスパースアテンション

XAttention: Block Sparse Attention with Antidiagonal Scoring

March 20, 2025

著者: Ruyi Xu, Guangxuan Xiao, Haofeng Huang, Junxian Guo, Song Han

cs.AI

要旨

長文脈Transformerモデル（LCTM）は実世界のアプリケーションにおいて重要ですが、アテンションの二次的な計算複雑性により高い計算コストが課題となっています。ブロックスパースアテンションは、重要な領域に計算を集中させることでこれを緩和しますが、既存の手法ではブロックの重要性を測定するコストが高く、精度と効率のバランスを取るのが困難です。本論文では、スパースアテンションを用いてTransformerモデルの長文脈推論を劇的に加速するプラグアンドプレイフレームワーク、XAttentionを紹介します。XAttentionの重要な革新点は、アテンション行列の対角線方向（左下から右上）の値の合計が、ブロックの重要性を表す強力な代理指標となるという洞察です。これにより、非必須のブロックを正確に特定し、削除することが可能となり、高いスパース性と劇的な推論の加速が実現します。言語処理のRULERやLongBench、ビデオ理解のVideoMME、ビデオ生成のVBenchといった厳しい長文脈ベンチマークでの包括的な評価を通じて、XAttentionは完全なアテンションと同等の精度を維持しながら、大幅な計算効率の向上を達成しました。アテンション計算において最大13.5倍の加速を実証しています。これらの結果は、XAttentionがブロックスパースアテンションの実用的な可能性を解き放ち、LCTMの実世界アプリケーションにおけるスケーラブルで効率的な展開の道を開く能力を強調しています。コードはhttps://github.com/mit-han-lab/x-attentionで公開されています。

English

Long-Context Transformer Models (LCTMs) are vital for real-world applications but suffer high computational costs due to attention's quadratic complexity. Block-sparse attention mitigates this by focusing computation on critical regions, yet existing methods struggle with balancing accuracy and efficiency due to costly block importance measurements. In this paper, we introduce XAttention, a plug-and-play framework that dramatically accelerates long-context inference in Transformers models using sparse attention. XAttention's key innovation is the insight that the sum of antidiagonal values (i.e., from the lower-left to upper-right) in the attention matrix provides a powerful proxy for block importance. This allows for precise identification and pruning of non-essential blocks, resulting in high sparsity and dramatically accelerated inference. Across comprehensive evaluations on demanding long-context benchmarks-including RULER and LongBench for language, VideoMME for video understanding, and VBench for video generation. XAttention achieves accuracy comparable to full attention while delivering substantial computational gains. We demonstrate up to 13.5x acceleration in attention computation. These results underscore XAttention's ability to unlock the practical potential of block sparse attention, paving the way for scalable and efficient deployment of LCTMs in real-world applications. Code is available at https://github.com/mit-han-lab/x-attention.