Introduction

Stable diffusion is a powerful (generative model) tool to create high-quality images from noise. Stable diffusion consists of two steps: a forward diffusion process and a reverse diffusion process. In the forward diffusion process, noise is progressively added to an image, effectively degrading its quality. This step is crucial for training the model, as it helps the model learn how images can transition from clarity to noise. We have covered the details of the forward diffusion process in our previous article.

In reverse diffusion, noise is progressively removed to generate a high-quality image. This article will focus on this process, exploring its mechanisms and mathematical foundations.

Overview

1. Stable diffusion uses forward and reverse processes to generate high-quality images from noise.
2. The forward diffusion process progressively adds noise to an image for training.
3. The reverse diffusion process removes noise iteratively to reconstruct the original image.
4. This article explores the reverse diffusion process and its mathematical foundations.
5. Training involves predicting noise at each step to enhance image quality.
6. The neural network architecture and loss function are key to effective training.

What is the Reverse Diffusion Process?

The reverse diffusion process aims to convert pure noise into a clean image by iteratively removing noise. Training a diffusion model is to learn the reverse diffusion process so that it can reconstruct an image from pure noise. If you guys are familiar with GANs, we’re trying to train our generator network, but the only difference is that the diffusion network does an easier job because it doesn’t have to do all the work in one step. Instead, it uses multiple steps to remove noise at a time, which is more efficient and easy to train, as figured out by the authors of this paper. 

Mathematical Foundation of Reverse Diffusion

• Markov Chain: The diffusion process is modeled as a Markov chain, where each step only depends on the previous state.
  
  Read about Markov Chain here: A Comprehensive Guide on Markov Chain
  
• Gaussian Noise: The noise removed (and added) is typically Gaussian, characterized by its mean and variance. 

What Does a Diffusion Model Do?

Many people think that a neural network (called a diffusion model for even more confusion) removes noise from an input image or predicts the noise to be removed from an input. Both are incorrect. What the diffusion model does is predict the entire noise to be removed at a particular timestep. This means that if we have timestep t=600, then our Diffusion model tries to predict the entire noise on which removal we should get to t=0, not t=599. 

Reverse Diffusion Algorithm

• Initialization: The Reverse Diffusion process starts with a noisy image, as you guys have guessed. This image acts as a sample for noise distribution. 
• Iterative Denoising: The model iteratively removes noise at each timestep to recover the original data. This is done by following a sequence of denoising steps, where the model predicts the noise present in the current noisy image. Usually, denoising steps are:
  • Estimate the noise in the current image (current timestep to timestep 0).
  • Subtract a portion of this estimated noise.
• Noise Addition: A small amount of noise is introduced back at each timestep to keep the process from becoming deterministic and to preserve generalization in the generated samples. This encourages exploration of the solution space and keeps the model from being trapped in local minima. The added noise is typically reduced as the process goes on to ensure that the final image is less noisy and more in line with the intended output.
• Final Output: The result after all iterations is the generated image.

Mathematical Formulation

This is the equation that we took from the paper Denoising Diffusion Probabilistic Models. 

It basically says that  𝑝𝜃(𝑥0:𝑇) is a chain of Gaussian transitions starting at  𝑝(𝑥𝑇) and iterating T times using the equation for one diffusion process step 𝑝𝜃(𝑥𝑡−1∣𝑥𝑡).

Now it’s time to explain how the single step works and how to get something to implement. 

𝑁(𝑥𝑡−1,𝜇𝜃(𝑥𝑡,𝑡),∑𝜃(𝑥𝑡,𝑡)) has 2 parts:

• 𝜇𝜃(𝑥𝑡,𝑡) (mean)
• ∑𝜃(𝑥𝑡,𝑡) which equals 𝜎𝑡2𝐼 (variance)

To know more about the mathematical foundations of the reverse diffusion process refer to this article.

Training the Model Using the Reverse Diffusion process

The generation of images using the reverse diffusion process relies highly on how well the model can predict the noise included in the forward diffusion process. This noise prediction capability is developed through a rigorous training process.

The main objective of training the model using reverse diffusion is to predict the noise at each diffusion process step. By minimizing the error between predicted and actual noise, the model learns to denoise the image effectively.

Training Data

The training data consists of pairs of noisy images and the corresponding noise added at each step during the forward diffusion process. This data is generated by applying the forward diffusion process to a set of clean images, progressively adding noise over multiple steps. 

Loss Function

A critical component of the training process is the loss function. The loss function quantifies the difference between predicted and actual noise. One commonly used loss function is the Mean Squared Error (MSE). The model is trained to minimize this MSE loss, thereby improving its ability to predict the noise accurately.

Neural Network Architecture

Convolutional neural networks (CNNs) are the most common type of neural network utilized in the reverse diffusion process for noise prediction. CNNs can record spatial hierarchies in images, making them ideal for image processing applications. Multiple convolutional layers, pooling layers, and activation functions may be used in the architecture to extract and learn complicated characteristics from noisy pictures. There are two common backbone architecture choices for diffusion models: U-Net and Transformer.

Training Procedure

• Initialization: Set random weights at the beginning of the neural network.
• Forward Pass: To obtain the predicted noise, send the noisy image through the neural network for each training sample.
• Loss Calculation: Determine the loss by comparing the expected and actual noise using the selected loss function (e.g., MSE).
• Backward Pass: Perform backpropagation to calculate the gradients of the loss with respect to the network’s weights.
• Weight Update: To minimize the loss, update the network’s weights using an optimization technique such as Adam or Stochastic Gradient Descent (SGD).
• Iteration: Until the model converges to an ideal set of weights, repeat the forward pass, loss computation, backward pass, and weight update for several epochs.

Evaluation

The model’s performance is assessed after training using a different validation dataset that wasn’t utilized for training. On this validation set, the model’s accuracy in predicting noise is an indication of its generalization ability. Metrics like mean squared error (MSE), root mean square error (RMSE), mean absolute error (MAE), and R-squared (coefficient of determination) are often used.

Conclusion

Stable diffusion models rely on both the forward and reverse diffusion processes. These processes work together to gradually reduce noise in an image, ultimately producing high-quality results. This iterative refining mechanism is rooted in strong mathematical foundations, making stable diffusion an effective tool in the generative model field. As research in this area progresses, we can anticipate even more advanced applications and developments in this intriguing field. 

Frequently Asked Questions

Badrinarayan M

Data science Trainee at Analytics Vidhya, specializing in ML, DL and Gen AI. Dedicated to sharing insights through articles on these subjects. Eager to learn and contribute to the field's advancements. Passionate about leveraging data to solve complex problems and drive innovation.