Summary: DiT (Diffusion Transformer, Peebles & Xie, 2023) replaces the U-Net backbone in diffusion models with a Transformer architecture. This enables better scaling, native conditioning via cross-attention, and SOTA results on ImageNet generation — forming the backbone of Diffusion Policy and modern video diffusion.
Motivation: Why Transformer for Diffusion?
| U-Net Limitation | DiT Solution |
|---|---|
| Fixed resolution (strided convs) | Patch-based, resolution-agnostic |
| Awkward conditioning (FiLM, concat) | Native cross-attention conditioning |
| Poor scaling beyond 1B params | Scales like Transformers (GPT-style) |
| Inductive bias for images only | General sequence modeling |
Architecture (Peebles & Xie, 2023)
Input latents (H×W×C) → Patchify → Sequence of patches
↓
Positional Embedding + Timestep Embedding + Class/Condition Embedding
↓
Transformer Blocks (DiT Blocks) × N
↓
Unpatchify → Output (same shape as input)
DiT Block: Standard Transformer block with adaLN-Zero (adaptive layer norm with zero-initialized scale/shift):
adaLN-Zero(x, c) = x + gamma(c) · Attn(LN(x)) + beta(c) · MLP(LN(x))
where (gamma, beta) are predicted from condition c and initialized to zero (residual starts as identity).
Conditioning (Native Cross-Attention)
| Condition Type | Method |
|---|---|
| Class labels | Learned embedding → adaLN |
| Text (CLIP/T5) | Cross-attention in DiT blocks |
| Images | Patchified → cross-attention |
| Action/Proprioception | Embedding → adaLN or cross-attn |
Diffusion Policy uses: Observation (image + proprio) → encoder → cross-attention in DiT.
Scaling Laws (DiT Paper)
| Param Count | FID (ImageNet 256×256, class-cond) |
|---|---|
| 33M | 8.95 |
| 131M | 4.51 |
| 459M | 2.78 |
| 675M | 2.27 (SOTA at the time) |
| 3.1B | 2.17 |
Key finding: DiT scales consistently with compute — no architectural changes needed.
DiT Variants
| Variant | Innovation | Use Case |
|---|---|---|
| DiT-XL/2 | 1152 dim, 28 layers, 16 heads | ImageNet SOTA |
| DiT-S/2 | 384 dim, 12 layers, 6 heads | Fast prototyping |
| Video DiT | Spatiotemporal patches + 3D pos emb | Sora, Gen-2, Stable Video |
| DiT for Robotics | Cross-attn to obs, action chunking | Diffusion Policy |
In Robotics / World Models
| Application | DiT Role |
|---|---|
| Diffusion Policy | ~100M params, obs-conditioned action diffusion |
| Video World Models (Gen-2, Sora) | Spatiotemporal DiT for video generation |
| DreamerV3 + Diffusion Decoder | DiT decoder for sharp reconstructions |
| UniSim | Trajectory diffusion with DiT backbone |
Related Concepts
- diffusion-model — Foundation
- diffusion-policy — Robotics application
- transformer-architecture — Backbone architecture
- latent-diffusion — DiT in latent space
- consistency-models — One-step distillation of DiT
- world-model — Video DiT for dynamics
Sources
- source-arxiv-2403.12022 — Diffusion Policy (DiT backbone)
- Peebles & Xie, "Scalable Diffusion Models with Transformers" (ICCV 2023)
- world-models-diffusion-ai-robotics — Decision guide