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 TransformersTwelve 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 with Gaussian noise according to a fixed schedule , which falls from 1 toward 0 as the timestep goes from 0 to 1000:
DiT uses ADM's linear schedule: 1000 steps, per-step noise variances rising linearly from to , and the running product of up to step . At (indexed from 0, as the code samples it) the signal coefficient is 0.279 and the noise coefficient is 0.960, so the noise already dominates; by the signal coefficient is 0.057.
The model learns the reverse step , a Gaussian whose mean and covariance a network predicts from and . Training maximizes a variational lower bound on the log-likelihood of , which splits into one KL divergence per step between the true posterior step and the model's step:
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 that was added in (1); with that substitution and the per-step weights dropped, the mean part of (2) becomes a mean-squared error:
The covariance is DiT's second output. Nichol and Dhariwal showed that learning improves likelihood and lets the sampler use fewer steps. They parameterize it with one number per dimension that interpolates, in the log domain, between two closed-form choices: the forward-step variance and the posterior variance .
So the network outputs two tensors the shape of its input, and . The noise prediction is trained with (3) and 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 and not , and the two losses are added. Sampling starts from and draws 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 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 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 (32×32×4 at 256px) into tokens. Following ViT, it cuts the latent into non-overlapping patches, flattens each patch into numbers, and applies one shared linear layer that maps them to the hidden width . The sequence length is
With and that is 256 tokens, each built from 2·2·4 = 16 numbers and embedded to in the XL config. With 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 .
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, per token for the query, key, value and output projections and for the MLP with its 4× hidden width, so per block. The attention itself, and then the weighted sum of values, compares every token with every other token and costs per block. For DiT-XL/2, with 28 blocks:
The conditioning layers, patch embedding and output head add about 0.2 more, and the paper reports 118.64. Halving quadruples , 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 are the patch embedding () 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.
At the attention term is 10% of the total for DiT-S/2 (narrow width, so the 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 . At 512px, with 1024 tokens, the XL/2 attention term reaches 68 Gflops of 524.6, still a minority. The 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 , 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 . 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 . The class is a row lookup in a table of 1001 learned embeddings, giving ; the code names the label , and the last row is the null class used for guidance. The two vectors are added, not concatenated, into one condition vector of width :
The paper tries four places to put that vector:
- In-context conditioning. Append and 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. The cost is two extra tokens out of 258.
- 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 decoder of the original Transformer and in latent diffusion's class conditioning. This costs about 15% more compute (137.6 against 118.6 Gflops for XL/2).
- Adaptive layer norm (adaLN). A transformer block normalizes its tokens twice, before attention and before the MLP, and a standard LayerNorm then applies a learned per-dimension scale and shift . adaLN computes those vectors from the condition instead: a SiLU and one linear layer map to 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 ADM used as AdaGN inside its U-Net. It costs multiply-adds per block per sample, and every token gets the same scale and shift.
- adaLN-Zero. adaLN plus a third regressed vector per residual branch: a per-dimension gate that multiplies the branch output just before it is added back to the token. The linear layer now outputs numbers and is initialized so every gate starts at zero.
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.
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 matrix, a third of the block's weights. Those weights run once per sample rather than once per token, so at 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
where is a LayerNorm with no learned affine parameters of its own and is a per-dimension product, applied identically to every token. Writing the scale as means a regressed leaves the normalized activations unchanged. One block computes
with and . The code sets both and 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 and for every input. The paper describes the initialization as zeroing the gates ; 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 returnedThe idea comes from residual networks. Goyal et al. (2017) set the scale 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 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.
With a gate at zero, the gradient of the loss with respect to that gate is the upstream gradient times the branch output , summed over tokens, and 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 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, numbers from the condition, no gate) and a linear map from to 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 and last four are .
Figure 4 walks one training step stage by stage. Pick a config and patch size and every shape updates; move and the noise coefficients in stage 2 follow the schedule. The color of each row says what this step does to its weights.
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 evaluatedEvery 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 , 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 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
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: 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.
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 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.
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 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:
At one latent element, if the unconditional guess is 0.20 and the conditional guess is 0.50, then returns 0.50 (ordinary conditional sampling), returns 0.20 + 1.5·0.30 = 0.65, and returns the unconditional 0.20. In Ho and Salimans's notation the weight is , so DiT's 1.5 is their 0.5. The reason (8) works: Bayes' rule gives the gradient of a classifier as
and each score at noise level equals . The difference of the two noise predictions is therefore the classifier gradient times . Rearranged, (8) is the conditional prediction plus 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 behaves about like four-channel guidance at : 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.
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 and 2.27 at , 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.
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
- 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.
- 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.
- Guidance: Ho and Salimans, Classifier-Free Diffusion Guidance (explainer).
- 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.
- 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.
- 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.
- 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.
- 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.
How could this explainer be improved? Found an error, or something unclear? I read every message.