Generative models II: Diffusion and flow matching

The idea

Take an image, add a little Gaussian noise, repeat $T$ times until nothing but noise is left. That forward process is fixed and needs no learning. The model learns the reverse process: given a noisy image and the step index, predict a slightly less noisy one. Chain the reverse steps from pure noise and you generate. Equivalently, the model learns the score $\nabla_{\mathbf{x}}\log p_t(\mathbf{x})$ — the direction toward higher data density — at every noise level.

The forward process, in closed form

With variances $\beta_1,\dots,\beta_T$, each step is $q(\mathbf{x}_t\mid\mathbf{x}_{t-1}) = \mathcal{N}\big(\sqrt{1-\beta_t}\,\mathbf{x}_{t-1},\;\beta_t\mathbf{I}\big)$. Define $\alpha_t = 1-\beta_t$ and $\bar\alpha_t = \prod_{s\le t}\alpha_s$. Then

Math check (two steps). $\mathbf{x}_t = \sqrt{\alpha_t}\,\mathbf{x}_{t-1} + \sqrt{1-\alpha_t}\,\boldsymbol{\epsilon}_{t-1}$ and $\mathbf{x}_{t-1} = \sqrt{\alpha_{t-1}}\,\mathbf{x}_{t-2} + \sqrt{1-\alpha_{t-1}}\,\boldsymbol{\epsilon}_{t-2}$. Substituting: $\mathbf{x}_t = \sqrt{\alpha_t\alpha_{t-1}}\,\mathbf{x}_{t-2} + \sqrt{\alpha_t(1-\alpha_{t-1})}\,\boldsymbol{\epsilon}_{t-2} + \sqrt{1-\alpha_t}\,\boldsymbol{\epsilon}_{t-1}$. The two independent Gaussian noises merge into one with variance $\alpha_t(1-\alpha_{t-1}) + (1-\alpha_t) = 1-\alpha_t\alpha_{t-1}$ (1.6 Probability I: random variables and distributions). Induction gives the formula.

A schedule sets the $\beta_t$: DDPM used linear from $10^{-4}$ to $0.02$ over $T=1000$; cosine schedules keep $\bar\alpha_t$ from collapsing too early. Because $\bar\alpha_T\approx0$, $\mathbf{x}_T$ is indistinguishable from $\mathcal{N}(\mathbf{0},\mathbf{I})$ — exactly the distribution we can sample at generation time. The demo below lets you drag $t$ and watch a dataset dissolve.

From the ELBO to a mean-squared error

The reverse model $p_\theta(\mathbf{x}_{t-1}\mid\mathbf{x}_t) = \mathcal{N}(\boldsymbol{\mu}_\theta(\mathbf{x}_t,t),\sigma_t^2\mathbf{I})$ is trained with an ELBO (8.8 Generative models I: VAEs and GANs) that decomposes into KL terms between Gaussians at each step. The target in each term is the true posterior $q(\mathbf{x}_{t-1}\mid\mathbf{x}_t,\mathbf{x}_0)$, which is Gaussian with mean

so each KL is a weighted squared distance between $\tilde{\boldsymbol{\mu}}_t$ and $\boldsymbol{\mu}_\theta$. Since $\mathbf{x}_0 = (\mathbf{x}_t - \sqrt{1-\bar\alpha_t}\,\boldsymbol{\epsilon})/\sqrt{\bar\alpha_t}$, predicting the mean is equivalent to predicting the noise: parametrize

Dropping the per-step weights (which empirically improves samples) leaves the celebrated simple loss:

Training is therefore: pick an image, a random step and random noise; noise the image in one shot; ask the network which noise was added; take an MSE gradient step. Alternatives predict $\mathbf{x}_0$ directly or the velocity $\mathbf{v} = \sqrt{\bar\alpha_t}\,\boldsymbol{\epsilon} - \sqrt{1-\bar\alpha_t}\,\mathbf{x}_0$, which is better behaved at both ends of the schedule.

def ddpm_loss(model, x0, alpha_bar):            # alpha_bar: tensor of length T
    t = torch.randint(0, len(alpha_bar), (x0.shape[0],), device=x0.device)
    ab = alpha_bar[t].view(-1, 1, 1, 1)
    eps = torch.randn_like(x0)
    x_t = ab.sqrt() * x0 + (1 - ab).sqrt() * eps
    return F.mse_loss(model(x_t, t), eps)

Sampling: DDPM and DDIM

DDPM (ancestral) sampling runs the learned reverse chain from $\mathbf{x}_T\sim\mathcal{N}(\mathbf{0},\mathbf{I})$:

with $\sigma_t^2 = \beta_t$ or the posterior variance $\tilde\beta_t = \frac{1-\bar\alpha_{t-1}}{1-\bar\alpha_t}\beta_t$. A thousand network evaluations per image is slow. DDIM observed that the same trained network defines a family of non-Markovian processes; the deterministic member ($\eta=0$) first estimates the clean image, $\hat{\mathbf{x}}_0 = (\mathbf{x}_t - \sqrt{1-\bar\alpha_t}\,\boldsymbol{\epsilon}_\theta)/\sqrt{\bar\alpha_t}$, then jumps:

Because it is deterministic, you can skip steps (20–50 instead of 1000) and the same noise always gives the same image — which enables interpolation and editing. Higher-order ODE solvers (DPM-Solver and relatives) push quality at 10–20 steps.

The score view

From the forward process, $\nabla_{\mathbf{x}_t}\log q(\mathbf{x}_t\mid\mathbf{x}_0) = -\boldsymbol{\epsilon}/\sqrt{1-\bar\alpha_t}$, so the noise predictor is a scaled score estimate: $\mathbf{s}_\theta(\mathbf{x}_t,t) \approx -\boldsymbol{\epsilon}_\theta/\sqrt{1-\bar\alpha_t}$. Song et al. unified diffusion as a stochastic differential equation; its probability-flow ODE has the same marginals and is the continuous version of DDIM's deterministic sampler. This view makes "noise levels" continuous and connects directly to flow matching below.

Conditioning and classifier-free guidance

To generate from text, feed a conditioning signal $\mathbf{c}$ (text-encoder embeddings) into the network — via cross-attention in a U-Net or as extra tokens in a transformer (8.10 Multimodal models). Classifier-free guidance (CFG) trains one network both conditionally and, by randomly dropping $\mathbf{c}$ (about 10% of the time), unconditionally. At sampling time, extrapolate:

$w=1$ is plain conditional sampling; $w\approx 5$–$8$ sharpens adherence to the prompt at the cost of diversity (in the score view it samples from $p(\mathbf{x})\,p(\mathbf{c}\mid\mathbf{x})^{w}$). Two network evaluations per step are the price.

Latent diffusion and backbones

Running diffusion on raw pixels at high resolution is expensive. Latent diffusion (Stable Diffusion) first trains a VAE (8.8 Generative models I: VAEs and GANs) that compresses a $512\times512\times3$ image into a $64\times64\times4$ latent — 48× fewer numbers — then diffuses in latent space and decodes at the end. The denoiser was a U-Net (5.6 Segmentation) with cross-attention to the text encoder; since 2023 Diffusion Transformers (DiT) patchify the latent and apply a transformer (7.4 Inside the block: residuals, LayerNorm, FFN and embeddings), scaling more predictably with compute. Modern image and video systems use transformer backbones with joint text–image attention.

Flow matching: straight lines instead of random walks

Instead of a noising Markov chain, define a straight interpolation between a data point $\mathbf{x}_0$ and noise $\boldsymbol{\epsilon}$ with $t\in[0,1]$:

The network learns a velocity field; sampling integrates the ODE $d\mathbf{x}/dt = \mathbf{v}_\theta$ from noise ($t=1$) to data ($t=0$) with a few Euler steps. Straight paths mean fewer steps and simpler schedules; the objective is equivalent to diffusion with a particular schedule and the $\mathbf{v}$-parametrization, so everything above (CFG, latents, DiT) carries over. Since 2024 most new image and video models — and many audio models — train this way (rectified flow). Distillation — consistency models, progressive and adversarial distillation — then compresses dozens of steps into one to four for interactive use.

Where this shows up later

• Text conditioning, CLIP/T5 encoders and multimodal generation — 8.10 Multimodal models.
• The backbone is a U-Net or a transformer — 5.6 Segmentation, 7.4 Inside the block: residuals, LayerNorm, FFN and embeddings; the compressor is a VAE — 8.8 Generative models I: VAEs and GANs.
• Safety questions of generative media (deepfakes, provenance) — 8.12 Safety, risks and responsible AI.