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(Piecewise Sparse Attention)를 제안합니다. 비중요 블록 정보를 직접 제거하는 기존의 유지-또는-삭제 패러다임과 달리, 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에서 확인할 수 있습니다.

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.