Build Your Own Diffusion Model — Starting from MNIST

A diffusion model built up from nothing on 28×28 handwritten digits: the two conditions a noise schedule has to satisfy, how the step number gets injected into a U-Net, and why the sampler adds noise back at the very end — the places you only discover by writing the code yourself.

Denoising Diffusion Probabilistic Models

Digging a buried sculpture out, one sweep at a time

Read an explanation of diffusion models and you walk away satisfied: add noise, then subtract it. Then you sit down to write one and your hands stop. Where does the array of $\sqrt{\bar{\alpha}_t}$ get built? How does the network find out which step number $t$ it is looking at? Once training finishes, where does the very first image come from? Code has a way of exposing exactly where "I understand this" turns into "I have seen this explained."

This article builds a diffusion model from a minimal skeleton on MNIST — 28×28 grayscale handwritten digits. MNIST is not a timid choice. At 28×28×1 it is the smallest subject that produces something digit-shaped in tens of minutes on a GPU you already own, and, just as importantly, its failures are legible. Gray fog means normalization or schedule; the same digit every time means diversity; a grainy finish means the last sampling step. Make those same mistakes at 512×512 in color and you will burn a night before you have even isolated the cause.

One metaphor, then we get to work. Bury a sculpture in a sandbox, one scoop at a time, hundreds of scoops, until nothing shows. Training a diffusion model means filming that burial and then practising, over and over, one single skill: look at the sand as it is now, and name the scoop that was thrown on last. Once you can name it, you can start from nothing but sand and take it off one scoop at a time. When you finish, a sculpture you never buried is standing there.

Choosing how to destroy it — the move that makes this work

Generative modeling is hard because the data distribution itself is too tangled to handle directly. GANs hand the "is this real" judgment to a second network; VAEs go through a latent variable and approximate (Build a VAE from scratch). Diffusion answers differently. Stop trying to solve the hard problem in one shot; replace it with hundreds of easy ones lined up in a row.

The trick is designing the destruction yourself. Add a whisper of noise to an image, repeat a few hundred times, and you end at pure noise. Because that process is a rule you wrote, you can compute the state at any point along it exactly. And when each individual change is small enough, the single step backwards can be approximated by an operation of the same shape — another Gaussian. So the thing being learned is pinned to "one step's worth of small discrepancy." Every time it is asked, the network solves the same easy regression problem.

The forward process: down a T-step staircase in a single jump

Start with the destruction. $\beta_t$ is a small positive number saying how much damage step $t$ does.

In words: shrink the previous image slightly, then add noise with variance $\beta_t$. The shrinking before the adding is what keeps the magnitudes from running away over hundreds of repetitions, and it is why the final state settles exactly onto a standard Gaussian.

Implemented literally, getting $x_t$ costs $t$ iterations. Training draws a random $t$ for every sample, so that is fatally slow. But set $\alpha_t = 1-\beta_t$ and $\bar{\alpha}_t = \prod_{s \le t}\alpha_s$, and the whole staircase collapses into one formula.

That is: dilute the original image to $\sqrt{\bar{\alpha}_t}$ of its strength and fill the remainder with noise. $\bar{\alpha}_t$ slides from 1 down toward 0, and how it slides is the mixing ratio between picture and noise. Implementation-wise this single line is the whole story — the loop disappears, and any $t$ you want is one expression away.

The noise schedule — effectively the only design decision

T = 1000
betas = torch.linspace(1e-4, 0.02, T)          # the linear schedule from the DDPM paper
alphas = 1.0 - betas
alphas_bar = torch.cumprod(alphas, dim=0)      # ᾱ_t: build it once, reuse everywhere

def q_sample(x0, t, eps):                      # t may differ across the batch
    ab = alphas_bar[t].view(-1, 1, 1, 1)
    return ab.sqrt() * x0 + (1 - ab).sqrt() * eps

How you pick $\beta$ is essentially the only thing in a diffusion model that feels like design. The DDPM paper used $T=1000$ and raised $\beta$ linearly from $10^{-4}$ to $0.02$. Two conditions are worth checking.

The first is that $\bar{\alpha}_T$ ends up close enough to zero. Break this and you get a real failure. Generation starts from pure noise, but during training the model only ever saw inputs with a faint ghost of the original still in them; its very first move at inference lands out of distribution, and what comes out is gray fog. Decide that MNIST is easy, drop $T$ to 200, and leave the $\beta$ range untouched, and this is precisely what happens. Shorten $T$ and you must raise the end of $\beta$ too.