Paper review: Denoising Diffusion Probabilistic Models

A summary of Ho et al. (2020) on high-quality image synthesis using diffusion models.

  1. Denoising diffusion probabilistic models
    Jonathan Ho, Ajay Jain, and Pieter Abbeel
    In Proceedings of the 34th International Conference on Neural Information Processing Systems, Jun 2020

The task

Neural network

Statistical Modeling: Reverse Process

$$p_{\theta}(x_0) := \int p_{\theta}(x_{0:T}) dx_{1:T}$$

Modeling 2: Forward process assumption

$$q(x_{1:T}|x_0) := \prod_{t=1}^T q(x_t|x_{t-1}), \quad q(x_t|x_{t-1}) := \mathcal{N}(x_t; \sqrt{1 - \beta_t}x_{t-1}, \beta_t\mathbf{I})$$

where $\beta_t$ is a variance schedule and we assume it is given. We will not learn it.

Training: KL Divergence

\begin{equation}L = \mathbb{E}_q \left[ D_{KL}(q(x_T |x_0) \Vert p(x_T )) + \sum_{t>1} D_{KL}(q(x_{t-1}|x_t, x_0) \Vert p_{\theta}(x_{t-1}|x_t)) - \log p_{\theta}(x_0|x_1) \right]\end{equation}

Gaussian Assumption: KL Divergence to Simple ERM

To minimize that summation of KL divergences, notice that both distributions inside the KL divergence are Gaussian:

\begin{equation}p_{\theta}(x_{t-1} \vert x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \Sigma_\theta(x_t, t))\end{equation}

\begin{equation}q(x_{t-1} \vert x_t, x_0) = \mathcal{N}(x_{t-1}; \tilde{\mu}_t(x_t, x_0), \tilde{\beta}_t I)\end{equation}

Because of the Markov property and properties of Gaussians, the forward process conditioned on $x_0$ has a beautifully tractable closed-form mean:

\begin{equation}\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\end{equation}

(where $\alpha_t = 1 - \beta_t$ and $\bar{\alpha}t = \prod{s=1}^t \alpha_s$)

Since the KL divergence between two Gaussians with fixed variances is simply proportional to the $L_2$ distance between their means, minimizing the KL divergence is mathematically equivalent to solving the following least squares problem:

\begin{equation}\arg\min_\theta \Vert \tilde{\mu}_t(x_t, x_0) - \mu_\theta(x_t, t) \Vert^2\end{equation}

Instead of predicting the mean directly, we can use the reparameterization trick. We know that $x_t$ is just a deterministic combination of $x_0$ and some pure noise $\epsilon \sim \mathcal{N}(0, I)$:

\begin{equation}x_t = \sqrt{\bar{\alpha}_t}x_0 + \sqrt{1 - \bar{\alpha}_t}\epsilon\end{equation}

If we substitute $x_0$ out of the $\tilde{\mu}t$ equation using the formula above, and we parameterize our neural network to predict the noise $\epsilon\theta(x_t, t)$ rather than predicting the mean $\mu_\theta$, the complex statistical KL divergence collapses into a remarkably simple empirical risk minimization problem:

\begin{equation}L_{simple}(\theta) = \mathbb{E}_{t, x_0, \epsilon} \left[ \Vert \epsilon - \epsilon_\theta(x_t, t) \Vert^2 \right]\end{equation}

Training in Practice

What the Pytorch Looks At

Unlike a standard image classifier that looks at a clean image, a diffusion model’s core network takes in two distinct pieces of information during a forward pass:

  1. The Noisy Data ($x_t$): This is the corrupted version of the original data.
  2. The Timestep ($t$): A scalar value telling the network how corrupted the data is (e.g., step 450 out of 1000). The network absolutely needs this context. Pulling noise out of a slightly blurry image requires a completely different mathematical transformation than pulling noise out of pure static.

The Target: What the Network Tries to Predict

The math in this paper shows that the target label is:

Building the Dataset on the Fly

One of the most elegant aspects of training a diffusion model is that you don’t pre-compute and save terabytes of noisy intermediate images. Instead, the training loop generates the neural network’s training pairs instantly in memory during every batch.

Here is the step-by-step anatomy of how this data is generated in a training iteration:

  1. Sample the Ground Truth: Draw a clean, real data point ($x_0$) from your original dataset.
  2. Select a Timestep: Randomly pick a timestep $t$ from a uniform distribution (e.g., $t \in [1, T]$).
  3. Generate the Target: Sample pure random noise ($\epsilon$) from a standard normal distribution. This will be the ground truth for our loss function.
  4. Create the Input: Mathematically mix $x_0$ and $\epsilon$ together using the closed-form forward process formula. Because of the properties of Gaussians, we don’t have to simulate every step; we can jump directly to step $t$:
  5. Feed the Network: Pass the newly minted noisy data ($x_t$) and the timestep ($t$) into the neural network.
  6. Calculate the Loss: The network outputs its prediction of the noise ($\epsilon_\theta$). The loss is simply the Mean Squared Error between the network’s prediction and the actual noise ($\epsilon$) drawn in step 3.

The Infinite Stream

By shifting our perspective to the neural network’s specific inputs and outputs, it becomes clear that the original dataset is merely a seed. The actual dataset being pushed through the optimizer is an infinite, constantly generated stream of (noisy data, timestep) -> actual noise mappings.

Because the noise and timesteps are sampled randomly, every time the model encounters the same underlying image $x_0$, it sees it corrupted by a completely different noise pattern at a different severity level. This dynamic generation is what prevents severe overfitting and allows the model to robustly learn the full reverse trajectory of generation.

My understanding

The names are forward and reverse. But the paper introduces reverse first, because reverse is the modeling of the density. forward is the modeling of training data generation.

References