22. Diffusion Transformers

Diffusion Transformers (DiT)

In 2023, Peebles and Xie¹ took the latent diffusion model of chapter 20, removed its U-Net and put a plain Vision Transformer in its place: cut the noisy latent into patches, run Transformer blocks over the tokens, turn the tokens back into a latent. The result beat the best U-Nets of the time on class-conditional ImageNet, and the paper's real finding was a different one: quality improved predictably with compute, the way it does for language models.

Three years later that swap has won where the money is. The leading open text-to-image and text-to-video models of 2024–2026 — SD3, FLUX, Qwen-Image, Z-Image, HunyuanVideo, Wan, LTX-2 — are descendants of that block, and so are the generators inside robot policies and 3-D asset models. The main exception is autoregressive image generation (chapter 23), and the closed systems (Sora 2, Veo 3, Midjourney) do not publish their architecture.

This chapter takes the block apart: how an image becomes tokens, how the timestep gets in (adaLN-Zero), how the prompt gets in (cross-attention, joint attention, one stream), how positions are encoded once the grid has a time axis, and where the compute actually goes — which, for video, is the whole story.

Slides

Click the slides and use ← → to move through them; F goes full screen and P opens the presenter view, with the notes. The simulators of this page run inside the slides. Open the slides in a new tab

From U-Net to Transformer

The denoiser of SD 1.x and SDXL is a U-Net: convolutions that halve the resolution level by level and bring it back, skip connections between matching levels, and attention — with the text entering by cross-attention — interleaved with the convolutions. A DiT has none of that hierarchy. It works at one resolution from the first block to the last, and every block sees every token.

Everything else in the pipeline stays: the VAE, the text encoder, the sampler, classifier-free guidance. Only the network that predicts the noise — or, since SD3, the velocity of flow matching — changes.

Step 1 — Patchify: a latent as a sequence of tokens

A DiT works on the latent of a VAE, not on pixels. A latent of shape \(C \times H \times W\) is cut into non-overlapping \(p \times p\) patches; each patch, \(p^2 C\) numbers, is flattened and projected by one linear layer to the model width \(d\):

It is the ViT patch embedding exactly — a Conv2d with kernel and stride \(p\). The new part is the way back. A generator has to output a latent of the input's shape, so a last linear layer maps each token to \(p^2 C\) numbers and unpatchify puts them back on the grid: a reshape and a permute, no parameters.

The patch size is the DiT's main compute dial, and the paper turned it on purpose. Halving \(p\) quadruples \(N\) while the parameter count barely moves (only \(E\) and the output layer depend on \(p\)). At 400K training steps, DiT-XL at \(p = 8, 4, 2\) costs 7.4, 29.1 and 118.6 Gflops per forward pass (billions of operations, counting one multiply-add as one) and reaches FID 106.4, 43.0 and 19.5¹ (the Fréchet Inception distance between 50,000 samples and the real images; lower is better): same parameters, more tokens, much better images. That is the observation the whole scaling argument rests on — what predicts quality is the compute per forward pass, whether it comes from width, depth or tokens. The lab reproduces the trend at \(28 \times 28\).

Step 2 — The block, and how the timestep gets in: adaLN-Zero

A denoiser must know how noisy its input is, and for a class-conditional model which class to draw. DiT folds both into one vector,

where \(\text{emb}(t)\) writes \(t\) as sines and cosines of many frequencies (as for word positions) and \(\text{emb}(y)\) is the row of a learned table for class \(y\). Every block reads \(c\) through its normalization layers. A LayerNorm normally ends in a learned scale and shift, the same for every input. adaLN (adaptive LayerNorm) removes them and lets a linear layer regress them from \(c\); adaLN-Zero also regresses a gate \(\alpha\) for each residual branch. One linear layer per block produces six vectors of size \(d\):

Here \(x\) holds the block's \(N\) tokens of width \(d\), \(\text{SiLU}(z) = z\,\sigma(z)\), LN is a LayerNorm with no parameters of its own, \(\odot\) multiplies elementwise and the same six vectors serve every token of the image. Step through the block: where \(c\) comes from, where each of the six vectors goes and what the initialization does to them.

The "Zero" is the initialization: \(W = 0\) and \(b = 0\), so at the start \(\alpha = 0\) and every block is the identity. The network begins as a stack of no-ops (the output layer is zero-initialized too, so it predicts zero) and each block learns to add something. It is the same safe-addition trick as the residual connections of chapter 9, ControlNet's zero convolutions and LoRA's \(B = 0\).

The same block as the lab writes it, run on numbers small enough to follow: a \(1 \times 4 \times 4\) latent cut into four \(2 \times 2\) patches, \(d = 4\), one line of DiTBlock.forward at a time. Step through the trained case, then pick at initialization and watch both gates hold \(z\) still:

The DiT paper tried four ways of feeding \(c\) to the block, at equal training budget: in-context (the embeddings of \(t\) and \(y\) appended to the sequence as two more tokens), cross-attention (the tokens read them in an extra attention layer), adaLN and adaLN-Zero.

adaLN-Zero won by a wide margin — FID 19.5 against 35.2 for in-context tokens at 400K steps — and at the lowest compute (tied with adaLN), because the modulation is computed once per image, not once per token. The lab runs three of the four designs on Fashion-MNIST and finds the same order: FD 26.66 for in-context tokens, 19.89 for adaLN and 17.15 for adaLN-Zero (FD is FID's formula computed on the features of a small Fashion-MNIST classifier instead of Inception's).

Step 3 — How the prompt gets in: cross-attention, joint attention, one stream

The DiT paper only had a class label. A text-to-image model needs a whole prompt — a sequence of hundreds of token embeddings from a text encoder — and three designs are in use.

Cross-attention. The image tokens attend to themselves, then read the text in a second attention layer whose keys and values come from the prompt: \(\text{softmax}(Q_{\text{img}} K_{\text{txt}}^\top / \sqrt{d_h})\, V_{\text{txt}}\). This is what the SD U-Net did, and what PixArt-α⁵, Wan⁷ and Movie Gen¹⁰ do. The text is a read-only input: it leaves every layer exactly as it came in.

Joint attention, two streams (MMDiT, multimodal DiT). SD3² concatenates image and text tokens into one sequence and runs one attention over it:

where \([\,\cdot\,;\,\cdot\,]\) concatenates along the sequence and \(d_h\) is the width of one attention head, but each modality keeps its own weights — its own Q, K, V projections, MLP and modulation. Text tokens now read the image and are rewritten by it, layer after layer. In the SD3 paper's comparison at equal size, MMDiT beat the cross-attention DiT, a U-ViT and a vanilla DiT that just appends the text tokens; a U-ViT is a ViT with long skip connections between shallow and deep blocks.

Joint attention, one stream. Once the two modalities have been aligned for a while, the separate weights can go: the concatenated sequence passes through a plain Transformer block with one set of weights. FLUX.1 runs 19 two-stream blocks and then 38 single-stream ones³; HunyuanVideo does 20 then 40⁸; FLUX.2 does 8 then 48⁶; Z-Image uses one stream from the first block to the last²¹.

Step through one two-stream block, and at the end the single-stream block that replaces it:

The panel below looks inside the attention: which blocks of the score matrix each design computes.

Two things the panel makes visible. First, joint attention is almost free in compute: the text is short next to the image, so the extra blocks of the score matrix add 12.5% to the scores of cross-attention at FLUX's 4,096 + 512 tokens and under 1% for a 720p video. What costs is the second set of weights of the two-stream blocks — which is exactly what the single-stream blocks save. Second, the choice is not settled: Wan keeps cross-attention at 14B parameters, Z-Image goes single-stream at 6B, and both are competitive.

Positions: from a sine table to RoPE in two and three dimensions

Attention is blind to order (chapter 13), so the tokens need positions. The original DiT adds a fixed 2-D sine-cosine table to the tokens after patchify — half the channels encode the row, half the column. It works at the training resolution and has to be interpolated for any other.

Since FLUX.1 most new models use rotary embeddings (RoPE²⁷) applied per axis instead (SD3 still adds a 2-D table). Each attention head's channels are split into groups, and each group rotates \(q\) and \(k\) by an angle proportional to one coordinate; the score \(q^\top k\) then depends only on the difference of positions along each axis.

Relative positions are what make a DiT indifferent to resolution and aspect ratio: a \(1024 \times 768\) image is just a different grid of coordinates, and a video is the same grid with a third axis. The text tokens get positions too — Qwen-Image places them on the diagonal of the image grid so the two share one coordinate system²⁰.

What DiT actually bought: a scaling law

The contribution of the DiT paper is not "a Transformer also works here". It is that swapping the backbone gave image generation the property language modelling already had: a curve you can plan against¹. Across four model sizes (S, B, L, XL) and three patch sizes, FID fell steadily as the Gflops of a forward pass rose, with no sign of saturation in the range tested. Parameters alone predicted it poorly: DiT-XL/8, with 676M parameters, ended at FID 106.4, far behind DiT-B/2 with 130M at 43.5; and DiT-S/2 (33M) and DiT-B/4 (130M), both at about 6 Gflops, reached 68.40 and 68.38. SD3 repeated the study at text-to-image scale, up to 8B parameters, and found validation loss falling smoothly with size and training steps, and tracking GenEval (an automatic benchmark of prompt following: counts, colours, positions) and human preference².

The twelve DiT models of the paper, plotted both ways (switch the horizontal axis):

That curve is what justified spending on the next generation:

Two things in the table. SD3's design ties width to depth (\(d = 64 \times\) blocks), so its 8B model is the same design at 38 blocks and 2,432 channels. And the token count did not grow with the models: at 1024² it has been 4,096 since 2024. The models got wider and deeper; the sequence was held fixed by 16-channel latents and \(p = 2\) (FLUX.2 moved to a 32-channel latent at the same token count).

Three properties came along with the swap, and each mattered more than the FID number:

• One architecture for everything. The same block, kernels and distributed-training machinery as an LLM transfer unchanged: FlashAttention (a fused attention kernel), sequence parallelism (one sequence split across GPUs), FSDP (weights sharded across GPUs), activation checkpointing (recomputing activations instead of storing them).
• A modality is more tokens. Text, video frames, audio, a reference image to edit: each is a sequence concatenated to the image tokens, not a new architecture.
• Resolution is a sequence length. With RoPE there is no architectural commitment to 512 or 1024; you change \(N\).

Video: the same block over spacetime patches

A video DiT changes the input, not the block. A 3-D VAE compresses the clip in time as well as in space — the common choice is 4× in time, 8 × 8 in space, 16 channels — and patchify cuts the latent into \(1 \times 2 \times 2\) spacetime patches. Sora's technical report described exactly this ("spacetime patches" fed to a diffusion transformer) without giving sizes¹¹; the open models give them:

The token count is \(T \cdot \frac{H}{2} \cdot \frac{W}{2}\); the latent has \(1 + (\text{frames} - 1)/4\) time steps (Mochi compresses time by 6 and Movie Gen by 8, both with the same \(1 \times 2 \times 2\) patch). Early video models factorized attention into a spatial pass inside each frame and a temporal pass along each pixel, to keep the cost down. The current ones all use full 3-D attention, and CogVideoX gives the reason: with factorized attention, the head of a person in frame \(i+1\) cannot attend directly to the head in frame \(i\) — it moved, so it is no longer at the same pixel — and consistency suffers⁹.

The same design now carries other signals in the same sequence: LTX-2 generates video and audio together with a 14B video stream and a 5B audio stream joined by cross-attention¹³, and Wan 2.2 splits denoising between two 14B experts, one for the high-noise steps and one for the low-noise ones, 27B in total with 14B active per step²⁴.

Where the compute goes

The cost of a Transformer block from chapter 13 applies unchanged:

multiply-adds per block, so attention's share of the work is \(N/(N + 6d)\).

An image DiT is not attention-bound; a video DiT is. At 1024² the projections and MLPs still do 80% of the work. At 720p video the attention scores take almost all of it — and the measurement agrees with the formula: generating a 5-second 720p clip with HunyuanVideo took 945 s, 800 of them in attention¹⁵. (The formula assumes a \(4d\) MLP; Wan's is \(2.7d\), which raises its share a little.)

That split decides which efficiency tricks matter where:

For images the first two rows dominate in practice; for video, the third is where the large gains are.

Where it is used, October 2026

And what replaced the U-Net where: for text-to-image, between SDXL (2023, U-Net) and SD3 (2024, MMDiT); for video, in 2024 — the U-Net video models of 2023 (Stable Video Diffusion) have no successors in U-Net form. The convolutional parts did not disappear: the VAE that every one of these models runs before and after the Transformer is a CNN, as are many super-resolution stages.

Where the design is going

Implementation reference

The lab builds this model and trains it. It is a complete class-conditional DiT with adaLN-Zero, written to be read; the comments give the shapes.

import math, torch, torch.nn as nn, torch.nn.functional as F

def t_embed(t, d):                                  # t in [0, 1] -> (B, d) sinusoidal features
    f = torch.exp(-math.log(10000) * torch.arange(d // 2, device=t.device) / (d // 2))
    a = 1000 * t[:, None] * f[None]
    return torch.cat([a.sin(), a.cos()], dim=1)

class DiTBlock(nn.Module):
    def __init__(self, d, heads):
        super().__init__()
        self.heads = heads
        self.ln1 = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)   # no learned γ, β:
        self.ln2 = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)   # they come from c
        self.qkv, self.proj = nn.Linear(d, 3 * d), nn.Linear(d, d)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(approximate='tanh'), nn.Linear(4 * d, d))
        self.mod = nn.Linear(d, 6 * d)                                   # c -> β1 γ1 α1 β2 γ2 α2
        nn.init.zeros_(self.mod.weight); nn.init.zeros_(self.mod.bias)   # the "Zero": α = 0

    def attn(self, x):                                                   # (B, N, d) -> (B, N, d)
        B, N, d = x.shape
        q, k, v = self.qkv(x).view(B, N, 3, self.heads, d // self.heads).permute(2, 0, 3, 1, 4)
        return self.proj(F.scaled_dot_product_attention(q, k, v).transpose(1, 2).reshape(B, N, d))

    def forward(self, x, c):                                             # x (B, N, d), c (B, d)
        b1, g1, a1, b2, g2, a2 = self.mod(F.silu(c))[:, None].chunk(6, dim=-1)
        x = x + a1 * self.attn(self.ln1(x) * (1 + g1) + b1)
        return x + a2 * self.mlp(self.ln2(x) * (1 + g2) + b2)

class DiT(nn.Module):
    def __init__(self, img=28, ch=1, p=4, d=96, depth=4, heads=4, n_classes=10):
        super().__init__()
        self.p, self.ch, self.g = p, ch, img // p                        # g × g patches
        self.patch = nn.Conv2d(ch, d, kernel_size=p, stride=p)           # patchify + projection E
        self.pos = nn.Parameter(torch.zeros(1, self.g ** 2, d), requires_grad=False)  # fill with 2-D sin-cos
        self.t_mlp = nn.Sequential(nn.Linear(d, d), nn.SiLU(), nn.Linear(d, d))
        self.y_emb = nn.Embedding(n_classes + 1, d)                      # last index = "no class", for CFG
        self.blocks = nn.ModuleList([DiTBlock(d, heads) for _ in range(depth)])
        self.ln = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)
        self.out_mod = nn.Linear(d, 2 * d)                               # final shift and scale
        self.out = nn.Linear(d, p * p * ch)                              # token -> one patch
        for m in (self.out_mod, self.out):
            nn.init.zeros_(m.weight); nn.init.zeros_(m.bias)             # the model starts predicting 0
        self.d = d

    def unpatchify(self, z):                                             # (B, N, p·p·C) -> (B, C, H, W)
        B, g, p, C = z.shape[0], self.g, self.p, self.ch
        z = z.view(B, g, g, p, p, C).permute(0, 5, 1, 3, 2, 4)           # B, C, row, p, col, p
        return z.reshape(B, C, g * p, g * p)

    def forward(self, x, t, y):                                          # x (B, C, H, W), t (B,), y (B,)
        z = self.patch(x).flatten(2).transpose(1, 2) + self.pos          # (B, N, d)
        c = self.t_mlp(t_embed(t, self.d)) + self.y_emb(y)               # (B, d)
        for blk in self.blocks:
            z = blk(z, c)
        shift, scale = self.out_mod(F.silu(c))[:, None].chunk(2, dim=-1)
        z = self.ln(z) * (1 + scale) + shift
        return self.unpatchify(self.out(z))                             # velocity, same shape as x

Training is the flow-matching loss of chapter 21 with \(t = 0\) as noise: draw \(x_0 \sim \mathcal{N}(0, I)\), \(t \sim U(0, 1)\), form \(x_t = (1-t)x_0 + t x_1\) and regress model(x_t, t, y) on \(x_1 - x_0\), replacing the label by the "no class" index 10% of the time so the same network can do classifier-free guidance.

Lab: a DiT, built and measured

Everything on this page that can be measured at small scale is measured in a Colab notebook: a class-conditional DiT built from scratch and trained with flow matching on Fashion-MNIST (28 × 28, ten classes of clothing). It needs a GPU: in Colab, Runtime → Change runtime type → T4 GPU (the free tier is enough; the reference run of every cell took 12 minutes on an Apple M2 laptop GPU). Seven sections, each ending in a checkpoint with a check value:

1. patchify and unpatchify by hand: shapes, the round trip and the strided convolution;
2. the adaLN-Zero block: identity at initialization, measured, and where its parameters go;
3. the whole DiT, its parameter count and the cost formula against the clock;
4. training with flow matching and sampling with classifier-free guidance, judged by a small classifier: how often the sample is the class you asked for, and a Fréchet distance in its feature space;
5. the conditioning ablation: in-context tokens, adaLN and adaLN-Zero at equal steps;
6. the patch-size ablation: \(p = 14, 7, 4\) — tokens, time per step and quality;
7. your own DiT under a budget, with a table to fill in.

Open the lab in Colab — save a copy to your Drive (File → Save a copy in Drive) to run and edit it.

Key takeaways

1. A DiT is a ViT over noisy latent patches: patchify, add positions, \(L\) identical blocks at one resolution, unpatchify to a velocity (or noise) of the input's shape.
2. The patch size is a compute dial. Halving \(p\) quadruples the tokens at almost the same parameter count, and quality follows the compute: DiT-XL went from FID 106 to 19.5 between \(p = 8\) and \(p = 2\).
3. adaLN-Zero conditions by regressing each block's normalization scale, shift and residual gate from \(c\), with the gate starting at zero so every block starts as the identity. It beat in-context tokens and cross-attention at the lowest compute; its parameter cost is why later models share the modulation across blocks.
4. Text gets in by cross-attention (read-only text: Wan, PixArt), joint attention with two streams (MMDiT: SD3, Qwen-Image) or one stream (FLUX's later blocks, Z-Image). Joint attention costs little compute; the second weight set is what costs.
5. Positions are RoPE per axis — row and column for images, plus time for video — which is what makes resolution, aspect ratio and duration a matter of sequence length.
6. The contribution is a scaling curve: quality follows Gflops, which is what justified 12B and 32B image models.
7. Images are MLP-bound, video is attention-bound: the attention share is \(N/(N + 6d)\), 20% for FLUX at 1024², about 87% for 720p video — which is why video is where sparse attention and caching matter.