DiTFastAttn: Compressão de Atenção para Modelos de Transformadores de Difusão

DiTFastAttn: Attention Compression for Diffusion Transformer Models

June 12, 2024

Autores: Zhihang Yuan, Pu Lu, Hanling Zhang, Xuefei Ning, Linfeng Zhang, Tianchen Zhao, Shengen Yan, Guohao Dai, Yu Wang

cs.AI

Resumo

Os Transformadores de Difusão (DiT) se destacam na geração de imagens e vídeos, mas enfrentam desafios computacionais devido à complexidade quadrática da autoatenção. Propomos o DiTFastAttn, um novo método de compressão pós-treinamento para aliviar o gargalo computacional dos DiT. Identificamos três redundâncias principais no cálculo da atenção durante a inferência dos DiT: 1. redundância espacial, onde muitos cabeçotes de atenção focam em informações locais; 2. redundância temporal, com alta similaridade entre as saídas de atenção de passos vizinhos; 3. redundância condicional, onde as inferências condicionais e incondicionais exibem similaridade significativa. Para abordar essas redundâncias, propomos três técnicas: 1. Atenção por Janelas com Cache Residual para reduzir a redundância espacial; 2. Redução de Similaridade Temporal para explorar a similaridade entre os passos; 3. Eliminação de Redundância Condicional para pular cálculos redundantes durante a geração condicional. Para demonstrar a eficácia do DiTFastAttn, aplicamos o método ao DiT, ao PixArt-Sigma para tarefas de geração de imagens e ao OpenSora para tarefas de geração de vídeos. Os resultados da avaliação mostram que, para a geração de imagens, nosso método reduz até 88% dos FLOPs e alcança uma aceleração de até 1,6x na geração de alta resolução.

English

Diffusion Transformers (DiT) excel at image and video generation but face computational challenges due to self-attention's quadratic complexity. We propose DiTFastAttn, a novel post-training compression method to alleviate DiT's computational bottleneck. We identify three key redundancies in the attention computation during DiT inference: 1. spatial redundancy, where many attention heads focus on local information; 2. temporal redundancy, with high similarity between neighboring steps' attention outputs; 3. conditional redundancy, where conditional and unconditional inferences exhibit significant similarity. To tackle these redundancies, we propose three techniques: 1. Window Attention with Residual Caching to reduce spatial redundancy; 2. Temporal Similarity Reduction to exploit the similarity between steps; 3. Conditional Redundancy Elimination to skip redundant computations during conditional generation. To demonstrate the effectiveness of DiTFastAttn, we apply it to DiT, PixArt-Sigma for image generation tasks, and OpenSora for video generation tasks. Evaluation results show that for image generation, our method reduces up to 88\% of the FLOPs and achieves up to 1.6x speedup at high resolution generation.