Source-linked AI summary

XAttention: Block Sparse Attention with Antidiagonal Scoring

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

arXiv:2503.16428v1cs.CLcs.CV

TL;DR

Long-context Transformers are hindered by quadratic attention costs, and existing block-sparse methods incur overhead when identifying important blocks. XAttention uses antidiagonal attention sums for thresholded block selection and reports accuracy comparable to full attention with up to 13.5× faster attention computation. Its evaluations span language, video understanding, and video generation benchmarks.

  • Problem

    Quadratic attention costs hinder long-context Transformer deployment, while existing sparse methods face accuracy-efficiency trade-offs from costly block-importance measurement.

  • Method

    XAttention uses antidiagonal attention sums as a block-importance proxy, then selects blocks whose normalized scores exceed a threshold.

  • Results

    Across RULER, LongBench, VideoMME, and VBench, XAttention achieves accuracy comparable to full attention and up to 13.5× acceleration in attention computation.

  • Takeaways & Limitations

    XAttention supports efficient and scalable deployment of long-context Transformers by unlocking practical block-sparse attention.

  • Takeaways & Limitations

    Applying XAttention from the beginning of HunyuanVideo denoising caused slight layout shifts and lower quantitative scores than full attention.

Abstract

from arXiv · show

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.

1. Introduction

Long-context Transformers face quadratic attention costs that hinder deployment, while existing block-sparse methods struggle to identify important blocks efficiently. XAttention uses antidiagonal sums to select sparse attention blocks, preserving accuracy while accelerating computation.

  • Long-Context Transformer Models support tasks such as video understanding and generation that require processing exceptionally long sequences.
  • Attention cost scales quadratically with sequence length, creating a substantial pre-filling bottleneck for long-context models.
  • Existing block-sparse methods face an accuracy-efficiency trade-off because identifying important attention blocks can introduce substantial overhead.
  • XAttention estimates block importance using sums of antidiagonal attention values, offering a computationally efficient alternative to token pooling.
  • 13.5× acceleration in attention computation is achieved during pre-filling while accuracy remains comparable to full attention across RULER, LongBench, VideoMME, and VBench.

2. Method

XAttention predicts block importance from antidiagonal attention sums, selects blocks whose normalized scores meet a threshold, and optionally tunes thresholds across attention heads. Its design targets robust pattern detection with reduced computation.

  • Method overview: XAttention comprises importance prediction, important-block selection, and minimum-threshold prediction for attention heads.
  • Antidiagonal scoring: Antidiagonal sums provide a block-importance proxy because every token contributes to an antidiagonal and the pattern intersects vertical and slash patterns.
  • Threshold Block selection: The selection procedure sums antidiagonal elements within S × S regions, applies softmax normalization, and retains the minimal blocks whose cumulative probability exceeds threshold τ.
  • Threshold Block selection: The resulting process determines important attention-map blocks from antidiagonal scores and the specified threshold, producing a sparse attention mask.
  • Minimum Threshold Prediction: Dynamic programming adjusts thresholds per attention head to balance accuracy and computation because heads can differ in sparsity and importance.
  • Minimum Threshold Prediction: Thresholds are reduced by 10% at each step to gradually reduce computation while preserving each head’s contribution to accuracy.
  • Minimum Threshold Prediction: Dynamic threshold prediction can further optimize sparsity but is not mandatory for XAttention.

3. Experiments

XAttention is evaluated across language, video understanding, and video generation benchmarks, where it generally preserves accuracy while reducing attention computation. Its efficiency comes from antidiagonal-based block selection, dynamic thresholds, and high sparsity, with warmup needed for stable video generation.

  • Cross-domain evaluation: XAttention achieves strong accuracy across RULER, LongBench, VideoMME, and VBench while substantially reducing attention computation.The reported maximum prefill attention speedup is 13.5×.
  • Language understanding: On RULER, XAttention remains robust at very long contexts, surpassing FlexPrefill and outperforming full attention at several sequence lengths.MInference and SeerAttention degrade more as context length increases.
  • Video understanding: On VideoMME, XAttention achieves the best average score among sparse methods and outperforms FlashAttention on long videos processed at 1 frame per second for up to 1 hour.The configuration uses stride S = 16 and threshold τ = 0.9.
  • Video generation: On VBench, both τ = 0.90 and τ = 0.95 achieve over 50% sparsity with high fidelity, while higher τ improves fidelity at slightly lower sparsity.The reported best values are PSNR 23.5, SSIM 0.822, and LPIPS 0.155.
  • Efficiency and ablations: At 256k context length, XAttention reaches 13.5× and 9.8× prefill attention speedups with densities of 7.32% and 6.89%, respectively.Its antidiagonal pattern also makes pattern selection up to 24.9× and 5.9× faster than MInference and FlexPrefill, respectively.

4. Related Work

Prior work accelerates long-context inference through sparse attention patterns, token selection, training, and system-level optimization, but practical efficiency remains constrained by selection overhead and training costs.

  • Sparse attention: Attention sparsity offers optimization opportunities, but its dependence on inputs, heads, and layers makes robust exploitation challenging.The passage notes that sparsity becomes more pronounced with longer contexts while remaining dynamic and input-dependent.
  • Sparse attention: Long-context inference methods use local or block-based attention, token discarding, or retention policies to reduce computation or stabilize latency and memory usage.Examples include Sparse Transformer, LongFormer, BigBird, Selective Attention, H2O, TOVA, and StreamingLLM.
  • Prefill acceleration: Prefill methods such as MInference and FlexPrefill achieve speedups, yet their pattern-selection overhead remains a bottleneck.SeerAttention avoids this overhead through trained gate parameters but requires costly training and has limited downstream performance.
  • System-level methods: System-level approaches improve attention efficiency through hardware-aware memory access or distributed computation across devices.FlashAttention optimizes memory access patterns, while RingAttention distributes attention computation across multiple devices.
  • Recent directions: Recent sparse-attention research also targets video generation and trainable long-context mechanisms without relying on strongly biased attention structures.Sparse VideoGen uses spatial and temporal heads, while NSA and MoBA pursue alternative sparse-attention designs.

5. Conclusion

XAttention is a plug-and-play framework that uses antidiagonal attention sums to identify and prune non-essential blocks. Across language, video understanding, and video generation benchmarks, it maintains performance comparable to full attention while achieving up to 13.5x attention-computation speedup.

  • Contribution: XAttention uses antidiagonal sums as a robust proxy for block importance to identify and prune non-essential attention blocks.The framework is designed for accelerating long-context inference in Transformer models.
  • Results: 13.5x speedup in attention computation is achieved while maintaining performance comparable to full attention across language, video understanding, and video generation benchmarks.Evaluations cover RULER and LongBench, VideoMME, and VBench.
  • Implication: XAttention is presented as a path toward efficient and scalable deployment of Long-Context Transformer Models in real-world applications.The conclusion frames this as unlocking the practical potential of block-sparse attention.
Loading 2503.16428v1…