XAttention: Atenção Esparsa em Blocos com Pontuação Antidiagonal

XAttention: Block Sparse Attention with Antidiagonal Scoring

March 20, 2025

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

cs.AI

Resumo

Modelos de Transformers de Contexto Longo (LCTMs) são essenciais para aplicações do mundo real, mas sofrem com altos custos computacionais devido à complexidade quadrática do mecanismo de atenção. A atenção esparsa por blocos mitiga isso ao focar a computação em regiões críticas, mas os métodos existentes lutam para equilibrar precisão e eficiência devido às medições custosas da importância dos blocos. Neste artigo, introduzimos o XAttention, uma estrutura plug-and-play que acelera drasticamente a inferência de contexto longo em modelos Transformers usando atenção esparsa. A inovação-chave do XAttention é a percepção de que a soma dos valores antidiagonais (ou seja, do canto inferior esquerdo ao superior direito) na matriz de atenção fornece um poderoso indicador da importância dos blocos. Isso permite a identificação precisa e a poda de blocos não essenciais, resultando em alta esparsidade e inferência drasticamente acelerada. Em avaliações abrangentes em benchmarks exigentes de contexto longo — incluindo RULER e LongBench para linguagem, VideoMME para compreensão de vídeo e VBench para geração de vídeo — o XAttention alcança precisão comparável à atenção completa enquanto oferece ganhos computacionais substanciais. Demonstramos uma aceleração de até 13,5x no cálculo da atenção. Esses resultados destacam a capacidade do XAttention de desbloquear o potencial prático da atenção esparsa por blocos, pavimentando o caminho para a implantação escalável e eficiente de LCTMs em aplicações do mundo real. O código está disponível em 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.