An Introduction to Diffusion Models
This post is an introduction to how diffusion models work from scratch, covering everything from the forward noising process to SDEs/ODEs.
The Generative Problem
The core of the diffusion generative problem is this: you have a dataset of examples, and you want to generate a new one that looks like it came from the same source. Formally, sample from the data distribution $p(x)$ without ever writing $p$ explicitly. Instead of generating in one shot, diffusion destroys structure gradually and learns to rebuild it.
What does generative mean? Two things are often conflated for generative.
- Density estimation: Given $x$, return $p(x)$
- Sampling: Produce a fresh $x \sim p_{data}$
We want sampling, though density could also be useful.
Any probability distribution can be written in terms of its energy density.
Why is Generation Hard? For high dimensional x, finding that $Z$ term is intractable. Even if you can get E(x), you can't get p(x) and sample from it easily. According to the manifold hypothesis, real data doesn't fill the full dimensionality on the space, but rather lives on a thin, curved, low dimensional sheet in the space. A model should find that thin sheet, not fill the entire space.
Generative Trilemma Until diffusion models, each family could only pick 2 out of the 3 following qualities: High sample quality, mode coverage / diversity, fast sampling.
- Autoregressive (like transformers): Good quality, good coverage, but sequential and slow
- VAEs: Fast and Diverse, but blurry and mode averaged output
- GANs: Suffer from mode collapse
How diffusion overcomes these issues Diffusion overcomes these issues via 2 main mechanisms:
By aggregating these mini steps in direction of steepest ascent, after a series of large N steps, the resulting design will be close to the distribution.
Forward Process (noising)
Noising is the part of the process that doesn't require any learning. It is a relatively straightforward process of injecting Gaussian noise into some training data.
The noising step looks like this: $X_t = \sqrt{1 - \beta_t} x_{t-1} + \sqrt{\beta_t} \epsilon$, $\epsilon \sim N(0, I)$
Notice how we scale down the previous data by an equivalent factor. This allows us to compose these steps without the data exploding to infinity. The reason we use Gaussians is because Gaussians compose, allowing us to do a simple trick.
This means that we can jump directly from the clean data ($x_0$) to any noise level t in a single line.
This allows for a ton of free training data very cheaply, without having to run the full chain. Every clean data point becomes a full set of noisy data at every noise level.
Reverse Process (denoising)
The only learned part of this whole process is the denoising.
The noising process can be modeled as: $x_0 \rightarrow x_1 \rightarrow \dots \rightarrow x_T$
The denoising is the learned model that can perform: $X_T \rightarrow \dots \rightarrow X_1 \rightarrow X_0$
$x_T$ is approximately Gaussian noise. The denoising model's task can be summarized as literally just predicting the noise that should be removed from one step to another, and subtracting it. At each step, additional random noise is added to add some stochasticity, but the first two steps are the core.
This is what the model is predicting: for some given sample $x_t$ at some step $t$, what was the noise vector added: $\epsilon_\theta(x_t, t)$
Next, we calibrate this to counteract both scales from the noising steps. This is to ensure unit variance again. $\mu_\theta(x_t, t) = \frac{1}{\sqrt{\alpha_t}} \left(x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}} \varepsilon_\theta(x_t, t)\right)$
Finally, after each denoising step, we add some noise.
Note that when training, $x_0$ is a term of the loss function, but never an input to the denoising model itself. This allows it to generate even at inference time when $x_0$ is inferred.