Module 6 — Diffusion models: adding noise, then learning to denoise

The GANs of modules 4 and 5 fight instability every step. Diffusion models solved the problem by trading a hard adversarial objective for a very stable regression objective, at the cost of slow sampling. In 2020 they were a curiosity, in 2022 they powered Stable Diffusion, and in 2026 they are the default for image, audio and even protein generation. This module builds a small one on CelebA faces cropped to 64 by 64 pixels.

The forward process: destroy the data

The forward diffusion process takes a real image $x_0$ and gradually adds Gaussian noise over $T$ steps, until the last state $x_T$ is indistinguishable from pure noise:

$q(x_t \mid x_{t-1}) = \mathcal{N}(x_t; \sqrt{1 - \beta_t}\, x_{t-1}, \beta_t I)$

The variance schedule $\beta_1, \ldots, \beta_T$ is small at the start and grows towards the end — typical values are $\beta_1 \approx 10^{-4}$ and $\beta_T \approx 0.02$ with $T = 1000$ steps. The nice property is that we can jump directly to any timestep without iterating:

$x_t = \sqrt{\bar{\alpha}_t}\, x_0 + \sqrt{1 - \bar{\alpha}_t}\, \varepsilon, \quad \varepsilon \sim \mathcal{N}(0, I)$

where $\bar{\alpha}_t = \prod_{s=1}^{t} (1 - \beta_s)$. This closed form is what makes training tractable.

import torch

def cosine_beta_schedule(T, s=0.008):
    t = torch.linspace(0, T, T + 1)
    f = torch.cos((t / T + s) / (1 + s) * torch.pi / 2) ** 2
    alphas_bar = f / f[0]
    betas = 1 - alphas_bar[1:] / alphas_bar[:-1]
    return torch.clip(betas, 1e-5, 0.999)

def q_sample(x_0, t, alphas_bar):
    eps = torch.randn_like(x_0)
    a_bar = alphas_bar[t].view(-1, 1, 1, 1)
    return a_bar.sqrt() * x_0 + (1 - a_bar).sqrt() * eps, eps

The forward process has no parameters and requires no training. It is deterministic once the schedule is fixed.

The reverse process: learn to undo one step

The generative model is a network that, given a noisy image $x_t$ and the timestep $t$, predicts the noise $\varepsilon$ that was added to reach $x_t$ from some cleaner state. The training loss is a plain squared error:

$\mathcal{L}(\theta) = \mathbb{E}_{x_0, \varepsilon, t} \Big[ \| \varepsilon - \varepsilon_\theta(x_t, t) \|^2 \Big]$

That is it. No adversarial game, no KL term, no reparameterisation subtlety — just a regression against a known target. This is why diffusion training is so stable compared to GANs.

Once trained, sampling starts from pure noise $x_T \sim \mathcal{N}(0, I)$ and takes $T$ small denoising steps back to a plausible $x_0$. Each step subtracts a scaled version of the predicted noise and adds a smaller amount of new noise — this is DDPM sampling, detailed below.

Why predict noise rather than the image?

You could also train the network to predict $x_0$ directly from $x_t$ and $t$. Mathematically the two formulations are equivalent. In practice, predicting noise gives a target of similar magnitude at every timestep: $\varepsilon$ is drawn from $\mathcal{N}(0, I)$ regardless of $t$, so the loss stays well-scaled. Predicting $x_0$ has targets that vary in magnitude across $t$, and requires per-timestep loss weighting to train evenly.

A U-Net conditioned on the timestep

The denoiser must accept the timestep $t$ as an input, because the amount of noise it has to remove depends on $t$. The standard architecture is a U-Net, borrowed from segmentation, with two additions.

Sinusoidal time embedding. The scalar $t$ is embedded into a vector via sinusoidal features of frequencies $\omega_k$, mirroring the positional encoding of transformers:

$e(t) = [\sin(\omega_1 t), \cos(\omega_1 t), \ldots, \sin(\omega_K t), \cos(\omega_K t)]$

That vector is passed through a small MLP and added to the feature maps at every U-Net block, so every layer knows which noise level it is looking at.

Skip connections. The U-Net's characteristic shape — a downsampling path, an upsampling path, and skip connections between symmetric levels — preserves fine spatial detail that a pure encoder-decoder would lose. Diffusion needs both the coarse "what is in this image" and the fine "where exactly are the edges", which is precisely what skips deliver.

import torch.nn as nn

class TimestepEmbedding(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        self.mlp = nn.Sequential(nn.Linear(dim, dim * 4), nn.SiLU(), nn.Linear(dim * 4, dim * 4))

    def forward(self, t):
        half = self.dim // 2
        freqs = torch.exp(torch.linspace(0, -8, half, device=t.device))
        args = t[:, None].float() * freqs[None, :]
        emb = torch.cat([args.sin(), args.cos()], dim=1)
        return self.mlp(emb)

A real implementation for CelebA at 64 by 64 uses roughly 60 million parameters — small by 2026 standards. It trains in a few hours on a single GPU.

The training loop

T = 1000
betas = cosine_beta_schedule(T).to(device)
alphas_bar = torch.cumprod(1 - betas, dim=0)

for epoch in range(50):
    for x_0, _ in celeba_loader:
        x_0 = x_0.to(device)
        t = torch.randint(0, T, (x_0.size(0),), device=device)
        x_t, eps = q_sample(x_0, t, alphas_bar)
        eps_pred = unet(x_t, t)
        loss = ((eps - eps_pred) ** 2).mean()
        opt.zero_grad(); loss.backward(); opt.step()

Note that each example draws a random $t$: the network trains simultaneously on every noise level. There is no schedule, no curriculum, no per-example timestep sequence — just uniform sampling of $t$.

DDPM and DDIM sampling

DDPM sampling (Ho 2020) reverses the forward process exactly. Starting from $x_T \sim \mathcal{N}(0, I)$, at each step:

$x_{t-1} = \frac{1}{\sqrt{1 - \beta_t}} \left( x_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \varepsilon_\theta(x_t, t) \right) + \sigma_t \, z, \quad z \sim \mathcal{N}(0, I)$

with $\sigma_t^2 = \beta_t$. Running this for $T = 1000$ steps produces a good sample. The problem is obvious: 1000 forward passes for one image. On a GPU that is seconds per image, not milliseconds.

DDIM sampling (Song 2020) exploits the fact that the exact reverse chain is not the only one whose marginal distributions match the forward chain. A deterministic version — $\sigma_t = 0$ — allows large jumps: 25 or 50 steps often suffice for near-identical quality. DDIM is deterministic given the initial noise, which is a bonus: identical noise yields identical output, useful for reproducibility.

@torch.no_grad()
def ddim_sample(unet, shape, n_steps=50):
    x = torch.randn(shape, device=device)
    step_indices = torch.linspace(T - 1, 0, n_steps + 1).long()
    for i in range(n_steps):
        t = step_indices[i].to(device).repeat(shape[0])
        t_prev = step_indices[i + 1].to(device).repeat(shape[0])
        a_t = alphas_bar[t].view(-1, 1, 1, 1)
        a_prev = alphas_bar[t_prev].view(-1, 1, 1, 1)
        eps = unet(x, t)
        x_0 = (x - (1 - a_t).sqrt() * eps) / a_t.sqrt()
        x = a_prev.sqrt() * x_0 + (1 - a_prev).sqrt() * eps
    return x

The role of $t$ is central, and it is easy to get wrong

The network's behaviour depends critically on the timestep it receives. Passing the wrong $t$ — off by one, or a scalar instead of a batch, or floats instead of ints — silently produces bad samples with no exception raised. When debugging diffusion, print the shape and range of $t$ at every call. It is the most common source of nonsense outputs.

Latent diffusion, which is what Stable Diffusion is

Running diffusion in pixel space at 512 by 512 is expensive: 262,144 pixels per image, hundreds of steps, tens of millions of parameters. Latent diffusion (Rombach 2021) reduces the cost by running the entire diffusion process in the latent space of a pretrained autoencoder. The autoencoder maps a $512 \times 512 \times 3$ image to a $64 \times 64 \times 4$ latent, and the U-Net runs on those latents. Perceptual quality is preserved, memory and compute drop by an order of magnitude, and the door opens to text conditioning at scale — the subject of module 7.

Stable Diffusion is a latent diffusion model with a text encoder attached. Everything else in module 7 builds on top of what we have here.

In summary

• The forward process deterministically adds Gaussian noise to $x_0$ over $T$ steps; there is a closed form to jump to any $t$, which is what makes training tractable.
• Training predicts the noise $\varepsilon$ added to reach $x_t$, from a squared-error loss on random $t$; there is no adversarial game, no likelihood computation, and training is exceptionally stable.
• The denoiser is a U-Net conditioned on the timestep via sinusoidal embeddings; sampling with DDPM takes $T$ steps, DDIM cuts it to 25 to 50 with near-identical quality.
• Latent diffusion runs the process in a pretrained autoencoder's latent space, which is the design of Stable Diffusion and the substrate for text conditioning in module 7.

Next module: text conditioning — how a diffusion model learns to obey a prompt like "a small dog wearing sunglasses" and what classifier-free guidance actually does.