Train a tiny diffusion model

You have the equations; this is running them. The dataset is a few thousand points arranged in two crescents on a plane, which trains in about a minute on a laptop and lets you watch structure appear out of noise in a scatter plot. Everything here transfers to images unchanged except the network. Starting in two dimensions is deliberate: on images a diffusion bug looks exactly like needing more training for several hours, while on points it is visible immediately.

Start in two dimensions#

Images take hours to train. Points on a plane take a minute on a laptop CPU and you can see every step. Everything here transfers to images; only the network changes.

import torch, torch.nn as nn
from sklearn.datasets import make_moons
torch.manual_seed(0)

X = torch.tensor(make_moons(10_000, noise=0.05, random_state=0)[0], dtype=torch.float32)

T = 1000
betas = torch.linspace(1e-4, 0.02, T)
alphas = 1 - betas
abar = torch.cumprod(alphas, dim=0)

class Denoiser(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(3, 128), nn.SiLU(),
            nn.Linear(128, 128), nn.SiLU(),
            nn.Linear(128, 2),
        )

    def forward(self, x, t):                        # x: [n, 2], t: [n]
        return self.net(torch.cat([x, (t / T).unsqueeze(-1)], dim=-1))

Train#

model = Denoiser()
opt = torch.optim.Adam(model.parameters(), lr=1e-3)

for step in range(10_000):
    x0 = X[torch.randint(len(X), (256,))]
    t = torch.randint(T, (256,))
    noise = torch.randn_like(x0)
    a = abar[t].unsqueeze(-1)
    xt = a.sqrt() * x0 + (1 - a).sqrt() * noise
    loss = ((model(xt, t.float()) - noise) ** 2).mean()
    opt.zero_grad(); loss.backward(); opt.step()
    if step % 2000 == 0 or step == 9999:
        print(f"step {step:5d}  loss {loss.item():.4f}")

# step     0  loss 1.0144
# step  2000  loss 0.1892
# step  4000  loss 0.2067
# step  6000  loss 0.1869
# step  8000  loss 0.2063
# step  9999  loss 0.1632     all 10,000 steps in a few seconds on a laptop CPU

The loss will not go to zero and should not. At large t the model is being asked to predict a noise vector from something that is almost entirely that noise vector, which is easy; at small t it is being asked to identify which specific noise was added to a nearly clean point, which is close to impossible, and the irreducible error there sets the floor. A converged run on this dataset sits around 0.16 to 0.20. The number itself tells you very little — that is why the next part measures the samples instead.

Sample#

@torch.no_grad()
def sample(n=2000):
    x = torch.randn(n, 2)
    for t in reversed(range(T)):
        eps = model(x, torch.full((n,), float(t)))
        x = (x - betas[t] / (1 - abar[t]).sqrt() * eps) / alphas[t].sqrt()
        if t > 0:
            x = x + betas[t].sqrt() * torch.randn_like(x)
    return x

points = sample()      # scatter plot these: two interleaved moons should appear

Score it, do not squint at it#

A scatter plot will tell you whether the crescents are there. It will not tell you whether a change helped, and by the third experiment you will be comparing plots from memory. One number fixes that: for each generated point, the distance to the nearest real point, averaged. It is cheap, it has an obvious floor, and it has an obvious ceiling, so a result is interpretable without any further apparatus.

real = X[torch.randperm(len(X))[:2000]]
fit = lambda p: torch.cdist(p, real).min(dim=1).values.mean().item()

print(f"generated  {fit(sample()):.4f}")
print(f"real data  {fit(X[torch.randperm(len(X))[:2000]]):.4f}")   # the floor
print(f"noise      {fit(torch.randn(2000, 2)):.4f}")               # the ceiling
# generated  0.0758
# real data  0.0107
# noise      0.4810

This metric only rewards being close to the data, so a model that produced two thousand copies of the same real point would score a perfect 0. Pair it with a spread measure — the mean distance from each generated point to its nearest other generated point — before you trust it. That pairing is the whole of generative evaluation in miniature: one number for fidelity, one for coverage, and neither is meaningful alone. Reported means and standard deviations agree too: generated points came out at mean (0.436, 0.286) with standard deviations (0.902, 0.496), against (0.500, 0.250) and (0.867, 0.498) for the data.

Fewer steps, and why the sampler is not a detail#

The obvious way to make sampling ten times faster is to run every tenth step of the same loop. Try it and the result is not ten times faster at slightly worse quality; it is ten times faster at no quality. The DDPM update assumes you are moving one step, and its coefficients are wrong for a jump. Swapping in the DDIM update, which converts the noise prediction into an estimate of the clean point and re-noises it to whatever level you ask for next, fixes it with no retraining at all.

@torch.no_grad()
def ddim(n, steps):
    ts = torch.linspace(T - 1, 0, steps).long().tolist()
    x = torch.randn(n, 2)
    for i, t in enumerate(ts):
        eps = model(x, torch.full((n,), float(t)))
        a = abar[t]
        x0 = (x - (1 - a).sqrt() * eps) / a.sqrt()               # estimate the clean point
        a_prev = abar[ts[i + 1]] if i + 1 < len(ts) else torch.tensor(1.0)
        x = a_prev.sqrt() * x0 + (1 - a_prev).sqrt() * eps       # re-noise to the next level
    return x

steps    DDPM loop, every kth step    DDIM
 1000        0.0717  (0.87 s)          0.0766  (0.67 s)
  100        0.3098  (0.07 s)          0.0787  (0.07 s)
   50        0.3933  (0.04 s)          0.0765  (0.11 s)
   20        0.4323  (0.02 s)          0.0744  (0.02 s)
   10        0.4608  (0.01 s)          0.0819  (0.01 s)

pure noise scores 0.4810, real data scores 0.0107

That table is the reason sampler names appear in every image-generation interface and why people argue about them. It is also the reason a diffusion result is not reproducible without saying which sampler produced it. If you take one habit from this section, take this: when generation is too slow, change the sampler before you change anything else, and measure rather than assume.

Conditioning and guidance, measured#

make_moons returns a label saying which crescent each point came from, so the same dataset supports conditional generation. Add an embedding for the label, reserve one extra index for no label, and drop the real label 10 per cent of the time during training. That single change is all classifier-free guidance requires: one network that has learned both the conditional and the unconditional prediction.

Xn, Yn = make_moons(10_000, noise=0.05, random_state=0)   # this time keep the labels
X = torch.tensor(Xn, dtype=torch.float32)
Y = torch.tensor(Yn, dtype=torch.long)
real = X[torch.randperm(len(X))[:2000]]

class Cond(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = nn.Embedding(3, 16)          # labels 0 and 1, plus 2 meaning no label
        self.net = nn.Sequential(
            nn.Linear(3 + 16, 128), nn.SiLU(),
            nn.Linear(128, 128), nn.SiLU(),
            nn.Linear(128, 2),
        )

    def forward(self, x, t, y):
        return self.net(torch.cat([x, (t / T).unsqueeze(-1), self.emb(y)], dim=-1))

cond = Cond()
opt = torch.optim.Adam(cond.parameters(), lr=1e-3)
for step in range(20_000):
    i = torch.randint(len(X), (256,))
    x0, y = X[i], Y[i].clone()
    y[torch.rand(256) < 0.1] = 2                # drop the label 10% of the time
    t = torch.randint(T, (256,))
    noise = torch.randn_like(x0)
    a = abar[t].unsqueeze(-1)
    xt = a.sqrt() * x0 + (1 - a).sqrt() * noise
    loss = ((cond(xt, t.float(), y) - noise) ** 2).mean()
    opt.zero_grad(); loss.backward(); opt.step()

@torch.no_grad()
def guided(n, cls, w, steps=100):
    ts = torch.linspace(T - 1, 0, steps).long().tolist()
    x = torch.randn(n, 2)
    for i, t in enumerate(ts):
        tt = torch.full((n,), float(t))
        eps_uncond = cond(x, tt, torch.full((n,), 2))
        eps_cond   = cond(x, tt, torch.full((n,), cls))
        eps = eps_uncond + w * (eps_cond - eps_uncond)          # the whole of guidance
        a = abar[t]
        x0 = (x - (1 - a).sqrt() * eps) / a.sqrt()
        a_prev = abar[ts[i + 1]] if i + 1 < len(ts) else torch.tensor(1.0)
        x = a_prev.sqrt() * x0 + (1 - a_prev).sqrt() * eps
    return x

Now the guidance scale stops being folklore. Train that model for 20,000 steps, sample 2,000 points asking for crescent 0 at each of several values of w, and score three things: what fraction of the samples a nearest-neighbour classifier assigns to the crescent you asked for, the fit-to-data number from above, and the spread.

   w   on-class   fit-to-data   spread
 0.0      57.4%        0.0753   0.0230     unconditional: it ignores what you asked for
 1.0      99.5%        0.0189   0.0119     the honest conditional model
 2.0     100.0%        0.0115   0.0046
 3.0     100.0%        0.0107   0.0026
 5.0     100.0%        0.0109   0.0014     8x less varied than at w = 1
10.0     100.0%        0.0130   0.0014
20.0     100.0%        0.2340   0.0017     off the data entirely
40.0     100.0%        0.7531   0.0066     worse than pure noise, which scores 0.4810

Three things fall out of that table that are usually stated as opinions. First, guidance buys adherence to the condition almost immediately: by w = 2 every sample is already on the requested crescent, so the adherence argument for a higher scale is finished there. Second, what continues past that point is the collapse in diversity, a factor of eight between w = 1 and w = 5, which on images is why heavily guided prompts return near-identical compositions across seeds. Third, fidelity holds for a while and then goes suddenly: fit is flat from w = 2 to w = 10 and then jumps to 0.2340 at w = 20 and 0.7531 at w = 40, worse than pure noise, because extrapolating that far past the conditional prediction pushes samples off the data entirely. On images that same failure is the oversaturated, high-contrast, plastic look, and now you know it is not a stylistic quirk of the model but points leaving the region the data occupies.

Things to try#

• Plot intermediate x at several t to watch structure appear from noise.
• Sample with only 100 of the 1,000 steps by skipping, then with DDIM at 100 steps. The gap is the table above, reproduced on your own run.
• Train a second model with the identical loop but the target set to x0 instead of noise, and compare fit-to-data. Noise prediction came out at 0.0767 against 0.0975 for image prediction.
• Break the schedule deliberately: set betas to linspace(1e-4, 0.002) so alphabar at T is 0.37 rather than 0.00004, retrain, and watch sampling fail because the starting noise is nothing like the end of the forward process.
• Move to small images such as MNIST with a small convolutional network.