XAttention: Attenzione Sparsa a Blocchi con Punteggio Antidiagonale

XAttention: Block Sparse Attention with Antidiagonal Scoring

March 20, 2025

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

cs.AI

Abstract

I modelli Transformer a contesto lungo (LCTMs) sono fondamentali per applicazioni nel mondo reale, ma soffrono di elevati costi computazionali a causa della complessità quadratica del meccanismo di attenzione. L'attenzione sparsa a blocchi mitiga questo problema concentrando il calcolo sulle regioni critiche, tuttavia i metodi esistenti faticano a bilanciare precisione ed efficienza a causa delle costose misurazioni dell'importanza dei blocchi. In questo articolo, introduciamo XAttention, un framework plug-and-play che accelera drasticamente l'inferenza a contesto lungo nei modelli Transformer utilizzando l'attenzione sparsa. L'innovazione chiave di XAttention è l'intuizione che la somma dei valori antidiagonali (cioè dall'angolo inferiore sinistro a quello superiore destro) nella matrice di attenzione fornisce un potente indicatore dell'importanza dei blocchi. Ciò consente un'identificazione precisa e una potatura dei blocchi non essenziali, ottenendo un'elevata sparsità e un'accelerazione significativa dell'inferenza. Attraverso valutazioni complete su benchmark impegnativi a contesto lungo, tra cui RULER e LongBench per il linguaggio, VideoMME per la comprensione video e VBench per la generazione video, XAttention raggiunge una precisione comparabile all'attenzione completa offrendo sostanziali vantaggi computazionali. Dimostriamo un'accelerazione fino a 13,5x nel calcolo dell'attenzione. Questi risultati sottolineano la capacità di XAttention di sbloccare il potenziale pratico dell'attenzione sparsa a blocchi, aprendo la strada a un'implementazione scalabile ed efficiente degli LCTMs nelle applicazioni reali. Il codice è disponibile all'indirizzo 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.