🌫️ Diffusion Model Lab: noise schedules, reverse sampling, DDPM vs DDIM

How image and video generators based on diffusion work: add noise step by step, then undo it with the score (the gradient of log-density). Here the "image" is a cloud of 2-D points, so everything can be computed exactly.

Honest scope: this is a 2-D toy that illustrates the mathematics of diffusion models. It is not an image generator. The "perfect denoiser" is the exact analytic score of a Gaussian mixture; real models only learn an approximation to it with a large neural network.

1. Particles travelling between noise and data

Reverse (generate): start from pure Gaussian noise at t = T and follow the sampler to t = 0. Background = target density at the current noise level; faint white dots = real data.

Settings

2. Noise schedule: ᾱt and signal-to-noise ratio

Top: fraction of signal variance left. Bottom: SNR(t) = ᾱt/(1−ᾱt) in dB. Dots = the S sampling steps; vertical line = current t.

3. Live histogram: generated vs true density

Bars = particles right now. Solid = target marginal at t = 0. Dashed = exact marginal pt at the current noise level.

4. Experiments (use the settings above)

The science

Forward (noising) process. With variances βt, αt = 1 − βt and ᾱt = ∏s≤t αs, each step is xt = √(1−βt) xt−1 + √βt z. Composing the steps gives the closed form q(xt|x0) = N(√ᾱt x0, (1−ᾱt) I) (Sohl-Dickstein et al., 2015; Ho, Jain and Abbeel, 2020). Experiment "Forward variance check" iterates the chain and compares the variance of xt − √ᾱt x0 with 1 − ᾱt.

Schedules and SNR. Linear: βt rises linearly from 10−4 to 0.02 (rescaled by 1000/T for other T), as in Ho et al. Cosine: ᾱt = f(t)/f(0), f(t) = cos2((t/T + s)/(1 + s) · π/2), s = 0.008, with βt clipped at 0.999 (Nichol and Dhariwal, 2021). SNR(t) = ᾱt/(1 − ᾱt). The linear schedule destroys the signal early (many steps are spent at almost zero SNR); the cosine schedule spreads the information loss more evenly, which matters most when you take few sampling steps.

Exact score for a Gaussian mixture. If the data are p0 = Σk wk N(μk, sk2 I), then pt is again a mixture: N(√ᾱt μk, vk I) with vk = ᾱt sk2 + 1 − ᾱt. Its score is ∇ log pt(x) = Σk rk(x) (√ᾱt μk − x)/vk where rk are the posterior responsibilities. Networks usually predict the noise ε, related by ε*(x,t) = −√(1−ᾱt) ∇ log pt(x). The page uses the exact value as a "perfect denoiser". The optional MLP (7 inputs, two hidden layers of 64 SiLU units, trained in your browser by Adam on the ε-prediction loss) is a deliberately small learned stand-in; its error against the exact ε is shown.

Reverse sampling (DDIM family). From ε one estimates x̂0 = (xt − √(1−ᾱt) ε)/√ᾱt and moves to an earlier time t': xt' = √ᾱt' x̂0 + √(1 − ᾱt' − σ2) ε + σ z with σ = η √((1−ᾱt')/(1−ᾱt)) √(1 − ᾱt/ᾱt') (Song, Meng and Ermon, 2021). η = 1 is DDPM ancestral sampling with the posterior variance (identical to Ho et al. when S = T); η = 0 is deterministic DDIM. Taking S < T steps just re-uses a sub-sequence of times (round(T·(S−i)/S)).

Probability-flow ODE. Song et al. (2021, "Score-Based Generative Modeling through SDEs") showed that every diffusion has a deterministic ODE with the same marginals pt: dx = [f(x,t) − ½ g(t)2 ∇ log pt(x)] dt. DDIM with η = 0 is an Euler-type discretisation of it, which is why it is smooth, reproducible from the same xT, and can be run backwards (the "ODE round trip" experiment: encode data to noise, decode again). With very few steps the discretisation error dominates: one step returns only the posterior mean of the data, a blur at the centre.

Classifier-free guidance. Ho and Salimans (2022) combine a conditional and an unconditional model: ε̃ = εu + w (εc − εu). Here w = 0 is unconditional, w = 1 is the plain conditional model and w > 1 extrapolates away from the unconditional score, which tilts the density towards p(x|c)w p(x)1−w ∝ p(x|c) p(c|x)w−1. In this toy the conditional and unconditional scores are both exact (class = a subset of mixture components). Measured in this page: class purity rises and the spread of the samples falls below that of p(x|c) as w grows (for overlapping classes such as the two blobs the samples are also pushed away from the shared region, so the distance to p(x|c) increases). The effect is the fidelity-versus-diversity trade-off reported for image generators. This follows the usual convention; some papers use w' = w − 1.

Quality metric. The energy distance D2 = 2E|X−Y| − E|X−X'| − E|Y−Y'| (Szekely and Rizzo) is zero if and only if the two distributions are equal. The page compares the generated points with an independent reference sample from the true target. With N points the value never reaches exactly zero; the "noise floor" tile is the same statistic between two independent true samples. Precision/recall use a fixed radius of 0.3: the fraction of generated points within 0.3 of a reference point, and the fraction of reference points having a generated point within 0.3.

Simplifications and honest limits. (1) 2-D data and a Gaussian mixture are what make the exact score possible; real images live in thousands of dimensions and the score must be learned. (2) A real model's error grows where data are rare and at small noise, and respacing the steps (S < T) is only valid because the network is conditioned on ᾱ. (3) Guidance here uses exact class-conditional scores, not a text encoder. (4) x̂0 is not clipped, so at t = T it divides by a tiny √ᾱT; the terms cancel exactly with the exact score, but a learned model may amplify errors there. (5) The ODE round trip uses the standard DDIM inversion approximation (the ε at the current point is re-used), which is exact only as S grows. (6) Estimates use a few hundred particles, so every metric has sampling noise of a few percent; it is shown by the noise-floor tile. (7) This is not a model of any specific product.

How this lab differs from the existing "Diffusion Model Simulator" (diffusion-models-basics). That page noises and denoises a single point-cloud picture with one beta-steepness knob. This lab compares linear and cosine schedules with SNR curves, runs genuine reverse samplers (DDPM and DDIM with adjustable η, step count and horizon T) with the exact analytic score of a mixture, adds classifier-free guidance, and measures sample quality against the true density with the energy distance. The "Langevin Dynamics and Score-Based Generative Models" lab shows annealed Langevin sampling; this one concerns the DDPM/DDIM updates and the schedules.

References. Sohl-Dickstein, Weiss, Maheswaranathan, Ganguli (2015); Ho, Jain, Abbeel (2020) DDPM; Song, Meng, Ermon (2021) DDIM; Song, Sohl-Dickstein, Kingma, Kumar, Ermon, Poole (2021) score SDEs and probability-flow ODE; Nichol and Dhariwal (2021) improved DDPM / cosine schedule; Ho and Salimans (2022) classifier-free guidance; Szekely and Rizzo (2013) energy distance.

Data: current samples as CSV (x_T, generated, reference)