Scalable Diffusion Models with Transformers

In plain words

Image generators built on diffusion start from pure noise and remove it step by step until a picture appears. The network that does the removing has so far almost always been a U-Net, a convolutional design. This paper asks whether the U-Net matters much at all, and replaces it with a transformer, the same kind of network that powers language models and image classifiers.

The image is first squeezed into a small grid of numbers by a separate autoencoder. That grid is cut into square patches, each patch becomes a token, and the tokens go through a stack of transformer blocks. The block receives the noise level and the class label as a learned conditioning signal that sets how each block scales and gates its work. The authors call these models Diffusion Transformers (DiTs).

The main finding is about scaling. Making the transformer deeper, wider, or giving it more tokens (smaller patches) raises the compute per image, and the generated images get consistently better. The largest model in the paper beats the earlier diffusion models on the class-conditional ImageNet benchmarks reported here. The results are the authors' own.

TL;DR

• DiT replaces the U-Net backbone of latent diffusion with a ViT-style transformer on latent patches §3.
• Conditioning uses adaLN-Zero: timestep and class regress a per-block scale, shift and gate, initialised so each block starts as the identity §3.2 app. A.
• Across 12 configs (S/B/L/XL, patch 8/4/2), FID-50K at 400K steps falls as forward-pass Gflops rise §5 tab. 4.
• DiT-XL/2 with classifier-free guidance (cfg=1.50) reports FID 2.27 at 256x256 and 3.04 at 512x512 on class-conditional ImageNet tab. 2 tab. 3.

Problem

Diffusion models for images had converged on convolutional U-Nets as the backbone §1 §2. The authors argue that the U-Net's inductive bias is not what makes diffusion models work, and that a transformer could replace it while keeping the scaling properties seen in other domains §1. They also argue that parameter count is a weak proxy for image-model cost, because it ignores input resolution, so they measure cost in forward-pass Gflops §2. Their experiments are class-conditional ImageNet at 256x256 and 512x512 §4 §5.

Method

Overview. The model is a latent diffusion model §3.1. A pretrained, frozen VAE encodes an image to a latent z = E(x). A DiT learns the reverse diffusion process on z, and the VAE decoder maps the sampled latent back to pixels §3.1 §4. The DiT itself follows ViT practice as closely as possible §3.2.

Forward noising. The paper uses the standard Gaussian forward process §3.1. Its equations are inline and unnumbered, so they are cited by section.

q(x_t \mid x_0) = \mathcal{N}\big(x_t;\ \sqrt{\bar{\alpha}_t}\,x_0,\ (1-\bar{\alpha}_t)\,\mathbf{I}\big) \quad [§3.1]

x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon_t, \quad \epsilon_t \sim \mathcal{N}(0, \mathbf{I}) \quad [§3.1]

Training objective. The network predicts the noise epsilon_theta(x_t) and a diagonal covariance. The mean is trained with the simple noise MSE. The covariance is trained with the full variational bound, which the authors take from Nichol and Dhariwal §3.1. The bound's term affects only the covariance, since its mean input is detached gaussian_diffusion.py.

\mathcal{L}_{simple}(\theta) = \lVert \epsilon_\theta(x_t) - \epsilon_t \rVert_2^2 \quad [§3.1]

Classifier-free guidance. The model is trained with the class label randomly dropped to a learned null embedding. At sampling time, the conditional and unconditional predictions are combined with scale s, where s = 1 gives standard sampling §3.1.

\hat{\epsilon}_\theta(x_t, c) = \epsilon_\theta(x_t, \varnothing) + s\,\big(\epsilon_\theta(x_t, c) - \epsilon_\theta(x_t, \varnothing)\big) \quad [§3.1]

Patchify. A latent of shape I x I x C becomes a sequence of T = (I/p)^2 tokens of width d through a linear patch embedding. Fixed 2-D sine-cosine positional embeddings are added §3.2 fig. 4. Halving the patch size p quadruples T, and at least quadruples transformer Gflops. The parameter count hardly changes §3.2 fig. 4. The design space includes p = 2, 4 and 8 §3.2.

Conditioning and blocks. Timesteps and class labels are embedded separately, and the paper compares four ways to feed them to a transformer block §3.2 fig. 5:

1. In-context: the two embeddings are appended as extra tokens, then removed after the final block §3.2.
2. Cross-attention: a length-two sequence of the embeddings feeds an added cross-attention layer after self-attention. The paper reports roughly 15% extra Gflops §3.2.
3. adaLN: layer-norm scale and shift are regressed from the sum of the two embeddings, costing the fewest Gflops §3.2.
4. adaLN-Zero: adaLN plus a per-branch gate alpha, regressed the same way, applied before each residual connection. The gate's MLP is zero-initialised, so each block starts as the identity §3.2.

The paper picks adaLN-Zero for all later experiments §5. Fig. 5 shows that it beats cross-attention and in-context conditioning at every training stage fig. 5.

Decoder. After the last block, a final layer norm (adaptive when adaLN is used) and a linear layer map each token to p x p x 2C values, which are rearranged into the noise and covariance maps §3.2.

# DiT forward pass with adaLN-Zero blocks  [§3.2] [app. A]
# Input:  latents (B, C, H, W), timesteps (B,), labels (B,)
# Output: (B, 2C, H, W): noise and covariance predictions

def dit_forward(latents, timesteps, labels):
    # Patchify into T = (H/p)^2 tokens of width d, add sine-cosine positions
    tokens = patch_embed(latents) + sincos_positions   # (B, T, d)
    # Sum of the timestep embedding and the class embedding  [app. A]
    cond = timestep_mlp(sinusoid(timesteps)) + class_table[labels]  # (B, d)
    # Run N adaLN-Zero blocks, all conditioned on the same vector  [§3.2]
    for block in blocks:
        tokens = adaln_zero_block(tokens, cond)
    # Final adaptive layer norm, linear to p*p*2C values per token  [§3.2]
    tokens = final_layer(tokens, cond)                  # (B, T, p*p*2C)
    # Rearrange tokens back into a spatial grid  [§3.2]
    return unpatchify(tokens)                           # (B, 2C, H, W)

# One adaLN-Zero DiT block  [§3.2] [app. A]
# Input:  tokens (B, T, d), cond (B, d)
# Output: tokens (B, T, d)

def adaln_zero_block(tokens, cond):
    # Regress shift, scale and gate for attention and MLP from cond  [app. A]
    shift_a, scale_a, gate_a, shift_m, scale_m, gate_m = linear(silu(cond))
    # Attention branch, gated; the gate starts at zero  [§3.2]
    normed = modulate(layer_norm(tokens), shift_a, scale_a)
    tokens = tokens + gate_a * self_attention(normed)
    # MLP branch with tanh-approximated GELU, same modulation  [app. A]
    normed = modulate(layer_norm(tokens), shift_m, scale_m)
    tokens = tokens + gate_m * mlp(normed)
    return tokens

# One DiT training step on latents  [§3.1] [§4]
# Input:  images (B, 3, H, W), labels (B,)
# Output: scalar loss

def training_step(images, labels):
    # Encode pixels with the frozen VAE, then scale the latents  [§4]
    latents = latent_scale * vae_encode(images)        # (B, 4, H/8, W/8)
    # Draw one timestep per example, uniformly over the schedule  [§4]
    timesteps = uniform_int(0, 1000, size=B)
    # Forward noising with the linear beta schedule  [§3.1] [§4]
    noise = normal(size=latents.shape)
    noisy = noise_schedule.add_noise(latents, noise, timesteps)
    # Predict noise and covariance from the noisy latents  [§3.2]
    raw = dit(noisy, timesteps, labels)                # (B, 2C, H/8, W/8)
    pred_noise, cov_values = split(raw, channels=C)
    # Simple MSE on the noise; the mean is the only part it trains  [§3.1]
    loss_mse = mean((pred_noise - noise) ** 2)
    # Variational term trains the covariance only  [§3.1]
    loss_vb = vlb_term(stop_grad(pred_noise), cov_values)
    return loss_mse + loss_vb

# Class-conditional sampling with classifier-free guidance  [§3.1] [§4]
# Input:  labels (B,), guidance scale s, number of respaced steps
# Output: images (B, 3, H, W)

def sample(labels, s):
    # Start from Gaussian noise in latent space  [§3.1]
    latents = normal(size=(B, C, H/8, W/8))
    # Walk the respaced reverse process, 250 steps in the reported runs  [§4]
    for t in sampling_timesteps:
        # Predict with the real label and with the learned null label  [§3.1]
        eps_cond, cov = dit(latents, t, labels)
        eps_null, _ = dit(latents, t, null_label)
        # Guide only the first three latent channels  [app. A]
        eps = eps_null + s * (eps_cond - eps_null)     # channels 0 to 2 only
        # Ancestral DDPM step using the learned covariance  [§3.1] [§4]
        latents = ddpm_step(latents, eps, cov, t)
    # Decode with the VAE and undo the latent scale  [§4]
    return vae_decode(latents / latent_scale)

Key constants: 1000 diffusion steps, a linear beta schedule from 1e-4 to 2e-2, and latent scale 0.18215 in the code §4 gaussian_diffusion.py train.py. The paper gives the schedule but not the scale factor (see Pitfalls).

Hyperparameters

The table covers the training and sampling settings the paper states, plus the code values where the paper is silent. Entries marked "code only" are not in the paper. Entries marked "not stated" are absent from both the paper and the paper's own text.

Model configurations tab. 1 and all trained variants tab. 4. Gflops are forward-pass counts; parameter counts exclude the 84M-parameter VAE app. A.

Source: tab. 4. Values are FID-50K without guidance, as the table states.

Results

Author-reported. Setup: class-conditional ImageNet at 256x256 and 512x512 §4 §5. FID-50K with 250 DDPM steps and the ADM evaluation suite §4. Numbers with a "-G" suffix use classifier-free guidance at the stated scale §5.1.

Headline numbers §5.1 tab. 2 tab. 3:

The 2.55 figure for 2.35M steps is in the text. The text compares it with ADM only as "similar" training length, gives no baseline number, and does not state the guidance scale. Table 4's no-guidance FID at 2352K steps is 10.67.

Full ImageNet 256x256 table tab. 2:

Full ImageNet 512x512 table tab. 3:

The paper notes that its Precision and Recall at 512x512 use 1000 real samples, following ADM tab. 3.

Scaling §5 fig. 8: FID at 400K steps falls with forward-pass Gflops across the 12 configs, with the paper reporting a strong negative correlation but no coefficient. Models with similar Gflops reach similar FID even when their size differs (e.g. S/2 and B/4) §5. Under a fixed training-compute estimate (Gflops x batch x steps x 3), larger models are more efficient, and XL/4 falls behind XL/2 at about 1e10 Gflops of training compute §5 fig. 9.

Sampling compute vs model compute §5.2 fig. 10: L/2 with 1000 steps uses 80.7 Tflops per image, while XL/2 with 128 steps uses 15.2 Tflops. XL/2 still has the better FID-10K, 23.7 against 25.9 §5.2. The paper's conclusion is that extra sampling steps do not make up for less model compute §5.2.

Compute of baselines tab. 6: ADM 256x256 is 1120 Gflops; ADM-U 256x256 is 742 (110 base + 632 upsampler); LDM-4 256x256 is 104; LDM-8 256x256 is 57. DiT-XL/2 at 256x256 is 118.64 and at 512x512 is 524.60 tab. 4 §5.1.

VAE decoder ablation tab. 5 app. D (DiT-XL/2-G, cfg=1.5, 256x256):

Paper vs code

Official code: facebookresearch/DiT, linked from the authors' project page [WebFetch: wpeebles.com/DiT.html] and the paper §1. The repository is under Attribution-NonCommercial 4.0 (see code_license). Only paths and line numbers are cited here, and no code is quoted.

Checked against the code at code_commit:

• The paper trains and samples in JAX on TPU §4. The released code is PyTorch with DDP on CUDA GPUs train.py. The README reports PyTorch-trained DiT-XL/2 reproducing the JAX result at 400K steps (19.5 JAX vs 18.1 PyTorch) [README: "PyTorch Training Results"].
• Latent scaling: the code multiplies VAE latents by 0.18215 before training and divides by it before decoding train.py sample.py. The paper does not state this constant.
• Classifier-free guidance on channels: the code applies guidance only to the first three latent channels by default, as the paper's appendix states models.py app. A. The code comment says the usual approach applies it to all four.
• Conditioning: models.py implements only adaLN-Zero blocks models.py. The in-context, cross-attention and plain-adaLN variants from the paper's block study tab. 4 are not in the released model file, so reproducing that study needs new code.
• Class dropout: the code drops labels to a null index with probability 0.1 by default models.py models.py. The paper gives no rate.
• Training augmentation: the code center-crops (ADM's routine) before the horizontal flip train.py train.py. The paper names only horizontal flips §4.
• Covariance: the code trains the mean on MSE and the covariance on the variational term, with the bound term left unscaled by default gaussian_diffusion.py [code: diffusion/init.py]. This matches the paper's description §3.1.
• Schedule and EMA: linear betas from 1e-4 to 2e-2 over 1000 steps and EMA decay 0.9999 match the paper gaussian_diffusion.py train.py.
• Optimizer and initialisation: AdamW with lr 1e-4 and no weight decay, zero initialised final layer and adaLN gates train.py models.py. These match the paper §4 §3.2.
• Sampling: the default is 250 steps, guidance scale 4.0, and the mse VAE decoder sample.py sample.py. The paper reports final numbers with the ft-EMA decoder app. D. The README says its FID figure uses the mse decoder and no guidance [README: "PyTorch Training Results"].
• Training length: the code's default is 1400 epochs with batch 256 train.py. The paper reports 7M steps for 256x256 §5. The code does not state a step count.
• TF32: the code enables TF32 matmuls train.py. The paper does not state numeric precision.

Pitfalls

• Latent scale is not in the paper. Without the 0.18215 factor, the diffusion model sees latents on the wrong scale. Use the code's value train.py and mark it as code-sourced.
• Which VAE decoder. The paper's scaling plots use ft-MSE and its final tables use ft-EMA app. D. The released sampler defaults to ft-MSE sample.py. Compare FID only with the decoder stated for that run.
• Guidance and reported FID. The 2.55 FID at 2.35M steps is in the text with no guidance scale given §5.1. The no-guidance FID in Table 4 for the same step count is 10.67 tab. 4.
• Cost hides in the token count. At 512x512 with patch 2, the model has 1024 tokens and 524.6 Gflops per forward pass §5.1 tab. 4. Halving the patch size raises tokens fourfold and at least quadruples the transformer cost §3.2.
• Cross-attention overhead. The text says about 15% extra Gflops §3.2. Table 4 gives 137.62 against 118.64 for adaLN-Zero, about 16%.
• Gate initialisation. adaLN-Zero starts each block as the identity by zeroing the gate MLP §3.2 models.py. Plain adaLN does not, and the paper reports it converging worse fig. 5 §5.
• Modulation form. The code modulates with x * (1 + scale) + shift models.py. The paper says the adaLN parameters are regressed as gamma and beta §3.2. With the zero-initialised regression, the 1 + form leaves the input unchanged at the start of training. Without it, a zero scale would zero the input. The gate zero-init is a separate guarantee that the residual branch adds nothing at the start.
• Guidance acts on a subset of channels. Three-channel guidance with scale 1 + x roughly matches four-channel guidance with 1 + 3x/4 app. A. For example, scale 1.5 on three channels gives FID-50K 2.27, and scale 1.375 on four channels gives 2.20 app. A. The paper leaves the cause to future work app. A.
• Sampling step selection. The paper says 250 DDPM steps but not how they are spaced. The code uses a respacing helper sample.py whose rules this card did not read.
• Learned covariance during sampling. Sampling uses the learned covariance, not a fixed one §4 gaussian_diffusion.py. The sampler also runs with clip_denoised=False sample.py.
• Scaling claims. The paper says FID correlates strongly with Gflops but gives no coefficient §5. Read the claim as a trend across 12 models, not as a fitted law.
• Hardware and precision. The paper trained on TPU-v3 pods §4. The numeric precision used is not stated, and the released code enables TF32 train.py.
• Parameter counts. They exclude the VAE's 84M parameters app. A tab. 4.

Lineage

• Builds on: Latent diffusion models (Rombach et al., CVPR 2022), which supply the VAE and latent-space framing §3.1. Denoising diffusion probabilistic models (Ho et al., 2020) and the ADM architecture and training recipe (Dhariwal and Nichol, NeurIPS 2021), whose hyperparameters are kept §4. Vision Transformer (Dosovitskiy et al., 2020) for the backbone and positional embedding §3.2. Classifier-free guidance (Ho and Salimans, 2021) §3.1. Adaptive normalisation (FiLM) and zero-initialised residual branches in ResNets (Goyal et al., 2017) for the adaLN-Zero block §3.2.
• Compared against: BigGAN-deep, StyleGAN-XL, ADM, ADM-U, ADM-G, CDM, LDM-4 and LDM-8 (with and without guidance) tab. 2 tab. 3.
• Superseded by: not stated in this wiki.