Score Matching, Diffusion, and Flow Matching

Dive into Deep Learning · §28.4

One regression, then one integral
score matching, diffusion, and flow matching.

Why learn a score?

Motivation

An energy model p_\theta = e^{-E_\theta}/Z_\theta needs the intractable Z_\theta at every step. The score sidesteps it:

\nabla_{\mathbf x}\log p_\theta = -\nabla_{\mathbf x} E_\theta, \qquad \nabla\log Z_\theta = 0.

Anything done with scores alone is normalizer-free.

image/svg+xml Matplotlib v3.10.8, https://matplotlib.org/

01

Learning the score

explicit → implicit → denoising

Score matching, made tractable

The objective

The Fisher divergence \tfrac12\mathbb E_p\|\mathbf s_\theta-\nabla\log p\|^2 still contains the unknown score. Hyvärinen integrates by parts:

J_{\mathrm{ESM}} = \mathbb E_p\bigl[\tfrac12\|\mathbf s_\theta\|^2 + \nabla\cdot\mathbf s_\theta\bigr] + C.

Tractable, but \nabla\cdot\mathbf s_\theta costs O(d) backward passes.

The regression lemma

The engine

\mathbb E\|\mathbf v(X)-Y\|^2 = \mathbb E\|\mathbf v(X)-\mathbf m(X)\|^2 + \text{const}, where \mathbf m(X)=\mathbb E[Y\mid X].

Proof. Insert \pm\mathbf m(X); the cross term vanishes by the tower rule. \blacksquare Least squares against a noisy target fits its conditional mean: used once for scores, once for velocities.

Denoising score matching & Tweedie

Denoising

Perturb \tilde{\mathbf x}=\mathbf x+\sigma\boldsymbol\epsilon; the conditional score is closed-form -\boldsymbol\epsilon/\sigma, and regressing on it (Vincent) recovers the marginal score. Rearranged, that is Tweedie:

\mathbb E[\mathbf x\mid\tilde{\mathbf x}] = \tilde{\mathbf x} + \sigma^2\,\nabla\log p_\sigma(\tilde{\mathbf x})

One step up the score lands exactly on the posterior mean.

Estimating the score and optimal denoising are the same function.

A score network in 1-D

Denoising

A tiny MLP trained by denoising score matching matches the analytic score, landing on the irreducible loss floor:

DSM loss 2.036 vs irreducible floor 2.038; max |s_theta - score| on [-4, 4]: 0.199

02

Score-based diffusion

forward noise, DDPM, Langevin, DDIM, guidance

One noise level → all of them

Forward process

Small \sigma approximates p but ignores empty space; large \sigma covers but blurs. The fix: a noise-conditional score \mathbf s_\theta(\mathbf x,t) trained along a forward SDE (VE or VP), then a reverse pass to generate:

A weighting \lambda(t) allocates effort across noise levels: \lambda = g^2 makes the loss a likelihood bound; DDPM’s 1-\bar\alpha_t trades that for sample quality.

Two clocks

Conventions

Diffusion runs data→noise and samples backward; flow matching runs noise→data and samples forward. To compare, substitute t\to 1-t:

DDPM is three propositions

DDPM

  1. The DDPM step is Euler–Maruyama on the VP-SDE (to O(\beta_t)).
  2. Exact marginal \mathbf x_t=\sqrt{\bar\alpha_t}\,\mathbf x_0+\sqrt{1-\bar\alpha_t}\,\boldsymbol\epsilon, \bar\alpha_t=\prod_s(1-\beta_s).
  3. The simple \|\boldsymbol\epsilon-\boldsymbol\epsilon_\theta\|^2 loss is denoising score matching, reweighted by \lambda(t)=1-\bar\alpha_t.
t =   10: Var(x_t) chain 1.000, formula 1.000, alpha_bar 0.9981
t =  100: Var(x_t) chain 0.994, formula 1.000, alpha_bar 0.8970
t = 1000: Var(x_t) chain 1.024, formula 1.000, alpha_bar 0.0000
chain vs one-shot at T: mean -0.027 vs +0.005, std 1.012 vs 1.000

Langevin: stationary but slow

Sampling

dX = \tfrac12\nabla\log p\,dt + dW has stationary density p (substitute \rho=p into Fokker–Planck → 0). But it mixes slowly across modes:

rng = np.random.default_rng(11)

def langevin(q, h, steps, rng):
    for _ in range(steps):
        q = q + 0.5 * h * mixture_score(q, 0.25) \
            + np.sqrt(h) * rng.standard_normal(q.shape)
    return q

warm = langevin(rng.normal(0.0, 3.0, 10000), 0.01, 2000, rng)
print(f'spread-out start: P(X > 0) = {(warm > 0).mean():.3f}, '
      f'E[X^2] = {(warm**2).mean():.2f} (truth 0.500, 4.25)')
cold = langevin(np.full(10000, -2.0), 0.01, 2000, rng)
print(f'one-mode start:   P(X > 0) = {(cold > 0).mean():.3f}  (slow mixing)')
spread-out start: P(X > 0) = 0.509, E[X^2] = 4.26 (truth 0.500, 4.25)
one-mode start:   P(X > 0) = 0.012  (slow mixing)

From one mode, almost no walker crosses (P(X>0)=0.012). Fix: anneal the noise, or use predictor–corrector.

DDIM: slide along your own curve

Sampling

Invert the marginal for \hat{\mathbf x}_0 = \bigl(\mathbf x_t - \sqrt{1-\bar\alpha_t}\,\boldsymbol\epsilon_\theta\bigr)/\sqrt{\bar\alpha_t}, then re-use the predicted noise instead of resampling:

\mathbf x_{t-1} = \sqrt{\bar\alpha_{t-1}}\,\hat{\mathbf x}_0 + \sqrt{1-\bar\alpha_{t-1}}\,\boldsymbol\epsilon_\theta.

Each sample slides deterministically along its own curve: strides can skip levels.

Ten strides for a thousand staggers

Sampling

On the closed-form mixture (no learning in the loop), ten strides land every sample in the same mode as the thousand-step reference; by fifty the terminal laws are statistically indistinguishable:

 10 strides vs 1000: mean |gap| 0.083, KS 0.080, mode fraction 0.495
 50 strides vs 1000: mean |gap| 0.015, KS 0.018, mode fraction 0.495
mode fraction at 1000 steps: 0.495; 5% KS threshold: 0.021

Same network; \eta controls reinjected noise. In the fine-step limit, deterministic DDIM is related to the probability-flow ODE; finite strides are numerical approximations, including for Gaussian marginals.

Guidance is Bayes on scores

Guidance

\nabla\log p_t(\mathbf x\mid y) = \nabla\log p_t(\mathbf x) + \nabla\log p_t(y\mid\mathbf x).

Classifier-free guidance trains one network with label dropout and extrapolates through the conditional: \tilde{\mathbf s} = (1-\gamma)\,\mathbf s_\varnothing + \gamma\,\mathbf s_y.

Measured on the closed-form mixture: \gamma=1 reproduces the exact conditional (mean 0.966 vs the analytic 0.970); at \gamma=3, 10 there is no more mass to move, so the mode distorts, drifting to 1.04, then 1.07, and narrowing.

For \gamma>1 the tilt p_t(\mathbf x)\,p_t(y\mid\mathbf x)^\gamma is the noised marginal of no clean distribution: a useful controlled distortion, not a consistent diffusion.

03

Flow matching

prescribe the path, regress the velocity, translate the targets

Probability paths and velocities

Flow matching

Prescribe a path (p_t) from noise to data; its velocity obeys the continuity equation. The intractable marginal velocity is again a posterior mean:

\mathbf u_t(\mathbf x) = \mathbb E\bigl[\mathbf u_t(\mathbf x\mid\mathbf z)\mid\mathbf x_t=\mathbf x\bigr].

Same disease as score matching, same cure.

The conditional flow-matching theorem

Flow matching

The tractable CFM loss (closed-form per-pair velocity) and the intractable FM loss have the same gradients.

Proof. Apply the regression lemma with target \mathbf u_t(\mathbf x\mid\mathbf z); its conditional mean is the marginal velocity. \blacksquare Identical structure to Vincent’s theorem.

Score, noise, velocity: one function

The dictionary

On a Gaussian path \mathbf x_t = \alpha_t\,\mathbf x_1 + \sigma_t\,\boldsymbol\epsilon, the marginal velocity and the marginal score determine each other:

\mathbf u_t(\mathbf x) = \frac{\dot\alpha_t}{\alpha_t}\,\mathbf x - \Bigl(\sigma_t\dot\sigma_t - \sigma_t^2\,\frac{\dot\alpha_t}{\alpha_t}\Bigr) \nabla\log p_t(\mathbf x)

Both are affine in the one posterior mean \hat{\mathbf x}_1 = \mathbb E[\mathbf x_1\mid\mathbf x_t] (Tweedie again):

max |u_posterior - u_dictionary| on [-4, 4]: 3.1e-15

Route one never mentions a score; route two never mentions a velocity.

One posterior mean, many targets

Parameterizations

Every target a practitioner meets is a t-dependent affine transformation of the score \mathbf s = \nabla\log p_t:

network predicts in terms of \mathbf s scaling
noise \hat{\boldsymbol\epsilon} -\sigma_t\,\mathbf s unit-scale near data (DDPM)
clean \hat{\mathbf x}_1 (\mathbf x + \sigma_t^2\,\mathbf s)/\alpha_t Tweedie; weak near noise
v-prediction \alpha_t\boldsymbol\epsilon - \sigma_t\mathbf x_1 affine in \mathbf s O(1) at both ends

The invariant clock is the log-SNR \lambda_t = \log(\alpha_t^2/\sigma_t^2): schedules covering the same \lambda range are re-clockings of the same model: EDM’s \sigma(t)=t is exactly such a re-clocking.

Rectified flow: straight paths

Flow matching

The simplest path is a straight line, \mathbf x_t=(1-t)\mathbf x_0+t\mathbf x_1, with constant target \mathbf x_1-\mathbf x_0. Conditional paths are straight; the marginal flow bends only where paths cross:

Gaussian → two moons

Flow matching

A small MLP trained by the rectified-flow loss, then Euler-integrated, sharpens the crescents as step count grows:

panels = [('data', moons[:2048])] + [
    (f'{K} step(s)', euler_sample(2048, K)) for K in (1, 2, 8, 32)]
fig, axes = d2l.plt.subplots(1, 5, figsize=(11, 2.4), sharex=True, sharey=True)
for ax, (title, s) in zip(axes, panels):
    ax.scatter(s[:, 0], s[:, 1], s=1)
    ax.set_title(title)
    ax.set_xlim(-2.5, 2.5), ax.set_ylim(-2.5, 2.5)

One reflow round, measured

Reflow

Integrate the trained ODE once, keep the couplings (\mathbf z, \hat{\mathbf x}_1(\mathbf z)), and retrain the same architecture on those pairs, which now almost never cross:

final reflow loss 0.003
 1 step(s): energy distance  CFM 0.582  ->  reflow 0.006
 2 step(s): energy distance  CFM 0.138  ->  reflow 0.005
32 step(s): energy distance  CFM 0.007  ->  reflow 0.005

One Euler step scores 0.016: within noise of the original’s 32-step quality, 40\times better than its one-step 0.676. The loss floor collapses too: the coupling’s posterior variance is gone.

04

Optimal transport and sampling

straightness, solver order, the unifying table

Straight = optimal

Benamou–Brenier

W_2^2(p_0,p_1) = \min_{(p_t,\mathbf v_t)}\int_0^1\!\!\int\|\mathbf v_t\|^2 p_t.

Any bridging flow costs at least W_2^2 (Jensen); the minimizer moves each particle in a straight line at constant speed.

Curvature = wasted kinetic energy = wasted solver steps. Reflow and OT couplings straighten the paths.

Steps buy quality; order sets the price

Sampling

Error falls with steps until model bias dominates; a higher-order solver buys orders of magnitude at equal cost (Heun at 20 steps beats Euler at 40):

steps_list = [1, 2, 4, 8, 16, 32, 64]
eds = [energy_distance(euler_sample(2048, K), held_out) for K in steps_list]
print('  '.join(f'{K}: {e:.3f}' for K, e in zip(steps_list, eds)))
d2l.plot(steps_list, eds, 'Euler steps', 'squared energy distance',
         xscale='log', yscale='log')

1: 0.582  2: 0.138  4: 0.027  8: 0.008  16: 0.006  32: 0.007  64: 0.007
K =  2 steps: endpoint error  Euler 0.7516 (2 NFE)   Heun 0.6659 (4 NFE)
K =  5 steps: endpoint error  Euler 0.1215 (5 NFE)   Heun 0.0808 (10 NFE)
K = 10 steps: endpoint error  Euler 0.0439 (10 NFE)   Heun 0.0282 (20 NFE)
K = 20 steps: endpoint error  Euler 0.0203 (20 NFE)   Heun 0.0085 (40 NFE)
K = 40 steps: endpoint error  Euler 0.0099 (40 NFE)   Heun 0.0023 (80 NFE)

One template, many names

Sampling

Every method here is the same template: a probability path, a closed-form conditional regression target, and a numerical integrator.

DDPM, score-SDE, PF-ODE, DDIM, and flow matching differ only in the path, the loss reweighting, and whether the sampler injects noise; the stochastic interpolants framework writes the whole family in one formalism.

Recap

Wrap-up

  • Score sidesteps Z; DSM (Vincent) regresses on -\boldsymbol\epsilon/\sigma; Tweedie = optimal denoising.
  • DDPM = VP-SDE discretized; \bar\alpha_t marginal; \boldsymbol\epsilon-loss = reweighted DSM.
  • Langevin mixes slowly; DDIM is deterministic; guidance is Bayes on scores.
  • Flow matching prescribes the path; CFM = FM by the same regression lemma.
  • Score, noise, \hat{\mathbf x}_1, velocity: one posterior mean in different parameterizations, on the log-SNR clock.
  • Straight = least energy (Benamou–Brenier); one reflow round makes 1 step match 32.
  • One thread: an intractable marginal average is a tractable conditional expectation, and regression finds it.