Diffusion models

The big idea

Destroying a picture is easy: add a little random noise, then a little more, and after enough steps nothing is left but static. A diffusion modelA generative model that learns to reverse a gradual noising process: starting from pure noise, a network removes a little predicted noise at each of many steps until a sample emerges. learns to run that film backwards. It is trained on one simple task, "here is a noisy image and how noisy it is; what noise was added?", and then generates by starting from pure static and removing a little predicted noise at a time, for hundreds of steps, until a clean image appears.

Think of a sculptor who removes material bit by bit. Each cut is small and easy to judge; no single step has to produce the whole statue. That is the contrast with a GAN, whose generator must produce a finished image in one forward pass and is trained through an unstable game. A diffusion model's training is plain regression on noise, and the hard work is split across many small denoising steps.

How it works

Step 1: The forward process

Fix a schedule of small noise levels ${\beta }_1,\dots ,{\beta }_T$. Each forward step shrinks the image a little and adds noise:

Because a sum of Gaussians is Gaussian, you can jump straight to any step. With ${\alpha }_t=1-{\beta }_t$ and ${\bar{\alpha }}_t={\alpha }_1{\alpha }_2\cdots {\alpha }_t$:

So $x_t$ is a blend: $\sqrt{{\bar{\alpha }}_t}$ of the picture plus $\sqrt{1-{\bar{\alpha }}_t}$ of one noise image. The DDPM paper (Ho, Jain and Abbeel, 2020) used $T=1000$ with $\beta$ rising linearly from $10^{-4}$ to $0.02$. With those numbers $\sqrt{{\bar{\alpha }}_t}$ is 0.724 at step 250, 0.280 at step 500 and 0.006 at step 1000, where almost nothing of the picture is left.

Step 2: The denoising network

A network ${ϵ}_{\theta }(x_t,t)$ is trained to predict the noise $ϵ$ that was mixed in. Training is remarkably simple (the paper's Algorithm 1):

1. Take a training image $x_0$, a random step $t$ between 1 and $T$, and random noise $ϵ$.
2. Make $x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ$ in one jump.
3. Take a gradient step on $\Vert ϵ-{ϵ}_{\theta }(x_t,t){\Vert }^2$.

No adversary, no sampling chain during training: just regression. DDPM derived this as a simplified, reweighted form of a variational bound like the ELBO of Variational autoencoders , and found the simple version gave better samples. For images the network is a U-Net ( Image segmentation ) that receives $t$ through a sinusoidal embedding, the same trick transformers use for positions.

Step 3: Sampling

To generate (the paper's Algorithm 2), start from pure noise $x_T\sim \mathcal{N}(0,I)$ and, for $t=T$ down to 1, remove a scaled amount of predicted noise and add back a little fresh noise:

That is $T$ network evaluations per sample, 1,000 in DDPM, which is the method's main cost. Later samplers take larger steps and need far fewer evaluations.

Step 4: Noise schedules

The schedule decides how much signal survives at each step, and so which noise levels the network spends its training on. Nichol and Dhariwal (2021) found that with the linear schedule the end of the forward process is too noisy to contribute much to sample quality, especially for 64 by 64 and 32 by 32 images, and proposed a cosine schedule, ${\bar{\alpha }}_t\propto {\cos }^2(\frac{t/T+s}{1+s}\cdot \frac{\pi }{2})$ with $s=0.008$, which removes signal more gradually.

Try it

• Slide $t$ from 0 to 1000 on the linear schedule. At step 250 the picture is multiplied by 0.724 and still obvious; at 500 (0.280) it is faint but still there; by 750 (0.058) it is gone, although a quarter of the steps remain.
• Switch to cosine at step 500: the picture is still multiplied by 0.703, about the linear schedule's value near step 250. The cosine schedule spends more of its steps at noise levels where there is still something to recover.
• Set $t$ to 0 and then 10. Even a tiny amount of noise is visible. Those first steps are where the network learns fine detail.

Going deeper

Code: a diffusion model with a perfect denoiser

import numpy as np

# Diffusion on 2-D points with the DDPM linear schedule (T = 1000, beta from 1e-4 to 0.02). Instead of a
# trained network we use the exact best noise predictor for a small dataset, so every number is exact.
rng = np.random.default_rng(0)
T = 1000
beta = np.linspace(1e-4, 0.02, T); alpha = 1 - beta; abar = np.cumprod(alpha)
s = 0.008; f = np.cos((np.arange(T + 1) / T + s) / (1 + s) * np.pi / 2) ** 2
abar_cos = f[1:] / f[0]                                         # cosine schedule (Nichol and Dhariwal)
for t in [100, 250, 500, 750, 1000]:
    print(f"t = {t:4d}: signal kept sqrt(abar) linear {np.sqrt(abar[t - 1]):.3f}   cosine {np.sqrt(abar_cos[t - 1]):.3f}")

ang = rng.uniform(0, 2 * np.pi, 200)
data = {0: np.c_[2 * np.cos(ang), 2 * np.sin(ang)],             # class 0: a ring of radius 2
        1: rng.normal(0, 0.6, (200, 2))}                         # class 1: a wide blob at the centre
both = np.r_[data[0], data[1]]

# 1. The forward process in one jump equals 1,000 small noising steps
x0 = np.tile([2.0, 0.0], (20000, 1)); x = x0.copy()
for t in range(T): x = np.sqrt(alpha[t]) * x + np.sqrt(beta[t]) * rng.normal(size=x.shape)
jump = np.sqrt(abar[-1]) * x0 + np.sqrt(1 - abar[-1]) * rng.normal(size=x0.shape)
print(f"x_T from 1,000 steps: mean {x[:, 0].mean():.3f} std {x.std():.3f}   one jump: mean {jump[:, 0].mean():.3f} std {jump.std():.3f}")

def eps_hat(xt, t, pts):                                        # exact E[noise | x_t] for these points
    d2 = ((xt[:, None, :] - np.sqrt(abar[t]) * pts[None]) ** 2).sum(-1)
    w = np.exp(-(d2 - d2.min(1, keepdims=True)) / (2 * (1 - abar[t])))
    x0_hat = (w @ pts) / w.sum(1, keepdims=True)
    return (xt - np.sqrt(abar[t]) * x0_hat) / np.sqrt(1 - abar[t])

def sample(n, cls=None, w=0.0):                                 # DDPM Algorithm 2, optional guidance w
    x = rng.normal(size=(n, 2))
    for t in reversed(range(T)):
        e = eps_hat(x, t, both if cls is None else data[cls])
        if cls is not None and w: e = (1 + w) * e - w * eps_hat(x, t, both)   # classifier-free guidance
        x = (x - beta[t] / np.sqrt(1 - abar[t]) * e) / np.sqrt(alpha[t])
        if t > 0: x += np.sqrt(beta[t]) * rng.normal(size=x.shape)
    return x

r = np.linalg.norm(sample(400), axis=1)
print(f"unconditional: {np.mean(np.abs(r - 2) < 0.3):.0%} on the ring, {np.mean(r < 0.9):.0%} in the blob")
for w in [0.0, 2.0, 6.0]:
    x = sample(400, cls=1, w=w); r = np.linalg.norm(x, axis=1)
    print(f"blob class, guidance w = {w}: {np.mean(r < 0.9):.0%} in the blob, spread (std) {x.std():.2f}")

t =  100: signal kept sqrt(abar) linear 0.947   cosine 0.986
t =  250: signal kept sqrt(abar) linear 0.724   cosine 0.920
t =  500: signal kept sqrt(abar) linear 0.280   cosine 0.703
t =  750: signal kept sqrt(abar) linear 0.058   cosine 0.380
t = 1000: signal kept sqrt(abar) linear 0.006   cosine 0.000
x_T from 1,000 steps: mean 0.024 std 0.999   one jump: mean 0.011 std 1.003
unconditional: 53% on the ring, 33% in the blob
blob class, guidance w = 0.0: 73% in the blob, spread (std) 0.57
blob class, guidance w = 2.0: 84% in the blob, spread (std) 0.47
blob class, guidance w = 6.0: 90% in the blob, spread (std) 0.44

Training a real network is out of reach in the browser, so the code uses the exact best noise predictor for a small 2-D dataset: 200 points on a ring of radius 2 (class 0) and 200 in a blob at the centre (class 1). For a finite dataset the best estimate of $x_0$ given $x_t$ is a weighted average of the data points, which gives $ϵ$ directly. Everything else is the DDPM recipe.

• Schedules. The linear schedule keeps 0.947, 0.724, 0.280, 0.058 and 0.006 of the signal at steps 100, 250, 500, 750 and 1,000; the cosine schedule keeps 0.986, 0.920, 0.703, 0.380 and essentially 0.
• One jump equals a thousand steps. Noising a point at (2, 0) through 1,000 small steps or in one jump gives the same result: mean 0.024 and 0.011 (the formula says 0.012), standard deviation 0.999 and 1.003.
• Sampling works. Starting from pure noise, 1,000 denoising steps put 53 percent of samples on the ring and 33 percent in the blob's core. With a perfect denoiser for a finite dataset the samples land on (or next to) training points: a real network generalizes only because it cannot fit that perfectly.
• Guidance trades diversity for fidelity. Sampling the blob class with classifier-free guidance weight 0, 2 and 6 puts 73, 84 and 90 percent of samples within 0.9 of the centre, while their spread shrinks from 0.57 to 0.47 to 0.44.

Classifier-free guidance

To generate "a photo of a red bicycle" rather than any image, give the network a condition $c$ (a class or a text embedding): ${ϵ}_{\theta }(x_t,t,c)$. On its own that often follows the condition only loosely. Classifier-free guidanceA sampling technique that combines a diffusion model's conditional and unconditional noise predictions, pushing samples further in the direction the condition implies; raises fidelity to the prompt at the cost of diversity. (Ho and Salimans) trains one network both with and without the condition, by replacing $c$ with a null token for a random fraction of training examples. At sampling time it extrapolates away from the unconditional prediction:

The guidance weight $w$ trades sample quality and prompt fidelity against diversity, as in the code. The authors stress how simple it is: one line of change in training (randomly drop the condition) and one in sampling (mix the two predictions).

Latent diffusion

Running a U-Net 1,000 times on every pixel of a large image is expensive; the latent diffusion paper notes that pixel-space diffusion models often need hundreds of GPU days to train. Rombach and colleagues (2021) moved diffusion into the latent space of a pretrained autoencoder ( Autoencoders ): the encoder downsamples the image by a factor $f$, the diffusion model works on the small latent grid, and the decoder turns the final latent back into pixels. They found downsampling factors of 4 and 8 gave the best balance of quality and cost, and added cross-attention layers so that text or other inputs can condition the U-Net. This is the design behind the open Stable Diffusion models.

Trade-offs

Common mistakes

• Forgetting to give the network the step $t$: it cannot know how much noise to remove.
• Iterating the forward process during training instead of jumping straight to $x_t$ with the closed form.
• Using a schedule designed for one resolution at another without checking how much signal survives at late steps.
• Turning guidance up to fix every problem; past a point it trades away diversity and naturalness.
• Comparing samplers at different numbers of steps without saying so; step count is part of the cost.

In the wild

Check yourself

Recap