Chapter 34 · Section 34.2 · 90 min

U-Nets, DiT, and MMDiT

This section is about the neural network inside an image generator, the denoiser. A diffusion or flow sampler makes a picture by starting from pure noise and calling this one network dozens of times. Each call looks at the current noisy picture and predicts the noise to remove (§33.2). So the denoiser decides what the generator can learn, how much it improves with more compute, and what every one of those calls costs.

Here is what the network must cope with. The denoiser of §33.2 was a small MLP that read a point in the plane. A real system denoises a latent instead of pixels (§34.1). An autoencoder compresses the image into a small grid of numbers, and the sampler works on that grid. For a 512 by 512 photograph the grid is $64\times64\times4$ , which is 16,384 numbers. An MLP on that grid would not know which entries are neighbors in the picture, and one layer from 16,384 numbers to 16,384 numbers already holds $16{,}384^2 \approx 268$ million weights. The same network must also work at every noise level. At high noise it has to invent a global layout from almost nothing. At low noise it has to fix fine detail without disturbing the rest. And it must follow a class label or a prompt.

So this section answers three concrete questions. What network should predict the noise? How does the noise level reach every one of its layers? How does a label or a prompt get in? Three designs have answered them, one after another.

A convolutional U-Net processes the latent at several resolutions with small filters. It feeds the noise level into every block, by adding a learned vector to its features or by rescaling them. The first strong image diffusion models used it, and so did Stable Diffusion 1 and SDXL.
A diffusion transformer, DiT, cuts the latent into patches and treats each patch as a token of a Transformer. It keeps the rescaling idea and adds one initialization detail, called adaLN-Zero, that makes every block start by passing its input through unchanged.
A multimodal diffusion transformer, MMDiT, lets image tokens and prompt tokens attend to each other in one attention operation. Stable Diffusion 3 uses it, and FLUX uses a variant. We will prove that this joint attention is exactly self-attention over the image and cross-attention to the text, blended by one softmax.

Along the way we count every parameter and multiply-add of the standard DiT model from its configuration and reproduce the published numbers. We close with what is known about scaling these networks. By the end you should be able to build each design, cost it, and defend its choices in an interview.

Notation. A latent is $\mathbf{z} \in \R^{c\times h\times w}$ , with $c$ channels, height $h$ , and width $w$ , written channels first as code stores it. The denoiser sees the noisy latent $\mathbf{z}_t$ at time $t \in [0, 1]$ . Clean data sits at $t = 0$ and pure noise at $t = 1$ , as in Chapter 33. The bold $\mathbf{c}$ is a condition vector, never the channel count $c$ . A transformer has width $d$ , depth $L$ , and $T$ tokens. Here $T$ always counts tokens, never sampling steps. The patch size is $p$ . Feature maps inside a U-Net are $\mathbf{h} \in \R^{C\times H_f\times W_f}$ , with $C$ channels.

What the denoiser has to do

Every denoiser has the same inputs and the same output, whatever is inside it. It reads a noisy latent, the time, and a condition, and it returns an array with the latent's shape,

\boldsymbol\epsilon_\theta : \R^{c\times h\times w}\times[0, 1]\times\mathcal{C}\to\R^{c\times h\times w}, \qquad (\mathbf{z}_t, t, \mathbf{c})\mapsto\hat{\boldsymbol\epsilon}.

Here $\mathcal{C}$ is the set of possible conditions, such as class labels, prompts, or nothing at all, and $\theta$ stands for the network's weights. In plain words, the network looks at a noisy picture, is told how noisy it is and what it should show, and returns one number for every number of the picture.

That output can mean three things. It can be the noise $\boldsymbol\epsilon$ that was added, the clean latent, or the velocity $\mathbf{v} = \boldsymbol\epsilon - \mathbf{z}_0$ that flow matching learns, which points from the data toward the noise. The choice changes the loss and the sampler (§33.4 and §33.6), not the network. Three demands on the network follow.

One output per input entry. Every output entry is a prediction about the matching input entry. So the network maps an image to an image, like a segmentation network. It never pools everything into one vector the way a classifier does.

Every scale at once. At large $t$ the input is mostly noise, and the best prediction depends on the global layout the network guesses. At small $t$ the input is almost clean, and the leftover noise is a fine pattern that must be found entry by entry. One network has to do both. So it must see the whole image and also keep a short path for fine detail.

A global input for the noise level. The time $t$ changes what a good feature is, everywhere in the image at once. So it should reach every layer, not enter as one more pixel. A class label is global in the same way, and so is a single vector that summarizes a prompt. A full prompt is a sequence of tokens and needs attention, which is the subject of §34.3.

The U-Net meets the first two demands with several resolutions, and the transformer meets them with attention. Both meet the third with the same idea, a rescaling of normalized features, which we build after the U-Net. If an interviewer asks why a classifier backbone cannot serve as a denoiser, these three demands are the answer.

The diffusion U-Net

A U-Net is an image-to-image network shaped like the letter U (Ronneberger et al. 2015). Its left half applies convolutions and halves the resolution, again and again. Its right half doubles the resolution and applies convolutions, again and again. A skip connection carries the features of each level on the left straight across to the level of the same resolution on the right. The U-Net was built to segment biomedical images, another task with one output per pixel. The first strong diffusion model for images, DDPM (Ho et al. 2020), used it, and almost every pixel and latent diffusion model until 2023 kept its shape.

Each part has a reason you can count. A $3\times3$ convolution computes each output from one position and its eight neighbors (§0.8). The receptive field of a feature is the patch of the input it depends on. Each $3\times3$ layer widens it by two positions, so at full resolution a feature would need about 32 layers to see across a $64\times64$ latent. Halve the resolution three times instead. The map is then $8\times8$ , and each of its positions summarizes an $8\times8$ block of the latent. One $3\times3$ convolution there spans 24 of the 64 latent positions on each side. That low-resolution middle, the bottleneck, is where the global layout is decided.

An $8\times8$ map cannot hold all the detail of a $64\times64$ input, and at small $t$ the noise to predict is exactly that detail. The skip connections give the detail a short path. Each level on the way up receives the features of the matching level on the way down, stacked along the channels. So the network only has to learn what to change at each scale.

DDPM's U-Net was later refined, one tested change at a time, into a stronger network that its authors called the ablated diffusion model, ADM (Dhariwal and Nichol 2021). It is built from four parts.

Residual blocks. A residual block computes a correction and adds it to its own input. Here the correction is GroupNorm (defined below), the activation $\operatorname{SiLU}(x) = x\cdot\operatorname{sigmoid}(x)$ , a $3\times3$ convolution, a second GroupNorm where the time enters (next section), SiLU, dropout, and a second $3\times3$ convolution. When the channel count changes, a $1\times1$ convolution resizes the input before the addition. The last convolution starts at zero, so every block starts as the identity, an idea DiT will reuse.
Down and up sampling. Each level has two or three residual blocks. A stride-2 convolution or average pooling then goes down a level, and a nearest-neighbor resize followed by a convolution goes up. The channel count grows as the resolution falls, for example $C, 2C, 4C, 4C$ .
Self-attention at low resolution. At chosen levels a residual block is followed by multi-head self-attention (§4.5), in which every position of the feature map is a token. DDPM used single-head attention at $16\times16$ and once more in the middle of the bottleneck. ADM used heads of 64 channels at $32\times32$ , $16\times16$ , and $8\times8$ . Its ablations found that more heads, attention at several resolutions, and residual blocks for the down and up sampling each improved FID. FID measures how far the statistics of generated images are from those of real images, and lower is better (§34.6 derives it).
Skip concatenation. On the way up, each block's input is the upsampled path stacked with the matching features from the way down. That is why blocks on the right side have more input channels than those on the left.

Why put attention only at low resolution? Self-attention compares every position with every other, so its cost grows with the square of the number of positions. The number of positions is the square of the side. So the number of query-key pairs grows with the fourth power of the side. A U-Net halves its channels each time it doubles the side, which only halves the cost of each pair. The next example puts real numbers on it.

✏️ Where SDXL puts its attention

SDXL is Stable Diffusion's 2023 model for $1024\times1024$ images (Podell et al. 2023, Table 1). Its U-Net works on a $128\times128$ latent and has three levels, with 320, 640, and 1,280 channels. The attention module at each level stacks 0, 2, and 10 transformer blocks. A transformer block here is self-attention, cross-attention to the prompt, and an MLP.

At the top level, a self-attention layer would see $128^2 = 16{,}384$ tokens, so $16{,}384^2 \approx 268$ million query-key pairs.
One level down there are $4{,}096$ tokens and 16.8 million pairs.
At the bottom there are $1{,}024$ tokens and about 1 million pairs.

Doubling the side multiplies the pairs by 16 and halves the channels. So one attention layer at the top level would cost 8 times one at the level below. SDXL therefore has no attention at the top level. It also dropped the lowest level of SD 1.5, which worked at one eighth of the latent's side, and put ten transformer blocks at the bottom, where they are cheapest. The whole U-Net has 2.6 billion parameters. SD 1.5 had 860 million, with four levels whose channels were in the ratio $1, 2, 4, 4$ and one transformer block per attention module.

Python · cell 1

In practice, the first question to ask of any image backbone is where it spends attention and how that cost grows with resolution. An interviewer who asks why U-Nets put attention only at low resolution wants the fourth-power argument.

Group normalization

A normalization layer rescales features so that each layer sees inputs of a steady size. BatchNorm computes each channel's mean and variance over all the examples of a batch (§0.7). That fails in a diffusion model for three reasons. Images are large, so each device holds a small batch and the statistics are noisy. Each example has its own noise level, so the statistics mix an almost clean image with pure noise, and each output depends on the other examples. And at sampling time the batch is whatever the user asked for. LayerNorm normalizes each example by itself (§5.2), but over a whole feature map it makes every channel share one mean and one scale. Group normalization (Wu and He 2018) sits between the two.

📐 Group normalization

Let $\mathbf{h} \in \R^{C\times H_f\times W_f}$ be the features of one example, and let the number of groups $G$ divide $C$ . Group $k \in \{1, \dots, G\}$ holds the $C/G$ consecutive channels $\mathcal{G}_k = \{(k-1)C/G + 1, \dots, kC/G\}$ . With $n = (C/G)H_fW_f$ entries per group,

\mu_k = \frac{1}{n}\sum_{c'\in\mathcal{G}_k}\sum_{u, v} h_{c'uv}, \qquad \sigma_k^2 = \frac{1}{n}\sum_{c'\in\mathcal{G}_k}\sum_{u, v}\big(h_{c'uv} - \mu_k\big)^2,

and for every channel $c' \in \mathcal{G}_k$ and position $(u, v)$ ,

\operatorname{GN}(\mathbf{h})_{c'uv} = \gamma_{c'}\,\frac{h_{c'uv} - \mu_k}{\sqrt{\sigma_k^2 + \varepsilon}} + \beta_{c'}.

The gain $\boldsymbol\gamma \in \R^C$ and the shift $\boldsymbol\beta \in \R^C$ are learned, one entry per channel. The small constant $\varepsilon$ is $10^{-5}$ in PyTorch and has nothing to do with the noise $\boldsymbol\epsilon$ . The variance is the biased one, divided by $n$ . A batch is normalized example by example. $G = 1$ is LayerNorm over $(C, H_f, W_f)$ , and $G = C$ is InstanceNorm, which normalizes each channel alone. DDPM, ADM, and Stable Diffusion use $G = 32$ .

In plain words, split the channels into groups. Give each group of each example its own mean and variance, computed over all its channels and positions together. The gain and the shift stay per channel.

Why pool the statistics over several channels? Because it keeps information that normalizing each channel alone would destroy. A channel that is constant across the image has zero variance on its own, so InstanceNorm maps it to 0. Inside a group, its offset from the other channels survives. The example shows this, and then previews how the time will act on the normalized features.

✏️ GroupNorm by hand

Take $C = 4$ channels on a $2\times2$ map, $G = 2$ , no $\varepsilon$ , $\boldsymbol\gamma = \mathbf{1}$ , and $\boldsymbol\beta = \mathbf{0}$ . Channel 0 is $(0, 4, 4, 8)$ in raster order. Channel 1 is constant 4, channel 2 is constant 1, and channel 3 is constant $-1$ .

Group 1 holds channels 0 and 1. Its 8 entries sum to 32, so $\mu_1 = 4$ . The deviations are $-4, 0, 0, 4$ and four zeros, so $\sigma_1^2 = 32/8 = 4$ and $\sigma_1 = 2$ . Channel 0 becomes $(-2, 0, 0, 2)$ and channel 1 becomes 0. Group 2 has four 1s and four $-1$ s, so $\mu_2 = 0$ and $\sigma_2^2 = 1$ , and channels 2 and 3 stay at $1$ and $-1$ . InstanceNorm would have sent channels 1, 2, and 3 all to 0, erasing the fact that channel 2 is on and channel 3 is off.

Now let the time rescale channel 0. Suppose that at $t = 0.9$ the network chooses a scale $1 + \gamma_0 = 0.5$ and a shift $\beta_0 = 0$ , and at $t = 0.1$ it chooses $1 + \gamma_0 = 2$ and $\beta_0 = 1$ . Channel 0 leaves as $(-1, 0, 0, 1)$ at high noise and as $(-3, 1, 1, 5)$ at low noise. The same feature is turned down in one case and turned up and offset in the other, by two numbers that depend only on $t$ . That is the whole idea of the next section.

Python · cell 2

Warm-up

GroupNorm for a batch of feature maps

Implement group_norm(h, G, gamma, beta, eps=1e-5). The input h has shape (B, C, H, W), and the output has the same shape. For each example and each group, compute one mean and one biased variance over the group's $C/G$ channels and all $H\times W$ positions, and normalize with them. Then apply a per-channel gain and shift from gamma and beta, each of shape (C,). If G does not divide C, raise a ValueError. Use no loops over the batch or the groups.

import numpy as np

def group_norm(h, G, gamma, beta, eps=1e-5):
    raise NotImplementedError

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

Telling every layer the noise level

The time reaches every layer in two steps. First it becomes a vector the network can use. Then that vector rescales and shifts the normalized features inside every block.

The embedding. One number is a poor input. A network finds it hard to respond sharply to small changes of a single scalar, and DDPM's time runs over a thousandfold range of steps. So we expand $t$ into sines and cosines at many frequencies, from fast to slow, the same features that §5.4 uses for token positions. Here is a tiny case. With 4 features there are two frequencies, 1 and 0.01. At step 50 the features are $(\cos 50, \cos 0.5, \sin 50, \sin 0.5) = (0.96, 0.88, -0.26, 0.48)$ . At step 51 the fast cosine has fallen to 0.74 while the slow one has barely moved, to 0.87. The fast features tell neighboring times apart, and the slow ones tell distant times apart.

In general, for an embedding of width $d_f$ (256 in DiT) with $k = d_f/2$ frequencies $\omega_i = 10000^{-i/k}$ , $i = 0, \dots, k - 1$ , the features are

\boldsymbol\psi(t) = \big(\cos(t\omega_0), \dots, \cos(t\omega_{k-1}),\ \sin(t\omega_0), \dots, \sin(t\omega_{k-1})\big) \in \R^{d_f}.

In plain words, all the cosines come first and then all the sines, the layout of ADM and DiT. At $t = 0$ every cosine is 1 and every sine is 0. In DDPM and DiT the number fed to these features is the step index, from 0 to 999. Flow models such as SD3, whose time runs over $[0, 1]$ , multiply it by 1,000 first, so the fast frequencies still turn many times over the range. A two-layer MLP then maps the features to the model width,

\operatorname{emb}(t) = \operatorname{SiLU}\big(\boldsymbol\psi(t)\mathbf{W}_1 + \mathbf{b}_1\big)\mathbf{W}_2 + \mathbf{b}_2 \in \R^d .

Libraries differ in the details. The diffusers function get_timestep_embedding can put the sines first, and it can divide the exponent by $k - 1$ instead of $k$ . A checkpoint only works with the convention it was trained with.

The class. A class label $y$ looks up a row of a learned table, $\operatorname{emb}(y) \in \R^d$ . The table has one extra row that means "no label". Classifier-free guidance (§33.5) needs that row. It sharpens samples by comparing the prediction made with the label against the prediction made without it, so one network must be able to make both. DiT's code replaces the label with the extra row 10 percent of the time during training. The condition vector is the sum,

\mathbf{c} = \operatorname{emb}(t) + \operatorname{emb}(y) \in \R^d .

The modulation. How should $\mathbf{c}$ act on a feature map? DDPM projected the time embedding to one number per channel and added it to the activations of each residual block. ADM instead rescaled and shifted the normalized activations. In its ablation on $128\times128$ images from ImageNet, the standard benchmark of 1,000 object classes, rescaling reached FID 13.06 against 15.08 for adding, at the same training length. The U-Nets of Stable Diffusion 1 and SDXL kept DDPM's addition. The general form of rescaling has a name.

📐 FiLM and adaptive normalization

Feature-wise linear modulation (FiLM, Perez et al. 2017) conditions features $\mathbf{h}$ on $\mathbf{c}$ by a per-channel affine map whose coefficients are functions of $\mathbf{c}$ ,

\operatorname{FiLM}(\mathbf{h}\mid\mathbf{c}) = \boldsymbol\gamma(\mathbf{c})\odot\mathbf{h} + \boldsymbol\beta(\mathbf{c}),

broadcast over all positions. Adaptive normalization applies FiLM right after a normalization whose own gain and shift are removed. ADM's adaptive group normalization is

\operatorname{AdaGN}(\mathbf{h}, \mathbf{c}) = \mathbf{y}_s\odot\operatorname{GN}(\mathbf{h}) + \mathbf{y}_b, \qquad [\mathbf{y}_s, \mathbf{y}_b] = \mathbf{c}\mathbf{W} + \mathbf{b},

and the adaptive LayerNorm of a transformer, applied to each token $\mathbf{x} \in \R^d$ , is

\operatorname{adaLN}(\mathbf{x}\mid\mathbf{c}) = \operatorname{LN}(\mathbf{x})\odot\big(\mathbf{1} + \boldsymbol\gamma(\mathbf{c})\big) + \boldsymbol\beta(\mathbf{c}).

Here $\operatorname{LN}$ has no learned gain or shift, and $\boldsymbol\gamma, \boldsymbol\beta \in \R^d$ come from a linear layer applied to $\operatorname{SiLU}(\mathbf{c})$ . Writing the scale as $\mathbf{1} + \boldsymbol\gamma$ makes $\boldsymbol\gamma = \boldsymbol\beta = \mathbf{0}$ the plain normalization. ADM's code uses the same $(1 + \text{scale})$ form.

In plain words, the condition is not one more input. It turns each channel up or down and moves its center, the same way at every position. One number per channel can change what the whole layer does.

Why does rescaling beat adding? A bias added before a normalization is partly undone by it. The next GroupNorm subtracts the group mean, and divides what survives by the activations' own spread. A scale applied after the normalization acts directly. It can also do what a bias cannot, which is to switch a channel off ( $1 + \gamma_k = 0$ ) or amplify it. A network that should rely on texture detectors at low noise and layout detectors at high noise needs exactly that kind of switch. Interviewers often ask why the scale is written as $1 + \boldsymbol\gamma$ . The answer is that a modulation layer that outputs zero then leaves the features unchanged.

Core

Timestep embedding and modulation

Implement two functions.

timestep_embedding(t, d, max_period=10000.0) takes times t of shape (B,) and an even width d. It returns the features $\boldsymbol\psi(t)$ above, of shape (B, d), with frequencies $\text{max\_period}^{-i/(d/2)}$ , all cosines first and then all sines.
modulate(x, shift, scale) takes tokens x of shape (B, T, d) and one shift and one scale per image, each of shape (B, d). It returns $\mathbf{x}\odot(\mathbf{1} + \text{scale}) + \text{shift}$ , so that every token of image $b$ gets image $b$ 's modulation.

import numpy as np

def timestep_embedding(t, d, max_period=10000.0):
    raise NotImplementedError

def modulate(x, shift, scale):
    raise NotImplementedError

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

The diffusion transformer

Peebles and Xie 2023 asked whether a denoiser needs convolutions at all. Their diffusion transformer, DiT, is a Vision Transformer with the fewest changes that make it a denoiser. A Vision Transformer cuts an image into square patches, flattens each patch into a vector, and maps every vector to a token with one shared linear layer (§27.1, §27.2). DiT does the same to the latent.

Here is the input in numbers. A $256\times256$ ImageNet image becomes a $32\times32\times4$ latent in Stable Diffusion's autoencoder. With patch size $p = 2$ , the latent is cut into $16\times16 = 256$ patches of $2\times2\times4 = 16$ numbers each. In general there are

T = \frac{h}{p}\cdot\frac{w}{p}

tokens, each made from $p^2c$ numbers. Fixed two-dimensional sine-cosine position embeddings are added, and there is no class token. DiT comes in four sizes, three of them taken from ViT. They are S (depth 12, width 384, 6 heads), B (12, 768, 12), L (24, 1024, 16), and a new XL (28, 1152, 16), each with patch sizes 2, 4, and 8. The name DiT-XL/2 means size XL with $p = 2$ .

The open question was how the condition $\mathbf{c}$ should enter. Peebles and Xie trained four DiT-XL/2 variants for 400,000 steps and compared them.

In context. Append $\operatorname{emb}(t)$ and $\operatorname{emb}(y)$ as two extra tokens and use unchanged ViT blocks, the way ViT uses a class token.
Cross-attention. Keep the two embeddings as a separate sequence of length two, and add a cross-attention layer after self-attention in every block (§4.6).
adaLN. Replace each LayerNorm with adaptive LayerNorm, with the shift and scale computed from $\mathbf{c}$ .
adaLN-Zero. adaLN plus a per-channel gate that multiplies each residual branch, also computed from $\mathbf{c}$ . The layer that computes all of them starts at zero.

block	G multiply-adds	parameters	FID-50K at 400K steps
in-context	119.37	449M	35.24
cross-attention	137.62	598M	26.14
adaLN	118.56	600M	25.21
adaLN-Zero	118.64	675M	19.47

The parameter column holds a surprise that the counting section will explain. adaLN-Zero has the most parameters of the four, because its modulation layers are large. Yet it costs almost no extra compute, because those layers run once per image instead of once per token.

📐 DiT block with adaLN-Zero

Let $\mathbf{X} \in \R^{T\times d}$ hold one image's tokens as rows and $\mathbf{c} \in \R^{1\times d}$ its condition. All vectors below are rows of length $d$ , $\mathbf{1}_T \in \R^{T\times1}$ is a column of ones, and $\operatorname{LN}$ normalizes each row with no gain or shift. A modulation layer with $\mathbf{W}_{\text{mod}} \in \R^{d\times 6d}$ and $\mathbf{b}_{\text{mod}} \in \R^{1\times6d}$ produces six vectors,

\big[\boldsymbol\beta_1, \boldsymbol\gamma_1, \mathbf{g}_1, \boldsymbol\beta_2, \boldsymbol\gamma_2, \mathbf{g}_2\big] = \operatorname{SiLU}(\mathbf{c})\mathbf{W}_{\text{mod}} + \mathbf{b}_{\text{mod}},

a shift, a scale, and a gate for each of the two sublayers. The block computes

\begin{aligned} \mathbf{U}_1 &= \operatorname{LN}(\mathbf{X})\operatorname{diag}(\mathbf{1} + \boldsymbol\gamma_1) + \mathbf{1}_T\boldsymbol\beta_1, & \mathbf{X}_1 &= \mathbf{X} + \operatorname{MHSA}(\mathbf{U}_1)\operatorname{diag}(\mathbf{g}_1),\\ \mathbf{U}_2 &= \operatorname{LN}(\mathbf{X}_1)\operatorname{diag}(\mathbf{1} + \boldsymbol\gamma_2) + \mathbf{1}_T\boldsymbol\beta_2, & \mathbf{X}_{\text{out}} &= \mathbf{X}_1 + \operatorname{MLP}(\mathbf{U}_2)\operatorname{diag}(\mathbf{g}_2). \end{aligned}

Here $\operatorname{MHSA}$ is multi-head self-attention with its output projection. The MLP is $\operatorname{MLP}(\mathbf{U}) = \operatorname{GELU}(\mathbf{U}\mathbf{W}_{\text{up}} + \mathbf{b}_{\text{up}})\mathbf{W}_{\text{down}} + \mathbf{b}_{\text{down}}$ , with hidden width $4d$ and GELU in its tanh approximation. Multiplying by $\operatorname{diag}(\mathbf{g})$ on the right scales column $j$ by $g_j$ . adaLN-Zero initializes $\mathbf{W}_{\text{mod}} = \mathbf{0}$ and $\mathbf{b}_{\text{mod}} = \mathbf{0}$ . The modulation depends only on $\mathbf{c}$ , so it is computed once per image and shared by all $T$ tokens.

In plain words, a DiT block is a pre-norm Transformer block (§5.6) with two changes. Its two LayerNorms get their shift and scale from the condition. Its two residual branches are multiplied by a gate, also computed from the condition, before being added back. The paper calls the gate $\boldsymbol\alpha$ . We write $\mathbf{g}$ to keep $\alpha_t$ for the noise schedule.

The output head. After the last block come one more adaptive LayerNorm and a linear layer, both initialized to zero. The linear layer maps each token to $p^2\cdot 2c$ numbers, twice what a patch holds. The first $c$ channels of each output position are the noise prediction $\hat{\boldsymbol\epsilon}$ . The other $c$ set the variance of each reverse sampling step, because DiT is trained like ADM, which also learns that variance. Each such number $\rho$ picks a variance between the two fixed choices of §33.2, the step's noise variance $\beta_t$ and the posterior variance $\tilde\beta_t$ , as $\exp\big(\rho\log\beta_t + (1 - \rho)\log\tilde\beta_t\big)$ (Nichol and Dhariwal 2021). A flow-matching DiT needs only $p^2c$ numbers per token, the velocity.

Finally the tokens are put back into their spatial layout. This step, unpatchify, is the inverse of the patch flattening. It is a place where a layout bug scrambles the output silently, because the loss still goes down a little. So it deserves a test.

Core

Unpatchify the output head

The setup defines patchify(z, p). It turns a latent of shape (B, C, H, W) into tokens of shape (B, T, p*p*C). The patches are in raster order, and each patch is flattened as (row in patch, column in patch, channel), with channels last. DiT's head emits tokens of width $p^2\cdot 2c$ in that layout. The first $c$ channels are the noise prediction and the next $c$ the variance output.

Implement unpatchify(tokens, h, w, p, c). It returns the pair (eps_pred, var_pred), each of shape (B, c, h, w). It must invert patchify exactly, including for latents that are not square.

import numpy as np

def unpatchify(tokens, h, w, p, c):
    raise NotImplementedError

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

Interview

A DiT block with adaLN-Zero, batched, in NumPy

Implement dit_block(X, c, P, n_heads) exactly as in the definition box. The tokens X have shape (B, T, d) and the conditions c have shape (B, d). The dict P holds these weights.

W_mod (d, 6d) and b_mod (6d,). Their output splits into six consecutive chunks of width d, in the order (shift, scale, gate) for attention, then (shift, scale, gate) for the MLP.
W_qkv (d, 3d) and b_qkv. The output columns are $Q$ , $K$ , $V$ in that order, and head $j$ uses columns j*d_h:(j+1)*d_h of each.
W_o (d, d) and b_o for the attention output.
The MLP's W_1 (d, 4d), b_1, W_2 (4d, d), and b_2.

The setup provides ln_noaffine (LayerNorm over the last axis with $\varepsilon = 10^{-6}$ and no gain or shift), gelu_tanh, silu, softmax_last, and make_block_params. Use no loops over the batch, the tokens, or the heads.

import numpy as np

def dit_block(X, c, P, n_heads):
    # X: (B, T, d) tokens, c: (B, d) condition, P: dict of weights (see the task)
    raise NotImplementedError

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

Why adaLN-Zero starts as the identity and still learns

Starting at zero sounds like a way to kill a network. If every weight were zero, every hidden unit would compute the same thing and receive the same gradient, so the units would stay identical forever. That is the symmetry problem of §0.7. adaLN-Zero zeros something much narrower, the one layer that produces the shifts, scales, and gates. The attention and MLP weights stay random.

So each branch still computes a rich random function. The zero gate in front of it hides that function from the residual stream, the running sum of branch outputs that flows through the network (§5.1). The question is whether the gate can open.

A tiny example shows why it can. Take one entry where the branch outputs $f = 0.7$ , the gate is $g = 0$ , and the gradient arriving from above is 2. The entry adds $g f = 0$ to the stream, so the block changes nothing. But the gradient with respect to the gate is $2 \cdot f = 1.4$ , which is not zero. The gradient with respect to the branch output is $2 \cdot g = 0$ . So the gate learns at once and the branch waits. In general, a product of a zero factor and a random factor is zero, yet its gradient with respect to the zero factor is not. The theorem states this for the whole block.

Write $\mathrm{d}\mathbf{Z} = \partial\mathcal{L}/\partial\mathbf{Z}$ for the gradient of the loss with respect to any array $\mathbf{Z}$ , as in §4.3. Let $\mathbf{A} = \operatorname{MHSA}(\mathbf{U}_1)$ and $\mathbf{M} = \operatorname{MLP}(\mathbf{U}_2)$ be the two branch outputs, each $T\times d$ .

∴ adaLN-Zero starts as the identity and still learns

Take one DiT block at the adaLN-Zero initialization ( $\mathbf{W}_{\text{mod}} = \mathbf{0}$ , $\mathbf{b}_{\text{mod}} = \mathbf{0}$ , any attention and MLP weights), and let $\mathrm{d}\mathbf{X}_{\text{out}}$ be the gradient arriving from above. Then

the block is the identity, $\mathbf{X}_{\text{out}} = \mathbf{X}$ , for every input and every condition;
the gate gradients are $\mathrm{d}\mathbf{g}_1 = \mathbf{1}_T^\top(\mathrm{d}\mathbf{X}_{\text{out}}\odot\mathbf{A})$ and $\mathrm{d}\mathbf{g}_2 = \mathbf{1}_T^\top(\mathrm{d}\mathbf{X}_{\text{out}}\odot\mathbf{M})$ , sums over tokens, generally nonzero;
the modulation layer receives $\mathrm{d}\mathbf{W}_{\text{mod}} = \operatorname{SiLU}(\mathbf{c})^\top\,\mathrm{d}\mathbf{m}$ and $\mathrm{d}\mathbf{b}_{\text{mod}} = \mathrm{d}\mathbf{m}$ , where $\mathrm{d}\mathbf{m} = [\mathbf{0}, \mathbf{0}, \mathrm{d}\mathbf{g}_1, \mathbf{0}, \mathbf{0}, \mathrm{d}\mathbf{g}_2]$ , so only the gate columns move;
the attention and MLP weights, the shift and scale columns, and everything that computes $\mathbf{c}$ (the timestep MLP and the class table) receive exactly zero gradient from this block;
the input gradient passes through unchanged, $\mathrm{d}\mathbf{X} = \mathrm{d}\mathbf{X}_{\text{out}}$ .

For a batch, the gradients of items 2 and 3 are summed over images. The proof is in the derivation below.

In plain words, at the start each block adds nothing to its input, yet each gate already receives a gradient. That gradient measures how well the branch's random output lines up with the gradient from above, summed over the tokens. Nothing inside the branch learns until its gate is no longer zero.

∂ Proof that adaLN-Zero starts as the identity and still learns

At initialization all six modulation vectors are zero. Then $\mathbf{X}_1 = \mathbf{X} + \mathbf{A}\operatorname{diag}(\mathbf{0}) = \mathbf{X}$ , and likewise $\mathbf{X}_{\text{out}} = \mathbf{X}_1 = \mathbf{X}$ , which is item 1.

Go backward through the MLP sublayer, $\mathbf{X}_{\text{out}} = \mathbf{X}_1 + \mathbf{M}\operatorname{diag}(\mathbf{g}_2)$ . Entry $(i, j)$ is $X_{1,ij} + M_{ij}g_{2,j}$ , so $\partial\mathcal{L}/\partial g_{2,j} = \sum_i \mathrm{d}X_{\text{out},ij}M_{ij}$ , which is item 2 for $\mathbf{g}_2$ . Also $\mathrm{d}\mathbf{M} = \mathrm{d}\mathbf{X}_{\text{out}}\operatorname{diag}(\mathbf{g}_2) = \mathbf{0}$ . Every weight inside the MLP, and its input $\mathbf{U}_2$ , receive gradient only through $\mathbf{M}$ , so they get zero. From $\mathbf{U}_2 = \operatorname{LN}(\mathbf{X}_1)\operatorname{diag}(\mathbf{1} + \boldsymbol\gamma_2) + \mathbf{1}_T\boldsymbol\beta_2$ , the gradients $\mathrm{d}\boldsymbol\gamma_2 = \mathbf{1}_T^\top(\mathrm{d}\mathbf{U}_2\odot\operatorname{LN}(\mathbf{X}_1))$ and $\mathrm{d}\boldsymbol\beta_2 = \mathbf{1}_T^\top\mathrm{d}\mathbf{U}_2$ are zero too. The path from $\mathbf{X}_1$ through the LayerNorm carries nothing, so $\mathrm{d}\mathbf{X}_1 = \mathrm{d}\mathbf{X}_{\text{out}}$ .

The attention sublayer has the same form, with $\mathrm{d}\mathbf{X}_1 = \mathrm{d}\mathbf{X}_{\text{out}}$ arriving from above. That gives item 2 for $\mathbf{g}_1$ , zero gradient for the attention weights, $\boldsymbol\gamma_1$ , and $\boldsymbol\beta_1$ , and $\mathrm{d}\mathbf{X} = \mathrm{d}\mathbf{X}_1 = \mathrm{d}\mathbf{X}_{\text{out}}$ , which is item 5.

The six vectors are one linear function of $\mathbf{s} = \operatorname{SiLU}(\mathbf{c})$ , namely $\mathbf{m} = \mathbf{s}\mathbf{W}_{\text{mod}} + \mathbf{b}_{\text{mod}}$ , and $\mathrm{d}\mathbf{m}$ is nonzero only in the gate slots. The linear-layer gradients of §0.2 give $\mathrm{d}\mathbf{W}_{\text{mod}} = \mathbf{s}^\top\mathrm{d}\mathbf{m}$ and $\mathrm{d}\mathbf{b}_{\text{mod}} = \mathrm{d}\mathbf{m}$ , which is item 3. They also give $\mathrm{d}\mathbf{s} = \mathrm{d}\mathbf{m}\,\mathbf{W}_{\text{mod}}^\top = \mathbf{0}$ , because $\mathbf{W}_{\text{mod}} = \mathbf{0}$ . So nothing upstream of the modulation layer learns from this block, which completes item 4. ∎

The same move appears all over the book. Each time, a new path starts with zero output but is not starved of gradient. ResNets often start the last normalization gain of each residual branch at zero (Goyal et al. 2017), as §0.8 notes, which is the closest cousin of this trick. LoRA starts one of its two matrices at zero (§14.5). Flamingo puts $\tanh$ gates that start at zero on its new layers (§31.2). ControlNet adds convolutions that start at zero (§34.7).

The whole network starts learning at its output. DiT also initializes its final adaptive LayerNorm and its output layer to zero. Follow the first three training steps.

Step 0. The output is identically zero. The output layer gets a gradient, because its input is not zero. The gradient that reaches the last block, $\mathrm{d}\mathbf{X}_L = \mathrm{d}\mathbf{O}\,\mathbf{W}_{\text{out}}^\top$ through the final norm, is zero because $\mathbf{W}_{\text{out}} = \mathbf{0}$ . So only the output layer learns.
Step 1. The output layer is now nonzero, so a gradient reaches every block and passes through each one unchanged (item 5). It moves every gate, the final modulation, and the patch embedding.
Step 2. The gates are now nonzero, so the attention and MLP weights, the shifts and scales, and the timestep and class embeddings receive gradient at last.

In practice, a deep stack starts as an identity map, which is trivially stable, and grows into a useful function one set of zeros at a time. The NumPy cell checks the theorem with finite differences on one sublayer. The PyTorch block further down checks the order of the three steps on a small DiT.

Python · cell 3

The widget runs a toy block of width 12 over 5 tokens, with random attention and MLP weights. Press Train to move its modulation layer from the zero initialization toward a fixed set of stand-in trained values. The gates fill with color first, then the shifts and scales. The grid on the right shows how much the block changes each token, and the curve below shows that change at every noise level $t$ .

The PyTorch block builds a small DiT the way the official code does. It has adaLN-Zero blocks, a final modulation and output layer that start at zero, a timestep MLP on 32 frequency features, and a class table with a null row. It takes three Adam steps and reports which groups of parameters received a nonzero gradient at each step.

🔥 In PyTorchfor a machine with PyTorch

MMDiT, two streams and one attention

A class label fits in one vector. A prompt does not. "A red cube on a blue sphere" and "a blue cube on a red sphere" contain the same words. So any summary that only records which words occur gives them the same vector, and only a model that reads the tokens in order can tell them apart (§30.6 showed the same failure in CLIP). A text-to-image denoiser therefore needs the whole sequence of text-token vectors, $\mathbf{C}_{\text{ctx}} \in \R^{T_c\times d_c}$ for $T_c$ tokens of width $d_c$ , not just one pooled vector. By 2023 there were three ways to read that sequence.

Cross-attention. Image tokens query the text tokens in an extra sublayer (§4.6). Stable Diffusion 1, SDXL, and PixArt-α work this way. The text is read, but the image never changes it, so the same fixed representation of "red" serves every image and every noise level.
One sequence with shared weights. Text tokens and image tokens enter one sequence and pass through the same weights, as in U-ViT (Bao et al. 2022) and DiT's in-context variant. Information flows both ways, but one set of projections must serve two kinds of token with very different statistics.
Joint attention with separate weights. Each modality keeps its own weights for everything except the attention operation itself, which runs over the two sequences stacked together. This is MMDiT (Esser et al. 2024).

📐 MMDiT joint attention

An MMDiT block holds image tokens $\mathbf{X} \in \R^{T\times d}$ and text tokens $\mathbf{Y} \in \R^{T_c\times d}$ . A linear layer at the input first projects the text sequence from width $d_c$ to $d$ . Each stream has its own adaLN-Zero modulation, driven by the same condition $\mathbf{c} = \operatorname{emb}(t) + \operatorname{MLP}(\mathbf{c}_{\text{vec}})$ , where $\mathbf{c}_{\text{vec}}$ is the pooled text vector. For one head of width $d_k$ , the image stream forms

\mathbf{Q}_x = \operatorname{RMSN}_{\mathbf{a}^x_q}\big(\mathbf{U}_x\mathbf{W}^x_Q\big), \quad \mathbf{K}_x = \operatorname{RMSN}_{\mathbf{a}^x_k}\big(\mathbf{U}_x\mathbf{W}^x_K\big), \quad \mathbf{V}_x = \mathbf{U}_x\mathbf{W}^x_V,

from its modulated, normalized tokens $\mathbf{U}_x$ , with $\mathbf{W}^x_Q, \mathbf{W}^x_K, \mathbf{W}^x_V \in \R^{d\times d_k}$ . The text stream forms $\mathbf{Q}_y, \mathbf{K}_y, \mathbf{V}_y$ from $\mathbf{U}_y$ with its own weights. Here $\operatorname{RMSN}_{\mathbf{a}}(\mathbf{q}) = \mathbf{a}\odot\mathbf{q}/\sqrt{\tfrac{1}{d_k}\lVert\mathbf{q}\rVert^2 + \varepsilon}$ is RMSNorm (§5.2) applied to each row of each head, with a learned gain $\mathbf{a} \in \R^{d_k}$ . This is called QK-norm. With $[\,\cdot\,;\,\cdot\,]$ stacking rows, joint attention is

\begin{bmatrix}\mathbf{O}_x\\ \mathbf{O}_y\end{bmatrix} = \operatorname{softmax}\!\left(\frac{1}{\sqrt{d_k}}\begin{bmatrix}\mathbf{Q}_x\\ \mathbf{Q}_y\end{bmatrix}\begin{bmatrix}\mathbf{K}_x\\ \mathbf{K}_y\end{bmatrix}^{\!\top}\right)\begin{bmatrix}\mathbf{V}_x\\ \mathbf{V}_y\end{bmatrix},

with the softmax over each row of the $(T + T_c)\times(T + T_c)$ score matrix. The heads of $\mathbf{O}_x$ are concatenated and projected by the image stream's $\mathbf{W}^x_O$ , and those of $\mathbf{O}_y$ by $\mathbf{W}^y_O$ . Each result is added to its own stream through its own gate, and each stream then runs its own MLP.

In plain words, an MMDiT block is two transformers, one for the picture and one for the prompt, that share a single attention. Every token can read every other token, while each modality keeps its own way of computing queries, keys, values, and features. SD3 compared this design on the CC12M dataset of image-caption pairs. It beat a shared-weight DiT, a DiT with cross-attention, and a U-Net with transformer blocks at its lowest levels, in validation loss, in FID, and in CLIP score, the similarity of each image to its prompt under CLIP (§34.6). A third set of weights, with separate streams for the tokens of SD3's two kinds of text encoder, added little (Esser et al. 2024, Section 5.2.3).

What exactly does an image token get from joint attention? A tiny example answers it before the general statement.

✏️ Joint attention with one query and three keys

An image query scores two image keys at $1$ and $0$ and one text key at $0.5$ , after the $1/\sqrt{d_k}$ scaling. The image values are the numbers 2 and 0, and the text value is 10.

Self-attention over the two image keys alone gives weights $e/(e + 1) = 0.731$ and $0.269$ , so it returns $1.462$ .
Cross-attention to the one text key alone returns 10.
Joint attention divides by all three exponentials, $e + 1 + e^{0.5} = 5.367$ , and returns $(2e + 10e^{0.5})/5.367 = 4.085$ .

Now let $\lambda = (e + 1)/5.367 = 0.693$ be the share of the softmax weight that lands on the image keys. Then $0.693 \cdot 1.462 + 0.307 \cdot 10 = 4.085$ , the same number. Joint attention returned a blend of the two separate attentions, weighted by where the softmax put its weight.

That is no coincidence, and the proof is two lines.

∴ Joint attention is self-attention and cross-attention sharing one softmax

Fix one head and one image query $\mathbf{q}$ (a row of $\mathbf{Q}_x$ ). Let $s_j = \mathbf{q}^\top\mathbf{k}^x_j/\sqrt{d_k}$ for the $T$ image keys, $r_l = \mathbf{q}^\top\mathbf{k}^y_l/\sqrt{d_k}$ for the $T_c$ text keys, $Z_x = \sum_j e^{s_j}$ , $Z_y = \sum_l e^{r_l}$ , and

\lambda = \frac{Z_x}{Z_x + Z_y} = \operatorname{sigmoid}\big(\log Z_x - \log Z_y\big) \in (0, 1),

with $\operatorname{sigmoid}(u) = 1/(1 + e^{-u})$ the logistic function (written out because $\sigma_t$ is the noise scale). Then the joint attention output for $\mathbf{q}$ is

\mathbf{o} = \lambda\,\underbrace{\sum_j \frac{e^{s_j}}{Z_x}\mathbf{v}^x_j}_{\text{self-attention over the image}} + (1 - \lambda)\,\underbrace{\sum_l \frac{e^{r_l}}{Z_y}\mathbf{v}^y_l}_{\text{cross-attention to the text}} .

The same holds for a text query, with the roles exchanged.

Proof. The softmax over all $T + T_c$ keys has denominator $Z_x + Z_y$ . Multiply and divide each partial sum by its own normalizer,

\mathbf{o} = \frac{\sum_j e^{s_j}\mathbf{v}^x_j + \sum_l e^{r_l}\mathbf{v}^y_l}{Z_x + Z_y} = \frac{Z_x}{Z_x + Z_y}\cdot\frac{\sum_j e^{s_j}\mathbf{v}^x_j}{Z_x} + \frac{Z_y}{Z_x + Z_y}\cdot\frac{\sum_l e^{r_l}\mathbf{v}^y_l}{Z_y}.

The two fractions on the right are ordinary softmax attentions over their own keys, and their weights are $\lambda$ and $1 - \lambda$ . Writing $\lambda = 1/(1 + Z_y/Z_x) = 1/(1 + e^{-(\log Z_x - \log Z_y)})$ gives the logistic form. ∎

In plain words, each image token runs self-attention and cross-attention side by side and blends the two results. It chooses the blend itself, through how strongly its query matches the image keys compared with the text keys. This is the same merge of partial softmaxes that FlashAttention uses (§12.2), with the pieces split by modality instead of by memory tile.

Two consequences matter in practice.

The text tokens are updated. The text queries produce $\mathbf{O}_y$ . So the prompt's representation changes from block to block in response to the image, and through its own modulation in response to the noise level. Cross-attention never does this.

The prompt's share starts small. At initialization the scores are roughly alike, so $\lambda$ is close to the share of keys that are image keys, $T/(T + T_c)$ . The next example puts SD3's numbers on it, and on the cost.

✏️ The text's share and the cost of joint attention in SD3

At $1024\times1024$ , SD3's latent is $128\times128$ , and $2\times2$ patches give $T = 4{,}096$ image tokens. The prompt arrives as $T_c = 154$ text tokens, 77 from its CLIP text encoders and 77 from its T5 encoder. At initialization each image query therefore puts about $4{,}096/4{,}250 = 0.964$ of its attention on the image, so the text starts with under 4 percent. A cross-attention sublayer gives the text all of its attention by construction. In joint attention the prompt has to earn its share through learned scores.

The cost changes little. Joint attention scores $(T + T_c)^2 = 4{,}250^2 = 18.06$ million pairs per head. Image self-attention plus cross-attention would score $T^2 + TT_c = 16.78 + 0.63 = 17.41$ million. Joint attention costs 3.8 percent more, and the extra $TT_c + T_c^2$ pairs are exactly what updates the text.

The cell checks the theorem for random image and text streams. It also prints the image share next to the token share.

Python · cell 4

QK-norm. Large transformers trained in 16-bit precision can diverge when their attention scores grow without limit. SD3 hit this problem. It pretrained at $256\times256$ . When it fine-tuned at higher resolution in mixed precision, which computes in a 16-bit format and keeps a 32-bit (fp32) copy of the weights (§8.5), the loss diverged. The cause was attention logit growth (§5.7). The largest attention scores grow until each softmax row puts nearly all its weight on one key, so the entropy of the attention collapses. SD3 saw it in the last transformer blocks.

Training fully in fp32 fixed it, at about half the speed. RMSNorm with a learned gain on the queries and keys of both streams fixed it in mixed precision with the bfloat16 format (bf16), together with AdamW at $\epsilon_{\text{Adam}} = 10^{-15}$ . It even worked when added to a model pretrained without it (Esser et al. 2024, Section 5.3.2). The norm acts on each head's $d_k$ numbers, so each head's scores are bounded by its own gains. In an interview about a run that diverges at high resolution, QK-norm is the first fix to name.

Single-stream blocks. FLUX.1 runs 19 MMDiT blocks and then 38 single-stream blocks, in which the text and image tokens form one sequence and share all weights. The box below describes them, and the counting section uses them.

🔬 Single-stream blocks in FLUX and video models

FLUX.1 runs 19 double-stream blocks, with separate weights and joint attention as above (the paper gives the 38 single-stream blocks, and the released code gives the 19). It then joins the two streams into one sequence and runs 38 single-stream blocks, ordinary transformer blocks with shared weights, before dropping the text tokens (Black Forest Labs 2025). Each single-stream block computes attention and the MLP in parallel from one modulated norm, the parallel layout of ViT-22B, a 22-billion-parameter Vision Transformer (Dehghani et al. 2023). The input projections of the two are fused into one matrix and the output projections into another. So the block needs one modulation of $3d$ numbers (a shift, a scale, and a gate) instead of $6d$ .

Positions are rotary embeddings (§5.5) over three axes, (time, row, column). The open video model HunyuanVideo uses the same plan with 20 double and 40 single blocks (Kong et al. 2024), and the image model Z-Image is single-stream after two small blocks per modality at its input (Z-Image Team 2025). Practitioners give this reasoning. Early blocks profit from separate weights while the two kinds of token still look very different. Later blocks can share weights once the streams are aligned, which buys more compute per parameter. This is a design argument, not a controlled result.

Interview

MMDiT joint attention with QK-norm

Implement joint_attention(X, Y, Px, Py, n_heads, eps=1e-6). The image tokens X have shape (B, T, d) and the text tokens Y have shape (B, T_c, d), both already normalized and modulated. Each stream's dict holds W_qkv of shape (d, 3d), W_o of shape (d, d), and per-head RMSNorm gains g_q and g_k of shape (d_h,). The columns of W_qkv are $Q$ , $K$ , $V$ , with heads contiguous, as in Exercise 34.2d. Follow these steps.

Project each stream with its own weights and split the heads.
Apply RMSNorm to each head's query and key over its $d_h$ entries.
Concatenate image tokens first, then text tokens, and attend with the $1/\sqrt{d_h}$ scale.
Split the output back into the two streams and project each with its own W_o.

Return (out_x, out_y, lam). Here lam has shape (B, n_heads, T) and holds the $\lambda$ of the theorem for every image query and head. Use no biases and no masks.

import numpy as np

def joint_attention(X, Y, Px, Py, n_heads, eps=1e-6):
    # X: (B, T, d) image tokens, Y: (B, T_c, d) text tokens, Px and Py: each stream's weights
    # returns (out_x of shape (B, T, d), out_y of shape (B, T_c, d), lam of shape (B, n_heads, T))
    raise NotImplementedError

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

The order of concatenation is a convention. The diffusers implementation of SD3 puts image tokens first, and FLUX puts text first. Attention without a mask treats its tokens as a set, so either order works. The split must undo the concatenation, and each token's position encoding must travel with it.

Counting parameters and multiply-adds

Counting a model's parameters and compute from its configuration is a standard interview test, and it is how you check any description of an architecture. We count multiply-adds (MACs). One MAC is one multiplication followed by one addition, the basic step of a matrix product. Vision papers call this unit "FLOPs", while this book counts two FLOPs per MAC (§27.2 made the distinction). §33.4 priced a sample with the cost of one DiT-XL/2 forward pass. Here we derive that cost term by term.

✏️ Counting DiT-XL/2 from its configuration

DiT-XL/2 has width $d = 1152$ , depth $L = 28$ , patch size $p = 2$ , $c = 4$ latent channels, and $T = 256$ tokens at $256\times256$ . It uses a 256-dimensional frequency embedding and has 1,000 classes.

Parameters per block. Attention has $\mathbf{W}_{QKV}$ and $\mathbf{W}_O$ with biases, $4d^2 + 4d$ , exactly as in §5.6. The MLP has $8d^2 + 5d$ . The modulation layer maps $d$ numbers to $6d$ , so it has $6d^2 + 6d$ . The LayerNorms have no parameters. A block therefore has $18d^2 + 15d = 23{,}905{,}152$ , and 28 blocks hold 669,344,256.

Outside the blocks. The patch embedding maps $p^2c = 16$ numbers to $d$ , so it has $16d + d = 19{,}584$ . The timestep MLP has $256d + d + d^2 + d = 1{,}624{,}320$ . The class table has $1001\,d = 1{,}153{,}152$ , counting the extra row for the null label. The final modulation has $2d^2 + 2d = 2{,}656{,}512$ , and the output layer $d\cdot p^2\cdot 2c + 32 = 36{,}896$ . The fixed position table is not trained, so we leave it out. The total is 674,834,720, which the paper rounds to 675M.

Multiply-adds per forward pass. Each token passes through $12d^2$ weights per block, $4d^2$ in attention and $8d^2$ in the MLP. The two attention products (scores, then the weighted sum of values) cost $2Td$ per token. So the blocks cost $L(12d^2T + 2T^2d) = 118.38$ G at $T = 256$ . The modulation layers run once per image, $L\cdot 6d^2 + 2d^2 = 0.23$ G. The patch embedding and output layer cost $T\cdot(16 + 32)\,d = 0.014$ G, and the timestep MLP 0.0016 G. The total is 118.62 G, against the paper's 118.6. At $512\times512$ the latent is $64\times64$ , so $T = 1{,}024$ , and the same formula gives 524.54 G, against the paper's 524.6.

Where it goes. Per token and per block, the weights cost $12d^2 = 15.9$ M MACs and attention $2Td = 0.59$ M at $T = 256$ , so attention is 3.6 percent of the block. At $T = 1{,}024$ attention is 12.9 percent. It equals the weights at $T = 6d = 6{,}912$ tokens, the crossover of §27.2. The modulation is a third of all parameters, $28(6d^2 + 6d) = 223$ M, but under 0.2 percent of the compute. Computed per token, it would cost $28\cdot 6d^2\cdot T = 57$ G, almost half again as much as everything else.

That last line explains the parameter column of the conditioning table, and PixArt-α acted on it. PixArt-α is a text-to-image DiT that reads the prompt through cross-attention (Chen et al. 2023), so its modulation only has to carry the time. Its adaLN-single computes one set of six modulation vectors from the time embedding, shared by all blocks, and adds a small learned offset in each block. That cut the model from 833M to 611M parameters, 26 percent fewer, and its GPU memory by 21 percent. Full adaLN kept a clearly lower FID in their ablation, 18.20 against 22.37, though the authors judged its samples visually on par.

The same arithmetic constrains every large model. The next cell counts the weights in the blocks of SD3, FLUX.1, and HunyuanVideo from their published configurations. SD3 sets its width to 64 times its depth. It also counts Movie Gen, Meta's 30-billion-parameter video generator (Polyak et al. 2024). Its blocks follow the layout of Meta's Llama language models, with self-attention, cross-attention to the text, and a gated SwiGLU MLP (§5.3), which has three weight matrices instead of two.

Python · cell 5

The formulas land within a few percent of every published size. SD3's depth-38 model has 8.09B weights in its blocks, matching the "8B" of its paper, and depth 24 gives the 2B model. FLUX.1's 11.83B plus its input layers is the advertised 12B, and HunyuanVideo's 12.46B plus its input layers is its 13B. The formula leaves out only small pieces. They are the text and pooled-vector projections, the timestep MLP, the final layer, the biases, and the norm gains.

The widget plots DiT's twelve models by compute and FID. Pick a size and a patch size, and read the computed cost next to the published one.

Research

The cost model of a diffusion transformer

Implement dit_cost(d, L, T, p, c=4, n_classes=1000, freq_dim=256). It returns (parameters, multiply_adds) for one forward pass of a class-conditional DiT with learned variance, whose output has $p^2\cdot 2c$ numbers per token. Return both counts as Python ints.

Count every trainable parameter. Every linear layer has a bias, and the LayerNorms have no gain or shift. Count the blocks of the definition box (MLP width $4d$ ), the patch embedding (a linear map from $p^2c$ numbers to $d$ ), the timestep MLP from freq_dim to $d$ to $d$ , the class table with n_classes + 1 rows (the extra one is the null label), the final modulation from $d$ to $2d$ , and the output layer. Leave out the frozen position table.
Count the multiply-adds of every matrix product and nothing else. That is the per-token weights, the two attention products, the per-image modulations and timestep MLP, the patch embedding, and the output layer. Biases, norms, activations, the softmax, and the class lookup cost nothing here.

The tests check both counts exactly, then compare your multiply-adds with all twelve rows of DiT's Table 4 and with the $512\times512$ model.

def dit_cost(d, L, T, p, c=4, n_classes=1000, freq_dim=256):
    # returns (parameters, multiply_adds) for one forward pass, both Python ints
    raise NotImplementedError

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

What scaling diffusion transformers has shown

Compute, not parameters, predicts quality. Across DiT's twelve models, FID after 400K steps fell steadily as the forward MACs rose, with a correlation of $-0.93$ between FID and compute, plotted on a log scale. Models of different shapes but similar compute mostly reached similar FID. DiT-S/2 at 6.06 G and DiT-B/4 at 5.56 G both reached 68.4. The exception is patch size 8, whose tokens are too coarse. DiT-B/8 and DiT-S/4 both cost 1.4 G, yet B/8 reached only 122.74 against 100.41. Halving the patch size at a fixed model size quadruples the tokens and the compute while leaving the parameters almost unchanged, and it improved FID at every size. For XL, FID went from 106.41 at $p = 8$ to 43.01 at $p = 4$ and 19.47 at $p = 2$ . Larger models also used training compute better, overtaking smaller models trained for longer (Peebles and Xie 2023, Figures 8 and 9).

Validation loss falls smoothly and predicts what people see. SD3 trained MMDiTs from depth 15 to 38 for 500,000 steps at $256\times256$ . Validation loss fell smoothly with model size and training steps, for images and for videos, without saturating. It correlated strongly with the GenEval and T2I-CompBench benchmarks of prompt following and with human preference. The largest model needed a learning-rate adjustment at 300,000 steps to avoid divergence. Larger models also needed fewer sampling steps to reach their best quality (Esser et al. 2024, Section 5.3.3).

U-Nets against transformers. U-ViT (Bao et al. 2022) is a plain ViT that treats the time and the condition as tokens and adds long skip connections between shallow and deep blocks. It matched a U-Net of similar size and reached FID 2.29 on ImageNet $256\times256$ . Its ablations showed that the long skips mattered while the down and up sampling did not. SDXL moved its U-Net's transformer blocks to the lowest resolutions, a step toward a transformer. In SD3's comparison, a U-Net with transformer blocks at its lowest levels (Hoogeboom et al. 2023) learned faster at first than the DiT with cross-attention. The cross-attention DiT overtook it, and MMDiT beat both.

Every frontier image and video model since 2024 whose architecture is published is a transformer, including SD3, FLUX, Movie Gen, Wan, and HunyuanVideo. That is evidence that transformers scale better, though not a controlled experiment. At small compute a well-tuned U-Net remains a strong baseline. EDM2 redesigned a convolutional U-Net's layers so that the sizes of activations and weights stay steady during training, and it reached FID 1.81 on ImageNet $512\times512$ (Karras et al. 2024).

🔬 Scaling laws and representation alignment

Power laws. Liang et al. 2024 fit power laws of pretraining loss against compute for text-to-image DiTs between $10^{17}$ and $6\times10^{18}$ FLOPs. They derived compute-optimal model and data sizes the way §9.1 does for language models, and predicted the loss of a 1B-parameter model trained with $1.5\times10^{21}$ FLOPs.

Representation alignment. Yu et al. 2024 observed that diffusion transformers learn semantic features slowly. They asked what happens if an early layer is told what good features look like. Their method, REPA, uses SiT, a DiT trained with a flow-matching objective (Ma et al. 2024). It takes the hidden states after block 8 and passes them through a small MLP. A loss then rewards them for matching, patch by patch, the features that a frozen DINOv2 encoder, a self-supervised vision transformer (§28.4), computes from the clean image. The loss is

\mathcal{L} = \mathcal{L}_{\text{velocity}} + \lambda_{\text{REPA}}\,\mathcal{L}_{\text{REPA}}, \qquad \mathcal{L}_{\text{REPA}} = -\mathbb{E}\Big[\frac{1}{T}\sum_{i=1}^{T}\cos\big(\mathbf{y}^{\star}_i, \operatorname{proj}(\mathbf{x}_{8,i})\big)\Big],

with weight $\lambda_{\text{REPA}} = 0.5$ . Here $\mathbf{y}^\star_i$ is the encoder's feature of patch $i$ , $\mathbf{x}_{8,i}$ is the hidden state of token $i$ after block 8, and $\operatorname{proj}$ is the small trainable MLP. In plain words, the usual flow-matching loss is joined by a term that pulls each token's early features toward a good image encoder's features for the same patch.

SiT-XL/2 with REPA matched in under 400K steps the FID that the plain model reached at 7M, more than 17.5 times faster. It reached FID 1.42 with guidance applied over a limited interval of noise levels. Aligning an early layer worked better than aligning deep ones, which the authors read as letting later layers specialize in detail. Whether alignment still pays at the largest scales and with text conditioning is open.

A narrow denoiser has a loss floor. One more result sets a hard rule for the width. A representation autoencoder uses a frozen DINOv2 encoder as its encoder, so each latent token is a 768-dimensional feature vector (§32.6, §34.1). Zheng et al. 2025 found that a DiT narrower than 768 could not even fit a single image. The reason is a short theorem. Each output token is a linear map of a $d$ -dimensional hidden state, so it can only move in $d$ directions, while the target varies in all $n$ .

∴ Width bounds the flow-matching loss from below

Let each token of the target be $\mathbf{w} = \boldsymbol\epsilon - \mathbf{z}_0 \in \R^n$ , the velocity of §33.6 restricted to one token, with covariance $\boldsymbol\Sigma_{\mathbf{w}}$ and eigenvalues $\nu_1 \ge \dots \ge \nu_n$ . Suppose the network's output for that token is $\hat{\mathbf{w}} = \mathbf{x}\mathbf{W}_{\text{out}}$ for some hidden state $\mathbf{x} \in \R^{1\times d}$ with $d < n$ and an output matrix $\mathbf{W}_{\text{out}} \in \R^{d\times n}$ . Then at every time $t$ , whatever the rest of the network,

\mathbb{E}\,\lVert\hat{\mathbf{w}} - \mathbf{w}\rVert^2 \;\ge\; \sum_{i=d+1}^{n}\nu_i \;\ge\; n - d ,

where the second inequality holds when $\boldsymbol\epsilon\sim\mathcal{N}(\mathbf{0}, \mathbf{I})$ is independent of $\mathbf{z}_0$ .

In plain words, a token's output can only point in $d$ directions, and the noise alone puts unit variance in all $n$ . So at least $n - d$ units of loss per token can never be removed. For a DiT on Stable Diffusion's latents, $n = p^2c = 16$ , or 64 for FLUX's 16 channels, far below any width, and the bound never bites. For 768-dimensional representation tokens, a DiT-S of width 384 is stuck at a loss of at least 384 per token, half of the 768 units of variance that the noise alone puts in each token. Zheng et al. widened the denoiser, or added a wide shallow head, to fix it. The interview lesson is that the width must exceed the number of values each token predicts.

∂ Proof of the width bound

The output always lies in the row space $\mathcal{S}$ of $\mathbf{W}_{\text{out}}$ , of dimension at most $d$ . Let $\mathbf{P}$ be the orthogonal projector onto $\mathcal{S}$ , so $\mathbf{I} - \mathbf{P}$ projects onto a subspace of dimension $m \ge n - d$ . Since $\hat{\mathbf{w}} \in \mathcal{S}$ , Pythagoras gives $\lVert\mathbf{w} - \hat{\mathbf{w}}\rVert^2 = \lVert(\mathbf{I} - \mathbf{P})\mathbf{w}\rVert^2 + \lVert\mathbf{P}\mathbf{w} - \hat{\mathbf{w}}\rVert^2 \ge \lVert(\mathbf{I} - \mathbf{P})\mathbf{w}\rVert^2$ .

Treat $\mathbf{w}$ as a column for the moment. Then $\mathbb{E}\lVert(\mathbf{I} - \mathbf{P})\mathbf{w}\rVert^2 = \operatorname{tr}\big((\mathbf{I} - \mathbf{P})\,\mathbb{E}[\mathbf{w}\mathbf{w}^\top]\big) \ge \operatorname{tr}\big((\mathbf{I} - \mathbf{P})\boldsymbol\Sigma_{\mathbf{w}}\big)$ . The inequality holds because $\mathbb{E}[\mathbf{w}\mathbf{w}^\top] = \boldsymbol\Sigma_{\mathbf{w}} + \boldsymbol\mu\boldsymbol\mu^\top$ and $\operatorname{tr}\big((\mathbf{I} - \mathbf{P})\boldsymbol\mu\boldsymbol\mu^\top\big) = \lVert(\mathbf{I} - \mathbf{P})\boldsymbol\mu\rVert^2 \ge 0$ .

Now expand $\boldsymbol\Sigma_{\mathbf{w}} = \sum_k\nu_k\mathbf{e}_k\mathbf{e}_k^\top$ in its orthonormal eigenvectors (the spectral theorem of §0.10). Then $\operatorname{tr}\big((\mathbf{I} - \mathbf{P})\boldsymbol\Sigma_{\mathbf{w}}\big) = \sum_k\nu_k a_k$ with $a_k = \lVert(\mathbf{I} - \mathbf{P})\mathbf{e}_k\rVert^2 \in [0, 1]$ and $\sum_k a_k = \operatorname{tr}(\mathbf{I} - \mathbf{P}) = m$ . A weighted sum of the $\nu_k$ with weights in $[0, 1]$ that add up to $m$ is smallest when weight 1 sits on the $m$ smallest eigenvalues. So it is at least $\sum_{i > n - m}\nu_i \ge \sum_{i > d}\nu_i$ .

For the last inequality, independence gives $\boldsymbol\Sigma_{\mathbf{w}} = \mathbf{I} + \operatorname{Cov}(\mathbf{z}_0)$ , so every $\nu_i \ge 1$ and the sum of $n - d$ of them is at least $n - d$ . An output bias changes nothing, since it moves $\hat{\mathbf{w}}$ within a shifted copy of $\mathcal{S}$ and only the mean term above absorbs it. ∎

That closes the theory of the backbone. What remains is practice, how these networks are trained at scale and where implementations of them go wrong.

🔬 Training recipes for diffusion transformers

DiT trained every model with AdamW at a constant learning rate of $10^{-4}$ , no weight decay, and batch 256. Horizontal flips were the only augmentation, there was no warmup, and sampling used an exponential moving average (EMA) of the weights with decay 0.9999. It reported no loss spikes at any size. Larger text-to-image models train in bf16 mixed precision (§8.5) with the norms and the softmax in fp32. They use QK-norm and tune the EMA horizon. EDM2's post-hoc EMA lets you choose that horizon after training from saved snapshots (§33.7).

At very high resolution the attention term dominates. SANA combined an autoencoder that downsamples 32 times, which leaves fewer tokens (§34.1), with linear attention (§21.2). Its 0.6B model is 20 times smaller and over 100 times faster in measured throughput than the 12B FLUX, with competitive quality (Xie et al. 2024).

The table puts the five backbones of this section side by side, with what each one does well and badly.

backbone	how $t$ and $\mathbf{c}$ enter	how text enters	strengths	weaknesses	used by
convolutional U-Net	in every residual block, added (DDPM, SD) or by AdaGN (ADM)	cross-attention at low resolutions	strong built-in assumptions about images, efficient at small compute, any resolution	hard to scale cleanly, many design choices to tune	DDPM, ADM, SD 1 and XL, EDM2
U-ViT	as tokens	as tokens, shared weights	simple, long skips help early training	shared weights for unlike tokens	U-ViT
DiT	adaLN-Zero	none (class labels) or added cross-attention	scales predictably with compute, identity start	quadratic in $T$ , sequence conditioning needs extra layers	DiT, SiT, PixArt-α (with cross-attention)
MMDiT	adaLN-Zero per stream	joint attention, separate weights	text and image update each other, best in SD3's comparison	twice the weights per block	SD3, FLUX double blocks
single-stream	adaLN, one shift, scale, and gate per block	joint sequence, shared weights	more compute per parameter, fused kernels	early blocks lose modality-specific weights	FLUX single blocks, HunyuanVideo, Z-Image

⚠️ Where diffusion backbones go wrong

Writing the modulation as $\mathbf{x}\odot\boldsymbol\gamma + \boldsymbol\beta$ instead of $\mathbf{x}\odot(\mathbf{1} + \boldsymbol\gamma) + \boldsymbol\beta$ , with zero initialization, multiplies every normalized token by zero. The branch then sees a constant input and learns slowly.
Broadcasting a (B, d) modulation against (B, T, d) tokens without a token axis is an error when $B \ne T$ and a silent bug when $B = T$ .
Zero-initializing the attention and MLP weights along with the gates leaves the branches with no gradient forever, since each factor's gradient is proportional to the other.
Unpatchifying with a channel-major layout, or splitting the noise and variance channels in the wrong order, scrambles the output while the loss still decreases somewhat. Test that unpatchify inverts patchify.
Splitting a text-first concatenation as if it were image-first hands text rows to the image stream and image rows to the text stream, with no shape error to warn you.
A GroupNorm whose $G$ does not divide the channel count after a skip concatenation crashes at the first block on the way up.
Loading a checkpoint with a timestep embedding of the wrong layout (sines first, or the $k - 1$ exponent) gives plausible but wrong outputs.
Sampling with the raw weights instead of the EMA weights costs FID even when everything else is right.

🧭 What to remember

A denoiser maps a noisy latent, the time, and a condition to an array of the same shape. It must see every scale, and the time must reach every layer.
The diffusion U-Net goes down to low resolution for global layout and back up with skip connections for detail. It uses GroupNorm (32 groups, per example) and puts attention only at low resolutions, because attention's cost grows with the fourth power of the side.
Time and class enter by adaptive normalization, a per-channel scale and shift computed from $\mathbf{c} = \operatorname{emb}(t) + \operatorname{emb}(y)$ and written $\operatorname{LN}(\mathbf{x})\odot(\mathbf{1} + \boldsymbol\gamma) + \boldsymbol\beta$ . In DiT, adaLN-Zero beat in-context and cross-attention conditioning (FID 19.47 against 35.24 and 26.14).
adaLN-Zero makes every block the identity at initialization. Only the gates receive gradient at first, so learning starts at the output and spreads backward.
MMDiT keeps separate weights per modality and one joint attention. For an image query it equals $\lambda$ self-attention plus $1 - \lambda$ cross-attention, with $\lambda = Z_x/(Z_x + Z_y)$ , and it also updates the text. QK-norm keeps it stable in bf16.
DiT-XL/2 is $L\cdot18d^2$ plus embeddings, 675M parameters, and $L(12d^2T + 2T^2d)$ plus small terms, 118.6 G MACs at 256 and 524.6 G at 512. The same counting gives SD3 8B, FLUX.1 12B, and HunyuanVideo 13B, and compute, not parameters, tracks FID.
These blocks have one slot for global conditions and, in MMDiT, a joint attention that can read a sequence. §34.3 asks what goes in those slots when the condition is a sentence.

Ronneberger, Fischer & Brox (2015). U-Net: Convolutional Networks for Biomedical Image Segmentation. The contracting and expanding paths with skip connections that every diffusion U-Net inherits.
Dhariwal & Nichol (2021). Diffusion Models Beat GANs on Image Synthesis. Section 3 is the ablation that built ADM, with attention at several resolutions, 64-channel heads, residual blocks for resampling, and adaptive group normalization.
Wu & He (2018). Group Normalization. The definition, its relation to LayerNorm and InstanceNorm, and why it does not depend on the batch size.
Perez, Strub, de Vries, Dumoulin & Courville (2017). FiLM: Visual Reasoning with a General Conditioning Layer. Feature-wise linear modulation, the idea behind every adaptive normalization in this section.
Peebles & Xie (2023). Scalable Diffusion Models with Transformers. DiT, the four conditioning designs, adaLN-Zero, and Table 4, whose twelve models this section counts.
Esser et al. (2024). Scaling Rectified Flow Transformers for High-Resolution Image Synthesis. Stable Diffusion 3, with MMDiT in Section 4, the backbone comparison and QK-normalization in Section 5, and the scaling study.
Bao et al. (2022). All are Worth Words: A ViT Backbone for Diffusion Models. U-ViT, which shows that long skips, not down and up sampling, are what a transformer denoiser needs from the U-Net.
Yu et al. (2024). Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think. REPA, its alignment depth and target-encoder ablations, and the 17.5 times speedup.
Chen et al. (2023). PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis. A cross-attention DiT for text, with adaLN-single, the shared modulation this section counts.
Zheng et al. (2025). Diffusion Transformers with Representation Autoencoders. Diffusion on DINOv2 features, and the finding that the denoiser must be at least as wide as its tokens.

Made for larger screens

U-Nets, DiT, and MMDiT

What the denoiser has to do

The diffusion U-Net

Group normalization

GroupNorm for a batch of feature maps

Telling every layer the noise level

Timestep embedding and modulation

The diffusion transformer

Unpatchify the output head

A DiT block with adaLN-Zero, batched, in NumPy

Why adaLN-Zero starts as the identity and still learns

MMDiT, two streams and one attention

MMDiT joint attention with QK-norm

Counting parameters and multiply-adds

The cost model of a diffusion transformer

What scaling diffusion transformers has shown

Read more

Watch next