The Diffusion Transformer, from first principles

In 2022 William Peebles and Saining Xie trained Diffusion Transformers (DiT): latent diffusion models whose denoiser, the network run at every sampling step, is a plain Transformer over patches of the latent instead of a convolutional U-Net. Their largest, DiT-XL/2, reached an FID of 2.27 on class-conditional ImageNet 256×256. Stable Diffusion 3 (Esser et al., 2024) scaled the idea for text-to-image with two changes: a text stream that attends jointly with the image, and a rectified-flow objective. The open image and video models this site has covered since, from Mage-Flow to Qwen-Image-2.1, descend from both. This page builds the block and then counts DiT-XL/2's 675,129,632 parameters and its multiply-adds.

What the denoiser is for

A diffusion model learns to undo noise. Training takes a clean sample $x_0$, draws Gaussian noise $\epsilon$ and a noise level $t$, mixes them, and asks the network to recover what was added. DiT used the DDPM form, with class label $y$:

There are 1,000 noise levels on a linear schedule. Sampling starts from pure noise and calls the network step after step, removing some predicted noise each time: 250 steps in the paper's evaluation. DiT also predicts a diagonal covariance, so its output has twice the input's channels.

Work in a latent space

A 256×256 RGB image is 196,608 numbers. DiT never sees one. Stable Diffusion's pretrained autoencoder (84M parameters) shrinks the image 8 times along each side into a 32×32×4 latent: 4,096 numbers, 48 times fewer. Diffusion runs entirely there; the decoder turns the final latent back into pixels once. This is latent diffusion; DiT kept it and replaced only the denoiser.

Latents have grown richer since. SD3 found 16 channels scale better than 4, and Qwen-Image-2.1's autoencoder has 64 channels at 16 times per side. The decoder, working at full resolution, dominates memory: on Qwen-Image-2.1 at 512², stable-diffusion.cpp's compute buffers were 234.03 MB for the denoiser and 2,450.50 MB for the VAE.

Patches become tokens

DiT cuts the latent into non-overlapping $p \times p$ patches and maps each one, $p \cdot p \cdot 4$ numbers, to a vector of width $d$ with one linear layer. A latent of side $I$ becomes $T = (I/p)^2$ tokens: at $I = 32$, that is 256 tokens for $p = 2$, 64 for $p = 4$ and 16 for $p = 8$. Fixed 2D sine-cosine position embeddings are added, as in a vision Transformer. A final linear layer maps each token back to a patch.

Patch size is a compute knob: halving $p$ quadruples $T$ and barely changes the weights. The "/2" in DiT-XL/2 is $p = 2$.

The block, and adaLN-Zero

A DiT block is the pre-norm Transformer block, self-attention and then a 4× MLP with GELU, each with a residual add, with one change: the norms take orders. The timestep becomes a vector through a sinusoidal embedding and a small MLP, the label through an embedding table, and their sum $c$ drives every block. From $c$, one linear layer per block regresses six vectors of width $d$:

$\gamma$ and $\beta$ replace the LayerNorm's learned scale and shift, which is adaptive LayerNorm (adaLN); the LayerNorms themselves hold no parameters. $\alpha$ gates each residual branch. The "Zero" is the initialisation, in the paper's words: "We initialize the MLP to output the zero-vector for all α; this initializes the full DiT block as the identity function." Of the four ways the paper tried to feed in the condition, adaLN-Zero scored best; cross-attention to it cost the most, "roughly a 15% overhead" in Gflops.

Where DiT-XL/2's parameters go

DiT-XL has 28 blocks of width $d = 1{,}152$ with 16 heads, and every linear layer has a bias.

The modulation is $6d^2$ per block against attention's $4d^2$ and the MLP's $8d^2$: a third of every block, 225,803,520 parameters in all, spent on turning one conditioning vector into scales. Successors have cut it. PixArt-α computes the modulation once and gives each block a learned offset, adaLN-single, for 26% fewer parameters; Wan shares one modulation MLP across its blocks; Qwen-Image-2.1 shares one projection of 67,108,864 parameters, 0.94% of its denoiser.

Multiply-adds. Per token per block, the linear layers cost $12d^2 = 15{,}925{,}248$ multiply-adds, and attention adds $2Td$: $T$ query-key products and $T$ weighted values, each $d$ wide. At 256 tokens, 28 blocks come to 114,152,177,664 in the linear layers and 4,227,858,432 in attention, 118.4 billion together; the embedders, the modulation and the norms bring it to the paper's 118.6 Gflops. So the paper counts a multiply-add as one flop, and in the two-per-multiply-add convention of the Transformer page a DiT-XL/2 step is about 237 GFLOPs.

How DiT replaced the U-Net

Latent diffusion's U-Net is convolutional: it downsamples, upsamples, joins matching resolutions with skips and interleaves attention. SD 1.5's has 859,520,964 parameters, 68.5% of them convolution weights. DiT has no convolution beyond its patch embedding, no downsampling and no long skips: one resolution, one block, repeated. The paper's claim was that this scales predictably: "DiTs with higher Gflops -- through increased transformer depth/width or increased number of input tokens -- consistently have lower FID." DiT-XL/2 at 118.6 Gflops beat the latent U-Net LDM-4 at 103.6 and the pixel-space U-Net ADM at 1,120.

It won because a Transformer denoiser is a Transformer: language-model kernels, parallelism and quantisers apply unchanged, and a sequence takes text or video frames as easily as patches. stable-diffusion.cpp's loader shows the result: a UNet for SD 1.x, 2.x and XL, and a DiT for everything since SD3.

Text: MM-DiT and single-stream blocks

DiT conditioned on one of 1,000 labels. A prompt is a sequence, and SD3 treats it as one. Its MM-DiT block uses "two separate sets of weights for the two modalities", in effect "two independent transformers for each modality, but joining the sequences of the two modalities for the attention operation". Text tokens, from CLIP-L, CLIP-G and T5-XXL, and image tokens get their own norms, projections, MLPs and modulation; the two sequences are concatenated inside attention, so image reads text and text reads image in every block. A pooled text vector joins the timestep in $c$. SD3 sizes the model by its depth $d$: width $64d$, $d$ heads, $d$ blocks. At $d = 38$ the width is 2,432, each block holds about $2 \times 18 \times 2{,}432^2$ weights, and 38 blocks come to about 8.09 billion, the paper's 8B. Queries and keys are RMS-normalised, which keeps attention logits from growing at high resolution.

Two streams double the weights per unit of width, and the text half serves a few hundred tokens at most. So designs have moved toward single-stream blocks: text is projected once into the image stream's width, and one set of weights serves every token. FLUX.1-dev mixes the two. Its 11,901,408,320 parameters decompose exactly as 19 two-stream blocks and 38 single-stream blocks at width 3,072, and 3,247,448,064 of them (27.3%) are modulation projections. Mage-Flow is a 4B MM-DiT that packs images of any aspect ratio into variable-length sequences with per-sample 2D RoPE. Qwen-Image-2.1 went all the way: 32 single-stream blocks at width 4,096, an MLP ratio of 3 and one shared modulation projection, 7,115,124,736 parameters, down from the 60 two-stream blocks and 20,430,401,088 parameters of Qwen-Image 1.0. Its attention is block-causal and its prompt is modulated as if at $t = 0$, so the prompt's keys and values can be cached across steps.

Rectified flow replaced noise prediction

SD3 also changed the target. Rectified flow (Liu et al.; see also flow matching) mixes data and noise along a straight line and regresses the velocity along it:

Sampling integrates from $t = 1$ to $t = 0$, most simply by Euler steps $x_{t-\Delta} = x_t - \Delta\,v_\theta(x_t, t)$. If each noise sample led to one image, the path would be straight and one step would do. It curves because at high noise the ideal velocity points at the average of every image the prompt allows, which is why Qwen-Image-2.1 samples in 40 steps and its distilled students in 5 to 8. SD3 compared 61 formulations and kept rectified flow with timesteps drawn from a logit-normal distribution, which spends training on the middle of the path, plus a shift toward high noise that grows with resolution, α = 3.0 at 1024². The network is the same DiT with a different target.

Tokens grow with the square of the resolution

Parameters do not depend on resolution; tokens do. DiT (8× autoencoder, 2×2 patches), SD3 and FLUX (8×, 2×2) and Qwen-Image-2.1 (16×, 1×1) all land on one token per 16×16 pixels. For DiT-XL/2, attention's share of the two main terms is $T/(6d + T)$:

Four times the tokens is four times the linear work and sixteen times the attention work, so each doubling of the side costs more than the last: 4.4 times from 256² to 512² (the paper reports 524.6 Gflops there) and 5.5 times to 1024². The last row is arithmetic; DiT was never trained at that size. Video multiplies again. Wan's autoencoder compresses 4 times in time and 8 in space, with 16 channels, and its patches are 1×2×2, so an 81-frame 832×480 clip is 21 × 30 × 52 = 32,760 tokens: about 8 times a 1024² image, with about 64 times the attention per layer.

What it is good and bad at

Good. One block at every scale; any conditioning is more tokens; parameters are decoupled from resolution.

Bad. Attention is quadratic in tokens, which bites at high resolution and in video. Every sampling step runs the whole network over every image token: 40 passes of a 7B model per Qwen-Image-2.1 image, which is why distillation to a handful of steps matters (Mage-Flow's 4-step Turbo renders 1024² in 0.59 s on an A100). Naive adaLN costs a third of a block, and the autoencoder's decoder dominates memory.

What changed since 2022

Latents went from 4 channels to 16, then 64 at 16× compression. Class labels became prompts read by large text encoders such as Qwen3-VL, and joint attention is giving way to single-stream blocks. Noise prediction became velocity under rectified flow, and 250 steps became 40, then 4 to 8 after distillation. Fixed sine-cosine positions became 2D rotary embeddings, packed at native resolution; queries and keys gained RMSNorm; and modulation is increasingly shared across blocks. The block itself, attention and an MLP on a residual stream steered by a scale, a shift and a gate, is the 2022 design.