XAttention: Blok Sparse Attention met Antidiagonale Scoring

XAttention: Block Sparse Attention with Antidiagonal Scoring

March 20, 2025

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

cs.AI

Samenvatting

Long-Context Transformer-modellen (LCTMs) zijn essentieel voor real-world toepassingen, maar lijden onder hoge computationele kosten vanwege de kwadratische complexiteit van aandacht (attention). Block-sparse attention vermindert dit door de rekenkracht te richten op kritieke regio's, maar bestaande methoden worstelen met het balanceren van nauwkeurigheid en efficiëntie vanwege kostbare metingen van blokbelangrijkheid. In dit artikel introduceren we XAttention, een plug-and-play framework dat de inferentie van lange context in Transformer-modellen aanzienlijk versnelt door gebruik te maken van sparse attention. De belangrijkste innovatie van XAttention is het inzicht dat de som van antidiagonale waarden (d.w.z. van linksonder naar rechtsboven) in de aandachtmatrix een krachtige proxy biedt voor blokbelangrijkheid. Hierdoor kunnen niet-essentiële blokken precies worden geïdentificeerd en verwijderd, wat resulteert in een hoge sparsity en aanzienlijk versnelde inferentie. Uit uitgebreide evaluaties op veeleisende lange-context benchmarks – waaronder RULER en LongBench voor taal, VideoMME voor videobegrip en VBench voor videogeneratie – blijkt dat XAttention een nauwkeurigheid bereikt die vergelijkbaar is met volledige aandacht, terwijl het aanzienlijke computationele voordelen biedt. We demonstreren een versnelling van tot 13,5x in de aandachtberekening. Deze resultaten onderstrepen het vermogen van XAttention om het praktische potentieel van block sparse attention te ontsluiten, waardoor de weg wordt geëffend voor schaalbare en efficiënte inzet van LCTMs in real-world toepassingen. Code is beschikbaar op 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.