XAttention: 대각선 반대 방향 점수화를 통한 블록 희소 어텐션

XAttention: Block Sparse Attention with Antidiagonal Scoring

March 20, 2025

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

cs.AI

초록

장문맥 트랜스포머 모델(LCTMs)은 실세계 애플리케이션에 필수적이지만, 어텐션의 2차 복잡도로 인해 높은 계산 비용이 발생합니다. 블록 희소 어텐션은 중요한 영역에 계산을 집중함으로써 이를 완화하지만, 기존 방법들은 비용이 많이 드는 블록 중요도 측정으로 인해 정확도와 효율성 간의 균형을 맞추는 데 어려움을 겪습니다. 본 논문에서는 희소 어텐션을 사용하여 트랜스포머 모델의 장문맥 추론을 획기적으로 가속화하는 플러그 앤 플레이 프레임워크인 XAttention을 소개합니다. XAttention의 핵심 혁신은 어텐션 행렬에서 대각선 값(즉, 좌하단에서 우상단으로)의 합이 블록 중요도를 측정하는 강력한 대리 지표로 사용될 수 있다는 통찰입니다. 이를 통해 비필수 블록을 정확하게 식별하고 제거할 수 있어 높은 희소성과 획기적으로 가속화된 추론을 가능하게 합니다. 언어를 위한 RULER 및 LongBench, 비디오 이해를 위한 VideoMME, 비디오 생성을 위한 VBench 등 까다로운 장문맥 벤치마크에 대한 포괄적인 평가에서 XAttention은 전체 어텐션과 비슷한 정확도를 유지하면서도 상당한 계산 이점을 제공합니다. 우리는 어텐션 계산에서 최대 13.5배의 가속화를 입증했습니다. 이러한 결과는 XAttention이 블록 희소 어텐션의 실용적 잠재력을 해제할 수 있는 능력을 보여주며, LCTMs의 확장 가능하고 효율적인 실세계 배포를 위한 길을 열어줍니다. 코드는 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.