Diffusion Transformer
(DiT)

Chapter 1: The U-Net Problem

To understand why DiT is a breakthrough, you first need to understand exactly how the U-Net worked in diffusion models — and precisely where it falls short.

The U-Net is an encoder-decoder architecture. The encoder repeatedly downsamples the feature map (halving spatial dimensions each level) while increasing channels. The decoder does the reverse — upsampling back to the original resolution. At each level, skip connections pass the encoder's features directly to the corresponding decoder layer, preserving fine-grained spatial detail.

In diffusion, the U-Net takes a noisy image [B, C, H, W] and a timestep embedding, and outputs a predicted noise map of the same shape. The timestep is injected via addition into each ResNet block. Cross-attention layers at low-resolution levels handle text conditioning in models like Stable Diffusion.

The Scaling Problem

Here's where the U-Net runs into trouble. To make a U-Net bigger, you have several options: increase channel counts, add more resolution levels, add more ResNet blocks per level, or add attention at more resolutions. But none of these is a clean, predictable scaling axis.

Increase channel counts? Compute grows as the square of channels. Add resolution levels? The model topology changes, requiring architectural redesign. Add attention everywhere? It becomes a different architecture. The ADM paper (Dhariwal & Nichol, 2021) found that more channels usually helps, more attention sometimes helps — but there was no clean formula predicting exactly how much.

The Architecture Monoculture Risk

There's also a less-obvious problem. When the entire community uses one backbone, all collective knowledge is tied to that architecture. Training tricks discovered for U-Nets don't transfer to language models. Insights from ViT pretraining don't apply. The community has to rediscover everything from scratch within the diffusion U-Net world.

Chapter 2: Patchify the Latent

DiT doesn't operate on raw pixels. It builds on the Latent Diffusion framework — the same one behind Stable Diffusion. First, a pretrained VAE compresses the image into a compact latent space. Then diffusion happens entirely in that smaller space.

Step 1: VAE Compression

A Variational Autoencoder (frozen during DiT training) maps a 256×256×3 RGB image to a 32×32×4 latent. That's an 8× spatial downsampling. The forward diffusion process adds Gaussian noise to this latent:

DiT's job is to predict the noise ε given z_t and the timestep t. After denoising, the clean latent z₀ is decoded back to pixels by the VAE decoder.

Step 2: Patchify

Now comes the key idea borrowed from Vision Transformers (ViT). The 32×32×4 latent is divided into non-overlapping patches of size p×p. Each patch is flattened and linearly projected to a d-dimensional embedding. This gives a sequence of T tokens that the transformer can process.

Let's work through the exact numbers. With patch size p=2:

• Number of tokens: T = (32/2)² = 256 tokens
• Each patch covers: 2×2 spatial positions × 4 channels = 16 values
• Each patch is linearly projected: 16 → d = 1152 (for DiT-XL)
• Result: a sequence of shape [256, 1152]

Chapter 3: The DiT Block

The core building block of DiT is a modified transformer layer. It looks almost identical to a standard ViT block — multi-head self-attention followed by a feed-forward network, with layer normalization — except for one critical change: how conditioning information enters.

In a standard transformer, the layer norm has fixed, learned scale (γ) and shift (β) parameters. In DiT, these parameters are computed dynamically from the conditioning vector (timestep + class label) at every forward pass. This is called adaptive layer normalization (adaLN).

The Block Equations

Given input x (the token sequence) and conditioning vector c:

Where adaLN(γ, β, x) = γ · LayerNorm(x) + β. The parameters γ₁, β₁, α₁, γ₂, β₂, α₂ are all regressed from the conditioning vector c via a small MLP. Each block gets its own output projection from the shared MLP.

What's Removed vs Standard ViT

DiT is intentionally minimal. Compared to a standard ViT block, it removes nothing structural — it only modifies the layer norm to be adaptive. There are no:

• Cross-attention layers between blocks (no extra Q/K/V projections per block)
• Skip connections between early and late blocks (unlike U-ViT)
• Convolutional layers of any kind (unlike hybrid architectures)
• Multi-scale processing at different resolutions

All 28 blocks (in DiT-XL) process the same 256-token sequence at the same 1152-dimensional hidden size. The same-resolution, same-width design is what gives DiT its clean scaling behavior.

Chapter 4: adaLN-Zero Deep Dive

The paper ablates four different ways to inject conditioning into transformer blocks. The winner — adaLN-Zero — achieves roughly half the FID of the simplest approach while adding negligible compute. Understanding why it wins requires understanding the zero initialization trick.

The Four Conditioning Strategies

1. In-Context: Append the timestep and class embeddings as two extra tokens. Process them alongside image tokens. Overhead: ~0%. Simple, but the conditioning has limited influence — it's just two tokens in a sea of 256.

2. Cross-Attention: Add a cross-attention layer after each self-attention. Image tokens attend to conditioning tokens. Overhead: ~15% extra Gflops. Borrows from language model conditioning (like Stable Diffusion's text conditioning).

3. adaLN: Replace fixed layer norm parameters with ones computed from the conditioning vector. Overhead: ~0% (the MLP is tiny). Better than in-context and cross-attention, because conditioning influences every single token's normalization.

4. adaLN-Zero: Same as adaLN, plus: regress additional scaling gates α₁, α₂ that multiply the residual branches. Initialize these gates to zero. Overhead: ~0%. Best by a large margin.

Why Zero Initialization?

Think about what happens at the start of training. Random weights produce random outputs. If the MHSA output is random noise, adding it back via the residual connection corrupts x. This makes the first few training steps chaotic — gradients are large and unstable.

With α initialized to zero: output = x + 0 × MHSA(x) = x. The block is exactly the identity function. Gradients flow cleanly through the residual path from the very first step. The network has a stable foundation to learn from, and α grows gradually as training teaches the block what to contribute.

Worked Numerical Example

Let's trace an example token x = [1.2, -0.3, 0.8] through a DiT block at initialization:

After a few training steps, α₁ and α₂ become small non-zero values (say 0.01, 0.02). The block begins to learn. By 100K steps, α values reach around 0.1-0.3, and the blocks are contributing meaningfully.

Chapter 5: Conditioning

Every DiT block needs to know two things: how noisy is this latent? (the timestep t) and what class should this image be? (the class label c). These aren't fed as pixel values — they go through a carefully designed embedding pipeline before influencing each block's normalization parameters.

Timestep Embedding

The timestep t is a single integer (say, 523 out of 1000). To make it useful for a neural network, DiT embeds it using sinusoidal frequency embeddings — the same technique from the original Transformer paper. Different frequencies capture different timescales:

This 256-dimensional sinusoidal embedding is then passed through a 2-layer MLP (with SiLU activation) to produce a d-dimensional vector: embed_t: int → R^d.

Class Label Embedding

The class label c (an integer 0–999 for ImageNet) is embedded via a standard learned embedding table: a matrix of shape [1000, d]. Each class gets its own d-dimensional vector, learned from scratch during training. embed_c: int → R^d.

For unconditional generation (needed for classifier-free guidance), there's one extra learned embedding: the null embedding ∅, used when the class label is dropped during training (10% of the time).

Combining Timestep and Class

The two embeddings are simply added:

This combined conditioning vector c (shape [d]) is shared across all 28 DiT blocks. Each block has its own linear projection layer that maps c to 6 scalars: γ₁, β₁, α₁, γ₂, β₂, α₂. This is a tiny linear layer — 6 outputs for each block — and it's all the architectural machinery needed to make conditioning work.

Parameter Count for Conditioning

Less than 1% of DiT-XL's parameters are used for conditioning. The overwhelming majority (99%+) live in the 28 transformer blocks. This is the point: conditioning is an afterthought in terms of compute. The transformer does the heavy lifting.

Chapter 6: The Full Pipeline

Now we can trace a complete forward pass through DiT — from a noisy latent z_t all the way to a predicted noise map. Let's walk through every operation with exact tensor shapes for DiT-XL/2 (patch size p=2, hidden dim d=1152, 28 blocks).

The Decode Step in Detail

The final linear layer maps each 1152-dim token to a 32-dim vector. Why 32? Each token corresponds to a 2×2 patch with 4 latent channels. The output predicts both noise and variance for that patch: 2×2×4 channels for ε + 2×2×4 channels for Σ = 2×2×8 = 32 values per token.

Unpatchifying is the exact reverse of patchifying: rearrange the [256, 32] output back to [32, 32, 8] spatial format, then split the last dimension into ε̂ [32, 32, 4] and Σ̂ [32, 32, 4].

python — complete DiT forward pass (simplified)
def forward(self, z_t, t, c):
    # z_t: [B, 4, 32, 32]  t: [B]  c: [B] (class labels)
    B = z_t.shape[0]

    # 1. Patchify: [B, 4, 32, 32] → [B, 256, 1152]
    x = self.patch_embed(z_t)          # linear: (p*p*4) → d
    x = x + self.pos_embed             # add fixed sine-cosine embed

    # 2. Conditioning vector: [B, d]
    cond = self.t_embed(t) + self.c_embed(c)

    # 3. 28 DiT blocks
    for block in self.blocks:
        x = block(x, cond)             # adaLN-Zero inside

    # 4. Final norm + linear decode
    x = self.final_norm(x, cond)       # one more adaLN
    x = self.final_linear(x)           # [B, 256, 32]

    # 5. Unpatchify: [B, 256, 32] → [B, 8, 32, 32]
    x = self.unpatchify(x)

    # 6. Split noise + variance
    eps_hat, sigma_hat = x.chunk(2, dim=1)
    return eps_hat, sigma_hat          # each [B, 4, 32, 32]

Chapter 7: Scaling Laws

The most important result in the DiT paper isn't the final FID score. It's what happens when you plot FID against Gflops across 12 different model variants.

Peebles and Xie train every combination of 4 model sizes (S, B, L, XL) × 3 patch sizes (2, 4, 8), all for 400K training steps on ImageNet 256×256. Then they ask: what predicts quality?

Finding 1: Gflops → FID, with correlation -0.93

Across all 12 models, there is a -0.93 correlation between model Gflops and FID-50K. On a log-log plot, the 12 points fall almost on a straight line. This is remarkable. It means you can predict a DiT model's image quality from its compute cost alone, with no other information needed.

This is the transformer's greatest strength: predictable scaling. For U-Nets, no equivalent relationship exists. For transformers, it's practically a law.

Finding 2: Gflops > Parameters

Here's a subtle but crucial point. When you decrease patch size from p=4 to p=2, you quadruple T (tokens) and thus quadruple the Gflops. But parameter count barely changes — the transformer weights are identical, you're just processing more tokens per forward pass.

Yet FID improves substantially. This means compute (Gflops) per image, not total parameter count, is the true driver of quality. DiT-S/2 (33M params, 5.8 Gflops) achieves better FID than DiT-L/8 (458M params, 5.3 Gflops) — despite having 14× fewer parameters — because it uses more compute per forward pass.

Finding 3: Larger Models are More Efficient

When you plot FID against total training compute (Gflops × steps), larger models reach any FID threshold faster. A large model trained for fewer steps beats a small model trained longer. This matters for production: if you have a fixed compute budget, train the largest model you can afford, not a small model for longer.

Model Configurations

*DiT-XL/2 achieves FID 9.62 without CFG, FID 2.27 with CFG (guidance scale 1.50) after 7M training steps.

Chapter 8: Showcase — The DiT Pipeline in Motion

Here's the payoff. Watch DiT denoise a latent from pure noise to structured signal, step by step. Choose a model size, a class label, and a guidance scale — then step through the denoising process and see how the image emerges.

The visualization shows the 16×16 patch grid (for p=2, a 32×32 latent becomes 256 patches arranged in a 16×16 grid). Color encodes signal-to-noise ratio: cool blue = mostly noise, warm orange = mostly signal. As t decreases from 250→0, patches progressively "lock in" their values.

Chapter 9: Connections

DiT (2023) didn't just introduce a new architecture. It changed the entire trajectory of visual generation. Here's how the ideas evolved — and where they landed.

The DiT Family Tree

What Each Successor Changed

Key Insight: What DiT Proved Was Unnecessary

The U-Net's core inductive bias was: you need to process images at multiple spatial scales simultaneously, with skip connections to preserve detail. DiT proved this is false. Self-attention over a flat patch sequence is sufficient — it implicitly learns multi-scale reasoning without the architectural scaffolding.

This matters because it means the transformer's core strength — attending to arbitrary pairs of tokens — is all you need. There's no special sauce in the U-Net that couldn't be learned from data.

Related Lessons

• → Diffusion Models — The foundation DiT builds on: DDPM, forward/reverse process, latent diffusion, CFG
• → Transformers — Self-attention, positional embeddings, the architecture DiT borrows from
• → VAE & VQ-VAE — The compression step DiT depends on: how images are mapped to latent space
• → Flow Matching — The training objective that superseded DDPM in FLUX and SD3

"The transformer didn't just replace the U-Net — it opened diffusion models to the same scaling playbook that transformed language modeling. We no longer have to ask 'will making it bigger help?' The answer is yes, quantifiably, predictably, always."
— Inspired by Peebles & Xie, 2023