XAttention: Atención Dispersa en Bloques con Puntuación Antidiagonal

XAttention: Block Sparse Attention with Antidiagonal Scoring

March 20, 2025

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

cs.AI

Resumen

Los modelos Transformer de contexto largo (LCTMs, por sus siglas en inglés) son fundamentales para aplicaciones del mundo real, pero sufren altos costos computacionales debido a la complejidad cuadrática de la atención. La atención dispersa por bloques mitiga este problema al enfocar el cálculo en regiones críticas, aunque los métodos existentes luchan por equilibrar precisión y eficiencia debido a las costosas mediciones de importancia de los bloques. En este artículo, presentamos XAttention, un marco plug-and-play que acelera drásticamente la inferencia de contexto largo en modelos Transformer utilizando atención dispersa. La innovación clave de XAttention es la idea de que la suma de los valores antidiagonales (es decir, desde la esquina inferior izquierda hasta la superior derecha) en la matriz de atención proporciona un poderoso indicador de la importancia de los bloques. Esto permite una identificación y poda precisa de bloques no esenciales, lo que resulta en una alta dispersión y una inferencia acelerada de manera significativa. En evaluaciones exhaustivas en puntos de referencia exigentes de contexto largo, incluyendo RULER y LongBench para lenguaje, VideoMME para comprensión de video y VBench para generación de video, XAttention logra una precisión comparable a la atención completa mientras ofrece ganancias computacionales sustanciales. Demostramos una aceleración de hasta 13.5x en el cálculo de la atención. Estos resultados subrayan la capacidad de XAttention para desbloquear el potencial práctico de la atención dispersa por bloques, allanando el camino para un despliegue escalable y eficiente de LCTMs en aplicaciones del mundo real. El código está disponible en 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.