Artificial Neural Networks and Deep Learning · Insper

Diffusion Transformers

A ViT where the U-Net was — and what it bought

One part of the pipeline changes

Stable Diffusion = VAE + text encoder + U-Net + sampler + guidance. You have 20× more compute. Which part do you replace?

  • The denoiser. Everything else stays: VAE, text encoder, sampler, classifier-free guidance.
  • Peebles & Xie (2023): cut the noisy latent into patches, run a plain ViT, turn the tokens back into a latent.
  • Beat the best U-Nets of the time on ImageNet — and showed something more useful: quality follows compute, predictably.
  • Since 2024: SD3, FLUX, Qwen-Image, HunyuanVideo, Wan, LTX-2 — the leading open image and video models are its descendants.

Simulator · U-Net, DiT, MMDiT

Patchify, and the way back

C × H × W = channels × height × width of the latent (the VAE shrinks each side by 8: 256/8 = 32) · p = patch side · d = model width · E = the learned projection

model latent tokens token
DiT-XL/2, 256² 2
FLUX.1-dev, 1024² 2
our lab, 28² pixels 4
  • In: a Conv2d with kernel = stride = — the ViT patch embedding.
  • Out: a linear layer to numbers per token, then unpatchify — a reshape and a permute, no parameters.

Exercise · follow the tensor through the lab's DiT

Input (B, 1, 28, 28) = (images, channels, height, width): a batch of grey Fashion-MNIST images. Ignore and fill in before turning the slide.

step output
Conv2d(1, 96, 4, stride=4) ?
flatten(2).transpose(1, 2) ?
+ positions, 4 adaLN-Zero blocks ?
Linear(96, 16) ?
unpatchify ?

Conv2d(channels in, number of filters, kernel side, stride)

With kernel = stride = , the output grid is .

flatten(2) merges height and width into one axis; transpose(1, 2) puts that axis before the channels.

Answer in channels × height × width, or in tokens × width.

Follow the tensor · answer

1×28×28patchifyconv 4×4, s = 496×7×7flatten7 · 7 = 4949×964 blocksadaLN-Zero49×96outputLinear(96, 16)49×16unpatchifyreshape1×28×28

step output where it comes from
Conv2d(1, 96, 4, stride=4) 96 × 7 × 7 96 filters = · 28/4 = 7 per side
flatten(2).transpose(1, 2) 49 × 96 7 · 7 = 49 tokens, each of numbers
+ positions, 4 adaLN-Zero blocks 49 × 96 each block adds to the residual: keeps the shape
Linear(96, 16) 49 × 16 numbers per token = one 4 × 4 patch
unpatchify 1 × 28 × 28 the 49 patches back on the 7 × 7 grid: 7 · 4 = 28 per side

Simulator · tokens and cost

Same parameters, more tokens

DiT-XL (675M parameters), ImageNet 256² (32 × 32 latent), 400K training steps:

patch tokens Gflops FID-50K ↓
8 16 7.4 106.4
4 64 29.1 43.0
2 256 118.6 19.5

Gflops = billions of operations in one forward pass · FID-50K = Fréchet distance, in Inception features, between 50,000 samples and the real images (lower is better) · DiT-XL/2 = size XL, patch 2

  • Only and the output layer depend on : the parameter count does not move. Half the , 4× the tokens, ~4× the Gflops.
  • Parameters alone predict badly: DiT-S/2 (33M) and DiT-B/4 (130M) both cost ~6 Gflops — FID 68.40 and 68.38.

Quality follows the compute per forward pass — from width, depth or tokens.

Twelve models, one curve

The 12 models of the DiT paper, sizes S, B, L and XL at patch sizes 8, 4 and 2: FID-50K after 400K training steps, on a log scale, against the Gflops of one forward pass or against the parameter count. Against Gflops the points fall on one band; against parameters each size is a vertical line. FID-50K, lower is better (log scale) MODEL SIZE S · 12 blocks, d = 384 B · 12 blocks, d = 768 L · 24 blocks, d = 1024 XL · 28 blocks, d = 1152 PATCH SIZE p (32 × 32 LATENT) p = 8 · 16 tokens p = 4 · 64 tokens p = 2 · 256 tokens POINT

The block reads the condition through its norms

t = noise level (timestep) · y = class · c, β, γ, α: vectors of size d — β shifts, γ scales, α is the gate · W maps d to 6d · emb = sines and cosines for t, a learned table for y · SiLU(z) = z·σ(z) · LN = LayerNorm with no parameters · ⊙ = elementwise product

  • LayerNorm loses its learned scale and shift; one linear layer regresses them from (adaLN).
  • adaLN-Zero adds a gate per residual branch — and initializes , .

Where c goes, in one block

adaLN-Zero. Left, computed once per image: the timestep t and the class y become one vector c, and a linear layer W of each block turns SiLU of c into six vectors. Right, one block run on every token: LayerNorm, scale by one plus gamma 1 and shift by beta 1, attention, gate alpha 1, residual sum; then the same around the MLP with gamma 2, beta 2 and alpha 2. With W equal to zero the gates are zero and the block is the identity. CONDITION · ONCE PER IMAGE ONE BLOCK · ON EVERY TOKEN x h Timestep tsin/cos features → MLP Class ya row of a learned table c: d numbersMLP(emb(t)) + emb(y) W SiLU(c) + bW: d → 6dits own W in each block six vectors of d γ₁, β₁ α₁ γ₂, β₂ α₂ Tokens x: N × d LN (no γ, β) (1 + γ₁) ⊙ · + β₁ Attention ⊙ α₁ LN (no γ, β) (1 + γ₂) ⊙ · + β₂ MLP ⊙ α₂ x′: N × d, next block

Why zero?

At initialization, what is ? Not "small" — which number?

  • Exactly 0: and give , so — and the same in the MLP. Every block starts as the identity.
  • The output layer is zero-initialized too: the network starts by predicting zero, and each block learns to add something.
  • Same trick as residual connections, ControlNet's zero convolutions and LoRA's .
  • The lab measures it (section 2): 0.0 at init, 0.224 as soon as moves.

Simulator · the block, line by line

x · latent 1×4×4 (one colour per patch)
c · t and the class (d = 4)

Simulator · four ways to give the block the condition

Measured in the lab: the conditioning

Fashion-MNIST, , , 4 blocks, 1,500 steps, guidance :

design parameters validation loss class accuracy FD ↓
in-context tokens 470,224 0.2890 0.850 26.66
adaLN 636,304 0.2733 0.894 19.89
adaLN-Zero 710,800 0.2619 0.901 17.15

class accuracy = fraction of the 1,000 samples a small classifier puts in the requested class · FD = Fréchet distance in that classifier's features (floor of 1.55 for real images) · w = classifier-free guidance scale

  • The order of the paper survives at 28 × 28: FD 26.66 → 19.89 → 17.15.
  • The extra parameters are the modulation: adaLN-Zero's gates are ; the in-context model has no modulation at all. But what matters is how reaches the tokens: one class token competes with 49 image tokens for attention.

The modulation is cheap in compute, not in parameters

  • maps : per block, against for attention (: Q, K, V and the output) + MLP (: ).
  • DiT-XL, 28 blocks of : 223M of its 675M parameters — a third of the model to read one vector.
who what they did
PixArt-α (2023) compute it once from , one learned offset per block — 833M → 611M
Wan (2025) one modulation MLP shared by all blocks
FLUX.2 (2025) one set per block type, reused in every block

And the pooled text vector left : the timestep goes through the modulation, the prompt through attention.

How the prompt gets in: three designs

design how text updated? weight sets used in
cross-attention image queries read the text in a second attention no 1 SD U-Net, PixArt-α, Wan, Movie Gen
joint, two streams (MMDiT) one attention over [image; text], separate weights per modality yes 2 SD3, Qwen-Image
joint, one stream a plain Transformer block over [image; text] yes 1 Z-Image

img, txt = image and text tokens · [ · ; · ] = concatenate along the sequence · dh = channels per head

Inside an MMDiT block

One MMDiT block. Image tokens on the left and text tokens on the right each go through their own modulation, their own Q, K, V projections, their own output projection and their own MLP; the two sequences meet only in one attention over the joined sequence. The condition c, timestep plus pooled text, feeds both modulations. In a single-stream block the same steps use one set of weights for all tokens. IMAGE WEIGHTSTEXT WEIGHTS ONE SET OF WEIGHTS FOR ALL TOKENS Image tokens: N × dpatches of the noisy latent Text tokens, width dfrom the text encoder c: timestep + pooled textone vector per image adaLN: LN, (1 + γ) ⊙ · + β adaLN: LN, (1 + γ) ⊙ · + β Q, K, V Q, K, V One attentionjoin, attend over all, splitover all tokens at once output proj., ⊙ α, + residual output proj., ⊙ α, + residual adaLN, MLP, ⊙ α, + residual adaLN, MLP, ⊙ α, + residual one adaLN for all tokens: LN, (1 + γ) ⊙ · + β one Q, K, V for all tokens one output projection, ⊙ α, + residual one MLP for all tokens (FLUX runs it beside the attention) Image tokens, updatedafter the last block: unpatchify Text tokens, rewrittenafter the last block: dropped

Simulator · who attends to whom

Two streams, then one

model two-stream one-stream
FLUX.1 19 38
HunyuanVideo 20 40
FLUX.2 8 48
Z-Image 0 30

Align the modalities with separate weights, then share them.

Where legible text came from. SD3's ablation: drop the T5 encoder (a language model; SD3 also has two CLIPs) and prompt following barely moves (46% win rate against the full model; 50% would be a tie) but typography drops (38%).

  • It was the language-model encoder first.
  • Since: a multimodal LLM (HunyuanVideo), Qwen2.5-VL (Qwen-Image), Mistral 24B (FLUX.2).

Positions: RoPE, one axis at a time

  • DiT: a fixed 2-D sine-cosine table added to the tokens — tied to the training grid.
  • Since FLUX.1: rotary embeddings per axis — the head's channels split into groups, each rotated by one coordinate.
model axes head channels per axis
FLUX.1 index, row, column 16 + 56 + 56 = 128
HunyuanVideo time, row, column 16 + 56 + 56 = 128
CogVideoX time, row, column 2/8, 3/8, 3/8 of the channels
FLUX.2 four axes 32 · 4 = 128

(the query of one token times the key of another) depends only on differences of position: resolution, aspect ratio and duration become a matter of sequence length.

The curve that paid for the next models

model year parameters blocks text in by tokens at 1024²
DiT-XL/2 2023 675M 28 1152 class, adaLN 256 at 256²
SD3 medium 2024 2B 24 1536 two streams 4,096
FLUX.1-dev 2024 12B 19 + 38 3072 two, then one 4,096
Qwen-Image 2025 20B 60 3072 two streams 4,096
Z-Image 2025 6B 30 3840 one stream 4,096
FLUX.2-dev 2025 32B 8 + 48 6144 two, then one 4,096
  • The models got wider and deeper; the sequence stayed at 4,096: an 8× VAE and give .
  • SD3 repeated the study to 8B: validation loss falls smoothly and tracks GenEval and human preference.

19 + 38 = two-stream blocks + one-stream blocks · d = model width · GenEval = an automatic benchmark of prompt following

Video: the same block over spacetime patches

3-D VAE (4× time, 8 × 8 space, 16 channels), patches (time × height × width):

model clip (width × height) tokens
CogVideoX-5B 720 × 480, 49 frames 13 · 30 · 45 17,550
Mochi 1 848 × 480, 163 frames 28 · 30 · 53 44,520
Movie Gen 768 × 768, 256 frames 32 · 48 · 48 73,728
Wan 2.1 14B 1280 × 720, 81 frames 21 · 45 · 80 75,600
HunyuanVideo 1280 × 720, 129 frames 33 · 45 · 80 118,800
  • In space, 8 from the VAE × 2 from the patch = 16 pixels per token side: .
  • In time, : for HunyuanVideo, . Mochi compresses by 6 (), Movie Gen by 8 ().

Full 3-D attention

Early video models attended inside each frame, then along each pixel through time. Cheaper. Why did CogVideoX, HunyuanVideo and Wan drop it?

  • A head that moves is not at the same pixel in the next frame.
  • With factorized attention its tokens in frame cannot attend to its tokens in frame directly — only through other patches.
  • Consistency suffers (CogVideoX's argument). So: one attention over all tokens.
  • Also new: LTX-2 generates video and audio together (14B + 5B streams); Wan 2.2 splits the steps between a high-noise and a low-noise expert (27B, 14B active).

Where the compute goes

multiply-adds per block, with N tokens of width d · 12Nd² = 4Nd² (Q, K, V and the output, d × d matrices) + 8Nd² (MLP d → 4d → d) · 2N²d = QKᵀ (N × N scores of d terms each) + the weighted sum of the V

FLUX at 1024²: (image + text), . Is it attention-bound?

  • No: . Projections and MLPs do 80% of the work.
  • The term only wins once passes : 18,432 tokens for FLUX. An image stays below; a video does not.

Image vs video: the attention share

model attention
DiT-XL/2, 256² 256 1152 256 / 7,168 4%
FLUX.1, 1024² 4,608 3072 4,608 / 23,040 20%
Wan 2.1, 720p 75,600 5120 75,600 / 106,320 ≈ 71%
HunyuanVideo, 720p 118,800 3072 118,800 / 137,232 ≈ 87%
  • The video are those of the 3-D VAE slide; the image ones, those of patchify.
  • Measured: a 5-second 720p HunyuanVideo clip took 945 s, 800 in attention — 85%, close to the formula's 87%.

So the efficiency tricks split

lever examples matters for
fewer steps LADD / SD3-Turbo (4 steps), FLUX.1 schnell (1–4), guidance distillation images and video
fewer tokens SANA's 32× autoencoder, token merging both
cheaper attention SANA (linear), sliding-tile: 945 → 685 s with no training video
reuse across steps TeaCache, up to 4.4× video
fewer active parameters Wan 2.2: two experts by noise level large models

Image DiTs are MLP-bound; video DiTs are attention-bound. That is why sparse attention is a video story.

Measured in the lab: the patch size

Same , 4 blocks, 1,500 steps, — only changes:

parameters tokens M multiply-adds class accuracy FD ↓
14 745,540 4 1.8 0.319 341.97
7 717,169 16 7.3 0.907 27.59
4 710,800 49 23.5 0.901 17.15
  • Only and the output depend on : — 37,924 at , 3,184 at . The rest is identical.
  • Compute : at , M — 13× that of , and the FD falls with it.
  • At each token is numbers squeezed into , and attention only moves information between tokens.
  • ( tokens): 116M, 4.94× — it is the open task.

Where this is used, October 2026

domain systems
text-to-image (open) SD3/3.5, FLUX.1, Qwen-Image, Z-Image, FLUX.2, Qwen-Image-2.1 MMDiT, two-then-one or one stream, rectified flow, LLM/VLM encoders
text-to-video (open) CogVideoX, HunyuanVideo (1.5), Wan 2.1/2.2, LTX-2 3-D VAE, 3-D RoPE, full or sparse 3-D attention
audio LTX-2 an audio stream of the same DiT
3-D assets TRELLIS rectified-flow Transformers over a 3-D latent
robotics GR00T N1 a DiT action module, flow matching
closed Sora, Sora 2, Veo 3, Midjourney Sora: a DiT over spacetime patches; others unpublished

rectified flow = flow matching on straight paths, the loss of the lab · MMDiT = multimodal DiT, two streams

What replaced the U-Net, where

  • Text-to-image: between SDXL (2023, U-Net) and SD3 (2024, MMDiT).
  • Video: in 2024 — Stable Video Diffusion (2023) has no U-Net successor.
  • Still convolutional: the VAE around every one of these models, many super-resolution stages.

The U-Net is not dead. EDM2 (a U-Net, 2024): FID 1.81 on ImageNet 512², better than DiT-XL/2's 3.04. The case for DiT was never "beats a convolution at a fixed budget" — it is that it keeps paying as the budget grows, and shares everything with LLMs.

Where the design is going

development what it changes
REPA align a DiT layer with a frozen DINOv2: up to 17.5× faster training
RAE diffuse in a frozen DINOv2 latent instead of a VAE's: FID 1.13 on ImageNet
one stream + big text encoder fewer separate weights in the generator, more language in the encoder
unified models one Transformer reads and generates text and images — next class

REPA = representation alignment · RAE = representation autoencoder · DINOv2 = a ViT trained without labels (self-supervised)

Lab · a DiT, built and measured

Colab notebook, ~2 h, T4 GPU (free tier), Fashion-MNIST. Each section ends in a checkpoint.

  1. patchify and unpatchify, by hand
  2. the adaLN-Zero block: identity, measured
  3. the whole DiT and its cost against the clock
  4. flow matching, guidance, and a judge
  1. in-context × adaLN × adaLN-Zero
  2. : tokens, time, quality
  3. your DiT under a budget: the table is yours

open in Colab →

Takeaways

  1. A DiT is a ViT over noisy latent patches; unpatchify gives back a velocity of the input's shape.
  2. The patch size is a compute dial: same parameters, 16× the Gflops (118.6/7.4), FID 106 → 19.5.
  3. adaLN-Zero: norms modulated by , gates at zero, every block starts as the identity — best and cheapest; its parameters are why later models share it.
  4. Text: cross-attention, joint with two streams, or one stream. Joint attention is nearly free; weights are what cost.
  5. RoPE per axis makes resolution and duration a sequence length.
  6. Images are MLP-bound, video is attention-bound: .

Outline: (1) the swap and patchify; (2) adaLN-Zero and the conditioning ablation; (3) how the prompt gets in: cross-attention, joint attention, one stream; (4) positions and the scaling curve; (5) video and where the compute goes; (6) where it is used in 2026. Closes with the Colab lab.

U-Net: resolution goes down and up, attention blocks between the convolutions. DiT: one resolution, every block sees every token, t and the class enter every block. MMDiT: text gets its own stream.

This is section 1 of the lab: patchify and unpatchify by hand, with the exact round trip.

Start on DiT-XL at p = 2: 256 tokens, 118 G multiply-adds — the paper's 118.6 Gflops. Move p to 4 and 8. Then FLUX at p = 1: 16,384 tokens and attention at almost half of the cost.

The 12 models of the paper's table 4. Against Gflops: one band, S/2 and B/4 on the same spot. Click "Against parameters": each size becomes a vertical line, FID 153.6 to 68.4 at the same 33M. Ask before clicking: what will the S points do?

Seven steps: t and y become c; one W per block makes six vectors, shared by all tokens; the attention half; the MLP half; W = 0 (the chips turn to 0 and only the residual path is lit — the next slide asks for the number); the cost: 0.2% of the compute, a third of the parameters.

Default case "a trained block": step through. Line 3 is the six vectors from c alone, the same for the 4 tokens; line 4 is LN (mean 0, variance 1); line 5 is where t enters. Then pick "at initialization: W = 0": the six vectors are 0, the attention still computes (line 6) but α1 = 0 lets none of it through; max |block(z) − z| = 0, out = 0 — the previous slide's answer. "the same block, another t": same x, other c, other output. "a flat patch": LN turns token 1 into zeros.

The paper's ablation on DiT-XL/2. Ask for the ranking before clicking through. adaLN-Zero is the best and, with adaLN, the cheapest: the modulation is per image, not per token.

Step through: two sequences in; one c, two modulations; own Q, K, V; join, one attention, split; own output and MLP — the text comes out rewritten; last step: the single-stream block, one set of weights, same work per token.

Cross-attention: the text rows are never computed. Joint: all four blocks in one softmax. Switch to FLUX and to Wan: joint attention costs 12.5% and 0.7% more scores. The difference between two streams and one is in the weights, not in the compute.