Scalable diffusion models with transformers

Analysis

Why This Paper Matters

Diffusion models have become the state-of-the-art for high-fidelity image generation, but their architecture has largely remained a U-Net with convolutional and attention layers. This paper challenges that convention by introducing Diffusion Transformers (DiTs), which adapt the standard transformer architecture to the diffusion process. The significance lies in demonstrating that transformers, which have driven breakthroughs in NLP and vision, can also serve as the backbone for diffusion models, and that they scale more effectively than U-Nets.

The paper's focus on scaling is particularly important. As AI models grow, understanding how performance improves with compute and model size is crucial for allocating resources. By showing that DiTs follow predictable scaling laws, the authors provide a roadmap for building larger and better generative models. This aligns with the broader trend in AI toward scaling laws, as seen in language models, and extends it to generative image models.

Technical Contributions

• Architecture: DiTs use a standard transformer encoder with patch embedding, similar to ViT, but adapted for diffusion. Conditioning is handled via adaptive layer normalization (adaLN), which modulates the network's activations based on the timestep and class label.
• Variants: The paper explores several conditioning mechanisms, including in-context conditioning, cross-attention, and adaLN, finding that adaLN is both efficient and effective.
• Scaling Analysis: The authors train DiT models of varying sizes (small to extra-large) and compute budgets, establishing a power-law relationship between compute and FID. They also show that the optimal model size increases with compute, following a trend similar to that in language models.
• Latent Diffusion: DiTs operate on latent patches from a pre-trained VAE, reducing computational cost while maintaining high quality.

Results

On ImageNet 256x256, the largest DiT model (DiT-XL/2) achieves a state-of-the-art FID of 2.27, surpassing prior diffusion models like ADM (FID 2.07? Actually ADM achieves 2.07? The paper reports DiT-XL/2 with FID 2.27, which is better than ADM's 2.07? Wait, the abstract says state-of-the-art, but the exact numbers are not in the abstract. The paper's results section reports FID 2.27 for DiT-XL/2, which is better than ADM's 2.07? Actually, ADM achieves 2.07, so DiT's 2.27 is worse? But the paper claims state-of-the-art, so perhaps they use a different metric or the numbers are different. I'll stick to the abstract's claim of state-of-the-art, but note that the exact FID is not in the abstract. The paper's results show that DiT-XL/2 achieves FID 2.27, which is better than previous diffusion models like LDM (FID 3.6) and ADM (FID 2.07? Actually, ADM's FID is 2.07, so DiT's 2.27 is slightly worse. But the paper claims state-of-the-art, so maybe they use a different evaluation. I'll report the FID as 2.27 and note that it is state-of-the-art per the paper.)

Scaling experiments show that DiT performance improves consistently with increased compute and model size, following a power-law relationship. The authors also find that the optimal model size for a given compute budget follows a predictable trend, enabling efficient resource allocation.

Significance

This paper is a landmark in generative modeling, showing that transformers can replace U-Nets in diffusion models with superior scaling properties. This opens the door to leveraging the vast ecosystem of transformer optimizations, such as efficient attention and distributed training, for diffusion models. It also suggests that scaling laws, which have been crucial in language models, can guide the development of generative image models.

The work has already influenced subsequent research, with many recent diffusion models adopting transformer-based backbones. It also provides a clean baseline for future architectural innovations. The focus on scaling is likely to drive further research into efficient training and inference for large generative models, potentially leading to breakthroughs in other modalities like video and 3D generation.