Diffusion models

Problem

A generative model has to turn simple randomness into samples that look like the data. A VAE does it in one decoder pass from a learned latent, and its samples tend to be blurry because the decoder is trained with a simple per-pixel likelihood. A GAN does it in one generator pass and gets sharp samples, but trains as an unstable two-player game that can drop whole modes of the data. An autoregressive model is exact and stable but makes one value at a time.

Diffusion models split the hard jump from noise to data into a long chain of small, easy steps. Each step only has to remove a little noise, which is a plain regression problem with a fixed target, so training is a single stable mean-squared-error loss with no adversary. The cost moves to sampling: making one sample runs the network once per step, tens to a thousand times. This trade, stable training and excellent samples for slow generation, is why diffusion became the standard for images, audio and video.

Intuition

Take a photo and add a little Gaussian noise. Do it again, and again, a thousand times, each time shrinking the image slightly so the values stay bounded. Nothing of the photo is left at the end: it is indistinguishable from a draw of pure noise. This forward process needs no learning, and because every step is Gaussian, the noisy image after any number of steps can be written directly as a mix of the original and one draw of noise.

Now run the film backwards. If a network can look at a noisy image and say which noise was added, subtracting a little of that noise gives a slightly cleaner image. Starting from pure noise and repeating this a thousand times walks a sample back to something that looks like data. The network never has to imagine a whole image at once; at high noise it only decides the rough layout, at low noise only fine texture.

Training is then self-supervised: pick a real image, pick a noise level, add the noise, and ask the network to predict it. The loss is the squared error between the noise it guessed and the noise that was actually added.

Mechanics

The forward process in closed form

Fix T steps and a noise schedule β₁, …, β_T, small numbers between 0 and 1. Each step scales the previous value down and adds fresh Gaussian noise:

Write α_t = 1 − β_t and ᾱ_t = α₁α₂⋯α_t. Then the distribution after t steps, given the data, is one Gaussian:

The proof is by induction, using how means and variances move under scaling and adding independent noise. The mean is multiplied by √α_t at each step, so after t steps it is √ᾱ_tx₀. The variance follows v_t = α_tv_t−1 + β_t; if v_t−1 = 1 − ᾱ_t−1, then v_t = α_t − ᾱ_t + 1 − α_t = 1 − ᾱ_t. The factor √(1 − β_t) is what keeps the total variance at 1 when the data has variance 1: signal variance ᾱ_t plus noise variance 1 − ᾱ_t. Training never runs the chain; it jumps straight to any t with one draw of ε.

The amount of signal left is measured by the signal-to-noise ratio SNR(t) = ᾱ_t/(1 − ᾱ_t), which falls from very large at t = 1 to almost zero at t = T. A schedule is good if x_T is close enough to N(0, I) that sampling can start there, and if the steps in between spread the work evenly.

The ε-prediction objective and the ELBO

The reverse model is a chain of Gaussians p_θ(x_t−1 | x_t) = N(μ_θ(x_t, t), σ_t² I) starting from p(x_T) = N(0, I). Like a VAE, with x₁, …, x_T as latents and the forward process as a fixed encoder, it is trained by minimizing a bound on −log p_θ(x₀):

L_T has no parameters, and L₀ is a reconstruction term for the last step. The middle terms compare the model with a distribution that is easy to write down: given both the noisy x_t and the clean x₀, the previous step is Gaussian by Bayes' rule,

Substitute x₀ = (x_t − √(1 − ᾱ_t) ε)/√ᾱ_t from the closed form and the mean simplifies to a function of x_t and the noise alone:

So give μ_θ the same form with a network ε_θ(x_t, t) in place of ε. Both Gaussians have the same fixed variance, and the KL divergence between two Gaussians with equal variance is their squared mean difference over 2σ²:

Every ELBO term is a weighted noise-prediction error. Ho, Jain and Abbeel found that dropping the weight trains better samples, giving the loss everyone uses:

The weight in front is largest at small t, where the noise is tiny and predicting it is easy; L_simple sets every weight to 1, which spends relatively more of the training on the large-t steps that decide global structure. It is also a mean squared error against a target that has unit variance at every t, which keeps the network's output scale the same across noise levels. Because a squared error is minimized by the conditional mean, the best possible ε_θ is E[ε | x_t], equivalently the best guess of the clean data, E[x₀ | x_t].

Noise schedules

The original DDPM schedule is linear: β grows evenly from 10⁻⁴ to 0.02 over T = 1000 steps. Nichol and Dhariwal noticed that this destroys the signal too early: by the middle of the chain the image is already almost pure noise, so many steps are spent where there is little to learn. Their cosine schedule defines ᾱ directly,

with β_t = 1 − ᾱ_t/ᾱ_t−1 clipped at 0.999 so the last step is not exactly 1. The small offset s keeps β tiny near t = 0. The cosine curve falls slowly at both ends and roughly linearly in between, so the SNR declines evenly. Later work reparameterizes the schedule entirely by SNR, and models trained with one schedule can be sampled with different step counts, which is where DDIM comes in.

DDPM sampling

To sample, draw x_T ~ N(0, I) and apply the learned reverse step T times:

DDPM uses σ_t² = β_t or β̃_t, which the paper found to give similar results. Every step injects fresh noise, so the sampler is stochastic and must take all T steps: the reverse Gaussians are only accurate when each step is small. A thousand network evaluations per sample is the main cost of DDPM.

DDIM sampling

Song, Meng and Ermon showed that the same trained network fits a whole family of non-Markovian forward processes that share the marginals q(x_t | x₀), and the deterministic member of that family can skip steps. At each step, use the predicted noise to estimate the clean data, then re-noise that estimate to an earlier, less noisy level t′ using the same predicted noise:

With no fresh noise, the map from x_T to x₀ is deterministic: the same starting noise gives the same sample, and nearby starting noises give similar samples, so the starting noise behaves like a latent code that can be interpolated. The step sequence t, t′, … can be any decreasing subsequence of 1…T, often 20 to 100 steps, with no retraining. DDIM also has a knob η that adds back a fraction of DDPM's noise; η = 0 is the deterministic sampler and η = 1 recovers DDPM. Many later samplers treat this update as a step of an ordinary differential equation solver.

Classifier-free guidance

A conditional model ε_θ(x_t, t, c) takes a class label or a text embedding c. To make samples follow c more strongly, Ho and Salimans train one network to be both conditional and unconditional: during training, c is replaced by a null token ∅ for a small random fraction of examples. At sampling time the network runs twice per step and the two predictions are extrapolated:

w = 0 is the plain conditional model. Larger w pushes away from what the unconditional model predicts, toward what is specific to c. In score terms the guided model follows p(x_t | c)^1+w / p(x_t)^w, which is proportional to p(x_t | c) · p(c | x_t)^w: the conditional distribution sharpened by an implicit classifier. Samples become more typical of c and less diverse. Many libraries write the same rule as ε_∅ + s · (ε_c − ε_∅) with scale s = 1 + w, so a "guidance scale of 7.5" means w = 6.5.

Latent diffusion

Running a large network a thousand times per sample over every pixel is expensive. Rombach and colleagues first train an autoencoder that compresses an image to a smaller latent, with a perceptual and adversarial loss so that the decoded image keeps its detail, and then train the diffusion model on the latents instead of the pixels. Every formula above is unchanged; x₀ is now the encoder's output, and the decoder runs once at the end. With a downsampling factor of 8 and 4 latent channels, a 512 × 512 RGB image becomes a 4 × 64 × 64 latent, 48 times fewer values. Text conditioning enters the denoiser through cross-attention from the latent to the text encoder's token embeddings. This is the design of Stable Diffusion.

PyTorch v2.15.0 has no diffusion modules; libraries such as Hugging Face Diffusers build schedulers and pipelines from these tensor and random-number operations.

Worked example

All numbers use T = 1000 steps. The schedule part follows one data value x₀ = 1 and one noise draw ε = -0.5. The sampling part uses a dataset small enough that the perfect denoiser can be written in closed form, so every sample below is what an ideally trained network would produce.

Implementation

Training with label dropout for guidance, and a guided DDIM sampler, for any denoiser called as model(x, t, labels). Labels are class indices, with one extra index reserved as the null label ∅.

import torch
import torch.nn.functional as F

T = 1000
betas = torch.linspace(1e-4, 0.02, T)
alpha_bars = torch.cumprod(1 - betas, dim=0)

def training_loss(model, x0, labels, null_label, p_uncond=0.1):
    t = torch.randint(0, T, (x0.size(0),))
    eps = torch.randn_like(x0)
    ab = alpha_bars[t].view(-1, *[1] * (x0.dim() - 1))       # broadcast over C, H, W
    xt = ab.sqrt() * x0 + (1 - ab).sqrt() * eps
    drop = torch.rand(x0.size(0)) < p_uncond                 # train the unconditional model too
    labels = torch.where(drop, torch.full_like(labels, null_label), labels)
    return F.mse_loss(model(xt, t, labels), eps)

@torch.no_grad()
def ddim_sample(model, shape, labels, null_label, steps=50, w=3.0, generator=None):
    ts = list(range(T - 1, -1, -(T // steps)))[:steps]
    x = torch.randn(shape, generator=generator)
    null = torch.full_like(labels, null_label)
    for k, t in enumerate(ts):
        ab = alpha_bars[t]
        ab_next = alpha_bars[ts[k + 1]] if k + 1 < len(ts) else torch.tensor(1.0)
        tt = torch.full((shape[0],), t)
        eps = (1 + w) * model(x, tt, labels) - w * model(x, tt, null)
        x0_hat = (x - (1 - ab).sqrt() * eps) / ab.sqrt()
        x = ab_next.sqrt() * x0_hat + (1 - ab_next).sqrt() * eps
    return x

Indices here run from 0 to T − 1, one less than the page's t. Real samplers usually clip x̂₀ to the data range or rescale it, batch the conditional and unconditional calls into one forward pass of twice the batch, and pass a torch.Generator so a seed reproduces a sample. The T // steps stride is one of several spacings; libraries also offer evenly spaced and "trailing" spacings that include t = T exactly. Some models predict v = √ᾱ_tε − √(1 − ᾱ_t)x₀ instead of ε (Salimans and Ho's progressive distillation), which behaves better when SNR is near zero; the sampler converts it back to ε and x̂₀.

Tradeoffs

• Versus GANs and VAEs. Diffusion trains with one stable regression loss and covers all the modes of the data; GANs train adversarially and can drop modes but sample in one pass; VAEs sample in one pass and give a likelihood bound but are blurrier. Diffusion pays with tens to a thousand network calls per sample.
• DDPM vs DDIM. DDPM's injected noise can correct earlier errors and gives the best quality at many steps; DDIM is deterministic, reproducible, invertible to a latent and good at 20 to 100 steps, but degrades faster at very few steps. Distillation and consistency models push to a handful of steps by training a new network.
• Schedule. Linear is the original and fine at 1000 steps; cosine spends steps more evenly and helps at low resolution and with few sampling steps. At high resolution the same schedule leaves more signal (neighbouring pixels share it), so large models shift the schedule toward more noise.
• Guidance strength. Higher w gives samples that match the prompt better and look cleaner, at the cost of diversity and, at high w, saturation and artifacts. It doubles the cost of every step.
• Pixel vs latent. Latent diffusion is many times cheaper and is the default for large images, but the autoencoder caps fine detail such as small text and faces, and is one more model to train.
• Likelihood. L_simple gives better samples than the true ELBO; models trained for likelihood need the weights back, learned variances or a different parameterization.

Connections

References

• Ho, Jain and Abbeel. Denoising Diffusion Probabilistic Models. NeurIPS 2020. arXiv:2006.11239
• Song, Meng and Ermon. Denoising Diffusion Implicit Models. ICLR 2021. arXiv:2010.02502
• Ho and Salimans. Classifier-Free Diffusion Guidance. 2022. arXiv:2207.12598
• Nichol and Dhariwal. Improved Denoising Diffusion Probabilistic Models. ICML 2021. arXiv:2102.09672
• Rombach, Blattmann, Lorenz, Esser and Ommer. High-Resolution Image Synthesis with Latent Diffusion Models. CVPR 2022. arXiv:2112.10752
• Sohl-Dickstein, Weiss, Maheswaranathan and Ganguli. Deep Unsupervised Learning using Nonequilibrium Thermodynamics. ICML 2015. arXiv:1503.03585
• PyTorch v2.15.0: F.mse_loss, torch.randn_like, torch.cumprod, torch.manual_seed, torch.random.fork_rng.