Blog post written by: Rohan Singh

Based on: J. Ho, A. Jain, P. Abbeel. Denoising Diffusion Probabilistic Models, DOI:10.48550/arXiv.2006.11239, (2020)   https://doi.org/10.48550/arXiv.2006.11239

1. Introduction

Diffusion Models are a class of probabilistic generative models, which can generate new images which are similar the original dateset that they train on. Denoising diffusion probabilistic models answers the problems faced by GAN's and VAE's, i.e: unstable training and not so realistic images, and they show great promise in image generation.

Diffusion Models consists of two processes:

1. Forward diffusion process
2. Reverse diffusion process

we will discuss both the aforementioned processes in depth, in the methodology.

2. Methodology

In this section, we will cover in depth how is a diffusion model able to create such realistic images. Now, we will first understand how the Forward diffusion process works.

2.1. Forward diffusion process

Given an image from a Dateset, which follows a data distribution $//<![CDATA[ \begin{array}{l}x_0 \sim q(x)\end{array} //]]>$. We define a forward diffusion process in which we add a small amount of Gaussian noise over the horizon $//<![CDATA[ \begin{array}{l}T\end{array} //]]>$ noise and perturb the image, which would produce a sequence of noisy samples $//<![CDATA[ \begin{array}{l}x_1, ....,x_T\end{array} //]]>$ .

The distribution $//<![CDATA[ \begin{array}{l}q\end{array} //]]>$ for the forward diffusion process can be defined as a Markov chain by:

where the conditional follows a Gaussian with mean as $//<![CDATA[ \begin{array}{l}\sqrt{1 - \beta_t }\end{array} //]]>$ and standard deviation as $//<![CDATA[ \begin{array}{l}\beta_t I\end{array} //]]>$, that means we rescale each pixel by mean and add the standard deviation to perturb it.

where the size of each step, is the amount of noise added at each time step is controlled by $//<![CDATA[ \begin{array}{l}\{\beta_t \in (0, 1)\}_{t=1}^T\end{array} //]]>$ 

Figure 1: Forward diffusion process.

Further a nice property of the above process is that we can sample a noisy image $//<![CDATA[ \begin{array}{l}x_t\end{array} //]]>$, at any arbitrary time step in a closed form using reparameterization trick . Let $//<![CDATA[ \begin{array}{l}\alpha = 1\end{array} //]]>$ and $//<![CDATA[ \begin{array}{l}\bar{\alpha}_t = \prod_{i=1}^t \alpha_i\end{array} //]]>$, using which we get:

The above is based on the fact that when we merge two gaussians $//<![CDATA[ \begin{array}{l}\mathcal{N}(\mathbf{0}, \sigma_1^2\mathbf{I})\end{array} //]]>$ and $//<![CDATA[ \begin{array}{l}\mathcal{N}(\mathbf{0}, \sigma_2^2\mathbf{I})\end{array} //]]>$, the new distribution $//<![CDATA[ \begin{array}{l}\mathcal{N}(\mathbf{0}, (\sigma_1^2 + \sigma_2^2)\mathbf{I})\end{array} //]]>$ will have the merged standard deviation, which for the above case could be explained as $//<![CDATA[ \begin{array}{l}\sqrt{(1 - \alpha_t) + \alpha_t (1-\alpha_{t-1})} = \sqrt{1 - \alpha_t\alpha_{t-1}}\end{array} //]]>$ .

It is important to take a note that at the end of the horizon the perturbed input image will follow a normal distribution: $//<![CDATA[ \begin{array}{l}\mathbf{q(x_T)} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})\end{array} //]]>$.

2.2. Reverse diffusion process

In the forward process we added noise over the input images over the horizon $//<![CDATA[ \begin{array}{l}T\end{array} //]]>$. In the reverse diffusion process we would like to reverse the forward process and sample from $//<![CDATA[ \begin{array}{l}q(\mathbf{x}_{t-1} \vert \mathbf{x}_t)\end{array} //]]>$, where:

In general $//<![CDATA[ \begin{array}{l}q(\mathbf{x}_{t-1} \vert \mathbf{x}_t)\end{array} //]]>$ is intractable, but it is shown that when during the forward process if the $//<![CDATA[ \begin{array}{l}\beta_t\end{array} //]]>$ is small then the posterior can be approximated using a normal distribution. Then the joint distribution of the reverse diffusion process is also a markov process

So, in the reverse process we sample a noise from $//<![CDATA[ \begin{array}{l}\mathbf{q(x_T)} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})\end{array} //]]>$, and then iterate backwards over the horizon and at each step we try to predict the mean of the gaussain noise that was added to the image in the forwards diffusion process using a neural network architecture(U-Net); we will discuss about the U-Net in brief in the training section

Figure 2: Reverse Diffusion process

3. Learning process in Denoising Diffusion Probabilistic Model

Since diffusion models are a class of latent variable models, therefore we need to approximate the Evidence Lower bound (ELBO) is a loss function that is used for variational inference. The Variational Lower bound is (VLB and ELBO are interchangeable terms):

This aforementioned objective can be further written as a combination of KL Divergence and entropy terms to make it analytically computable.

Every term $//<![CDATA[ \begin{array}{l}L_V_L_B\end{array} //]]>$ besides $//<![CDATA[ \begin{array}{l}L_t_-_1\end{array} //]]>$,  can be computed analytically, so we can ignore them. Further $//<![CDATA[ \begin{array}{l}q(x_t_-_1 \vert x_t, x_0)\end{array} //]]>$ is  the tractable posterior distribution,  which can be defined as:

 where $//<![CDATA[ \begin{array}{l}\tilde{\mu}(x_t,x_0) := \frac{\sqrt{\bar{\alpha}}_{_t_-_1}\beta_t }{1 - \alpha_t} x_0 + \frac{\sqrt{1 - \beta_t}(1 - \tilde{\alpha}_{_t_-_1})}{1 - \tilde{\alpha_t}} x_t\end{array} //]]>$ (mean )and $//<![CDATA[ \begin{array}{l}\bar{\beta} := \frac{1 - \bar{\alpha}_t_-_1}{1 - \bar{\alpha}_t} \beta_t\end{array} //]]>$ (variance). Further since both  $//<![CDATA[ \begin{array}{l}q(x_t_-_1 \vert x_t, x_0)\end{array} //]]>$ and $//<![CDATA[ \begin{array}{l}p_\theta(x_t_-_1 \vert x_t)\end{array} //]]>$ are normal distribution, the KL divergence has a simpler form

Further, it was observed that $//<![CDATA[ \begin{array}{l}\tilde{\mu}(x_t,x_0) = \frac{1}{\sqrt{1 - \beta_t}}(x_t - \frac{\beta_t}{\sqrt{1 - \alpha_t}}\epsilon\end{array} //]]>$ and output of neural network model (U-Net) is  $//<![CDATA[ \begin{array}{l}\mu_\theta = \frac{1}{1 - \beta_t}\left(x_t - \frac{\beta_t}{\sqrt{1 - \tilde{\alpha}_t}}\epsilon_\theta(x_t,t)\right)\end{array} //]]>$, and the above equation can be further simplified into

where $//<![CDATA[ \begin{array}{l}x_t = \sqrt{\tilde{\alpha}_t}x_0 + \sqrt{1 -\tilde{\alpha}_t}\epsilon\end{array} //]]>$, Also it was found that when $//<![CDATA[ \begin{array}{l}\lambda_t = 1\end{array} //]]>$ the model produces high quality images, which the simplifies above equation to:

so after performing the forward diffusion process, in the reverse diffusion process the model takes noisy image and learns to predict the noise that was used to generate that specific noisy image.

4. Training

4.1. Training

During training while iterating through the epoch loop and the batch loop of the dataset, we randomly choose $//<![CDATA[ \begin{array}{l}t\end{array} //]]>$ in the horizon $//<![CDATA[ \begin{array}{l}H\end{array} //]]>$ and generate noisy images using $//<![CDATA[ \begin{array}{l}x_t =\sqrt{\bar{\alpha}_t}\mathbf{x}_0+\sqrt{1 - \bar{\alpha}_t}{\epsilon}\end{array} //]]>$, where $//<![CDATA[ \begin{array}{l}\epsilon \sim \mathcal{N}(0, \mathbb{I})\end{array} //]]>$ (forward process) , then we head to the reverse process and we first sample a noise $//<![CDATA[ \begin{array}{l}\mathbf{x_T} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})\end{array} //]]>$  and pass the noisy sample through the neural network (U-Net), which predicts the noise that must have been added to generate that noisy sample during the forward diffusion process. Then we compute the Mean Squared Error between the noise we added in the forward part of training and the noise that was predicted by the neural network (U-net) model and we then repeat this process until we converge

The U-Net architecture used is shown below:

But there are few changes made to the U-Net architecture, the model has to be conditioned on the current time step. To do so in practice, we can incorporate sinusoidal embedding and one-layer MLPs. Self-attention Blocks are also added at the 16 × 16 feature map resolution.

4.2. Sampling

During sampling, we sample a noise $//<![CDATA[ \begin{array}{l}\mathbf{x_T} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})\end{array} //]]>$ and iterate backward over the horizon $//<![CDATA[ \begin{array}{l}H\end{array} //]]>$ and compute less noisy image $//<![CDATA[ \begin{array}{l}x_{t-1} = \frac{1}{\sqrt\alpha_{t}} \left( x_t - \frac{1 - \alpha_t}{\sqrt{1 - \alpha_{t}}} \epsilon_\theta(x_t, t) + \sigma_tz \right)\end{array} //]]>$ where $//<![CDATA[ \begin{array}{l}z \sim \mathcal{N}(0,\mathbb{I})\end{array} //]]>$ until we reach$//<![CDATA[ \begin{array}{l}x_0\end{array} //]]>$, which would be the sampled or generated image from the Denoising diffusion probabilistic model.

5. Results

After implementing the all the aforementioned information using PyTorch and training the model for 35 epoch using ADAM optimizer, I was able to achieve the following results on the MNIST dataset:

6. Conclusion 

1. Neuroscience Data Augmentation: Machine learning models in neuroscience often require large amounts of labeled data, which can be difficult and expensive to obtain. DDPMs can generate synthetic brain imaging or activity data that can be used to augment existing datasets, improving the training of machine learning models for tasks like classification, segmentation, and anomaly detection.

2. Denoising Brain Scans: Brain imaging techniques like MRI, fMRI, and EEG are often affected by noise due to the complexities of data acquisition and the environment. DDPMs can be used to denoise these brain scans.

7. References

1. J. Ho, A. Jain, P. Abbeel. Denoising Diffusion Probabilistic Models, DOI:10.48550/arXiv.2006.11239, (2020)   https://doi.org/10.48550/arXiv.2006.11239
2. Ronneberger, Olaf, Fischer, Philipp and Brox, Thomas. “U-Net: Convolutional Networks for Biomedical Image Segmentation.” (2015). URL https://arxiv.org/abs/1505.04597.