Parameterizing the Reverse Process with Neural Networks

As established previously, the true reverse transition probability $p(x_{t-1}|x_t)$ is what we need to sample from to generate data, moving backward from noise $x_T$ towards a clean sample $x_0$. However, calculating this distribution requires knowing the original data distribution $q(x_0)$, making it intractable.

Our solution is to approximate this reverse transition using a neural network. We define a parameterized distribution $p_\theta(x_{t-1}|x_t)$ that our model will learn. Since the forward process adds Gaussian noise, a reasonable choice for the reverse transition is also a Gaussian distribution:

$$p_\theta(x_{t-1}|x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \Sigma_\theta(x_t, t))$$

Here, $\mu_\theta(x_t, t)$ is the mean and $\Sigma_\theta(x_t, t)$ is the covariance matrix of the Gaussian distribution for $x_{t-1}$, given $x_t$ at timestep $t$. Both the mean and covariance could, in principle, depend on the noisy data $x_t$ and the timestep $t$, and they are determined by the parameters $\theta$ of our neural network.

The task of the neural network is therefore to predict the parameters of this Gaussian distribution for each timestep $t$.

Predicting the Mean $\mu_\theta(x_t, t)$

There are two common ways to parameterize the mean $\mu_\theta(x_t, t)$:

1. Predict the Mean Directly: The network could be designed to directly output $\mu_\theta(x_t, t)$.
2. Predict the Noise $\epsilon$: The network predicts the noise component $\epsilon$ that was added during the corresponding forward step $t$. This approach is more common and often leads to more stable training and a simpler objective function (as we'll see in the next chapter).

Let's focus on the second approach, predicting the noise. We design a neural network, often denoted as $\epsilon_\theta$, which takes the current noisy sample $x_t$ and the timestep $t$ as input and outputs a predicted noise component $\epsilon_\theta(x_t, t)$.

Why is predicting the noise useful? Recall from the forward process (Chapter 2, "Sampling from Intermediate Steps") that the posterior distribution $q(x_{t-1}|x_t, x_0)$ is tractable and also Gaussian. Its mean has a specific form. While we don't know $x_0$ during the reverse process, we can use the relationship derived for $q(x_{t-1}|x_t, x_0)$ and substitute $x_0$ with an estimate derived from $x_t$ and the predicted noise $\epsilon_\theta(x_t, t)$. This leads to the following expression for the mean of our parameterized reverse distribution $p_\theta(x_{t-1}|x_t)$:

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

Where:

• $\epsilon_\theta(x_t, t)$ is the output of our neural network.
• $\alpha_t = 1 - \beta_t$ and $\bar{\alpha}_t = \prod_{i=1}^t \alpha_i$ are derived from the forward process noise schedule ($\beta_t$).

This equation provides a direct way to calculate the mean of the denoising step $x_{t-1}$ using the current state $x_t$ and the network's noise prediction $\epsilon_\theta(x_t, t)$. The network's primary job is to learn the correct noise $\epsilon$ that was added at step $t$, given the result $x_t$.

Fixing the Variance $\Sigma_\theta(x_t, t)$

What about the covariance matrix $\Sigma_\theta(x_t, t)$? While the network could also learn the variance, Ho et al. (authors of DDPM) found that fixing the variance works well and simplifies the model. The covariance is typically set to be diagonal, $\Sigma_\theta(x_t, t) = \sigma_t^2 \mathbf{I}$, where $\mathbf{I}$ is the identity matrix. The variance $\sigma_t^2$ is often chosen based on the forward process variances:

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

Both options connect the reverse process variance to the noise schedule used in the forward process. Using $\tilde{\beta}_t$ corresponds to the variance of the true posterior $q(x_{t-1}|x_t, x_0)$ when $x_0$ is known, while using $\beta_t$ performs well empirically, especially for the $L_{simple}$ objective function we'll discuss later. For simplicity, $\sigma_t^2 = \beta_t$ is a common choice. Fixing the variance means the neural network only needs to predict the noise $\epsilon_\theta(x_t, t)$ to define the reverse transition.

The Neural Network Architecture

The neural network $\epsilon_\theta(x_t, t)$ needs to process an input $x_t$ (which usually has the same dimensions as the original data, e.g., an image) and the timestep $t$ (a scalar). It must then output a tensor $\epsilon_\theta$ of the same shape as $x_t$, representing the predicted noise.

A very successful architecture for this task is the U-Net, which we will explore in detail in Chapter 4. The U-Net architecture is well-suited for image-like data and allows for effective integration of the timestep information $t$.

Diagram illustrating the role of the neural network $\epsilon_\theta$ in parameterizing the reverse process. It takes the noisy data $x_t$ and timestep $t$ to predict the noise $\epsilon_\theta$. This prediction is then used, along with fixed parameters from the noise schedule ($\alpha_t, \beta_t, \bar{\alpha}_t$), to calculate the mean $\mu_\theta$ of the approximate reverse transition $p_\theta(x_{t-1}|x_t)$. The variance $\sigma_t^2$ is typically fixed.

In summary, we approximate the intractable reverse distribution $p(x_{t-1}|x_t)$ with a learned Gaussian distribution $p_\theta(x_{t-1}|x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \sigma_t^2 \mathbf{I})$. We use a neural network $\epsilon_\theta(x_t, t)$ to predict the noise component, which then allows us to compute the mean $\mu_\theta(x_t, t)$. The variance $\sigma_t^2$ is usually fixed based on the forward process noise schedule. This setup forms the basis for training the diffusion model, which we will cover next.