22. Diffusion Transformers

Diffusion Transformers (DiT)

In 2023, Peebles & Xie¹ demonstrated something simple and impactful: the U-Net is not necessary for diffusion models. By replacing it with pure Transformer blocks, the model not only maintained quality — it started to scale predictably with more parameters and data, exactly like language models.

Today, all state-of-the-art image and video generation uses DiT:

From U-Net to Transformer

The classic U-Net uses convolutions with hierarchical skip connections — good for capturing local details, but difficult to scale. DiT replaces all of this with global attention blocks.

Step 1 — Patchify: Images as Token Sequences

Just like ViT divides images into patches, DiT operates in the latent space (after the VAE encoder). A latent of shape \(H \times W \times C\) is divided into patches of size \(p \times p\):

Each patch is flattened and projected to dimension \(d_{\text{model}}\) — becoming a "visual token".

Step 2 — DiT Block with AdaLN

DiT uses Adaptive Layer Normalization (AdaLN) to inject timestep and class/text information directly into the normalization parameters:

where \(c = \text{MLP}(\text{emb}(t) + \text{emb}(\text{class}))\) is the conditioning vector.

The parameters \(\gamma\) and \(\beta\) are predicted — not statically learned — making normalization sensitive to the diffusion step and the prompt.

Step 3 — MMDiT: Multi-Modal Bidirectional Attention

MMDiT (SD3, FLUX) goes beyond cross-attention conditioning. Text and image participate in the same attention operation:

Image tokens see text tokens and vice versa — much richer conditioning than injecting text only via cross-attention.

FLUX uses a "dual stream" design: separate weights for image and text in Q/K/V/FFN blocks, but shared attention:

Img stream:  x_img → W_q^img·x  ─┐
                                   ├─→ concat → Attention(Q,K,V) → split
Txt stream:  x_txt → W_q^txt·x  ─┘

Visualization: Complete Generation Process

What DiT actually bought

The DiT paper's real contribution is not "a Transformer works here too". It is that swapping the backbone gave image generation the property language modelling already had: a scaling law you can plan against¹. Peebles and Xie showed FID falling smoothly and predictably with training compute across four model sizes and three patch sizes, with no sign of the plateau U-Nets hit.

That is what made the next generation possible. You cannot justify a 12B-parameter image model without a curve telling you what 12B buys.

The token counts are for each model's own training resolution, and they are the quantity that matters: SD3's depth follows its width by \(d = 64 \times \text{blocks}\), so the 8B configuration is the same design at 38 blocks and 2432 channels.

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

• One architecture for everything. The same block, the same kernels, the same distributed-training machinery as an LLM. Sequence parallelism, FlashAttention, activation checkpointing, FSDP — all of it transfers unchanged.
• Modality is just more tokens. Adding video frames, audio, or a second image is concatenation, not a new architecture. This is the whole reason video models are DiTs.
• Resolution is a sequence length. No architectural commitment to 512 or 1024; you change \(N\). The cost is the \(O(N^2)\) from chapter 13, which is why high-resolution DiTs lean on 16-channel latents and larger patches rather than more pixels.

Where the design is going

Simplified Implementation

import torch
import torch.nn as nn

class AdaLN(nn.Module):
    def __init__(self, d_model, d_cond):
        super().__init__()
        self.norm = nn.LayerNorm(d_model, elementwise_affine=False)
        self.proj = nn.Linear(d_cond, 2 * d_model)  # → γ, β

    def forward(self, x, c):
        gamma, beta = self.proj(c).chunk(2, dim=-1)
        return (1 + gamma.unsqueeze(1)) * self.norm(x) + beta.unsqueeze(1)

class DiTBlock(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, d_cond):
        super().__init__()
        self.adaln1 = AdaLN(d_model, d_cond)
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.adaln2 = AdaLN(d_model, d_cond)
        self.ff = nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model))

    def forward(self, x, c):
        h = self.adaln1(x, c)
        x = x + self.attn(h, h, h)[0]      # self-attention
        x = x + self.ff(self.adaln2(x, c)) # FFN
        return x

class DiT(nn.Module):
    def __init__(self, in_channels, patch_size, d_model, n_heads, d_ff, n_layers, d_cond):
        super().__init__()
        self.patch_size = patch_size
        p = patch_size
        self.patchify = nn.Conv2d(in_channels, d_model, p, stride=p)
        self.blocks = nn.ModuleList([DiTBlock(d_model, n_heads, d_ff, d_cond) for _ in range(n_layers)])
        self.norm_out = nn.LayerNorm(d_model)
        self.depatchify = nn.Linear(d_model, p*p*in_channels)

    def forward(self, x, t_emb, cond):
        # x: (B, C, H, W) noisy latent
        B, C, H, W = x.shape
        tokens = self.patchify(x)                       # (B, d, H/p, W/p)
        tokens = tokens.flatten(2).transpose(1, 2)      # (B, N, d)
        c = t_emb + cond                                # combine conditioning
        for block in self.blocks:
            tokens = block(tokens, c)
        tokens = self.norm_out(tokens)
        patches = self.depatchify(tokens)               # (B, N, p*p*C)
        # reshape back to (B, C, H, W)
        p = self.patch_size
        patches = patches.view(B, H//p, W//p, p, p, C).permute(0,5,1,3,2,4).reshape(B,C,H,W)
        return patches  # predicted velocity field v_θ(x_t, t)

Key takeaways

1. A DiT is a ViT operating on noisy latent patches. Patchify, add position, run \(L\) Transformer blocks, un-patchify to a velocity or noise prediction.
2. AdaLN-Zero injects conditioning by modulating normalization rather than by concatenating a token, and the zero-initialized gate makes each block start as the identity — the same safe-addition trick as ControlNet and LoRA.
3. The contribution is a scaling law for image generation. Smooth, predictable FID versus compute is what justified 12B-parameter image models.
4. MMDiT makes text a full participant with its own weights and joint attention. That is where legible text in generated images came from.
5. One architecture now covers language, vision, image generation and video, and every infrastructure investment transfers between them.
6. DiT wins at scale; a U-Net still wins when data and compute are scarce. And \(O(N^2)\) in tokens is the constraint that shapes every high-resolution design.