An annotated companion · AI Primer

Scalable Diffusion Models with Transformers, annotated

About this page. This is a companion, not a copy. It follows the paper (version 2, March 2023) section by section, quotes only a sentence or two per section (clearly marked), and explains everything in its own words. The paper is released under the Creative Commons Attribution 4.0 licence; rows of its tables are reproduced here with attribution, its charts are replotted from those published numbers, and its diagrams are redrawn from scratch. Equations are reproduced with every symbol decoded; a few formulas that write out, in symbols, something the paper says in words are labelled as this page's own. Made-up numbers are labelled illustrative. Read the original alongside: every section links to it.

How to read this page

Nothing on this page assumes you already know the jargon. Three things help:

  • Any dotted word explains itself when you hover it, tab to it, or tap it.
  • Every symbol inside an equation does the same, and each equation is followed by a table of its symbols.
  • The diagrams are live: hover or tap any part to see what it does. The charts read out their values as you hover or use the arrow keys.

Each idea is explained in the same order: an everyday picture, a tiny example you could check by hand, a diagram, then the math, and finally why it still matters. This paper stands on three others with companions of their own: DDPM (the noise-guessing diffusion model), latent diffusion (diffusing in an autoencoder's small latent) and ViT (a transformer that reads an image as patches). Most of the arithmetic uses one running example:

The running example. A 256 × 256 photo is encoded by the latent diffusion autoencoder into a latent of 32 × 32 × 4. The transformer cuts that latent into p × p patches: with p = 2 that is a 16 × 16 grid, 256 tokens of 2 × 2 × 4 = 16 numbers each. The diffusion lesson builds diffusion from scratch in NumPy and counts these tokens; the transformer lesson builds the blocks.

Abstract

“In addition to possessing good scalability properties, our largest DiT-XL/2 models outperform all prior diffusion models on the class-conditional ImageNet 512 × 512 and 256 × 256 benchmarks, achieving a state-of-the-art FID of 2.27 on the latter.”Peebles and Xie (2022), Abstract. Read the original

Everyday picture

For years every diffusion model's denoiser was the same shape of network, a U-Net, the way every kitchen once had the same stove. This paper asks whether the stove matters, installs the general-purpose appliance everyone else was already using, a transformer, and finds that the food gets better the more power you give it.

What the paper claims

  • A diffusion transformer (DiT), a plain transformer over patches of the latent, can replace the U-Net inside a latent diffusion model.
  • Quality tracks forward-pass compute, measured in Gflops (billions of floating-point operations): deeper, wider or more tokens all lower FID.
  • The largest model, DiT-XL/2, sets a new state of the art on class-conditional ImageNet: FID 2.27 at 256 × 256, and 3.04 at 512 × 512 among diffusion models.

Why it matters today

Once the denoiser is a transformer, everything learned about scaling transformers for language carries over. Many image and video generators since are built around a transformer denoiser over latent patches, the design this paper benchmarked.

1 Introduction · original

“We show that the U-Net inductive bias is not crucial to the performance of diffusion models, and they can be readily replaced with standard designs such as transformers.”Peebles and Xie (2022), §1

Everyday picture

The U-Net came to diffusion almost by inheritance: DDPM borrowed it from a pixel-by-pixel image model, and later work tuned it without changing its outline. It is a convolutional network, mostly ResNet blocks, that shrinks the image and grows it back, with some self-attention at low resolutions. Its built-in assumptions (inductive bias) about images were widely assumed to be the secret of diffusion's success. The paper tests that assumption by removing it.

Tiny example: why a common architecture helps

A transformer is a stack of identical blocks, described by three numbers: depth N, width d, and heads. Making it bigger is a matter of turning those dials, and the language and vision literature already knows which settings work. DiT borrows ViT's S, B and L settings unchanged and adds one, XL. Doubling width at fixed depth roughly quadruples the matrix work: DiT-S (d = 384) to DiT-B (d = 768), both 12 layers, goes from 1.4 to 5.6 Gflops at the same patch size, exactly 4 times.

In Python:

# Table 1: Gflops at a 32 × 32 latent with 4 × 4 patches
gflops = {"S": 1.4, "B": 5.6}
d = {"S": 384, "B": 768}
# doubling the width d, depth fixed at 12
d["B"] / d["S"], gflops["B"] / gflops["S"]  # → (2.0, 4.0)

Why it matters

The paper frames its goal as demystifying the role of architecture in diffusion, and offering baselines: a standard architecture lets diffusion inherit the training recipes, scaling behaviour and tooling of the rest of the field. The ViT companion tells the parallel story for image classification, where a plain transformer also caught up with convolutions once given enough data.

2 Related work · original

Everyday picture

Three threads meet here. Transformers had already taken over language and vision and shown clean scaling laws. Diffusion models had overtaken GANs on images, mostly through better sampling and classifier-free guidance, all with U-Nets. And the question of how to measure a network's size: by parameter count, or by the work it does per input?

Tiny example: parameters are not compute

A convolution's weights do not depend on the image size, but its work does. The same with DiT's patch size: DiT-XL/8 and DiT-XL/2 have almost the same parameters (676 million against 675 million) but 7.39 against 118.64 Gflops, because p = 2 makes 16 times as many tokens and every token passes through every weight. That is why the paper measures complexity in Gflops.

In Python:

# Table 4: parameters (millions) and Gflops for DiT-XL at two patch sizes
params = {8: 676, 2: 675}
gflops = {8: 7.39, 2: 118.64}
# tokens for a 32 × 32 latent: (32 / p)²
tokens = {p: (32 // p) ** 2 for p in (8, 2)}
tokens  # → {8: 16, 2: 256}
round(gflops[2] / gflops[8], 1), round(params[2] / params[8], 3)  # → (16.1, 0.999)

Why it matters

Counting Gflops puts DiT in line with how architectures are compared elsewhere, and it is the axis on which the paper's main result (quality follows compute) is stated. The transformer lesson counts a transformer's flops per token.

3 Diffusion transformers · original

Everyday picture

Section 3 first recalls the diffusion recipe DiT is trained with (§3.1), then lays out the menu of design choices the experiments will sweep (§3.2): how finely to cut the latent into patches, how the noise step and class label get into the network, and how big to make it.

Why it matters

Almost nothing about diffusion itself changes; the paper's contribution is entirely in the network. That makes the comparison with U-Net models clean: same autoencoder, same noise schedule, same training settings as ADM, a strong U-Net baseline.

3.1 Preliminaries · original

Everyday picture

The recipe is DDPM's: fade a real picture into static a little at a time, and train a network to guess the static that was added. To generate, start from pure static and remove a little of the guessed static, step after step.

Tiny example

One latent number with clean value x0 = 2.0, at a step where the surviving signal is ᾱt = 0.64, and noise εt = 0.5: the noisy value is √0.64 × 2.0 + √0.36 × 0.5 = 1.6 + 0.3 = 1.9. If the network guesses the noise in a two-number latent as (0.4, −0.8) when the truth was (0.5, −1.0), the squared miss is 0.05.

In words: “after t steps a clean input is shrunk by the square root of the surviving signal share and topped up with noise; the network is trained on the squared distance between its noise guess and the noise that was really added.”

With the numbers: xt = 0.8 × 2.0 + 0.6 × 0.5 = 1.9, and ‖(0.4, −0.8) − (0.5, −1.0)‖² = 0.01 + 0.04 = 0.05.

In Python:

import math
x0, alpha_bar, eps_t = 2.0, 0.64, 0.5
# x_t = √ᾱ_t x_0 + √(1 − ᾱ_t) ε_t
round(math.sqrt(alpha_bar) * x0 + math.sqrt(1 - alpha_bar) * eps_t, 2)  # → 1.9
# L_simple = ‖ε_θ(x_t) − ε_t‖²
round(sum((g - e) ** 2 for g, e in zip([0.4, -0.8], [0.5, -1.0])), 2)  # → 0.05

A learned spread as well. DDPM fixed the width of each reverse step; DiT follows Nichol and Dhariwal and also learns it, as a diagonal covariance Σθ. The noise guess is trained with Lsimple and Σθ with the full variational bound, which the paper writes as the reconstruction term plus a sum of KL divergences between the true backward step and the model's. (The printed bound reads “−p(x0|x1)”; it should be −log p(x0|x1), as in DDPM.) This is why DiT's output layer produces two things, a noise guess and a covariance.

Classifier-free guidance

Everyday picture: ask the same painter twice, once “paint a dog” and once “paint anything”, and push the canvas further along the difference between the two answers. Classifier-free guidance does this at every sampling step; the network learns both answers because the label is randomly replaced by a learned “null” label during training.

Tiny example: the guess with no label is 0.2 and with the label “dog” is 0.5. At guidance scale s = 1.5 the guided guess is 0.2 + 1.5 × (0.5 − 0.2) = 0.65; at s = 4 it is 1.4. Why this works comes from Bayes' rule: if p(c) = 0.25, p(x|c) = 0.3 and p(x) = 0.1, then p(c|x) = 0.75, and log p(c|x) = log p(x|c) − log p(x) + log p(c); the last term does not depend on x, so its slope with respect to x is zero.

In words: “how strongly an input says ‘class c’ is how much more likely it is under class c than in general; so start from the unlabelled noise guess and move s times as far as the label moves it.”

With the numbers: 0.2 + 1.5 × 0.3 = 0.65 and 0.2 + 4 × 0.3 = 1.4; s = 1 gives the plain labelled guess 0.5. For the Bayes line, log 0.3 − log 0.1 + log 0.25 = −0.288 = log 0.75.

In Python:

import math
eps_null, eps_c = 0.2, 0.5
# ε̂ = ε(x_t, ∅) + s (ε(x_t, c) − ε(x_t, ∅)), for s = 1, 1.5 and 4
[round(eps_null + s * (eps_c - eps_null), 2) for s in (1, 1.5, 4)]  # → [0.5, 0.65, 1.4]
# Bayes: log p(c|x) = log p(x|c) − log p(x) + log p(c)
p_c, p_x_given_c, p_x = 0.25, 0.3, 0.1
round(math.log(p_x_given_c) - math.log(p_x) + math.log(p_c), 3), round(math.log(0.75), 3)  # → (-0.288, -0.288)

Latent diffusion

DiT runs inside the latent diffusion framework: a frozen encoder E turns an image x into z = E(x), the diffusion model learns to generate z, and the decoder turns a generated z back into pixels, x = D(z). The paper calls the result a hybrid: an off-the-shelf convolutional autoencoder around a transformer denoiser. DiT could run on pixels unchanged; latents are simply cheaper.

Why it matters today

Nothing on this list is new, which is the point: the comparison with U-Nets isolates the backbone. The diffusion lesson builds the loss, the sampler and guidance on dots in a plane.

In code: add_noise is the one-jump noising, train_noise_predictor trains on Lsimple (hiding labels at random when given them), and guided_noise is the guidance formula.

3.2 Diffusion transformer design space · original

“We aim to be as faithful to the standard transformer architecture as possible to retain its scaling properties.”Peebles and Xie (2022), §3.2

Everyday picture

A transformer reads a sequence of tokens. An image latent is a grid. So cut the grid into square tiles, flatten each tile into a row of numbers, and treat each row as a token, exactly as ViT does with photos. The tile size is the main dial: small tiles give many tokens and a detailed but expensive read; big tiles give few tokens and a cheap, coarse one.

Latent diffusion transformer DiT block with adaLN-Zero noised latent32 × 32 × 4 timestep t label y Patchify Embed N × DiT block Layer norm Linear and reshape noise32 × 32 × 4 Σ32 × 32 × 4 γ₁, β₁ α₁ γ₂, β₂ α₂ input tokens conditioning MLP Layer norm Scale, shift Multi-headself-attention Scale + Layer norm Scale, shift Pointwisefeedforward Scale + to next block

Hover or tap any part. Start at the bottom left with the noised latent and read upwards; the right-hand panel opens up one DiT block.

Figure 3 of the paper, redrawn: the DiT architecture (left) and the adaLN-Zero block (right). The paper's two other block designs, cross-attention and in-context conditioning, are described below. Based on Peebles and Xie (2022), Figure 3.

Reading it: read both columns from the bottom up. On the left, the noised latent is cut into patch tokens (Patchify) while the step t and the class label y are turned into vectors (Embed). N identical DiT blocks process the tokens; the dashed lines open one block up on the right. A final layer norm and a linear layer turn each token back into its patch of numbers, which are rearranged into two latent-sized grids: the noise guess and the covariance Σ. On the right, the main path is an ordinary transformer block: layer norm, self-attention, add back the input (the first ⊕), layer norm, feedforward, add back again (the second ⊕). The difference is the right-hand rail: the conditioning passes through a small MLP that produces six vectors per block, a scale and a shift after each layer norm (γ, β) and a gate before each addition (α). The long wires on the left of the panel are the residual connections, which carry the input around each sub-layer untouched.

Patchify: the tile size sets the token count

Tiny example: a latent with side I = 32 and patch size p = 2 is cut into a (32/2) × (32/2) = 16 × 16 grid: T = 256 tokens, each holding 2 × 2 × 4 = 16 numbers that a linear layer maps to width d. With p = 4 there are 64 tokens, with p = 8 only 16. Halving p quadruples T.

In words: “the number of tokens is the number of tiles along one side, squared; each tile holds p × p positions of C channels, the numbers a token carries before the linear embedding.”

With the numbers: I = 32, p = 2: T = 256 tokens of 16 numbers. At 512 × 512 the latent is 64 × 64 × 4, and p = 2 gives 1,024 tokens. Self-attention compares every token with every other, so its score grid grows as T²: 65,536 entries per head at 256 tokens, 1,048,576 at 1,024.

In Python:

C = 4
for I, p in [(32, 8), (32, 4), (32, 2), (64, 2)]:
    # T = (I / p)², tile size p·p·C, and attention's T² scores
    T = (I // p) ** 2
    print(I, p, T, p * p * C, T * T)  # → 32 8 16 256 256 32 4 64 64 4096 32 2 256 16 65536 64 2 1024 16 1048576

Reading it: the square is the 32 × 32 latent of a 256 × 256 image, seen from above; the lines cut it into p × p patches, one token each. Drag p from 8 down to 2: the patch count goes 16, 64, 256. The readout gives the token count, the width of each patch before embedding, and the measured cost of DiT-XL at that patch size from the paper's Table 4. Notice the cost column: it rises about 4 times per halving of p (7.39, 29.05, 118.64 Gflops) while the parameter count stays at about 675 million. The paper uses p = 2, 4 and 8, following each patchify with ViT's fixed sine-cosine positional embeddings.

Four ways to feed in the step and the label

The network must know the noise step t and the class label c. The paper tries four block designs:

  • In-context conditioning: add the embeddings of t and c as two extra tokens, like ViT's class token, and use unmodified blocks. Almost no extra compute.
  • Cross-attention: keep t and c as a separate length-two sequence and add a cross-attention layer after each self-attention, like the latent diffusion U-Net. The most expensive, about 15% more Gflops.
  • Adaptive layer norm (adaLN): instead of learning one fixed scale γ and shift β for each layer norm, compute them from the sum of the t and c embeddings. The cheapest, and the only design that applies the same conditioning function to every token.
  • adaLN-Zero: adaLN plus a per-channel gate α just before each residual addition, with the MLP that produces α initialized to output zero, so every block starts out as the identity function.

Everyday picture for adaLN: a layer norm is a volume leveller that sets every token's numbers to a standard level. adaLN lets the conditioning turn each channel's volume knob (γ) and offset (β) after levelling: at a very noisy step, “turn up the coarse-shape channels”; for the class “dog”, “turn up whatever the dog features are”. The gate α is a fader on each block's contribution, which starts fully down.

Tiny example: one token h = (2, 0, −2). Layer norm makes it (1.225, 0, −1.225) (mean 0, spread 1). If the conditioning MLP outputs γ = (2, 0.5, 1) and β = (0, 1, 0), the modulated token is (2.449, 1, −1.225). With the gate α = 0, whatever the sub-layer then computes is multiplied by zero, and the block returns h unchanged.

In words: “add the step and label embeddings and let a small network turn them into a scale, a shift and a gate; level the token, rescale and shift it with those, run the sub-layer, and add its output back to the token only as much as the gate allows.” These three lines are this page's notation for what the paper says in words and draws in its Figure 3; each block has two such sub-layers (attention and feedforward), each with its own γ, β and α.

With the numbers: LN(2, 0, −2) = (1.225, 0, −1.225); γ ⊙ LN + β = (2 × 1.225 + 0, 0.5 × 0 + 1, 1 × (−1.225) + 0) = (2.449, 1, −1.225); with α = 0 the update is h + 0 = (2, 0, −2).

In Python:

import math
h = [2.0, 0.0, -2.0]
# LN(h): subtract the mean, divide by the spread
mean = sum(h) / len(h)
spread = math.sqrt(sum((v - mean) ** 2 for v in h) / len(h))
ln = [(v - mean) / spread for v in h]
[round(v, 3) for v in ln]  # → [1.225, 0.0, -1.225]
# adaLN(h) = γ ⊙ LN(h) + β, with γ and β from the conditioning MLP
gamma, beta = [2, 0.5, 1], [0, 1, 0]
modulated = [g * v + b for g, v, b in zip(gamma, ln, beta)]
[round(v, 3) for v in modulated]  # → [2.449, 1.0, -1.225]
# h ← h + α ⊙ F(...): at initialization α = 0, so the block is the identity
alpha = [0, 0, 0]
F_out = [0.7, -0.3, 1.1]
[v + a * f for v, a, f in zip(h, alpha, F_out)]  # → [2.0, 0.0, -2.0]

What the gate costs. The conditioning MLP's last layer outputs 4d numbers per block for adaLN (γ and β for both sub-layers) and 6d for adaLN-Zero (adding the two α). For DiT-XL (d = 1,152, N = 28 blocks) the two extra d-wide outputs need 28 × (1,152 × 2,304 + 2,304) ≈ 74.4 million more weights, which matches Table 4's 675 million against 600 million. The same arithmetic for cross-attention (four d × d projections per extra attention layer, 28 × 4 × 1,152² ≈ 148.6 million) matches its 598 million against in-context conditioning's 449 million. (This accounting is this page's, not the paper's.)

In Python:

d, N = 1152, 28
# two extra d-wide outputs (the gates) per block, weights plus biases
extra_zero = N * (d * 2 * d + 2 * d)
print(f"{extra_zero:,}", 675 - 600)  # → 74,382,336 75
# one extra attention layer per block: four d × d projections
extra_xattn = N * 4 * d * d
print(f"{extra_xattn:,}", 598 - 449)  # → 148,635,648 149

Model size, and the decoder

Table 1 of Peebles and Xie (2022), reproduced with attribution (CC BY 4.0): DiT configurations, with Gflops for a 32 × 32 latent and p = 4
ModelLayers NHidden size dHeadsGflops
DiT-S1238461.4
DiT-B12768125.6
DiT-L2410241619.7
DiT-XL2811521629.1

A model is named by size and patch: DiT-XL/2 is the XL configuration with p = 2. After the last block, a final layer norm (adaptive, with adaLN) and a linear layer turn each token into a p × p × 2C patch: C channels of noise guess and C of covariance. With p = 2 and C = 4 that is 32 numbers per token; 256 tokens × 32 = 8,192 = 32 × 32 × 8, rearranged into the two 32 × 32 × 4 outputs.

In Python:

p, C, tokens = 2, 4, 256
# each token decodes to a p × p × 2C patch
per_token = p * p * 2 * C
per_token, tokens * per_token, 32 * 32 * 8  # → (32, 8192, 8192)

Why it matters today

adaLN-style conditioning, with its zero-initialized gates, became a common way to feed the noise step (and often more) into transformer denoisers. The whole design space is small on purpose: patch size, block type and model size, three dials with clear costs.

In code: primer.ml.cnn_rnn.patchify cuts an image into flattened patches and patch_tokens counts a diffusion transformer's tokens; primer.ml.deep_nets.layer_norm takes the scale γ and shift β as arguments, so passing values computed from the conditioning makes it adaLN; TransformerBlock is the unconditioned block DiT modifies.

4 Experimental setup · original

Everyday picture

A fair race needs identical conditions. Every DiT is trained with the same settings, borrowed almost unchanged from the ADM U-Net, so any difference in quality comes from the architecture.

The settings

Settings stated in §4 of Peebles and Xie (2022), collected into one table
SettingValue
Dataclass-conditional ImageNet at 256 × 256 and 512 × 512; horizontal flips only
AutoencoderStable Diffusion's pretrained VAE, downsampling 8: a 256 × 256 × 3 image becomes 32 × 32 × 4
Diffusion1,000 steps, linear variance schedule from 10−4 to 2 × 10−2, ADM's covariance and embeddings
OptimizerAdamW, constant learning rate 10−4, no weight decay, batch 256, no warmup, no regularization
Initializationfinal linear layer zero; otherwise ViT's standard initialization
Weights reportedan exponential moving average with decay 0.9999
EvaluationFID-50K with 250 DDPM sampling steps, using ADM's evaluation code; no guidance unless stated
ComputeJAX on TPU v3 pods; DiT-XL/2 trains at about 5.7 iterations per second on a v3-256 pod

The paper stresses that training was “highly stable” across every configuration, with none of the loss spikes common in transformer training, even without warmup or regularization.

The moving average, decoded

Tiny example: a weight's running average stands at 1.0 when a training step moves the weight itself to 2.0. The average moves only a ten-thousandth of the way: to 1.0001. Roughly, the average remembers the last 1/(1 − 0.9999) = 10,000 steps, which smooths out the jitter of individual updates.

In words: “after each training step, keep almost all of the averaged weights and mix in a tiny share of the current ones.” This line is this page's notation for the paper's “EMA of DiT weights with a decay of 0.9999”.

With the numbers: 0.9999 × 1.0 + 0.0001 × 2.0 = 1.0001; the effective memory is 1/0.0001 = 10,000 steps.

In Python:

decay = 0.9999
theta_bar, theta = 1.0, 2.0
# θ̄ ← decay θ̄ + (1 − decay) θ
round(decay * theta_bar + (1 - decay) * theta, 6)  # → 1.0001
# how many recent steps the average roughly remembers
round(1 / (1 - decay))  # → 10000

Why it matters

Using ADM's recipe without tuning is a strong statement: a transformer needed no special care to train as a denoiser. The averaged weights are the ones that sample; the raw weights are only used to keep training. The optimizers lesson builds the same kind of running average inside Adam.

5 Experiments · original

Everyday picture

Three questions in turn: which block design, how does quality change with size and patch, and is the key quantity parameters or compute?

Block design

Four DiT-XL/2 models, one per block design, are trained for 400,000 steps (Figure 5 of the paper):

Reading it: each bar is one block design's FID-50K after 400,000 steps, without guidance, drawn from 0 to 40; shorter is better, and the Gflops of each model are listed under its name. The striped bar, adaLN-Zero, is best at 19.47, and among the cheapest (plain adaLN costs about the same). In-context conditioning is worst at 35.24, nearly twice adaLN-Zero's FID. The most expensive design, cross-attention at 137.62 Gflops, lands in the middle. And the gate's zero start matters: plain adaLN, with the same compute, scores 25.21. Values from the paper's Table 4, reproduced with attribution.

Every later model in the paper uses adaLN-Zero blocks.

Scaling model size and patch size

Twelve models: four sizes (S, B, L, XL) times three patch sizes (8, 4, 2), each trained 400,000 steps. Making the model deeper and wider improves FID at every patch size; making the patches smaller improves FID at every model size, with parameters essentially unchanged.

Hover or use the arrow keys to read a point.

Reading it: the x-axis is the model's forward-pass compute in Gflops on a logarithmic scale; the y-axis is FID-50K after 400,000 steps, without guidance (lower is better). Each line is one model size; its three points are p = 8, 4 and 2 from left to right. Every line slopes down, and, more strikingly, the four lines fall on one shared curve: at similar Gflops, different models reach similar FID. DiT-S/2 (6.06 Gflops, FID 68.40) and DiT-B/4 (5.56 Gflops, FID 68.38) are the paper's own example. This is Figure 8 of the paper, replotted from its Table 4 with attribution.

Tiny example: DiT-S/2 and DiT-B/4 differ four-fold in parameters (33 million against 130 million) but land within 0.02 FID, because their Gflops are about the same. Parameters alone do not predict quality here; compute does.

In Python:

# Table 4: (params in millions, Gflops, FID-50K at 400K steps)
S2, B4 = (33, 6.06, 68.40), (130, 5.56, 68.38)
round(B4[0] / S2[0], 1), round(abs(S2[2] - B4[2]), 2)  # → (3.9, 0.02)

Larger models use training compute better

Everyday picture: given a fixed budget of hours, a small team working long can lose to a large team working briefly. The paper estimates training compute as model Gflops × batch size × steps × 3, the 3 because a backward pass costs about twice a forward pass.

In words: “the work to train is the work of one forward pass per image, times images per step, times steps, times three for the forward and backward passes.”

With the numbers: DiT-XL/2 at 400,000 steps: 118.64 × 256 × 400,000 × 3 ≈ 3.6 × 1010 Gflops; DiT-XL/4 at the same step count, 29.05 × 256 × 400,000 × 3 ≈ 8.9 × 109. The paper reports XL/2 overtaking XL/4 at matched training compute from about 1010 Gflops on.

In Python:

batch, steps = 256, 400_000
for name, gflops in [("XL/2", 118.64), ("XL/4", 29.05)]:
    # Gflops · batch · steps · 3
    print(name, f"{gflops * batch * steps * 3:.1e}")  # → XL/2 3.6e+10 XL/4 8.9e+09

Why it matters

“Quality follows compute, however you spend it” is the transformer scaling story from language, now shown for diffusion. It means a DiT can be scaled predictably, and that cutting patches finer is as good a way to spend compute as adding layers. The scaling laws companion tells the language version.

5.1 State-of-the-art diffusion models · original

Everyday picture

Having found the best dial settings, the paper trains the biggest model much longer and enters it in the standard competition: class-conditional ImageNet.

Tiny example

DiT-XL/2, trained for 7 million steps and sampled with guidance scale 1.5, reaches FID 2.27, against the previous best diffusion model, LDM-4-G, at 3.60: a 37% lower FID. Without guidance it is 9.62; guidance at 1.5 cuts that by more than three quarters.

In Python:

# FID-50K, ImageNet 256 × 256
dit_g, ldm_g, dit_plain = 2.27, 3.60, 9.62
round(1 - dit_g / ldm_g, 2)  # → 0.37
round(1 - dit_g / dit_plain, 2)  # → 0.76
0 2 4 6 8 FID (lower is better) bubble area ∝ Gflops per forward pass ADM-U-G LDM-8-G LDM-4-G DiT-XL/2-G

Hover or tap a bubble to see its FID and compute.

Figure 2 (right) of the paper, replotted from its Tables 2 and 6 with attribution: guided class-conditional ImageNet 256 × 256 models, FID against model, with bubble area proportional to Gflops.

Reading it: height is FID, so lower is better; a bubble's area is the model's compute for one forward pass. The biggest, faintest bubble is ADM-U-G, a pixel-space U-Net with an upsampler: 742 Gflops for FID 3.94. The two latent U-Nets from the latent diffusion paper are far cheaper: LDM-8-G at 57 Gflops (FID 7.76) and LDM-4-G at 104 Gflops (FID 3.60). DiT-XL/2-G, the highlighted bubble, sits lowest, at FID 2.27, with a bubble about the size of LDM-4-G's (118.6 Gflops): a transformer of similar cost to the best latent U-Net, and better. The ADM-U figure is the combined guided model's FID from Table 2 (“ADM-G, ADM-U”), with the Gflops Table 6 gives for ADM-U.

Rows of Tables 2 and 3 of Peebles and Xie (2022), reproduced with attribution (CC BY 4.0): class-conditional ImageNet
ModelFID ↓sFID ↓IS ↑Precision ↑Recall ↑
256 × 256
BigGAN-deep6.957.36171.40.870.28
StyleGAN-XL2.304.02265.120.780.53
ADM-G, ADM-U3.946.14215.840.830.53
LDM-4-G (cfg = 1.50)3.60247.670.870.48
DiT-XL/29.626.85121.500.670.67
DiT-XL/2-G (cfg = 1.25)3.225.28201.770.760.62
DiT-XL/2-G (cfg = 1.50)2.274.60278.240.830.57
512 × 512
StyleGAN-XL2.414.06267.750.770.52
ADM-G, ADM-U3.855.86221.720.840.53
DiT-XL/2-G (cfg = 1.50)3.045.02240.820.840.54

What the rows say

  • At 256 × 256, DiT-XL/2-G has the lowest FID of every model listed, including the GAN StyleGAN-XL (2.30). Trained only 2.35 million steps, like ADM, it already reaches 2.55.
  • DiT-XL/2 has higher recall than LDM-4 and LDM-8 at each guidance scale tested: it covers more of ImageNet's variety.
  • At 512 × 512 the latent is 64 × 64 × 4, so p = 2 gives 1,024 tokens and 524.6 Gflops, still far below ADM's 1,983 and ADM-U's 2,813. FID 3.04 beats ADM's 3.85, though not StyleGAN-XL's 2.41.
  • The secondary columns: sFID, the Inception score (IS), and precision and recall.

Why it matters today

This was the result that made transformer denoisers the default to try. Note the numbering slip in the paper: §5.1 refers to the 256 × 256 results as “Table 3”, but the table carrying them is captioned Table 2 (Table 3 is the 512 × 512 one).

5.2 Scaling model compute against sampling compute · original

Everyday picture

A diffusion model can be given more thinking time after training: take more sampling steps. Can a small model that thinks longer match a big model that thinks briefly? This section says no.

Tiny example

DiT-L/2 with 1,000 sampling steps spends 80.71 Gflops per step, 80.7 Tflops per image. DiT-XL/2 with 128 steps spends 118.64 × 128 = 15.2 Tflops, about a fifth as much. XL/2 still wins: FID-10K 23.7 against 25.9.

In words: “the work to draw one image is one forward pass per sampling step.”

With the numbers: 80.71 × 1,000 = 80,710 Gflops = 80.7 Tflops; 118.64 × 128 = 15,186 Gflops = 15.2 Tflops; the ratio is 5.3.

In Python:

# Gflops per forward pass × sampling steps, in Tflops
l2 = 80.71 * 1000 / 1000
xl2 = 118.64 * 128 / 1000
round(l2, 1), round(xl2, 1), round(l2 / xl2, 1)  # → (80.7, 15.2, 5.3)

Why it matters today

“Scaling-up sampling compute cannot compensate for a lack of model compute”: for diffusion, the capacity of the network sets a ceiling that more steps cannot lift. Much later engineering went the other way, keeping large models and cutting steps; the diffusion lesson measures quality against steps for a fixed network.

6 Conclusion · original

What the paper concludes

DiT, a simple transformer backbone for diffusion, outperforms prior U-Net models and inherits transformers' scaling behaviour. The authors suggest scaling further, in model size and tokens, and trying DiT as a drop-in backbone for text-to-image models such as DALL·E 2 and Stable Diffusion.

Why it matters today

The diffusion lesson's closing section describes the pipeline that followed: a transformer over latent patches that reads a text encoder's tokens, with video as the same idea over more tokens. Its reading list ends with Stable Diffusion 3, a rectified-flow transformer trained at scale.

Appendices A, C and D · original

A: embeddings, activations, and guidance on three channels

The timestep is embedded as a 256-dimensional frequency vector (sines and cosines, as in the sinusoidal positional encoding), then passed through a two-layer MLP of width d with SiLU activations. Each adaLN layer feeds the sum of the timestep and class embeddings through a SiLU and a linear layer with 4d (adaLN) or 6d (adaLN-Zero) outputs. The transformer's own activations are GELU.

Everyday picture: a curious detail. With guidance, the authors applied it to only the first three of the four latent channels, and found the effect equivalent to four-channel guidance at a smaller scale.

Tiny example: three-channel guidance at 1.5 (x = 0.5) behaves like four-channel guidance at 1 + ¾ × 0.5 = 1.375; the reported FIDs are 2.27 and 2.20.

In words: “guiding three of four channels by an extra x is about the same as guiding all four by three quarters of that extra.”

With the numbers: x = 0.5 gives 1 + 0.375 = 1.375.

In Python:

x = 0.5
# 1 + ¾ x
1 + 3 / 4 * x  # → 1.375

C: scaling beyond FID, and the training loss

The compute trend holds for sFID, Inception score, precision and recall too, with Inception score and precision benefiting most. Larger DiTs also reach lower training loss, faster, as larger language models do.

D: the decoder, and the U-Nets' compute

The VAE used everywhere is the latent diffusion paper's f = 8 autoencoder; Stable Diffusion's two fine-tuned decoders (ft-MSE and ft-EMA) share its encoder, so they can be swapped in after training. The choice barely matters:

Table 5 of Peebles and Xie (2022), reproduced with attribution (CC BY 4.0): DiT-XL/2-G (cfg = 1.5), ImageNet 256 × 256, with three decoders
DecoderFID ↓sFID ↓IS ↑Precision ↑Recall ↑
original (LDM)2.465.18271.560.820.57
ft-MSE2.304.73276.090.830.57
ft-EMA2.274.60278.240.830.57

Even with the original LDM decoder, 2.46 beats every earlier diffusion model. Table 6 of the appendix gives the U-Net baselines' compute: ADM at 256 × 256 costs 1,120 Gflops, ADM-U 742, LDM-4 104 and LDM-8 57. One reporting wrinkle: Table 4 says its FIDs use the ft-MSE decoder, yet its 7-million-step DiT-XL/2 row shows 9.62, the same as the ft-EMA figure in Table 2.

In Python:

# Gflops per forward pass at 256 × 256 (Table 6), against DiT-XL/2
gflops = {"ADM": 1120, "ADM-U": 742, "LDM-4": 104, "LDM-8": 57}
dit = 118.64
{k: round(v / dit, 2) for k, v in gflops.items()}  # → {'ADM': 9.44, 'ADM-U': 6.25, 'LDM-4': 0.88, 'LDM-8': 0.48}

Why it matters

Swapping decoders without retraining shows the two-stage design paying off: the autoencoder and the denoiser are separate parts. And a DiT-XL/2 forward pass costs about a tenth of ADM's, which is how a transformer at 118.6 Gflops beats a U-Net at 1,120.

What changed since 2022

The transformer denoiser stayed; much around it moved on:

Choice in the paperCommon todayWhyWhere to read
Class labels through adaLN-ZeroText prompts, which are sequences rather than one vector: read by cross-attention, or joined to the patch tokens, the paper's other two block designsFollow a sentence, not a classlatent diffusion companion
DDPM noise prediction, 250 stepsFlow matching and rectified flow, fewer stepsStraighter paths, cheaper samplingflow matching companion
ImagesVideo, with patches spanning space and timeThe same design with more tokensmultimodal lesson

Glossary

Every term with hover guidance on this page, in one place.