Mathematical Formulation of the Denoising Step

Okay, we've established that the reverse diffusion process aims to reverse the noising effects, starting from a noisy sample $x_t$ and trying to estimate the slightly less noisy sample $x_{t-1}$. We also know that the true reverse probability $q(x_{t-1}|x_t)$ is intractable, but we can approximate it using a neural network, $p_\theta(x_{t-1}|x_t)$. Our network $\epsilon_\theta(x_t, t)$ is trained to predict the noise $\epsilon$ that was added to the original data $x_0$ to get $x_t$. Now, let's see how we use this predicted noise to define the mathematical operation for one step of denoising.

The goal is to define the distribution $p_\theta(x_{t-1}|x_t)$. It turns out that if we knew the original data point $x_0$, the true posterior distribution $q(x_{t-1}|x_t, x_0)$ can be calculated analytically and is a Gaussian distribution. While we won't go through the full derivation here (it involves applying Bayes' theorem to the Gaussian distributions of the forward process, as detailed in Appendix B of the original DDPM paper by Ho et al. 2020), the result is important. The posterior $q(x_{t-1}|x_t, x_0)$ is given by:

$$q(x_{t-1}|x_t, x_0) = \mathcal{N}(x_{t-1}; \tilde{\mu}_t(x_t, x_0), \tilde{\beta}_t I)$$

where the mean $\tilde{\mu}_t(x_t, x_0)$ and variance $\tilde{\beta}_t$ are:

$$\tilde{\mu}_t(x_t, x_0) = \frac{\sqrt{\bar{\alpha}_{t-1}} \beta_t}{1 - \bar{\alpha}_t} x_0 + \frac{\sqrt{\alpha_t} (1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} x_t$$ $$\tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \beta_t$$

Remember that $\alpha_t = 1 - \beta_t$ and $\bar{\alpha}_t = \prod_{i=1}^t \alpha_i$ are derived from the noise schedule $\beta_t$.

Now, during the reverse (generation) process, we don't have access to the original $x_0$. This is where our trained neural network $\epsilon_\theta(x_t, t)$ comes in. We use it to approximate $x_0$. Recall the forward process equation that allows sampling $x_t$ directly from $x_0$:

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

where $\epsilon \sim \mathcal{N}(0, I)$. We can rearrange this to express $x_0$ in terms of $x_t$ and $\epsilon$:

$$x_0 = \frac{1}{\sqrt{\bar{\alpha}_t}} (x_t - \sqrt{1 - \bar{\alpha}_t} \epsilon)$$

Since our network $\epsilon_\theta(x_t, t)$ is trained to predict $\epsilon$, we can get an estimate of $x_0$, let's call it $\hat{x}_0$, using the network's output:

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

This gives us an approximation of the original image based on the current noisy image $x_t$ and the predicted noise $\epsilon_\theta(x_t, t)$.

Now we substitute this estimate $\hat{x}_0$ back into the equation for the posterior mean $\tilde{\mu}_t(x_t, x_0)$. This gives us the mean for our parameterized reverse step, $\mu_\theta(x_t, t)$:

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

This looks complicated, but after some algebraic simplification (using the relationships $\alpha_t = 1 - \beta_t$ and $\bar{\alpha}_t = \alpha_t \bar{\alpha}_{t-1}$), it reduces to a much cleaner form:

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

This equation is the heart of the denoising step. It tells us how to calculate the mean of the distribution for the previous timestep $x_{t-1}$, given the current state $x_t$ and the noise $\epsilon_\theta(x_t, t)$ predicted by our U-Net model for that timestep.

For the variance of the reverse step $p_\theta(x_{t-1}|x_t)$, we need to choose a value $\sigma_t^2$. The DDPM paper proposes using the variance derived from the posterior $q(x_{t-1}|x_t, x_0)$, which is $\sigma_t^2 = \tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \beta_t$. Another option explored later is simply setting $\sigma_t^2 = \beta_t$. This variance term introduces stochasticity into the generation process; sampling from this distribution involves adding Gaussian noise scaled by $\sigma_t$.

So, a single step in the reverse diffusion process is defined by sampling $x_{t-1}$ from the Gaussian distribution:

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

Where:

• $x_t$ is the noisy data at the current timestep.
• $\epsilon_\theta(x_t, t)$ is the noise predicted by the neural network.
• $\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)$ is the predicted mean for $x_{t-1}$.
• $\sigma_t^2$ is the variance, often chosen as $\tilde{\beta}_t$ or $\beta_t$.
• $\alpha_t$, $\beta_t$, $\bar{\alpha}_t$ are constants determined by the noise schedule.

By repeatedly applying this formula, starting from pure noise $x_T \sim \mathcal{N}(0, I)$ and stepping backwards from $t=T$ down to $t=1$, we can generate a sample $x_0$ that should resemble the data the model was trained on. This iterative process, guided at each step by the neural network's noise prediction, is how diffusion models generate new data. We will explore the full sampling algorithms in Chapter 5.