Diffusion Models: Foundations

Diffusion models now dominate generative AI for images, audio, and video — DALL-E 2, Stable Diffusion, Imagen, and Sora are all built on this foundation. Understanding why diffusion beat GANs requires engaging with the actual math, not just the high-level narrative. This file builds the derivation from scratch: forward process, reparameterization trick, reverse process, loss derivation, and the UNet backbone. The goal is that after working through this, you can implement DDPM (arXiv:2006.11239) from scratch on MNIST without referencing anyone else’s code.

---

1. Why Diffusion Won Over GANs

This is not triumphal storytelling. It is a concrete list of failure modes that diffusion addresses:

Problem	GANs	Diffusion
Training stability	Adversarial minimax → frequent collapse	MSE loss on noise → stable gradient signal
Mode collapse	Generator ignores modes the discriminator ignores	Full distribution coverage via denoising chain
Likelihood evaluation	Not tractable (no encoder)	Tractable via variational lower bound
Sample diversity	Low — tuned for realism over coverage	High — stochastic reverse process explores modes
Sample quality (FID)	Competitive on single domains	SOTA across text-to-image, audio, video
Architecture coupling	Generator + discriminator must stay balanced	Single network, no adversary

The tradeoff diffusion loses on: inference speed. GANs generate in one forward pass. Vanilla DDPM requires 1000 denoising steps. DDIM (arXiv:2010.02502) reduced this to 50 without retraining. Flow matching (arXiv:2210.02747) reduces it further, and is the backbone of Stable Diffusion 3.

---

2. The Forward Process (Noising)

Define a Markov chain that gradually adds Gaussian noise to a data sample x₀ over T steps.

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

Where:

• β_t is a noise schedule — a sequence of small positive numbers (e.g., linear β₁=1e-4 to β_T=0.02, or cosine schedule from arXiv:2102.09672 which avoids overshooting at the end)

• At each step t, we keep √(1-β_t) of the signal and add √(β_t) of Gaussian noise

• After T=1000 steps: x_T ≈ N(0, I) — pure noise, no information about x₀ remaining

The signal term √(1-β_t) shrinks the mean toward zero. The noise term √(β_t) maintains unit variance. This is a variance-preserving process.

The Reparameterization Trick — Critical for Training Efficiency

Iterating T steps of the chain to sample x_t from x₀ would make training O(T) per step. Instead, we can sample x_t at any arbitrary timestep t directly from x₀:

Define:

αt = 1 - βt
ᾱt = ∏_{i=1}^{t} αi    (cumulative product of alphas)

Then:

q(x_t | x_0) = N(x_t; √ᾱ_t · x_0, (1-ᾱ_t) · I)

Which means we can sample directly:

x_t = √ᾱ_t · x_0 + √(1-ᾱ_t) · ε,    where ε ~ N(0, I)

This is the single most important equation for understanding DDPM training. It means:

• At t=0: x_t ≈ x_0 (clean image, ᾱ_0 ≈ 1)

• At t=T: x_t ≈ ε (pure noise, ᾱ_T ≈ 0)

• Training can sample any t uniformly and compute x_t in O(1) time

---

3. The Reverse Process (Denoising)

The reverse process undoes the noising. The true posterior q(x_{t-1} | x_t) is intractable — it requires knowing the full data distribution. So we learn an approximation:

p_θ(x_{t-1} | x_t) = N(x_{t-1}; μ_θ(x_t, t), Σ_θ(x_t, t))

Ho et al. (DDPM, 2020) showed empirically that predicting the noise ε rather than the mean directly works best. The parameterization:

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

So the model ε_θ(x_t, t) predicts the noise ε that was added to x_0 to produce x_t.

Why Predict Noise, Not x_0?

The original DDPM paper tried both (Section 4, ablation). Noise prediction (ε-prediction) consistently outperformed x_0 prediction. The intuition: predicting x_0 directly requires the model to simultaneously handle all scales of structure (coarse composition + fine texture) at every timestep. Noise prediction decomposes this — at high t (heavy noise), the model learns coarse structure; at low t (light noise), it refines fine details.

---

4. The Training Loss

The full ELBO derivation telescopes into a surprisingly clean objective. After simplification (Ho et al. Section 3.4):

L_simple = E_{t, x_0, ε} [ ||ε - ε_θ(x_t, t)||² ]

That’s it. MSE between the actual noise ε and the model’s predicted noise. The timestep t is sampled uniformly from [1, T]. This is the entire training loss.

The full derivation (skipping for brevity, but you should read it once):

1. Start with ELBO: log p_θ(x_0) ≥ E_q[log p_θ(x_0|x_1)] - KL terms

2. Rewrite KL divergences as denoising score matching objectives

3. Weight all terms equally (the “L_simple” simplification — this is a departure from the mathematically correct weighting but works better empirically)

---

5. The UNet Backbone

The model ε_θ(x_t, t) needs to:

1. Accept a noisy image x_t as input

2. Know the current noise level t

3. Return a noise estimate of the same spatial dimensions

A UNet satisfies all three requirements naturally.

UNet Architecture for Diffusion

Input: x_t (noisy image, e.g., 1×28×28 for MNIST)
       t (timestep scalar, embedded → 256-dim vector)

Encoder (downsampling path):
  ResBlock(64) → MaxPool → [skip connection]
  ResBlock(128) → MaxPool → [skip connection]
  ResBlock(256) → MaxPool → [skip connection]

Bottleneck:
  ResBlock(512)
  [Optional: Self-Attention for global context]

Decoder (upsampling path):
  Upsample → Concat(skip) → ResBlock(256)
  Upsample → Concat(skip) → ResBlock(128)
  Upsample → Concat(skip) → ResBlock(64)

Output: Conv(1) → ε̂ (predicted noise, same shape as input)

How t is injected: Sinusoidal embedding of t (same formula as transformer positional encoding) → MLP → added to the feature maps inside each ResBlock. This is how the network knows which noise level it’s operating at. Without this, the same UNet weights must handle radically different noise levels — the network would fail.

For text-conditional generation (Stable Diffusion): Cross-attention layers are inserted at each decoder block. The query comes from the image features; key/value come from CLIP text embeddings. This is why the UNet in Stable Diffusion is much heavier than a vanilla DDPM UNet.

---

6. Sampling (Inference)

Given a trained ε_θ, generate a new sample:

# DDPM sampling (Algorithm 2 from Ho et al. 2020)
x_T ~ N(0, I)
for t in reversed(range(1, T+1)):
    z ~ N(0, I)  if t > 1  else  z = 0
    eps_pred = model(x_t, t)
    x_{t-1} = (1/sqrt(alpha_t)) * (x_t - (beta_t/sqrt(1-alphabar_t)) * eps_pred)
               + sqrt(beta_t) * z
return x_0

This runs T=1000 UNet forward passes. For MNIST this takes seconds on CPU. For 512×512 images it takes minutes — hence why DDIM matters.

---

7. DDIM: Deterministic Fast Sampling (arXiv:2010.02502)

DDIM rewrites the reverse process as a deterministic ODE instead of a stochastic SDE. The key insight: if you remove the stochastic term (the z in the sampling loop), the process still converges to a valid sample — and you can take bigger steps.

Result: same trained model weights, but sample in 50 steps instead of 1000. Quality barely degrades. This is a pure inference trick — no retraining needed.

DDIM also enables latent space interpolation and image editing because the forward process (encoding) is now deterministic: given x_0, you can find the exact x_T that maps back to it.

---

8. Latent Diffusion / Stable Diffusion (arXiv:2112.10752)

Running DDPM at 512×512 pixel resolution is expensive: each of 1000 UNet passes operates on a 512×512×3 tensor. The fix: diffuse in latent space.

Architecture:
1. VAE Encoder: 512×512×3 → 64×64×4 (8× spatial compression, ×3→4 channels from KL-VAE)
2. DDPM in latent space: denoise 64×64×4 tensors (64× cheaper per step)
3. VAE Decoder: 64×64×4 → 512×512×3

The VAE is pretrained separately and frozen during diffusion training. The CLIP text encoder is also frozen. Only the UNet is trained.

Why this works: The VAE learns a perceptually rich latent space where semantic structure is preserved but pixel-level noise is compressed away. Diffusion in this space learns to generate semantically coherent latents, not pixel noise.

---

9. Code: DDPM on MNIST

This is a minimal but complete DDPM implementation. Target: after ~50 epochs on a GPU (or overnight on CPU), generated samples are visually recognizable as handwritten digits.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np

# ─── Noise Schedule ───────────────────────────────────────────────────────────
def make_beta_schedule(T=1000, beta_start=1e-4, beta_end=0.02):
    """Linear schedule from Ho et al. 2020."""
    betas = torch.linspace(beta_start, beta_end, T)
    alphas = 1.0 - betas
    alphabar = torch.cumprod(alphas, dim=0)
    return betas, alphas, alphabar

T = 1000
betas, alphas, alphabar = make_beta_schedule(T)

# Precompute for sampling
sqrt_alphabar = torch.sqrt(alphabar)
sqrt_one_minus_alphabar = torch.sqrt(1.0 - alphabar)
sqrt_recip_alphas = torch.sqrt(1.0 / alphas)
betas_over_sqrt_one_minus_alphabar = betas / sqrt_one_minus_alphabar

# ─── Forward Process ──────────────────────────────────────────────────────────
def q_sample(x0, t, noise=None):
    """
    Sample x_t from x_0 directly using reparameterization:
        x_t = sqrt(alphabar_t) * x0 + sqrt(1 - alphabar_t) * eps
    """
    if noise is None:
        noise = torch.randn_like(x0)
    sqrt_ab = sqrt_alphabar[t].view(-1, 1, 1, 1).to(x0.device)
    sqrt_one_minus_ab = sqrt_one_minus_alphabar[t].view(-1, 1, 1, 1).to(x0.device)
    return sqrt_ab * x0 + sqrt_one_minus_ab * noise

# ─── Sinusoidal Time Embedding ────────────────────────────────────────────────
class SinusoidalTimeEmbedding(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, t):
        device = t.device
        half = self.dim // 2
        freqs = torch.exp(
            -np.log(10000) * torch.arange(half, device=device) / half
        )
        args = t[:, None].float() * freqs[None]
        return torch.cat([args.sin(), args.cos()], dim=-1)

# ─── Residual Block ───────────────────────────────────────────────────────────
class ResBlock(nn.Module):
    def __init__(self, in_ch, out_ch, time_dim):
        super().__init__()
        self.norm1 = nn.GroupNorm(8, in_ch)
        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
        self.norm2 = nn.GroupNorm(8, out_ch)
        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
        self.time_proj = nn.Linear(time_dim, out_ch)
        self.skip = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()

    def forward(self, x, t_emb):
        h = self.conv1(F.silu(self.norm1(x)))
        h = h + self.time_proj(F.silu(t_emb))[:, :, None, None]
        h = self.conv2(F.silu(self.norm2(h)))
        return h + self.skip(x)

# ─── UNet ─────────────────────────────────────────────────────────────────────
class UNet(nn.Module):
    """Minimal UNet for MNIST (1×28×28)."""
    def __init__(self, time_dim=256):
        super().__init__()
        self.time_emb = nn.Sequential(
            SinusoidalTimeEmbedding(time_dim),
            nn.Linear(time_dim, time_dim),
            nn.SiLU(),
        )
        # Encoder
        self.down1 = ResBlock(1, 64, time_dim)
        self.down2 = ResBlock(64, 128, time_dim)
        self.pool = nn.MaxPool2d(2)
        # Bottleneck
        self.bot = ResBlock(128, 256, time_dim)
        # Decoder
        self.up1 = ResBlock(256 + 128, 128, time_dim)
        self.up2 = ResBlock(128 + 64, 64, time_dim)
        self.upsample = nn.Upsample(scale_factor=2, mode='nearest')
        self.out_conv = nn.Conv2d(64, 1, 1)

    def forward(self, x, t):
        t_emb = self.time_emb(t)
        # Encoder
        x1 = self.down1(x, t_emb)          # 1×28×28 → 64×28×28
        x2 = self.down2(self.pool(x1), t_emb)  # 64×14×14 → 128×14×14
        # Bottleneck
        x3 = self.bot(self.pool(x2), t_emb)    # 128×7×7 → 256×7×7
        # Decoder
        x = self.upsample(x3)               # 256×14×14
        x = self.up1(torch.cat([x, x2], dim=1), t_emb)  # (256+128)×14×14 → 128×14×14
        x = self.upsample(x)                # 128×28×28
        x = self.up2(torch.cat([x, x1], dim=1), t_emb)  # (128+64)×28×28 → 64×28×28
        return self.out_conv(x)             # 64×28×28 → 1×28×28

# ─── Training Loop ────────────────────────────────────────────────────────────
def train_ddpm(epochs=50, batch_size=128, lr=2e-4, device='cuda'):
    dataset = datasets.MNIST(
        './data', train=True, download=True,
        transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5,), (0.5,))  # scale to [-1, 1]
        ])
    )
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=4)

    model = UNet().to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)

    for epoch in range(epochs):
        total_loss = 0.0
        for x0, _ in loader:
            x0 = x0.to(device)
            # Sample random timesteps
            t = torch.randint(0, T, (x0.shape[0],), device=device).long()
            # Sample noise
            noise = torch.randn_like(x0)
            # Forward process: produce x_t
            x_t = q_sample(x0, t, noise)
            # Predict noise
            noise_pred = model(x_t, t)
            # L_simple: MSE
            loss = F.mse_loss(noise_pred, noise)
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            total_loss += loss.item()
        print(f"Epoch {epoch+1}/{epochs} | Loss: {total_loss/len(loader):.4f}")
    return model

# ─── Sampling ─────────────────────────────────────────────────────────────────
@torch.no_grad()
def sample_ddpm(model, n_samples=16, device='cuda'):
    """DDPM reverse process sampling."""
    model.eval()
    x = torch.randn(n_samples, 1, 28, 28, device=device)

    betas_ = betas.to(device)
    alphabar_ = alphabar.to(device)
    sqrt_recip_alphas_ = sqrt_recip_alphas.to(device)
    betas_over_sqrt_one_minus_ab = betas_over_sqrt_one_minus_alphabar.to(device)

    for t_idx in reversed(range(0, T)):
        t_tensor = torch.full((n_samples,), t_idx, device=device, dtype=torch.long)
        eps_pred = model(x, t_tensor)

        # Compute mean (no noise at final step)
        mean = sqrt_recip_alphas_[t_idx] * (
            x - betas_over_sqrt_one_minus_ab[t_idx] * eps_pred
        )
        if t_idx > 0:
            z = torch.randn_like(x)
            x = mean + torch.sqrt(betas_[t_idx]) * z
        else:
            x = mean

    return x.clamp(-1, 1)

# ─── Visualize Forward Process ────────────────────────────────────────────────
def visualize_forward_process():
    """Show how noise is added across timesteps."""
    import matplotlib.pyplot as plt
    dataset = datasets.MNIST('./data', train=True, download=True,
                             transform=transforms.Compose([
                                 transforms.ToTensor(),
                                 transforms.Normalize((0.5,), (0.5,))
                             ]))
    x0 = dataset[0][0].unsqueeze(0)  # 1×1×28×28

    fig, axes = plt.subplots(1, 8, figsize=(16, 2))
    timesteps = [0, 100, 200, 300, 500, 700, 900, 999]
    for ax, t in zip(axes, timesteps):
        t_tensor = torch.tensor([t])
        x_t = q_sample(x0, t_tensor)
        ax.imshow(x_t[0, 0].numpy(), cmap='gray', vmin=-1, vmax=1)
        ax.set_title(f't={t}')
        ax.axis('off')
    plt.tight_layout()
    plt.savefig('forward_process.png', dpi=100)
    print("Saved: forward_process.png")

# ─── Main ─────────────────────────────────────────────────────────────────────
if __name__ == '__main__':
    import matplotlib.pyplot as plt

    visualize_forward_process()

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"Training on: {device}")

    model = train_ddpm(epochs=50, device=device)
    torch.save(model.state_dict(), 'ddpm_mnist.pt')

    samples = sample_ddpm(model, n_samples=64, device=device)
    samples = (samples + 1) / 2  # scale to [0, 1]

    fig, axes = plt.subplots(8, 8, figsize=(8, 8))
    for i, ax in enumerate(axes.flat):
        ax.imshow(samples[i, 0].cpu().numpy(), cmap='gray')
        ax.axis('off')
    plt.tight_layout()
    plt.savefig('ddpm_mnist_samples.png', dpi=100)
    print("Generated samples saved: ddpm_mnist_samples.png")

Expected outcome: After 50 epochs on MNIST (CPU: ~2 hours, GPU: ~20 minutes), samples should be visually recognizable as digits. Loss typically starts around 0.08 and converges below 0.015.

Acceptance check: If a naive observer (non-ML person) looking at ddpm_mnist_samples.png identifies the majority as handwritten digit-like shapes, the implementation is correct. This is the exit criterion for this phase.

---

10. What Most People Get Wrong

The UNet does not generate images. It predicts noise. This distinction matters enormously when you’re debugging. If your samples look like noise, the UNet learned wrong. If your samples look blurry but structured, the noise schedule is wrong (likely ᾱ_T never reaches near zero). If samples look like a single averaged digit, the stochastic term z was removed or zeroed — you collapsed the distribution.

The loss going down does not mean samples look good. L_simple is a denoising score matching loss. A model can achieve low loss but poor sample quality by learning to predict the rough direction of noise but not its fine-grained structure. Always visually inspect intermediate samples during training (every 5 epochs).

DDIM is not a different model. It is a different sampler on the same weights. If you train a DDPM model and want 50-step inference, switch to DDIM at inference time — no retraining.

---

Key Papers

Paper	arXiv	Year	What It Contributes
DDPM	2006.11239	2020	Foundation: forward/reverse process, L_simple
DDIM	2010.02502	2020	Deterministic 50-step sampling
Improved DDPM	2102.09672	2021	Cosine noise schedule, learned Σ
Latent Diffusion (Stable Diffusion)	2112.10752	2021	Diffusion in VAE latent space
Classifier-free guidance	2207.12598	2022	Conditional generation without a classifier
Flow Matching	2210.02747	2022	ODE-based generation, Stable Diffusion 3 backbone

---

Return to [README.md] · Previous: [03_vision_transformers.md] · Next: [05_graph_neural_networks.md]