A DDPM trained from scratch in one 512-line PyTorch file, on a 6GB laptop GPU, in 20 minutes — and then used to answer the question most "diffusion from scratch" repos skip:
How few network evaluations can you get away with, and what actually buys you quality — a better solver, or better-placed steps?
One trained model, four sampler configurations, compared on a fair axis — NFE (network evaluations), not steps, so the 2nd-order solver gets no free 2×.
| config | solver | step placement | NFE per step |
|---|---|---|---|
ancestral (t) |
DDPM posterior, stochastic | uniform in t |
1 |
ddim (t) |
η=0 deterministic — literally Euler | uniform in t |
1 |
ddim (karras) |
same Euler | ρ=7 power law in σ | 1 |
heun (karras) |
2nd-order Heun | ρ=7 power law in σ | 2 (minus 1 on the last step) |
FMD, 10 000 generated samples vs the full 10 000-image MNIST test set. Lower is better; the parenthesised number is the exact NFE spent, and bold marks the winner at each budget.
| NFE (approx) | ancestral (t) | ddim (t) | ddim (karras) | heun (karras) |
|---|---|---|---|---|
| ~5 | 51.28 (5) | 69.07 (5) | 105.93 (5) | 1521.08 (5) |
| ~10 | 11.89 (10) | 28.86 (10) | 18.75 (10) | 8.76 (9) |
| ~20 | 1.85 (20) | 9.04 (20) | 4.60 (20) | 1.91 (19) |
| ~50 | 1.27 (50) | 1.92 (50) | 1.54 (50) | 0.88 (49) |
| ~100 | 0.97 (100) | 1.18 (100) | 1.10 (100) | 0.87 (99) |
Reading it:
- 2nd-order Heun wins from ~9 NFE upward, and it wins by enough to matter: FMD 0.88 at 49 NFE beats every other config's 100-NFE result. That is a genuine 2x saving over stochastic ancestral sampling at equal quality.
- Below that it collapses catastrophically — 1521 at 5 NFE, because 5 NFE buys Heun only 3 steps. A 2nd-order method with too few steps is far worse than a 1st-order method with twice as many.
- Step placement alone is worth up to 2x, and it pays most where you care. Karras spacing beats
uniform-in-
tfor the identical Euler solver at every budget from 10 NFE up, by 1.54x at 10 NFE (18.75 vs 28.86) and 1.96x at 20 (4.60 vs 9.04), narrowing to 1.25x at 50 and 1.07x at 100. Changing where you put the steps costs nothing and needs no new solver. - Stochastic ancestral sampling is a much stronger baseline than its reputation suggests. It beats both deterministic 1st-order configs at every single budget. If you only ever read that "DDIM is the fast one", that is not what happens here.
- The interesting crossover is between 5 and 10 NFE, where the ranking completely inverts. Any claim of the form "sampler X is better than Y" that does not name an NFE budget is not saying anything.
FMD looked non-monotone in the number of samples (17.5 at n=1000, 14.1 at 2000, 11.2 at 5000, then 1.9 at 10 000) and I spent a while convinced the estimator was fragile — the feature covariance is singular, after all, with 14 dead ReLU units and a condition number of 1e32.
It wasn't the estimator. The reference set was test[:n], and the two halves of the MNIST test set are
different populations: 89.1% classifier accuracy on the first 5000 versus 93.6% on the last 5000, and
FMD 22.0 across the halves versus 0.72 within one. Growing n was quietly changing which
distribution I was measuring against. The reference is now pinned to the full test set. Dropping the
dead dimensions and ridge-regularising the covariances — the fix for the bug I thought I had — changed
results in the 4th decimal place.
Print the two σ schedules side by side and the reason is obvious:
t-uniform: 91.66 15.63 8.52 5.78 4.38 ... 0.176 0.120 0.066 0.006
karras: 91.66 76.40 63.36 52.28 42.91 ... 0.041 0.023 0.013 0.006
Uniform-in-t on a cosine schedule is badly conditioned at the top: it leaps σ 91.7 → 15.6 in a
single step, then spends half its remaining budget below σ=1 where the trajectory is nearly straight
and cheap. Karras spacing removes that leap, and that alone is worth ~2x at matched NFE for an
otherwise identical Euler solver.
It also has to be fixed before solver order is even measurable. Heun on uniform-in-t spacing was
not merely worse, it was broken: it produced posterized images with 67.5% of pixels pinned at exactly
±1 (versus 9.5% for Euler), which look crisper to the eye and score 13x worse.
The same conditioning problem is what broke Heun outright in the first place: averaging the derivative at σ=2e4 with the one at σ=12 is meaningless, so a 2nd-order method degrades worse than 1st-order when steps are huge. See INSIGHTS.md for the full debugging trail.
Write the reverse process in the σ parameterisation, where x̃ = x/√ᾱ = x₀ + σ·ε and σ = √((1-ᾱ)/ᾱ). The probability-flow ODE collapses to
dx̃/dσ = ε_θ(x̃, σ)
so deterministic DDIM is exactly the Euler discretisation of it, and DDIM at η=1 is exactly ancestral sampling. A repo that benchmarks "DDIM vs Euler" is benchmarking the same method twice. That is why the fourth config here is 2nd order rather than a differently-named 1st-order method, and why selfcheck asserts both identities numerically instead of asking you to take the algebra on faith:
ancestral == DDIM(eta=1) agrees to 7.0e-05 relative
sigma-space Euler == DDIM(eta=0) agrees to 2.2e-04 relative
One caveat the assert made visible: the equivalence only holds before the x₀ clip. Clipping helps real samples at low NFE, but it breaks the identity between the posterior-in-x form and the DDIM-in-x₀ form, so the check runs with clamp=False.
Quality is measured as a Fréchet distance in the 64-d feature space of a small MNIST CNN trained for about a minute. These numbers are comparable within this repo only — never against published FIDs. Inception-v3 features on upscaled grayscale digits produce authoritative-looking numbers that mean very little, so this repo doesn't pretend otherwise.
The trace term uses the symmetric-PSD form tr((Σ₁^½ Σ₂ Σ₁^½)^½) via eigvalsh. The tempting eigvals(Σ₁ @ Σ₂) returns complex garbage on noisy finite-sample covariances.
python3 napkin_diffusion.py selfcheck # schedule, sampler algebra, metric, one-batch overfit (~1 min)
python3 napkin_diffusion.py train # ~20 min, 2.4GB VRAM
python3 napkin_diffusion.py sweep # 4 configs x 5 NFE budgets x 10k samples, ~35 min
python3 napkin_diffusion.py gif # denoise.gif + samples.png--dataset fashion swaps in Fashion-MNIST. --nfe 5 10 20 50 100 sets the sweep grid.
Everything the run produces lands in out/ (gitignored). The committed copies of the chart, the GIF and the raw sweep numbers are in assets/ — assets/sweep.json is the exact data behind the table above.
torch 2.2 is compiled against numpy 1.x. Shadow it locally instead of touching your global env:
pip install --target .deps "numpy<2"then prefix commands with PYTHONPATH=.deps.
No config system, no trainer class, no package layout, no Hydra, no CIFAR-10. Honest unconditional CIFAR quality is many GPU-hours on a 6GB laptop — this repo would have to either ship bad samples or stop being reproducible in an hour. End to end here is about 55 minutes: 20 min training, 1 min for the metric CNN, ~35 min for the full sweep.
2.81M U-Net, 32/64/128 channels at 32/16/8 px, self-attention at 8 px, sinusoidal time embedding, cosine ᾱ schedule (T=1000), ε-prediction, EMA 0.999, AMP. 30 epochs, batch 128, AdamW 2e-4.
Sampling starts at T_START = 992, the largest t whose ᾱ is still ≥ 1e-4, rather than at t=999. The cosine schedule drives ᾱ(999) down to ~2.4e-9, putting σ at ~2e4 — a step no 2nd-order solver survives. The discarded region holds 0.01% of the signal.