Diffusion models — the mathematics

Diffusion rests on an imbalance: destroying an image is trivial, recreating it is hard. So the destruction is defined exactly and the model is taught to reverse it, one small step at a time.

x₀ ──noise──→ x₁ ──noise──→ ... ──noise──→ x_T ≈ pure noise
   ←─denoise──   ←─denoise──    ←─denoise──

Three things make the construction practical:

1. The forward process has a closed form — $x_t$ can be computed directly from $x_0$ without iterating.
2. The training objective reduces to a simple MSE on the added noise.
3. Every backward step is small, so a Gaussian approximation is enough.

It is the second that is surprising: after a fairly long derivation you are left with mse_loss(the_model's_guess, the_actual_noise) — and nothing more.

The forward process is defined as a Markov chain:

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

with a small noise schedule $\beta_1 < \dots < \beta_T$ (typically linear from $10^{-4}$ to $0.02$ over $T = 1000$ steps).

The closed form. Set $\alpha_t = 1-\beta_t$ and $\bar\alpha_t = \prod_{s\le t}\alpha_s$. Since the sum of two independent Gaussian variables is Gaussian with the variances added, the steps can be merged:

$q(x_t \mid x_0) = \mathcal{N}\!\left(x_t;\ \sqrt{\bar\alpha_t}\,x_0,\ (1-\bar\alpha_t)I\right)$ $\Longrightarrow\quad x_t = \sqrt{\bar\alpha_t}\,x_0 + \sqrt{1-\bar\alpha_t}\,\varepsilon, \qquad \varepsilon \sim \mathcal{N}(0, I)$

This is the key to the training being practical: you can jump straight to a random $t$ without simulating the chain.

The backward step. The true posterior is intractable, but conditioned on $x_0$ it is Gaussian in closed form:

$q(x_{t-1}\mid x_t, x_0) = \mathcal{N}\!\left(x_{t-1};\ \tilde\mu_t(x_t, x_0),\ \tilde\beta_t I\right)$

$\tilde\mu_t = \frac{\sqrt{\bar\alpha_{t-1}}\beta_t}{1-\bar\alpha_t}x_0 + \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})}{1-\bar\alpha_t}x_t, \qquad \tilde\beta_t = \frac{1-\bar\alpha_{t-1}}{1-\bar\alpha_t}\beta_t$

The simplification. The variational bound (the ELBO) gives a sum of KL terms between Gaussian distributions. Ho et al. (2020) showed that if the model is parameterised to predict the noise instead of the mean, and the weighting factors are then dropped, the whole objective reduces to

$\boxed{\mathcal{L}_{\text{simple}} = \mathbb{E}_{x_0,\varepsilon,t}\left[\left\|\varepsilon - \varepsilon_\theta\!\left(\sqrt{\bar\alpha_t}x_0 + \sqrt{1-\bar\alpha_t}\varepsilon,\ t\right)\right\|^2\right]}$

The dropped weights mean the objective is no longer an exact ELBO — but in practice it trains better, since the weighting otherwise gives the noisiest steps disproportionate weight.

The sampling step then becomes:

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

The connection to score-based models: $\nabla_{x}\log q(x_t) = -\varepsilon_\theta(x_t,t)/\sqrt{1-\bar\alpha_t}$. Predicting the noise is therefore the same thing as estimating the score function, and DDPM and score matching are two views of the same model. Song et al. (2021) showed that both are discretisations of a stochastic differential equation, which explains why different samplers can be swapped freely after training.

import torch, torch.nn.functional as F

T = 1000
beta = torch.linspace(1e-4, 0.02, T)
alpha = 1.0 - beta
alpha_bar = torch.cumprod(alpha, dim=0)

def q_sample(x0, t, noise=None):
    """The closed form: jump straight to step t without iterating."""
    noise = torch.randn_like(x0) if noise is None else noise
    a = alpha_bar[t].view(-1, 1, 1, 1)
    return a.sqrt() * x0 + (1 - a).sqrt() * noise, noise

def loss(model, x0):
    t = torch.randint(0, T, (x0.size(0),), device=x0.device)
    xt, noise = q_sample(x0, t)
    return F.mse_loss(model(xt, t), noise)        # the whole training objective

@torch.no_grad()
def sample(model, shape, device="cpu"):
    x = torch.randn(shape, device=device)
    for t in reversed(range(T)):
        tt = torch.full((shape[0],), t, device=device, dtype=torch.long)
        eps = model(x, tt)
        a, ab = alpha[t], alpha_bar[t]
        mean = (x - beta[t] / (1 - ab).sqrt() * eps) / a.sqrt()
        x = mean + beta[t].sqrt() * torch.randn_like(x) if t > 0 else mean
    return x

# Check the closed form numerically against the step-by-step definition
torch.manual_seed(0)
x0 = torch.randn(1, 1, 8, 8)
t_target = 200

# Iteratively
x = x0.clone()
for t in range(t_target + 1):
    x = alpha[t].sqrt() * x + beta[t].sqrt() * torch.randn_like(x)
var_iterative = float(x.var())

# The closed form — 4000 samples to compare the distributions
samples = torch.stack([q_sample(x0, torch.tensor([t_target]))[0] for _ in range(4000)])
print(f"theoretical variance: {float(1 - alpha_bar[t_target]):.4f}")
print(f"empirical variance:   {float((samples - alpha_bar[t_target].sqrt() * x0).var()):.4f}")

# The noise schedule: how much signal remains at each step
for t in (0, 100, 300, 500, 800, 999):
    print(f"  t={t:>4}: signal {float(alpha_bar[t].sqrt()):.4f}  "
          f"noise {float((1 - alpha_bar[t]).sqrt()):.4f}")
#   t=   0: signal 0.9999  noise 0.0100
#   t= 500: signal 0.3403  noise 0.9403
#   t= 999: signal 0.0063  noise 1.0000    ← practically pure noise

# The cosine schedule (Nichol & Dhariwal): keeps more signal for longer
def alpha_bar_cosine(t, T=1000, s=0.008):
    import math
    f = lambda u: math.cos((u / T + s) / (1 + s) * math.pi / 2) ** 2
    return f(t) / f(0)

for t in (0, 300, 500, 800):
    print(f"  t={t:>3}: linear {float(alpha_bar[t]):.4f}  "
          f"cosine {alpha_bar_cosine(t):.4f}")