DiTFastAttn: Aandachtscompressie voor Diffusion Transformer-modellen

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

Samenvatting

Diffusion Transformers (DiT) blinken uit in beeld- en videogeneratie, maar kampen met computationele uitdagingen vanwege de kwadratische complexiteit van self-attention. Wij stellen DiTFastAttn voor, een nieuwe post-trainingscompressiemethode om het computationele knelpunt van DiT te verlichten. We identificeren drie belangrijke redundanties in de aandachtberekening tijdens DiT-inferentie: 1. ruimtelijke redundantie, waarbij veel aandachtkoppen zich richten op lokale informatie; 2. temporele redundantie, met een hoge gelijkenis tussen de aandachtuitvoer van opeenvolgende stappen; 3. conditionele redundantie, waarbij conditionele en unconditionele inferenties aanzienlijke overeenkomsten vertonen. Om deze redundanties aan te pakken, stellen we drie technieken voor: 1. Window Attention met Residual Caching om ruimtelijke redundantie te verminderen; 2. Temporele Gelijkenisreductie om de gelijkenis tussen stappen te benutten; 3. Conditionele Redundantie-eliminatie om overbodige berekeningen tijdens conditionele generatie over te slaan. Om de effectiviteit van DiTFastAttn aan te tonen, passen we het toe op DiT, PixArt-Sigma voor beeldgeneratietaken en OpenSora voor videogeneratietaken. Evaluatieresultaten tonen aan dat onze methode voor beeldgeneratie tot 88\% van de FLOPs reduceert en tot 1,6x versnelling bereikt bij generatie op hoge resolutie.

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.