Efficient Scaling of Diffusion Transformers for Text-to-Image Generation

TL;DR

This paper investigates the scaling behavior of Diffusion Transformers for text-to-image generation, demonstrating that a pure self-attention model scales effectively and outperforms some existing models when scaled properly.

Contribution

It introduces the U-ViT model, a self-attention based DiT that scales better and is simpler than cross-attention variants, and provides empirical insights on data scaling effects.

Findings

U-ViT outperforms SDXL UNet and other DiT variants at 2.3B parameters.

Scaling dataset size and caption quality improves performance and learning efficiency.

Pure self-attention DiT models are more effective and scalable than cross-attention variants.

Abstract

We empirically study the scaling properties of various Diffusion Transformers (DiTs) for text-to-image generation by performing extensive and rigorous ablations, including training scaled DiTs ranging from 0.3B upto 8B parameters on datasets up to 600M images. We find that U-ViT, a pure self-attention based DiT model provides a simpler design and scales more effectively in comparison with cross-attention based DiT variants, which allows straightforward expansion for extra conditions and other modalities. We identify a 2.3B U-ViT model can get better performance than SDXL UNet and other DiT variants in controlled setting. On the data scaling side, we investigate how increasing dataset size and enhanced long caption improve the text-image alignment performance and the learning efficiency.