Diffusion Models in generative AI

Module 12 · Generative AI · 18 concepts

In depth / 15 lines

Diffusion Models, in more detail

Diffusion models learn to undo a process of gradual destruction.

Training takes a real image and adds a little noise, then a little more, over many steps.

After enough steps the image is pure noise with nothing left of the original.

That forward process needs no learning, since it is just adding noise on a schedule.

The model learns the reverse: given a noisy image, predict the noise that was added.

Subtract the predicted noise and you have a slightly cleaner image.

Generation starts from pure random noise and repeats that step many times.

Each step removes a little noise and the image emerges gradually.

The genius is turning one impossible jump into many easy small ones.

Predicting a whole image from noise is hard, and predicting a little noise is not.

Training is far more stable than a GAN, since there is no adversary to balance.

The cost is that generation needs many passes, historically fifty or more.

Faster samplers and distillation have cut that to a handful of steps.

Latent diffusion runs the whole process in a compressed space, which made it affordable.

Stable Diffusion, DALL-E and Midjourney are all built on this idea.

Want the full technical article? Wikipedia ↗

Deep dive

What is underneath Diffusion Models

Destroy the image in known steps, then learn to undo one

q(\mathbf{x}_t \mid \mathbf{x}_{t-1}) = \mathcal{N}\big(\sqrt{1 - \beta_t}\, \mathbf{x}_{t-1}, \; \beta_t I\big)

$\beta_t$ : the variance of the noise added at step $t$ , following a fixed schedule.

The forward process is not learned. It adds a little Gaussian noise at a time until, after $T$ steps, nothing of the image remains.

q(\mathbf{x}_t \mid \mathbf{x}_0) = \mathcal{N}\big(\sqrt{\bar{\alpha}_t}\, \mathbf{x}_0, \; (1 - \bar{\alpha}_t) I\big), \qquad \bar{\alpha}_t = \prod_{s=1}^{t} (1 - \beta_s)

Because Gaussians compose, any step can be reached in one jump, so training never has to run the chain.

\mathbf{x}_t = \sqrt{\bar{\alpha}_t}\, \mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\, \boldsymbol{\epsilon}

That line is the whole training data generator: pick an image, pick a $t$ , pick a noise, and you have an input and its answer.

The objective is a regression

\mathcal{L} = \mathbb{E}_{\mathbf{x}_0, \, t, \, \boldsymbol{\epsilon}} \Big[ \big\| \boldsymbol{\epsilon} - \boldsymbol{\epsilon}_{\theta}(\mathbf{x}_t, \, t) \big\|^{2} \Big]

The network is given a noisy image and the timestep, and asked to predict the noise that was added.

It is mean squared error. There is no adversary, no divergence and no sampling inside the loss, which is why diffusion trains stably where GANs do not.

Predicting the noise rather than the clean image is a reparameterisation of the same problem, and it happens to be far better conditioned across the range of $t$ .

Sampling, in order

Start from pure noise $\mathbf{x}_T \sim \mathcal{N}(\mathbf{0}, I)$ .
Predict the noise in the current image with the network.
Subtract a scaled portion of it to get a slightly cleaner image.
Add a small amount of fresh noise, except at the last step.
Repeat down to $t = 0$ .

\mathbf{x}_{t-1} = \frac{1}{\sqrt{\alpha_t}}\left( \mathbf{x}_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\epsilon}_{\theta}(\mathbf{x}_t, t) \right) + \sigma_t \mathbf{z}

Step four is not a mistake. The reverse process is stochastic, and removing all the noise at once lands on a blurred average rather than a sample.

The cost, and what removed it

T = 1000 \implies 1000 \text{ forward passes per image}

Generation is a loop over the network, which is why a diffusion model is thousands of times slower per sample than a GAN.

Faster samplers such as DDIM need only 20 to 50 steps.

Deterministic samplers reinterpret the same trained model as an ordinary differential equation and solve it in far fewer, larger steps.

\text{latent diffusion} : 512 \times 512 \times 3 \to 64 \times 64 \times 4

The other halving of cost is running the whole process inside an autoencoder's latent space, which is what Stable Diffusion is and why it fits on one card.

Conditioning, in one line

\tilde{\boldsymbol{\epsilon}} = \boldsymbol{\epsilon}_{\theta}(\mathbf{x}_t, \varnothing) + s \big( \boldsymbol{\epsilon}_{\theta}(\mathbf{x}_t, c) - \boldsymbol{\epsilon}_{\theta}(\mathbf{x}_t, \varnothing) \big)

Classifier-free guidance runs the network twice, once with the prompt and once without, and extrapolates away from the unconditional prediction.

s \approx 7

The guidance scale is the dial every image tool exposes. Higher follows the prompt harder and loses diversity, and past about 15 it produces oversaturated, broken images.

Test yourself

2 questions on this card. Every option explains itself once you pick, right or wrong.

Question 1 of 2

How does a diffusion model generate an image?

Question 2 of 2

Why is latent diffusion faster than pixel-space diffusion?

Generative AI

18 concepts in this module, then a graded test with every answer explained.

Study this module →