What are Diffusion Models?

Imagine watching a drop of ink slowly spread across a blank page, its color slowly diffusing through the paper until it becomes a beautiful, intricate pattern. This natural process of diffusion, where particles move from areas of high concentration to low concentration, is the inspiration behind diffusion models in machine learning. Just as the ink spreads and blends, diffusion models work by gradually adding and then removing noise from data to generate high-quality results. In this article, you will get to know all about the diffusion models , how do they works, its differences and its benefits.

What are Diffusion Models?

Diffusion models are inspired by the natural process where particles spread from areas of high concentration to low concentration until they are evenly distributed. This principle is seen in everyday examples, like the gradual dispersal of perfume in a room.

In the context of machine learning, diffusion models use a similar idea by starting with data and progressively adding noise to it. They then learn to reverse this process, effectively removing the noise and reconstructing the data or creating new, realistic versions. This gradual transformation results in detailed and high-quality outputs, useful in fields such as medical imaging, autonomous driving, and generating realistic images or text.

The unique aspect of diffusion models is their step-by-step refinement approach, which allows them to achieve highly accurate and nuanced results by mimicking natural processes of diffusion.

How Do Diffusion Models Work?

Diffusion models operate through a two-phase process: first, a neural network is trained to add noise to data (known as the forward diffusion phase), and then it learns to systematically reverse this process to recover the original data or generate new samples. Here’s an overview of the stages involved in a diffusion model’s functioning.

Data Preparation

Before starting the diffusion process, the data must be prepared correctly for training. This preparation includes steps like cleaning the data to remove anomalies, normalizing features to maintain consistency, and augmenting the dataset to enhance variety—especially important for image data. Standardization is used to ensure a normal distribution, which helps manage noisy data effectively. Different types of data, such as text or images, may require specific adjustments, such as addressing imbalances in data classes. Proper data preparation is crucial for providing the model with high-quality input, allowing it to learn significant patterns and produce realistic outputs during use.

Forward Diffusion Process : Transforming Images to Noise

The forward diffusion process starts by drawing from a simple distribution, typically Gaussian. This initial sample is then progressively altered through a sequence of reversible steps, each adding a bit more complexity via a Markov chain. As these transformations are applied, structured noise is incrementally introduced, allowing the model to learn and replicate the intricate patterns present in the target data distribution. The purpose of this process is to evolve the basic sample into one that closely resembles the complexity of the desired data. This approach demonstrates how beginning with simple inputs can result in rich, detailed outputs.

Mathematical Formulation 

Let x0​ represent the initial data (e.g., an image). The forward process generates a series of noisy versions of this data x1,x2,…,xT​ through the following iterative equation:

Here,q is our forward process, and xt is the output of the forward pass at step t. N is a normal distribution, 1-txt-1 is our mean, and tI defines variance.    

Reverse Diffusion Process : Transforming Noise to Image

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. 

Read More about the this article Image Generation with Stable Diffusion

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.
• Gaussian Noise: The noise removed (and added) is typically Gaussian, characterized by its mean and variance. 

The reverse diffusion process aims to reconstruct x0 ​ from xT, the noisy data at the final step. This process is modeled by the conditional distribution:

where:

• μθ(xt,t)is the mean predicted by the model,
• σθ2(t) is the variance, which is usually a function of t and may be learned or predefined.

The above image depicts the reverse diffusion process often used in generative models.

Starting from noise xT​, the process iteratively denoises the image through time steps T to 0. At each step t, a slightly less noisy version xt−1​ is predicted from the noisy input xt​ using a learned model pθ​(xt−1​∣xt​).

The dashed arrow labeled q(xt​∣xt−1​) represents the forward diffusion process, while the solid arrow pθ​(xt−1​∣xt​) represents the reverse process that is modeled and learned.

Implementation of How diffusion Model Works

We will now look into the steps of how diffusion model works.

Step1: Import Libraries

import torch
import torch.nn as nn
import torch.optim as optim

Step2: Define the Diffusion Model

class DiffusionModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(DiffusionModel, self).__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, output_dim)

    def forward(self, noise_signal):
        x = self.fc1(noise_signal)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.relu(x)
        x = self.fc3(x)
        return x

Defines a neural network model for the diffusion process with:

• Three Linear Layers: fc1, fc2, and fc3 for transforming the input through the network.
• ReLU Activations: Applied after the first and second linear layers to introduce non-linearity.

Step3: Initialize the Model and Optimizer

input_dim = 100
hidden_dim = 128
output_dim = 100
batch_size = 64
num_epochs = 5

model = DiffusionModel(input_dim, hidden_dim, output_dim)
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()
data_loader = [(torch.randn(batch_size, input_dim), torch.randn(batch_size, output_dim))] * 10
target_data = torch.randn(batch_size, output_dim)

• Sets dimensions for input, hidden, and output layers.
• Creates an instance of the DiffusionModel.
• Initializes the Adam optimizer with a learning rate of 0.001.

Training Loop:

for epoch in range(num_epochs):
    epoch_loss = 0
    for batch_data, target_data in data_loader:
        # Generate a random noise signal
        noise_signal = torch.randn(batch_size, input_dim)
        
        # Forward pass through the model
        generated_data = model(noise_signal)
        
        # Compute loss and backpropagate
        loss = criterion(generated_data, target_data)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
    # Print the average loss for this epoch
    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {epoch_loss / len(data_loader):.4f}')

Epoch Loop: Runs through the specified number of epochs.

Batch Loop: Processes each batch of data.

• Noise Signal: Generates random noise as input.
• Forward Pass: Passes the noise through the model to generate data.
• Compute Loss: Calculates the loss between generated data and target data.
• Backpropagation: Computes gradients and updates model parameters.
• Accumulate Loss: Adds the loss for each batch to compute the average loss per epoch.

Diffusion Model Techniques

Let us now discuss diffusion model techniques.

Denoising Diffusion Probabilistic Models (DDPMs)

DDPMs are one of the most widely recognized types of diffusion models. The core idea is to train a model to reverse a diffusion process, which gradually adds noise to data until all structure is destroyed, converting it to pure noise. The reverse process then learns to denoise step-by-step, reconstructing the original data.

Forward Process

 This is a Markov chain where Gaussian noise is sequentially added to a data sample over a series of time steps. This process continues until the data becomes indistinguishable from random noise.

Reverse Process

The reverse process, which is also a Markov chain, learns to undo the noise added in the forward process. It starts from pure noise and progressively denoises to generate a sample that resembles the original data.

Training

 The model is trained using a variant of a variational lower bound on the negative log-likelihood of the data. This involves learning the parameters of a neural network that predicts the noise added at each step. 

Score-Based Generative Models (SBGMs)

Score-based generative models use the concept of a “score function,” which is the gradient of the log probability density of data. The score function provides a way to understand how the data is distributed.

Score Matching

The model is trained to estimate the score function at different noise levels. This involves learning a neural network that can predict the gradient of the log probability at various scales of noise.

Langevin Dynamics

Once the score function learns, the process generates samples by starting with random noise and gradually denoising it using Langevin dynamics. This Markov Chain Monte Carlo (MCMC) method uses the score function to move towards higher-density regions.

Stochastic Differential Equations (SDEs)

In this approach, diffusion models are treated as continuous-time stochastic processes, described by SDEs.

Forward SDE

The forward process is described by an SDE that continuously adds noise to data over time. The drift and diffusion coefficients of the SDE dictate how the data evolves into noise.

Reverse-Time SDE

The reverse process is another SDE that goes in the opposite direction, transforming noise back into data by “reversing” the forward SDE. This requires knowing the score (the gradient of the log density of data).

Numerical Solvers

Numerical solvers like Euler-Maruyama or stochastic Runge-Kutta methods are used to solve these SDEs for generating samples.

Noise Conditional Score Networks (NCSN)

NCSN implements score-based models where the score network conditions on the noise level.

Noise Conditioning

 The model predicts the score (i.e., the gradient of the log-density of data) for different levels of noise. This is done using a noise-conditioned neural network.

Sampling with Langevin Dynamics

 Similar to other score-based models, NCSNs generate samples using Langevin dynamics, which iteratively denoises samples by following the learned score.

Variational Diffusion Models (VDMs)

VDMs combine the diffusion process with variational inference, a technique from Bayesian statistics, to create a more flexible generative model.

Variational Inference

 The model uses a variational approximation to the posterior distribution of latent variables. This approximation allows for efficient computation of likelihoods and posterior samples.

Diffusion Process

The diffusion process adds noise to the latent variables in a way that facilitates easy sampling and inference.

Optimization

The training process optimizes a variational lower bound to efficiently learn the diffusion process parameters.

Implicit Diffusion Models

Unlike explicit diffusion models like DDPMs, implicit diffusion models do not explicitly define a forward or reverse diffusion process.

Implicit Modeling

These models might leverage adversarial training techniques (like GANs) or other implicit methods to learn the data distribution. They do not require the explicit definition of a forward process that adds noise and a reverse process that removes it.

Applications

They are useful when the explicit formulation of a diffusion process is difficult or when combining the strengths of diffusion models with other generative modeling techniques, such as adversarial methods.

Augmented Diffusion Models

Researchers enhance standard diffusion models by introducing modifications to improve performance.

Modifications

Changes could involve altering the noise schedule (how noise levels distribute across time steps), using different neural network architectures, or incorporating additional conditioning information (e.g., class labels, text, etc.).

Goals

 The modifications aim to achieve higher fidelity, better diversity, faster sampling, or more control over the generated samples.

GAN vs. Diffusion Model

Difference Between GAN Vs Diffusion Model

Applications of Diffusion Models

We will now explore applications of diffusion model in detail.

Image Generation

Diffusion models excel in generating high-quality images. Artists have used them to create stunning, realistic artworks and generate images from textual descriptions.

Import Libraries

import torch
from diffusers import StableDiffusionPipeline

Set Up Model and Device

model_id = "CompVis/stable-diffusion-v1-4"
device = "cuda"

Load and Configure the Model

pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to(device)

Generate an Image

prompt = "a landscape with rivers and mountains"
image = pipe(prompt).images[0]

Save the Image

image.save("Image.png")

Image-to-Image Translation

From changing day scenes to night to turning sketches into realistic images, diffusion models have proven their worth in image-to-image translation tasks.

Install Necessary Libraries