DiTFastAttn: Aufmerksamkeitskompression für Diffusions-Transformer-Modelle

DiTFastAttn: Attention Compression for Diffusion Transformer Models

June 12, 2024

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

cs.AI

Zusammenfassung

Diffusions-Transformer (DiT) zeichnen sich bei der Bild- und Videogenerierung aus, stehen jedoch vor Rechenschwierigkeiten aufgrund der quadratischen Komplexität der Selbst-Aufmerksamkeit. Wir schlagen DiTFastAttn vor, eine neuartige Methode zur Nachbearbeitungskompression, um das Rechenengpass von DiT zu lindern. Wir identifizieren drei Schlüsselredundanzen bei der Aufmerksamkeitsberechnung während der DiT-Inferenz: 1. räumliche Redundanz, bei der viele Aufmerksamkeitsköpfe sich auf lokale Informationen konzentrieren; 2. zeitliche Redundanz, mit hoher Ähnlichkeit zwischen den Ausgaben der benachbarten Schritte der Aufmerksamkeit; 3. bedingte Redundanz, bei der bedingte und unbedingte Inferenzen eine signifikante Ähnlichkeit aufweisen. Um diesen Redundanzen zu begegnen, schlagen wir drei Techniken vor: 1. Fensteraufmerksamkeit mit Restcaching zur Reduzierung der räumlichen Redundanz; 2. Reduzierung der zeitlichen Ähnlichkeit zur Nutzung der Ähnlichkeit zwischen den Schritten; 3. Beseitigung der bedingten Redundanz zur Überspringung redundanter Berechnungen während der bedingten Generierung. Um die Wirksamkeit von DiTFastAttn zu demonstrieren, wenden wir es auf DiT, PixArt-Sigma für Bildgenerierungsaufgaben und OpenSora für Videogenerierungsaufgaben an. Die Evaluierungsergebnisse zeigen, dass unser Verfahren bei der Bildgenerierung bis zu 88\% der FLOPs reduziert und bei der Generierung mit hoher Auflösung bis zu 1,6-fache Beschleunigung erreicht.

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.