The DDPM Sampling Algorithm

Diffusion models, specifically the U-Net, are trained to predict the noise $\epsilon$ present in a noisy input $x_t$ at timestep $t$. This trained model can then be used to generate new data. The Denoising Diffusion Probabilistic Models (DDPM) algorithm provides the original, foundational procedure for this generative process.

As outlined in the chapter introduction, the core idea is to start with pure noise and progressively refine it using the learned denoising steps. We begin by sampling an initial tensor $x_T$ from a standard Gaussian distribution, $x_T \sim \mathcal{N}(0, \mathbf{I})$. This $x_T$ represents maximum entropy, pure noise, corresponding to the end state of the forward diffusion process.

The DDPM sampling algorithm then iteratively applies the learned reverse process, stepping backward from timestep $t=T$ down to $t=1$. In each step, the goal is to sample a slightly less noisy version $x_{t-1}$ given the current state $x_t$.

The Denoising Step

Recall from Chapter 3 that the reverse transition $p_\theta(x_{t-1} | x_t)$ is approximated by a Gaussian distribution whose mean $\mu_\theta(x_t, t)$ depends on $x_t$ and the predicted noise $\epsilon_\theta(x_t, t)$, and whose variance $\sigma_t^2$ is related to the noise schedule $\beta_t$.

The model $\epsilon_\theta(x_t, t)$ takes the current noisy image $x_t$ and the timestep $t$ as input and outputs its prediction of the noise component that was added to get from $x_0$ to $x_t$. Using this prediction, we can estimate the mean of the distribution for the previous state $x_{t-1}$. The equation for the mean $\mu_\theta(x_t, t)$ is derived from the properties of the forward and reverse processes:

$$\mu_\theta(x_t, t) = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \epsilon_\theta(x_t, t) \right)$$

Here, $\alpha_t = 1 - \beta_t$ and $\bar{\alpha}_t = \prod_{i=1}^t \alpha_i$ are parameters derived from the noise schedule $\beta_t$ used during the forward process. This equation essentially takes the current noisy sample $x_t$ and subtracts the scaled predicted noise $\epsilon_\theta$ to estimate the mean of the previous, less noisy state $x_{t-1}$.

The variance of the reverse step, $\sigma_t^2$, is also determined by the noise schedule. A common choice, derived to match the variance of the forward process posterior $q(x_{t-1} | x_t, x_0)$, is:

$$\sigma_t^2 = \tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \beta_t$$

Note that $\sigma_t^2 = 0$ for $t=1$. To perform one denoising step and sample $x_{t-1}$ from $x_t$, we calculate the mean $\mu_\theta(x_t, t)$ using the model's prediction and then add Gaussian noise scaled by the standard deviation $\sigma_t$:

$$x_{t-1} = \mu_\theta(x_t, t) + \sigma_t z$$

where $z \sim \mathcal{N}(0, \mathbf{I})$ is standard Gaussian noise. This added noise $z$ introduces stochasticity into the generation process, allowing the model to produce diverse samples even when starting from the same initial $x_T$ (though typically we start from different $x_T$ samples). However, for the very last step ($t=1$), we usually set $z=0$ to obtain the final deterministic mean prediction as our output $x_0$.

The Full DDPM Algorithm

Putting these steps together, the complete DDPM sampling algorithm proceeds as follows:

1. Initialization: Sample the initial noise tensor $x_T \sim \mathcal{N}(0, \mathbf{I})$. This tensor should have the same dimensions as the desired output data (e.g., image height, width, channels).
2. Iterative Denoising: Loop backward through the timesteps, from $t = T, T-1, \dots, 1$.
   • Sample standard Gaussian noise $z \sim \mathcal{N}(0, \mathbf{I})$. If $t=1$, set $z=0$.
   • Use the trained neural network $\epsilon_\theta$ to predict the noise present in the current state: $\epsilon_{pred} = \epsilon_\theta(x_t, t)$.
   • Calculate the mean of the reverse distribution using the predicted noise: $$\mu_\theta(x_t, t) = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \epsilon_{pred} \right)$$
   • Calculate the variance $\sigma_t^2 = \tilde{\beta}_t$ (pre-computed from the noise schedule).
   • Sample the next state (less noisy) using the calculated mean and variance: $$x_{t-1} = \mu_\theta(x_t, t) + \sigma_t z$$
3. Output: After the loop finishes (at $t=1$), the resulting tensor $x_0$ is the generated sample.

This process is visualized below:

Diagram illustrating the iterative DDPM sampling process. Starting from noise $x_T$, each step uses the noise predictor $\epsilon_\theta$ to compute the mean $\mu_\theta$ for the previous state, then samples $x_{t-1}$ by adding scaled noise $\sigma_t z$ (except for the last step where $z=0$).

The quality of the final sample $x_0$ depends heavily on the accuracy of the trained noise predictor $\epsilon_\theta$ and the chosen noise schedule ($\beta_t$) and number of diffusion steps ($T$). DDPM typically requires a large number of steps (e.g., $T=1000$) to achieve high-quality results, which can make sampling relatively slow. We will explore faster alternatives like DDIM in the subsequent sections.