Capstone: a one-dimensional diffusion sampler

The plan

Everything in this course fits into one small program. Pick a dataset, choose a schedule, fit a denoiser, run the reverse loop, and look at what comes out. The dataset here is one-dimensional on purpose: two points at minus one and plus one. A denoiser for that data has a closed form, which means every number the sampler makes can be checked against the truth. The same loop, with a UNet instead of a table, is what an image model runs.

The denoiser will be a table. For each noise level, bin the noisy values and average the clean values that produced them. That is a learned conditional mean, the exact thing the network is trained to approximate, and it needs no gradients. The table is coarse, and that is useful. You can see the discretisation error separately from the sampling error, which is much harder with a neural network. A table is also easy to inspect. You can print it, compare it with the exact curve, and see where the data is too sparse to support an estimate.

Fit, then sample

Training and sampling are two loops over the same schedule. The training loop draws a clean point, draws a timestep, adds the closed-form noise, and updates the denoiser toward the clean point or the noise. The sampling loop starts from pure noise and applies the reverse update at each timestep, using the denoiser it fitted in training. Neither loop knows what the other is doing. They only share the schedule and the shape of the tensor.

Keep the two loops separate in the code. A sampler that reaches into the training data is a sampler that cannot ship, and a training loop that calls the sampler will be slow. The interface between them is one function: given a noisy sample and a timestep, return an estimate of the clean sample. Everything else is scheduling. That separation is what lets you swap a table for a network without rewriting the loop.

The training step, in torch

import torch
from diffusers import DDPMScheduler

scheduler = DDPMScheduler(num_train_timesteps=1000)

for batch in loader:
    x0 = batch.to(device)
    t = torch.randint(0, 1000, (x0.shape[0],), device=device)
    noise = torch.randn_like(x0)
    xt = scheduler.add_noise(x0, noise, t)
    pred = model(xt, t).sample
    loss = torch.nn.functional.mse_loss(pred, noise)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

What to measure

Three checks tell you whether the sampler works. The first is the distribution of the output: does it have the modes of the data, with the right mass? The second is the error along the path: does the estimate of x_0 improve steadily as t falls? The third is the effect of step count: does the output settle as you add steps, or does it change character? A sampler that only looks right at one step count is a coincidence.

Compare against a baseline you trust. For the two-point dataset the exact denoiser is available. The fitted table can be scored against it directly, and the sampler can be scored against the true distribution. For a real model there is no exact answer. The baseline is then a slower sampler, a fixed evaluation set and a metric like FID or CLIP score. The discipline is the same: change one thing and measure.

Where to go next

The same loop scales to images, audio and video with a bigger denoiser and a better autoencoder. The choices that matter are the ones this course kept returning to. They are the schedule, the parameterization, the number of steps, the guidance scale and the space the diffusion runs in. None of them is universal. Every one of them is recorded in the model card of a system you have used.

Build the small version first. A one-dimensional sampler that you can debug end to end teaches more than a pipeline you can only call. Once the loop is yours, every paper is a variation on a theme you already know: a different schedule, a different solver, a different condition. The vocabulary changes. The chain does not.