Scalable Diffusion Models with Transformers

Problem

Framing

U-Net backbones had driven diffusion progress, but scaling laws for transformer backbones in diffusion were unclear. DiT closes this gap by patchifying latent inputs and scaling a class-conditional transformer directly with compute, reaching ImageNet 256$\times$256 FID 2.27 and 512$\times$512 FID 3.04 with guidance.

Currently Used Methods

Foundational

• @DenoisingDiffusionProbabilisticModels2020 — Gaussian diffusion with learned reverse denoising objective.
  • Limitation in context: backbone scaling focused on convolutional U-Nets, not transformers.
• @nicholImprovedDDPM2021 — stronger DDPM training and faster sampling refinements.
  • Limitation in context: keeps the same U-Net-centered design space.
• @dhariwalDiffusionBeatGANs2021 — classifier guidance pushes diffusion to top ImageNet FID.
  • Limitation in context: high quality still relies on large convolutional U-Nets.
• @rombachLatentDiffusion2022 — latent diffusion cuts pixel-space cost with a VAE.
  • Limitation in context: uses latent U-Nets, not transformer backbones.
• @dosovitskiyViT2020 — patchified transformers scale cleanly for vision tasks.
  • Limitation in context: does not show they work inside diffusion denoisers.

Proposed Method

Architecture

DiT runs diffusion in Stable-Diffusion-style latent space. A latent of size $32 \times 32 \times 4$ is patchified into tokens, processed by $N$ transformer blocks at hidden size $d$, then projected back to predict noise and diagonal covariance. Conditioning on timestep $t$ and class label $y$ enters each block through adaLN-Zero; this beats cross-attention and token concatenation.

Loss / Objective

The model keeps the standard diffusion training target in latent space.

Sampling Rule / Algorithm

For class-conditional sampling, DiT uses classifier-free guidance.

Training Procedure

• Latent VAE downsample factor: $8$.
• Input latent shape at 256 resolution: $32 \times 32 \times 4$.
• Learning rate: $1 \times 10^{-4}$.
• Weight decay: $0$.
• Batch size: $256$.
• Sampling for reported FID-50K: $250$ DDPM steps.
• Largest 256 and 512 DiT-XL/2 runs continue to $7$M steps.

Evaluation

Datasets

• ImageNet class-conditional, 256$\times$256.
• ImageNet class-conditional, 512$\times$512.

Metrics

• FID-50K.
• sFID.
• Inception Score.
• Precision.
• Recall.

Headline results

• ImageNet 256$\times$256 (DiT-XL/2-G): FID 2.27.
• ImageNet 512$\times$512 (DiT-XL/2-G): FID 3.04.
• ImageNet 256$\times$256: DiT-XL/2-G beats ADM-U-G and LDM variants in Figure 2.
• Scaling from DiT-S to DiT-XL monotonically improves FID at 400K steps.

Ablations

• Conditioning block: adaLN-Zero gives the best FID among tested conditioning schemes.
• Model size: larger backbones improve FID across training.
• Patch size: smaller patches improve quality by increasing token count and Gflops.
• Training compute: larger models are more compute-efficient at a fixed target FID.

Method Strengths and Weaknesses

Strengths

• Clean result: transformer Gflops strongly predict diffusion quality.
• Strong headline numbers: FID 2.27 at 256 and 3.04 at 512.
• Simple swap: replaces U-Net with a standard ViT-style backbone.
• adaLN-Zero improves conditioning without cross-attention overhead.

Weaknesses

• Reported wins depend on classifier-free guidance at sampling time.
• Best quality needs very large DiT-XL/2 compute budgets.
• Evaluation centers on ImageNet; cross-domain robustness is untested.
• Latent VAE quality constrains the final sample ceiling.

Suggestions from the authors

• Study larger transformer backbones and longer training regimes.
• Test DiT scaling beyond ImageNet class-conditional generation.
• Improve latent-space diffusion components beyond the fixed VAE.
• Explore transformer design choices specialized for diffusion.

Links

Prior Papers

• @DenoisingDiffusionProbabilisticModels2020 — establishes the base diffusion objective DiT keeps while changing the backbone.
• @nicholImprovedDDPM2021 — supplies the improved DDPM setup and faster sampling conventions used for comparison.
• @dhariwalDiffusionBeatGANs2021 — the main U-Net diffusion baseline that DiT surpasses on ImageNet.
• @rombachLatentDiffusion2022 — provides the latent-diffusion formulation and VAE compression DiT adopts.
• @dosovitskiyViT2020 — motivates patchified transformer scaling, which DiT imports into diffusion.

Further Papers

• @karrasEDM2022 — studies improved diffusion design choices and remains a strong comparison point for generative quality.
• @albergoStochasticInterpolants2023 — extends continuous-time generative transport methods that overlap with diffusion-style modeling.
• @songConsistencyModels2023 — targets faster diffusion-family sampling, a direct downstream concern for large DiT models.