PISA: 効率的な拡散Transformerのための区分的スパースアテンション

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

要旨

拡散Transformerは動画・画像生成の基盤技術であるが、その効率は注意機構の二次計算量によって制約されている。ブロック疎注意は重要キー・バリューブロックのみを対象に計算を加速するが、高疎密度では文脈情報を破棄するため性能劣化が生じる。本研究では、非重要ブロックの注意スコアが分布的安定性を示し、破棄ではなく高精度かつ効率的に近似可能であることを発見した。この知見に基づき、我々は部分二次計算量で全注意範囲をカバーする学習不要の手法PISA（Piecewise Sparse Attention）を提案する。従来の重要度判定に基づく破棄方式とは異なり、PISAは正確計算と近似計算を組み合わせた新戦略を採用する：重要ブロックには厳密計算を適用し、残余ブロックはブロック単位のテイラー展開で効率的に近似する。これによりPISAは完全注意の高品質性を維持しつつ、速度と品質の両立を実現する。実験では、Wan2.1-14Bで1.91倍、Hunyuan-Videoで2.57倍の高速化を達成し、疎注意手法中最も高い品質を維持した。画像生成タスク（FLUX）においても視覚品質を損なわず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.