Diffusion & generative models, section 3 of 5, 4h

The maths of denoising diffusion, just enough

By the end: You can write down the forward process, the training objective and one sampling step, and explain each symbol.

  • The noise schedule
  • Jumping straight to any noise level
  • Why predicting the noise is enough
  • One step of sampling
  • Guidance as arithmetic

New to diffusion & generative models? Start with a square of random static, the kind an untuned television shows. A diffusion model has been trained to do one narrow thing: look at a noisy picture and produce a slightly less noisy one. Run that step fifty or a hundred times, each pass removing a little more, and a picture is what is left. Training works the process backwards — take real photographs, add noise in small steps until they are static, and teach the model to undo one step at a time. Why it matters

Diffusion models are trained by destroying data with noise and learning to undo it. This section is the four equations that say exactly that, with every symbol explained: the schedule of noise, the closed form that lets you jump to any noise level in one line, the training loss, which is plain squared error on the noise that was added, and a single step of sampling. The closed form is the piece everything else rests on, because it is why training is cheap: you never simulate a thousand noising steps to make one training example.

A schedule of small noise steps#

beta_t        small amount of noise added at step t, e.g. from 0.0001 up to 0.02 over T = 1000 steps
alpha_t       1 - beta_t, how much signal survives step t
alphabar_t    alpha_1 x alpha_2 x ... x alpha_t, how much signal survives all steps up to t

One step of the forward process is x_t = sqrt(alpha_t) * x_(t-1) + sqrt(beta_t) * z. The square roots are there so that variance is preserved: if x_(t-1) has unit variance, alpha_t of it survives, beta_t of fresh noise is added, and alpha_t + beta_t = 1, so x_t has unit variance too. That is the whole reason the schedule is written in terms of beta rather than a plain noise standard deviation, and it is why data is normalised to roughly zero mean and unit variance before any of this starts.

Skip to any step directly#

x_t = sqrt(alphabar_t) * x_0  +  sqrt(1 - alphabar_t) * noise        noise ~ Normal(0, 1)

Because adding Gaussian noise repeatedly is still Gaussian, you do not need to run t noising steps to get x_t. One line produces a noisy version of a clean image at any level. This is what makes training cheap: pick a random image, a random t, and noise it in one go.

Do not take that on trust. Running the chain 500 times on 200,000 copies of a single value and comparing with the one-line version is four lines of code, and it is the check that catches a wrong schedule before it costs you a training run.

import torch
torch.manual_seed(0)

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

N, t_check, x0 = 200_000, 500, 0.8
x = torch.full((N,), x0)
for t in range(t_check):                                   # the long way
    x = alphas[t].sqrt() * x + betas[t].sqrt() * torch.randn(N)
print(f"chain, {t_check} steps : mean {x.mean():.4f}  std {x.std():.4f}")

a = abar[t_check - 1]                                      # the one-line way
y = a.sqrt() * x0 + (1 - a).sqrt() * torch.randn(N)
print(f"closed form, 1 line  : mean {y.mean():.4f}  std {y.std():.4f}")
print(f"what it should be    : mean {a.sqrt()*x0:.4f}  std {(1-a).sqrt():.4f}")
# chain, 500 steps : mean 0.2222  std 0.9617
# closed form, 1 line  : mean 0.2253  std 0.9622
# what it should be    : mean 0.2243  std 0.9599
Same distribution, 500 times cheaper. If your two numbers disagree, the schedule or the cumulative product is wrong, and every downstream symptom will look like the model needing more training.

The numbers the schedule actually produces are worth having in your head, because they are what makes the reverse process possible at all. With the linear schedule above, alphabar is 0.971 at t = 50, 0.659 at t = 200, 0.0786 at t = 500 and 0.00004 at t = 1000. So at the last step the original image is multiplied by 0.0064 — it is gone, to any precision that matters, which is exactly what lets sampling begin from a fresh draw of noise that has never seen a dataset.

The training objective#

loss = mean( (noise - model(x_t, t))^2 )

The model sees the noisy image and the step number and predicts the noise that was mixed in. It is plain regression with a squared error. If you know the noise, you know the clean image, since the equation above can be solved for x_0.

x_0 = ( x_t  -  sqrt(1 - alphabar_t) * eps )  /  sqrt(alphabar_t)
The rearrangement. Predicting the noise and predicting the clean image are the same request; they differ only in how the error is scaled at each t, and that is what makes noise the better target.

That scaling is the whole argument for the noise target. At t = 1000, sqrt(alphabar_t) is 0.0064, so recovering x_0 means dividing by 0.0064 — amplifying the model's error by a factor of 156. A network asked for x_0 directly at that step can do no better than output the dataset mean, and the squared error there swamps the loss from every step where the task was learnable. Predicting the noise keeps the target at unit scale for every t, so no part of the schedule dominates. Trained identically on the same data, noise prediction beat image prediction by about a quarter on sample quality in the 2-D experiment in the next section.

The full derivation from the variational bound puts a t-dependent weight in front of this squared error. The result that made diffusion practical is that dropping the weight entirely and using the plain mean squared error above works better in practice. So the loss you implement is simpler than the loss the theory hands you, and knowing that saves you looking for the weighting term in someone's code and concluding you have misread it.

One sampling step#

x_(t-1) = ( x_t - beta_t / sqrt(1 - alphabar_t) * model(x_t, t) ) / sqrt(alpha_t)  +  sqrt(beta_t) * z

z is fresh noise (none on the final step)

Subtract the predicted noise, scaled for this step, rescale, and add a little new noise. Repeat from pure noise down to step one. Faster samplers take bigger, smarter steps using the same trained model.

Two details in that line trip people up. The first is that the coefficient on the predicted noise is beta_t / sqrt(1 - alphabar_t) and not sqrt(1 - alphabar_t): you are removing only the noise this step added, not all the noise accumulated so far, so the subtraction is small. At t = 500 that coefficient is 0.0105, against a total accumulated noise scale of 0.96. Each step nudges. The second is the fresh noise z added back afterwards. Removing it makes the sampler deterministic and, with the DDPM formula, makes the samples worse, because the update is an approximation of a distribution rather than a point and the variance is part of what is being modelled. Deterministic sampling is a real option, but it needs the DDIM update rather than this one.

That DDIM update is worth writing down alongside, because it is what every fast sampler is a variant of. Predict the noise, convert it to an estimate of the clean image, then re-noise that estimate to the next noise level you want — which need not be t - 1. It is the freedom to choose the next level that lets a sampler run 20 steps instead of 1000.

x0_hat  = ( x_t - sqrt(1 - alphabar_t) * model(x_t, t) ) / sqrt(alphabar_t)
x_prev  = sqrt(alphabar_prev) * x0_hat  +  sqrt(1 - alphabar_prev) * model(x_t, t)

alphabar_prev is the noise level of the next step you have chosen, not necessarily t - 1
The deterministic DDIM step. Same trained network, no retraining, and the step count becomes a free parameter at sampling time.

Classifier-free guidance#

guided = model(x_t, t, empty_prompt) + w * ( model(x_t, t, prompt) - model(x_t, t, empty_prompt) )

Train one model that sometimes sees the prompt and sometimes an empty one. At sampling time, take the direction the prompt pushes the prediction, and exaggerate it by a factor w. It costs two model evaluations per step.

Read the endpoints of that formula. At w = 0 the prompt term cancels and you get the unconditional model: a sample from the data distribution, ignoring what you asked for. At w = 1 the two unconditional terms cancel and you get the plain conditional model, which is what an honest sample from p(image given prompt) looks like. Every value above 1 is extrapolation past the model's own conditional prediction, which is why the usual range of 5 to 8 is not a mild setting: the samples produced there are not draws from the conditional distribution at all. They are draws from a sharpened version of it, which is exactly why they look better and exactly why diversity falls.

The name is historical and actively misleading. The original method, classifier guidance, ran a separately trained classifier on the noisy image and pushed the sample along the gradient of the class score. Classifier-free guidance gets the same effect from one model trained with the condition dropped at random, typically 10 per cent of the time, and needs no classifier. Nothing in the formula above involves a classifier; the word only records what it replaced.

Where the field went next#

Several recent image models use flow matching, a closely related formulation that learns a velocity carrying noise to data along straighter paths, which allows fewer sampling steps. The intuitions here carry over directly.

Concretely, flow matching interpolates linearly between noise and data, x_t = (1 - t) * x_0 + t * noise with t running from 0 to 1, and trains the network to predict the velocity noise - x_0 rather than the noise alone. Sampling is then integrating an ordinary differential equation, and because the paths are close to straight, a crude integrator with few steps gets close to the answer. The forward process, the regression objective and guidance all keep their shape; what changes is the interpolation and what the network is asked to output. If you have understood the four equations above you can read a flow-matching paper the same afternoon.

Where people go wrong

Memorising the sampling equation without the closed form of the forward process. The closed form is the idea everything else rests on.

If this comes up in an interview

Write down the diffusion training objective and explain it.

Sample a clean example, a random step t and Gaussian noise, form x_t as sqrt(alphabar_t) times the clean example plus sqrt(1 minus alphabar_t) times the noise, and train the model to predict that noise with a squared error loss.

Then they push further

Why can training jump straight to any step?

The sum of Gaussian noise steps is itself Gaussian with a known variance, so x_t has a closed form given x_0. There is no need to simulate the chain during training.

How does classifier-free guidance work?

The model is trained with and without the condition. At sampling, the unconditional prediction is pushed further in the direction of the conditional one by a guidance weight, strengthening adherence to the prompt at the cost of diversity.

What gets marked up. Marked up for the closed-form forward process, the noise-prediction objective and the guidance formula.

Build this: Implement the forward process on one image and plot it at ten values of t. Then check numerically that the one-line jump and running the steps one by one give samples with the same mean and variance.

This is one section of Diffusion & generative models, 20 hours across 5 sections. Read it in full to get the whole topic in order.