Core Concepts
Introducing U-shaped Diffusion Transformers (U-DiTs) that leverage downsampled self-attention to achieve state-of-the-art performance on latent-space image generation tasks while significantly reducing computational costs compared to isotropic Diffusion Transformers (DiTs).
Abstract
The paper investigates the potential of combining the U-Net architecture with Diffusion Transformers (DiTs) for latent-space image generation tasks. The authors first conduct a toy experiment by comparing a U-Net-style DiT (DiT-UNet) with an isotropic DiT, finding that the U-Net inductive bias is not fully leveraged when simply combining U-Nets and transformer blocks.
Inspired by the observation that the U-Net backbone features are low-frequency-dominated, the authors propose to downsample the tokens for self-attention in the U-Net-style DiT. This simple yet effective modification, called U-DiT, brings significant performance improvements despite a considerable reduction in computation costs.
The authors then scale up the U-DiT models and conduct extensive experiments, demonstrating the extraordinary performance and scalability of U-DiTs compared to DiTs and their improvements. Specifically, U-DiT-B can outperform the much larger DiT-XL/2 model with only 1/6 of its computation cost. The authors also show that U-DiTs consistently outperform their isotropic counterparts across different training iterations.
The key innovations of the U-DiT models include:
Downsampling the query-key-value tuple for self-attention in the U-Net-style DiT, which leverages the low-frequency dominance of the U-Net backbone features and reduces computation costs.
Incorporating additional techniques such as cosine similarity attention, RoPE2D, depthwise conv FFN, and re-parameterization to further improve the performance of U-DiTs.
Scaling up the U-DiT models to achieve state-of-the-art results on latent-space image generation tasks, outperforming isotropic DiTs by large margins.
The paper demonstrates the effectiveness of the proposed U-DiT models and highlights the potential of combining the U-Net architecture with Diffusion Transformers for efficient and high-quality latent-space image generation.
Stats
The FID-50K of U-DiT-B is 16.64, which is significantly better than the 42.84 of DiT-B/2 with similar computation costs.
The FID-50K of U-DiT-L is 10.08, outperforming the 20.05 of DiT-XL/2 with only 1/6 of its computation cost.
Quotes
"Amazingly, the proposed U-DiT model could perform better than DiT-XL/2 which is 6 times larger in terms of FLOPs."
"U-DiTs could outperform DiTs by large margins."