Denoising diffusion

What it is and when you reach for it

Denoising diffusion probabilistic models (DDPM, Ho et al. 2020) are likelihood-based generative models that sit next to VAEs, GANs, and autoregressive models in the family tree, and since about 2021 they have won most of the territory for continuous data. Against GANs they trade single-step sampling for training stability and mode coverage: there is no adversary, no collapse, just a regression loss that goes down. Against VAEs they trade a one-shot latent decode for an iterative refinement that reaches far higher sample fidelity. Against autoregressive models they generate the whole signal in parallel at each step rather than one element at a time. The price is inference cost, hundreds of network evaluations in the original formulation, and a large slice of the field since 2020 (DDIM, distillation, consistency models, flow matching) is the project of paying that price down. You reach for diffusion when the data is continuous, the distribution is multimodal, and you care about sample quality and diversity more than about millisecond latency: images, video, audio, molecular conformations, and robot action trajectories are all in that regime.

The math

The forward process, and its closed form

Fix a horizon T (1000 in the original paper) and a schedule of small variances β₁..β_T. The forward process is a Markov chain that mixes a little fresh noise into the data at every step:

q(x_t | x_t−1) = N(x_t; √(1−β_t) x_t−1, β_tI).

The √(1−β_t) shrink is what makes the chain a variance-preserving one: if x_t−1 has unit variance, then x_t has variance (1−β_t) + β_t = 1, so the signal decays while the total scale stays fixed, and the chain converges to a standard normal rather than blowing up. Writing α_t = 1−β_t and ᾱ_t = α₁α₂···α_t, a composition of Gaussians collapses the whole chain into one jump. Two steps compose as x_t = √α_t(√α_t−1x_t−2 + √(1−α_t−1)ε′) + √(1−α_t)ε, and the two independent noise terms add in variance, α_t(1−α_t−1) + (1−α_t) = 1−α_tα_t−1. Induction gives the identity the entire training algorithm rests on:

q(x_t | x₀) = N(x_t; √ᾱ_t x₀, (1−ᾱ_t)I), equivalently x_t = √ᾱ_t x₀ + √(1−ᾱ_t) ε with ε ~ N(0, I).

Any noise level is one sample away from clean data: no simulation of the chain, just interpolate between the data and a fresh Gaussian with coefficients read off a precomputed table. A worked number: at whatever step gives ᾱ_t = 0.5, the noisy sample is 0.707·x₀ + 0.707·ε, equal parts signal and noise. And at the far end of the standard linear schedule (β from 10⁻⁴ to 0.02 over 1000 steps), ᾱ₁₀₀₀ ≈ 4×10⁻⁵, so √ᾱ ≈ 0.006: the surviving signal is under one percent of a standard deviation, and x_T is indistinguishable from pure noise, which is exactly what lets sampling start from N(0, I).

forward q (fixed):   x0 ─► x1 ─► x2 ─► ... ─► xT ≈ N(0, I)
                        add β1      β2              noise wins

reverse p_θ (learned): x0 ◄─ x1 ◄─ x2 ◄─ ... ◄─ xT ~ N(0, I)
                        each step: predict noise, step toward x0

The reverse process and the objective

Generation means sampling x_T ~ N(0, I) and inverting the chain. The true reversal q(x_t−1 | x_t) is intractable (it depends on the whole data distribution), but for small β_t it is close to Gaussian, so we learn a Gaussian approximation p_θ(x_t−1 | x_t) = N(μ_θ(x_t, t), σ_t²I), with variances kept fixed to the schedule in plain DDPM. The crucial tractable object is the posterior conditioned on the clean image, which Bayes gives in closed form: q(x_t−1 | x_t, x₀) = N(μ̃_t, β̃_tI) with

μ̃_t = (√ᾱ_t−1β_t / (1−ᾱ_t)) x₀ + (√α_t(1−ᾱ_t−1) / (1−ᾱ_t)) x_t,    β̃_t = ((1−ᾱ_t−1) / (1−ᾱ_t)) β_t.

Training maximizes a variational lower bound exactly as in a VAE (the diffusion chain is a VAE with a fixed, depth-T encoder; the general framework is on the VAE page). The ELBO splits into per-step KL terms KL(q(x_t−1 | x_t, x₀) ‖ p_θ(x_t−1 | x_t)), and since both are Gaussians with matched variance, each term is just ‖μ̃_t − μ_θ‖² / 2σ_t²: the objective is mean-matching. Now the parameterization insight. Solve the closed form for the clean data, x₀ = (x_t − √(1−ᾱ_t)ε)/√ᾱ_t, substitute into μ̃_t, and the algebra simplifies to

μ̃_t = (1/√α_t) (x_t − (β_t/√(1−ᾱ_t)) ε).

The mean the network must produce is a fixed affine function of x_t, which it already has, and ε, the noise that was mixed in. So instead of predicting the mean, predict the noise: let the network output ε_θ(x_t, t) and define μ_θ by the same formula. The per-step KL becomes (β_t² / 2σ_t²α_t(1−ᾱ_t)) ‖ε − ε_θ(x_t, t)‖², a weighted noise-regression loss. Ho et al.'s empirical move, and this is the honest statement of what the famous "simple loss" is, was to drop the weight:

L_simple = E_{t, x0, ε} ‖ε − ε_θ(√ᾱ_tx₀ + √(1−ᾱ_t)ε, t)‖².

This is not the ELBO; it is a reweighted ELBO that de-emphasizes the low-noise steps (where the schedule-derived weights are large but denoising is nearly trivial) and spends relatively more of the model's capacity on the heavily-noised steps where the real structure-from-noise work happens. The reweighting costs a little log-likelihood and buys a lot of sample quality, and it is what everyone trains in practice. So the entire training algorithm is: draw a clean sample, draw a uniform random t, draw Gaussian noise, form x_t in one line, and take an MSE step on the noise prediction.

Schedules, sampling, and DDIM

The linear schedule (β from 10⁻⁴ to 0.02) was the original choice and works, but it destroys information front-heavily: ᾱ_t falls so fast that the last few hundred steps of a 1000-step chain are spent shuffling nearly pure noise. The cosine schedule of Nichol and Dhariwal defines ᾱ_t directly as a squared cosine, ᾱ_t = f(t)/f(0) with f(t) = cos²(((t/T + s)/(1 + s))·π/2), s = 0.008, recovers β_t = 1 − ᾱ_t/ᾱ_t−1, and clips the result away from 1. The effect is a slower, more even decay of signal, which measurably helps on low-resolution images where the linear schedule is most wasteful.

Sampling runs the learned mean formula with fresh noise at every step except the last: x_t−1 = (1/√α_t)(x_t − (β_t/√(1−ᾱ_t))ε_θ(x_t, t)) + σ_tz, with z ~ N(0, I), σ_t² either β_t or the posterior variance β̃_t (both are used in practice), and z = 0 at t = 1 so the final output is the mean. That is T network evaluations. DDIM (Song et al.) is the standard answer to the cost: it observes that many non-Markovian forward processes share the same marginals q(x_t | x₀), so a model trained with L_simple is secretly a model of all of them. Choosing the deterministic member of that family gives a sampler that, at each step, estimates x̂₀ = (x_t − √(1−ᾱ_t)ε_θ)/√ᾱ_t and jumps directly toward it along the schedule, with no noise injection. Determinism makes big jumps coherent, so you can keep every 20th timestep and sample in 50 steps instead of 1000 with modest quality loss, using the same trained weights. Conceptually DDIM turns the chain into an ODE integrator, which is the door through which the whole fast-sampler literature walked.

Implementation, twice

Everything below is parameterized over an arbitrary ε-model, any callable taking (x_t, t) and returning a tensor shaped like x_t. For images the right ε-model is a U-Net with the timestep injected via an embedding into every residual block; that architecture has its own page at /ml/unet and is imported as an idea here rather than duplicated. The toy example further down uses a three-layer MLP so the whole page runs on a laptop CPU. First, the schedules and the training step:

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

def linear_beta_schedule(T, lo=1e-4, hi=0.02):
    return torch.linspace(lo, hi, T)

def cosine_beta_schedule(T, s=0.008):
    # Nichol & Dhariwal: define alpha_bar as a squared cosine,
    # then recover betas from consecutive ratios.
    t = torch.linspace(0, T, T + 1) / T
    f = torch.cos((t + s) / (1 + s) * math.pi / 2) ** 2
    alpha_bar = f / f[0]
    betas = 1 - alpha_bar[1:] / alpha_bar[:-1]
    return betas.clamp(1e-8, 0.999)     # keep every step invertible

class Diffusion:
    """Schedule tables + the two operations that matter.
    Works for any eps_model(x_t, t) whose output matches x_t."""

    def __init__(self, betas):
        self.T = len(betas)
        self.betas = betas
        self.alphas = 1.0 - betas
        self.alpha_bar = torch.cumprod(self.alphas, dim=0)

    def q_sample(self, x0, t, eps):
        # closed form q(x_t | x_0): one jump to any noise level
        ab = self.alpha_bar[t].view(-1, *([1] * (x0.dim() - 1)))
        return ab.sqrt() * x0 + (1 - ab).sqrt() * eps

    def loss(self, eps_model, x0):
        # L_simple: sample t and eps, form x_t, regress the noise
        t = torch.randint(0, self.T, (x0.size(0),), device=x0.device)
        eps = torch.randn_like(x0)
        xt = self.q_sample(x0, t, eps)
        return F.mse_loss(eps_model(xt, t), eps)

import math
import jax
import jax.numpy as jnp

def linear_beta_schedule(T, lo=1e-4, hi=0.02):
    return jnp.linspace(lo, hi, T)

def cosine_beta_schedule(T, s=0.008):
    # Nichol & Dhariwal: define alpha_bar as a squared cosine,
    # then recover betas from consecutive ratios.
    t = jnp.linspace(0, T, T + 1) / T
    f = jnp.cos((t + s) / (1 + s) * math.pi / 2) ** 2
    alpha_bar = f / f[0]
    betas = 1 - alpha_bar[1:] / alpha_bar[:-1]
    return jnp.clip(betas, 1e-8, 0.999)  # keep every step invertible

def make_diffusion(betas):
    alphas = 1.0 - betas
    return dict(betas=betas, alphas=alphas,
                alpha_bar=jnp.cumprod(alphas))

def q_sample(diff, x0, t, eps):
    # closed form q(x_t | x_0): one jump to any noise level
    ab = diff['alpha_bar'][t].reshape(-1, *([1] * (x0.ndim - 1)))
    return jnp.sqrt(ab) * x0 + jnp.sqrt(1 - ab) * eps

def loss_fn(params, apply_fn, diff, x0, key):
    # L_simple: sample t and eps, form x_t, regress the noise
    kt, ke = jax.random.split(key)
    t = jax.random.randint(kt, (x0.shape[0],), 0,
                           diff['betas'].shape[0])
    eps = jax.random.normal(ke, x0.shape)
    xt = q_sample(diff, x0, t, eps)
    pred = apply_fn(params, xt, t)
    return jnp.mean((pred - eps) ** 2)

Then the full ancestral sampling loop. The PyTorch version is a plain Python loop under no_grad; the JAX version expresses the same loop as a lax.scan over timesteps so the whole sampler jit-compiles into one program.

class Diffusion(Diffusion):          # continued from above

    @torch.no_grad()
    def sample(self, eps_model, shape, device='cpu'):
        x = torch.randn(shape, device=device)      # x_T ~ N(0, I)
        for i in reversed(range(self.T)):
            t = torch.full((shape[0],), i, device=device)
            eps = eps_model(x, t)
            a, ab, b = self.alphas[i], self.alpha_bar[i], self.betas[i]
            # posterior mean in the eps-parameterization
            mean = (x - b / (1 - ab).sqrt() * eps) / a.sqrt()
            if i > 0:
                ab_prev = self.alpha_bar[i - 1]
                var = b * (1 - ab_prev) / (1 - ab)   # beta_tilde
                x = mean + var.sqrt() * torch.randn_like(x)
            else:
                x = mean       # last step: no noise, return the mean
        return x

from functools import partial

@partial(jax.jit, static_argnums=(1, 3))
def sample(params, apply_fn, diff, shape, key):
    T = diff['betas'].shape[0]
    key, k0 = jax.random.split(key)
    x_T = jax.random.normal(k0, shape)             # x_T ~ N(0, I)

    def step(x, inp):
        i, k = inp
        t = jnp.full((shape[0],), i)
        eps = apply_fn(params, x, t)
        a, ab, b = (diff['alphas'][i], diff['alpha_bar'][i],
                    diff['betas'][i])
        # posterior mean in the eps-parameterization
        mean = (x - b / jnp.sqrt(1 - ab) * eps) / jnp.sqrt(a)
        # beta_tilde; ab_prev := 1 at i == 0 (guarded by where)
        ab_prev = jnp.where(i > 0, diff['alpha_bar'][i - 1], 1.0)
        var = b * (1 - ab_prev) / (1 - ab)
        z = jax.random.normal(k, shape) * (i > 0)  # no noise at i == 0
        return mean + jnp.sqrt(var) * z, None

    steps = jnp.arange(T - 1, -1, -1)              # T-1 ... 0
    keys = jax.random.split(key, T)
    x0, _ = jax.lax.scan(step, x_T, (steps, keys))
    return x0

Using it on a real shape of problem

The smallest problem where diffusion visibly works is a 2-D point cloud with structure a single Gaussian cannot capture. A swiss-roll spiral is the classic: multimodal along its length, curved, and plottable. The ε-model is a tiny MLP whose only architectural obligation is taking t seriously, here via a learned embedding concatenated onto the point.

def swiss_roll(n):
    # 2-D spiral, scaled to roughly unit variance: the closed
    # form assumes data on the same scale as N(0, I).
    theta = 1.5 * math.pi * (1 + 2 * torch.rand(n))
    x = torch.stack([theta * theta.cos(), theta * theta.sin()], -1)
    return x / 6.0

class TimeMLP(nn.Module):
    """eps-model for 2-D points; t enters via a learned embedding."""
    def __init__(self, T, width=128):
        super().__init__()
        self.emb = nn.Embedding(T, width)
        self.net = nn.Sequential(
            nn.Linear(2 + width, width), nn.SiLU(),
            nn.Linear(width, width), nn.SiLU(),
            nn.Linear(width, 2))
    def forward(self, x, t):
        return self.net(torch.cat([x, self.emb(t)], dim=-1))

T = 200                                # short chain is plenty in 2-D
diff = Diffusion(linear_beta_schedule(T))
model = TimeMLP(T)
opt = torch.optim.AdamW(model.parameters(), lr=2e-3)

for step in range(5000):
    loss = diff.loss(model, swiss_roll(512))
    opt.zero_grad(); loss.backward(); opt.step()
    if step % 1000 == 0:
        print(step, float(loss))       # ~1.0 down to ~0.3-0.5

pts = diff.sample(model, (2000, 2))    # scatter this: a spiral

import optax                            # standard JAX optimizer library

def swiss_roll(key, n):
    # 2-D spiral, scaled to roughly unit variance
    theta = 1.5 * jnp.pi * (1 + 2 * jax.random.uniform(key, (n,)))
    x = jnp.stack([theta * jnp.cos(theta),
                   theta * jnp.sin(theta)], -1)
    return x / 6.0

def init_mlp(key, T, width=128):
    ks = jax.random.split(key, 4)
    dims = [(2 + width, width), (width, width), (width, 2)]
    ws = [jax.random.normal(k, d) / jnp.sqrt(d[0])
          for k, d in zip(ks[:3], dims)]
    return dict(emb=jax.random.normal(ks[3], (T, width)) * 0.02,
                ws=ws, bs=[jnp.zeros(d[1]) for d in dims])

def apply_mlp(p, x, t):
    h = jnp.concatenate([x, p['emb'][t]], axis=-1)
    for w, b in zip(p['ws'][:-1], p['bs'][:-1]):
        h = jax.nn.silu(h @ w + b)
    return h @ p['ws'][-1] + p['bs'][-1]

T = 200                                 # short chain is plenty in 2-D
diff = make_diffusion(linear_beta_schedule(T))
key = jax.random.PRNGKey(0)
params = init_mlp(key, T)
opt = optax.adamw(2e-3)
opt_state = opt.init(params)

@jax.jit
def train_step(params, opt_state, x0, key):
    loss, grads = jax.value_and_grad(loss_fn)(
        params, apply_mlp, diff, x0, key)
    updates, opt_state = opt.update(grads, opt_state, params)
    return optax.apply_updates(params, updates), opt_state, loss

for step in range(5000):
    key, kd, kl = jax.random.split(key, 3)
    params, opt_state, loss = train_step(
        params, opt_state, swiss_roll(kd, 512), kl)
    if step % 1000 == 0:
        print(step, float(loss))        # ~1.0 down to ~0.3-0.5

pts = sample(params, apply_mlp, diff, (2000, 2),
             jax.random.PRNGKey(1))     # scatter this: a spiral

What to expect, hedged as machine-dependent where it is: the loss starts near 1.0 (the variance of the noise you are trying to predict) and settles somewhere around 0.3 to 0.5 within a few thousand steps; exact values depend on seed and hardware, but it will not approach zero, for reasons the traps section makes precise. Scatter-plot the 2000 sampled points against a fresh draw of the data and you should see the spiral reproduced with slightly soft edges; if you instead see a single Gaussian blob, the time conditioning is broken, and if you see the data mean only, the sampler is adding noise at the last step or using the wrong coefficient table. Watching intermediate x_t during sampling is worth the ten lines of plotting: the spiral condenses out of the fog over the last quarter of the steps.

Applications

Image generation is where diffusion earned its reputation: Stable Diffusion, Imagen, DALL·E 2 and 3, Midjourney, and Flux are all diffusion (or diffusion-descended flow) models. The dominant architecture is latent diffusion, which runs the chain not on pixels but in the latent space of a pretrained autoencoder, an 8× spatial compression that makes high-resolution generation affordable; that autoencoder is the VAE described at /ml/vae, and the text-to-image conditioning enters through cross-attention from U-Net feature maps to text embeddings. Video generation (Sora, Veo, Runway) is diffusion over spacetime latents with transformer backbones. In audio, diffusion drives text-to-music and text-to-sound systems (Stable Audio, AudioLDM) and neural vocoders. In science, RFdiffusion designs novel protein backbones by denoising in structure space. And in robotics, diffusion policies treat a short horizon of future actions as the thing to denoise, conditioned on observations: the iterative refinement handles multimodal demonstrations (go-left and go-right both present in the data) that regression policies average into failure. I train exactly that class of model in my robot imitation lab project. The common thread across every domain: wherever "the answer" is a continuous object with many valid modes, denoising beats regressing.

Against the real libraries

huggingface/diffusers is the production home of this algorithm. The pieces map onto this page one to one: DDPMScheduler is the Diffusion class above, holding the β table, ᾱ products, add_noise (our q_sample), and step (one iteration of our sampling loop); UNet2DModel is the image-scale ε-model; and DDPMPipeline wires the two into a loop that matches sample line for line. What the library adds over this reference is breadth and hardening: a zoo of interchangeable schedulers (DDIM, DPM-Solver++, Euler and Heun variants) behind one interface, the prediction_type switch between ε, x₀, and v-prediction, classifier-free guidance, latent-space plumbing, mixed precision, and attention-slicing memory tricks, plus tested weights for every major released model. The from-scratch version is enough whenever the model is yours end to end, low-dimensional scientific data, robot action spaces, or any research idea that touches the schedule or loss, where a 60-line implementation you fully control beats configuring a framework.

Two other repositories are worth knowing as references. The original hojonathanho/diffusion is Ho et al.'s TensorFlow code for the DDPM paper, still the ground truth for what the paper actually did, down to the schedule constants. And lucidrains/denoising-diffusion-pytorch is the community's standard PyTorch reimplementation, close enough to the math to read alongside the paper and the usual starting point for research forks; its sibling repos track most published improvements (v-prediction, min-SNR weighting, self-conditioning) within weeks of publication.

Verification against diffusers is direct because the scheduler exposes its tables. With the same linear schedule, our constants and its constants must agree to float tolerance, and our q_sample must reproduce its add_noise:

from diffusers import DDPMScheduler

sched = DDPMScheduler(num_train_timesteps=1000,
                      beta_start=1e-4, beta_end=0.02,
                      beta_schedule='linear')
mine = Diffusion(linear_beta_schedule(1000))

assert torch.allclose(mine.alpha_bar, sched.alphas_cumprod, atol=1e-6)

x0 = torch.randn(4, 2)
eps = torch.randn_like(x0)
t = torch.tensor([0, 10, 500, 999])
assert torch.allclose(mine.q_sample(x0, t, eps),
                      sched.add_noise(x0, eps, t), atol=1e-6)

import numpy as np
import torch                          # only to drive the reference
from diffusers import DDPMScheduler   # tables are framework-neutral

sched = DDPMScheduler(num_train_timesteps=1000,
                      beta_start=1e-4, beta_end=0.02,
                      beta_schedule='linear')
mine = make_diffusion(linear_beta_schedule(1000))

np.testing.assert_allclose(np.asarray(mine['alpha_bar']),
                           sched.alphas_cumprod.numpy(), atol=1e-6)

rng = np.random.default_rng(0)
x0 = rng.standard_normal((4, 2)).astype('float32')
eps = rng.standard_normal((4, 2)).astype('float32')
t = np.array([0, 10, 500, 999])

ours = q_sample(mine, jnp.array(x0), jnp.array(t), jnp.array(eps))
ref = sched.add_noise(torch.tensor(x0), torch.tensor(eps),
                      torch.tensor(t))
np.testing.assert_allclose(np.asarray(ours), ref.numpy(), atol=1e-5)

If those two asserts pass, everything downstream (training loss values, sampler behavior) is comparable against DDPMPipeline outputs step by step, since both samplers are deterministic given the same noise draws.

Traps and misconceptions

Expecting the loss to reach zero. The target ε is fresh randomness, and at high noise levels x_t simply does not contain enough information to recover it exactly; the loss has a large irreducible floor. A diffusion loss curve flattens early and stays flat while sample quality keeps improving for a long time afterwards. Judge these models by samples (or a held-out metric like FID), never by how close the training MSE gets to zero.

Weak or missing time conditioning. One network must denoise at every noise level, from nearly-clean to pure noise, and it can only do that if t reaches it with enough capacity: an embedding injected into every block for a U-Net, not a scalar appended to the input once. Break the conditioning and the model learns an average denoiser that fails at both ends of the schedule, typically producing blurry mean-seeking samples.

Data scale mismatch. The closed form interpolates between the data and N(0, I), so the data must live on a comparable scale (images in [−1, 1], point clouds standardized). Feed raw unscaled data and x_T is nowhere near the prior you sample from at generation time, so the reverse chain starts from a distribution the model never saw.

Noise at the final step, or the wrong variance. The sampling loop adds σ_tz at every step except t = 1, where the mean is returned; leaving noise in the last step visibly grains the output. Similarly, mixing up β_t and β̃_t as the sampling variance, or mismatching the ε / x₀ / v prediction convention between training and sampling (the prediction_type field in diffusers configs), produces samples that are wrong in quietly systematic ways rather than obviously broken.

"Diffusion is inherently 1000× slower than a GAN." Plain DDPM sampling is, but the trained model is not welded to its sampler: DDIM cuts to ~50 steps, DPM-Solver++ to ~20, and distillation and consistency training reach 1-4 steps with the diffusion model as teacher. Slowness is a property of the ancestral sampler, not of the modeling framework.