XAttention: Блочно-разреженное внимание с антидиагональным подсчетом

XAttention: Block Sparse Attention with Antidiagonal Scoring

March 20, 2025

Авторы: Ruyi Xu, Guangxuan Xiao, Haofeng Huang, Junxian Guo, Song Han

cs.AI

Аннотация

Трансформеры с длинным контекстом (LCTMs) играют ключевую роль в реальных приложениях, но сталкиваются с высокими вычислительными затратами из-за квадратичной сложности механизма внимания. Разреженное блочное внимание смягчает эту проблему, концентрируя вычисления на критически важных областях, однако существующие методы испытывают трудности с балансом между точностью и эффективностью из-за дорогостоящих измерений важности блоков. В данной статье мы представляем 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.