Diffusion Models
A from-scratch implementation of time- and class-conditioned diffusion models for image generation — implementing and experimenting with a range of U-Net backbones.
Overview
DiffusionLab is a from-scratch implementation of Denoising Diffusion Probabilistic Models (DDPM) for image generation. Rather than fine-tuning an existing library, the entire pipeline — forward noising, reverse sampling, the training loop, noise schedules, and evaluation — is implemented directly in PyTorch.
The project is organized around one central decision: a generic diffusion wrapper owns the diffusion mechanics and the conditioning (time step plus class label), while a pluggable backbone only has to predict the noise. Because the wrapper builds the conditioning and the backbone stays agnostic to it, architectures can be swapped in and out while keeping the exact same training and inference pipeline — making controlled architecture and scaling experiments directly comparable.
At a glance
Motivation
Understanding by construction
Most diffusion code is either a black box (a high-level library call) or a large research codebase that is hard to trace end to end. The goal here was to build it from first principles and to treat the U-Net backbone as a variable under study.
- →Implement q(x_t | x_0), the reverse process, and the noise schedules by hand
- →Isolate conditioning — show that time-only and time+class are just different embeddings for the same backbone
- →Explore architecture scaling empirically (depth, width, attention, schedule, EMA, training length)
- →Reuse one pipeline across 2D toy data, MNIST, and CIFAR-10
Architecture
One wrapper, swappable backbones
x0 ──▶ forward diffusion q(x_t|x_0) (training: random t per sample)
│
▼
xt = √ᾱ_t·x0 + √(1−ᾱ_t)·ε
│
▼
wrapper (cond_emb = time [+ class]) ──▶ backbone ──▶ ε̂
│
▼
loss = MSE(ε̂, ε) ──▶ backprop (AdamW + grad clip + EMA update)
│
▼
reverse sampling: x_T ~ N(0,I) ──▶ iterative denoising ──▶ x0Every backbone implements a single interface — forward(x, cond_emb) → ε̂ — where the wrapper supplies the conditional embedding and the backbone predicts the noise. The backbones are built on shared DownBlock / UpBlock / AttentionBlock primitives (GroupNorm, SiLU, skip connections, residual projections).
Backbones (the variable under study)
| Backbone | Notes | Params |
|---|---|---|
| MLPDenoiser | Proof of concept for 2D toy data | tiny |
| CNNDenoiser | Small U-Net, 2 down/up stages | ~3.3M / ~13M |
| CNNDenoiserLarge | Wider, 3 stages | ~53M |
| CNNDenoiserLargeAttention | + self-attention at 4×4 / 8×8 | ~58M |
| CNNDenoiser15 | 15 conv blocks (2/stage) + attention | ~172M |
| CNNDenoiser21 | 21 conv blocks (3/stage) + attention | ~237M |
| CNNDenoiser27 | 27 conv blocks (4/stage), not yet trained | ~302M |
Results
From a proof-of-concept to a single-digit FID
Each stage is measured by FID on CIFAR-10. Rows 3–5 are the same architecture with only the training/eval technique changing (36.75 → 29.54 → 13.09), which isolates the effect of EMA and the cosine schedule. The final model reaches FID 8.64 with class-conditional generation.
FID vs. epoch — the two headline runs
Loss keeps falling while FID rises (CNNDenoiser21 · EMA · cosine)
| Epoch | Train loss | FID |
|---|---|---|
| 100 | 0.0541 | 31.08 |
| 200 | 0.0528 | 9.40 |
| 300 | 0.0501 | 8.64 |
| 400 | 0.0466 | 11.75 |
| 500 | 0.0432 | 18.02 |
Challenges
What the experiments taught me
Exploding gradients on deep models
During CNNDenoiser21 training the loss jumped from ~0.3 to 5.2e14 in one epoch. Fixed with gradient clipping (max_norm=1.0), a reduced learning rate, and AdamW with weight decay — plus a guard to skip non-finite-loss batches.
The loss / FID divergence
Across every run the training loss kept falling while FID plateaued. The noise-prediction objective is not the only proxy for perceptual quality.
Deterministic sampling collapses
Without injecting noise during the reverse process, all samples converge to a blurry mean image. Stochastic sampling is what buys diversity.
Normalization matters for diffusion
BatchNorm mixes noise states across a batch; switching to GroupNorm keeps each sample's noise state independent.
Memory is the real scaling constraint
Batch size dropped from 128 to 64 for larger models, and FID sample counts were reduced (2048 → 1500) to fit vRAM — motivating batched FID computation.
Generality was worth it
The same trainer/sampler runs on 2D toy data with an MLP and on RGB images with a U-Net — the only change is the architecture plugged into the wrapper.
Takeaway