DiTFastAttn: Compressione dell'Attenzione per Modelli di Trasformatori a Diffusione

DiTFastAttn: Attention Compression for Diffusion Transformer Models

June 12, 2024

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

cs.AI

Abstract

I Diffusion Transformers (DiT) eccellono nella generazione di immagini e video, ma affrontano sfide computazionali a causa della complessità quadratica del self-attention. Proponiamo DiTFastAttn, un innovativo metodo di compressione post-addestramento per alleviare il collo di bottiglia computazionale dei DiT. Identifichiamo tre ridondanze chiave nel calcolo dell'attenzione durante l'inferenza dei DiT: 1. ridondanza spaziale, dove molti head di attenzione si concentrano su informazioni locali; 2. ridondanza temporale, con un'elevata somiglianza tra gli output di attenzione di passi vicini; 3. ridondanza condizionale, dove le inferenze condizionate e non condizionate mostrano una significativa somiglianza. Per affrontare queste ridondanze, proponiamo tre tecniche: 1. Window Attention con Residual Caching per ridurre la ridondanza spaziale; 2. Riduzione della Somiglianza Temporale per sfruttare la somiglianza tra i passi; 3. Eliminazione della Ridondanza Condizionale per saltare calcoli ridondanti durante la generazione condizionata. Per dimostrare l'efficacia di DiTFastAttn, lo applichiamo a DiT, PixArt-Sigma per compiti di generazione di immagini e OpenSora per compiti di generazione di video. I risultati di valutazione mostrano che, per la generazione di immagini, il nostro metodo riduce fino all'88\% delle FLOP e raggiunge un accelerazione fino a 1.6x nella generazione ad alta risoluzione.

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.