Denoising Diffusion Model

The diffusion model is a type of latent variable model that can be viewed as an evolved version of the hierarchical Variational Autoencoder (VAE). Implementing a diffusion model requires only a few modifications to the hierarchical VAE, yet it yields remarkably superior generative performance. The key changes include:

1. Dimensional Parity Between Latent and Observed Variables: Both latent and observed variables share the same number of dimensions.

2. Simplified Encoder as a Noise-Adding Module: The encoder functions merely as a mechanism to add Gaussian noise, sourced from a fixed normal distribution.

The fundamental concept behind the diffusion model is to learn the distribution of observed data through a denoising process. The diffusion model operates in two main phases: diffusion and denoising.

• During the diffusion process, the observed data is incrementally corrupted by repeatedly adding noise from a fixed normal distribution until it is transformed into pure noise.

• In the denoising process, a neural network is trained to predict the noise that was added to the data. Through this predictive denoising process, the original data is progressively reconstructed.

Diffusion

In the diffusion process, noise is progressively added to the data at each time step, building on the noise added in the previous step. The noise addition is calibrated such that by the final time step, the data resembles latent variables, essentially being destroyed. Typically, the latent variable at the final step, \(x_T\), is assumed to follow a fixed normal distribution, denoted as \(\mathcal{N}(x_T; 0, \text{I})\). This approach mirrors the assumption in Variational Autoencoders (VAEs) that latent variables are derived from a normal distribution.

The modeling of the diffusion process is defined for each time step \(1 \leq t \leq T\) as follows.

\[ q(x_t|x_{t-1}) = \mathcal{N}(x_t;\sqrt{1-\beta_t}x_{t-1}, \beta_t\text{I}) \]

If \(\beta_t\) is set to a small value such as 0.001 and \(T\) is set sufficiently large, the distribution of \(x_t\) can be made close to a normal distribution.

In fact, \(x_t\) can be obtained using the following variable transformation trick.

\[\begin{split}\epsilon\sim\mathcal{N}(\epsilon;0,\text{I})\\x_t = \sqrt{1-\beta_t}x_{t-1}+\sqrt{\beta_t}\epsilon\end{split}\]

Denoising

In the denoising process of the diffusion model, a neural network is employed. Given that all latent and observed variables share the same dimensional space in the diffusion model, the same neural network architecture can be utilized for denoising at each time point. The denoising process can be modeled as follows.

\[\begin{split}\hat{x}_{t-1} = \text{NeuralNetwork}(x_t,t;\theta)\\p_\theta(x_{t-1}|x_t) =\mathcal{N}(x_{t-1};\hat{x}_{t-1},\text{I})\end{split}\]

Parameters Estimation

Parameters estimation in a diffusion model is also performed by maximizing the log-likelihood function. In the diffusion model, the parameters are those of the neural network used in the denoising process. However, since it is difficult to compute the log-likelihood of a diffusion model, similar to the VAE, ELBO is used.

First, let’s calculate the ELBO of the diffusion model.

\[\begin{split}\begin{align*}\text{ELBO}(x_0;\theta) &= \int q(x_{1:T}|x_0)\log\frac{p_{\theta}(x_{0:T})}{q(x_{1:T}|x_0)}dx_{1:T} \\ &= \mathbb{E}_{q(x_{1:T}|x_0)}\left[\log\frac{p_{\theta}(x_{0:T})}{q(x_{1:T}|x_0)}\right] \\&= \mathbb{E}_{q(x_{1:T}|x_0)}\left[\log\frac{p(x_{T})\prod_{t=1}^{T}{p_{\theta}(x_{t-1}|x_t)}}{\prod_{t=1}^{T}{q(x_{t}|x_t-1)}}\right] \\&= \mathbb{E}_{q(x_{1:T}|x_0)}\left[\log{\prod_{t=1}^{T}{p_{\theta}(x_{t-1}|x_t)}} + \log{\frac{p(x_T)}{\prod_{t=1}^{T}{q(x_{t}|x_t-1)}}}\right]\end{align*}\end{split}\]

Here, since the last time step data is a complete noise which follow a normal distribution, \(p(x_T) \) is not a function of \(\theta\). Therefore, maximizing ELBO is equivalent to maximizing the following function.

\[\begin{split}\begin{align*} J(\theta) &= \mathbb{E}_{q(x_{1:T}|x_0)}\left[\log{\prod_{t=1}^{T}{p_{\theta}(x_{t-1}|x_t)}} \right] \\ &= \mathbb{E}_{q(x_{1:T}|x_0)}\left[{\sum_{t=1}^{T}\log{p_{\theta}(x_{t-1}|x_t)}} \right] \\ &= \sum_{t=1}^{T}\mathbb{E}_{q(x_{1:T}|x_0)}\left[{\log{p_{\theta}(x_{t-1}|x_t)}} \right] \\ &= \sum_{t=1}^{T}\mathbb{E}_{q(x_{t-1}, x_t|x_0)}\left[{\log{p_{\theta}(x_{t-1}|x_t)}} \right] \\ \end{align*}\end{split}\]

\(\mathbb{E}_{q(x_{t-1}, x_t|x_0)}[\cdot]\) can be approximated with Monte Carlo method and together with the modeling of denoising process, this function can be further approximated as follows.

\[\begin{split}\begin{align*} J(\theta) &\approx \sum_{t=1}^{T}\log p_\theta(x_{t-1}|x_t) \\ \\ \hat{x}_{t-1} &= \text{NeuralNet}(x_t,t;\theta) \\ p_\theta(x_{t-1}|x_t) &= \mathcal{N}(x_{t-1};\hat{x}_{t-1},\mathbf{I})\\ \\ J(\theta) &\approx -\frac{1}{2}\sum_{t=1}^{T}||x_{t-1}-\hat{x}_{t-1}||^2 \end{align*}\end{split}\]

Maximizing this \(J(\theta)\) is equivalent to minimizing loss between \(x_{t-1}\) and denoised \(\hat{x}_{t-1}\). This means we train the neural network to conduct the denoising process given time step \(t\) and corresponding data \(x_t \) as inputs. However, as you may have noticed, each time of \(J(\theta)\) calculation, we need to do \(T\) times diffusion process sampling. This leads to high computational cost when we have large value of \(T\) set. To make the process more efficient, another method of approximation method can be considered by using only one sampling of the diffusion process.

Noise Prediction

The idea of this method is to train the neural network in denoising process to predict the noise \(\hat\epsilon\) added to the data, instead of predicting the data \(\hat{x}_{t-1}\).

We will mathematical show that training the neural network to predict the noise is equivalent to the previous approximation with \(T\) times samplings in the previous section.

Note that the distribution \(p_\theta(x_{t-1}|x_t)\) can be modeled as

\[\begin{align*} p_\theta(x_{t-1}|x_t) &= \mathcal{N}(x_{t-1};\mu_\theta(x_t,t),\sigma_{q}^2(t)\text{I}) \end{align*}\]

\[\begin{align*} \mu_\theta(x_t,t) &= \text{NeuralNetwork}(x_t, t) \end{align*}\]

Let’s start from the modeling of the diffusion process.

\[\begin{split}\begin{align*} q(x_t|x_{t-1}) &= \mathcal{N}(x_t;\sqrt{{\alpha}_t}x_{t-1},(1-\alpha_t)\text{I}) ~~~,\text{where}~\alpha_t = 1-\beta_t \\ q(x_t|x_0) &= \mathcal{N}(x_t;\sqrt{\bar{\alpha}_t}x_0,(1-\bar\alpha_t)\text{I}) ~~~,\text{where}~\bar\alpha_t = \prod_{t=1}^{T}(1-\beta_t) \\ \end{align*}\end{split}\]

Thus, \(x_t\) can be sampled as follows.

\[\begin{split}\begin{align*} \epsilon &\sim \mathcal{N}(0, \text{I}) \\ x_t &= \sqrt{\bar\alpha_t}x_0 + \sqrt{1-\bar\alpha_t}\epsilon \end{align*}\end{split}\]

Now, we can express the original observed data \(x_0\) as

\[\begin{align*} x_0 &= \frac{x_t-\sqrt{1-\bar\alpha_t}\epsilon}{\sqrt{\bar\alpha_t}} \end{align*}\]

Given \(x_0\) and \(x_t\), the distribution of \(x_{t-1}\) can be obtained as follows.

\[\begin{align*} q(x_{t-1}|x_t,x_0) &= \frac{q(x_t|x_{t-1},x_0)q(x_{t-1}|x_0)}{q(x_t|x_0)} \end{align*}\]

With Markov assumption, \(q(x_t|x_{t-1},x_0) = q(x_t|x_{t-1})\)

\[\begin{split}\begin{align*} q(x_{t-1}|x_t,x_0) &= \frac{q(x_t|x_{t-1})q(x_{t-1}|x_0)}{q(x_t|x_0)} \\ &=\frac{\mathcal{N}(x_t;\sqrt{{\alpha}_t}x_{t-1},(1-\alpha_t)\text{I}) ~ \mathcal{N}(x_{t-1};\sqrt{\bar{\alpha}_{t-1}}x_0,(1-\bar\alpha_{t-1})\text{I}) }{\mathcal{N}(x_t;\sqrt{\bar{\alpha}_t}x_0,(1-\bar\alpha_t)\text{I}) } \end{align*}\end{split}\]

By expanding the right side of this equation, we obtain the following distribution.

\[\begin{split}\begin{align*} q(x_{t-1}|x_t,x_0) &= \mathcal{N}\left(x_{t-1}; \mu_q(x_t,x_0), \sigma_q^2(t)\text{I}\right) \\ \mu_q(x_t,x_0) &= \frac{1}{\sqrt{\alpha_t}}\left(x_t - \frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}}\epsilon\right) \\ \sigma_q^2(t) &= \frac{(1-\alpha_t)(1-\bar\alpha_{t-1})}{1-\bar\alpha_t} \end{align*}\end{split}\]

Here, let’s rewrite \(\mu_\theta(x_t,t)\) to match the form of \(\mu_q(x_t,x_0)\) as

\[\begin{align*} \mu_\theta(x_t,t) &= \frac{1}{\sqrt{\alpha_t}}\left(x_t - \frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}}\epsilon_\theta(x_t, t)\right) \end{align*}\]

, where now we set the output of the neural network to \(\epsilon_\theta(x_t, t)\) instead.

Now, let’s look back to the KL divergence between the two normal distributions \(q(x_{t-1}|x_t,x_0)\) and \(p_\theta(x_{t-1}|x_t)= \mathcal{N}(x_{t-1};\hat{x}_{t-1},\mathbf{I})\).

\[\begin{split}\begin{align*} D_{\text{KL}}(q(x_{t-1}|x_t,x_0)\|p_\theta(x_{t-1}|x_t)) &= \frac{1}{2\sigma_q^2(t)}\|\mu_\theta(x_t,t) - \mu_q(x_t,x_0)\|^2 \\ &= \frac{1}{2\sigma_q^2(t)}\|\epsilon_\theta(x_t,t) - \epsilon\|^2 \end{align*}\end{split}\]

This shows that by train the neural network to predict the noise, which is minimizing the above KL divergence, lead to maximizing the EBOL of the diffusion model. Thus, this optimization is equivalent to the parameter estimation with \(T\) times sampling in the previous section.

To sum up, the training algorithm of the diffusion model is as follows.

Training Algorithm

Repeat the following steps:

1. Randomly sample a training data instance

2. Sample a time step \(t\) from a uniform distribution : \(t \sim U[1,T]\)

3. Generate Gaussian noise for diffusion process: \(\epsilon \sim \mathcal{N}(0,\text{I})\)

4. Diffusion process: \(x_t = \sqrt{\bar\alpha_t}x_0 + \sqrt{1-\bar\alpha_t}\epsilon\)

5. Calculate loss: \(\text{Loss}(x_0, \theta) = \|\epsilon_\theta(x_t,t)-\epsilon\|^2\)

6. Update parameters \(\theta\) with gradient decent

Generating New Data

As shown in previous sections, the denoising process is modeled as follows.

\[p_\theta(x_{t-1}|x_t) = \mathcal{N}(x_{t-1};\mu_\theta(x_t,t),\sigma_q^2(t)\text{I})\]

Using the following trick, we can sample every time step data with the following process.

\[\begin{split}\begin{align*} \epsilon &\sim \mathcal{N}(0, I) \\ x_{t-1} &= \mu_\theta(x_t,t)+\sigma_q(t)\epsilon \end{align*}\end{split}\]

Here,

\[\begin{split}\begin{align*} \mu_\theta(x_t,t) &= \frac{1}{\sqrt{\alpha_t}}\left(x_t - \frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}}\epsilon_\theta(x_t, t)\right) \\ \sigma_q(t)&=\sqrt{\frac{(1-\alpha_t)(1-\bar\alpha_{t-1})}{1-\bar\alpha_t}} \end{align*}\end{split}\]

. Thus, the process of new data generation can be summarized as follows.

Generation Process

\(x_T \sim \mathcal{N}(0,\text{I})\)

for \(t\) in \([T,...,1]\):

\(\epsilon\sim\mathcal{N}(0,\text{I})\)

if \(t=1\) then \(\epsilon=0\)

\(\sigma_q(t)=\sqrt{\frac{(1-\alpha_t)(1-\bar\alpha_{t-1})}{1-\bar\alpha_t}}\)

\(x_{t-1} = \frac{1}{\sqrt{\alpha_t}}\left(x_t - \frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}}\epsilon_\theta(x_t, t)\right) +\sigma_q(t)\epsilon\)

return \(x_0\)