27 Diffusion Transformers
Depends on: The Transformer · Generative Models
27.1 Why this matters
Generative Models traced why diffusion won the contest against VAEs and GANs: a simple, stable denoising objective at the cost of multi-step sampling. But the original diffusion models (DDPM (Ho et al., 2020)) used a convolutional U-Net as the denoising network, the same architecture family as CNN Vision Lineage. That leaves an obvious question: does the Transformer, which displaced recurrence for sequences, also displace convolution for image generation?
DiT (Peebles & Xie, 2023) answers yes. Swapping the U-Net for a Transformer block operating on image patches scales the same way LLMs scale: more compute and more parameters reliably produce better samples. It also means one architecture, not two, underlies both text and image generation at the frontier.
27.2 The mechanism
27.2.1 The forward process: adding noise
Generative Models introduced the forward noising step. To make this note self-contained: fix a total number of steps \(T\) (e.g. \(T=1000\)) and a noise schedule \(\beta_1, \dots, \beta_T\), a sequence of small numbers controlling how much noise is added at each step. Define \(\alpha_t = 1-\beta_t\) and \(\bar\alpha_t = \prod_{s=1}^t \alpha_s\) (the cumulative fraction of signal remaining after \(t\) steps). The forward process lets you jump directly to step \(t\) in one equation:
\[ \mathbf{x}_t = \sqrt{\bar\alpha_t}\,\mathbf{x}_0 + \sqrt{1-\bar\alpha_t}\,\boldsymbol\epsilon, \qquad \boldsymbol\epsilon \sim \mathcal{N}(0, I). \]
As \(t \to T\), \(\bar\alpha_t \to 0\) and \(\mathbf{x}_t\) becomes indistinguishable from pure Gaussian noise. This process is fixed (no learned parameters); it exists purely to generate training pairs \((\mathbf{x}_t, \boldsymbol\epsilon)\).
27.2.2 The network’s job: predicting the noise
The network \(\epsilon_\theta(\mathbf{x}_t, t)\) takes the noisy image and the timestep and predicts the noise \(\boldsymbol\epsilon\) that was added:
\[ \mathcal{L} = \mathbb{E}_{t, \mathbf{x}_0, \boldsymbol\epsilon} \big[\lVert \boldsymbol\epsilon - \epsilon_\theta(\mathbf{x}_t, t) \rVert^2 \big]. \]
This connects to score matching: the noise prediction \(\epsilon_\theta(\mathbf{x}_t, t)\) is a rescaled estimate of the score, \(\nabla_{\mathbf{x}_t} \log p(\mathbf{x}_t)\), the direction in which the data density increases fastest at that noise level. Predicting the noise and predicting the score are mathematically equivalent up to a known scaling factor; DDPM’s contribution was showing the noise-prediction parameterization trains more stably in practice.
27.2.3 Generating: running the process backward
Training only fits \(\epsilon_\theta\); generation runs the chain in reverse. Start from pure noise \(\mathbf{x}_T \sim \mathcal{N}(0, I)\) and, for \(t = T, \dots, 1\), apply the learned denoising step:
\[ \mathbf{x}_{t-1} = \frac{1}{\sqrt{\alpha_t}} \left(\mathbf{x}_t - \frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}}\,\epsilon_\theta(\mathbf{x}_t, t)\right) + \sigma_t \mathbf{z}, \qquad \mathbf{z} \sim \mathcal{N}(0, I). \]
Each step removes the amount of noise the network predicts was added at that timestep, plus a small amount of fresh noise \(\sigma_t \mathbf{z}\) to keep the process stochastic (this is what makes repeated generations from the same starting noise still vary slightly). After \(T\) such steps, \(\mathbf{x}_0\) is a sample from the learned data distribution. This multi-step loop, not a single forward pass, is diffusion’s main cost.
27.2.4 The Diffusion Transformer block (DiT)
Nothing above constrains \(\epsilon_\theta\) to be a convolutional U-Net; it is any function mapping a noisy input and a timestep to a predicted noise map of the same shape. DiT’s contribution is showing a Transformer does this job at least as well, and scales better. Concretely:
- Patchify: the (usually latent-space, VAE-encoded) noisy image is cut into a grid of patches and flattened into a token sequence, exactly like ViT.
- N Transformer blocks: standard self-attention + FFN, residual connections, layer norm (The Transformer), process the token sequence.
- Conditioning via adaLN: the timestep \(t\) and a class/text condition are each embedded, summed, and passed through a small MLP that outputs a per-block scale and shift. Every block applies this scale/shift after its layer norm, so the network is modulated by “how noisy” and “what to generate” without adding extra tokens to the sequence.
- Unpatchify: a final linear layer maps each output token back to a pixel patch, and the patches are reassembled into the predicted noise map \(\epsilon_\theta(\mathbf{x}_t, t)\).
27.2.5 From images to video: Sora and beyond
Extending DiT from single images to video is architecturally almost free: patchify space and time into spatio-temporal tokens, and the same attention mechanism that mixes information across image patches now mixes across frames too. OpenAI’s Sora (Brooks et al., 2024) applies exactly this recipe at scale, a diffusion Transformer over spacetime patches, to generate minutes-long coherent video, describing itself explicitly as a step toward “world simulators”: models that must implicitly learn physics, object permanence, and camera geometry to denoise video convincingly.
27.2.6 What to observe
- The denoising objective from Generative Models is unchanged; only the network architecture computing \(\epsilon_\theta\) changed from U-Net to Transformer.
- Generation is a loop of \(T\) reverse steps, not one forward pass. This is the direct source of diffusion’s inference cost.
- Conditioning (timestep, class/text) enters through modulation (adaLN), not through the attention mechanism directly, a different pattern from the encoder-decoder cross-attention in The Transformer.
- Video is “just” more patches along a time axis; no new mechanism, only a larger token sequence, which is why scaling laws transfer directly from image to video generation.
27.3 Application & impact
| Component | Role | Lives on as |
|---|---|---|
| U-Net denoiser | Original DDPM backbone | Displaced by DiT in frontier models |
| DiT block | Transformer as denoiser | Stable Diffusion 3, FLUX |
| adaLN conditioning | Inject timestep + class/text | Standard conditioning method in diffusion Transformers |
| Spatio-temporal patchify | Image → video, same mechanism | Sora and video-diffusion successors |
- One architecture, two modalities: DiT means the same block (attention + FFN + residual + norm) that reads and writes text is now also the state-of-the-art image/video generator, the Transformer’s second independent proof of scaling, after language.
- Latent-space diffusion is standard practice: production systems diffuse in the VAE latent space (small, compressed), not raw pixels, directly reusing the VAE from Generative Models as a compression front-end; DiT then operates on that compressed representation.
→ Next: Embodied AI & VLA