Diffusion models, visually

When DALL·E 2 came out I wrote that I understood diffusion about as well as I’d understood backprop before I did it by hand, which is to say I could explain the picture but not the details. That stayed true for almost two years, mostly because the picture was enough to use the tools. Then on the 15th OpenAI showed Sora, and I watched the clip of the woman walking down a wet Tokyo street at night about six times, and decided the picture wasn’t enough anymore.

So I spent a couple of evenings with the 2020 DDPM paper (Ho, Jain and Abbeel, “Denoising Diffusion Probabilistic Models”), a notebook, and a lot of scratch paper. This is what I ended up with. It’s the image version first, because that’s where all the ideas are, and then what Sora seems to add for video, with the caveat that OpenAI hasn’t said much.

The trick in one sentence

Making an image from nothing is hard. Making an image slightly less noisy is much easier. A diffusion model is a network trained to do the easy thing, applied a few dozen to a thousand times in a row.

Everything else is bookkeeping about how much noise, at which step.

Going forward: destroying an image

Start with a real image. Think of it as a long list of numbers, one per pixel per color channel, scaled to between −1 and 1. Call it x₀.

The forward process adds a little Gaussian noise, many times. In the DDPM paper it’s 1,000 steps, with the amount added at each step growing slowly from 0.0001 to 0.02 (that’s the variance, β, if you’re reading along). After 1,000 steps there’s nothing left of the image and it’s indistinguishable from random noise.

The nice part, and the part that makes the whole thing practical, is that you don’t have to actually run the 1,000 steps to get to step t. Because adding Gaussian noise to Gaussian noise gives you Gaussian noise, there’s a closed form that jumps straight there:

x_t = √ᾱ_t · x₀ + √(1 − ᾱ_t) · ε

where ε is fresh standard Gaussian noise (one random number per pixel), and ᾱ_t is a number between 0 and 1 that you get by multiplying (1 − β) for every step up to t. It’s just a schedule, worked out once, and it falls from almost 1 at the start to almost 0 at the end. With the DDPM schedule I get ᾱ ≈ 0.90 at step 100, about 0.66 at step 200, 0.079 at step 500, and 0.00004 at step 1,000.

The two coefficients always have squares that add to 1, so the result keeps roughly the same scale while the mix shifts from image to noise. Let me do one pixel. Say x₀ = 0.5 (a fairly bright red channel), and pick a step where ᾱ is 0.64, which is close to step 200 and makes the square roots round: √0.64 = 0.8 and √0.36 = 0.6. The random draw for this pixel is ε = −1.2. Then

x_t = 0.8 × 0.5 + 0.6 × (−1.2) = 0.4 − 0.72 = −0.32

The pixel that was fairly bright is now fairly dark, because this particular draw of noise pushed it down. Do that for every pixel and you get a recognizable but grainy image. At step 500 the image coefficient is about 0.28, and at 750 it’s about 0.06, which is a ghost.

Nothing in the forward direction is learned. It’s arithmetic and a random number generator.

Training: guess the noise

The whole training loop is short enough to list, which stunned me a bit:

1. Take an image x₀ from the training set.
2. Pick a random step t between 1 and 1,000.
3. Draw random noise ε, and make x_t with the formula above.
4. Give the network x_t and t, and ask it to predict ε.
5. The loss is the squared difference between the noise it predicted and the noise you actually added.

That’s it. The network isn’t asked to produce an image. It’s asked, “here’s a noisy picture and here’s how noisy I told you it is, which part is the noise?” You can make an infinite number of training examples from each image, one for every combination of step and noise, and the loss is plain mean squared error, the same thing I fitted linear regressions in R with. The paper shows that this simple version, which drops a weighting term from the more principled objective, actually works better.

Going back to my pixel: if the network sees −0.32 at that step and predicts the noise was −1.2, you can invert the formula to estimate the original:

x̂₀ = (x_t − 0.6 × ε̂) / 0.8 = (−0.32 + 0.72) / 0.8 = 0.5

Perfect recovery, if the guess is perfect. The guess is never perfect.

Sampling: why it takes many steps

If you can estimate x₀ from any x_t, why not start from pure noise and jump straight to the answer in one step?

Because at high noise the network can’t know. If all you have is static, the best guess for “what was the image” in the squared-error sense is the average of all the images that could have produced that static, and the average of many images is a brown blur. Ask a well-trained model to go from pure noise to an image in one step and that’s roughly what you get.

So you take a small step instead. Starting from pure noise at t = 1,000, you predict the noise, remove a portion of it to get to something like step 999, add back a little fresh random noise (which feels wrong, but it keeps the sampler exploring instead of collapsing to that average), and repeat. At each step the estimate is a little sharper, so the next prediction is better. Early steps settle the big decisions, where the horizon is and roughly what colors go where. Late steps fill in texture and edges.

The DDPM paper did all 1,000 steps at sampling time, which is slow. A lot of the work since has been about doing it in fewer. DDIM, from later in 2020, showed you could take big strides through the schedule deterministically, and in practice tools like Stable Diffusion use something like 20 to 50 steps. Each step is a full pass through a network with hundreds of millions of parameters, which is why generating an image takes seconds, not milliseconds.

Where the words come in

None of that says anything about text. For “an astronaut riding a horse”, the network gets a second input: an embedding of the prompt from a text model, fed in through cross-attention layers, so at every step the noise prediction depends on the prompt as well as the image.

The surprising part is a trick called classifier-free guidance (Ho and Salimans, 2021). During training, you sometimes blank out the prompt, so the same network learns both “predict the noise given this prompt” and “predict the noise given nothing”. At sampling time you run it both ways at every step and push past the prompted answer, away from the unprompted one:

ε̂ = ε̂_uncond + w × (ε̂_cond − ε̂_uncond)

With w = 1 you just get the prompted prediction. Stable Diffusion’s default is 7.5. So if, for some pixel, the unprompted prediction is 0.10 and the prompted one is 0.30, the guided prediction is 0.10 + 7.5 × 0.20 = 1.60, way past either. It exaggerates whatever the prompt changed. Images get more on-topic and more saturated, and if you push w too far they look burnt, which anyone who’s played with the slider has seen.

Latents: doing it somewhere smaller

Running this on raw pixels is expensive. A 512 × 512 color image is 786,432 numbers, and you’d be pushing all of them through a big network 50 times.

The idea behind Stable Diffusion (the latent diffusion paper, Rombach and others, 2022) is to first train a separate autoencoder that compresses an image into a smaller grid and back. For Stable Diffusion that grid is 64 × 64 with 4 channels, 16,384 numbers, 48 times fewer. The diffusion happens entirely in that compressed space, and the decoder turns the final result back into pixels once at the end. The autoencoder handles the fine pixel detail, and the diffusion model only has to learn composition, which is where the interesting part is anyway.

And then video

Which brings me back to Sora. What OpenAI published is a technical report, titled “Video generation models as world simulators”, that describes the approach at a high level and explicitly leaves out the model and implementation details. From that report:

It’s a diffusion model, working in a compressed latent space like Stable Diffusion, except the compression network squeezes video in time as well as space. The compressed video is then cut into “spacetime patches”, little blocks a few latent pixels wide and a few frames deep, and each patch becomes a token. The denoising network is a transformer, not the U-Net that most image models use, which they say scales well with more compute. (The idea of a diffusion transformer comes from a 2022 paper by William Peebles and Saining Xie, and Peebles is one of the Sora leads.)

The patch idea is the bit I like best. To make the numbers concrete, with entirely invented sizes since OpenAI doesn’t give any: a latent video 16 frames long and 32 × 32 latent pixels per frame, cut into patches 2 frames deep and 2 × 2 across, is 8 × 16 × 16 = 2,048 tokens. A transformer doesn’t care whether those tokens came from a square image, a wide video or a tall one, so they can train on clips of different lengths and shapes without cropping everything to one size, which they say helps composition. An image is just a video one frame long.

Text prompts get the DALL·E 3 treatment: they trained a captioning model to write detailed descriptions of the training videos, and at use time GPT expands your short prompt into a long, detailed one before Sora sees it.

What I can’t tell you is how it does the thing that impressed me, which is that the woman’s reflection in the puddles stays consistent as she walks, and things that go off screen come back looking the same. The report says those properties emerge from scale, without anything built in for 3D or objects. I’d want to see a lot more unfiltered output before I believed that completely. OpenAI’s announcement and the report are both frank about failures, with examples where physics doesn’t hold, like glass not shattering properly, or a bite taken out of something that doesn’t leave a mark. Nobody outside OpenAI and a group of red teamers and artists can use it yet, so every clip anyone’s seen was chosen by someone at OpenAI, the same caveat I had about DALL·E 2.

That caveat aged reasonably well, for what it’s worth. I also guessed back then that the garbled text in images would get fixed “in a version or two”, and DALL·E 3 is noticeably better at it, though still far from perfect. I’m not going to guess about video yet. Stability AI announced Stable Diffusion 3 last week, also built on a diffusion transformer, so the architecture at least seems to be where everyone’s heading.