DiTFastAttn : Compression de l'attention pour les modèles de transformateurs de diffusion

DiTFastAttn: Attention Compression for Diffusion Transformer Models

June 12, 2024

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

cs.AI

Résumé

Les Transformers de Diffusion (DiT) excellent dans la génération d'images et de vidéos, mais rencontrent des défis computationnels en raison de la complexité quadratique de l'auto-attention. Nous proposons DiTFastAttn, une nouvelle méthode de compression post-entraînement pour atténuer le goulot d'étranglement computationnel des DiT. Nous identifions trois redondances clés dans le calcul de l'attention lors de l'inférence des DiT : 1. la redondance spatiale, où de nombreuses têtes d'attention se concentrent sur des informations locales ; 2. la redondance temporelle, avec une forte similarité entre les sorties d'attention des étapes voisines ; 3. la redondance conditionnelle, où les inférences conditionnelles et non conditionnelles présentent une similarité significative. Pour résoudre ces redondances, nous proposons trois techniques : 1. l'Attention par Fenêtrage avec Mise en Cache Résiduelle pour réduire la redondance spatiale ; 2. la Réduction de Similarité Temporelle pour exploiter la similarité entre les étapes ; 3. l'Élimination de la Redondance Conditionnelle pour sauter les calculs redondants lors de la génération conditionnelle. Pour démontrer l'efficacité de DiTFastAttn, nous l'appliquons à DiT, PixArt-Sigma pour les tâches de génération d'images, et à OpenSora pour les tâches de génération de vidéos. Les résultats d'évaluation montrent que pour la génération d'images, notre méthode réduit jusqu'à 88 % des FLOPs et permet une accélération jusqu'à 1,6x pour la génération en haute résolution.

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.