Diffusion Transformers (DiT)

Take one 675M-parameter transformer and run it on a 32×32×4 latent twice: once cut into 4×4 patches, once into 2×2 patches. The parameter count barely moves. The forward pass goes from 29.1 to 118.6 Gflops, and the sample quality improves by a wide margin. That observation, from Peebles & Xie, 2023, Scalable Diffusion Models with Transformers, ICCV, arXiv:2212.09748, goes a long way to explaining why so many recent image and video generators left the U-Net behind: the transformer turned the backbone into something you scale by a known rule rather than by hand-tuned architecture search.

The architecture, stripped down

DiT operates inside the latent space of a pretrained autoencoder, which the neighbouring concept on latent diffusion covers. A 256×256 image becomes a latent \(z\) of shape \(I \times I \times C = 32 \times 32 \times 4\). Patchify cuts it into non-overlapping \(p \times p\) squares and linearly embeds each one, giving

tokens of width \(d\). With \(p = 8, 4, 2\) that is 16, 64 or 256 tokens. Standard sine-cosine positional embeddings are added, a stack of \(N\) transformer blocks processes the sequence, and a linear decoder maps each token back to a \(p \times p\) patch of predicted noise (and a diagonal covariance). Apart from patchify and the 2D positional embeddings, nothing in the stack knows it is looking at an image.

Halving \(p\) quadruples \(T\), so the attention-matrix term grows sixteenfold and the per-token MLP cost fourfold, while the embedding layer is the only parameter that changes. This decoupling of compute from parameters is what let the authors ask a clean question: which of the two predicts quality?

Conditioning: adaLN-Zero

A diffusion network needs the timestep \(t\) and, for class-conditional models, a label \(c\). DiT tested four ways to inject them: append them as extra tokens (in-context), add a cross-attention layer, replace layer norm with adaptive layer norm, and adaptive layer norm with a zero-initialised residual scale.

Adaptive layer norm regresses the scale and shift from the summed embedding \(e = e_t + e_c\) instead of learning them as free parameters:

adaLN-Zero also regresses a per-dimension scale \(\alpha(e)\) applied just before each residual addition, and initialises the MLP producing \(\alpha\) to output zeros. Every block therefore starts as the identity function, which is the same trick as zero-initialising the last batch-norm scale in a ResNet. In the paper's ablation at 400K steps, adaLN-Zero reached roughly half the FID of in-context conditioning, while cross-attention added about 15% more Gflops and still lost.

There is a limit hiding in that result. Adaptive normalisation applies the same function to every token, conditioned on one vector. That suffices for a class label. It cannot carry a 77-token prompt.

Scaling in Gflops

Twelve models, sizes S, B, L and XL crossed with \(p \in \{8, 4, 2\}\), produced a correlation of \(-0.93\) between transformer Gflops and FID-50K at 400K steps. Models with matching compute but different shapes, such as DiT-S/2 and DiT-B/4, landed at similar FID. Parameters did not uniquely determine quality; compute per forward pass did.

Two numbers make the claim concrete. DiT-XL/2 trained for 7M steps reached FID 2.27 on class-conditional ImageNet 256×256 with classifier-free guidance, below the 3.60 of the latent U-Net LDM, at 118.6 Gflops against pixel-space ADM's 1120. And extra sampling steps did not rescue a small model: DiT-L/2 at 1000 steps spent 80.7 Tflops per image for FID-10K 25.9, while DiT-XL/2 at 128 steps spent 15.2 Tflops and scored 23.7. Compute spent in the network beat compute spent in the sampler, which should temper any plan to cover a weak backbone with a better solver from the samplers concept.

MM-DiT: text as a second stream

Stable Diffusion 3 extended the recipe to text-to-image (Esser et al., 2024, Scaling Rectified Flow Transformers for High-Resolution Image Synthesis, arXiv:2403.03206). The timestep and a pooled text vector still drive adaptive modulation, but because a pooled vector keeps only coarse information, text tokens and 2×2 image patches are also concatenated into one sequence for attention. Each modality keeps its own weights for the projections and MLPs; only the attention operation is joint. On CC12M the authors found this MM-DiT beat plain DiT with concatenation, a cross-attention variant and the U-Net/transformer hybrid UViT, and they scaled it to 8B parameters. They also added RMSNorm on queries and keys after attention logits grew uncontrollably in the last blocks during high-resolution fine-tuning with mixed precision.

The literature does not fully agree on the lesson. DiT concluded that the U-Net's inductive bias is not crucial to diffusion performance. The SD3 comparison found UViT, which keeps long skip connections, learning faster early in training even though MM-DiT won in the end, so whether the skips help depends on how long you plan to train.

When it breaks

Token count is a quadratic tax. At 512×512 the same XL/2 design processes 1024 tokens and 524.6 Gflops. Video multiplies tokens by frame count, so patch size and latent compression become the main levers on cost, and aggressive latent compression moves quality risk into the autoencoder.

Compute scaling is not a free lunch on memory. The Gflops law says small patches win; attention memory at small patches is what stops you.

Global conditioning cannot localise. adaLN modulates every token identically, so any per-token instruction (layout, spatial text rendering, region prompts) needs attention over a sequence, which reintroduces the cost adaLN-Zero was chosen to avoid.

Instability arrives late. The attention-logit growth SD3 reports appeared at scale and at higher resolution, after smaller runs looked healthy, which makes such problems expensive to find.