XAttention: Block-Sparse-Aufmerksamkeit mit Antidiagonal-Bewertung

XAttention: Block Sparse Attention with Antidiagonal Scoring

March 20, 2025

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

cs.AI

Zusammenfassung

Long-Context Transformer-Modelle (LCTMs) sind für reale Anwendungen von entscheidender Bedeutung, leiden jedoch unter hohen Rechenkosten aufgrund der quadratischen Komplexität der Aufmerksamkeitsmechanismen. Block-sparse Attention mildert dies, indem sie die Berechnung auf kritische Regionen konzentriert, doch bestehende Methoden kämpfen mit der Balance zwischen Genauigkeit und Effizienz aufgrund kostspieliger Messungen der Blockbedeutung. In diesem Artikel stellen wir XAttention vor, ein Plug-and-Play-Framework, das die Langkontext-Inferenz in Transformer-Modellen durch spärliche Aufmerksamkeit erheblich beschleunigt. Die zentrale Innovation von XAttention ist die Erkenntnis, dass die Summe der antidiagonalen Werte (d.h. von der unteren linken zur oberen rechten Ecke) in der Aufmerksamkeitsmatrix ein leistungsstarker Proxy für die Blockbedeutung darstellt. Dies ermöglicht die präzise Identifikation und das Beschneiden nicht-essentieller Blöcke, was zu hoher Sparsity und einer drastisch beschleunigten Inferenz führt. In umfassenden Bewertungen anspruchsvoller Langkontext-Benchmarks – einschließlich RULER und LongBench für Sprache, VideoMME für Videoverständnis und VBench für Videogenerierung – erreicht XAttention eine Genauigkeit, die der vollen Aufmerksamkeit vergleichbar ist, bei gleichzeitig erheblichen Rechengewinnen. Wir demonstrieren eine Beschleunigung der Aufmerksamkeitsberechnung um bis zu das 13,5-fache. Diese Ergebnisse unterstreichen die Fähigkeit von XAttention, das praktische Potenzial block-sparser Aufmerksamkeit zu erschließen und den Weg für eine skalierbare und effiziente Bereitstellung von LCTMs in realen Anwendungen zu ebnen. Der Code ist verfügbar unter 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.