Diffusion and flow matching

Imagine a sculptor who cannot carve a statue in one go. Hand them a shapeless block and ask for a horse, and they freeze. But hand them a slightly rough horse, and they can always make it a little less rough: smooth this bump, deepen that line. Now chain that one modest skill. Start from a shapeless block, make it a bit more horse-like, then a bit more, a hundred times over, and a horse comes out.

A diffusion model is that sculptor. It never learns "draw a picture from nothing", which is very hard. It learns "here is a picture with a little too much noise in it: take some of the noise away", which is much easier. And the practice material is free: take any real picture, add noise yourself, and you know exactly what the answer should be.

flowchart LR
  D["Real data<br/>(a photo)"] -- "add a little noise,<br/>T times (fixed, no learning)" --> N["Pure noise<br/>(TV static)"]
  N -- "remove a little noise,<br/>T times (learned)" --> S["A new sample<br/>(a photo no one took)"]

The tiny world this lesson works in. A real image is a long list of numbers, one per pixel colour. To keep everything small enough to see, our "images" have just two numbers, so each one is a dot on a flat plane. The real data is four round blobs of dots, one in each corner (north-east, north-west, south-west, south-east), 600 dots in all. The job: learn to make new dots that land in the blobs, starting from random noise. Every idea below works unchanged on a 512 × 512 colour image; it just has 786,432 numbers per "dot" instead of 2.

Notice one thing about the blobs now, because it matters later: their average is the empty middle. A model that plays safe and outputs "the average answer" puts its dots where no real data lives. The same failure makes blurry images: the average of many sharp faces is a blur.

Everyday picture Think of an old television slowly losing its signal. Turn the static dial up one notch: the picture fades a little and a little snow creeps in. Turn it a hundred notches and you see only snow, and nothing in it tells you what the programme was.

Tiny worked example Follow one pixel whose clean value is . Each step keeps most of the signal and mixes in a little fresh random noise. The noise is a random number drawn from a standard normal distribution: the bell curve centred on 0 whose typical size is 1 (values like 0.5, −1.2 and 0.1 are common, 3 is rare). Say step 1 mixes in a share of noise, and the random draw is .

import math
x_prev, beta, eps = 2.0, 0.1, 0.5
# √(1 − β) x_{t−1}: shrink the old value
round(math.sqrt(1 - beta) * x_prev, 3)  # → 1.897
# √β ε: the noise mixed in
round(math.sqrt(beta) * eps, 3)  # → 0.158
round(math.sqrt(1 - beta) * x_prev + math.sqrt(beta) * eps, 2)  # → 2.06

Why the square roots? A value's variance (its typical squared size) is what matters, and variances of independent things add. If the old value has variance 1, the new one has variance : exactly 1 again. The square roots keep every step at the same overall size, so after many steps the value is noise of size 1, not a number that has blown up or faded to zero.

The shortcut: jump straight to any step

Running steps one by one is slow, and training needs noisy examples at every step. Luckily, adding noise twice is the same as adding it once in a bigger dose, so there is a formula that jumps straight to step . First, multiply up how much signal survives every step so far:

import math
betas = [0.1, 0.2]
alpha_bar = 1.0
# ∏ (1 − β_s): multiply up what each step keeps
for beta in betas:
    alpha_bar *= 1 - beta
round(alpha_bar, 2)  # → 0.72
# the jump to step 30, where ᾱ = 0.64
x0, eps, alpha_bar_30 = 2.0, 0.5, 0.64
round(math.sqrt(alpha_bar_30) * x0 + math.sqrt(1 - alpha_bar_30) * eps, 2)  # → 1.9

flowchart LR
  X0["x₀<br/>clean"] --> X1["x₁"] --> X2["x₂"] --> DOTS["…"] --> XT["x_T<br/>pure noise"]
  X0 -. "shortcut: √ᾱ_t · x₀ + √(1 − ᾱ_t) · ε" .-> XM["x_t<br/>any step, in one jump"]

The noise schedule is the list of . This lesson uses steps with rising in a straight line from 0.0001 to 0.1: tiny doses first, while there is fine detail to lose, bigger ones once little is left.

In code: linear_schedule builds the list of as a NoiseSchedule, whose NoiseSchedule.alpha_bars is the running product; noise_step is one step and add_noise is the shortcut.

Why it matters in practice. The forward process has no parameters and needs no training. Its only job is to manufacture unlimited practice material, and thanks to the shortcut, any clean example can be turned into a practice question at any noise level in one line.

Everyday picture A photo restorer wants to practise removing specks of dust. They take clean photos and sprinkle dust on them themselves, keeping a note of exactly where every speck went. Now they can practise all day: guess where the dust is, then check against the note. The homework is unlimited and every answer is known.

Tiny worked example Take the pixel from before: , noise , step 30, so . The network is shown only and the step number 30, and asked: "what noise was added?" The right answer is 0.5. If it guesses 0.4, it is off by 0.1 and scores . Lower is better.

Why guess the noise rather than the clean value? The two are equivalent: given , and the noise, the clean value follows by undoing the shortcut, . Guessing the noise simply works better in practice, because the target always has the same size (a standard-normal draw), whatever the step.

flowchart LR
  D["pick a real point x₀"] --> MIX
  T["pick a random step t"] --> MIX
  E["draw noise ε"] --> MIX["mix with the shortcut<br/>x_t = √ᾱ_t x₀ + √(1 − ᾱ_t) ε"]
  MIX --> NET["network ε_θ<br/>sees x_t and t"]
  NET --> G["guess ε̂"]
  G --> L["loss ‖ε − ε̂‖²"]
  E --> L
  L --> U["nudge the weights<br/>(backprop + Adam)"]

import random
eps = [0.5, -1.0]
guess = [0.4, -0.8]
# ‖ε − ε̂‖²: square each error and add
round(sum((e - g) ** 2 for e, g in zip(eps, guess)), 2)  # → 0.05
# the baseline: always guessing 0, averaged over many draws of 2-D noise
random.seed(0)
draws = [random.gauss(0, 1) ** 2 + random.gauss(0, 1) ** 2 for _ in range(100_000)]
round(sum(draws) / len(draws), 1)  # → 2.0

This is plain regression (predicting numbers, graded by squared error), the most stable kind of training there is. There is no opponent to balance and no trick: that is a big part of why diffusion displaced GANs.

The network. Our denoiser is a small multi-layer network: the noisy point (2 numbers), plus the step turned into 8 numbers of sines and cosines (a single number "30" is hard for a small network to use; a spread of waves at different speeds is easy, the same trick as the sinusoidal positions in primer.ml.positional), go through two hidden layers of 64 ReLU units and come out as 2 numbers, the noise guess. It trains for 1,500 steps of Adam (primer.ml.optimizers) on batches of 256, in about half a second.

What the noise guess really is: the score

Flip the noise guess around and it becomes an arrow pointing toward the data. If the noise pushed the point up and to the right, the way back is down and to the left. That arrow has a name, the score: at any point, the direction in which the data gets more crowded, fastest. Learning to guess noise is secretly learning the score at every noise level, which is why diffusion models are also called score-based models, and why the training trick is called denoising score matching.

import math
x_t, alpha_bar = 1.5, 0.64
# the best noise guess for standard-normal data: √(1 − ᾱ) · x_t
eps_hat = math.sqrt(1 - alpha_bar) * x_t
round(eps_hat, 2)  # → 0.9
# −ε̂ / √(1 − ᾱ): the score, which for the bell curve is −x
round(-eps_hat / math.sqrt(1 - alpha_bar), 2)  # → -1.5

In code: DenoiserMLP is the network with its hand-written backward pass (denoiser_gradient_check compares it with finite differences), time_features turns the step into waves, train_noise_predictor is the training loop in the diagram and returns a NoisePredictor, and noise_to_score is the score formula.

Why it matters in practice. This loss, a squared error on guessed noise, is the training loss of the original DDPM paper and of the first Stable Diffusion models. The network is far bigger and the data is images, but the training loop is the boxes above.

Everyday picture Back to the sculptor. Look at the block, guess which bits are "not horse", chip away a small part of them, look again. Never chip away everything you think is wrong at once: the guess is rough when the block is rough, and it gets better as the shape emerges.

Tiny worked example A point sits at at a step with and . The network guesses the noise in it is . One step back removes a scaled share of that guess, undoes the shrinking, and (except on the very last step) adds a little fresh noise.

import math
x_t, eps_hat = 1.0, 0.5
beta, alpha_bar = 0.19, 0.75
# remove this step's share of the guessed noise, then undo the shrink
mean = (x_t - beta / math.sqrt(1 - alpha_bar) * eps_hat) / math.sqrt(1 - beta)
round(mean, 4)  # → 0.9
# add a small fresh wobble z
z = 0.5
round(mean + math.sqrt(beta) * z, 4)  # → 1.1179

Why add noise while trying to remove it? The network's guess is an average over every clean point that could have produced . Stepping to that average would drift every sample toward the safe, blurry middle. The fresh wobble keeps each sample committed to one specific possibility, so the samples spread out over the whole data, as the real data does. This recipe is DDPM (denoising diffusion probabilistic models).

flowchart LR
  N["x_T: draw pure noise"] --> G["network guesses<br/>the noise ε̂ at step t"]
  G --> R["remove this step's share,<br/>undo the shrink"]
  R --> W["add a small fresh wobble<br/>(not on the last step)"]
  W --> Q{"t = 0?"}
  Q -- "no: t ← t − 1" --> G
  Q -- "yes" --> OUT["x₀: a new sample"]