DiT: a Transformer inside latent diffusion

An interactive field guide

DiT changes the denoising network. The original system still uses a VAE to convert between images and latents.

Before this article: Diffusion and flow matching · ViT · Variational autoencoders

What was already possible

Latent diffusion already reduced sampling cost by denoising a compressed image space. U-Nets provided the denoising backbone. DiT studies a scalable Transformer backbone in that role, with a controlled comparison of model size, patch size, and conditioning designs.

What changes, and why

Patchify the noisy VAE latent, add fixed 2D sine-cosine positions, and process all positions with Transformer blocks. Timestep and class embeddings modulate normalization and residual gates through adaLN-Zero. Zero-initialized gates initially close residual branches. The output head predicts noise and variance terms, then unpatchifies the result.

Read from left to right. Select a module to connect its role, tensor shape, and calculation. Snippets are teaching pseudocode. Sources provide the paper method or audited implementation.

Latent patches

Patch size 2 reduces the latent grid, not the original pixel grid.

z = patch_embed(noisy_latent) + fixed_2d_position

adaLN-Zero blocks

Time and class control shift, scale, and residual gates.

shift, scale, gate = modulation(time + class_embedding)
h = layer_norm(z)*(1+scale)+shift
z = z + gate*attention(h)

Denoising then pixels

The original output carries noise and variance channels; the separate VAE outputs pixels.

noise, variance = unpatchify(head(z)).chunk(2,dim=1)
clean_latent = diffusion_sampler(model)
image = vae.decode(clean_latent)

Inspect the pinned source

h\leftarrow h+\alpha(c)\,F((1+\gamma(c))\operatorname{LN}(h)+\beta(c))

Work through the mechanism

For 256² pixels, a common VAE produces 32 × 32 × 4 latents. DiT patch size 2 gives 16 × 16 Transformer positions. These are patches of latents, not 2 × 2 RGB patches. One denoising call updates all noisy positions; it is not left-to-right autoregression.

The gate demo starts with α=0. Change the condition scale γ: the branch remains closed. Then increase α to let the condition affect the residual output. The example isolates gating, not full layer normalization or learned timestep embeddings.

Change a parameter, inspect the computed values, then compare this small example with the model described above.

Analytical teaching model · controls change the example, not a trained research model.

Residual gate α 0Condition scale γ 1

At α = 0 the residual branch is closed, even when its condition changes. This is a scalar illustration of gated residuals.

Computed values for the current example
Quantity	Value
x	-1, 0, 1
x + α(γx + 1)	-1.000, 0.000, 1.000

Training is not sampling

Train the denoiser on noisy latents from an existing VAE. During generation, sample a latent through repeated denoising calls, then call the VAE decoder for RGB. The decoder comparison in Appendix D uses the same encoder and DiT-XL/2 with CFG=1.5; it tests final generation quality, not tokenizer reconstruction.

Reported evidence, not a local reproduction

Setting	Metric	Value	Interpretation
ImageNet 256²; LDM / ft-MSE / ft-EMA decoder	gFID ↓	2.46 / 2.30 / 2.27	Appendix D, same DiT-XL/2 and CFG 1.5

What this does not establish

Original DiT is not a replacement image tokenizer, and it does not use RoPE. To use a Transformer for reconstruction itself, supply the image code as a condition to a pixel denoiser, as FlowMo does. The pinned DiT repository uses CC-BY-NC-4.0; the snippets here are independently written teaching pseudocode.

Check your understanding

After DiT produces a clean latent, is the RGB image ready?

No. The VAE decoder must still map the latent to pixels.

Sources and next steps

Return to the learning map

Diagrams and snippets are teaching reconstructions. Paper reports, implementation observations, and analytical examples are distinguished. No research model runs in this page.