Scalable Diffusion Models with Transformers

The problem

Every diffusion model up to this point used a U-Net denoiser, and for pixel-to-pixel prediction its inductive biases are genuinely useful: an encoder compresses through a bottleneck, a decoder reconstructs at full resolution, and skip connections let fine spatial detail bypass the bottleneck. The trouble is scaling. A U-Net does not grow gracefully — the skip connections fix where information may flow, and the multi-resolution structure bakes in assumptions about how features at different scales combine, so adding parameters means widening hand-designed stages rather than scaling a single uniform primitive. Meanwhile the transformer had repeatedly shown that a stack of identical blocks with no architectural opinion about its input scales more predictably than anything purpose-built. The question DiT asks is whether diffusion is an exception. It is not.

Patchify

Turning a latent into something a transformer can read borrows one step from Vision Transformers. The 8×-downsampling VAE from latent diffusion turns a $256 \times 256$ image into a $32 \times 32 \times 4$ latent, which is cut into non-overlapping patches. With patch size $p = 2$, the $32 \times 32$ grid becomes a $16 \times 16$ array of patches, each flattened and linearly projected to the model width $d$, giving a sequence of tokens.

$$N = \left(\frac{32}{p}\right)^2 = 256 \quad \text{tokens for } p = 2$$

Add a 2D positional embedding and from there it is a standard transformer — self-attention and feed-forward blocks, no skip connections, no resolution hierarchy, every token attending to every other from the first layer.

Conditioning via adaLN-Zero

A transformer has no cross-attention by default, so DiT must route in the two conditioning signals every diffusion model needs: the timestep $t$ and the class label $y$. Both are embedded and summed into a single conditioning vector $c$, and instead of learning fixed scale and shift inside each LayerNorm, a small per-block MLP predicts them from $c$ — adaptive layer normalization, the transformer analogue of the AdaGN that Dhariwal and Nichol introduced for U-Nets.

$$\mathrm{AdaLN}(h, c) = \gamma(c) \odot \frac{h - \mu(h)}{\sigma(h)} + \beta(c)$$

The best variant goes one step further. adaLN-Zero also uses $c$ to predict a per-block residual scale $\alpha(c)$ that gates each sublayer's contribution, and it zero-initializes that scale so every block starts as an identity function.

$$h = h + \alpha_1(c)\,\mathrm{MSA}\big(\mathrm{AdaLN}(h, c)\big), \qquad h = h + \alpha_2(c)\,\mathrm{FFN}\big(\mathrm{AdaLN}(h, c)\big)$$

Starting each residual block at the identity makes the full network an identity map at initialization, which stabilizes training of deep diffusion transformers — the same trick that zero-initialized residual gates provide elsewhere. The adaLN-Zero ablation was the single largest architectural win in the paper.

Scaling and results

With the architecture fixed, scaling is the experiment. DiT runs it across four sizes — DiT-S at 33M parameters up to DiT-XL at 675M — holding everything else constant. Trained on ImageNet at $256 \times 256$ and measured by FID, the result is the curve the field has seen in language modeling over and over: FID falls monotonically as the model grows and as it trains longer, with no sign of the architecture fighting back. Holding parameters fixed and instead shrinking the patch size — which raises the token count and compute per forward pass — improves FID the same way, so the gains track total compute (GFLOPs) rather than any single knob.

The headline number lands where it matters: DiT-XL/2 reaches FID 2.27 on ImageNet 256×256, the best reported at the time, beating the prior U-Net diffusion models while doing all its work in the VAE's compressed latent space. That figure is reported with classifier-free guidance — DiT trains the class conditioning with the same conditioning dropout CFG requires — and 2.27 is the best point on the guidance sweep, at scale $w = 4.0$. None of the diffusion mathematics changed; only the function that predicts the noise is different.

Why it matters

DiT reframes what diffusion progress looks like. The earlier advances each found a better idea — a faster sampler, a cheaper space, a stronger way to condition. DiT's contribution is to show that once the backbone is a transformer, the next gains come from the same lever that drives language models: make it bigger and train it longer. That settled the architecture question, and the lineage followed. The transformer backbone slots directly into the infrastructure the field had spent years building — long context, flash and sparse attention, the whole tooling stack — which is why Sora built its video architecture on a diffusion transformer, and why Stable Diffusion 3 and Flux dropped the U-Net for an MMDiT, a multimodal DiT in which image and text tokens attend to each other jointly. DiT is the hinge between the U-Net era of diffusion and the transformer era that now dominates.