An Overview of Denoising Diffusion Probabilistic Models

Score matching network

Diffusion models are said to have emerged from score-matching networks. So what exactly is a score-matching network?

The figure above gives a simple introduction to score-matching networks. Let us define the gradient of a PDF (probability density function) as its score. This definition presupposes ideas from energy-based models and normalized probabilities, which I will explain in more detail below. For now, recall that a gradient is defined at a point on a function and indicates the direction in which the function value increases most rapidly.

That is why gradient-based learning, commonly used in deep learning, computes the derivative (gradient) at a particular point of the loss function and optimizes the loss by moving in the opposite direction—the direction of steepest decrease. This method is called gradient descent.

Viewed instead from the perspective of sampling from a PDF, the predicted score at a point—the gradient of the PDF—points toward a “more plausible sample.” The central idea of score matching is that repeatedly generating a sample in this direction will eventually produce a plausible one.

The example above uses MNIST handwritten digits. You can think of the process as starting from noise and gradually following the gradient uphill until a clean sample appears.

I will explain where this idea of score matching came from and, ultimately, how its history led to DDPM.

Score matching to NCSN

In energy-based machine learning, energy is defined as an unnormalized negative log-likelihood. Conversely, a normalized Gibbs–Boltzmann distribution can be treated as a probability distribution.

In other words, because the $p_\theta(x)$ that we truly want is difficult to determine, we approach it as a normalized probability over an accessible function $f$: $p_\theta(x) = \frac{e^{-f_\theta(x)}}{z_\theta}$. Here, $z_\theta$ is the integral of $e^{-f_\theta(x)}$ over every $x$.

$$z_\theta = \int_{x\sim X} e^{-f_\theta(x)}dx$$

We adopted this formulation because the probability is hard to obtain directly, yet computing $z_\theta$ remains difficult.

The same kind of problem appears in neural networks. We want to optimize a network over a vast number of parameter weights, but cannot express the “function” represented by that network analytically. We therefore optimize it with gradients, defining the network as a differentiable function and training it on enough data.

What ultimately matters is not $z_\theta$, integrated over all $x$, but the direction in which $p_\theta(x)$ should be optimized. If we take the negative logarithm of the expression above and compute its gradient with respect to $x$,

$$\nabla_x \log p(x) = -\nabla_x f_\theta(x)$$

the intractable term disappears from the gradient of the probability's log-likelihood. More precisely, the score in score matching is the derivative of a function defined through a probability distribution's log-likelihood, and estimating it becomes the objective of generation. A method for score estimation was introduced in 2005; I have included the paper below for anyone interested.

Paper

In brief, the paper proves the analytically calculated result. After the derivation, the terms involving the parameter $\theta$ become

$$\int_x p_{data}(x)\left(\frac{1}{2} (\nabla_x \log p_{\theta} (x))^\top (\nabla_x \log p_{\theta} (x)) - Tr(\nabla_x^2 \log p_{\theta} (x))\right)$$

as shown above. The problem is that the cost of the gradient depends on the dimensionality of $x$, making computation prohibitively slow for high-dimensional modalities such as images.

Consequently, score matching attracted less attention than methods that were easier to sample from, such as VAEs, flow-based models, and GANs. It nearly faded into obscurity in deep learning, where training speed is critical.

NCSN (Noise-Conditioned Score Network) then appeared to address this problem by adopting a denoising-autoencoder approach. It is also described as a score-matching network based on annealed Langevin dynamics.

That paper is also packed with equations, so I will give only a brief overview. Because computing the score's Jacobian and then its trace was too expensive for large-scale data, the authors reformulated the original expression.

They add a noise distribution—specifically Gaussian noise—to the original data and recover the original distribution through denoising. For this approximation to hold, of course, the amount of noise added at each step must be very small.

$$q_\sigma(\tilde{x}) = \int_x q_\sigma(\tilde{x} \vert x) p_{data}(x) dx$$

Assuming that a very small noise distribution has been added,

$$\nabla_x \log q_\sigma(x) \approx \nabla_x \log p_{data}(x)$$

Sliced score matching was also proposed, but I will not cover it separately because its projection costs four times as much as the denoising method.

In short, NCSN with a denoising autoencoder approximates the original score-matching objective by adding noise from a predefined Gaussian distribution to the data and training the model to remove it.

Sampling uses a Markov process. The score predicted from the noisy data $\tilde{x}_{t-1}$ gives the direction from that point to the next sample—as explained earlier, toward a more plausible sample—and Langevin dynamics adds the diffusion term $z_t$.

$$x_t = x_{t-1}+\frac{\epsilon}{2}\nabla_x \log p(x_{t-1}) + \sqrt{\epsilon}z_t$$

Langevin dynamics has the same form as an SDE solution for Brownian motion. The equation above omits a lengthy derivation; in short, it is one of the best-known discrete approximations of an SDE solution. We can regard $z_t$ as a normal distribution with mean 0 and variance 1, while $\sqrt{\epsilon}$ serves as the standard deviation in a reparameterization. This dynamics, however, suffers from

1. Manifold hypothesis
2. Inaccurate score matching in low density region
3. Slow mixing of Langevin dynamics

so in practice, annealed Langevin sampling is used instead.

The manifold problem means that, as illustrated above, data in a three-dimensional space may actually lie on a two-dimensional manifold (a particular lower-dimensional topology). Score matching is defined in the higher-dimensional ambient space, which is assumed to make proper sampling difficult. The low-density-region problem is that scores are predicted well around the dominant parts of the training distribution, where samples are concentrated, but inaccurately elsewhere. Slow mixing of Langevin dynamics means that when distinct distributions are mixed with weights at different scales, sampling must be able to distinguish them; conventional Langevin dynamics cannot do so with its small step size.

The details are rather involved, so the key idea is simply that sampling proceeds while the noise is gradually reduced.

There is much more to say, but let us move on after noting that a score estimator can be trained through this kind of denoising.

What Does Diffusion Actually Learn?

A diffusion model has the following training objective. To keep the explanation simple, it helps to look at the diagram comparing types of generative models.

A GAN trains its generator by having a discriminator determine whether an image produced by the generator is fake or comes from the real data distribution. A VAE works differently: with an encoder that maps into embedding space $Z$, it sends data $x$ to a manifold space and uses that representation to train the decoder, which acts as the generator. Because the real data distribution $p(x)$ cannot be obtained, it instead learns conditionally on a latent-space distribution over $Z$ that is easy to define. It also maximizes probability indirectly through the ELBO (Evidence Lower Bound) seen in the VAE objective. A flow-based model defines a flow function that encodes a sample into a latent representation, based on the assumption that if this function can be defined, it can be inverted to sample from the latent space.

Diffusion is not fundamentally different from these broader ideas in generative modeling. The word diffusion evokes a drop of ink dispersing in water or perfume spreading through the air. This is related to Brownian motion (a Wiener process). Diffusion models call this the “forward process”; sampling from noise in the opposite direction with a learned parameterized network is the “reverse process.”

Forward process in DDPM

The DDPM paper defines the gradual addition of noise to a sample as the forward process. One could view noise addition as something learnable through a reparameterization trick, much like training an encoder. The paper mentions this only as a possibility, however; its experiments use a fixed noise schedule in the forward process.*

*The experimental details describe constant, linear, and quadratic $\beta_t$ schedules ($T = 1000$, with $\beta_1 = 10^{-4}$ through $\beta_T = 0.02$).

The forward process in DDPM adds noise to the next state ($X_t$) based on the immediately preceding state ($X_{t-1}$), following a Markov process. In a Markov process, the next state depends only on the current state. As the figure shows, a tiny amount of noise is added to the original image ($X_0$) to produce the next state, $X_1$; this is repeated $T$ times until it becomes pure noise ($X_T$). The forward-process term used in the paper can be written as $q(X_t \vert X_{t-1})$. Using a predetermined variance schedule $\beta$, the paper defines $q$ as the following sequence of Gaussian-noise additions.

$$q(x_t \vert x_{t-1}) := N(x_t; \sqrt{1-\beta_t}x_{t-1},~\beta_t\rm{I})$$

Reverse process in DDPM

As discussed later, the noisy samples generated at each $t$ in DDPM serve as conditions for learning the reverse process. Expressed like a flow-based model, the generative model we want would be $q(X_{t-1} \vert X_t)$. Because we cannot obtain it directly, we instead learn $p_{\theta}(X_{t-1} \vert X_t)$. In contrast to the forward process, the reverse process removes noise. Repeatedly adding or subtracting Gaussian noise that is very small relative to a given signal does not differ greatly in terms of the learned distribution, allowing the model to imitate the process of creating samples from noise. In the paper's notation,

$$p_{\theta}(x_{0:T}) := p(x_T)\prod_{t=1}^T p_{\theta}(x_{t-1} \vert x_t)$$

The reverse-process network $p_\theta$, parameterized by $\theta$, samples the next state conditioned on the latent at each time step $(x_{1,~2,~\cdots,~T})$. Sampling each state resembles Langevin dynamics in denoising score matching; for applications of Langevin dynamics, see the score-based-model literature. Put simply, the network looks at the noise and decides the next step, whose direction should make the result resemble the samples observed during diffusion-model training as closely as possible.

How to define Loss function?

From the discussion above, it is clear that our objective is to optimize $p_\theta$. Every deep-learning model has a loss (cost) function suited to its training setup and purpose; here it can be derived relatively easily from the ELBO used by VAEs.

$$\begin{aligned} \log p(x) =& E_{z\sim q(z \vert x)}(\log p(x)) \newline =& E_{z} \left(\log \frac{p(x \vert z)p(z)}{p(z \vert x)} \right) \newline =& E_{z} \left( \log \frac{p(x \vert z)p(z)}{p(z \vert x)} \frac{q(z \vert x)}{q(z \vert x)} \right) \newline =& E_z (\log p(x \vert z))-E_z \left( \log \frac{q(z \vert x)}{p(z)} \right) + E_z \left( \log \frac{q(z \vert x)}{p(z \vert x)} \right) \end{aligned}$$

Because $p(x)$ cannot be obtained directly, we assume it can be estimated through a known $z$. We use $p(x \vert z)$—the likelihood of $p(x)$ conditioned on the prior $z$—and divide by the posterior. The $x \to z$ path can be defined as encoding and the $z \to x$ path as decoding. In this formulation, $p$ as the encoder handles the forward process and $q$ as the decoder handles the reverse process, allowing us to construct a KL-divergence expression. If we know the distribution of $z$, the decoder's starting point for the otherwise hard-to-obtain reverse process, the encoder can be trained to follow that distribution. Unlike DDPM, a VAE defines both encoder and decoder with learnable parameters $\phi$ and $\theta$.

$$=E_z ( \log p_{\theta}(x \vert z)) - D_{KL} (q_{\phi}(z \vert x) \vert\vert p(z)) + D_{KL} (q_{\phi}(z \vert x) \vert\vert p(z \vert x))$$

The main objective of the variational Bayes method in a Variational Autoencoder is to obtain the intractable posterior $p_\theta(z \vert x)$ with the help of an encoder. From the generation perspective the objective is the maximum likelihood (MLE) $P_\theta(x \vert z)$, though one could imagine other algorithms for this training, such as the adversarial approach used by GANs, rather than VAE training specifically. The key point here is that a VAE learns the decoder's posterior with help from the encoder.

Transferring this perspective to a denoising diffusion probabilistic model suggests a similar optimization.

Define the latent variable in the conventional ELBO as the time-dependent variable $x_t$, and the image domain as $x_0$.

Here, time $t$ indicates “how much” noise has been added relative to $x_0$.

$$=E_z(\log p_{\theta}(x_0 \vert x_t)) - D_{KL} (q(x_t \vert x_0) \vert\vert p_{\theta}(x_t)) + D_{KL} (q(x_t \vert x_0) \vert\vert p_{\theta}(x_t \vert x_0))$$

As noted earlier, a likelihood model cannot obtain the posterior, so $p(x_t \vert x_0)$ in the expression cannot be handled. We therefore use $D_{KL} \ge 0$ for the final term and define a lower bound with the remaining terms; this is the variational bound. A score-matching-based DDPM can handle it through conditioning. Let us use a small algebraic trick to derive the actual expression.

We can change both the part that averages the leading $\log p_\theta(x_0 \vert x_t)$ over $z$ and the denominator of the following KL divergence. For now, ignore the final $D_{KL} \ge 0$ term.

$$\ge E_{x_T}(\log p_\theta(x_t)) - D_{KL} (q(x_t \vert x_0) \vert\vert p_{\theta}(x_0 \vert x_t))$$

This yields the expression above. Generalizing it to the final time $T$ gives

$$\begin{aligned} \mathcal{L} :=& E\_{x\_T}(-\log p\_\theta(x\_0)) \newline \mathcal{L} \le& \mathbb{E}\_{x\_T}\left(-\log \left(\frac{p(x\_0 \vert x\_T)}{q(x\_T \vert x\_0)} \right) \right) \newline \mathcal{L} \le& \mathbb{E}{x\_T}\left(-\log \left(\frac{p\_\theta(x\_T \prod_{t=1}^T p\_\theta(x\_{t-1}\vert x\_t))}{\prod\_{t=1}^T q(x\_t \vert x\_{t-1})} \right) \right) \end{aligned}$$

which admits a graphical-model formulation. Then, after expanding and rearranging it,

$$\mathcal{L} \le \mathbb{E}{x_T}\left(-\log (p_\theta(x_T))-\sum_{t=2}^T \log \frac{p_\theta(x_{t-1} \vert x_t)}{q(x_t \vert x_{t-1})} -\log \frac{p_\theta(x_0 \vert x_1)}{q(x_1 \vert x_0)} \right)$$

we obtain the expression above.

The crucial part of this rearrangement begins at $t=2$, because from that point onward $x_t$ can be conditioned independently of $x_0$. The predefined posterior in the expression above can therefore be replaced by the likelihood.

$$\mathcal{L} \le \mathbb{E}{x_T}\left(-\log (p_\theta(x_T))-\sum_{t=2}^T \log \left( \frac{p_\theta(x_{t-1} \vert x_t)}{q(x_{t-1} \vert x_t, x_0)} \times \frac{q(x_{t-1} \vert x_0)}{q(x_t \vert x_0)} \right) -\log \frac{p_\theta(x_0 \vert x_1)}{q(x_1 \vert x_0)} \right)$$

Rearranging this expression one final time gives the following result and completes the derivation of the loss.

$$\mathcal{L} \le \mathbb{E}{x_T}\left(-\log \left(\frac{p_\theta(x_T)}{q(x_T \vert x_0)} \right)-\sum_{t=2}^T \log \left( \frac{p_\theta(x_{t-1} \vert x_t)}{q(x_{t-1} \vert x_t, x_0)} \right) -\log p_\theta(x_0 \vert x_1) \right)$$

The leading negative log-likelihood makes the sequence that creates noise from $x_0$ follow a Gaussian distribution, so that the reverse process can begin sampling from a Gaussian distribution. It can be understood as the VAE's regularization term (forward process $\to$ Gaussian).

The second expression, containing $\sum$, corresponds to the green middle section: it aligns the prediction at each reverse-process step with the predefined forward process. The final term reconstructs $x_0$, the real image, from $x_1$.

$$q(x\_{t-1} \vert x_t,~x_0) = q(x_t \vert x_{t-1}) \times \frac{q(x_{t-1} \vert x_0)}{q(x_t \vert x_0)} = N(x_{t-1}; \tilde{\mu}\_t(x_t, x_0), \tilde{\beta}\_t I)$$ $$\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$$

The tractable terms are summarized above. As the paper's derivation shows, the formula can be obtained simply by conditioning Gaussian distributions.

$$L_{t-1}:= \mathbb{E}\_{x_T} \left( \log \left( \frac{q(x_{t-1} \vert x_t, x_0)}{p\_{\theta}(x_{t-1} \vert x_t)} \right) \right) = \mathbb{E}\_{x_T \sim q(x_t \vert x_0)} \left( \frac{1}{2\sigma^2_t} \parallel \tilde{\mu}\_t (x_t, x_0) - \mu\_\theta (x_t, t) \parallel^2 \right) + C$$

The objective is therefore to optimize the expression above, predicting at each step $t$ a Gaussian distribution in the direction of the next step. The original loss consists of KL divergences between Gaussian distributions. By minimizing the noise variance and increasing the number of time steps $T$, the reverse process can follow the Gaussian distribution effectively. The model must predict the mean for its input, which can also be simplified as follows.

$$L_{t-1} -C := \mathbb{E}\_{x_0, \epsilon} \left( \frac{\beta_t^2}{2\sigma_t^2 \alpha_t (1- \bar{\alpha}\_t)} \parallel \epsilon - \epsilon\_\theta (\bar{\alpha}\_t x_0 + \sqrt{1 - \bar{\alpha}\_t}\epsilon, t) \parallel^2 \right)$$ $$\mathcal{L}\_{simple} := \mathbb{E}\_{x_0, \epsilon} \left( \parallel \epsilon - \epsilon\_\theta(\sqrt{\bar{\alpha}\_t}x_0 + \sqrt{1- \bar{\alpha}\_t}\epsilon, t) \parallel^2 \right)$$

As shown above, the objective can be written in terms of $\epsilon \sim N(0, I)$. The authors report that either the full expression or the simple objective (which removes the normalization term $\gamma$) can be optimized. This differs from NCSN, which normalizes every step to obtain a weighted loss; DDPM demonstrates empirically that this is unnecessary.

This is the change from the baseline, and it reportedly produces better performance.

Limitations

As an early paper, DDPM has many limitations; one is pixel-level interpolation. In generative models such as GANs, interpolating in latent space generally preserves the consistency of generated images. In DDPM, however, interpolation at a low level can reportedly produce an entirely different image. My interpretation is that because DDPM samples by removing noise little by little, its implicit model can optimize only a single step ($t \to t-1$). A low-level latent therefore has no awareness of the ultimately generated image $x_0$. DDIM later claimed to solve this problem through implicit, non-Markovian sampling.

Other limitations are sampling and training speed. Training is time-consuming in any deep-learning task, but slow sampling itself is a major drawback for a generative model. DDIM later addressed this issue as well, and most diffusion samplers now use DDIM. I plan to review DDIM in a later post.