Sampling from Intermediate Steps

In the previous sections, we established the forward diffusion process as a step-by-step Markov chain where a small amount of Gaussian noise is added at each timestep $t$, governed by the variance schedule $\beta_t$. The transition is defined as $q(x_t | x_{t-1})$. While this defines the process, simulating it step-by-step to get a noisy sample $x_t$ from an initial data point $x_0$ can be computationally expensive, especially for large $T$. Fortunately, there's a more direct way.

A significant property of this specific noising process is that we can derive a closed-form equation to sample $x_t$ directly from $x_0$ for any timestep $t$, without needing to compute all the intermediate states $x_1, x_2, ..., x_{t-1}$. This is extremely useful, particularly during the training phase of the diffusion model.

Let's derive this relationship. Recall the single-step transition:

$$x_t = \sqrt{1 - \beta_t} x_{t-1} + \sqrt{\beta_t} \epsilon_{t-1} \quad \text{where} \quad \epsilon_{t-1} \sim \mathcal{N}(0, I)$$

For convenience, let's define $\alpha_t = 1 - \beta_t$. The equation becomes:

$$x_t = \sqrt{\alpha_t} x_{t-1} + \sqrt{1 - \alpha_t} \epsilon_{t-1}$$

Now, we can recursively expand this expression. Let's look at $x_{t-1}$ in terms of $x_{t-2}$:

$$x_{t-1} = \sqrt{\alpha_{t-1}} x_{t-2} + \sqrt{1 - \alpha_{t-1}} \epsilon_{t-2}$$

Substituting this into the equation for $x_t$:

$$\begin{align*} x_t &= \sqrt{\alpha_t} (\sqrt{\alpha_{t-1}} x_{t-2} + \sqrt{1 - \alpha_{t-1}} \epsilon_{t-2}) + \sqrt{1 - \alpha_t} \epsilon_{t-1} \\ &= \sqrt{\alpha_t \alpha_{t-1}} x_{t-2} + \sqrt{\alpha_t(1 - \alpha_{t-1})} \epsilon_{t-2} + \sqrt{1 - \alpha_t} \epsilon_{t-1} \end{align*}$$

Notice a pattern emerging. The term multiplying $x_{t-2}$ is the product of the square roots of the $\alpha$ values. The noise terms are also accumulating. We can leverage a property of Gaussian distributions: adding two independent Gaussian variables results in another Gaussian variable. Specifically, if $Z_1 \sim \mathcal{N}(0, \sigma_1^2 I)$ and $Z_2 \sim \mathcal{N}(0, \sigma_2^2 I)$ are independent, then $Z_1 + Z_2 \sim \mathcal{N}(0, (\sigma_1^2 + \sigma_2^2) I)$.

In our expansion, $\epsilon_{t-1}$ and $\epsilon_{t-2}$ are independent standard Gaussian noises ($\mathcal{N}(0, I)$). The combined noise term $\sqrt{\alpha_t(1 - \alpha_{t-1})} \epsilon_{t-2} + \sqrt{1 - \alpha_t} \epsilon_{t-1}$ is also Gaussian. Its variance is $(\alpha_t(1 - \alpha_{t-1})) Var(\epsilon_{t-2}) + (1 - \alpha_t) Var(\epsilon_{t-1}) = \alpha_t(1 - \alpha_{t-1}) + (1 - \alpha_t) = \alpha_t - \alpha_t\alpha_{t-1} + 1 - \alpha_t = 1 - \alpha_t\alpha_{t-1}$. So, we can rewrite the combined noise term as $\sqrt{1 - \alpha_t \alpha_{t-1}} \bar{\epsilon}_{t-2}$, where $\bar{\epsilon}_{t-2} \sim \mathcal{N}(0, I)$.

This leads to:

$$x_t = \sqrt{\alpha_t \alpha_{t-1}} x_{t-2} + \sqrt{1 - \alpha_t \alpha_{t-1}} \bar{\epsilon}_{t-2}$$

If we continue this expansion all the way back to $x_0$, we introduce the cumulative product notation $\bar{\alpha}_t = \prod_{i=1}^t \alpha_i$. The general form becomes:

$$x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon \quad \text{where} \quad \epsilon \sim \mathcal{N}(0, I)$$

This elegant equation is a cornerstone of diffusion models. It tells us that any noisy version $x_t$ can be obtained directly from the original data $x_0$ by scaling $x_0$ down by $\sqrt{\bar{\alpha}_t}$, generating a standard Gaussian noise vector $\epsilon$, scaling it by $\sqrt{1 - \bar{\alpha}_t}$, and adding the two results.

Equivalently, we can say that the conditional distribution $q(x_t | x_0)$ is a Gaussian distribution:

$$q(x_t | x_0) = \mathcal{N}(x_t; \sqrt{\bar{\alpha}_t} x_0, (1 - \bar{\alpha}_t) I)$$

The mean of this distribution is $\sqrt{\bar{\alpha}_t} x_0$, and the variance is $(1 - \bar{\alpha}_t) I$.

Since $\alpha_t = 1 - \beta_t$ and $\beta_t$ is typically small and positive, $\alpha_t$ is slightly less than 1. The cumulative product $\bar{\alpha}_t$ starts at $\bar{\alpha}_0 = 1$ (by convention) and decreases monotonically towards 0 as $t$ increases towards the total number of steps $T$. Consequently, $\sqrt{\bar{\alpha}_t}$ (the "signal rate") decreases from 1 to 0, while $\sqrt{1 - \bar{\alpha}_t}$ (the "noise rate") increases from 0 to 1. This matches our intuition: as $t$ increases, the contribution of the original data $x_0$ diminishes, and the sample $x_t$ becomes increasingly dominated by noise, eventually approaching a standard Gaussian distribution $\mathcal{N}(0, I)$ when $\bar{\alpha}_T \approx 0$.

This plot shows the typical evolution of the signal rate ($\sqrt{\bar{\alpha}_t}$) and noise rate ($\sqrt{1 - \bar{\alpha}_t}$) over 1000 timesteps using a linear variance schedule from $\beta_1 = 10^{-4}$ to $\beta_{1000} = 0.02$. As $t$ increases, the influence of the original data decreases while the influence of noise increases.

Importance for Training

This closed-form expression $x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon$ is important for training diffusion models efficiently. The goal of training is to learn a neural network $\epsilon_\theta(x_t, t)$ that can predict the noise $\epsilon$ that was added to get $x_t$. To do this, we need training examples consisting of noisy data $x_t$ and the corresponding noise $\epsilon$ that was used to generate it.

Using the formula, we can create these training pairs quickly:

1. Choose a data point $x_0$ from your dataset.
2. Sample a random timestep $t$ uniformly from $\{1, ..., T\}$.
3. Sample a standard Gaussian noise vector $\epsilon \sim \mathcal{N}(0, I)$.
4. Calculate the noisy sample $x_t$ using the formula: $x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon$.
5. Provide $(x_t, t)$ as input to the network $\epsilon_\theta$ and train it to output a prediction $\epsilon_\theta(x_t, t)$ that is close to the original noise $\epsilon$ (usually using a mean squared error loss).

This ability to directly jump to any timestep $t$ allows for parallelized and efficient generation of training batches, which is much faster than simulating the Markov chain step-by-step for each training sample. This sampling strategy forms the basis of the training loop we will discuss in Chapter 4.