27  Diffusion Transformers

TipTL;DR

Diffusion models generate by learning to reverse a noising process: denoising step by step from pure noise back to data. The Diffusion Transformer (DiT) (Peebles & Xie, 2023) replaces the convolutional U-Net denoiser with a plain Transformer block operating on image patches, unifying the generative-image backbone with the same architecture that powers every LLM in this book.

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:

  1. 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.
  2. N Transformer blocks: standard self-attention + FFN, residual connections, layer norm (The Transformer), process the token sequence.
  3. 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.
  4. 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)\).
Noisy latent image x_t Patchify → tokens DiT block × N Self-attention (adaLN-modulated) Feed-forward (adaLN-modulated) + residual around each sublayer (same block as the LLM Transformer) Timestep t Class / text condition MLP → scale, shift (per block) modulates every block Unpatchify → Predicted noise ε_θ(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.
TipGoing deeper

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.
NoteKey takeaway

Diffusion Transformers are not a new generative idea. They are the DDPM denoising objective with its backbone swapped from convolution to the Transformer block already built in this book. The result is the same architecture generating text, images, and now video, scaling by the same recipe of more data and more compute.

→ Next: Embodied AI & VLA