DiTFastAttn: Attention Compression for Diffusion Transformer Models
Authors: name: Zhihang Yuan; name: Pu Lu; name: Hanling Zhang; name: Xuefei Ning; name: Linfeng Zhang; name: Tianchen Zhao; name: Shengen Yan; name: Guohao Dai; name: Yu Wang
Published: 2024-06-12 18:00:08

Abstract: 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.