PISA: Piecewise Sparse Attention Is Wiser for Efficient Diffusion Transformers

2026-02-01

AI Summary

Diffusion Transformers suffer from efficiency bottlenecks due to the quadratic complexity of attention mechanisms, and existing block-sparse approaches often degrade performance at high sparsity levels by discarding non-critical context. This work proposes PISA, a training-free segmented sparse attention method that transcends the conventional “keep-or-discard” paradigm by introducing a novel “exact-or-approximate” strategy: critical attention blocks are computed exactly, while non-critical blocks are efficiently approximated via block-level Taylor expansion, leveraging the stability of their attention score distributions. This approach achieves sub-quadratic complexity while preserving full attention coverage. PISA accelerates inference by 1.91× and 2.57× on Wan2.1-14B and Hunyuan-Video, respectively, outperforming all existing sparse methods in generation quality, and achieves a 1.2× speedup in FLUX image generation without perceptible visual degradation.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

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.

Problem

• Diffusion Transformers
• attention efficiency
• sparse attention
• quadratic complexity
• context degradation

Innovation

• Piecewise Sparse Attention
• Distributional Stability
• Block-wise Taylor Expansion
• Training-free Approximation
• Sub-quadratic Complexity