VerifiedarXiv:2212.0974829 min
Diffusion · Architecture

DiT: Scalable Diffusion Models with Transformers

Cut the noised latent into patches and run a plain transformer over them: across twelve model sizes, image quality follows the compute of a forward pass.

DiT replaces the convolutional U-Net inside a latent diffusion model with a Vision Transformer. The timestep and class label enter through layer-norm scales, shifts and gates whose weights start at zero, and the forward-pass Gflops, rather than the parameter count, predict image quality (FID).

Explaining the paperScalable Diffusion Models with TransformersPeebles, Xie · UC Berkeley / NYU · ICCV 2023 · arXiv:2212.09748 ↗

Twelve diffusion transformers from 0.36 to 118.6 Gflops, one training recipe for all of them, and an FID of 2.27 on class-conditional ImageNet 256×256, below every earlier diffusion model and below the best GAN.

Why replace the U-Net with a transformer

By 2022 transformers had taken over language modeling, image classification through the Vision Transformer, and autoregressive image generation. Image diffusion models had not switched. DDPM, ADM and latent diffusion (the model behind Stable Diffusion) all used a convolutional U-Net: ResNet blocks at several resolutions, skip connections between matching resolutions, and self-attention only at the coarse levels. Dhariwal and Nichol had tuned that U-Net's widths, attention heads and normalization for ADM, but its overall layout was the one Ho et al. adapted from PixelCNN++ in 2020.

Peebles and Xie test whether the convolutional structure matters. Their Diffusion Transformers (DiTs) keep ADM's diffusion training recipe and Stable Diffusion's latent space, and replace only the network in the middle with a ViT that reads patches of the noised latent. The paper reports two results. DiT-XL/2 reaches an FID of 2.27 on 256×256 ImageNet, lower than every prior diffusion model and lower than the GAN record of 2.30. And across twelve DiTs trained for the same number of steps, FID falls steadily as the forward-pass compute rises, whether the compute comes from a wider and deeper transformer or from cutting the latent into more tokens.

FID (Fréchet Inception distance) is the image-quality score used throughout: generate 50,000 images, embed them and real ImageNet images with a pretrained Inception network, fit a Gaussian to each set of embeddings, and measure the distance between the two Gaussians. Lower is better; 0 would mean the two sets of embeddings have the same mean and covariance. On the benchmark used here, 256×256 ImageNet, the best earlier models scored between 2.30 and 3.94; DiT's smallest model scores 153.6.

The paper measures model size in Gflops per forward pass, not in parameters, because the parameter count of a transformer does not depend on how many tokens it processes. The same DiT-XL weights can run on 16 tokens (patch size 8) or 256 tokens (patch size 2). The two runs differ by 16× in compute (7.39 against 118.64 Gflops) and by 5.5× in FID (106.41 against 19.47). The paper's Gflops count one fused multiply-add as one flop; a per-token cost model on that convention reproduces all twelve entries of its appendix table to within 0.6%, and the figures on this page use it.

The diffusion recipe DiT keeps

DiT changes nothing about the diffusion process, so this section collects the inherited recipe in one place. The forward process mixes a clean sample x0x_0 with Gaussian noise according to a fixed schedule αˉt\bar{\alpha}_t, which falls from 1 toward 0 as the timestep tt goes from 0 to 1000:

q(xt∣x0)=N ⁣(xt; αˉt x0, (1−αˉt)I),xt=αˉt x0+1−αˉt ϵt,ϵt∼N(0,I)q(x_t \mid x_0) = \mathcal{N}\!\left(x_t;\ \sqrt{\bar{\alpha}_t}\,x_0,\ (1-\bar{\alpha}_t)\mathbf{I}\right), \qquad x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon_t, \quad \epsilon_t \sim \mathcal{N}(\mathbf{0},\mathbf{I})q(xt∣x0)=N ⁣(xt; αˉt x0, (1−αˉt)I),xt=αˉt x0+1−αˉt ϵt,ϵt∼N(0,I)\begin{gathered} q(x_t \mid x_0) = \mathcal{N}\!\left(x_t;\ \sqrt{\bar{\alpha}_t}\,x_0,\ (1-\bar{\alpha}_t)\mathbf{I}\right), \\ x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon_t, \\ \epsilon_t \sim \mathcal{N}(\mathbf{0},\mathbf{I}) \end{gathered}
(1)

DiT uses ADM's linear schedule: 1000 steps, per-step noise variances βt\beta_t rising linearly from 10−410^{-4} to 0.020.02, and αˉt\bar{\alpha}_t the running product of 1−βs1-\beta_s up to step tt. At t=500t = 500 (indexed from 0, as the code samples it) the signal coefficient αˉt\sqrt{\bar{\alpha}_t} is 0.279 and the noise coefficient is 0.960, so the noise already dominates; by t=750t = 750 the signal coefficient is 0.057.

The model learns the reverse step pθ(xt−1∣xt)p_\theta(x_{t-1} \mid x_t), a Gaussian whose mean μθ\mu_\theta and covariance Σθ\Sigma_\theta a network predicts from xtx_t and tt. Training maximizes a variational lower bound on the log-likelihood of x0x_0, which splits into one KL divergence per step between the true posterior step q∗(xt−1∣xt,x0)q^*(x_{t-1} \mid x_t, x_0) and the model's step:

L(θ)=−log⁡pθ(x0∣x1)+∑tDKL ⁣(q∗(xt−1∣xt,x0) ∥ pθ(xt−1∣xt))\mathcal{L}(\theta) = -\log p_\theta(x_0 \mid x_1) + \sum_t \mathcal{D}_{KL}\!\left(q^*(x_{t-1} \mid x_t, x_0)\ \|\ p_\theta(x_{t-1} \mid x_t)\right)L(θ)=−log⁡pθ(x0∣x1)+∑tDKL ⁣(q∗(xt−1∣xt,x0) ∥ pθ(xt−1∣xt))\begin{aligned} \mathcal{L}(\theta) = &-\log p_\theta(x_0 \mid x_1) \\ &+ \sum_t \mathcal{D}_{KL}\!\left(q^*(x_{t-1} \mid x_t, x_0)\ \|\ p_\theta(x_{t-1} \mid x_t)\right) \end{aligned}
(2)

The paper prints the first term without its logarithm; (2) is the form in the DDPM paper the bound comes from. Both distributions inside each KL are Gaussian, so each term is a closed-form function of two means and two covariances. Ho et al. rewrote the mean as a prediction of the noise ϵt\epsilon_t that was added in (1); with that substitution and the per-step weights dropped, the mean part of (2) becomes a mean-squared error:

Lsimple(θ)=∥ϵθ(xt)−ϵt∥22\mathcal{L}_{\text{simple}}(\theta) = \left\lVert \epsilon_\theta(x_t) - \epsilon_t \right\rVert_2^2
(3)

The covariance is DiT's second output. Nichol and Dhariwal showed that learning Σθ\Sigma_\theta improves likelihood and lets the sampler use fewer steps. They parameterize it with one number vv per dimension that interpolates, in the log domain, between two closed-form choices: the forward-step variance βt\beta_t and the posterior variance β~t\tilde{\beta}_t.

Σθ(xt,t)=exp⁡ ⁣(vlog⁡βt+(1−v)log⁡β~t)\Sigma_\theta(x_t, t) = \exp\!\left(v \log \beta_t + (1 - v)\log \tilde{\beta}_t\right)
(4)

So the network outputs two tensors the shape of its input, ϵθ\epsilon_\theta and vv. The noise prediction is trained with (3) and vv with the full bound (2). In DiT's code, inherited from ADM, the bound is computed with the predicted mean detached, so its gradient reaches vv and not ϵθ\epsilon_\theta, and the two losses are added. Sampling starts from x1000∼N(0,I)x_{1000} \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) and draws xt−1x_{t-1} from the predicted Gaussian; every FID-50K number uses 250 such steps (the sampling-compute study in Section 5.2 varies the step count from 16 to 1000).

The last inherited piece is the latent space. Diffusion on 256×256×3 pixels is expensive, so latent diffusion first trains an autoencoder and runs diffusion on its codes. DiT uses Stable Diffusion's pretrained VAE with downsampling factor 8. The encoder EE maps a 256×256×3 image to a 32×32×4 latent, the code multiplies that latent by 0.18215 to give its channels roughly unit variance, and the decoder DD maps a sampled latent back to pixels. Both are frozen. Everything below operates on the 32×32×4 tensor, which holds 4,096 numbers against the image's 196,608.

How DiT turns a latent into tokens

A transformer reads a sequence of vectors, so DiT's first layer converts the noised latent of shape I×I×CI \times I \times C (32×32×4 at 256px) into tokens. Following ViT, it cuts the latent into non-overlapping p×pp \times p patches, flattens each patch into p2Cp^2 C numbers, and applies one shared linear layer that maps them to the hidden width dd. The sequence length is

T=(Ip)2T = \left(\frac{I}{p}\right)^2
(5)

With I=32I = 32 and p=2p = 2 that is 256 tokens, each built from 2·2·4 = 16 numbers and embedded to d=1152d = 1152 in the XL config. With p=8p = 8 it is 16 tokens of 256 numbers each. A fixed two-dimensional sine-cosine positional embedding is added to every token, half of its dimensions encoding the row and half the column, as in the masked autoencoder rather than the learned table of the original ViT. The paper tries p∈{2,4,8}p \in \{2, 4, 8\}.

Every transformer block does two kinds of arithmetic, and (5) sets both. The attention projections and the MLP act on each token separately: in multiply-adds, 4d24d^2 per token for the query, key, value and output projections and 8d28d^2 for the MLP with its 4× hidden width, so 12 d2 T12\,d^2\,T per block. The attention itself, QK⊤QK^\top and then the weighted sum of values, compares every token with every other token and costs 2 T2d2\,T^2 d per block. For DiT-XL/2, with 28 blocks:

28×(12⋅11522⋅256+2⋅2562⋅1152)≈114.2+4.2=118.4 Gflops28 \times \left(12 \cdot 1152^2 \cdot 256 + 2 \cdot 256^2 \cdot 1152\right) \approx 114.2 + 4.2 = 118.4\ \text{Gflops}28×(12⋅11522⋅256+2⋅2562⋅1152)≈114.2+4.2=118.4 Gflops\begin{aligned} &28 \times \left(12 \cdot 1152^2 \cdot 256 + 2 \cdot 256^2 \cdot 1152\right) \\ &\approx 114.2 + 4.2 = 118.4\ \text{Gflops} \end{aligned}

The conditioning layers, patch embedding and output head add about 0.2 more, and the paper reports 118.64. Halving pp quadruples TT, so the first term grows 4× and the second 16×; the paper says the total at least quadruples. The parameter count barely moves, because the only layers whose size depends on pp are the patch embedding (p2C×dp^2 C \times d) and the output head: DiT-XL/8 has 676M parameters and DiT-XL/2 has 675M. In Figure 1, move the patch-size slider and the config buttons and watch the token count, the two cost bars and the unchanged parameter readout.

Figure 1 · patchify and its cost
Drag the patch size from 8 to 2 and watch the token count go 16 → 64 → 256 while the parameter readout stays put. The dense bar grows with T, the attention bar with T². The 512px toggle uses the 64×64×4 latent, where XL/2 processes 1024 tokens at 524.6 Gflops. Each highlighted patch is one token: p·p·4 numbers in, d numbers out.

At p=2p = 2 the attention term is 10% of the total for DiT-S/2 (narrow width, so the d2d^2 term is small) and 3.6% for DiT-XL/2. On the paper's 32×32 latents the cost is therefore close to linear in TT. At 512px, with 1024 tokens, the XL/2 attention term reaches 68 Gflops of 524.6, still a minority. The T2T^2 term would dominate only at resolutions well past those the paper trains.

Four ways to feed in the timestep and class

A ViT block takes tokens and returns tokens. A diffusion network also needs the timestep tt, because the right amount of noise to remove depends on it, and a class-conditional one needs the class label. Both are first turned into vectors of width dd. The timestep goes through a 256-dimensional sinusoidal frequency embedding (the same sine-cosine table the original transformer used for positions, with 10,000 as the longest period) and a two-layer MLP with a SiLU between the layers, giving tembt_{\text{emb}}. The class is a row lookup in a table of 1001 learned embeddings, giving yemby_{\text{emb}}; the code names the label yy, and the last row is the null class used for guidance. The two vectors are added, not concatenated, into one condition vector of width dd:

cemb=temb+yembc_{\text{emb}} = t_{\text{emb}} + y_{\text{emb}}

The paper tries four places to put that vector:

The paper trains DiT-XL/2 with each block design and tracks FID over 400K steps. Figure 2 draws each design and carries its numbers. adaLN-Zero has the lowest FID at every checkpoint against cross-attention and in-context conditioning, and it ends 5.74 points below plain adaLN at 400K steps.

Figure 2 · where the condition enters
Click through the four designs. The amber vector is the summed timestep and class embedding; the arrows show what it feeds. The bars along the bottom are FID-50K at 400K steps for XL/2: in-context 35.24, cross-attention 26.14, adaLN 25.21, adaLN-Zero 19.47, with each design's Gflops and parameter count.

adaLN-Zero ties adaLN for the lowest compute of the four designs (118.6 Gflops) and reaches an FID 45% below in-context conditioning (19.47 against 35.24), which costs about the same. Its parameter count is the highest: 675M against 449M for in-context, because each block's modulation layer is a d×6dd \times 6d matrix, a third of the block's 18d218d^2 weights. Those weights run once per sample rather than once per token, so at T=256T = 256 they account for 0.2% of the block's compute.

adaLN-Zero: every block starts as the identity

adaLN-Zero differs from adaLN in two ways that act together: a gate on each residual branch, and an initialization that sets every gate to zero. Together they account for the 5.74-point gap. In the released code the modulation is

modulate(x,β,γ)=LN(x)⊙(1+γ)+β\text{modulate}(x, \beta, \gamma) = \text{LN}(x)\odot(1 + \gamma) + \beta
(6)

where LN\text{LN} is a LayerNorm with no learned affine parameters of its own and ⊙\odot is a per-dimension product, applied identically to every token. Writing the scale as 1+γ1 + \gamma means a regressed γ=0\gamma = 0 leaves the normalized activations unchanged. One block computes

(β1,γ1,α1,β2,γ2,α2)=W SiLU(cemb)+bx←x+α1⊙Attn ⁣(modulate(x,β1,γ1))x←x+α2⊙MLP ⁣(modulate(x,β2,γ2))\begin{aligned} (\beta_1, \gamma_1, \alpha_1, \beta_2, \gamma_2, \alpha_2) &= W\,\text{SiLU}(c_{\text{emb}}) + b \\ x &\leftarrow x + \alpha_1 \odot \text{Attn}\!\left(\text{modulate}(x, \beta_1, \gamma_1)\right) \\ x &\leftarrow x + \alpha_2 \odot \text{MLP}\!\left(\text{modulate}(x, \beta_2, \gamma_2)\right) \end{aligned}(β1,γ1,α1,β2,γ2,α2)=W SiLU(cemb)+bx←x+α1⊙Attn ⁣(modulate(x,β1,γ1))x←x+α2⊙MLP ⁣(modulate(x,β2,γ2))\begin{aligned} (\beta_1, \gamma_1, \alpha_1, &\beta_2, \gamma_2, \alpha_2) \\ &= W\,\text{SiLU}(c_{\text{emb}}) + b \\ x &\leftarrow x + \alpha_1 \odot \text{Attn}\!\left(\text{modulate}(x, \beta_1, \gamma_1)\right) \\ x &\leftarrow x + \alpha_2 \odot \text{MLP}\!\left(\text{modulate}(x, \beta_2, \gamma_2)\right) \end{aligned}
(7)

with W∈R6d×dW \in \mathbb{R}^{6d \times d} and b∈R6db \in \mathbb{R}^{6d}. The code sets both WW and bb to zero at initialization, so all six vectors start at zero for every condition: the scales and shifts leave the LayerNorm output as it is, and the gates multiply both branch outputs by zero. Each block returns its input, and the 28-block DiT-XL passes the patch embedding through to the final layer unchanged. The final layer's linear projection also starts at zero, so the network's first prediction is ϵ^=0\hat{\epsilon} = 0 and v=0v = 0 for every input. The paper describes the initialization as zeroing the gates α\alpha; zeroing the whole layer, as the code does, gives the same identity block.

# one adaLN-Zero block (models.py lines 19-20, 113-122)
def modulate(x, shift, scale):            # x: [B,T,d]; shift, scale: [B,d]
    return x * (1 + scale[:, None]) + shift[:, None]

def block(x, c):                           # c = t_emb + y_emb, shape [B,d]
    m = linear_6d(silu(c))                 # [B, 6d]; weight AND bias start at 0
    sh1, sc1, g1, sh2, sc2, g2 = m.chunk(6, dim=1)
    x = x + g1[:, None] * attn(modulate(ln1(x), sh1, sc1))
    x = x + g2[:, None] * mlp(modulate(ln2(x), sh2, sc2))
    return x                               # at init g1 = g2 = 0, so x is returned

The idea comes from residual networks. Goyal et al. (2017) set the scale γ\gamma of the last BatchNorm in every ResNet block to zero, so that at the start the signal passes only through the skip connections. On ImageNet that lowered top-1 error from 23.84% to 23.60% at batch size 256 and from 24.11% to 23.74% at 8k. ADM's U-Net does the same by zero-initializing the last convolution in each residual block. adaLN-Zero puts the zero in the conditioning path, so after training the gate is a function of tt and the class: a block can scale its attention output up at one noise level and down at another, which plain adaLN cannot do because it has no per-branch multiplier.

Figure 3 runs one six-dimensional token through (6) and (7) with real numbers. With the gate at zero the output column matches the input column exactly, whatever the condition slider does to the scale and shift; as the gate rises the branch starts to change the token.

Figure 3 · one token through an adaLN-Zero branch
Six values of one token flow left to right: input x, the layer-normed and modulated vector, the branch output f, the gated output α·f, and the sum. Drag α to 0 and the last column equals the first, the state every block is in at initialization. The condition slider changes the regressed scale and shift and reshapes the middle columns, but with α = 0 none of it reaches the output.

With a gate at zero, the gradient of the loss with respect to that gate is the upstream gradient times the branch output ff, summed over tokens, and ff is nonzero because the attention and MLP weights are Xavier-initialized. So the modulation layer moves at the first step. The attention and MLP weights get zero gradient at that step, because their contribution is multiplied by a zero gate, and start moving at the second.

One training step, with shapes

After the last block, the TT tokens have to become a noise prediction and a variance prediction the shape of the latent. The final layer applies one more adaptive LayerNorm (scale and shift only, 2d2d numbers from the condition, no gate) and a linear map from dd to p×p×2Cp \times p \times 2C numbers per token. For XL/2 that is 1152 → 32: a 2×2 patch with 8 channels. Unpatchify puts the patches back at their grid positions, giving an 8×32×32 tensor whose first four channels are ϵ^\hat{\epsilon} and last four are vv.

Figure 4 walks one training step stage by stage. Pick a config and patch size and every shape updates; move tt and the noise coefficients in stage 2 follow the schedule. The color of each row says what this step does to its weights.

Figure 4 · one training step, stage by stage
Press Play, drag the stage slider, or click a row. Teal rows hold weights that this step's gradient updates; grey rows are frozen (the VAE) or have no weights (noising, reshapes); amber rows are data and the loss. At t = 0 the latent is almost clean (coefficients 0.9999 and 0.010); at t = 999 the signal coefficient is 0.006.

The same step for DiT-XL/2 as code, batch size 1 so the shapes are readable:

# one training step of DiT-XL/2 on one 256x256 image (batch 1 for shapes)
z0  = vae.encode(img).sample() * 0.18215   # [4,32,32], VAE frozen, no grad
t   = randint(0, 1000)                     # say t = 500
eps = randn_like(z0)
z_t = 0.279 * z0 + 0.960 * eps             # sqrt(abar_500), sqrt(1 - abar_500)

x = patchify(z_t, p=2) @ W_embed + pos     # [256,16] -> [256,1152] + 2D sin-cos
c = t_mlp(sincos(t)) + y_table[label]      # [1152]; label 1000 = "no class"
for blk in blocks: x = blk(x, c)           # 28 blocks, each [256,1152] -> same
out = unpatchify(final(x, c))              # [256,32] -> [8,32,32]
eps_hat, v = out[:4], out[4:]              # noise guess, variance knob

loss = mean((eps_hat - eps) ** 2)          # L_simple, drives everything
loss += vlb_term(eps_hat.detach(), v, t)   # trains v only (mean detached)
loss.backward(); adamw.step()              # lr 1e-4, no weight decay
ema.update(model, decay=0.9999)            # the EMA weights are what gets evaluated

Every DiT weight receives a gradient from the mean-squared error (3): the patch embedding, the 28 blocks including their modulation layers, the timestep MLP, the class table and the final layer. The variance term adds gradient only along the paths that produce vv, because the noise prediction is detached before it enters. The VAE receives nothing. The class label is replaced by the null row for 10% of samples, which trains the unconditional prediction that guidance needs at sampling time.

The optimizer settings are ADM's, untuned: AdamW at a constant learning rate of 10−410^{-4} with no weight decay, batch size 256, horizontal flips as the only augmentation, no warmup, and an exponential moving average of the weights with decay 0.9999. The EMA copy is the model that gets evaluated. All twelve models use the same settings, and the paper reports none of the loss spikes common in transformer training. DiT-XL/2 runs 5.7 iterations per second on a TPU v3-256 pod, so 400K steps take about 19.5 hours and the 7M steps of the final model about two weeks. Counting the backward pass as twice the forward pass, 400K steps of XL/2 cost

118.64 Gflops×256×400,000×3≈3.6×1019 flops118.64\ \text{Gflops} \times 256 \times 400{,}000 \times 3 \approx 3.6 \times 10^{19}\ \text{flops}118.64 Gflops×256×400,000×3≈3.6×1019 flops\begin{aligned} &118.64\ \text{Gflops} \times 256 \times 400{,}000 \times 3 \\ &\approx 3.6 \times 10^{19}\ \text{flops} \end{aligned}

which is the training-compute unit the paper uses when it compares models by total training cost.

Why Gflops predict FID and parameters do not

The design space is patch size crossed with model config: p∈{8,4,2}p \in \{8, 4, 2\} and four configs, three borrowed from ViT and one larger. DiT-S has 12 blocks of width 384 with 6 heads, DiT-B 12 blocks of width 768 with 12 heads, DiT-L 24 blocks of width 1024 with 16 heads, and DiT-XL 28 blocks of width 1152 with 16 heads. That gives twelve models, from DiT-S/8 at 0.36 Gflops and 33M parameters to DiT-XL/2 at 118.64 Gflops and 675M. All twelve train for 400K steps with the recipe above, and the paper scores them with FID-50K, 250 sampling steps and no guidance.

Read Figure 5 on its default Gflops axis first. FID falls along one band from 153.6 down to 19.47, and it does not depend on how the Gflops were obtained: DiT-S/2 (6.06 Gflops, 33M parameters) and DiT-B/4 (5.56 Gflops, 130M parameters) reach 68.40 and 68.38. Then switch the axis to parameters. Each config becomes a vertical stack of three points with nearly identical weights and FIDs up to five times apart: DiT-XL at 675M parameters scores 106.41, 43.01 or 19.47 for patch size 8, 4 or 2.

Figure 5 · twelve models, two x axes
FID-50K after 400K steps, no guidance. Points are colored by config (grey S, violet B, amber L, teal XL) and joined along decreasing patch size. On the Gflops axis the dashed line is this page's least-squares fit of FID against log Gflops across all twelve. Switch to the parameter axis and each config collapses into a vertical stack. Hover a point for its numbers.

The training curves agree at every checkpoint: a larger config beats a smaller one at the same patch size, and a smaller patch beats a larger one at the same config. The paper also plots FID against total training compute (Gflops × batch × steps × 3) and finds the larger models more compute-efficient. Small DiTs trained longer fall behind large DiTs trained for fewer steps, and DiT-XL/2 overtakes DiT-XL/4 after about 101010^{10} Gflops of training compute even though each of its steps costs four times as much. The appendix repeats the analysis for sFID (FID computed on spatial features from an earlier Inception layer, more sensitive to layout), Inception Score (higher when a classifier labels each sample confidently and the labels cover many classes), precision and recall (both defined in the guidance section), and the training loss, and all of them improve with Gflops. Samples drawn from the twelve models with the same starting noise and label show a visible gain in fidelity as Gflops rise.

The paper fits no scaling law. It reports a strong negative correlation between Gflops and FID across twelve models at one training length, on one dataset and one resolution, and gives no exponent. The 400K-step FIDs are also far from converged: DiT-XL/2 goes from 19.47 at 400K steps to 9.62 at 7M steps with nothing else changed. For a broader treatment of compute-based scaling in language models, see the scaling laws explainer.

Sampling steps versus model compute

A diffusion model can spend more compute after training by taking more sampling steps, so a small DiT might catch a larger one by sampling longer. The paper tests this by evaluating all twelve models at 16, 32, 64, 128, 256 and 1000 sampling steps. The cost per image is the forward-pass Gflops times the step count (twice that with guidance, which this experiment does not use). DiT-L/2 at 1000 steps spends 80.7 Tflops per image. DiT-XL/2 at 128 steps spends 15.2 Tflops, 5.3× less, and still has the better FID-10K, 23.7 against 25.9. Figure 6 places the twelve models on a log-log chart of Gflops against step count, where a fixed per-image budget is a straight line.

Figure 6 · model compute versus sampling compute
15 T
Each column is one model at its forward-pass Gflops; the dots are the six step counts. Drag the budget and the dashed line moves; dots below it are affordable at that cost per image. The two circled points are the paper's comparison: at 5.3× the per-image compute, L/2 at 1000 steps still has the worse FID-10K.

On log axes the budget line has slope −1, so a model 1.5× more expensive per pass has to give up only a third of its steps to stay on budget, and past a few hundred steps extra DDPM steps change FID very little. The paper concludes that more sampling compute cannot make up for less model compute. The experiment is at 400K training steps and uses FID-10K (10,000 samples), which is noisier than the FID-50K used everywhere else in the paper.

Classifier-free guidance and the headline numbers

The headline numbers use classifier-free guidance. Training already replaced the class label with the null embedding ∅\varnothing 10% of the time, so one network can predict noise with or without the label. At sampling time both predictions are computed and the output is pushed further in the direction the label moved it:

ϵ^θ(xt,c)=ϵθ(xt,∅)+s⋅(ϵθ(xt,c)−ϵθ(xt,∅))\hat{\epsilon}_\theta(x_t, c) = \epsilon_\theta(x_t, \varnothing) + s\cdot\left(\epsilon_\theta(x_t, c) - \epsilon_\theta(x_t, \varnothing)\right)ϵ^θ(xt,c)=  ϵθ(xt,∅)+s⋅(ϵθ(xt,c)−ϵθ(xt,∅))\begin{aligned} \hat{\epsilon}_\theta(x_t, c) = \;&\epsilon_\theta(x_t, \varnothing) \\ &+ s\cdot\left(\epsilon_\theta(x_t, c) - \epsilon_\theta(x_t, \varnothing)\right) \end{aligned}
(8)

At one latent element, if the unconditional guess is 0.20 and the conditional guess is 0.50, then s=1s = 1 returns 0.50 (ordinary conditional sampling), s=1.5s = 1.5 returns 0.20 + 1.5·0.30 = 0.65, and s=0s = 0 returns the unconditional 0.20. In Ho and Salimans's notation the weight is w=s−1w = s - 1, so DiT's 1.5 is their 0.5. The reason (8) works: Bayes' rule gives the gradient of a classifier as

∇xlog⁡p(c∣x)=∇xlog⁡p(x∣c)−∇xlog⁡p(x)\nabla_x \log p(c \mid x) = \nabla_x \log p(x \mid c) - \nabla_x \log p(x)

and each score ∇xlog⁡p\nabla_x \log p at noise level tt equals −ϵ/1−αˉt-\epsilon / \sqrt{1-\bar{\alpha}_t}. The difference of the two noise predictions is therefore the classifier gradient times −1−αˉt-\sqrt{1-\bar{\alpha}_t}. Rearranged, (8) is the conditional prediction plus s−1s - 1 times that difference, so guidance adds a classifier gradient without training a classifier. Each guided step costs two forward passes, which the code runs as one call at twice the batch size.

Every guided number in the paper applies (8) to only the first three of the four latent channels, a detail stated only in the appendix and the code. The fourth channel keeps the conditional prediction unchanged. The code comment says this is kept "for exact reproducibility reasons". The appendix reports that three-channel guidance at scale 1+x1 + x behaves about like four-channel guidance at 1+34x1 + \tfrac{3}{4}x: three-channel 1.5 gives FID 2.27, four-channel 1.375 gives 2.20. Figure 7 shows both modes on one latent position, with the scales the paper measured marked underneath.

Figure 7 · guidance, channel by channel
Drag the scale s. In each channel the guided value (amber) is the unconditional guess (grey) plus s times the gap from grey to the conditional guess (teal); the white ticks on the amber bar mark the grey level and each whole multiple of that gap. In the code's default mode channel 3 keeps the teal value. The strip marks the scales with reported results for DiT-XL/2 at 256px; land on one to read its FID, precision and recall. The two guesses are illustrative values.

Across the three measured scales, precision rises from 0.67 to 0.76 to 0.83 and recall falls from 0.67 to 0.62 to 0.57. Precision is the fraction of generated images that land near some real image in Inception-feature space (fidelity); recall is the fraction of real images that have a generated image nearby (coverage). Guidance trades coverage for fidelity, and FID, which penalizes both, is lowest at 1.5 of the scales tested.

For the headline run the authors train DiT-XL/2 to 7M steps. On class-conditional ImageNet 256×256 it reaches FID 9.62 without guidance, 3.22 at s=1.25s = 1.25 and 2.27 at s=1.5s = 1.5, with sFID 4.60, Inception Score 278.24, precision 0.83 and recall 0.57. The prior diffusion best was LDM-4 with guidance 1.5 at 3.60; ADM with classifier guidance and its upsampler had 3.94; StyleGAN-XL, the best GAN, had 2.30. Trained only 2.35M steps, about as long as ADM, XL/2 already reaches 2.55. Its forward pass costs 118.6 Gflops against 103.6 for the LDM-4 U-Net in the same latent space, 742 for ADM-U and 1120 for pixel-space ADM. DiT-XL/2's recall is higher than LDM-4's and LDM-8's at every guidance scale tested (0.57 against LDM-4's 0.48 at 1.5), so at matched guidance its samples cover more of the real distribution.

At 512×512 the same architecture runs on the 64×64×4 latent: 1024 tokens, 524.6 Gflops per pass. Trained for 3M steps with the same hyperparameters, it reaches FID 3.04 at guidance 1.5, below the previous diffusion best of 3.85 (ADM-G with ADM-U), at about a quarter of ADM's per-pass compute (1983 Gflops) and under a fifth of ADM-U's (2813). StyleGAN-XL's 2.41 at this resolution is still lower, so at 512×512 the claim is the best diffusion model, short of the best generative model.

The VAE decoder can be swapped without retraining, because the three Stable Diffusion decoders share one encoder. With guidance 1.5 the original LDM decoder gives 2.46, the fine-tuned MSE decoder 2.30 and the fine-tuned EMA decoder 2.27, which is the one in the results tables; the scaling study uses the MSE decoder throughout. The sampling loop that produces those numbers, with the three-channel detail written out:

# guided sampling of one image of class 207 (sample_ddp.py + forward_with_cfg)
z = randn(4, 32, 32)                        # x_1000 ~ N(0, I)
for t in respaced_steps(250):               # 250 of the 1000 steps
    out = model(stack[z, z], t, [207, 1000])   # one call, batch 2: label, null
    e_c, e_0 = out[0, :3], out[1, :3]       # first 3 channels only
    eps = e_0 + 1.5 * (e_c - e_0)           # guidance at s = 1.5
    eps = cat[eps, out[0, 3:4]]             # channel 3 keeps the conditional guess
    z = p_sample(z, eps, v=out[0, 4:], t=t) # Gaussian step, learned variance
img = vae.decode(z / 0.18215)               # [3,256,256]

What the paper does not show

Every experiment is class-conditional ImageNet at 256 or 512 pixels in one frozen latent space, trained in JAX on TPU v3 pods. There is no text conditioning, no pixel-space DiT and no other dataset. The scaling claim rests on twelve models at one training length, with a correlation and no fitted law. FID moves with implementation details such as image resizing (Parmar, Zhang and Zhu, 2022), so all numbers come from ADM's TensorFlow evaluation suite run on exported samples; the sampling-compute comparison uses FID-10K rather than FID-50K.

The authors conclude that the U-Net's inductive bias is not needed for diffusion models to perform well, and propose scaling DiT further and using it as a drop-in backbone for text-to-image models such as DALL·E 2 and Stable Diffusion. Stable Diffusion 3 (Esser et al., 2024) did the second: its paper says its MM-DiT architecture builds upon DiT, with text tokens joined to the image tokens and the same modulated attention and MLP blocks.

Provenance Verified against primary literatureHow we verify
Peebles & Xie (2022)The paper: design space (patch size, block, config), Table 1 configs, appendix Table 4 Gflops/params/FID at 400K steps, Section 5 ablations, Section 5.2 sampling study, Tables 2-3 headline results, appendix note on guidance over three channels.
facebookresearch/DiT, models.pymodulate (lines 19-20), adaLN-Zero block (113-122), Xavier init and zero-init of the modulation and final layers (186-216), c = t + y (243), forward_with_cfg guiding model_out[:, :3] only (262).
train.py, sample_ddp.py, gaussian_diffusion.pyAdamW lr 1e-4 and weight decay 0 (train.py 155), EMA 0.9999 (40), latent scale 0.18215 (203); sampling defaults s = 1.5 and 250 steps with null label 1000 (sample_ddp.py); VLB term on the detached mean, added to the MSE (gaussian_diffusion.py 757-781); linear beta schedule 1e-4 to 0.02 (106-113).
Ho, Jain & Abbeel (2020)Forward process (1), the variational bound (2) with its −log p_θ(x₀|x₁) term, and the ε-prediction loss (3).
Nichol & Dhariwal (2021)Learned covariance (4) as a log-domain interpolation between β_t and β̃_t; hybrid objective with a stop-gradient on the mean.
Ho & Salimans (2022)Classifier-free guidance. DiT's scale s equals their w + 1, so DiT's best s = 1.5 is their w = 0.5.
Rombach et al. (2022)Latent diffusion with the f = 8 autoencoder: 256×256×3 → 32×32×4. LDM-4 is the prior best diffusion FID (3.60) DiT beats.
Dosovitskiy et al. (2021); He et al. (2022)ViT patchify and the S/B/L configs; MAE's fixed 2D sine-cosine positional embedding, which the DiT code imports.
Goyal et al. (2017), Section 5.1 and Table 2bZero-γ in the last BatchNorm of each ResNet block: top-1 error 23.84% → 23.60% at batch 256 and 24.11% → 23.74% at 8k.
Esser et al. (2024)Stable Diffusion 3, Section 4: "Our architecture builds upon the DiT (Peebles & Xie, 2023) architecture."
correctionThree points where the paper's text and its math or code disagree. (a) Its variational bound prints the first term as −p(x₀|x₁); the DDPM bound it takes the term from has −log p_θ(x₀|x₁), which (2) restores. (b) Its guidance derivation writes the guidance term as s·∇ log p(x|c) and folds the score-to-ε factor into a proportionality sign; the term Bayes' rule produces is the classifier gradient ∇ log p(c|x) = ∇ log p(x|c) − ∇ log p(x), and the final formula (8) is correct as printed. (c) The paper says adaLN-Zero initializes the modulation MLP "to output the zero-vector for all α"; the code zeroes the weight and bias of the whole modulation layer, so the scales and shifts also start at zero. Both versions make every block the identity at initialization. Two conventions to keep in mind when comparing numbers: the paper's Gflops count one fused multiply-add as one flop, and every guided number in the paper applies guidance to only three of the four latent channels.

Questions you might still have

?

Is this just a ViT with the timestep and class as extra tokens?
That design is the in-context variant, and the paper trained it: FID 35.24 at 400K steps against 19.47 for adaLN-Zero at the same Gflops. The block is a standard ViT block in every variant; the ablation changes only how the timestep and class reach it.

?

Does DiT need the VAE?
No. The paper says DiTs could run in pixel space without modification but never trains one; the latent space is a cost choice. A 256×256×3 image cut into 2×2 patches would be 16,384 tokens instead of 256, and the attention term grows with the square of the token count. The latent diffusion explainer covers the autoencoder.

?

Why does a smaller patch help if the parameter count is the same?
The weights are almost the same, but each forward pass does more work: halving p gives four times the tokens, so four times the projection and MLP compute and sixteen times the token-to-token attention compute. Each token also covers a 2×2 patch of latent cells instead of 8×8, so the output is predicted at a finer grain. The paper reads the FID gain as a compute effect because models with matched Gflops and different shapes land at matched FID (S/2 at 6.06 Gflops: 68.40; B/4 at 5.56 Gflops: 68.38).

?

Is a Gflop here a multiply or a multiply-add?
A multiply-add. A cost model of 12·d² multiply-adds per token per block plus 2·T²·d per block for attention reproduces all twelve entries of the paper's Table 4 to within 0.6%. Doubling every number for the two-flops-per-multiply-add convention leaves every comparison in the paper unchanged.

?

How does the guidance scale 1.5 relate to the guidance scale in Stable Diffusion?
Same convention: s = 1 means no guidance, as in the diffusers guidance_scale. In Ho and Salimans's notation w = s − 1, so DiT's best FID is at w = 0.5. Stable Diffusion's default of 7.5 is far above anything the paper measured, and the DiT repository's picture-making script defaults to 4.0, while its FID script uses 1.5.

?

If every block starts as the identity, does the first gradient step do anything?
Yes. The branch output f is nonzero because the attention and MLP weights are Xavier-initialized, so the derivative of the loss with respect to a gate (the upstream gradient times f, summed over tokens) is nonzero and the zero-initialized modulation layer moves first. The attention and MLP weights get zero gradient at step one because their output is multiplied by a zero gate; from step two the gates are nonzero and every weight trains.

?

Where did the architecture go after this paper?
Into most large image and video diffusion systems that followed. The Stable Diffusion 3 paper (2024) says its MM-DiT architecture builds upon DiT, and OpenAI's Sora report (2024) describes a diffusion model built on a transformer over spacetime patches of video latents. The DiT paper itself ends by proposing DiT as a drop-in backbone for text-to-image models.

Footnotes & further reading

  1. The paper: Peebles and Xie, Scalable Diffusion Models with Transformers (ICCV 2023). Code and project page. Line numbers in the Provenance panel refer to the main branch of the repository.
  2. The diffusion recipe: Ho, Jain and Abbeel, Denoising Diffusion Probabilistic Models (explainer); Nichol and Dhariwal, Improved Denoising Diffusion Probabilistic Models for the learned covariance; Dhariwal and Nichol, Diffusion Models Beat GANs on Image Synthesis (explainer) for the ADM U-Net, AdaGN, the hyperparameters and the evaluation suite DiT reuses.
  3. Guidance: Ho and Salimans, Classifier-Free Diffusion Guidance (explainer).
  4. The latent space: Rombach et al., High-Resolution Image Synthesis with Latent Diffusion Models (explainer). The decoders ablated in the appendix are sd-vae-ft-mse and sd-vae-ft-ema.
  5. The backbone: Dosovitskiy et al., An Image is Worth 16x16 Words (explainer) for patchify and the S/B/L configs; He et al., Masked Autoencoders Are Scalable Vision Learners (explainer), whose fixed 2D sine-cosine positional embedding the DiT code imports.
  6. Adaptive normalization and zero initialization: Perez et al., FiLM; Huang and Belongie, Adaptive Instance Normalization; Karras et al., StyleGAN; Goyal et al., Accurate, Large Minibatch SGD, Section 5.1 and Table 2b, for the zero-γ residual initialization.
  7. The metrics: Heusel et al., GANs Trained by a Two Time-Scale Update Rule introduced FID; Kynkäänniemi et al., Improved Precision and Recall Metric for Assessing Generative Models defined the precision and recall in the results tables; Parmar, Zhang and Zhu, On Aliased Resizing and Surprising Subtleties in GAN Evaluation is the sensitivity result the paper cites for using one evaluation suite.
  8. Where the architecture went: Esser et al., Scaling Rectified Flow Transformers for High-Resolution Image Synthesis (Stable Diffusion 3, MM-DiT), and OpenAI's Sora technical report.