Generative AI papers

Part 8 of 8 in Generative AI papers

Diffusion models and generation by denoising

Start with an image that is pure noise. Train a network to predict and remove the noise, one step at a time. That is how diffusion models generate images, and why they are slow but scalable.

Language models are one branch of generative AI. Diffusion models are the other, and they work by a different principle. Instead of predicting the next token, they start with noise and predict what was added to an image, so they can subtract it and get something that looks more like a real image. Iterate that enough times and you have generated a new image from nothing but noise and a learned model.

The forward process: adding noise

Take an image and add a small amount of Gaussian noise to it. The image is still mostly recognizable. Add more noise, and gradually the image becomes harder to pick out. Keep adding noise, carefully scaled so the changes are smooth, and after enough steps the image is indistinguishable from pure random noise. That is the forward process, described in Denoising Diffusion Probabilistic Models by Ho, Jain and Abbeel.

The paper adds noise in a carefully designed way. Each step scales the signal by a coefficient less than 1 and adds noise by a coefficient close to 1. After a thousand steps, the original image has been multiplied down to nearly nothing and the noise has been scaled up, so what remains is almost pure Gaussian noise. The forward process is not learnable. It is a fixed mathematical transformation.

The reverse process: training to predict noise

Now train a neural network to undo it. The training step is simpler than it sounds, and getting it exactly right matters, because the obvious guess is wrong.

Take a clean image. Draw a step number t at random, say 500. Draw a noise sample. Use the forward process to jump straight to the noised image at step 500 in one shot, which the mathematics allows: you do not have to walk through 500 steps to get there. Hand the network that noisy image and the number 500, and ask it to output the noise that was added.

Note what is being predicted. It is not the small increment between step 499 and step 500. It is the whole noise component separating this image from the clean original, and the network is told which step it is looking at so it knows how much to expect. The loss is the squared difference between the noise it predicted and the noise actually drawn.

The training signal is effectively unlimited. Every image, crossed with every step count, crossed with every noise draw, is another valid training pair, all of it self-supervised, with no labels required anywhere.

At generation time, start from pure noise. Ask the network what noise it sees, remove a scaled fraction of it rather than all of it, and add a small amount of fresh noise back in. Repeat.

Both of those details carry weight. Subtracting the whole predicted noise in one move lands on the network's blurred average of every image consistent with that input, which looks like a smudge rather than a photograph. Taking a small step and re-noising keeps the sample on the path the forward process actually traced. The added noise is also what makes generation stochastic: run it twice from different starting noise and you get two different images, which is the point of a generative model.

This is the whole idea. The network learns the structure of image distribution by learning to reverse a specific noise process. At test time, you sample from the noise distribution and let the network walk you backward to realistic images. The power of the approach comes from the fact that the noise prediction task, though seemingly simple, requires the network to learn all the structure in the image distribution. If you train a network only to remove random Gaussian noise, it has to learn what images look like in order to do that task well.

Why it works

The forward process destroys information gradually. A network that learns to undo that destruction step by step learns what images look like, because reversing noise destruction is equivalent to learning the distribution of images. The denoising task is a way of exploring the space of realistic images.

Compare this to a language model, which predicts one token at a time by learning to complete any prefix. A diffusion model predicts noise, one step at a time, learning to complete any partially-noised image. The problems have the same form: learn a distribution by predicting the next piece. The difference is that a language model predicts forward in sequence space and a diffusion model predicts backward in noise space.

Latent diffusion and computational affordability

The original diffusion model works on pixels. A 512 by 512 image has 262,144 pixels. A diffusion step over all 262,144 pixels is expensive. Each step runs a neural network over the entire image, a thousand times. Rombach and colleagues put the cost plainly when motivating their own work: optimising powerful pixel-space diffusion models "often consumes hundreds of GPU days". That is out of reach for most people who might want to train one.

Rombach and colleagues introduced latent diffusion models, which run diffusion in a compressed space instead of pixel space. Take a powerful pretrained autoencoder, a network that has already learned to compress images to a small representation and decompress them back to pixels. The autoencoder is trained separately, on a large dataset, to minimize reconstruction error. Once trained, freeze its weights and use it as a fixed tool.

Compress a 512 by 512 image by a factor of 8 in each spatial direction using the autoencoder encoder. The result is a 64 by 64 latent representation, which is 4,096 positions. Divide the original 262,144 pixels by 4,096: the result is 64. The compressed space has 64 times fewer positions per step. When the diffusion model runs in latent space, each step touches 4,096 positions instead of 262,144. After a thousand steps, the computation savings are enormous.

That factor of 64 in positions per step is the difference between a model most people cannot train and one they can. Rombach and colleagues report that latent diffusion achieves state-of-the-art performance on image inpainting and competitive results on unconditional generation, scene synthesis and super-resolution.

The cost is that compression discards information. The autoencoder compresses to save compute, so detail finer than what the latent code can represent is gone. You cannot recover detail the compressed code never stored. High-frequency details in images are the first casualty. A latent diffusion model trained to generate realistic 512-by-512 images will never generate pixel-perfect fine-grained textures because the latent space does not carry that information.

Whether that trade is worth making depends on the use. For generating an illustration it plainly is. For work where fine texture is the content, such as medical imaging or scientific measurement, discarding high-frequency detail to save compute is not a neutral choice, and a generated detail that was never in the latent code is an invention rather than a reconstruction.

The autoencoder discards information, but it does so in a way that preserves the semantic content of images. The latent code still carries enough information about objects, composition, colors and high-level structure for a diffusion model trained to predict noise in latent space to generate recognizable images. When you decompress the final latent code back to pixels using the decoder, you get a complete 512 by 512 image. Some fine detail is lost because the latent code is compressed, but the important structure is preserved.

Steering with text: cross-attention

A diffusion model that takes only noise at input cannot be steered by a user. Rombach and colleagues introduce cross-attention layers into the diffusion network. Cross-attention works the same way as the self-attention described in Attention is a lookup you can learn, but instead of a position attending to other positions in the same image, the denoising network attends to embeddings of the text prompt.

The text is encoded into a set of embeddings, one per word or token. The denoising network, at each step, uses cross-attention to look up which parts of the text prompt are relevant to each region of the image it is generating. A prompt like "a red cat sitting on a blue chair" becomes embeddings for red, cat, sitting, blue, chair, and other words. When the network is denoising the top-left region of the image, it can attend to the embeddings for red and cat and chair to figure out what should be there.

The mechanism is not mysterious. The query is generated from the image being denoised. The keys and values are generated from the text embeddings. Attention scores all keys and weights all values, so the image generation process mixes information from every word. A region of the image can attend to multiple words, and a word can be attended to by multiple regions. The result is that the text guides the generation without scripting exactly which pixels come from which words.

Without cross-attention, a diffusion model generates images but has no way to accept user input about what to generate. With cross-attention, the model can be steered by any conditioning signal that can be embedded: text, semantic layouts, edge maps, or anything else. The flexibility comes from the generality of attention as a mechanism. Because attention is just a learned lookup, it works for any type of key and value, not just image positions looking at other image positions.

Speed and quality trade-offs

Generation is slow, and the slowness is structural rather than an implementation detail. Every step is a full forward pass through the denoising network, and the original formulation used on the order of a thousand of them for a single image. A language model producing a sentence of thirty tokens runs thirty forward passes. A diffusion model producing one image may run thirty times that.

Later work reduced the step count substantially, using samplers that take larger steps without falling off the path, and by distilling a many-step model into a few-step one. The count came down by more than an order of magnitude. It did not come down to one, because the iteration is the method.

The slowness buys something specific: a stable training objective. Diffusion training minimises a squared error against a known target, which is about as well behaved as an objective gets. The previous leading approach, adversarial training, optimised a generator against a discriminator that was itself still learning, and that moving target made training fragile. Diffusion moved the difficulty out of training, where it costs researcher time and failed runs, and into sampling, where it costs compute that can be optimised later. That is a good trade, and it is much of why the field switched.

Why diffusion matters for scale

The latent diffusion work connects to Scaling laws and the Chinchilla correction because the compute saving makes scale feasible. Training a diffusion model at scale requires affordable compute per step. Compressing to latent space divided by 64 is what enabled billion-parameter diffusion models trained on hundreds of millions of images. Without that compression, the cost would have been prohibitive.

Latent diffusion also connects to Embeddings are coordinates for meaning, because the power of the cross-attention steering comes from embeddings. A language model tokenizes text and embeds tokens into vectors. Diffusion models tokenize text the same way and use those embeddings to steer the image generation process. The embedding space, learned from a large language model or a text encoder, is what carries the semantic meaning from text into the image generation loop.

Forward process: add noise at each step until the image is unrecognizable. Reverse process: predict and remove noise at each step, guided by text embeddings, until you have a new image.

What diffusion does not do

Diffusion models do not think step by step like Chain of thought and what it does not prove. Each denoising step is not a reasoning step. It is a mechanical transformation guided by learned weights. The network has no goals or intermediate reasoning states. It predicts noise and the noise is subtracted, purely mechanically.

Diffusion models also do not store information in the order dimension the way language models do. A language model's tokens are ordered; a prompt is a sequence. An image has spatial structure, not sequence structure. The diffusion network attends to different spatial regions in parallel, not sequentially. This makes diffusion useful for images and problematic for text, where order matters.

The slowness is real and matters. Generating an image takes seconds to minutes on a consumer GPU. Generating text takes milliseconds. That is not a minor difference for user-facing applications. Systems built on diffusion have to plan for latency where text systems do not.

But the stability and scalability have made diffusion the standard for image generation. The mechanism is learnable, the architecture is straightforward, and the process works reliably at any size of the model or dataset. Those properties outweigh the speed cost for most applications.

References

  1. Denoising Diffusion Probabilistic Models. Jonathan Ho, Ajay Jain and Pieter Abbeel, arXiv, 2020.
  2. High-Resolution Image Synthesis with Latent Diffusion Models. Robin Rombach and colleagues, arXiv, 2021.
  3. Attention Is All You Need. Ashish Vaswani and colleagues, arXiv, 2017.

All insights

Working on something like this?

If this is close to something you are trying to solve, tell us where you have got to and we will say what we would test first.

Book a discovery call