Scalable Diffusion Models with Transformers (DiT)

The problem

Diffusion models generate images by learning to reverse a gradual noising process: a network is trained to undo noise step by step, and sampling starts from pure noise and removes it incrementally. Most image diffusion systems of the time relied on a U-Net backbone — a convolutional network that first shrinks the spatial resolution of its input, then expands it back, using skip connections to carry fine detail from shallow to deep layers. Transformers work differently: they treat the input as a sequence of tokens and let attention mix information between all of them. The question the paper poses is whether the U-Net is actually necessary, and whether a transformer backbone would inherit the property that made transformers so effective in language modelling — performance that improves smoothly and predictably as the model is made larger and given more compute.^[1]

How the model works

The paper keeps the diffusion recipe and the latent-space training setup — the model denoises compressed image representations rather than raw pixels — and swaps out only the backbone. Instead of a U-Net, a transformer operates on latent patches: the latent representation is cut into a grid of patches, each patch is embedded as a token, and transformer blocks process the resulting sequence.^[1] Everything downstream of the backbone, the noisy-input handling and the noise prediction, stays the same, which is what makes the comparison clean.

To study scaling, the authors measure model cost as forward-pass complexity in Gflops, that is, the number of floating point operations a single forward pass requires, counted in billions. They then vary three knobs: how deep the transformer is, how wide its layers are, and how many input tokens it receives, which is set by the patch size.^[1]

Results

The paper reports two things. First, scaling works. Diffusion Transformers (DiTs) with higher Gflops — obtained through increased transformer depth, increased width, or a larger number of input tokens — consistently have lower FID.^[1] FID, the Fréchet Inception Distance, compares the distribution of generated images with that of real ones; lower is better. The relationship is not a lucky jump at one particular size but a steady trend across the models the authors trained, which is exactly the behaviour one wants if architecture size is to be treated as a dial.

Second, the largest models beat what came before. The largest DiT-XL/2 models outperform all prior diffusion models on the class-conditional ImageNet 512×512 and 256×256 benchmarks, reaching a state-of-the-art FID of 2.27 on the 256×256 benchmark.^[1]

The compute cost is stated concretely enough to be reproduced: DiT-XL/2, the most compute-intensive model in the family, trains at roughly 5.7 iterations per second on a TPU v3-256 pod with a global batch size of 256.^[1]

Why it matters

A diffusion model bundles several largely independent decisions: how noise is scheduled, what objective is optimised, and which network predicts the noise. This paper isolates the last of these and shows that a generic, well-understood sequence architecture can serve as the backbone, with quality tracking compute in a measurable way. That turns "which architecture should the denoiser use" from a matter of design taste into a question that can be answered by measuring Gflops against FID. It also means diffusion research can draw directly on the tooling and scaling experience built up around transformers, rather than on architectures specific to image generation. The stated training throughput gives other groups a concrete reference point for budgeting their own runs.^[1]

^[1]: arXiv, "Scalable Diffusion Models with Transformers (DiT)".