Diffusion Model

Diffusion models are probabilistic generative models that define a nonlinear mapping from latent variables to the observed data.

• Like normalizing flows, the latent variables and observed data have the same dimensions.
• Like VAEs, they approximate the data likelihood using a lower bound based on an encoder that maps to the latent variable. However, in diffusion models, the encoder is pre-determined; the goal is to learn a decoder that is the inverse of this process and can be used to produce samples.

Diffusion models are easy to train and can produce very high-quality samples that exceed the realism of GANs.

Overview

A diffusion model consists of an encoder and a decoder:

• The encoder takes a data sample $x$ and maps it through a series of intermediate latent variables $z_1,\dots ,z_T$.
• The decoder reverses this process; it starts with $z_T$ and maps back through $z_{T-1},\dots ,z_1$ until it finally re-creates the data point $x$.

In both the encoder and the decoder, the mappings are stochastic rather than deterministic.

The encoder is pre-specified; it gradually blends the input with samples of white noise. With enough steps, the conditional distribution $q(z_T|x)$ and marginal distribution $q(z_T)$ of the final latent variable both become the standard normal distribution. Since this is all pre-specified, all the learned parameters are in the decoder.

In the decoder, a series of netwoks are trained to map backward between each adjacent pairs of latent variables $z_t$ and $z_{t-1}$. The loss function encourages each network to invert the corresponding encoder step. The result is that the noise is gradually removed from the representation until a realistic-looking data example remains. To generate a new data example $x$, we draw a sample from $q(z_T)$ and pass it through the decoder.

Encoder

The Diffusion Encoder maps the data example $x$ through a series of intermediate variables:

$$\begin{array}{rl}{z}_1& =\sqrt{1-{\beta }_1}\cdot x+\sqrt{{\beta }_1}\cdot {ϵ}_1\\ z_t& =\sqrt{1-{\beta }_t}\cdot z_{t-1}+\sqrt{{\beta }_t}\cdot {ϵ}_t\forall t\in 2,\dots ,T\end{array}$$

We can find $z_t$ without the intermediate variables using a diffusion kernel:

$$\begin{array}{rl}{z}_t& =\sqrt{{\alpha }_t}\cdot x+\sqrt{1-{\alpha }_t}\cdot ϵ\\ {\alpha }_t& =\prod _{s=1}^t(1-{\beta }_s)\end{array}$$

Or in probabilistic form:

$$q(z_t|x)={\text{Norm}}_{z_t}[\sqrt{a_t}\cdot x,(1-{\alpha }_t)I]$$

With this, we can marginalize to find the marginal distributions

$$q(z_t)=\int q(z_t|x)Pr(x)dx$$

We can also Bayes’ rule to find the conditional distribution $q(z_{t-1}|z_t)$:

$$q(z_{t-1}|z_t)=\frac{q(z_t|z_{t-1})q(z_{t-1})}{q(z_t)}$$

but this is intractable so we just approximate it using a normal distribution.

Similarly, we can also use Bayes’ rule to find the conditional diffusion distribution $q(z_{t-1}|z_t,x)$:

$$\begin{array}{rl}q(z_{t-1}|z_t,x)& =\frac{q(z_t|z_{t-1},x)q(z_{t-1}|x)}{q(z_t|x)}\\ & \propto {\text{Norm}}_{z_{t-1}}\left[\frac{1}{\sqrt{1-{\beta }_t}}{z}_t,\frac{{\beta }_t}{1-{\beta }_t}I\right]{\text{Norm}}_{z_{t-1}}[\sqrt{{\alpha }_{t-1}}\cdot x,(1-{\alpha }_{t-1})I]\end{array}$$

Decoder

The Diffusion Decoder learns the reverse of the encoder. It maps from latent variables $z_T,z_{T-1},\dots$ back to the data $x$. We saw that the true reverse distributions $q(z_{t-1}|z_t)$ are intractable, so we approximate them as normal distributions:

$$\begin{array}{rl}Pr(z_T)& ={\text{Norm}}_{z_T}[0,I]\\ Pr(z_{t-1}|z_t,{\varphi }_t)& ={\text{Norm}}_{z_{t-1}}[f_t[z_t,{\varphi }_t],{\sigma }_t^{2}I]\\ Pr(x|z_1,{\varphi }_1)& ={\text{Norm}}_x[f_1[z_1,{\varphi }_1],{\sigma }_1^{2}I]\end{array}$$

where we use a neural network $f_t[z_t,{\varphi }_t]$ to compute the mean. We can generate new samples by ancestral sampling.

Training

In Diffusion Training, we maximize the log-likelihood of the training data $\{x_i\}$ with respect to parameters $\varphi$:

$${\hat{\varphi }}_{1\dots T}=\underset{{\varphi }_{1\dots T}}{\mathrm{argmax}}\left[\sum _{i=1}^J\log [Pr(x_i|{\varphi }_{1\dots T})]\right]$$

However this is intractable because of $Pr(x_i|{\varphi }_{1\dots T})$. Thus, we compute an evidence lower bound (ELBO) for the log-likelihood, which ends up being:

$$\text{ELBO}[{\varphi }_{1\dots T}]={\mathbb{E}}_{q(z_1|x)}[\log [Pr(x|z_1,{\varphi }_1)]]-\sum _{t=2}^T{\mathbb{E}}_{q(z_{t,t-1}|x)}[D_{KL}[q(z_{t-1}|z_t,x)\mid \mid Pr(z_{t-1}|z_t,{\varphi }_t)]]$$

where:

• The first term is equivalent to the reconstruction term in the VAE. The ELBO will be larger if the model prediction matches the observed data.
• The second term is the KL divergence between exact reverse posterior of the forward diffusion process $q(z_{t-1}|z_t,x)$, and the decoder’s learned reverse transition $Pr(z_{t-1}|z_t,{\varphi }_t)$.

The loss function is then a minimization version of this, where we also approximate the expectations with sampling:

$$\begin{array}{r}L[{\varphi }_{1\dots T}]=\sum _{i=I}^I\left(\underset{\text{reconstruction term}}{\underbrace{-\log [{\text{Norm}}_{x_i}[f_1[z_{i_1},{\varphi }_1],{\sigma }_1^{2}I]]}}+\sum _{t=2}^T\frac{1}{2{\sigma }_t^2}{\left|\left|\underset{\text{target, mean of }q(z_{t-1}|z_t,x)}{\underbrace{\frac{1-{\alpha }_{t-1}}{1-{\alpha }_t}\sqrt{1-{\beta }_t}{z}_{it}+\frac{\sqrt{{\alpha }_{t-1}{\beta }_t}}{1-{\alpha }_t}{x}_i}}-\underset{\text{predicted }{z}_{t-1}}{\underbrace{f_t[z_{it},{\varphi }_t]}}\right|\right|}^2\right)\end{array}$$

We then use this loss function to train a network for each diffusion timestep, minimizing the difference between the estimate $z_{t-1}=f_t[z_t,{\varphi }_t]$ of the hidden variable at the previous timestep and the most likely value that it took given the ground truth $x$.

In practice, we often use a different reparameterized diffusion loss which has been found to work better. Here, the loss function is modified so that the model aims to predict the noise that was mixed with the original data example to create the current variable. It gives a simple formulation:

$$\begin{array}{rl}L[{\varphi }_{1\dots T}]& =\sum _{i=1}^I\sum _{t=1}^T\left|\right|g_t[z_{it},{\varphi }_t]-{ϵ}_{it}|{|}^2\\ & =\sum _{i=1}^I\sum _{t=1}^T\left|\right|g_t[\sqrt{{\alpha }_t}\cdot x_i+\sqrt{1-{\alpha }_t}\cdot {ϵ}_{it},{\varphi }_t]-{ϵ}_{it}|{|}^2\end{array}$$

where we have rewritten $z_t$ using the diffusion kernel in the second line.

Implementation