This post gives a gentle introduction to the diffusion models, a family of the state of art generative model.

Theory

Diffusion model can be formulated as a Hierarchichal Variational Encoder (HVAE). To optimize variational auto-encoders, the critical technique is the use of evidence lower bound (ELBo) to estimate otherwise intractable distribution.

Eveidence Lower Bound (ELBo) and Vartiational Auto-encoder (VAE)

Given a variable $x$ and its latent variable $z$, $p(x) = \int p(x, z) dz$ However, the integration over all latent variables is intractable, so we introduce an encoder parametrized by $\phi$, $q_\phi(z| x)$, to derive the evidence lower bound,

For the variational auto-encoder, we also introduce a decoder parametrized by $\theta$, $p_\theta(z|x)$, so the training objective becomes:

In the training procedure, given a sample $x$, we obtain the encoding $z$ by sampling from $q_{\phi}(z|x)$, and reconstruct a $x'$ by sampling from $p_\theta(z|x)$. The first term in ELBo is a reconstruction term, and the second term is a prior matching term. We usually assume $z$ follows a simple distribution like a Gaussian distribution. In actual optimization, we use Montecarlo samples to estimate the expectation value.

Hierarchical VAE (HVAE)

One problem of VAE is that its prior distribution is usually very simple. To model more complex distributions, we can introduce $T$ latent variables $z_1, \dots, z_T$.

We can obtain a similar ELBo with Jensen’s inequality,

Diffusion Model as a HVAE

The most intuitive way to understand diffusion models is the process of gradually denoising a data sample into Gaussian noise, and the reverse process generates new data samples from Gaussian noise. Mathematically, the easist formulation of diffusion models is a special HVAE. In particular, instead of usual dimensionaly reduction in HVAE, each $z_1, \dots, z_T$ has the same dimension as $x$. Therefore, we denote $x$ with $x_0$ instead, and $z_1, \dots, z_T$ with $x_1, \dots, x_T$. Further, we assume the encoder follows a fixed Gaussian distribution,

where $\alpha_t$ is a hyper-parameter.

Noising Encoder

We will show that after multiple encoding steps, the encoded sample is simply the original sample plus some Gaussian noise. Using the reparamatrization trick, we can write

where $\epsilon_0 \sim \mathcal N(\epsilon_0; 0, 1)$

By applying the encoder repeatedly, one can show that

ELBo for Diffusion Model

Now we rewrite ELBo to decompose it into more intuitive terms,

Now the first term is the reconstruction term, which can be approximated and optimized using a Monte Carlo estimate. The second term is the prior matching term, but since our encoder and prior distribution is fixed, there are no parameters to optimize here. The new third term is the consistency term. With smarter term collection and cancellation, we decompose the ELBo differently so that each term is only an expectation over one random variable.

With the new derivation, we can understand the third term as a noise matching term, where $p_{\theta}(x_{t-1}| x_t)$ is estimating the noise added by $x_t$ and approximating the GT denoising fucntion $q(x_{t-1}| x_t, x_0)$ which has access to the data. We will now rewrite $q(x_{t-1}| x_t, x_0)$ to derive different training objectives of diffusion models.

1st Training Objective: Reconstruct Data Sample

We’ve already derived $q(x_t | x_0)$, which is

By Bayes rule and using the above fact with a lot of algebra, we can show that

Set our decoder to be $p_\theta(x_{t-1}|x_t) = \mathcal N(x_{t-1}; \mu_\theta(x_t, t), \sigma_q^2(t)I)$, then we can use the DL divergence for Gaussians to derive,

Finally, instead of predicting the mean $\mu_\theta(x_t, t)$, we can directly predict the sample $\hat x_\theta(x_t, t)$ and set the mean to match $\mu_q(x_t, x_0)$

Now substitue this assumption into the DL divergence term,

Thus, the training objective becomes to predict the original data sample from a noisified version after $t$ steps.

2nd Training Objective: Predicting Error

The most popular interpretation and the training objective in DDPM (Denoising Diffusion Probablistic Model), however, is to predict the noise in each denoising step. To derive this equivalent training objective, we need to rewrite $\mu_q(x_t, x_0)$ by substituting in $x_0 = \frac{x_t-\sqrt{1-\bar\alpha_t}\epsilon_0}{\sqrt{\bar\alpha_t}}$.

Now instead of predicting the sample $\hat x_\theta(x_t, t)$, we predict the noise given a noised sample at time step $t$, $\hat\epsilon(x_t, t)$, and set

Now substitue this assumption into the DL divergence term,

Thus, the training objective becomes to predict the noise at a time step $t$.

3rd Training Objective: Predict the Score Function

In my opinion, the most profound understanding of diffusion models is that it is training a score function. Roughly speaking, the score function predicts the gradient field of a data distribution, so given any point, it predicts the direction to move into dense regions. To derive this formulation, we make use of the Tweedie’s Formula. Given $z \sim \mathcal N(z; \mu_z, \Sigma_z)$,

Applying this formula, we have that

We’ve drived that $x_t \sim \mathcal N(x_t; \sqrt{\bar\alpha_t} x_0, (1-\bar\alpha_t)I)$, which gives us $E[\mu_{x_t}|x_t] = \sqrt{\bar\alpha_t}x_0$. Thus, we can rewrite $x_0$ as

Substitute this back into $\mu_q(x_t, x_0)$,

Now instead of predicting the sample $\hat x_\theta(x_t, t)$ or the noise $\hat\epsilon(x_t, t)$, we predict the score function $s_\theta(x_t, t)$ to match $\nabla_{x_t} \log p(x_t)$,

Now substitue this assumption into the DL divergence term,

Guidance for Generation

Now everthing is great. We can either train a neural network to predict sample from Gaussian samples, or to predict noise and denoise, or to predict the score function. What if we want to guide the generation given some prior information, say we want to generate a cat but we train on all kinds of animals? The two main guidance methods are classifer guidance and classifer-free guidance.

Classifer Guidance

More concretely, given some prior $y$ to condition on, we want to generate a sample according to $p(x_0 | y)$. Let’s decompose this using the Bayes’ rule,

and the score function is

The first term is the usual unconditional score, and the second term is an adversarial gradient of a classifer $p(y|x_t)$. Thus, in practice, we can train a classifier for each class $y$ to guide the generation process.

To gain fine-grained control over the guidance, we can introduce a hyperparameter $\gamma$ to scale the adversarial gradient,

Classifer-free Guidance

Nevertheless, it’s costly and even impossible to train a classifer for each new condition. To introduce classifer-free guidance, we rewrite the scaled conditional score,

The first term is the conditional score, and the second term is the unconditional score. In practice, we can train a neural network that always takes in a conditional imformation (which is empty for unconditional score).

Toy Implementation

For a toy impelmentation, you can check this colab notebook.

References

[1] Luo, Calvin. “Understanding diffusion models: A unified perspective.” arXiv preprint arXiv:2208.11970 (2022).

[2] Song, Yang, et al. “Score-based generative modeling through stochastic differential equations.” arXiv preprint arXiv:2011.13456 (2020).

[3] Song, Jiaming, Chenlin Meng, and Stefano Ermon. “Denoising diffusion implicit models.” arXiv preprint arXiv:2010.02502 (2020).

[4] Ho, Jonathan, Ajay Jain, and Pieter Abbeel. “Denoising diffusion probabilistic models.” Advances in neural information processing systems 33 (2020): 6840-6851.

[5] Ho, Jonathan, and Tim Salimans. “Classifier-free diffusion guidance.” arXiv preprint arXiv:2207.12598 (2022).