Replace the U-Net backbone in diffusion models with a Vision Transformer. Patchify latent representations, condition with adaptive layer norm, and watch FID drop as you scale compute — just like language models.
By 2022, diffusion models had become the dominant paradigm for image generation. DALL-E 2, Imagen, Stable Diffusion — all of them produced stunning images. And all of them used the same backbone: a U-Net.
The U-Net was inherited from early pixel-level models (PixelCNN++) and adapted by Ho et al. for DDPM. It's convolutional, with ResNet blocks at multiple resolutions and self-attention sprinkled in at lower resolutions. Dhariwal and Nichol (ADM) ablated some architectural choices — channel counts, normalization layers — but the high-level design remained essentially unchanged from 2020 to 2022.
Meanwhile, a different story was playing out everywhere else in deep learning. Transformers had taken over NLP, vision (ViT), reinforcement learning, and even protein folding. The reason? Transformers scale. Double the compute, get predictably better performance. Clean scaling laws. No architectural ceilings.
But diffusion models were stuck with U-Nets. Nobody had seriously asked: what happens if you replace the U-Net with a transformer?
There were practical concerns too. U-Nets are awkward to scale — you can add channels or attention heads, but there's no clean "just make it bigger" knob like there is for transformers (where you simply increase depth and width). The architecture is also hard to unify across modalities, making it difficult to share training recipes or transfer insights from language modeling.
Peebles and Xie's insight is beautifully simple: use a standard Vision Transformer as the denoising backbone in a latent diffusion model.
The recipe:
That's it. No multi-scale processing. No skip connections across resolutions. No convolutional layers (except inside the frozen VAE). Just a plain transformer operating on a flat sequence of patches.
| Component | Parameters | Notes |
|---|---|---|
| Patch embedding | ~74K | Linear: (p×p×4) → 1152 = 16 × 1152 + bias |
| Positional embedding | ~295K | 256 positions × 1152 dims (sine-cosine, not learned) |
| 28 DiT blocks | ~668M | Per block: MHSA (4 × 1152²) + FFN (2 × 1152 × 4608) + adaLN MLP |
| adaLN conditioning MLP | ~5M shared | Maps t+c embedding to 6 scalars per block (shared MLP, block-specific outputs) |
| Final linear decode | ~37K | 1152 → p×p×2C = 1152 → 32 |
| Label embedding | ~1.2M | 1000 classes × 1152 dim lookup table |
| Timestep MLP | ~2.7M | Sinusoidal embed → 2-layer MLP → 1152 dim |
| Total DiT-XL/2 | 675M | All trained from scratch |
The overwhelming majority (99%+) of parameters sit in the 28 transformer blocks. The conditioning, embedding, and decode layers are negligible. This is what makes DiT so clean: scale the blocks and everything else is noise.
The paper calls this architecture DiT: Diffusion Transformer. The key finding is that DiTs exhibit the same clean scaling behavior that makes transformers so powerful in language: more Gflops = lower FID, with strong correlation (-0.93). This means you can improve image quality simply by making the model bigger, with no architectural changes needed.
The best model, DiT-XL/2, achieves a state-of-the-art FID of 2.27 on class-conditional ImageNet 256x256, outperforming all prior diffusion models while using fewer Gflops than pixel-space U-Net models like ADM.
DiT doesn't operate on raw pixels. It uses the Latent Diffusion Model (LDM) framework from Rombach et al. — the same framework behind Stable Diffusion. This is a two-stage approach:
A variational autoencoder is trained to compress images into a compact latent space. For a 256x256x3 RGB image, the VAE encoder E produces a latent z = E(x) with shape 32x32x4 — an 8x spatial downsampling with 4 channels. DiT uses the off-the-shelf VAE from Stable Diffusion, frozen during DiT training.
The forward diffusion process adds Gaussian noise to z:
The DiT model learns to reverse this process — given a noisy latent zt and timestep t, predict the noise ε. After denoising, the clean latent z0 is decoded back to an image via the VAE decoder: x = D(z).
Let's trace a single image through the full DiT pipeline, tracking every tensor shape:
Here's the complete forward pass of DiT, step by step. This is the core of the paper.
The noised latent zt (shape 32x32x4) is divided into non-overlapping patches of size p x p. Each patch is linearly embedded into a d-dimensional token vector. For patch size p, the number of tokens is:
Where I = 32 (the spatial dimension of the latent). With p=2, you get T = 256 tokens. With p=8, just T = 16 tokens. Halving p quadruples T and thus quadruples the Gflops (since self-attention is quadratic in sequence length).
Standard sine-cosine frequency-based positional embeddings are added to all tokens, exactly as in the original ViT. Nothing fancy.
The token sequence passes through N DiT blocks. Each block contains multi-head self-attention and a pointwise feedforward network (MLP), with layer normalization. The key modification is how conditioning information (timestep t and class label c) enters the block — via adaptive layer normalization with zero initialization (adaLN-Zero). We'll detail this in the next chapter.
After the final block, a final adaptive layer norm is applied, followed by a linear projection that maps each d-dimensional token to a p x p x 2C tensor (predicting both noise and diagonal covariance). The decoded tokens are rearranged back into spatial layout to produce the final noise prediction εθ and covariance Σθ.
Why patchify instead of pixel-level tokens? A 32×32 latent has 1024 spatial positions. Self-attention is O(N²), so 1024 tokens would cost 16× the compute of 256 tokens (p=2). Patchification is a compute-quality tradeoff. With p=2, you get 256 tokens — manageable for attention — while still preserving fine-grained spatial information within each patch via the linear projection.
Why predict both noise AND variance? Standard DDPM predicts only noise ε and uses a fixed variance schedule. But Nichol & Dhariwal (2021) showed that learning the variance improves log-likelihood and sample quality, especially with fewer sampling steps. DiT outputs 2C = 8 channels: 4 for noise prediction, 4 for an interpolation parameter v that blends between the fixed DDPM upper and lower variance bounds. The variance prediction adds zero parameters (just doubles the output projection) but meaningfully helps quality.
Why adaLN-Zero instead of cross-attention for conditioning? Cross-attention adds 15% Gflops overhead (extra QKV projections + attention computation) for a length-2 conditioning sequence. adaLN-Zero adds negligible overhead (a single MLP that regresses 6 scalars per block: γ1, β1, α1, γ2, β2, α2) yet achieves lower FID. The insight: conditioning on timestep and class doesn't need per-token flexibility (every patch should denoise the same amount), so a global modulation suffices.
| Model | Layers N | Hidden dim d | Heads | Gflops (p=4) |
|---|---|---|---|---|
| DiT-S | 12 | 384 | 6 | 1.4 |
| DiT-B | 12 | 768 | 12 | 5.6 |
| DiT-L | 24 | 1024 | 16 | 19.7 |
| DiT-XL | 28 | 1152 | 16 | 29.1 |
With patch size p=2, DiT-XL/2 reaches 118.6 Gflops. The naming convention is DiT-{size}/{patch_size}.
A diffusion model needs to know two things beyond the noised input: the timestep t (how noisy is this?) and the class label c (what should this image be?). The question is: how do you inject this conditioning information into each transformer block?
DiT explores four approaches. All of them first embed t and c into vector representations using learned MLPs, then combine them (usually by summing). The difference is how this combined conditioning vector enters each transformer block.
Simply append the conditioning embeddings as two extra tokens in the input sequence. The transformer processes them alongside image tokens with no architectural changes. After the final block, remove the extra tokens. Gflops overhead: negligible.
Add a cross-attention layer after self-attention in each block. The image tokens attend to the conditioning tokens (a length-2 sequence of t and c embeddings). This is similar to how the original transformer decoder attends to encoder outputs, and how LDM conditions on text. Gflops overhead: ~15%.
Replace the standard learnable scale (γ) and shift (β) parameters in layer norm with ones regressed from the conditioning vector. Instead of learning fixed normalization parameters, the model computes them as a function of t + c. This applies the same transformation to all tokens (unlike cross-attention, which can apply different weights per token). Gflops overhead: minimal.
Same as adaLN, but with a critical addition: regress additional scaling parameters α that are applied immediately before each residual connection. The α parameters are initialized to zero, which means each DiT block starts as the identity function — it passes the input straight through. The network gradually learns to "turn on" each block during training.
Where γ, β, α are all regressed from the conditioning embedding via a shared MLP.
At 400K training steps, adaLN-Zero achieves roughly half the FID of in-context conditioning. Cross-attention is better than in-context but worse than adaLN-Zero, despite costing 15% more Gflops. Vanilla adaLN (without zero init) is also worse than adaLN-Zero, confirming that the zero initialization matters. All subsequent DiT models use adaLN-Zero.
The most important result in the paper isn't the final FID number — it's the scaling behavior. Peebles and Xie train 12 DiT models spanning 4 model sizes (S, B, L, XL) and 3 patch sizes (8, 4, 2), then plot FID against Gflops.
Across all 12 models, there is a -0.93 correlation between model Gflops and FID-50K at 400K training steps. This is remarkably clean — almost a straight line on a log-log plot. You can predict a DiT model's quality from its Gflops alone.
Here's a subtle but crucial point. When you decrease the patch size (say from p=4 to p=2), you quadruple the number of tokens and thus the Gflops. But the parameter count barely changes — the transformer weights are the same, you're just processing more tokens. Yet FID improves substantially. This means compute (Gflops), not parameter count, is the true driver of quality.
Models with similar Gflops achieve similar FID regardless of how they get there (bigger model vs. smaller patches). For example, DiT-S/2 and DiT-B/4 have similar Gflops and similar FID.
When you plot FID against total training compute (Gflops x batch size x steps x 3), larger models reach any given FID threshold with less total compute. A small model trained for a long time is eventually overtaken by a large model trained for fewer steps. This mirrors the compute-optimal scaling behavior seen in language models (Chinchilla).
A key subtlety: larger DiT models are more compute-efficient at reaching a given quality level. Consider what it takes to reach FID 50:
DiT-XL/2 uses more compute per step but reaches the target in fewer steps, with similar total compute. And unlike DiT-S, the XL model keeps improving beyond FID 50 — it hasn't saturated. This mirrors the "Chinchilla" finding in language models: it's better to train a large model for fewer steps than a small model for many steps.
After the scaling analysis, Peebles and Xie train their best model — DiT-XL/2 — for 7 million steps (up from the 400K used in ablations). The results speak for themselves.
With classifier-free guidance (cfg scale = 1.50):
| Model | FID ↓ | sFID ↓ | IS ↑ | Precision | Recall |
|---|---|---|---|---|---|
| ADM-G | 4.59 | 5.25 | 186.7 | 0.82 | 0.52 |
| ADM-G + ADM-U | 3.94 | 6.14 | 215.8 | 0.83 | 0.53 |
| LDM-4-G (cfg=1.50) | 3.60 | — | 247.7 | 0.87 | 0.48 |
| DiT-XL/2-G (cfg=1.50) | 2.27 | 4.60 | 278.2 | 0.83 | 0.57 |
DiT-XL/2 achieves 2.27 FID, beating the previous best of 3.60 from LDM-4 by a large margin. It also achieves the highest Inception Score (278.2) and a strong balance between Precision (0.83) and Recall (0.57).
DiT-XL/2 also sets a new state-of-the-art at 512x512 resolution with an FID of 3.04, outperforming ADM-G + ADM-U (3.85 FID) while being substantially more compute-efficient.
DiT-XL/2 uses 118.6 Gflops per forward pass. Compare this to:
DiT achieves better FID than all of these despite comparable or lower compute cost.
The scaling analysis reveals clear degradation patterns:
DiT's strong results depend on classifier-free guidance (CFG), a technique that dramatically improves sample quality at the cost of some diversity. Let's understand how it works.
During sampling, you want images x where the class probability p(c|x) is high — if you asked for "golden retriever," the image should clearly look like a golden retriever, not an ambiguous blob. By Bayes' rule:
The gradient of the class-conditional score minus the unconditional score points toward images that strongly belong to class c.
Classifier-free guidance modifies the noise prediction during sampling:
Where s > 1 is the guidance scale (s = 1 recovers standard sampling). Each sampling step requires two forward passes: one conditioned on c, one unconditional (with a learned null embedding ∅). The unconditional pass is enabled by randomly dropping the class label during training (replacing it with ∅).
Higher guidance scale s pushes the model toward higher-fidelity but lower-diversity samples. For DiT-XL/2:
CFG reduces FID by over 4x but slightly reduces diversity (recall drops from 0.67 to 0.57). This tradeoff is well-known and consistent across all diffusion model architectures.
With CFG, every sampling step requires two forward passes (one conditioned, one unconditional). The full cost:
Compare to ADM (pixel-space U-Net): 250 steps × 2 × 1120 = 560,000 Gflops — nearly 10x more expensive for worse results. This is the fundamental efficiency win of latent diffusion: the 9.4x compute reduction from working in latent space compounds across all 250 steps.
DiT's success isn't just about one paper's results. It's about why transformers are fundamentally better suited as a backbone for scaling generative models. Let's unpack the structural advantages.
Transformers have two orthogonal axes for scaling: depth (number of layers N) and width (hidden dimension d). Doubling either increases Gflops predictably. U-Nets have channels, resolution levels, attention layers at specific resolutions — but these interact in complex, non-linear ways. There's no clean "make it 2x bigger" knob.
Modern GPUs and TPUs are optimized for the dense matrix multiplications that dominate transformer computation (attention QKV projections, FFN layers). U-Nets mix convolutions at various spatial resolutions with attention at lower resolutions — this heterogeneous compute profile is harder to optimize and often leaves hardware underutilized.
If your image generator, text model, and video model all use transformers, you can share training recipes, optimization tricks, and infrastructure. The community's collective knowledge about transformer training (learning rate schedules, initialization, regularization, mixed precision) transfers directly. With U-Nets, every insight had to be re-discovered within the diffusion community.
Transformers naturally handle variable-length sequences and cross-attention over conditioning tokens. This makes it straightforward to condition on text embeddings, multiple images, or any other modality. U-Nets require bespoke injection points for each conditioning type.
DiT sits at a pivotal point in the evolution of generative models. Here's how it connects to the broader landscape.