PISA: Поэлементно-разреженное внимание — более эффективный подход для диффузионных трансформаторов

PISA: Piecewise Sparse Attention Is Wiser for Efficient Diffusion Transformers

February 1, 2026

Авторы: Haopeng Li, Shitong Shao, Wenliang Zhong, Zikai Zhou, Lichen Bai, Hui Xiong, Zeke Xie

cs.AI

Аннотация

Диффузионные трансформеры являются основой для генерации видео и изображений, но их эффективность ограничена квадратичной сложностью механизма внимания. Хотя блочно-разреженное внимание ускоряет вычисления, обрабатывая только ключевые блоки ключ-значение, оно страдает от деградации качества при высокой степени разреженности из-за потери контекста. В данной работе мы обнаружили, что оценки внимания некритических блоков демонстрируют распределительную стабильность, что позволяет точно и эффективно аппроксимировать их вместо полного отбрасывания, что принципиально важно для проектирования разреженного внимания. Руководствуясь этим ключевым наблюдением, мы предлагаем PISA — не требующее дообучения кусочно-разреженное внимание, которое охватывает полный диапазон внимания с субквадратичной сложностью. В отличие от традиционной парадигмы «сохранить-или-отбросить», которая напрямую удаляет информацию некритических блоков, PISA вводит новую стратегию «точно-или-аппроксимировать»: она сохраняет точные вычисления для критических блоков, одновременно эффективно аппроксимируя остальные с помощью покомпонентного разложения Тейлора. Такой подход позволяет PISA служить точным приближением полного внимания, эффективно устраняя разрыв между скоростью и качеством. Результаты экспериментов показывают, что PISA обеспечивает ускорение в 1.91 и 2.57 раза на моделях Wan2.1-14B и Hunyuan-Video соответственно, при этом стабильно сохраняя наивысшее качество среди методов разреженного внимания. Примечательно, что даже для генерации изображений на FLUX PISA достигает ускорения в 1.2 раза без ущерба для визуального качества. Код доступен по адресу: https://github.com/xie-lab-ml/piecewise-sparse-attention.

English

Diffusion Transformers are fundamental for video and image generation, but their efficiency is bottlenecked by the quadratic complexity of attention. While block sparse attention accelerates computation by attending only critical key-value blocks, it suffers from degradation at high sparsity by discarding context. In this work, we discover that attention scores of non-critical blocks exhibit distributional stability, allowing them to be approximated accurately and efficiently rather than discarded, which is essentially important for sparse attention design. Motivated by this key insight, we propose PISA, a training-free Piecewise Sparse Attention that covers the full attention span with sub-quadratic complexity. Unlike the conventional keep-or-drop paradigm that directly drop the non-critical block information, PISA introduces a novel exact-or-approximate strategy: it maintains exact computation for critical blocks while efficiently approximating the remainder through block-wise Taylor expansion. This design allows PISA to serve as a faithful proxy to full attention, effectively bridging the gap between speed and quality. Experimental results demonstrate that PISA achieves 1.91 times and 2.57 times speedups on Wan2.1-14B and Hunyuan-Video, respectively, while consistently maintaining the highest quality among sparse attention methods. Notably, even for image generation on FLUX, PISA achieves a 1.2 times acceleration without compromising visual quality. Code is available at: https://github.com/xie-lab-ml/piecewise-sparse-attention.