XAttention : Attention parcimonieuse par blocs avec score antidiagonal

XAttention: Block Sparse Attention with Antidiagonal Scoring

March 20, 2025

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

cs.AI

Résumé

Les modèles Transformer à contexte long (LCTMs) sont essentiels pour les applications réelles, mais ils souffrent de coûts de calcul élevés en raison de la complexité quadratique de l'attention. L'attention parcimonieuse par blocs atténue ce problème en concentrant les calculs sur les régions critiques, mais les méthodes existantes peinent à équilibrer précision et efficacité en raison des mesures coûteuses de l'importance des blocs. Dans cet article, nous présentons XAttention, un framework plug-and-play qui accélère considérablement l'inférence à contexte long dans les modèles Transformer grâce à une attention parcimonieuse. L'innovation clé de XAttention réside dans l'idée que la somme des valeurs antidiagonales (c'est-à-dire du coin inférieur gauche au coin supérieur droit) dans la matrice d'attention constitue un puissant indicateur de l'importance des blocs. Cela permet une identification et une élimination précises des blocs non essentiels, entraînant une forte parcimonie et une accélération significative de l'inférence. À travers des évaluations approfondies sur des benchmarks exigeants à contexte long — incluant RULER et LongBench pour le langage, VideoMME pour la compréhension vidéo, et VBench pour la génération vidéo — XAttention atteint une précision comparable à celle de l'attention complète tout en offrant des gains de calcul substantiels. Nous démontrons une accélération allant jusqu'à 13,5 fois dans le calcul de l'attention. Ces résultats soulignent la capacité de XAttention à libérer le potentiel pratique de l'attention parcimonieuse par blocs, ouvrant la voie à un déploiement scalable et efficace des LCTMs dans des applications réelles. Le code est disponible à l'adresse 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.