Annotated Diffusion Model – detailed walkthrough of DDPM implementation

TL;DR

Hugging Face published a complete, annotated PyTorch notebook that implements the original Denoising Diffusion Probabilistic Model (DDPM) – a neural network that learns to denoise images from pure Gaussian noise – and demonstrates training on low‑resolution datasets and sampling high‑quality images. This resource demystifies the mathematics, network architecture, training loop, and inference procedure, making diffusion models accessible to practitioners.

What a diffusion model does (key takeaway)

A diffusion model learns a reverse denoising process that transforms random Gaussian noise into realistic data; the forward process adds noise in a predefined schedule, while a U‑Net‑style neural network predicts the added noise at each timestep, enabling image generation after a fixed number of steps.

Forward diffusion: adding noise with a variance schedule

Conclusion: The forward process is a closed‑form Gaussian transition that can sample any noisy timestep directly, eliminating the need for iterative noise addition.

The forward diffusion distribution is

q(x_t | x_{t-1}) = N(x_t ; sqrt(1-β_t)·x_{t-1}, β_t·I)

where the variance schedule (β_1,…,β_T) is monotonic (e.g., linear, cosine, quadratic, or sigmoid). By repeatedly applying this transition, the marginal

q(x_t | x_0) = N(x_t ; sqrt(\bar α_t)·x_0, (1-\bar α_t)·I)

can be sampled directly using pre‑computed (\bar α_t = \prod_{s=1}^t (1-β_s)). This “nice property” lets the training algorithm draw a random timestep (t) and corrupt a real image (x_0) in one step.

Reverse diffusion: learning the mean via noise prediction

Conclusion: The model is trained to predict the exact Gaussian noise (ε) added at timestep (t); the predicted noise is then used to compute the mean of the reverse Gaussian distribution.

The reverse conditional is assumed Gaussian:

p_θ(x_{t-1} | x_t) = N(x_{t-1} ; μ_θ(x_t, t), σ_t^2·I)

DDPM fixes (σ_t^2) to the known forward variance and learns only (μ_θ). By re‑parameterising the mean in terms of the noise predictor (ε_θ):

μ_θ(x_t, t) = (1/√α_t)·(x_t - (β_t/√(1-\bar α_t))·ε_θ(x_t, t))

the training loss reduces to a simple mean‑squared error between the true noise (ε) and the network output (ε_θ):

L_t = || ε - ε_θ(x_t, t) ||^2

Randomly sampling (t) each batch yields an unbiased estimator of the variational lower bound.

Network architecture: a conditional U‑Net

Conclusion: The DDPM uses a time‑conditioned U‑Net with sinusoidal position embeddings, residual blocks, attention, and group normalization to predict noise at every spatial resolution.

Key components:

  • Sinusoidal position embeddings encode the timestep (t) and are injected via a small MLP into each ResNet block.
  • Weight‑standardized convolutions improve training stability when combined with group norm.
  • ResNet blocks (two convolution‑norm‑SiLU layers with optional scale‑shift from the time embedding) provide the main feature transformation.
  • Attention modules (either full multi‑head or linear attention) capture long‑range dependencies.
  • Group normalization is applied before attention (PreNorm).
  • Down/upsampling paths halve or double spatial resolution while preserving channel depth, mirroring the classic U‑Net design.
  • Final head maps the concatenated bottleneck features back to the image shape.

The full PyTorch class Unet assembles these pieces, accepts a noisy image tensor and a timestep tensor, and returns a noise tensor of identical shape.

Training loop: stochastic timestep sampling and Huber loss

Conclusion: Training proceeds by sampling a random timestep per batch, corrupting the input with the closed‑form forward process, and minimizing a Huber loss between true and predicted noise.

Pseudo‑code (simplified):

for epoch in range(num_epochs):
    for batch in dataloader:
        t = torch.randint(0, T, (batch_size,)).to(device)          # uniform timestep
        loss = p_losses(model, batch, t, loss_type='huber')        # MSE/Huber on noise
        loss.backward()
        optimizer.step()

The helper p_losses calls q_sample to obtain (x_t) and then computes the loss against the network output model(x_t, t). Periodic sampling (using the reverse process) visualizes training progress.

Sampling (inference): reversing the diffusion chain

Conclusion: Generation starts from pure Gaussian noise and iteratively applies the learned denoising step; each step uses the predicted noise to compute the posterior mean and adds calibrated Gaussian noise according to the known variance schedule.

Algorithm (Algorithm 2 in the DDPM paper):

x_T = torch.randn(shape)                     # start from noise
for t in reversed(range(T)):
    pred_noise = model(x_t, t)
    mean = (1/√α_t) * (x_t - β_t/√(1-\bar α_t) * pred_noise)
    if t > 0:
        x_{t-1} = mean + √posterior_variance_t * torch.randn_like(x_t)
    else:
        x_0 = mean

The provided notebook implements this loop in p_sample_loop and visualises the denoising trajectory as a GIF.

End‑to‑end example on Fashion‑MNIST

Conclusion: The tutorial trains a DDPM on the 28×28 Fashion‑MNIST dataset, achieving recognizable clothing items after a few thousand training steps.

  • Data loading uses the 🤗 datasets library with on‑the‑fly transforms (random horizontal flip, scaling to ([-1,1])).
  • Model hyper‑parameters: dim=image_size, dim_mults=(1,2,4), channels=1 for grayscale images.
  • Optimizer: Adam with learning rate 1e‑3.
  • Training runs for 6 epochs; loss quickly drops below 0.05.
  • Sampling after training produces clear T‑shirt shapes, confirming that the implementation works on low‑resolution data.

How this fits into the diffusion literature

Conclusion: The annotated implementation reproduces the original DDPM (Ho et al., 2020) and serves as a foundation for many later advances such as variance‑learning (Nichol et al., 2021), cascaded diffusion (Ho et al., 2021), classifier‑free guidance, and large‑scale text‑to‑image models (DALL‑E 2, Imagen).

Key follow‑up works listed in the blog post:

  • Improved DDPM – learns both mean and variance, improving sample quality.
  • Cascaded Diffusion Models – stack multiple diffusion models for high‑resolution synthesis.
  • Diffusion Models Beat GANs – demonstrates superior FID scores with architectural tweaks and classifier guidance.
  • Classifier‑Free Guidance – removes the need for an external classifier during conditional generation.
  • DALL‑E 2 & ImageGen – combine diffusion with CLIP embeddings or large language models for text‑conditional image synthesis.

The main drawback remains the need for many denoising steps (typically 1000), though recent research (e.g., high‑order solvers) reduces this to as few as 10 steps.

Practical takeaways for developers

  • The notebook provides a ready‑to‑run reference implementation; replace the dataset, increase T, or swap the U‑Net for a larger backbone to scale up.
  • Adjust the variance schedule (linear_beta_schedule, cosine_beta_schedule, etc.) to trade off sample fidelity versus speed.
  • Use the Huber loss for robustness; the code also supports L1/L2.
  • Periodic sampling (save_and_sample_every) is essential for monitoring mode collapse or training instability.
  • For faster inference, consider implementing DDIM or stochastic sampler variants, or adopt recent “few‑step” solvers.

Sources