Diffusion from scratch

A class-conditional DDPM with classifier-free guidance and DDIM sampling, written from scratch in PyTorch (no diffusers or other diffusion libraries) and trained entirely on a laptop CPU, on MNIST and FashionMNIST.

Try it

• In the browser: interactive demo. Switch between MNIST and FashionMNIST (direct link), drag the guidance and sampler sliders over precomputed samples (every setting starts from the same noise), and watch the samples denoise.
• Locally, with live sampling: clone the repo, then uv sync && make serve and open http://127.0.0.1:8000. Pick a dataset, class, seed, guidance weight, sampler, step count, and η; the model runs on one CPU thread. The trained weights for both datasets ship with the repo (models/*.pt, 1.7 MB each, EMA weights only), so no training is needed; a checkpoint you train yourself (make train) takes precedence. API docs are at /api/docs.

Results (MNIST)

Class-conditional samples, DDIM 50 steps, guidance w = 2. Each row is the requested class:

DDIM matches 1000-step DDPM at 20× fewer steps. At guidance w = 2, DDIM with 50 steps reaches the same class accuracy as DDPM with 1000 (99.9%) at a lower FID (19.0 vs 28.0; real images score 3.1), while sampling 18–21× faster in wall-clock time. Even 10 steps keeps 99.9% accuracy.

sampler	steps	class acc ↑	FID ↓	KID×1000 ↓	wall-clock speedup
DDPM	1000	99.9%	28.0	132.1	1×
DDIM	100	99.8%	18.7	78.8	—
DDIM	50	99.9%	19.0	80.7	18–21×
DDIM	20	99.9%	19.4	81.4	51×
DDIM	10	99.9%	18.6	72.2	113×
real images	—	99.0%	3.1	0.4	—

N = 1,000 samples per config. Speedups come from a sequential benchmark on one CPU thread (batch 10); DDIM 50 was timed before and after DDPM to bracket changes in background load, and the step ratio is exactly 20×.

Why DDIM beats DDPM here: it's the noise, not the step count. Turning DDIM's stochasticity back on (η) at a fixed 50 steps walks FID back up to DDPM's level; at 10 steps, full stochasticity hurts badly, because each large step injects noise the model has too few remaining steps to remove:

sampler (w = 2)	class acc ↑	FID ↓
DDIM 50, η = 0	99.9%	19.0
DDIM 50, η = 0.5	99.7%	22.6
DDIM 50, η = 1	100.0%	30.1
DDPM 1000	99.9%	28.0
DDIM 10, η = 0	99.9%	18.6
DDIM 10, η = 1	100.0%	36.4

Same starting noise, same label, different samplers. Deterministic DDIM (η = 0) gives nearly the same image at every step count; DDPM's stochastic steps take a different path and, at this guidance weight, produce visibly thinner strokes:

Classifier-free guidance trades diversity for fidelity. w = 1 (the plain conditional model) gives by far the best FID (8.5) at 92.7% class accuracy; w = 2 buys 99.9% accuracy at FID 19.0; past that, samples stay on-class but FID climbs as strokes thicken and variety collapses. w = 0 is the unconditional model (chance accuracy), which saw only 10% of training batches:

guidance w (DDIM 50)	class acc ↑	FID ↓
0	10.7%	76.4
1	92.7%	8.5
2	99.9%	19.0
3	100.0%	45.4
5	100.0%	79.0

Longer training shifted this trade-off: after 1,812 steps the plain conditional model reached only 75% accuracy and needed w = 2 to look right; after 15,351 steps it reaches 93%, so the FID-optimal guidance weight dropped from 2 to 1.

Denoising trajectory (DDIM 50, x_t from pure noise to the final sample) and the training loss of the 0.42M-parameter UNet (15,351 steps, about 16 epochs, 3 hours in total on a shared laptop CPU):

Results (FashionMNIST)

The same model and code trained on FashionMNIST (15,921 steps at batch 52, about 14 epochs, 2.5 hours on one CPU core). Each row is the requested class (T-shirt, trouser, pullover, dress, coat, sandal, shirt, sneaker, bag, boot):

The MNIST findings hold on a harder, more varied dataset. At w = 2, DDIM 50 matches DDPM 1000 on class accuracy (92.0% vs 93.4%) with a lower FID (13.8 vs 24.3; real images score 2.3) at 19–21× the speed, and adding DDIM noise (η) again moves FID back toward DDPM's:

sampler (w = 2)	class acc ↑	FID ↓	wall-clock speedup
DDPM 1000	93.4%	24.3	1×
DDIM 100	91.8%	13.2	—
DDIM 50	92.0%	13.8	19–21×
DDIM 20	92.4%	14.4	47×
DDIM 10	92.9%	15.6	114×
DDIM 50, η = 0.5	93.0%	16.9	—
DDIM 50, η = 1	93.0%	24.9	—
DDIM 10, η = 1	94.0%	26.6	—
real images	86.7%	2.3	—

N = 1,000 samples per config. The evaluation classifier is only 86.5% accurate on the test set (shirt, T-shirt, pullover, and coat are easily confused), so real images score about that; generated samples score higher because guidance pulls them toward unambiguous, prototypical examples of each class.

Guidance behaves as on MNIST, but the plain conditional model (w = 1) is weaker here (78.4% accuracy), so w = 2 buys 14 points of accuracy for about 1 FID point before larger weights collapse variety:

guidance w (DDIM 50)	class acc ↑	FID ↓
0	10.3%	49.6
1	78.4%	12.9
2	92.0%	13.8
3	93.9%	37.1
5	94.7%	73.0

Ablation: cosine vs linear noise schedule (MNIST)

Two MNIST models that differ only in the noise schedule: linear (Ho et al. 2020) and cosine (Nichol & Dhariwal 2021), each trained with the same UNet, optimizer, learning rate, batch size (52), and step count (15,351), then scored with the same sampler configs, seeds, and classifier at N = 1,000. Both were trained and evaluated on a Kaggle T4 GPU (scripts/gpu_ablation.py); each training run took about 7.5 minutes, against about 3 hours on the laptop CPU. As a check, the GPU run also re-scored the CPU-trained model from the sections above and reproduced every number to within 0.0003, because sampling noise is drawn on the CPU from the same seeds either way.

sampler	FID, linear	FID, cosine	class acc, linear	class acc, cosine
DDPM 1000, w = 2	25.4	22.5	100.0%	99.6%
DDIM 100, w = 2	16.9	11.6	99.8%	99.6%
DDIM 50, w = 2	16.9	11.1	99.7%	99.8%
DDIM 20, w = 2	17.2	9.5	99.8%	99.3%
DDIM 10, w = 2	16.4	10.2	99.9%	98.4%
DDIM 50, w = 2, η = 1	28.7	24.7	100.0%	100.0%
DDIM 50, w = 1	9.5	18.8	91.8%	88.1%
DDIM 50, w = 0	76.9	91.1	10.4%	10.9%
DDIM 50, w = 3	41.7	32.6	100.0%	99.9%

The cosine schedule helps guided sampling at every step count. With guidance w = 2 it cuts DDIM FID by 32–45% (DDIM 20: 17.2 → 9.5), and DDPM improves too (25.4 → 22.5). The linear schedule destroys signal quickly, so many of its 1000 timesteps sit at almost pure noise; the cosine schedule spreads them more evenly, which gives a coarse DDIM subsequence more useful steps.

It hurts unguided sampling. The plain conditional model (w = 1) roughly doubles its FID with cosine (9.5 → 18.8), the unconditional model (w = 0) is also worse, and few-step class accuracy drops slightly (98.4% at 10 steps vs 99.9%). This is not a training-budget artifact: an earlier comparison against the CPU-trained linear model, which saw 18% more images, showed the same pattern, and this matched rerun keeps it. Training losses are not comparable across schedules, because each schedule weights the timesteps differently.

(The CPU-trained linear model in the sections above scores slightly worse than the matched GPU one with guidance, e.g. DDIM 50 at w = 2: 19.0 vs 16.9, a reminder that training setup and run-to-run variation move FID by a few points, which is why both models in this table were trained the same way.)

What's implemented

Piece	Where	Notes
Noise schedule + closed-form forward process	ddpm/schedule.py	linear (default) and cosine β schedules, T = 1000
Noise-prediction UNet	ddpm/unet.py	28→14→7, ResBlocks + GroupNorm, attention at 7×7, sinusoidal time + learned class embeddings, 0.42M params
Training	ddpm/train.py	ε-prediction MSE, 10% label dropout for CFG, EMA weights, data-parallel CPU training via torchrun
DDPM sampler	ddpm/sampling.py	ancestral sampling from the true posterior, 1000 steps
DDIM sampler	ddpm/sampling.py	any step count, η ∈ [0, 1]; η = 1 with consecutive steps reduces to DDPM (tested)
Classifier-free guidance	ddpm/sampling.py	conditional and null-label passes batched into one forward call
Evaluation	ddpm/evaluate.py, ddpm/classifier.py, ddpm/metrics.py	class accuracy, FID and KID in a dataset-specific classifier's feature space, wall-clock timing
2D toy experiments	ddpm/toy.py	8 Gaussians, two moons, spiral with an MLP denoiser
Interactive demo	web/, ddpm/export_web.py, ddpm/serve.py	one vanilla-JS page: precomputed samples on GitHub Pages, live sampling through a local FastAPI server (design)

NOTES.md walks through the math (forward process, ε-prediction, DDIM's non-Markovian reverse process, why CFG works) and maps each equation to the code.

2D toy distributions

The same schedule, samplers, and guidance code with a small MLP denoiser. Colors are classes; each model is class-conditional.

Sliced Wasserstein distance to held-out data (lower is better; "real" is the floor from two independent draws of real data):

dataset	DDPM 1000	DDIM 50	DDIM 10	DDIM 5	real	DDPM 1000 time	DDIM 50 time
8gaussians	0.029	0.031	0.033	0.043	0.017	34.2s	1.4s
moons	0.041	0.044	0.092	0.170	0.020	40.1s	0.9s
spiral	0.038	0.079	0.139	0.233	0.015	32.6s	0.7s

DDIM holds quality well down to ~10 steps on simple geometry (8 Gaussians) but degrades sooner on the spiral, where the trajectory curves more and large deterministic steps cut corners.

Classifier-free guidance on 8 Gaussians, asking for one mode. w = 0 is the unconditional model; w > 1 over-sharpens and pushes samples past the mode, the usual diversity-for-fidelity trade:

Reproduce

uv sync                         # CPU-only PyTorch
make test                       # sampler math, schedule, CFG, UNet, metrics
make toy                        # 2D experiments (~5 min)
make train MINUTES=60           # MNIST, data-parallel over single-threaded CPU workers
make eval N=100                 # samples for every sampler config -> metrics, figures, timing
make web                        # regenerate the static demo in docs/ (GitHub Pages)
make serve                      # local demo with live sampling at http://127.0.0.1:8000

Variables: DATASET=mnist|fashion, WORKERS, MINUTES, N, JOBS; make resume continues from the last checkpoint. Long jobs run under nice with OMP_NUM_THREADS=1: on a contended CPU, single-threaded DDP workers sync once per step and scale far better than intra-op threads.

With a GPU, python scripts/gpu_ablation.py --repo . --out out trains and evaluates the cosine-schedule model end to end in about 20 minutes (--device on ddpm.train and ddpm.evaluate generate selects the device).

Larger evaluations can split the expensive DDPM run across processes; shards reproduce the unsharded samples exactly and are merged automatically:

for i in 0 1 2 3; do uv run python -m ddpm.evaluate generate --method ddpm --n 1000 --batch 50 --shard $i --num-shards 4 & done; wait

Roadmap

• DDPM, DDIM, classifier-free guidance, UNet, CPU data-parallel training
• 2D toy experiments and MNIST evaluation (class accuracy, FID/KID, guidance sweep, timing)
• Longer MNIST training run (15,351 steps, about 16 epochs)
• FashionMNIST model and evaluation
• Larger evaluation (N = 1,000 per config) for tighter FID/KID estimates
• DDIM η ablation (η ∈ {0, 0.5, 1})
• Cosine vs linear noise schedule (trained and evaluated on a Kaggle GPU)

Evaluation caveats

• FID here is not comparable to published FIDs. Inception-v3 features suit 28×28 grayscale digits poorly, so FID and KID are computed on the 64-d penultimate features of a small CNN trained on the same dataset. Scores are meaningful only relative to each other and to the "real images" floor at the same sample size.
• Sample size. Each config uses N = 1,000 samples (the 1000-step DDPM run alone is 2M network evaluations with guidance, about 2.5 hours on one CPU core). FID is biased at finite N, so every config, including the real-image floor, uses the same N. An earlier N = 100 run was too small: on FashionMNIST it scored real images worse than generated ones.
• Wall-clock timings were measured on a machine running other jobs. Step counts are the load-independent measure; the timing table confirms the ratio holds in practice.

References

Ho et al. 2020 (DDPM) · Song et al. 2021 (DDIM) · Ho & Salimans 2022 (classifier-free guidance) · Nichol & Dhariwal 2021 (cosine schedule) · Bińkowski et al. 2018 (KID)