Minimal Diffusion

This repo is designed to a minimal implementation of 2D image diffusion. There are four key pieces:

Piece	File	What it does
Data	[datasets.py](datasets.py)	MNIST images in [-1, 1]
Model	[unet.py](unet.py)	Defines the model architecture (U-Net).
Sampler	[ddpm.py](ddpm.py) NoiseScheduler + sample()	Performs the diffusion utilities: Forward: add noise. Reverse: denoise step by step. This is where the math behind the technique lives.
Training loop	[ddpm.py](ddpm.py) if __name__ == "__main__"	Run the training loop. For each step: Sample t, noise the image to that level, MSE on predicted vs real noise and backprop.

Data

For this implementation, we use the MNIST dataset, a set of 28x28 black and white images of digits. There are 60,000 train and 10,000 test images.

To pre-process the data, use ToTensor() to map pixels to [0, 1], then Normalize((0.5,), (0.5,)) maps that to [-1, 1].

def get_mnist(root="data"):
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,)),  # [0, 1] -> [-1, 1]
    ])
    return datasets.MNIST(root, train=True, download=True, transform=transform)

Why [-1, 1]? The forward process mixes the image with N(0, I). Matching the data scale to the noise scale keeps the math honest, and sampling clamps reconstructed x0 to that same range.

The loader returns (image, label). Unconditional training ignores the label; --conditional passes it into the U-Net.

Wired up in the training script at ddpm.py L194–200.

---

Model architecture

Everything lives in [unet.py](unet.py). The U-Net is a noise predictor: (noisy image, t) → noise with the same shape as the image (B, 1, 28, 28).

The U-Net is made up of a series of residual blocks, each made up of convolutional layers, normalization layers, and activation layers. The structure of each residual block is (unet.py L22–36):

1. GroupNorm → GELU → 3×3 conv
2. add Linear(time_dim → out_ch) broadcast as (B, C, 1, 1)
3. GroupNorm → GELU → 3×3 conv
4. skip: identity, or a 1×1 conv if the channel count changes

The overall shape of the U-Net is:

The input starts as 28x28x1, and then as we go through the model, we first downsample in space while upsampling in channels, then start upsampling in space and downsampling in channels to get back a 28x28x1 output. Features from the early layers are also concatenated and passed as input to later layers (the arrows above called "skip connections")

Time embedding

The network has to behave differently at t=0 (almost clean) and t=99 (almost pure noise). [SinusoidalEmbedding](unet.py) (L8–19) turns the integer timestep into a 32-length vector of sines and cosines. A two-layer MLP in [SimpleUNet.__init__](unet.py) (L63–68) expands that to time_dim=128. Every residual block then adds a projected copy of this vector onto its feature maps (L34).

Class-conditional generation is the same trick: an nn.Embedding(10, time_dim) is added to the time embedding (unet.py L69, L91–93). Digit identity rides along the same pathway as t without any extra parameters.

Constructor: unet.py L71–88. Forward with skips: unet.py L95–107.

• Down: Conv2d(stride=2) ([Downsample](unet.py) L39–45)
• Up: nearest-neighbor ×2 then a 3×3 conv ([Upsample](unet.py) L48–55)

No attention. MNIST is 28×28; local convs are enough.

---

Sampler

The scheduler is the DDPM math. The U-Net never calls it; the training loop and sample() do.

Forward process (add_noise)

The forward process is computed in a closed form. For a random t:

xt = sqrt(ᾱ_t) * x0 + sqrt(1 - ᾱ_t) * ε,    ε ~ N(0, I)

ᾱ_t (alphas_cumprod) is precomputed in [NoiseScheduler.__init](ddpm.py)__ (L55–61). The actual mix is [add_noise](ddpm.py) (L123–126). _broadcast (L81–82) reshapes the per-timestep scalar to (B, 1, 1, 1) so it multiplies a batch of images.

The figure below is 2,000 MNIST digits projected to 2D using PCA, colored by label. Ones (orange) start as their own lobe; as ᾱ_t drops, class clouds overlap and become one Gaussian. Each panel is zoomed to its own cloud. This matches the "expanding the hills" analogy we discussed in meetings.

python plot_forward_pca.py

Same ε is reused at every t, so each point is one forward trajectory.

Reverse process (step)

For one denoising step, Ho et al. 2020:

1. From the model's noise prediction, reconstruct x0 ([reconstruct_x0](ddpm.py) L90–93)
2. Clamp that reconstruction to [-1, 1] (L113) — pixels live there; unclipped x0 blows up into high-contrast blobs
3. Posterior mean q(x_{t-1} | x_t, x0) ([q_posterior](ddpm.py) L95–98)
4. If t > 0, add posterior noise ([get_variance](ddpm.py) L100–108)

All of that is [NoiseScheduler.step](ddpm.py) L110–121.

Sampling loop

When we want to inference our model, we start with pure noise, then iteratively "denoise". Each iteration the U-Net predicts noise, and then step produces x_{t-1}. At the end, map [-1, 1] back to [0, 1] for saving.

---

Training loop

Each step of the training loop takes these key steps in ddpm.py L226–242:

ε ~ N(0, I) # Sample noise
t ~ Uniform({0, …, T-1}) # Sample a timestep
xt = add_noise(x0, ε, t) # Noise the ground truth image to the timestep level
ε̂ = model(xt, t)          # or model(xt, t, y) if conditional -- call the model to predict noise added
loss = MSE(ε̂, ε) # Compare GT and predicted noise
loss.backward() # Update the weights according to loss

That is ε-prediction DDPM. The network does not predict the clean image or the next step in the iteration directly; the scheduler converts a noise prediction into those when sampling.

Other setup is done just above the loop: Setting the device to train on (L17–22, L191), Getting the MNIST loader, Creating a SimpleUNet, NoiseScheduler, Setting up the Optimizer: AdamW at 2e-4 (L202–211).

--conditional only changes three things: num_classes=10 on the U-Net, pass labels into model(...), and sample a labeled grid. The residual blocks do not change.

Suggested Path

1. Look through the code.
2. Train unconditionally and look at the samples. Make a virtualenv first so the deps stay local to this repo:

python3 -m venv .venv
source .venv/bin/activate   # Windows: .venv\Scripts\activate
pip install -r requirements.txt
python ddpm.py --experiment_name base

This runs training for 10 epochs with batch 128 and T = 100. It should take no more than ~15 min on a laptop. During training, it will output sampled grids in exps/base/images/ (one PNG per epoch) and the loss curve in exps/base/loss.png. The MSE should flatten quickly; judge progress from the grids. Digits should start to look recognizable well before epoch 10.

1. Train conditionally and look at the labeled grid. With the venv still activated:

python ddpm.py --conditional --experiment_name cond

Same loop, but the U-Net gets digit labels. The sample grid in exps/cond/images/ has one row per digit 0–9. Check that a row of sevens actually looks like sevens.

1. Play in vis.ipynb. After you have checkpoints, the notebook is for poking at the trained models: the forward process on a real digit, the reverse process from noise, and a conditional grid if you trained --conditional. Point the notebook kernel at .venv (or source .venv/bin/activate then jupyter notebook). Load exps/base/model.pth or exps/cond/model.pth in the corresponding cells.