Forward Diffusion Process Simulation and Verification

1. The Forward Process: An Intuitive Overview

At its core, the forward diffusion process is conceptually simple: it's a fixed procedure that systematically and gradually destroys structure in data by adding noise. We start with a clean image, ${\mathbf{x}}_0$, and over a series of $T$ timesteps, we add a small amount of Gaussian noise at each step. This continues until, at the final timestep $T$, the image is indistinguishable from pure random noise.

This is a fixed process, meaning we don't need a neural network to learn it. The amount of noise added at each step is determined by a pre-defined variance schedule. The real goal of a diffusion model, which we'll explore in the next lesson, is to learn how to reverse this process.

2. The Mathematics of Gradual Noising

The forward process is mathematically defined as a Markov chain, where the state at timestep $t$, denoted ${\mathbf{x}}_t$, depends only on the state at the previous timestep, ${\mathbf{x}}_{t-1}$.

The distribution of ${\mathbf{x}}_t$ given ${\mathbf{x}}_{t-1}$ is a Gaussian distribution:

Let's break this down:

• ${\beta }_t$: This is the variance schedule. It's a set of pre-defined small positive constants, $\{{\beta }_1,{\beta }_2,...,{\beta }_T\}$, that control how much noise is added at each step $t$. A common choice is a linear schedule, where ${\beta }_t$ increases from a small value (e.g., 0.0001) to a larger one (e.g., 0.02).
• Mean ($\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1}$): The mean of the distribution is the previous image ${\mathbf{x}}_{t-1}$, scaled down slightly.
• Variance (${\beta }_t\mathbf{I}$): The variance determines the "strength" of the noise we add. $\mathbf{I}$ is the identity matrix, meaning we add independent noise to each pixel.

To get a more intuitive feel for this process and the role of the variance schedule, watch the following segment.

3. The "Nice Property": Jumping to Any Timestep

The step-by-step Markovian process is intuitive, but it's computationally inefficient. If we want to generate a noisy sample at timestep $t=500$, we would need to perform 500 iterative sampling steps. This is impractical for training.

Fortunately, a key property of this process is that we can sample ${\mathbf{x}}_t$ directly from the original image ${\mathbf{x}}_0$ in a single step. This is possible because the sum of Gaussian distributions is also a Gaussian. By defining ${\alpha }_t=1-{\beta }_t$ and ${\bar{\alpha }}_t={\prod }_{s=1}^t{\alpha }_s$ (the cumulative product of ${\alpha }_s$), we can derive a closed-form expression for $q({\mathbf{x}}_t|{\mathbf{x}}_0)$.

The derivation involves recursively expanding the expression for ${\mathbf{x}}_t$. If you're interested in the step-by-step math, the following resource provides an excellent walkthrough.

The final result of this derivation is the most important equation for the forward process:

This formula tells us that any noised sample ${\mathbf{x}}_t$ can be expressed as a simple combination of the original image ${\mathbf{x}}_0$ and some Gaussian noise. Using the reparameterization trick, we can write this as a direct sampling equation:

where $\mathbf{ϵ}\sim \mathcal{N}(0,\mathbf{I})$ is a random noise tensor with the same dimensions as ${\mathbf{x}}_0$.

This equation is powerful. It shows that ${\mathbf{x}}_t$ is just a weighted sum of the original "signal" (${\mathbf{x}}_0$) and "noise" ($\mathbf{ϵ}$). The term ${\bar{\alpha }}_t$ controls the balance between the two.

4. Implementation and Verification in PyTorch

Now, let's translate this theory into code. Your task is to implement the forward diffusion process and visualize its effects on an image. We'll use PyTorch.

Follow these steps to build your simulation:

Step 1: Define the Variance Schedule (${\beta }_t$)

First, create a linear variance schedule for a chosen number of timesteps, T. This schedule will contain the values of ${\beta }_t$ for $t=0,\dots ,T-1$.

import torch
import torchvision
import matplotlib.pyplot as plt

def linear_beta_schedule(timesteps, start=0.0001, end=0.02):
    return torch.linspace(start, end, timesteps)

# Example usage
T = 300
betas = linear_beta_schedule(timesteps=T)

Step 2: Pre-compute Coefficients

Using the betas tensor, calculate the ${\alpha }_t$ and ${\bar{\alpha }}_t$ values. The cumulative product function torch.cumprod is perfect for calculating ${\bar{\alpha }}_t$.

alphas = 1. - betas
alphas_cumprod = torch.cumprod(alphas, axis=0)

The resource Denoising Diffusion Probabilistic Models (DDPM) from labml.ai provides a clean implementation of this setup within a class structure, which is a common practice. You can find these lines of code in the __init__ method.

Step 3: Implement the q_sample Function

This is the core of your implementation. Write a function that takes an initial image x0, a timestep t, and the pre-computed alphas_cumprod tensor, and returns the noised image ${\mathbf{x}}_t$ using the direct sampling formula.

def q_sample(x0, t, alphas_cumprod):
    """
    Sample a noisy image x_t from x_0 at a specific timestep t.
    """
    # Get the value of alpha_bar_t for the given timestep t
    # Note: t is a tensor of timesteps
    alpha_bar_t = alphas_cumprod[t].view(-1, 1, 1, 1) # Reshape for broadcasting
    
    # Generate random noise
    epsilon = torch.randn_like(x0)
    
    # Apply the formula
    mean = torch.sqrt(alpha_bar_t) * x0
    std_dev = torch.sqrt(1. - alpha_bar_t)
    
    xt = mean + std_dev * epsilon
    
    return xt

This implementation is central to any DDPM. For a detailed code walkthrough that links each line back to the mathematical formulas, check out the following videos.

Step 4: Visualize the Process

Finally, let's verify that your implementation works as expected. Load an image, apply your q_sample function for several different timesteps, and plot the results.

# (Assuming you have defined T, betas, alphas, alphas_cumprod)

# Load an image (e.g., from torchvision)
# Note: You'll need to preprocess it to the range [-1, 1]
# For simplicity, we'll create a dummy tensor here.
# In a real scenario, you'd load an image.
from PIL import Image
import requests
from io import BytesIO
import torchvision.transforms as transforms

# Let's use an example image URL
url = 'https://upload.wikimedia.org/wikipedia/commons/5/57/A_cat_on_a_motorcycle_in_the_medina_of_Tunis_2017.jpg'
response = requests.get(url)
img = Image.open(BytesIO(response.content))

# Preprocess the image
img_size = 64
transform = transforms.Compose([
    transforms.Resize((img_size, img_size)),
    transforms.ToTensor(), # Scales to [0, 1]
    transforms.Lambda(lambda t: (t * 2) - 1) # Scales to [-1, 1]
])
x0 = transform(img).unsqueeze(0) # Add a batch dimension

# Visualize at different timesteps
timesteps_to_show = [0, 50, 100, 150, 200, 299]
fig, axes = plt.subplots(1, len(timesteps_to_show), figsize=(15, 3))

for i, t_int in enumerate(timesteps_to_show):
    t = torch.tensor([t_int])
    xt = q_sample(x0, t, alphas_cumprod)
    
    # Reverse the transformation for plotting
    img_plot = (xt.squeeze(0).clamp(-1, 1) + 1) / 2
    axes[i].imshow(img_plot.permute(1, 2, 0))
    axes[i].set_title(f"Timestep {t_int}")
    axes[i].axis('off')

plt.show()

Running this code should produce a series of images, starting with your original image and progressively becoming noisier, eventually looking like static. This visually confirms that you have successfully implemented the forward diffusion process.

Conclusion

In this lesson, you've established the theoretical and practical foundation of the diffusion process. Here are the key takeaways:

• The forward process is a fixed Markov chain that gradually adds Gaussian noise to data according to a variance schedule ${\beta }_t$.
• A crucial "nice property" allows us to sample a noised image ${\mathbf{x}}_t$ directly from the original ${\mathbf{x}}_0$ using the formula ${\mathbf{x}}_t=\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}\mathbf{ϵ}$.
• This direct sampling equation is efficient and forms the basis for training diffusion models. You've now implemented and verified it in PyTorch.

You're now perfectly positioned for the next lesson, where we'll tackle the more complex and exciting part: the reverse process. We will introduce a neural network designed to predict the noise that was added, and by doing so, learn to reverse the diffusion process one step at a time to generate new, high-quality images from pure noise.