import Comments from '../../components/Comments'

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)$. 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

In practice, we can further simplify the loss function Eq. (\ref{eq:lt_before_simple_sigma}) by predefine the variancen 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. Therefore, we can write:

Furthermore, ${\mu }_{\theta }$ is further parameterized in section 3.2 of Ho et al. (Ho et al., 2020) to be corresponded with the form of ${\tilde{\mu }}_t$ in Eq. (\ref{eq:standard_form_mu_t}) as follows,

In this parameterization, the neural network will be used to approximate the ${ϵ}_{\theta }$ instead of ${\mu }_{\theta }$ directly. This parameterization further simplify $L_{t-1}$ in Eq. (\ref{eq:lt_mu_t_not_param}) into

To this point, every term in $L_{\text{VLB}}$ can be computed in explicit closed form and is ready for training. However, empirically, Ho et al. (Ho et al., 2020) found that training the diffusion model works better with a simplified objective that ignores the weighting term:

This simplification enable us to compute arbitrary time steps for each sample $x_0$, instead of computing the entire series as in $L_{\text{VLB}}$. The entire DDPM algorithm is show as follows.

Side notes: The timestep is also an input to the neural network ${ϵ}_{\theta }(x_t,t)$ and it is typically encoded into some vector. For instance, in DDPM, integer timesteps are encoded into floats vector through sinusoidal function (opens in a new tab).

Denoising Diffusion Implicit Models (DDIM)

Though diffusion models like DDPM already demonstrated the ability to produce high quality samples that are comparable with the state-of-the-art generative model, such as GANs, the computation complexity of the sampling process is a critical drawback. These diffusion-based models typically require many iterations to produce a high quality sample, whereas models like GANs only need one iteration. A quantitative experiment in Song et al. (Song et al., 2020) shows that, with same GPU setup and similar neural network complexity, it takes around 20 hours to sample 50k images of size $32\times 32$ from a DDPM, but less than a minute to do so from a GAN model. To resolve this high computational cost without lossing too much generation quality, Song et al. (Song et al., 2020) proposed Denoising Diffusion Implicit Models(DDIM). This algorithm is based on two observation/intuitives.

1. The deduction of the loss function only depends on $q(x_t\mid x_0)$ and the sampling process only depends on $p_{\theta }(x_{t-1}\mid x_t)$. To be more specific, the loss function remains the same form as long as the relation in Eq. (\ref{eq:xt_x0_relation}) still hold.

2. A DDPM trained on $\{{\alpha }_t{\}}_{t=1}^N$ has, in fact, included the "knowledge" for training a DDPM with $\{{\alpha }_{\tau }\}\subset \{{\alpha }_t{\}}_{t=1}^N$. This can be naturally observed from training process of the simplified version of loss function. It gives us a intuition that we can use a subset of parameters during the sampling process and reduce the computational cost.

Based on the first observation, we can build different conditional distributions $q(x_t\mid x_{t+1},x_0)$ that has the same $q(x_t\mid x_0)$ distribution. Same marginal distribution $q(x_t\mid x_0)$ results in the same loss function and different choices of conditional distribution $q(x_t\mid x_{t+1},x_0)$ results in different sampling choices. In fact, without the constraint of $q(x_{t+1}\mid x_t)$ as in Eq. (\ref{eq:diff_model_origin}), we have a broader choice(i.e. a larger solution space) of $q(x_t\mid x_{t+1},x_0)$.

Based on the second observation, the sampling process can only use a subset of steps used in training process. By reducing the updating steps, the sampling process can greatly speed-up.

Non-Markovian Forward Processes

The key observation here is that the DDPM loss function only deppends on the marginals $q(x_t\mid x_0)$, but not directly on the joint distribution $q(x_{1:T}\mid x_0)$ or transition distribution $q(x_t\mid x_{t-1})$. Follow the deduction in Spaces.Ac.cn (opens in a new tab), we can use undetermined coefficient method to compute the form of $q(x_t\mid x_{t+1},x_0)$ and we also assume it takes a normal distribution. We first summarize the condition as follows:

1. To maintain the same loss function, we need the same marginal distribution $q(x_t\mid x_0)=\mathcal{N}(\sqrt{{\bar{\alpha }}_t}{x}_0,(1-{\bar{\alpha }}_t)\mathbf{I})$. The corresponded sampling process formula is $x_t=\sqrt{{\bar{\alpha }}_2}{x}_0+\sqrt{1-{\bar{\alpha }}_t}{ϵ}_1$.

2. Assume $q(x_{t-1}\mid x_t,x_0)=\mathcal{N}(x_{t-1};k_t{x}_t+{\lambda }_t{x}_0,{\sigma }^2\mathbf{I})$, where $k_t$ and ${\lambda }_t$ are coefficient to be decided. The sampling process with $q(x_{t-1}\mid x_t,x_0)$ is $x_{t-1}=k_t{x}_t+{\lambda }_t{x}_0+{\sigma }_t{ϵ}_2$.

By combining the marginal distribution $q(x_t\mid x_0)$ and assumed form of $q(x_{t-1}\mid x_t,x_0)$, we can compute the marginal distribution of $x_{t-1}$ as follows,

Comparing Eq. (\ref{eq:sample_with_undeter_coefs}) with Eq. (\ref{eq:xt_x0_relation}), remember that we need to let the marginal distribution to be the same and $q(x_{t-1}\mid x_0)={\int }_{x_t}q(x_{t-1}\mid x_t,x_0)q(x_t\mid x_0)\mathrm{d}{x}_t$, we can have the following relation,

There are three variables and only two equation, therefore, we can view ${\sigma }_t^2$ as a independent variable and solve that

Therefore, we can obtain a family $\mathcal{Q}$ of inference distribution indexed by $\{{\sigma }_t{\}}_{1:T}$,

As a result, in the sampling procedure, the updating formula is,

Comparing with DDPM, this is generalized form of generative processes. Since the marginal distribution remains the same, the loss function did not change and the training process is identical. This means that we can use this new generative process with a diffusion model trained in DDPM way and, with different level of ${\sigma }_t$, we can generate different image with same initial noise. Among different choices, ${\sigma }_t=0$ is a special case in which the generation process is deterministic given the initial noise. This model is called denoising diffusion implicit model since it is an implicit probabilistic model and it is trained with the DDPM objective.

Accelerate generation processes

We need to point out that, in our previous discussion, we did not start with $q(x_t\mid x_t-1)$ and the sequence $\{{\alpha }_t\}$ determined the model. The key observation here is that the training process of DDPM, in its essence, contained the data/processes of training over any subsequence $\{{\alpha }_{\tau }\}\subset \{{\alpha }_t\}$. This can be observed from the loss functions. The training process over a set of parameter $\{{\alpha }_{\tau }\}$ is

Therefore, DDPM trained on $\{{\alpha }_t\}$ already incorporated information used to train DDPM on $\{{\alpha }_{\tau }\}\subset \{{\alpha }_t\}$. When the size of $\{{\alpha }_{\tau }\}$ is much smaller than $\{{\alpha }_t\}$, generating samples with the former parameters set will be much faster.

Remarks on DDIM

1. Why don't we just directly train on $\{{\alpha }_{\tau }\}$ and sample from the model?\ There might be two considerations for training on T steps but sampling in $\dim (\tau )$ steps. Firstly, diffusion model trained on more sophisticated setup might improve the model's capability of generalization. Secondly, use subsequence to speed up is one way of acceleration and there might be other acceleration method with this more sophisticate model.
2. Can we use DDPM and sample with subset of parameters $\{{\alpha }_{\tau }\}$? What is the purpose of choosing this new family of conditional distribution?\ For purpose of accelerating sample generation process, one can certainly use DDPM and skip some steps during generation. However, clearly, the newly proposed distribution family is more flexible and has the potential of generate more diversified samples without any additional cost other than DDPM. As a matter of fact, letting ${\sigma }_t^2=\sqrt{\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}}{\beta }_t$, the DDIM is equivalent to DDPM's sampling process.
3. Additional benefit comes with DDIM?\ DDIM has "consistency" property since the generative process is deterministic. It means that multiple samples conditioned on the same latent variable should have similar hig-level features. Due to this consistency, DDIM can do semantically meaningful interpolation in the latent variable.

Result comparison between DDPM and DDIM

Experiments in DDPM and DDIM paper have quantitatively and qualitatively examined the images generated. Here we review two aspects: the sampling quality and the interpolation result.

Sampling quality

In DDIM's experiment, as the following screenshot taken from it shows, the authors compared results of different number of diffusion steps $\text{dim}(\tau )$ and different level of noise. The empirical result is that lower noise level ${\sigma }_t^2$, the better image quality generated with accelerated diffusion process.

Interpolation

Both DDIM and DDPM examined their performance on interpolation of images. They use the forward process as stochastic encoder to generate embeddings $x_0\to x_T,x_0^{\prime }\to x_T^{\prime }$. Then decoding the interploated latent ${\bar{x}}_T=f(x_T,x_T^{\prime },\alpha )$ where $\alpha$ represent interpolation parameter(s).

1. In DDPM, the authors simply use a linear interpolation, i.e. ${\bar{x}}_T=\alpha x_T+(1-\alpha )x_T^{\prime }$.
2. n DDIM, the authors use a spherical linear interpolation,\

where $\theta =\arccos \left(\frac{({\mathbf{x}}_T{)}^T{\mathbf{x}}_T^{\prime }}{\Vert {\mathbf{x}}_T\Vert \Vert {\mathbf{x}}_T^{\prime }\Vert }\right)$.

Interesting Reading

• Lilian Weng's post on diffusion model (opens in a new tab)
• Spaces.Ac.cn's post on DDIM (opens in a new tab)

References