Diffusion Workbench

Diffusion and Flow Models on One Loop

Generative AI · Oct 2026 · Completed

DDPM, score SDEs, flow matching and EDM with one network and one loop, so only the formulation changes.

FID 1.06 at 50 NFE
CFG: 58.5% → 96.0% label accuracy
45 tests

Built withPyTorchtorchvisionMatplotlibpytestKaggle

Overview

DDPM, score-based SDEs, flow matching and EDM are usually presented as different models, but they are one idea: corrupt data with Gaussian noise at many levels, train a network to undo a little of it, then start from pure noise and undo it step by step. They differ in three choices: how the noise is added, what the network predicts, and how the reverse process is integrated. I wanted to see those three choices and nothing else, so I wrote each formulation as one class with two functions, loss() and sample(), in plain PyTorch.

They all share one small U-Net (or a DiT), one training loop with AdamW and EMA weights, one sampler interface (DDPM ancestral, DDIM, Euler-Maruyama, the probability-flow ODE, Euler and Heun) and one evaluation. The network never knows which method it serves: it receives an image and a noise level and returns an image-shaped tensor whose meaning the method decides. Classifier-free guidance works with every method. The models are 0.21 M parameters by default so that they train in minutes on a CPU; a Kaggle/Colab notebook trains them bigger on CIFAR-10.

After 500 steps on MNIST, x0-prediction, v-prediction, flow matching and EDM all produce digits, and with the same starting noise they produce the same digits. ε-prediction is the slowest to learn and deterministic samplers expose it, which is exactly the instability that motivated v-prediction and EDM's preconditioning. Guidance at w = 3 takes the conditional model from following the label 58.5% of the time to 96.0%. The repo replaces a coursework notebook that called Stable Diffusion through a pipeline with num_inference_steps and guidance_scale as black-box knobs; it builds what is behind those two numbers.

The question

Every diffusion method adds Gaussian noise and learns to remove it; with the network, data and budget fixed, how much do the three remaining choices (how noise is added, what is predicted, how the reverse process is integrated) change what you get per sampling step?

Key features

Four formulations behind one flag: DDPM with ε, x0 or v prediction and linear or cosine schedules, score SDEs in VP and VE form, flow matching (rectified flow) and EDM with Karras preconditioning.
Six samplers through one interface, with the number of network evaluations counted by wrapping the network rather than assumed (Heun costs 2N - 1).
Classifier-free guidance for every method: labels dropped 10% of the time during training, then uncond + w (cond - uncond) at sampling.
Two backbones with the same call: a convolutional U-Net with attention at 7x7, and a patch Transformer (DiT) with adaLN-Zero.
One evaluation: a classifier-feature FID from a small MNIST CNN, label accuracy for conditional models, and an NFE sweep at 5, 10, 20, 50 and 100 evaluations per image.
Four commands (diffusion-data, diffusion-train, diffusion-sample with denoising trajectories, diffusion-eval), five example scripts on an offline shapes dataset, and a Kaggle/Colab notebook that derives each method before running it.
45 tests, including exact-denoiser tests: for data made of a single image the optimal network output is known in closed form, and every sampler must land on that image.

How it works

Add noise
Each method defines a noisy sample at a random level: sqrt(ā) x0 + sqrt(1 - ā) ε for DDPM, m(t) x0 + s(t) ε for the SDEs, the straight line (1 - t) x0 + t ε for flow matching, x0 + σ ε for EDM. Every image in a batch gets its own level.
Regress a target
The method decides what the network output means: the noise ε, the clean image x0, v = sqrt(ā) ε - sqrt(1 - ā) x0, the velocity ε - x0, or EDM's preconditioned output. Given any one of them the others follow by algebra; they differ in how hard they are to learn at each noise level.
One network, one loop
The U-Net or DiT receives the noisy image and a time embedding (plus a label) and returns an image-shaped tensor. The loop draws a batch from memory, calls method.loss(), clips gradients, steps AdamW with warm-up and updates an EMA copy that is used for sampling.
Integrate the reverse process
method.sample() starts from pure noise and runs a sampler: DDPM ancestral or DDIM for DDPM, Euler-Maruyama or the probability-flow ODE for score SDEs, Euler or Heun for flow matching and EDM. The number of network evaluations (NFE) is the real cost.
Guide and measure
With guidance, one evaluation runs the network on a doubled batch and extrapolates from the unconditional output. Each checkpoint is scored at several NFE budgets with the classifier-feature FID against the test set and, for conditional models, label accuracy.

Examples

Untrained versus trained, same noise and same sampler, for x0-prediction DDPM with DDIM, flow matching with Euler, EDM with Heun and flow matching with guidance w = 3. The untrained x0 model returns flat grey; after 500 steps the digits appear and the guided model follows the labels 0, 1, 2, 3, ...

The reverse process step by step: the same 16 noise images through DDPM ancestral (100 steps), DDIM (50), flow matching with Euler (50) and EDM with Heun (25), each frame at the same fraction of its steps. The three deterministic samplers end on nearly the same digits; DDPM injects fresh noise each step and stays grainy until the end.

Learning checkpoint by checkpoint: 36 fixed noise images sampled with Euler at 20 steps from flow-matching snapshots saved every 25 training steps. Noisy strokes appear at step 75, clean strokes by 150, and from step 300 each noise image keeps its digit and only gets smoother.

Classifier-feature FID against the number of network evaluations for every run. Flow matching with Euler is usable at 5 NFE; EDM's Heun needs about 19; ε-prediction with deterministic samplers stays above 300.

Denoising trajectory of flow matching with Euler at 20 NFE. Rows alternate x_t and the network's current guess of the final image; the guess has the digit's shape long before x_t is clean.

Input

diffusion-eval runs/flow-matching-cfg/last.pt --guidance 0 1 3 --nfe 20

Output

w = 0: FID 3.20, label accuracy 11.0% (chance) w = 1: FID 2.07, label accuracy 58.5% w = 3: FID 1.10, label accuracy 96.0%

The guidance sweep of the conditional flow-matching run, 200 samples at 20 NFE. At this budget guidance improves both label accuracy and FID because the extrapolation sharpens an undertrained model's blurry samples.

Tech stack

PyTorch
Methods, backbones, samplers, the training loop and the evaluation
torchvision
MNIST, FashionMNIST and CIFAR-10 downloads
Matplotlib
Sample grids, denoising trajectories, loss and NFE curves
pytest
45 tests, including the exact-denoiser sampler checks
Kaggle
Notebook that derives each method and trains bigger models on CIFAR-10 on a free GPU

What I learned

Deterministic samplers of different formulations learn nearly the same noise-to-image map: the same starting noise gives the same digits
The prediction target decides how hard the noisiest steps are: a small ε error near pure noise leaves noise that nothing removes later
Testing samplers against an exact denoiser separates sampler bugs from training effects; it showed that the score-SDE samplers needed a final denoising step

Limitations

Short single-thread CPU runs on 28-px MNIST with 200-sample FIDs from a small classifier, not Inception: the numbers rank the runs in this repo and nothing else
The DiT backbone is undertrained after 500 steps (FID 13.8-17.6, blocky patches), and no pretrained weights are published

Keywords

diffusion models · DDPM · DDIM · v-prediction · x0-prediction · score-based generative models · score SDE · probability-flow ODE · flow matching · rectified flow · EDM · classifier-free guidance · DiT · U-Net · NFE · FID · MNIST · PyTorch · generative models