Introduction

A central challenge in machine learning, particularly in generative modeling, is to model complex datasets using highly flexible families of probability distributions that maintain analytical or computational tractability for learning, sampling, inference, and evaluation.

This post summarizes the fundamental concepts of diffusion models, their optimization strategies, and applications, focusing on the mathematical foundations and practical implications.

Denoising Diffusion Probabilistic Models (DDPM)

Diffusion models are a class of generative models first proposed by Sohl-Dickstein et al. (Sohl-Dickstein et al., 2015). Inspired by nonequilibrium thermodynamics, the method systematically and gradually destroys data structure through a forward diffusion process, then learns a reverse process to restore structure and yield a highly flexible and tractable generative model.

Forward Process

A diffusion model formulates the learned data distribution as $p_{\theta }(x_0):=\int p_{\theta }(x_{0:T})d{x}_{1:T}$, where $x_1,x_2,\dots ,x_T$ are latent variables with the same dimensionality as the real data $x_0\sim q(x_0)$.

The forward process is a Markov chain where transitions from $x_t$ to $x_{t+1}$ follow multivariate Gaussian distributions. The joint distribution of latent variables ($x_{1:T}$) given the real data $x_0$ is:

Key properties:

1. The coefficients $\{{\beta }_t\}$ are pre-defined and determine the "velocity" of diffusion.
2. After sufficient diffusion steps, the final state $x_T$ approaches an isotropic Gaussian distribution when $T$ is large enough.
3. For clarity, we use $p_{\theta }$ to represent the learned distribution and $q$ to represent the real data distribution.

A notable property of this forward process, as mentioned in Ho et al. (Ho et al., 2020) (Section 2), is that $x_t$ has a closed-form expression derived using the reparameterization trick (opens in a new tab). Let ${\alpha }_t=1-{\beta }_t$ and ${\bar{\alpha }}_t={\prod }_{i=1}^t{\alpha }_i$:

${}^{\ast }$ The sum of two uncorrelated multivariate normal distributions $\mathcal{N}({\mathbf{\mu }}_1,{\sigma }_1^2\mathbf{I})$ and $\mathcal{N}({\mathbf{\mu }}_2,{\sigma }_2^2\mathbf{I})$ is also a multivariate normal distribution $\mathcal{N}({\mathbf{\mu }}_1+{\mathbf{\mu }}_2,({\sigma }_1^2+{\sigma }_2^2)\mathbf{I})$ (proof details (opens in a new tab)).

Typically, we can afford larger update steps when the sample becomes noisier, so ${\beta }_1<{\beta }_2<\cdots <{\beta }_T$ and therefore ${\bar{\alpha }}_1>{\bar{\alpha }}_2>\cdots >{\bar{\alpha }}_T$.

Reverse Process

Ideally, if we knew $q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)$, we could gradually remove noise from corrupted samples to recover the original image. However, this conditional distribution is not readily available and its computation requires the entire dataset. Specifically:

Computing $q(x_{t-1}\mid x_t)$ requires evaluating integrals in Eq. (\ref{eq:imprac_cond_prob_expr}), which is computationally expensive. Instead, we use the diffusion model $p_{\theta }(x_{t-1}\mid x_t)$ to learn and approximate the true conditional distribution. When ${\beta }_t$ is sufficiently small, $q(x_{t-1}\mid x_t)$ is also Gaussian (details (opens in a new tab)).

The joint distribution of the diffusion model is:

With this distribution, we can sample from an isotropic Gaussian distribution and expect the reverse process to gradually transform it into samples that follow $p_{\theta }(x_0)\approx q(x_0)$.

Loss Function

The loss function for training the diffusion model is the standard variational bound on negative log likelihood:

Therefore, the expected negative log likelihood is lower bounded by the variational lower bound:

To convert each term in the equation to be analytically computable, the objective can be further rewritten to be a combination of several KL-divergence and entropy terms (See the detailed step-by-step process in Appendix B in Sohl-Dickstein et al. (Sohl-Dickstein et al., 2015)):

In summary, we can split the variational lower bound into $(T+1)$ components and label them as follows:

In the above decomposition:

1. $q(x_T\mid x_0)$ can be computed from Eq. (\ref{eq:xt_x0_relation})
2. $p_{\theta }(x_t\mid x_{t+1})$ are parameterized and learned

Next, we show that $q\left({\mathbf{x}}_t\mid {\mathbf{x}}_{t+1},{\mathbf{x}}_0\right)$ can be computed in closed form even though $q\left({\mathbf{x}}_t\mid {\mathbf{x}}_{t+1}\right)$ can't.

From the above derivation, we observe that the conditional distribution $q\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0\right)$ is also Gaussian and can be written in standard multivariate normal form as $q\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0\right)=\mathcal{N}\left(x_{t-1};\tilde{\mu }(x_t,x_0),{\tilde{\beta }}_t\mathbf{I}\right)$. Notice that $a{x}^2+bx+c=\frac{1}{1/a}\left(x^2+\frac{b}{a}x+\frac{c}{a}\right)=\frac{1}{1/a}\left((x+\frac{b}{2a}{)}^2+\text{ const }\right)$, the $\tilde{\mu }$ and ${\tilde{\beta }}_t$ can be computed as follows,

Recall the relation between $x_t$ and $x_0$ deduced from Eq. \ref{eq:xt_x0_relation}, the Eq. \ref{eq:standard_form_mu_t} can be further rewrited as follows:

Parameterization of reverse diffusion process and $L_t$

Recall our previous decomposition of variational lower bound loss, we have closed form computation of real data distribution (i.e. $q(x_T\mid x_0)$ and $q(x_t\mid x_{t+1},x_0)$), we still need a parameterization of $p_{\theta }(x_t\mid x_{t+1})$. As we discussed previously, when ${\beta }_t$ is small enough, we can approximate $q(x_t\mid x_{t+1})$ by the Gaussian distribution $p_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)=\mathcal{N}\left({\mathbf{x}}_{t-1};{\mathbf{\mu }}_{\theta }\left({\mathbf{x}}_t,t\right),{\mathbf{\Sigma }}_{\theta }\left({\mathbf{x}}_t,t\right)\right)$. In practice, we can further simplify the loss function Eq. (\ref{eq:lt_before_simple_sigma}) by predefine the variance matrix as ${\mathbf{\Sigma }}_{\theta }\left({\mathbf{x}}_t,t\right)={\sigma }_t^2\mathbf{I}$ and, experimentally, ${\sigma }_t^2={\tilde{\beta }}_t=\frac{1-{\bar{\alpha }}_{t-1}}{1-\overline{{\alpha }_t}}{\beta }_t$ or ${\sigma }_t^2={\beta }_t$ had similar results. We expect that the training process can let ${\mu }_{\theta }$ to predict ${\tilde{\mu }}_t$. With this parameterization, each component of loss function $L_t=D_{\mathrm{KL}}\left(q\left({\mathbf{x}}_t\mid {\mathbf{x}}_{t+1},{\mathbf{x}}_0\right)\Vert p_{\theta }\left({\mathbf{x}}_t\mid {\mathbf{x}}_{t+1}\right)\right)$ is the KL divergence between two multivariate Gaussian distributions and has a relatively simple closed form (opens in a new tab). The loss term $L_t$ become