Diffusion Lab

About

Diffusion models generate images by learning to reverse a process that slowly turns images into noise. The forward process is fixed and needs no learning. The reverse process is a neural network that, given a noisy input and the noise level, predicts the noise so it can be removed a little at a time.

The first half of this lab runs the real forward process on a 64 × 64 image. The second half runs the real reverse process on 2D points, where you can watch every sample. Use the exact denoiser first: it is the mathematically optimal answer for the training points, so any imperfection you see comes from the sampler, not from training. Then train the small network and compare.

The Generative AI lesson puts diffusion next to language models, and the Transformer Explorer shows the architecture most modern diffusion models now use as their denoiser.

Forward process

Each step adds a little Gaussian noise: q(x_t | x_(t−1)) = N(√(1 − β_t)·x_(t−1), β_t·I). Because Gaussians compose, you can jump straight to any step:

That closed form is what makes training cheap: pick a random image, a random t and a random ε, build x_t in one line, and ask the network to recover ε. No simulation of the chain is needed.

The schedule decides how fast signal disappears. With the linear schedule from Ho et al. (2020), ᾱ drops below 0.08 by t = 500, so the second half of the chain is almost pure noise. Nichol and Dhariwal (2021) proposed the cosine schedule so that information is destroyed more evenly; compare the SNR curves.

Reverse process

To generate, start from x_T ~ N(0, I) and step backwards. At each step the denoiser predicts ε, which gives an estimate of the clean point x̂₀ = (x_t − √(1 − ᾱ_t)·ε̂) / √ᾱ_t. DDPM then samples x_(t−1) from the Gaussian posterior q(x_(t−1) | x_t, x̂₀), adding a little fresh noise. DDIM (Song et al. 2021) takes the deterministic path through the same x̂₀ and works well with far fewer steps.

Predicting ε is equivalent to estimating the score, the gradient of log-density: ∇ log p(x_t) = −ε̂ / √(1 − ᾱ_t). That is the arrow field you can switch on: at high t it points to the centre of the data; at low t it points to the nearest part of the shape.

The memorisation lesson. For a finite training set the best possible denoiser is known exactly, and it only ever reproduces training examples. Diffusion models generalise because neural networks can’t represent that exact solution; Kadkhodaie et al. (2024) show generalisation emerges from the network’s inductive biases. When a model is large relative to its data it can drift toward memorisation, and Carlini et al. (2023) extracted near-copies of training images from Stable Diffusion.

Guidance

A conditional model learns ε(x, c) for a class or text prompt c. Classifier-free guidance (Ho and Salimans 2022) trains the same network to also work without c (by dropping the condition some of the time) and at sampling time extrapolates:

w = 0 ignores the condition, w = 1 is ordinary conditional sampling, and w > 1 pushes samples further toward what distinguishes the class. Try class A on the moons with w = 1 and then w = 5: samples crowd into the most class-typical region and variety shrinks. Text-to-image systems use the same trade: higher guidance follows the prompt more literally, at the cost of diversity and, when pushed hard, saturated images.

Latent diffusion and today

Running diffusion on every pixel of a large image is expensive. Latent diffusion (Rombach et al. 2022, the basis of Stable Diffusion) first trains an autoencoder that compresses an image into a much smaller latent grid, runs diffusion there, and decodes at the end.

The denoiser has shifted from U-Nets to transformers: Peebles and Xie (2023) introduced diffusion transformers (DiT), which treat latent patches as tokens. Training objectives have also moved toward flow matching (Lipman et al. 2023) and rectified flow, which learn a velocity field along straighter paths from noise to data; Stable Diffusion 3 (Esser et al. 2024) uses rectified flow with a transformer backbone. The core idea you see in this lab, learning to turn noise into data step by step, is unchanged.

References

• Sohl-Dickstein, J. et al. (2015). Deep unsupervised learning using nonequilibrium thermodynamics. ICML.
• Ho, J., Jain, A., Abbeel, P. (2020). Denoising diffusion probabilistic models. NeurIPS.
• Song, J., Meng, C., Ermon, S. (2021). Denoising diffusion implicit models. ICLR.
• Nichol, A., Dhariwal, P. (2021). Improved denoising diffusion probabilistic models. ICML.
• Ho, J., Salimans, T. (2022). Classifier-free diffusion guidance.
• Rombach, R. et al. (2022). High-resolution image synthesis with latent diffusion models. CVPR.
• Peebles, W., Xie, S. (2023). Scalable diffusion models with transformers. ICCV.
• Lipman, Y. et al. (2023). Flow matching for generative modeling. ICLR.
• Esser, P. et al. (2024). Scaling rectified flow transformers for high-resolution image synthesis. ICML.
• Carlini, N. et al. (2023). Extracting training data from diffusion models. USENIX Security.
• Kadkhodaie, Z. et al. (2024). Generalization in diffusion models arises from geometry-adaptive harmonic representations. ICLR.