DiTFastAttn: Attention Compression for Diffusion Transformer Models
Diffusion Transformers (DiT) have emerged as a powerful tool for image and video generation tasks. However, their quadratic computational complexity due to the self-attention mechanism poses a significant challenge, particularly for high-resolution and long video tasks. This paper mitigate the computational bottleneck in DiT models by introducing a novel post-training model compression method. We identify three key redundancies in the attention computation during DiT inference and we propose three techniques.
- Window Attention with Residual Caching - Reduces spatial redundancy.
- Temporal Similarity Reduction - Exploit the similarity between steps.
- Conditional Redundancy Elimination - Skips redundant computations during conditional generation.
Generation Speed Comparasion
![DiTFastAttn demo](../assets/posts_images/DiTFastAttn_Demo_v3_complete.gif)
Image Generation Results
![pixart1k_result](../assets/posts_images/pixart1k_result.png)
![DiTFastAttn_overview](../assets/posts_images/pixart2k_result.png)
Video Generation Results
![raw](../assets/posts_images/raw.gif)
![D1](../assets/posts_images/D1.gif)
![D2](../assets/posts_images/D2.gif)
![D3](../assets/posts_images/D3.gif)
![D4](../assets/posts_images/D4.gif)
![D5](../assets/posts_images/D5.gif)
![D6](../assets/posts_images/D6.gif)
![raw](../assets/posts_images/raw2.gif)
![D1](../assets/posts_images/D1_2.gif)
![D2](../assets/posts_images/D2_2.gif)
![D3](../assets/posts_images/D3_2.gif)
![D4](../assets/posts_images/D4_2.gif)
![D5](../assets/posts_images/D5_2.gif)
![D6](../assets/posts_images/D6_2.gif)
FLOPs Reduction
![DiTFastAttn_overview](../assets/posts_images/DiTFastAttn_overview.jpg)
You can find the code for DiTFastAttn on GitHub at DiTFastAttention. Feel free to check out the repository for more details and to access the code.