Scalable Diffusion Models with Transformers

Cut the noised latent into patches, run a plain transformer over them, and image quality follows the Gflops.

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 that start at zero, and across twelve models the forward-pass compute, not the parameter count, predicts sample quality.

Twelve diffusion transformers, from 0.36 to 118.6 Gflops per forward pass, trained with one recipe, and a state-of-the-art FID of 2.27 on class-conditional ImageNet at 256×256.

Why swap the U-Net for a transformer

By 2022 transformers had taken over language modeling, image classification through the Vision Transformer, and autoregressive image generation. Diffusion models were the holdout. Every strong image diffusion model, from DDPM through ADM to Stable Diffusion, used a convolutional U-Net that Ho et al. had adapted from PixelCNN++: 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 channel counts, attention heads and normalization, but its overall shape had not changed since 2020.

Peebles and Xie ask whether the U-Net's convolutional structure is doing anything a transformer cannot. Their answer is a family of models called Diffusion Transformers (DiTs) that keep ADM's diffusion recipe, keep the latent space of Stable Diffusion, and replace only the network in the middle with a ViT that operates on patches of the noised latent. The paper has two results. The architecture works: DiT-XL/2 reaches FID 2.27 on 256×256 ImageNet, below every prior diffusion model and below the GAN record of 2.30. And the architecture scales in a measurable way: across twelve DiTs the FID at a fixed training length is a near-linear function of the log of the forward-pass Gflops, whether those Gflops come from a wider and deeper transformer or from more input tokens.

The paper measures model size in Gflops rather than parameters, because a parameter count does not change with the number of tokens the network processes. The same DiT-XL weights run on 16 tokens (patch size 8) or on 256 tokens (patch size 2), and the two runs differ by a factor of 16 in compute and by a factor of 5.5 in FID. The counting convention is part of the number: the paper's Gflops count one fused multiply-add as one flop. A per-token cost model built on that convention reproduces all twelve entries of the paper's 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; this section is the recipe it inherits, in one place. The forward process adds Gaussian noise to a clean sample $x_0$ according to a fixed schedule of constants $\bar{\alpha}_t$:

DiT uses ADM's schedule: 1000 steps, with per-step variances $\beta_t$ rising linearly from $10^{-4}$ to $0.02$ and $\bar{\alpha}_t = \prod_{s \le t}(1-\beta_s)$. At $t = 500$ the coefficients are $\sqrt{\bar{\alpha}_t} = 0.279$ on the signal and $0.960$ on the noise; by $t = 750$ the signal coefficient is 0.057. The model learns the reverse step $p_\theta(x_{t-1} \mid x_t) = \mathcal{N}(\mu_\theta(x_t), \Sigma_\theta(x_t))$, a Gaussian whose mean and covariance a network predicts from $x_t$ and $t$. Training maximizes a variational lower bound on the log-likelihood of $x_0$, which reduces to a sum of KL divergences between the true posterior step $q^*(x_{t-1} \mid x_t, x_0)$ and the model's step:

The paper prints the first term without its logarithm; (2) is the form from the DDPM paper the bound comes from. Since both distributions are Gaussian, each KL term is a closed-form function of two means and two covariances. Ho et al. reparameterized the mean as a prediction of the noise $\epsilon_t$ that was added in (1), which turns the mean part of (2), up to per-step weights they dropped, into a mean-squared error:

The covariance gives DiT a second output. Nichol and Dhariwal showed that a learned $\Sigma_\theta$ improves likelihood and lets the sampler take fewer steps, and they parameterized it as one number $v$ per dimension that interpolates, in the log domain, between the two analytic extremes $\beta_t$ and the posterior variance $\tilde{\beta}_t$:

The network therefore outputs two tensors the shape of its input: $\epsilon_\theta$ and $v$. The noise prediction is trained with (3) and the variance with the full bound (2), computed with the predicted mean detached so the KL term cannot pull on $\epsilon_\theta$. DiT's training code inherits this from ADM verbatim: the mean is detached in the variational term and the two losses are added. Sampling starts at $x_{1000} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ and draws $x_{t-1}$ from the predicted Gaussian 250 times for the reported numbers.

Class conditioning uses classifier-free guidance. During training the class label $c$ is replaced, 10% of the time, by a learned null embedding $\varnothing$ (in the code, row 1000 of a 1001-row table). At sampling time the two predictions are extrapolated:

With $s = 1$ this is the plain conditional prediction; $s = 1.5$ gives the paper's best FID. In Ho and Salimans's notation the weight is $w = s - 1$, so the paper's 1.5 is their 0.5. The derivation in the paper motivates (5) through Bayes' rule: the classifier gradient $\nabla_x \log p(c \mid x)$ equals $\nabla_x \log p(x \mid c) - \nabla_x \log p(x)$, and each score is $-\epsilon_\theta / \sqrt{1-\bar{\alpha}_t}$, so the difference of the two noise predictions is a scaled classifier gradient. The paper's intermediate line writes this gradient as $\nabla \log p(x \mid c)$ and hides the scale factor in a proportionality sign; the final formula (5) is the standard one. Each guided sampling step costs two forward passes, one with $c$ and one with $\varnothing$, which the code batches into a single call of twice the batch size.

The last inherited piece is the latent space. Running diffusion on 256×256×3 pixels is expensive, so latent diffusion first trains an autoencoder and then runs diffusion on its codes. DiT uses Stable Diffusion's pretrained VAE with downsampling factor 8: the encoder $E$ maps a 256×256×3 image to a 32×32×4 latent $z = E(x)$, the code multiplies it by 0.18215 to give the channels roughly unit variance, and the decoder $D$ maps a sampled latent back to pixels. $E$ and $D$ are frozen throughout. Everything below happens on that 32×32×4 tensor.

Patchify: a latent becomes tokens

A transformer takes a sequence of vectors, so the first layer of DiT converts the noised latent $x_t$ of shape $I \times I \times C$ (32×32×4 for 256px images) into tokens. Following ViT, it cuts the latent into non-overlapping $p \times p$ patches, flattens each patch into a vector of $p^2 C$ numbers, and multiplies by one shared linear layer to the hidden size $d$. The sequence length is

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

Two costs follow from (6). Every transformer block does two kinds of arithmetic. The attention projections and the MLP touch each token independently: in multiply-adds, $4d^2$ for the query, key, value and output projections plus $8d^2$ for the MLP with its 4× hidden width, so $12\,d^2\,T$ per block. The attention matrix itself costs $2\,T^2 d$ per block for $QK^\top$ and $AV$. For DiT-XL/2:

The conditioning layers, patch embedding and output head add about 0.2 more, and the paper reports 118.64. Halving $p$ quadruples $T$, so the first term grows 4× and the second 16×, which is why the paper says the total at least quadruples. The parameter count moves the other way, slightly: the only layers whose size depends on $p$ are the patch embedding ($p^2 C \times d$) and the output head, so DiT-XL/8 has 676M parameters and DiT-XL/2 has 675M. Figure 1 lets you set the patch size and the config and watch the two costs separate.

At $p = 2$ the attention term is 10% of the total for DiT-S/2 and 3.6% for DiT-XL/2, so on the paper's 32×32 latents the cost is close to linear in $T$. At 512px, with 1024 tokens, the attention term for XL/2 reaches 68 Gflops of 524.6, still a minority. The T² term dominates the cost only at resolutions well past the ones the paper trains.

Four ways to feed in t and c

A ViT block takes tokens and returns tokens. A diffusion network also needs the timestep $t$ and, for a conditional model, the class $c$; the network cannot know how much noise to remove without $t$. Both are first turned into vectors of size $d$. The timestep goes through a 256-dimensional sinusoidal frequency embedding (the same table used for positions in the original transformer, with 10,000 as the longest period) and then a two-layer MLP with a SiLU between, giving $t_{\text{emb}}$; the class is a row lookup in a table of 1001 embeddings, giving $y_{\text{emb}}$ (the code calls the label $y$), the last row being the null class for guidance. The two vectors are added, not concatenated, to give one condition vector $c_{\text{emb}} = t_{\text{emb}} + y_{\text{emb}}$ of size $d$. The design question is where that vector enters the block, and the paper tries four answers:

• In-context conditioning. Append $t_{\text{emb}}$ and $y_{\text{emb}}$ to the token sequence as two extra tokens, the way ViT appends a class token, use unmodified ViT blocks, and drop the two tokens after the last block. Extra cost: two tokens out of 258, negligible.
• Cross-attention. Keep the two condition vectors as a separate length-2 sequence and add a multi-head cross-attention layer after the self-attention in every block, with the image tokens as queries and the condition sequence as keys and values, as in the original encoder-decoder transformer and in LDM's class conditioning. This is the most expensive option, about 15% more Gflops (137.6 versus 118.6 for XL/2).
• Adaptive layer norm (adaLN). A transformer block normalizes its tokens twice, once before attention and once before the MLP, and a standard LayerNorm then applies a learned per-dimension scale $\gamma$ and shift $\beta$. adaLN regresses those vectors from the condition instead: a SiLU followed by one linear layer maps $c_{\text{emb}}$ to $4d$ numbers, a scale and a shift for each of the two norms. This is the FiLM and AdaIN idea that StyleGAN used for its style vector and that ADM used as AdaGN inside its U-Net. It costs $4d^2$ multiply-adds per block per sample, not per token, and it is the only design in which every token receives the same function of the condition.
• adaLN-Zero. The same, plus a third regressed vector per residual branch: a per-dimension gate $\alpha$ that multiplies the branch output immediately before it is added back to the residual stream. The linear layer now outputs $6d$ numbers, and it is initialized so that every gate starts at zero.

The paper trains DiT-XL/2 with each block and tracks FID over 400K steps. Figure 2 draws the four designs and carries the numbers; adaLN-Zero leads cross-attention and in-context conditioning at every checkpoint.

The cheapest mechanism wins: adaLN-Zero ties plain adaLN for the lowest Gflops of the four and reaches an FID 45% below in-context conditioning, which costs the same. The parameter counts run the other way from the compute: adaLN-Zero has 675M parameters against 449M for in-context, because the modulation layer in each block is a $d \times 6d$ matrix, a third of the block's $18d^2$ weights. Those weights are applied once per sample rather than once per token, so at $T = 256$ they account for 0.2% of the block's compute.

adaLN-Zero: every block starts as the identity

The gap between adaLN (25.21) and adaLN-Zero (19.47) comes from initialization. In the released code the modulation is

where $\text{LN}$ is a LayerNorm with no learned affine parameters and $\odot$ is a per-dimension product broadcast over tokens. The $1 + \gamma$ form means a regressed scale of zero leaves the normalized activations unchanged. One block then computes

with $W \in \mathbb{R}^{6d \times d}$ and $b \in \mathbb{R}^{6d}$. The code sets both $W$ and $b$ to zero at initialization, so all six vectors start at zero: the scales and shifts leave the LayerNorm alone and the gates zero out both branches. Every block returns its input unchanged, and the whole 28-block DiT-XL is the identity map from the patch embedding to the final layer. The final layer's linear projection is zero-initialized as well, so the network's first prediction is $\hat{\epsilon} = 0$ and $v = 0$ for every input. The paper describes this as zero-initializing the gates $\alpha$; the code zeroes the whole modulation layer, which also zeroes the scales and shifts, and both descriptions give the 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 lineage is the residual-network trick from Goyal et al. (2017): setting the scale $\gamma = 0$ in the last BatchNorm of every ResNet block so that signal initially flows only through the skip connection, which cut ImageNet top-1 error from 23.84% to 23.60% at batch 256 and helped more at batch 8192. ADM's U-Net does the same by zero-initializing the last convolution in each residual block. adaLN-Zero moves the zero into the conditioning path, where it doubles as a per-sample gate: after training, a block can scale its attention contribution up or down depending on $t$ and $c$, which is a capability the plain adaLN block lacks.

Figure 3 runs one token through (7) and (8) with real numbers. Set the gate to zero and the output column matches the input column exactly, whatever the condition slider does to the scale and shift; raise it and the branch starts to move the token.

The identity initialization does not stall training. With a gate at zero, the loss gradient with respect to the gate vector is the upstream gradient multiplied element-wise by the branch output $f$, and $f$ is nonzero because attention and MLP weights are initialized normally (Xavier uniform). The gate weights move at the first step; the attention and MLP weights, whose gradients carry a factor of the gate, begin to move at the second. The paper's training curves show adaLN-Zero below cross-attention and in-context conditioning from the earliest checkpoint, and 5.74 FID points below plain adaLN at 400K steps.

One training step, with shapes

After the last block, the sequence of $T$ tokens has 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, $2d$ regressed numbers, no gate) and a linear map from $d$ to $p \times p \times 2C$ numbers per token. For XL/2 that is 1152 → 32: a 2×2 patch with 8 channels. Unpatchify rearranges the $T$ outputs back into the spatial grid, giving an $8 \times 32 \times 32$ tensor whose first four channels are $\hat{\epsilon}$ and last four are $v$. The step, for one image:

# 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 in (3): the patch embedding, the 28 blocks including their modulation layers, the timestep MLP, the class table and the final layer. The VLB term adds gradient only to the paths that produce $v$, because the code detaches the noise prediction before computing it. The VAE receives nothing and never changes. The class label is swapped for the null row on 10% of samples so that the unconditional prediction in (5) is trained too.

The optimizer settings are copied from ADM without tuning: AdamW with a constant learning rate of $10^{-4}$, 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 that is the model actually evaluated. The same settings train every one of the twelve models. The paper reports that none of them showed the loss spikes common in transformer training. DiT-XL/2 runs at 5.7 iterations per second on a TPU v3-256 pod, so the 400K steps of the scaling study 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

which is the training-compute axis the paper uses in its efficiency plot.

Gflops predict FID; parameters do not

The design space is patch size times config: $p \in \{8, 4, 2\}$ crossed with three transformer configs borrowed from ViT (S, B, L) plus one larger one (XL). 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 is 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 measures FID-50K with 250 sampling steps and no guidance.

Read Figure 4 on its default axis first. FID falls with Gflops along one line from 153.6 down to 19.47, and the FID 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 almost identical weights and FIDs a factor of five apart: DiT-XL at 675M parameters scores 106.41, 43.01 or 19.47 depending on whether the patch size is 8, 4 or 2.

The paper draws the same conclusion from its training curves: at every checkpoint, a deeper and wider config beats a smaller one at the same patch size, and a smaller patch beats a larger one at the same config. It also plots FID against total training compute (Gflops × batch × steps × 3) and finds that larger models are more compute-efficient: small DiTs trained longer fall behind larger DiTs trained for fewer steps, and DiT-XL/2 overtakes DiT-XL/4 after roughly $10^{10}$ Gflops of training compute despite costing four times as much per step. The appendix repeats the analysis for sFID, Inception Score, precision and recall and for the training loss itself, and every one of them improves with Gflops. Samples from the twelve models drawn from the same starting noise and label get visibly sharper as Gflops increase.

The scaling study contains no fitted law. The twelve models are one training length, one dataset and one resolution, and the paper reports a strong negative correlation between log Gflops and FID rather than an 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 no other change.

More sampling steps do not substitute

A diffusion model can spend extra compute after training by taking more sampling steps, so a smaller DiT could in principle 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 per-image cost 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 wins on FID-10K, 23.7 against 25.9. Figure 5 puts the twelve models on a log-log chart of Gflops against steps, where a fixed per-image budget is a straight line.

The paper's reading is that sampling compute cannot make up for missing model compute, and the figure shows why the accounting favors the larger model: the budget line has slope −1, so a model 1.5× more expensive per pass only has to give up a third of its steps, and past a few hundred steps additional DDPM steps buy very little FID. The experiment is at 400K training steps and FID-10K, a noisier metric than the FID-50K used elsewhere in the paper.

The numbers, and what they leave out

For the headline result the authors keep training DiT-XL/2 to 7M steps and sample with guidance. On class-conditional ImageNet at 256×256, without guidance the model reaches FID 9.62; with guidance scale 1.25 it reaches 3.22; with 1.5 it reaches 2.27, with sFID 4.60, Inception Score 278.24, precision 0.83 and recall 0.57. The prior best diffusion model, LDM-4 with guidance scale 1.5, had 3.60; ADM with classifier guidance and an upsampler had 3.94; StyleGAN-XL, the best GAN, had 2.30. At a training length matched to ADM (2.35M steps) XL/2 already reaches 2.55. The 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. At every guidance scale tested, DiT-XL/2's recall is above LDM-4's and LDM-8's, which means 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, it reaches FID 3.04 with guidance 1.5, below the previous diffusion best of 3.85 (ADM-G with ADM-U) at a quarter or less of the compute (ADM 1983 Gflops, ADM-U 2813). StyleGAN-XL's 2.41 at this resolution stays ahead, so the paper's claim at 512×512 is the best diffusion model, not the best generative model.

Two details of the evaluation are in the appendix. The guidance in every reported number is applied to only the first three of the four latent channels, a quirk the authors kept for reproducibility; three-channel guidance at scale $1 + x$ behaves like four-channel guidance at $1 + \tfrac{3}{4}x$, and four-channel guidance at 1.375 gives FID 2.20, slightly better than the reported 2.27. And the VAE decoder can be swapped without retraining, because the three Stable Diffusion decoders share one encoder: 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 tables. The scaling plots use the MSE decoder throughout.

The scope is narrower than the title suggests. Every experiment is class-conditional ImageNet at two resolutions 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. FID is known to move with implementation details, which is why all numbers come from ADM's TensorFlow evaluation suite on exported samples, and the sampling-compute comparison uses FID-10K rather than FID-50K. The authors conclude that the U-Net's inductive bias is not needed, a standard transformer scales as it does elsewhere, and future work should keep scaling it and try it as the backbone of text-to-image models.