Diffusion from scratch — a 2D generative modelling library

A compact but complete implementation of denoising diffusion — the model family behind Stable Diffusion and Midjourney — written from first principles and studied on 2D toy distributions so every experiment trains in seconds to minutes on a CPU.

The point of working in 2D is not that the problems are easy; it is that the maths and architecture are identical to the image case, while the low dimensionality lets you see the forward process, the reverse process, and the learned score field directly, and measure sample quality exactly. Swap the 2D dataset for images and the MLP for a U-Net and this is DDPM.

One unchanged library learns six qualitatively different distributions. Top: targets. Bottom: samples drawn from pure noise.

---

What's implemented

Beyond a vanilla DDPM, this library includes the pieces that separate a demo from a real implementation:

Area	What's here
Forward process	Closed-form q(x_t | x_0); linear and cosine (Nichol & Dhariwal) noise schedules with derived posterior variances and SNR
Model	Residual MLP noise predictor with sinusoidal time embedding and class conditioning (with a null token)
Objectives	epsilon-prediction and v-prediction (Salimans & Ho) parameterisations
Samplers	Ancestral DDPM; DDIM with eta (η=0 is the deterministic probability-flow path) and few-step subsequencing
Guidance	Classifier-free guidance for conditional generation
Training	Cosine-annealed Adam, gradient clipping, EMA weight averaging
Evaluation	Sliced-Wasserstein, energy distance, and RBF-MMD implemented from scratch
Engineering	Typed configs, mypy-clean, ruff-clean, 42 passing tests, narrated notebook

---

Results

Every model below is small (≤ ~0.4 M parameters) and trained on CPU. Quality is measured with sliced-Wasserstein distance (lower is better; two independent draws from the same distribution score ≈ 0.02–0.03, so these are near the noise floor).

Dataset	Sliced-Wasserstein ↓	Energy distance ↓	MMD ↓
moons	0.036	0.0014	0.0002
spirals	0.036	0.0015	0.0002
circles	0.029	0.0014	0.0002
checkerboard	0.036	0.0014	0.0002
s_curve	0.036	0.0022	0.0005
eight_gaussians	0.058	0.0022	0.0003

The "hero" moons model (200 steps, cosine schedule, 12 k training steps) reaches sliced-Wasserstein 0.023 — visually indistinguishable from the target.

---

How diffusion works (the 30-second version)

Forward process. Take a data point x_0 and add Gaussian noise over T steps until it is indistinguishable from N(0, I). Thanks to the reparameterisation, you can jump to any noise level in one step:

$$x_t = \sqrt{\bar\alpha_t},x_0 + \sqrt{1-\bar\alpha_t},\varepsilon, \qquad \varepsilon \sim \mathcal{N}(0, I)$$

where $\bar\alpha_t = \prod_{s\le t}(1-\beta_s)$ is set by the noise schedule. No learning is involved — this direction is fixed.

Reverse process. Train a network $\varepsilon_\theta(x_t, t)$ to predict the noise that was added. The training objective is just a mean-squared error:

$$\mathcal{L} = \mathbb{E}_{x_0,,t,,\varepsilon}\big[,\lVert \varepsilon - \varepsilon_\theta(x_t, t)\rVert^2,\big]$$

To generate, start from pure noise and repeatedly denoise. The predicted noise is equivalent to the score $\nabla_x \log p_t(x) = -\varepsilon_\theta / \sqrt{1-\bar\alpha_t}$ — the gradient field pointing toward the data. Here it is, learned:

Arrows point toward the data manifold; they are largest where there is no data and vanish on it. This is the vector field the sampler integrates.

---

Beyond the baseline

Cosine vs linear schedule. The cosine schedule keeps signal alive for longer and noises more completely by the final step. At an equal training budget it wins on sample quality (sliced-Wasserstein 0.034 cosine vs 0.041 linear on moons).

Fast sampling with DDIM. Ancestral DDPM needs all T steps. DDIM (eta=0, deterministic) produces comparable samples in 20–50 steps by sampling on a subsequence of timesteps.

Classifier-free guidance. Training a conditional model with occasional label dropout lets us steer generation at sample time by extrapolating between the conditional and unconditional predictions, $\varepsilon = \varepsilon_\varnothing + w,(\varepsilon_c - \varepsilon_\varnothing)$. Increasing the guidance weight w trades diversity for fidelity — each requested class collapses onto its mode.

---

Repository layout

ddpm-from-scratch/
├── diffusion/               # the library
│   ├── config.py            # typed DiffusionConfig / ModelConfig / TrainConfig
│   ├── schedules.py         # linear & cosine schedules + derived quantities
│   ├── datasets.py          # 7 standardised 2D datasets (labelled + unlabelled)
│   ├── model.py             # residual MLP noise predictor (time + class conditioning)
│   ├── process.py           # GaussianDiffusion: q_sample, loss, eps/v predictions
│   ├── ema.py               # exponential moving average of weights
│   ├── sampling.py          # DDPM & DDIM samplers with CFG
│   ├── metrics.py           # sliced-Wasserstein, energy distance, MMD
│   ├── training.py          # training loop
│   └── viz.py               # every figure
├── experiments.py           # reproduce all figures + metrics into outputs/
├── minimal_ddpm.py          # a 150-line single-file reference (start here)
├── notebook/                # narrated walkthrough
├── tests/                   # 42 pytest tests
└── outputs/                 # generated figures, GIFs, metrics, checkpoints

New to diffusion? Read minimal_ddpm.py first — it is the whole idea in one file. Then the diffusion/ package generalises it.

---

Running it

pip install -r requirements.txt

# reproduce every figure and metric (a few minutes on CPU)
python experiments.py

# run the test suite
pytest

# lint + type-check
ruff check . && mypy diffusion

The narrated notebook (notebook/ddpm_from_scratch.ipynb) walks through the theory and runs the library end to end.

---

References

• Ho, Jain, Abbeel. Denoising Diffusion Probabilistic Models. NeurIPS 2020.
• Song, Meng, Ermon. Denoising Diffusion Implicit Models. ICLR 2021.
• Nichol, Dhariwal. Improved Denoising Diffusion Probabilistic Models. ICML 2021.
• Ho, Salimans. Classifier-Free Diffusion Guidance. 2022.
• Salimans, Ho. Progressive Distillation for Fast Sampling of Diffusion Models. ICLR 2022.
• Song et al. Score-Based Generative Modeling through Stochastic Differential Equations. ICLR 2021.