Approximating the Reverse Transition

The forward process in diffusion models involves gradually adding noise to data through a fixed procedure. For generative AI, the objective is to reverse this process: starting from pure noise $x_T \sim \mathcal{N}(0, \mathbf{I})$, the aim is to sequentially sample $x_{T-1}$, then $x_{T-2}$, and so on, until a clean data sample $x_0$ is obtained that resembles the original data distribution. Achieving this requires calculating the reverse transition probabilities, $p(x_{t-1}|x_t)$.

If we could compute $p(x_{t-1}|x_t)$ exactly, we could sample from it iteratively to generate data. However, calculating this distribution directly poses a significant challenge. Why? Because the probability of transitioning back to a less noisy state $x_{t-1}$ depends not just on the current noisy state $x_t$, but implicitly on the entire distribution of possible starting data points $q(x_0)$. Mathematically, evaluating $p(x_{t-1}|x_t)$ involves marginalizing over all possible initial data points $x_0$:

$$p(x_{t-1}|x_t) = \int p(x_{t-1}|x_t, x_0) q(x_0|x_t) dx_0$$

The term $q(x_0|x_t)$ represents the probability of a specific starting data point $x_0$ given the noisy version $x_t$. Calculating this requires knowledge of the (unknown) true data distribution $q(x_0)$ and involves complex integration, making the true reverse probability $p(x_{t-1}|x_t)$ intractable to compute for complex datasets.

The Tractable Conditional Reverse Distribution

Interestingly, the situation changes if we know the starting point $x_0$ that led to $x_t$. The reverse conditional probability $p(x_{t-1}|x_t, x_0)$ is tractable. Using Bayes' theorem and the properties of the Gaussian noise added during the forward process $q(x_t|x_{t-1})$ (defined by the noise schedule $\beta_t$), we can show that this distribution is also a Gaussian:

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

Here, $\mathcal{N}(x; \mu, \sigma^2 \mathbf{I})$ denotes a Gaussian distribution with mean $\mu$ and diagonal covariance $\sigma^2 \mathbf{I}$. The variance $\tilde{\beta}_t$ and the mean $\tilde{\mu}_t(x_t, x_0)$ depend only on the known forward process noise schedule parameters ($\beta_t$ or equivalently $\alpha_t = 1 - \beta_t$ and $\bar{\alpha}_t = \prod_{s=1}^t \alpha_s$) and the specific values of $x_t$ and $x_0$. Specifically:

$$\tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \beta_t$$ $$\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$$

These equations tell us that if we knew the original image $x_0$ corresponding to a noisy image $x_t$, we could precisely calculate the distribution of the slightly less noisy image $x_{t-1}$ it came from.

The Need for Approximation

But here's the catch: during the actual generative process (sampling), we don't have $x_0$. Our goal is precisely to generate $x_0$ starting from pure noise $x_T$. Knowing $x_0$ would defeat the purpose.

Since the true reverse transition $p(x_{t-1}|x_t)$ is intractable, and the conditional reverse transition $p(x_{t-1}|x_t, x_0)$ requires the unknown $x_0$, we need an alternative. The solution is to approximate the true reverse transition using a learned model.

We introduce a parameterized distribution, typically a neural network, denoted as $p_\theta(x_{t-1}|x_t)$, to approximate the intractable true posterior $p(x_{t-1}|x_t)$. The parameters $\theta$ of this network will be learned from data.

Diagram illustrating the relationship between the intractable true reverse transition, the tractable conditional reverse transition (requiring $x_0$), and the learned approximation $p_\theta(x_{t-1}|x_t)$ using a neural network.

This neural network takes the current noisy state $x_t$ and the current timestep $t$ as input and outputs the parameters of the approximate reverse distribution $p_\theta(x_{t-1}|x_t)$. Since the tractable conditional reverse distribution $p(x_{t-1}|x_t, x_0)$ is Gaussian, it's convenient and effective to also model our approximation $p_\theta(x_{t-1}|x_t)$ as a Gaussian:

$$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, both predicted by the neural network $\theta$. The network's task is therefore to learn functions $\mu_\theta$ and $\Sigma_\theta$ that accurately predict the parameters of the distribution for $x_{t-1}$ given $x_t$.

In many successful diffusion models, like the original Denoising Diffusion Probabilistic Models (DDPM), the covariance matrix $\Sigma_\theta(x_t, t)$ is not learned directly. Instead, it's often fixed to a value related to the forward process variance, such as $\Sigma_\theta(x_t, t) = \tilde{\beta}_t \mathbf{I}$ or $\Sigma_\theta(x_t, t) = \beta_t \mathbf{I}$. This simplifies the learning problem significantly: the neural network only needs to learn the mean $\mu_\theta(x_t, t)$ of the reverse transition.

The next step is to understand how we parameterize this neural network, specifically how it predicts the mean $\mu_\theta(x_t, t)$, and how we train it to effectively reverse the diffusion process. As we will see, a common and effective strategy is to train the network to predict the noise that was added at timestep $t$.