Diffusion Models

Original paper: Denoising Diffusion Probabilistic Models
Survey: Understanding Diffusion Models A Unified Perspective

Intuition

A diffusion model is trained through a diffusion process that progressively adds noise to the original data. The model then learns how to reconstruct the original data from this noisy input.

graph LR;
A[Original Data] -->|Forward Diffusion| B[Noisy Data];
B -->|Reverse Denoising| A

Once the model has sufficiently learned to reconstruct the data distribution from a typically Gaussian noise distribution, it gains the capability to generate new, novel data.

Forward Process

At each time step $t$, the transition probability from input ${\mathbf{x}}_{t-1}$ to output ${\mathbf{x}}_t$ is defined as a Gaussian:

where ${\beta }_t\in (0,1)$.

Reparameterization

We can generate ${\mathbf{x}}_t$ using the reparameterization trick:

Furthermore, we can generate ${\mathbf{x}}_t$ from ${\mathbf{x}}_0$ with one reparameterization step:

where

Asymptotics

Since ${\alpha }_i\in (0,1)$, as $t\to \mathrm{\infty }$

Therefore, regardless of ${\mathbf{x}}_0$,

Reverse Process

The posterior distribution, $q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)$, is generally intractable:

where

Therefore, we learn a network parameterized by $\theta$ to approximate $q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)$ with a Gaussian distribution:

We approximate the log-likelihood using the ELBO:

The ${\mathcal{L}}_T$ Term

Since ${\mathbf{x}}_T$ is directly sampled from a standard Gaussian distribution, ${\mathcal{L}}_T$ has no learnable parameters for fixed ${\beta }_{1:T}$, thus ignored from the training objective.

The ${\mathcal{L}}_{1:T-1}$ Term

The posterior conditioned on ${\mathbf{x}}_0$, $q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0)$, is a tractable Gaussian:

where

For simplicity, ${\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t)$ is set to untrained time-dependent constants ${\sigma }_t^2\mathbf{I}$, where ${\sigma }_t^2={\tilde{\beta }}_t$. Since $q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0)$ and $p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)$ are both tractable Gaussians, the KL divergence has a closed form solution (see KL divergence between Gaussians):

The ${\mathcal{L}}_0$ Term

To obtain discrete log-likelihood, this term is set to an independent discrete decoder derived from Gaussian $\mathcal{N}({\mathbf{x}}_0;{\mathit{\mu }}_{\theta }({\mathbf{x}}_1,1),{\sigma }_1^2\mathbf{I})$. For discrete ${\mathbf{x}}_t$ of dimension $D$,

where $F(x)$ is a discretization function that maps $x\in \mathbb{R}$ into discrete $({\mathbf{x}}_0{)}_i$.

Training

The formulation of ${\tilde{\mathit{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0)$ indicates that instead of directly predicting a mean ${\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t)$, it is better to predict a ${\mathit{ϵ}}_{\theta }({\mathbf{x}}_t,t)$, such that

Combining with the reparameterization step, we have the training objective:

We thus have the training algorithm:

\begin{algorithm}
\caption{Training}
\begin{algorithmic}
\Repeat
    \State $\mathbf{x}_0 \sim q(\mathbf{x}_0)$
    \State $t \sim \text{Uniform}(\{1, \ldots, T\})$
    \State $\boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$
    \State Take gradient descent step on $\nabla_\theta \| \boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta (\sqrt{\bar\alpha_t} \mathbf{x}_0 + \sqrt{1-\bar\alpha_t} \boldsymbol{\epsilon}, t) \|^2$
\Until{converged}
\end{algorithmic}
\end{algorithm}

Sampling

The sampling algorithm also uses the reparameterization trick to sample ${\mathbf{x}}_{t-1}$ given ${\mathbf{x}}_t$:

where $\mathbf{z}\sim \mathcal{N}(\mathbf{0},\mathbf{I})$ if $t>1$, and $\mathbf{z}$ is set to zero when $t=1$ to generate deterministic results.

\begin{algorithm}
\caption{Sampling}
\begin{algorithmic}
\State $\mathbf{x}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$
\For{$t = T, \cdots, 1$}
    \State $\mathbf{z} \sim \begin{cases} 
    \mathcal{N}(\mathbf{0}, \mathbf{I}) & \text{if } t > 1 \\
    \mathbf{0} & \text{otherwise}
    \end{cases}$
    \State $\mathbf{x}_{t-1} = \frac{1}{\sqrt{\alpha_t}} \left(\mathbf{x}_t - \frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}} \boldsymbol\epsilon_\theta(\mathbf{x}_t, t)\right) + \sigma_t \mathbf{z}$
\EndFor
\State \Return $\mathbf{x}_0$
\end{algorithmic}
\end{algorithm}

Network Architecture

The network for estimating ${\mathit{ϵ}}_{\theta }({\mathbf{x}}_t,t)$ is typically a U-Net architecture, where parameters are shared across time steps, and the time step $t$ is taken as input.

Extensions and Applications

Conditional Generation

• Additional conditioning information $y$ is applied to control generation.
• The network now samples ${\mathbf{x}}_{t-1}$ from $p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,t,y)$.
• The condition $y$ is taken as input in the similar fashion to $t$.
• ControlNet

Guidance Methods

• Gradient guidance added to shift the predicted mean.
• For more information, see Guidance Methods.

Latent Diffusion Models

• Train a VAE to map the input data to latent space.
• Diffusion and denoising are performed within the latent space.
• Lower computational cost.