PISA通过分段近似提升扩散Transformer效率,兼顾速度与质量。
PISA: Piecewise Sparse Attention Is Wiser for Efficient Diffusion Transformers
- 用分段泰勒展开近似非关键块,替代传统丢弃策略。
- 在Wan2.1-14B上提速1.91倍,Hunyuan-Video上提速2.57倍。
- 无需训练,适配图像与视频生成任务,性能领先现有稀疏注意力方法。
扩散Transformer是图像与视频生成的核心模型,但其注意力机制存在二次复杂度瓶颈。现有块稀疏注意力虽通过仅关注关键的键值块加速计算,但在高稀疏度下因丢弃上下文信息导致性能下降。本文发现非关键块的注意力得分具有分布稳定性,可被精确高效地近似而非直接丢弃,这对稀疏注意力设计至关重要。受此启发,我们提出PISA——一种无需训练的分段稀疏注意力机制,以次二次复杂度覆盖完整注意力范围。不同于传统的保留或丢弃范式,PISA采用精确或近似策略:对关键块保持精确计算,其余部分通过块级泰勒展开高效近似。该设计使PISA成为全注意力的忠实代理,有效弥合速度与质量之间的差距。实验表明,PISA在Wan2.1-14B和Hunyuan-Video上分别实现1.91倍和2.57倍加速,且始终优于其他稀疏注意力方法。值得注意的是,在FLUX图像生成任务中,PISA实现1.2倍加速且视觉质量无损。代码已公开于https://github.com/xie-lab-ml/piecewise-sparse-attention。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。