XAttention: Block Sparse Attention with Antidiagonal Scoring

2025-03-20

AI Summary

Long-context Transformer inference is hindered by the quadratic computational complexity of standard attention. Existing block-sparse methods struggle to balance efficiency and accuracy due to the high overhead of explicit block importance scoring. This paper proposes a plug-and-play block-sparse attention framework that, for the first time, identifies the sum of elements along the anti-diagonal (bottom-left to top-right) of the attention matrix as an efficient, parameter-free proxy for block importance—enabling rapid, dynamic block pruning without explicit scoring. Our method preserves full-attention accuracy across diverse long-context multimodal benchmarks, including RULER, LongBench, VideoMME, and VBench, while achieving up to 13.5× speedup in attention computation. It significantly enhances inference efficiency for long-context modeling without compromising performance.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Computer Vision: Large Vision Models
• Natural Language Processing: (Large) Language Models

Application Category

• Security and Privacy: Large-scale security measurements
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs

Abstract

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.

Problem

• Reduces computational costs in Long-Context Transformer Models
• Improves accuracy and efficiency in block-sparse attention
• Accelerates long-context inference using antidiagonal scoring

Innovation

• Antidiagonal scoring for block importance
• Plug-and-play sparse attention framework
• High sparsity with accelerated inference