PISA: Atenção Esparsa por Partes é Mais Sábia para Transformadores de Difusão Eficientes
PISA: Piecewise Sparse Attention Is Wiser for Efficient Diffusion Transformers
February 1, 2026
Autores: Haopeng Li, Shitong Shao, Wenliang Zhong, Zikai Zhou, Lichen Bai, Hui Xiong, Zeke Xie
cs.AI
Resumo
Os Transformadores de Difusão são fundamentais para a geração de vídeos e imagens, mas sua eficiência é limitada pela complexidade quadrática do mecanismo de atenção. Embora a atenção esparsa por blocos acelere o cálculo ao atender apenas aos blocos chave-valor críticos, ela sofre degradação em alta esparsidade ao descartar contexto. Neste trabalho, descobrimos que os escores de atenção de blocos não críticos exibem estabilidade distribucional, permitindo que sejam aproximados de forma precisa e eficiente em vez de descartados, o que é essencial para o projeto de atenção esparsa. Motivados por esta percepção fundamental, propomos o PISA, uma Atenção Esparsa por Partes sem necessidade de treinamento que cobre todo o intervalo de atenção com complexidade subquadrática. Diferente do paradigma convencional de manter-ou-descartar que elimina diretamente a informação de blocos não críticos, o PISA introduz uma nova estratégia de exato-ou-aproximado: mantém o cálculo exato para blocos críticos enquanto aproxima eficientemente o restante através de expansão de Taylor por blocos. Este projeto permite que o PISA sirva como um proxy fiel da atenção completa, efetivamente preenchendo a lacuna entre velocidade e qualidade. Resultados experimentais demonstram que o PISA alcança acelerações de 1,91 vezes e 2,57 vezes em Wan2.1-14B e Hunyuan-Video, respectivamente, mantendo consistentemente a mais alta qualidade entre os métodos de atenção esparsa. Notavelmente, mesmo para geração de imagens no FLUX, o PISA alcança uma aceleração de 1,2 vezes sem comprometer a qualidade visual. O código está disponível em: 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.