Understanding Diffusion Models: Math and Intuition

Diffusion models (like Stable Diffusion and Midjourney) have revolutionized generative AI. Unlike GANs, which use adversarial training, diffusion models are grounded in non-equilibrium thermodynamics.

While the math can seem daunting, the core intuition is surprisingly elegant.

The Core Intuition

Imagine a drop of ink falling into a glass of water. Over time, the ink diffuses until the water is a uniform light blue. This process is easy to simulate: just add random noise at each step.

Now, imagine reversing time. You start with light blue water and somehow extract the ink drop. This is incredibly difficult.

Diffusion models learn to reverse this process. They learn to take pure noise and iteratively remove it to reveal a structured image.

The Forward Process (Adding Noise)

The forward process (or diffusion process) is a Markov chain that gradually adds Gaussian noise to the data over $T$ steps.

Let $x_0$ be the original image. We define the forward process as:

$q(x_t | x_{t-1}) = \mathcal{N}(x_t; \sqrt{1 - \beta_t} x_{t-1}, \beta_t I)$

Here, $\beta_t$ is the variance schedule. It determines how much noise is added at each step. By the end of $T$ steps, $x_T$ is nearly an isotropic Gaussian distribution (pure noise).

A key property of this process is that we can jump directly to any timestep $t$ without simulating all previous steps, using the reparameterization trick:

$x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon$

Where $\epsilon \sim \mathcal{N}(0, I)$ and $\alpha_t = 1 - \beta_t$, and $\bar{\alpha}_t$ is the cumulative product of $\alpha_s$ up to $t$.

The Reverse Process (Denoising)

The goal of our neural network (usually a U-Net) is to learn the reverse process: starting from $x_T \sim \mathcal{N}(0, I)$, we want to sample $x_{t-1} \sim p_\theta(x_{t-1} | x_t)$.

Because the forward steps are small, the reverse steps can also be modeled as Gaussian distributions. Our neural network learns to predict the parameters of these Gaussians.

Specifically, it's often formulated such that the neural network $\epsilon_\theta(x_t, t)$ tries to predict the noise $\epsilon$ that was added to $x_0$ to get $x_t$.

The Loss Function

The training objective is surprisingly simple. We sample a random image $x_0$, a random timestep $t$, and random noise $\epsilon$. We create the noisy image $x_t$, and train the network to predict $\epsilon$.

The simplified loss function is Mean Squared Error (MSE) between the actual noise and the predicted noise:

$L_{simple}(\theta) = \mathbb{E}_{t, x_0, \epsilon} [ || \epsilon - \epsilon_\theta(x_t, t) ||^2 ]$

Code Snippet: The U-Net Objective

python
import torch
import torch.nn.functional as F

def diffusion_loss(model, x_0, t, noise_schedule):
    # 1. Sample random noise
    noise = torch.randn_like(x_0)
    
    # 2. Get scheduling parameters for timestep t
    sqrt_alpha_bar = noise_schedule.sqrt_alpha_bar[t]
    sqrt_one_minus_alpha_bar = noise_schedule.sqrt_one_minus_alpha_bar[t]
    
    # 3. Add noise to the original image (Forward Process)
    # Broadcasting magic assumes shape alignment
    x_t = sqrt_alpha_bar * x_0 + sqrt_one_minus_alpha_bar * noise
    
    # 4. Predict the noise using the U-Net model
    predicted_noise = model(x_t, t)
    
    # 5. Calculate the MSE loss
    loss = F.mse_loss(predicted_noise, noise)
    
    return loss

By repeatedly applying the reverse process, we can generate stunning, high-resolution images from pure static. The beauty of diffusion models lies in this stable, iterative refinement compared to the volatile training dynamics of GANs.