Diffusion Models
Table of Contents
Autoregressive models generate one element at a time, which makes sampling sequential. VAEs trained with simple pixel-space likelihoods may produce blurry samples, while GANs can be difficult to train and may suffer from mode collapse. Diffusion models offer another approach: they learn to transform noise into data through a sequence of denoising steps.
1. Diffusion Models - An Overview from Likelihood-Based Perspective
Suppose \( \mathbf{x}_{0} \sim q_{\mathrm{data}} \) is a data sample and \( \mathbf{x}_{T} \) is approximately standard Gaussian noise. The forward process \(q\) is fixed and gradually adds noise, while the learned reverse process \(p_{\theta}\) gradually removes it.
1.1. Forward process
The forward process is more akin to training. It transforms a clean image into noise.
\[ \boxed{\mathbf{x}_{t-1}}\xrightarrow{q(\mathbf{x}_{t}\mid\mathbf{x}_{t-1})}\boxed{\mathbf{x}_{t}} \]
We choose each forward transition to be Gaussian:
\begin{equation} q(\mathbf{x}_{t}\mid\mathbf{x}_{t-1}) := \mathcal{N}(\mathbf{x}_{t};\sqrt{1-\beta_{t}}\mathbf{x}_{t-1}, \beta_{t}\mathbf{I}) \end{equation}Since we have modeled the diffusion as Markovian, we can represent the overall probability given original image:
\begin{equation} q(\mathbf{x}_{1:T}\mid\mathbf{x}_{0}) = \prod_{t=1}^{T} q(\mathbf{x}_{t}\mid\mathbf{x}_{t-1}) \end{equation}The purpose of selection of \(\sqrt{1-\beta_{t}}\) and \( \beta_{t} \) for mean and variance is solely for convenience in computation.
The reason lies in when we evaluate the variance. First, with the reparameterization trick, we can write
\[ \mathbf{x}_{t} = \sqrt{1-\beta_{t}} \mathbf{x}_{t-1} + \sqrt{\beta_{t}} \boldsymbol{\epsilon} \]
Then evaluate the variance. Applying the property of variance of linear combination.
\[ \mathrm{Var}[\mathbf{x}_{t}] = (1-\beta_{t}) \mathrm{Var}[\mathbf{x}_{t-1}] + \beta_{t} \]
If the variance of \( \mathbf{x}_{t-1} \) is \( 1 \), then the variance of \( \mathbf{x}_{t} \) is also \(1\). Thus, we may preserve the property of unit variance.
1.2. Reverse process
Similarly, the reverse process is more akin to inference. It tries to reconstruct images from noise.
\[ \boxed{\mathbf{x}_{t}}\xrightarrow{p_{\theta}(\mathbf{x}_{t-1}\mid\mathbf{x}_{t})}\boxed{\mathbf{x}_{t-1}}\]
The exact marginal reverse conditional \(q(\mathbf{x}_{t-1}\mid\mathbf{x}_{t})\) is generally intractable because it depends on the unknown data distribution. When each forward step is small, it can be approximated with a learned Gaussian transition:
\begin{equation} p_{\theta}(\mathbf{x}_{t-1}\mid\mathbf{x}_{t}) := \mathcal{N}\big(\mathbf{x}_{t-1}; \boldsymbol{\mu}_{\theta}(\mathbf{x}_{t},t), \boldsymbol{\Sigma}_{\theta}(\mathbf{x}_{t}, t) \big) \end{equation}The backward process can be modeled as
\begin{equation} p_{\theta}(\mathbf{x}_{0:T}) = p(\mathbf{x}_{T}) \prod_{t=1}^{T} p_{\theta}(\mathbf{x}_{t-1}\mid\mathbf{x}_{t}) \end{equation}2. Likelihood Modeling
Maximum Likelihood Learning would maximize \( \log p_{\theta}(\mathbf{x}_{0}) \) (so as to make \( p_{data} \) and \( p_{\theta} \) closer).
Note that, here what we want is \( p_{\theta} \), the reverse process, i.e., \( p_{\theta}(\mathbf{x}_{t}\mid \mathbf{x}_{t-1}) \). We can know \( p(\mathbf{x}_{T}) \) because it’s sampled Gaussian noise. (This is reason supporting the first equal sign.) That is, by definition
\[ p_{\theta}(\mathbf{x}_{0:T}) = p(\mathbf{x}_{T}) \prod_{t=1}^{T} p_{\theta}(\mathbf{x}_{t-1}\mid \mathbf{x}_{t}) \]
However, since during the reverse process, all \( x_{1:T} \) are randomly sampled. So we have to integral the joint probability \( p_{\theta}(\mathbf{x}_{0:T}) \) w.r.t. \( \mathbf{x}_{1:T} \), kinda like enumerating all possible transformation sequence \( \mathbf{x}_{T}\to\mathbf{x}_{T-1}\to\dots\to\mathbf{x}_{1}\to\mathbf{x}_{0} \)
But, direct marginalization over the latent variables \( \mathbf{x}_{1:T} \) is intractable, so we optimize a variational lower bound instead, using a trackable distribution over noisy trajectory. So we introduce \( q(\mathbf{x}_{1:T}\mid \mathbf{x}_{0}) \) to sample trajectories in a controlable manner.
\begin{equation} \begin{split} \log p_{\theta}(\mathbf{x}_{0}) &= \log \int p_{\theta}(\mathbf{x}_{0:T}) \,\mathrm{d}\mathbf{x}_{1:T} \\ &=\log \int q(\mathbf{x}_{1:T}\mid\mathbf{x}_{0}) \frac{p_{\theta}(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T}\mid\mathbf{x}_{0})} \,\mathrm{d}\mathbf{x}_{1:T}\\ &=\log \mathbb{E}_{q(\mathbf{x}_{1:T}\mid\mathbf{x}_{0})} \left[ \frac{p_{\theta}(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T}\mid\mathbf{x}_{0})} \right] \\ &\ge \mathbb{E}_{q(\mathbf{x}_{1:T}\mid\mathbf{x}_{0})} \left[ \log \frac{p_{\theta}(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T}\mid\mathbf{x}_{0})} \right] \\ \end{split} \end{equation}Denote the negative evidence lower bound (negative ELBO) by \(L = \mathbb{E}_{q} \left[ -\log \frac{p_{\theta}(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T}\mid\mathbf{x}_{0})} \right] \). It is an upper bound on \(-\log p_{\theta}(\mathbf{x}_{0})\). Using Bayes’ rule,
\[ q(\mathbf{x}_{t}\mid\mathbf{x}_{t-1}) = \frac{q(\mathbf{x}_{t-1}\mid\mathbf{x}_{t},\mathbf{x}_{0}) q(\mathbf{x}_{t}\mid\mathbf{x}_{0})} {q(\mathbf{x}_{t-1}\mid\mathbf{x}_{0})}, \]
the intermediate terms telescope, giving
\[ L = \mathbb{E}_{q} \left[\underbrace{D_{\mathrm{KL}}\big(q(\mathbf{x}_{T}\mid\mathbf{x}_{0}) \,\|\, p(\mathbf{x}_{T})\big)}_{L_T:\ \text{prior matching}} + \underbrace{\sum_{t=2}^{T}D_{\mathrm{KL}}\big(q(\mathbf{x}_{t-1}\mid\mathbf{x}_{t},\mathbf{x}_{0}) \,\|\, p_{\theta}(\mathbf{x}_{t-1}\mid\mathbf{x}_{t})\big)}_{L_{t-1}:\ \text{denoising matching}} - \underbrace{\log p_{\theta}(\mathbf{x}_{0}\mid\mathbf{x}_{1})}_{L_0:\ \text{reconstruction}} \right] \]