Diffusion and flow matching in 2D

A diffusion model creates new data (images, sounds, molecules) from pure noise. It learns to undo a fixed process that destroys data by adding noise step by step. On this page you see the whole idea on points in the plane:

• Forward process. Add Gaussian noise to the data until only noise is left.
• Score. At every noise level, the score $\nabla_x \log p_t(x)$ is a vector field that points toward regions of higher probability, that is, back toward the data.
• Reverse process. Start from noise and follow the score backward in time. You can do this with random steps (the reverse SDE, or DDPM) or with deterministic steps (the probability-flow ODE, or DDIM). Flow matching is a close relative that trains on straight lines from data to noise.

A real model learns the score with a neural network. Here the data distribution is a mixture of small Gaussians around the data points, so the noised density and its score have a closed form. The score in every figure is exact, not learned. This lets you see the sampling methods themselves, without training errors, and see clearly why few sampling steps give blurry results.

1. The playground #

Interactive figure: loading…

2. Paths through time #

Interactive figure: loading…

3. Few steps or many #

Interactive figure: loading…

4. Try this #

1. When is the shape gone? In the playground, choose two moons and Forward, and move $t$ slowly from 0 to 1. Question: at which $t$ can you no longer see two moons, and where do the arrows point at $t = 1$? What to see: the schedule destroys the shape fast. At $t = 0.1$ the signal scale is $a_t = 0.95$ and the moons are clear; at $t = 0.3$, $a_t = 0.63$ and $\sigma_t = 0.78$, and the moons have merged into one cloud. At $t = 1$, $p_1 \approx \mathcal N(0, I)$, so the score is about $-x$: every arrow points to the origin.
2. One step gives the average. Choose Reverse, the sampler DDIM and $N = 1$, then $N = 2, 4, 8, 16$. Question: where do the samples land with one step? What to see: with $N = 1$ all samples land on one point, the mean of the data (here the origin). With 2 and 4 steps the samples form a blurry cloud between the moons. With 16 steps they lie on the moons (the badge says close to the data). In Figure 2, choose 80 / 20, DDIM and $N = 1$: every path ends at $x \approx -0.72$, which is the mean $0.8 \cdot (-1.2) + 0.2 \cdot 1.2$.
3. Same start, same end? Set DDIM with $N = 40$ and press Sample again twice. Then do the same with DDPM. Question: which sampler is a fixed map from noise to data? What to see: DDIM draws exactly the same paths and samples again; DDPM draws new paths every time from the same starting noise, because it adds fresh noise at every step. In Figure 2 the ODE paths never cross; the SDE paths cross often.
4. Too much noise. Choose Euler–Maruyama with $N = 2$, then $N = 10$, then $N = 40$, and compare with DDPM at the same $N$. Question: both methods sample the same reverse SDE; why do they fail in different ways with few steps? What to see: at $N = 2$ the Euler–Maruyama points fly out of the frame: the first step has $\beta(1)\Delta t = 10$, far too large for an explicit method. DDPM stays in the frame, because each step draws from an exact Gaussian; its error with few steps is blur. At $N = 10$ the Euler–Maruyama samples are a noisy band around the moons, and the DDPM samples are still too smooth. At $N = 40$ both are close to the data: the distance is near that of true samples. Figure 3 shows that Euler–Maruyama needs the most steps of the four samplers.
5. Straight lines in training, curves in sampling. Choose Flow matching, turn on training pairs and sample with $N = 40$. Question: do the paths follow the grey dashed lines from noise to data? What to see: no. Each training pair $(x_0, \varepsilon)$ is joined by a straight line, but these lines cross. The model learns the average velocity of all lines through a point, so the sampling paths bend and never cross. They are still straighter than the DDIM paths: compare both samplers at $N = 4$ in Figure 3.
6. The exact score copies the data. Choose draw and click five separate points. Set the kernel width $s$ to $0.03$ and sample with DDIM and $N = 64$. Question: does the model create new points? What to see: every sample lands on one of your five points. The exact score of a small data set reproduces the training data; this is memorization. Increase $s$ and the five points become five blobs. A trained network does not learn this exact score, and that is one reason why real diffusion models can create new data.

5. The math in short #

The forward process and DDPM #

DDPM (denoising diffusion probabilistic models, Ho et al. 2020) adds a little Gaussian noise in each of $T$ steps:

A sum of independent Gaussians is Gaussian, so we can jump to any step at once: $x_t = \sqrt{\bar\alpha_t}\, x_0 + \sqrt{1 - \bar\alpha_t}\, \varepsilon$ with $\bar\alpha_t = \prod_{s \le t} (1 - \beta_s)$ and $\varepsilon \sim \mathcal N(0, I)$. With many small steps this becomes the variance-preserving (VP) SDE of Song et al. (2021), with continuous time $t \in [0, 1]$:

Here $w$ is Brownian motion. The figures use $\beta(t) = 0.1 + 19.9\, t$, so $a_1 \approx 0.007$: at $t = 1$ almost no signal is left. "Variance preserving" means that data with unit variance keep unit variance. The signal-to-noise ratio is $\mathrm{SNR}(t) = a_t^2 / \sigma_t^2$.

DDPM trains a network $\hat\varepsilon_\theta(x_t, t)$ to predict the noise that was added, with the loss $\mathbb E\, \|\varepsilon - \hat\varepsilon_\theta(x_t, t)\|^2$. The best possible prediction is the conditional mean $\mathbb E[\varepsilon \mid x_t]$, and the next section shows that this is the score up to a factor.

The score and the denoiser #

The score of the noised density $p_t$ is $\nabla_x \log p_t(x)$. The denoiser $D(x) = \mathbb E[x_0 \mid x_t = x]$ is the best guess of the clean point (best in mean squared error). Tweedie's formula connects the two:

So "predict the noise", "predict the clean point" and "learn the score" are the same task.

The exact score of these figures. Each data set is a mixture $p_{\text{data}} = \sum_k \pi_k\, \mathcal N(\mu_k, s^2 I)$: one small Gaussian of width $s$ (the kernel width) around each data point $\mu_k$. Noising maps $\mathcal N(\mu_k, s^2 I)$ to $\mathcal N(a_t \mu_k, v_t I)$ with $v_t = a_t^2 s^2 + \sigma_t^2$. With the responsibilities $w_k(x) \propto \pi_k\, \mathcal N(x;\ a_t \mu_k,\ v_t I)$ and their weighted mean $m(x) = \sum_k w_k(x)\, \mu_k$:

Near the data (small $t$) one weight $w_k$ is close to 1, so the score points to the nearest data point. When the noise is large, all $w_k$ are about equal, $m(x)$ is the mean of the data, and the score points toward $a_t m$, the scaled mean.

The reverse SDE and the probability-flow ODE #

For a forward SDE $dx = f(x, t)\, dt + g(t)\, dw$, the process run backward in time is again an SDE (Anderson 1982):

where $\bar w$ is a Brownian motion in reverse time. Song et al. (2021) showed that the deterministic probability-flow ODE

has the same marginal densities $p_t$ for every $t$. Both start from $p_1 \approx \mathcal N(0, I)$ and end at $p_0 = p_{\text{data}}$. The SDE paths are random and cross each other; the ODE paths are smooth, never cross, and define a fixed, invertible map from noise to data. For the VP SDE, $f = -\tfrac12\beta x$ and $g = \sqrt\beta$.

The four samplers use the time grid $t_i = 1 - i/N$ and take $N$ steps, one score evaluation per step:

• Euler–Maruyama (reverse SDE): $x_{t - \Delta t} = x_t + \Delta t\, \beta(t)\big[\tfrac12 x_t + \nabla \log p_t(x_t)\big] + \sqrt{\beta(t)\,\Delta t}\; z$, with fresh $z \sim \mathcal N(0, I)$.
• DDPM (reverse SDE): compute $\hat x_0 = D(x_t)$, then draw $x_{t'}$ from the Gaussian $q(x_{t'} \mid x_t, x_0 = \hat x_0)$ of the forward process. The noise of each step then has the right size, even for large steps.
• DDIM (probability-flow ODE): with $\hat\varepsilon = (x_t - a_t \hat x_0)/\sigma_t$, set $x_{t'} = a_{t'}\, \hat x_0 + \sigma_{t'}\, \hat\varepsilon$. This is the Euler method for the ODE written in the variables $x/a_t$ and $\sigma_t / a_t$.
• Flow matching (its own ODE, next section): $x_{t - \Delta t} = x_t - \Delta t\; v(x_t, t)$.

Flow matching #

Flow matching (Lipman et al. 2023; also called rectified flow, Liu et al. 2023) uses the straight line between a data point and a noise point:

A network $v_\theta(x, t)$ is trained with $\mathbb E\, \|v_\theta(x_t, t) - (\varepsilon - x_0)\|^2$. The best solution is the average velocity of all pairs whose line passes through $x$:

Sampling solves $dx/dt = v(x, t)$ from $t = 1$ to $t = 0$. The figures use the convention $t = 0$ for data and $t = 1$ for noise, as in DDPM and in Stable Diffusion 3; Lipman et al. use the opposite direction.

Two facts are useful. First, an Euler step $x_{t'} = D(x_t) + \tfrac{t'}{t}\big(x_t - D(x_t)\big)$ moves on the straight line from $x_t$ toward the current guess $\hat x_0 = D(x_t)$, as the yellow lines in Figure 2 show. Second, this path is also a Gaussian path $x_t = a_t x_0 + \sigma_t \varepsilon$ (with $a_t = 1 - t$, $\sigma_t = t$). With the exact score, its ODE and the probability-flow ODE of the VP SDE send the same noise point to (almost exactly) the same data point; only the path in between differs. The flow-matching path bends less, so Euler steps make smaller errors.

Why few steps blur the result #

Every DDPM, DDIM and flow-matching step moves toward the guess $\hat x_0 = D(x_t) = \mathbb E[x_0 \mid x_t]$. This guess is an average over all clean points that could have produced $x_t$. When the noise is large, many very different clean points are possible, so their average lies between them: between the two moons, or in the middle of the ring.

• With one step from $t = 1$, DDIM, DDPM and flow matching all return $D(x_1) \approx \mathbb E[x_0]$: every sample is the mean of the data.
• With a few steps, each step asks the denoiser again at a lower noise level. The guess then depends more on the current point, and the samples separate into the parts of the data. Too few steps leave them between the parts: the result is too smooth, it is blurred. For images, the average of many possible images is a blurry image.
• The ODE samplers have no noise to correct the errors, but their paths are smooth, so they need fewer steps. Euler–Maruyama has the opposite problem: with large steps it adds too much noise and its drift overshoots, so the samples are too noisy or fly away.

This is why fast samplers follow straighter paths (flow matching, rectified flow), use higher-order ODE solvers, or train a model to jump along the ODE in one step (distillation and consistency models).

Distances between point clouds

The panels report the sliced Wasserstein-2 distance between the samples $\{y_i\}$ and true samples $\{x_i\}$ of the same size $n$. Project both sets onto a direction $u$; in one dimension the Wasserstein-2 distance pairs the sorted values: $W_2^2 = \frac1n \sum_i \big(u^\top x_{(i)} - u^\top y_{(i)}\big)^2$. The sliced distance averages $W_2^2$ over 24 evenly spaced directions and takes the square root. It sees both errors: samples that are too concentrated (blur) and samples that are too spread (noise). With a finite $n$ even two sets of true samples have a distance above 0; this is the grey band in Figure 3.

6. Learn more on this site #

• Generative Models: the lessons on DDPM, diffusion parameterizations, reverse SDE and PF-ODE samplers, flow matching, guidance and few-step generation (being written).
• Probability and Statistics: Gaussians, conditional expectation and Bayes' rule, the tools behind the denoiser $\mathbb E[x_0 \mid x_t]$ (being written).
• Calculus, lesson 10: integrals, expectations and the change of variables: how a map moves a density, the idea behind the probability-flow ODE.
• Calculus, lesson 11: Gaussian integrals: why a noised Gaussian mixture is again a Gaussian mixture.
• Calculus, lesson 12: score-function and pathwise gradients: a different "score". There it is the gradient of $\log p_\theta$ with respect to the parameters $\theta$; here it is the gradient with respect to the point $x$.
• Information Theory, lesson 3: the KL divergence: the divergence that maximum-likelihood training of generative models minimizes.
• Optimization, lesson 5: stochastic gradient descent: a noisy gradient step, the same kind of update as an Euler–Maruyama step of the reverse SDE.
• ML for Science: stochastic differential equations and generative models for weather and physics.