Predicting the Noise Component

As we established, directly computing the true reverse probability $p(x_{t-1}|x_t)$ is generally intractable. Our strategy is to approximate this reverse transition using a parameterized distribution, typically a Gaussian, learned by a neural network: $p_\theta(x_{t-1}|x_t)$. The network needs to estimate the parameters of this Gaussian distribution, specifically its mean $\mu_\theta(x_t, t)$ and variance $\Sigma_\theta(x_t, t)$.

While the network could be trained to directly predict the mean $\mu_\theta(x_t, t)$ or even the denoised sample $x_{t-1}$, a different approach has proven remarkably effective in practice: predicting the noise component $\epsilon$ that was added during the forward process at timestep $t$.

Let's see why this is useful. Recall the closed-form expression for sampling $x_t$ directly from $x_0$ in the forward process:

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

Here, $\epsilon$ is a standard Gaussian noise sample, and $\bar{\alpha}_t$ is derived from the noise schedule. This equation links the original data $x_0$, the noisy version $x_t$, and the noise $\epsilon$ added to reach that state.

If our neural network, let's call it $\epsilon_\theta$, can accurately predict the noise $\epsilon$ based on the noisy input $x_t$ and the timestep $t$, we can use this prediction, $\epsilon_\theta(x_t, t)$, to inform our estimate of the previous state $x_{t-1}$.

How does predicting noise help estimate the mean $\mu_\theta(x_t, t)$ of the reverse step $p_\theta(x_{t-1}|x_t)$? The original Denoising Diffusion Probabilistic Models (DDPM) paper showed that the mean of the reverse transition $p(x_{t-1}|x_t, x_0)$ can be expressed as:

$$\tilde{\mu}_t(x_t, x_0) = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \epsilon \right)$$

where $\alpha_t = \bar{\alpha}_t / \bar{\alpha}_{t-1}$. Notice this expression depends on the original data $x_0$ (through $\epsilon$) which we don't have during generation.

However, we can rearrange the first equation to get an estimate of $x_0$ if we know $x_t$ and $\epsilon$:

$$x_0 \approx \hat{x}_0 = \frac{1}{\sqrt{\bar{\alpha}_t}}(x_t - \sqrt{1 - \bar{\alpha}_t} \epsilon_\theta(x_t, t))$$

By substituting our network's noise prediction $\epsilon_\theta(x_t, t)$ for $\epsilon$ in the equation for $\tilde{\mu}_t$, we arrive at an expression for the mean of our approximate reverse transition $p_\theta(x_{t-1}|x_t)$:

$$\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)$$

This establishes a direct link: if our network $\epsilon_\theta(x_t, t)$ successfully predicts the noise added at step $t$, we can calculate the mean required for the denoising step $x_t \rightarrow x_{t-1}$. The variance $\Sigma_\theta(x_t, t)$ is often fixed to a value related to the forward process variances, or sometimes also learned, but predicting the noise is primarily used to determine the mean.

Why Predict Noise?

Parameterizing the reverse process by predicting noise offers several advantages:

1. Simpler Learning Target: The noise $\epsilon$ is typically drawn from a simple distribution (e.g., standard Gaussian). Predicting this noise might be an easier task for the neural network compared to directly predicting the complex structure of the data $x_0$ or $x_{t-1}$, especially at high noise levels (large $t$).
2. Alignment with Loss Function: As we will see in the next chapter, the training objective for diffusion models often simplifies to a mean squared error loss between the true noise $\epsilon$ added during the forward process and the network's predicted noise $\epsilon_\theta(x_t, t)$. Training the network to directly output $\epsilon_\theta$ perfectly aligns the network's output with this common loss function, potentially leading to more stable training.
3. Empirical Success: This noise prediction formulation has been central to the success of many high-performing diffusion models, particularly in image generation.

Therefore, the standard approach involves training a neural network $\epsilon_\theta$ that takes the noisy data $x_t$ and the timestep $t$ as input and outputs a prediction of the noise $\epsilon$ that was used to generate $x_t$ from $x_0$. This predicted noise $\epsilon_\theta(x_t, t)$ then allows us to compute the parameters (specifically the mean) of the approximate reverse distribution $p_\theta(x_{t-1}|x_t)$, enabling the step-by-step generation process.

The neural network $\epsilon_\theta$ takes the current noisy sample $x_t$ and the timestep $t$ as input. Its goal is to predict the noise $\epsilon_\theta$ that was likely added to the original data to produce $x_t$. This prediction is the core component used to guide the reverse denoising step.