DDPM Forward and Reverse Processes

Hello! Welcome to your next lesson in our journey through generative models.

In our last session, we dove into advanced GAN architectures like Conditional GANs and StyleGAN, learning how to gain explicit control over generation and push image quality to photorealistic levels. We saw how architectural innovations like the U-Net generator, PatchGAN, and StyleGAN's mapping network created a new standard for generative modeling.

Today, we pivot to an entirely different and incredibly powerful paradigm: Denoising Diffusion Probabilistic Models (DDPMs). While GANs learn to generate in a single step, diffusion models learn to generate images by gradually refining them from pure noise. This iterative approach has proven to be remarkably stable and capable of producing state-of-the-art results, forming the backbone of modern marvels like DALL-E 2 and Stable Diffusion.

Your learning outcome for this lesson is to implement the forward and reverse processes of Denoising Diffusion Probabilistic Models (DDPM).

We will break this down into three core parts:

1. The Forward Process: A fixed procedure of systematically adding noise to an image.
2. The Reverse Process: The learned process of iteratively removing that noise.
3. The Training Objective: A surprisingly simple objective where a neural network learns to predict the noise itself.

By the end of this lesson, you will understand the theory behind DDPMs and have a clear map for implementing one from scratch.

1. The Core Idea of Diffusion

Diffusion models consist of two opposing processes:

• Forward Process (Diffusion): We start with a clean image from our dataset, $x_0$. Over a large number of timesteps $T$ (e.g., 1000), we progressively add a small amount of Gaussian noise at each step. By the final step $T$, the image is indistinguishable from pure Gaussian noise. This is a fixed, non-learned process.
• Reverse Process (Denoising): We train a neural network to reverse this process. Starting with a sample of pure noise, $x_T$, the model learns to iteratively denoise it, step by step, until it produces a clean image $x_0$.

This elegant concept of destruction and creation is the heart of diffusion models. The video below provides a fantastic animated overview of this core idea.

A conceptual overview of the forward and reverse diffusion processes. The forward process q adds noise, while the reverse process p learns to remove it. Source: LearnOpenCV.

2. The Forward Process: A Mathematical Walkthrough

The forward process is defined as a Markov chain where we add Gaussian noise at each timestep $t$ according to a variance schedule $\{{\beta }_t{\}}_{t=1}^T$. The distribution of a noisy image $x_t$ given the previous one $x_{t-1}$ is:

$q(x_t|x_{t-1})=\mathcal{N}(x_t;\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_t\mathbf{I})$

Here, $\sqrt{1-{\beta }_t}$ is a scaling factor that keeps the variance from exploding. The values of ${\beta }_t$ are typically small and increase over time (e.g., from 0.0001 to 0.02).

A crucial property of this process is that we don't need to iterate $t$ times to get $x_t$. We can sample $x_t$ directly from the original image $x_0$ in one step using a closed-form equation.

Let ${\alpha }_t=1-{\beta }_t$ and ${\bar{\alpha }}_t={\prod }_{s=1}^t{\alpha }_s$. Then, the distribution of $x_t$ given $x_0$ is:

$q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{{\bar{\alpha }}_t}{x}_0,(1-{\bar{\alpha }}_t)\mathbf{I})$

This equation is fundamental for efficient training. It means we can create a noisy version of an image for any timestep $t$ instantly. Using the reparameterization trick, we can write the sampling operation as:

$x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ$

where $ϵ\sim \mathcal{N}(0,\mathbf{I})$ is a random noise sample.

The following resources provide a clear derivation of this process and show how it translates directly into code.

3. The Reverse Process & The Simplified Loss

The goal of the reverse process is to learn the distribution $p_{\theta }(x_{t-1}|x_t)$, where $\theta$ represents the parameters of our neural network. The key insight from the DDPM paper is that if the ${\beta }_t$ steps are small enough, this reverse transition is also a Gaussian. Our network, therefore, needs to learn the mean and variance of this Gaussian.

Deriving the full training objective involves maximizing the Evidence Lower Bound (ELBO) on the log-likelihood, which can be mathematically intensive. However, it leads to a remarkably simple and intuitive result. After several steps of simplification, the authors found that instead of training the network to predict the mean of the denoised image, it's more effective to reparameterize the mean and train the network to predict the noise $ϵ$ that was added to the image at timestep $t$.

This transforms the training objective into a simple Mean Squared Error (MSE) loss:

$L_{\text{simple}}(\theta )={\mathbb{E}}_{t,x_0,ϵ}\left[\Vert ϵ-{ϵ}_{\theta }(x_t,t){\Vert }^2\right]$

Where:

• $ϵ$ is the true Gaussian noise added at timestep $t$.
• ${ϵ}_{\theta }(x_t,t)$ is the noise predicted by our neural network.
• The network, ${ϵ}_{\theta }$, is typically a U-Net which takes the noisy image $x_t$ and the timestep $t$ as input.

The following video masterfully walks through this complex derivation, building intuition at each step.

4. Implementation: Putting It All Together

Now that we have the theory for both processes, let's see how they are implemented in a full training and sampling loop. Your background in software engineering and PyTorch will make following this code-centric explanation straightforward.

The Training Algorithm

For each step in training:

1. Load a batch of real images $x_0$.
2. Sample a random timestep $t$ for each image in the batch.
3. Sample a batch of Gaussian noise $ϵ$.
4. Create the noisy images $x_t$ using the q_sample function: $x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ$.
5. Pass $x_t$ and $t$ to the U-Net model to get the predicted noise ${ϵ}_{\theta }$.
6. Calculate the MSE loss between the real noise $ϵ$ and predicted noise ${ϵ}_{\theta }$.
7. Update the model weights using backpropagation.

The Sampling (Inference) Algorithm

To generate a new image:

1. Start with a pure Gaussian noise tensor $x_T$.
2. Iterate backwards from $t=T$ down to $1$:
   a. Feed the current noisy image $x_t$ and the timestep $t$ into the U-Net to predict the noise ${ϵ}_{\theta }$.
   b. Use this predicted noise to calculate the mean of the reverse distribution.
   c. Sample $x_{t-1}$ from this distribution, usually by taking the mean and adding a small amount of noise (if $t>1$).
3. The final result $x_0$ is the generated image.

The video below is an exceptional, detailed walkthrough of a real-world DDPM implementation. It connects every line of code back to the formulas in the papers.

Conclusion

Fantastic work today! We've unpacked the theory and implementation of Denoising Diffusion Probabilistic Models, a cornerstone of modern generative AI. You've moved beyond the adversarial training of GANs to understand a new, iterative approach to generation.

Key Takeaways:

• Two-Process System: Diffusion models are defined by a fixed forward (noising) process and a learned reverse (denoising) process.
• Efficient Forward Sampling: The forward process allows for efficient, one-step sampling of a noisy image $x_t$ for any timestep $t$ from the original image $x_0$.
• Noise Prediction Objective: The model (a U-Net) is trained on a simple objective: to predict the noise $ϵ$ that was added to an image, using an MSE loss.
• Iterative Generation: During inference, new images are generated by starting with pure noise and iteratively applying the trained model to denoise the image step-by-step.

Preview of the next lesson:
While DDPMs are powerful, they have one significant drawback: the diffusion process operates directly in the high-dimensional pixel space. Generating a 512x512 image requires running a large U-Net on a 512x512x3 tensor for hundreds of steps, which is computationally immense.

In our next lesson, we will tackle this challenge by exploring Latent Diffusion Models (LDMs). You will learn how to implement the architecture that made models like Stable Diffusion possible by moving the entire diffusion process from pixel space into a much smaller, more manageable latent space.