Reproduce “DiTFastAttn: Attention Compression for Diffusion Transformer Models” on OpenResearch

About this paper

Zhihang Yuan, Hanling Zhang, Pu Lu, Xuefei Ning, Linfeng Zhang, Tianchen Zhao, Shengen Yan, Guohao Dai and 1 more

Researchers from Tsinghua University and Infinigence AI developed DiTFastAttn, a post-training attention compression framework for Diffusion Transformer (DiT) models. The framework identifies and exploits spatial, temporal, and conditional redundancies inherent in DiT inference, achieving up to 1.8x end-to-end speedup for 2Kx2K image generation and reducing attention FLOPs by up to 76% while preserving generative quality.

Key insights

• Self-attention computations in Diffusion Transformer (DiT) models exhibit three specific types of redundancy during inference: spatial (local focus within attention heads), temporal (similarity of attention outputs across adjacent denoising steps), and conditional (similarity between conditional and unconditional attention outputs in Classifier-Free Guidance).
• Exploiting spatial redundancy with window attention is feasible without significant quality loss if crucial long-range dependencies are preserved, which can be achieved by caching and reusing a 'residual' from earlier full attention computations.
• The inherent gradual refinement process of diffusion models means attention patterns can be highly similar across neighboring timesteps and between conditional/unconditional CFG passes, allowing for direct reuse of computed attention outputs.

Summary by alphaXiv.