This blog post is about image generation with diffusion and flow matching. As the first step, to build intuition, I’ll use points on a 2D plane to represent a “image”, let you draw a shape to represent the manifold for valid images. It will use a simple 3-layer MLP that can train and inference right inside your browser.

Everything on this page is live. The widgets share one state: the shape you pick in section 1 is what the networks in sections 3 and 4 train on, and the models you train for diffusion are the ones the sampling widgets use. If you change the data after training, the page asks before throwing the models away. If you just want to play, there is a standalone playground with all the controls in one place.

We go in two passes. First the classic recipe, diffusion (DDPM and DDIM). Then flow matching, which keeps the same idea but replaces the noise schedule with a straight line, and ends with reflow, a trick that makes those lines straighter still.

1. Pick a distribution

A generative model learns to produce new samples from some unknown distribution $p_\text{data}$. All we get to see are samples. Here the distribution is a thin curve, so we can see at a glance whether a generated point lands on it.

Draw your own shape, or pick a preset. “Color” paints each point by its position, which is helpful later when we follow points through the process.

2. Add noise (the forward process)

The core trick is to define an easy direction first: slowly destroy the data by mixing in Gaussian noise. For a clean point $x_0$ and a standard Gaussian noise vector $\varepsilon$, the noisy point at time $t \in [0, 1]$ is

$$ x_t = \alpha(t)\,x_0 + \sigma(t)\,\varepsilon $$

where $\alpha(t)$ is the signal weight (1 at $t = 0$, 0 at $t = 1$) and $\sigma(t)$ is the noise weight (0 at $t = 0$, 1 at $t = 1$). Different choices of $\alpha$ and $\sigma$ give different methods. The usual diffusion choice is variance preserving: $\alpha^2 + \sigma^2 = 1$, so the total variance stays constant while signal turns into noise. Here $\alpha = \sqrt{1-t}$ and $\sigma = \sqrt{t}$.

Drag $t$ to watch the shape dissolve. Notice that every point keeps its own noise vector $\varepsilon$ as $t$ changes, so a point moves along a smooth path rather than jumping around.

At $t = 1$ nothing of the shape is left: we have a plain Gaussian blob. The goal of generation is to run this film backwards, starting from fresh Gaussian noise.

3. Learn to undo the noise

Running the film backwards needs one thing: at each noise level, which way is “towards the data”? Diffusion models learn this by denoising. Take a data point, pick a random $t$, add noise to get $x_t$, and train a network $\hat\varepsilon(x_t, t)$ to guess the noise $\varepsilon$ that was added. The loss is just a mean squared error:

$$ L = \mathbb{E}_{x_0,\,\varepsilon,\,t}\,\big\| \hat\varepsilon(x_t, t) - \varepsilon \big\|^2 $$

The same network handles every $t$; time is simply one more input. Concretely, the network is a plain MLP. The position is expanded into sines and cosines at a few frequencies (Fourier features, so a small network can draw sharp curves), $t$ is embedded similarly, and everything is concatenated into one input vector:

And the training loop is as short as it sounds:

for step in range(steps):
    x0  = data[randint(len(data), size=64)]      # a minibatch of clean points
    t   = tmax * rand(64) ** 2                   # random noise levels, a bit more weight on low noise
    eps = randn(64, 2)
    xt  = alpha(t) * x0 + sigma(t) * eps         # noisy points
    loss = ((net(xt, t) - eps) ** 2).mean()      # predict the noise
    loss.backward(); adam.step()

Press Train. Within a few hundred steps the picture below shows what the network has learned at the noise level $t$ you pick. Two views are available:

  • Score. The predicted noise is a score estimate: $\nabla \log p_t(x) \approx -\hat\varepsilon(x, t) / \sigma(t)$, the direction in which the noisy density $p_t$ increases. We draw the arrow $-\hat\varepsilon$ (the score times $\sigma$, which just rescales it).
  • Sampling direction. The direction the deterministic sampler (section 4, $\eta = 0$) actually moves points when going from $t$ to $t - dt$. It mixes the predicted noise with the predicted clean point $\hat x_0$.

At small $t$ the arrows hug the curve and point onto it; at large $t$ they point towards the middle of the data, since from far away the best guess is “somewhere in the cloud”.

4. Generate: walk the field backwards

Now we sample. Start with Gaussian noise at $t \approx 1$ and take $N$ steps down to $t = 0$. In each step the network predicts the noise, we use it to estimate the clean point $\hat x_0$, and then we re-noise that estimate to the next (lower) noise level. A single parameter $\eta$ decides how much fresh randomness goes into the re-noising:

  • $\eta = 1$ is DDPM: every step adds new noise. The paths are jittery, especially at high noise, but the process is a proper stochastic sampler.
  • $\eta = 0$ is DDIM: no new noise. The whole path is determined by the starting point, so the sampler is an ODE that follows the sampling direction from the previous section. It needs fewer steps for a given quality.

Both use the same trained network. Generate, then try Few paths to follow individual points, and slide the number of steps down to see where each sampler starts to break.

Same network, $\eta = 0$. Press Generate and compare the paths with the widget above: smooth curves instead of random walks, and the final shape is just as good (or better) at low step counts.

The start-points menu is a fun one to try: a ring or a grid of starting points is not what the model was trained for, but you can still watch the deterministic sampler carry each point to the data.

5. Flow matching: a straight line from noise to data

Diffusion’s noise schedule was a design choice. Flow matching asks: why not take the simplest possible path between a noise sample $\varepsilon$ and a data sample $x_0$? A straight line:

$$ x_t = (1-t)\,x_0 + t\,\varepsilon $$

This is the same family as before, with $\alpha = 1 - t$ and $\sigma = t$. Because the path is straight, the velocity along it is a constant, $v = \varepsilon - x_0$, and that is what the network learns to predict.

Training changes in exactly one line: the target is the velocity instead of the noise.

    xt  = (1 - t) * x0 + t * eps
    loss = ((net(xt, t) - (eps - x0)) ** 2).mean()

The model learns the average velocity over all (noise, data) pairs that pass through a point $x_t$. At a given $(x, t)$ many different pairs cross, so the average points roughly towards where the density is, and the arrows below look similar to the diffusion sampling direction. Train it on the same data you picked in section 1.

6. Sample with a few Euler steps

Generating is now just solving an ODE: start from noise and follow the learned velocity backwards from $t = 1$ to $t = 0$, one Euler step at a time,

$$ x \leftarrow x - \Delta t \cdot \hat v(x, t). $$

With many steps this works well. With few steps, things get interesting. The learned paths are slightly curved, and an Euler step is a straight segment along the current tangent. Try 1, 2, 3 steps and see where the points land.

7. Reflow: straighten the paths

If the paths were exactly straight, a single Euler step would be exact. Reflow gets closer to that with a neat self-distillation:

  1. Take the trained model and run it from many fresh noise samples $z$ to the end of the ODE, giving pairs $(z, x_1)$.
  2. Train a second model (initialised from the first) to follow the straight line between each pair: target velocity $z - x_1$.

The new model is trained on a coupling where each noise sample is matched to the data point the original model would send it to. Lines between such pairs rarely cross each other, so the learned velocity field has less averaging to do and the paths come out straighter.

First generate the pairs, then train the reflow model.

8. Original vs. reflowed

The payoff is at low step counts. Both models below start from the same points and use the same number of Euler steps. Look at the overlay at 1–3 steps: the original model smears points between the arms of the shape, while the reflowed one lands much closer to the data. The table puts numbers on this: the mean distance from each sample to the nearest data point, and path straightness (path length divided by chord, 1 = perfectly straight).

Wrapping up

  • The forward process mixes data with noise; the network learns to undo it.
  • Diffusion predicts the noise; flow matching predicts a constant velocity along a straight line. In both cases the sampler follows an arrow field, and the only question is how many steps you can afford.
  • DDPM adds noise at each step, DDIM doesn’t; flow matching is an ODE from the start.
  • Reflow pays one extra round of training to buy much cheaper sampling.

The standalone playground has all of these in one page if you want to explore without reading.