DiTFastAttn：拡散Transformerモデルのための注意機構圧縮

DiTFastAttn: Attention Compression for Diffusion Transformer Models

June 12, 2024

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

cs.AI

要旨

Diffusion Transformers (DiT)は画像および動画生成において優れた性能を発揮しますが、セルフアテンションの二次的な計算複雑性により計算上の課題に直面しています。本研究では、DiTの計算ボトルネックを軽減するための新しいポストトレーニング圧縮手法であるDiTFastAttnを提案します。DiTの推論過程におけるアテンション計算には、以下の3つの主要な冗長性が存在することを特定しました：1. 空間的冗長性 - 多くのアテンションヘッドが局所的な情報に集中している、2. 時間的冗長性 - 隣接するステップ間のアテンション出力が高い類似性を示す、3. 条件的冗長性 - 条件付き推論と無条件推論の間に顕著な類似性が存在する。これらの冗長性に対処するため、以下の3つの技術を提案します：1. 空間的冗長性を低減するためのWindow Attention with Residual Caching、2. ステップ間の類似性を活用するTemporal Similarity Reduction、3. 条件付き生成における冗長な計算をスキップするConditional Redundancy Elimination。DiTFastAttnの有効性を実証するため、画像生成タスクにおけるDiTとPixArt-Sigma、および動画生成タスクにおけるOpenSoraに適用しました。評価結果は、画像生成において最大88%のFLOPsを削減し、高解像度生成では最大1.6倍の高速化を達成したことを示しています。

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.