Understanding Diffusion Models

Behind Diffusion Models

Diffusion models were first introduced in the seminal work by Sohl-Dickstein et al. (2015). The core idea involves reversing a Markov chain-based forward diffusion process, which gradually degrades the structure of the data $\mathbf{z}_0$ from the real data distribution $q(\mathbf{z}_0)$, by adding noise over a sufficient number of steps. When this noise is Gaussian, as is commonly assumed in practice, the cumulative effect transforms the data distribution towards a standard normal distribution $\mathcal{N}(0, I)$ as the forward process progresses. We can then sample from this distribution and use a learned reverse process to generate new samples that match the real data distribution.

Specifically, as illustrated in Figure 1, we define the forward diffusion process $q$ that gradually applies Gaussian noise to the data distribution $\mathbf{z}_0 \sim q(\mathbf{z}_0)$ with a variance schedule $\beta_t \in (0, 1)$ as follows:

Figure 1: Illustration of the forward and reverse diffusion process. Figure adapted from Ho et al. (2020).

Here, $t$ and $T$ refer to the current and total timesteps, respectively. With a sufficiently large $T$ and a properly defined variance schedule $\beta$, the final $\mathbf{z}_T$ will approximate an isotropic Gaussian distribution, where variances are equal in all directions. We can then reverse the forward diffusion process by starting from $\mathbf{z}_T \sim \mathcal{N}(0, I)$ and gradually obtain a sample $\mathbf{z}_0$ that resembles the real data distribution $q(\mathbf{z}_0)$, defined as follows:

Where all $\mathbf{z}t$ for $t \in {0, 1, \ldots, T }$ share the same dimensionality, $p{\theta}(\mathbf{z}{t-1} \vert \mathbf{z}{t})$ is practically a neural network learned to approximate the actual reverse distribution $q(\mathbf{z}{t-1} \vert \mathbf{z}{t})$, as it depends on the entire dataset which is not easily accessible. The term $p_{\theta}(\mathbf{z}_{0:T})$ denotes the entire learned reverse process.

An interesting property of the forward process is the existence of a closed-form sampling solution for $\mathbf{z}t$ at an arbitrary $t \in [0, T]$ using the reparameterization trick. By defining $\alpha_t := 1 - \beta_t$ and $\bar{\alpha}_t := \prod{i=1}^{t} \alpha_i$, given Equation 2, and recall that $\mathcal{N}(z; \mu, \sigma^2\mathbf{I})$ is equivalent to $z = \mu + \sigma \epsilon$ where $\epsilon \sim \mathcal{N}(0, 1)$, we obtain:

Which leads to:

Recall that combining two Gaussian distributions with different variances, $\mathcal{N}(z; \mu, \sigma_1^2\mathbf{I})$ and $\mathcal{N}(z; \mu, \sigma_2^2\mathbf{I})$, results in $\mathcal{N}(z; \mu, (\sigma_1^2+\sigma_2^2)\mathbf{I})$. Thus, in the case of Equation 7, for transitioning from $t$ to $t-1$, the variance becomes $\sqrt{(1 - \alpha_t) + \alpha_t (1-\alpha_{t-1})} = \sqrt{1 - \alpha_t\alpha_{t-1}}$, with subsequent derivations following the same intuition.

The goal is to learn a network $p_\theta$ that can approximate the actual reverse process. We start by looking at the objective that we want to minimize: the negative log-likelihood, $- \log p_\theta(\mathbf{z}_0)$, which involves maximizing the likelihood of the real data. However, directly optimizing this objective is practically infeasible due to its dependence on the entire sequence of previous steps. As a solution (Sohl-Dickstein et al. (2015)), we can derive an variational lower bound (vlb) on the data log-likelihood to achieve the same optimization effect:

Note that the numerator and denominator of $L_{vlb}$ refer to the forward process and the learned reverse process, respectively. Therefore, we can derive the following expressions (Sohl-Dickstein et al. (2015), Ho et al. (2020)):

Here, given the intrinsically high variance in the intermediate samples that introduces intractability in modeling the reverse process, we can condition on $\mathbf{z}_0$ due to its accessibility and feasibility to reduce the variance and make it tractable:

Here, since the term $\sum_{t=2}^T \log \frac{q(\mathbf{z}t \vert \mathbf{z}_0)}{q(\mathbf{z}{t-1} \vert \mathbf{z}_0)}$ can be unrolled as:

Where the intermediate terms cancel out, resulting in $\log \frac{q(\mathbf{z}T \vert \mathbf{z}_0)}{q(\mathbf{z}{1} \vert \mathbf{z}_0)}$, we can then continue the derivation as follows:

Recall that in Equation 2, the forward process requires a variance term $\beta_t$, which can either be learned through reparameterization (Kingma and Welling (2013)) or treated as hyperparameters fixed as time-dependent constants, both retaining the same functional form (Sohl-Dickstein et al. (2015)). The subsequent groundbreaking work by Ho et al. (2020) adopts the latter approach. Specifically, $\beta_t$ is set according to a linearly increasing schedule within a specific range based on the timesteps, although subsequent work (Nichol and Dhariwal (2021)) has proposed other scheduling schemes.

Given that we have fixed the forward process variances $\beta_t$, in Equation 30, since $\mathbf{z}T$ is Gaussian noise sampled from $\mathcal{N}(0, I)$ and $q$ has no learnable parameters, while $p(\mathbf{z}_T)$ converges to the normal distribution $\mathcal{N}(0, I)$ given a sufficiently large $T$. Therefore, the $L_T$ term is constant and typically small, allowing it to be ignored, and we only need to focus on minimizing the $L{t-1}$ and $L_0$ terms.

Recall that conditioning on $\mathbf{z}_0$ in Equation 24 makes the reverse posterior tractable, let:

Also recall that the Gaussian density function $\mathcal{N}(z; \mu, \sigma^2)$ can be written as:

We can then derive the followings:

Where the function $\Lambda(\mathbf{z}t, \mathbf{z}_0)$ is irrelevant of $\mathbf{z}{t-1}$, and hence is treated as a constant being ignored. By identifying the coefficients in Equation 37 as the form of:

With reference to Equation 31, we have:

Recall Equation 10, which allows us to directly sample $\mathbf{z}_t$ from $\mathbf{z}_0$ at any arbitrary $t$. Similarly, we can perform the reverse to obtain $\mathbf{z}_0$ directly from $\mathbf{z}_t$ at any arbitrary $t$, as follows:

Plugging into $\tilde{\mu}_t (\mathbf{z}_t, \mathbf{z}_0)$ in Equation 42, we can further simplify to:

Revisiting the reverse posterior in Equation (4), following the configuration as (Ho et al. (2020)), the reverse process variance $\boldsymbol{\Sigma}\theta(\mathbf{z}_t, t) = \sigma^2_t \mathbf{I}$ is also set to be untrained time-dependent constants. Here, $\sigma^2_t$ is either $\beta_t$ or $\tilde{\beta}_t = \frac{1 - \bar{\alpha}{t-1}}{1 - \bar{\alpha}t} \cdot \beta_t$, both of which are claimed to yield similar empirical results. It is noteworthy that subsequent work (Nichol and Dhariwal (2021)) has proposed learning the variance $\boldsymbol{\Sigma}\theta(\mathbf{z}_t, t)$ as an interpolation between $\beta_t$ and $\tilde{\beta}_t$, providing an alternative scheme.

Returning to Equation (30), where we aim to minimize the $L_{t-1}$ term, given Equation (4) and Equation (31), recall that the $\beta_t$ term, which their variances depend on, is fixed to constants. This leaves only the mean $\mu_{\theta}(\mathbf{z}{t}, t)$ to match. We can then parameterize $L{t-1}$ as follows:

Recall Equation (45), which we want the trained $\mu_{\theta}(\mathbf{z}_{t}, t)$ to match given $\mathbf{z}_t$. Since $\mathbf{z}_t$ is available as input during training, we can instead choose to fit the model on the noise, as derived below:

Here, $\epsilon_\theta$ refers to the model predicting the noise given $z_t$. Furthermore, (Ho et al. (2020)) found that simplifying the objective to the following form, without the coefficient term, works even better during training. This simplified objective is termed $L_{\text{simple}}$:

For the $L_0$ term in Equation (30), Ho et al. (2020) use a separate discrete decoder, which results in simply avoiding noise addition when $t=1$ during sampling. Consequently, we arrive at $L_{\text{simple}}$ being the sole objective to train diffusion models, termed denoising diffusion probabilistic models (DDPM) (Ho et al. (2020)).

Algorithm 1: DDPM Training \(\begin{aligned} &\\ &1. \text{ Repeat:} \\ &2. \quad \mathbf{z}_0 \sim q(\mathbf{z}_0) \\ &3. \quad t \sim \text{Uniform}(\{1, \ldots, T\}) \\ &4. \quad \epsilon \sim \mathcal{N}(0, \mathbf{I}) \\ &5. \quad \text{Take gradient descent step on } \nabla_\theta \left\Vert \epsilon - \epsilon_\theta \left( \sqrt{\bar{\alpha}_t} \mathbf{z}_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon, t \right) \right\Vert^2 \\ &6. \text{ Until converged} \end{aligned}\)

Algorithm 2: DDPM Sampling $$ \begin{aligned}

1. & \quad \mathbf{z}_T \sim \mathcal{N}(0, \mathbf{I}) \
2. & \quad \text{For } t = T, \ldots, 1: \
3. & \quad \quad \mathbf{z} \sim \mathcal{N}(0, \mathbf{I}) \text{ if } t > 1, \text{ else } \mathbf{z} = 0 \
4. & \quad \quad \mathbf{z}{t-1} = \frac{1}{\sqrt{\alpha_t}} \left( \mathbf{z}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \epsilon\theta (\mathbf{z}_t, t) \right) + \sigma_t \mathbf{n} \quad n \in \mathcal{N}(0, \mathbf{I}) \
5. & \quad \text{Return } \mathbf{z}_0 \end{aligned} $$

Algorithm 1 and Algorithm 2 outline the training and sampling processes of DDPM (Ho et al. (2020)), respectively. However, this sampling scheme is significantly less efficient compared to other generative models such as Generative Adversarial Networks (GANs). Subsequent work by Song et al. (2021), termed as DDIM (Song et al. (2021)) addresses this inefficiency by generalizing the forward and reverse diffusion processes to non-Markovian ones. Specifically, with the exact same training objective as in DDPM, another sampling scheme is proposed to sample $\mathbf{z}_{s}$ given $\mathbf{z}_t$, where $s < t$:

Here, $\epsilon \in \mathcal{N}(0, \mathbf{I})$. Note that when $\sigma_t = \sqrt{\frac{(1 - \bar{\alpha}{t-1})}{(1 - \bar{\alpha}_t)} \cdot \frac{1 - \bar{\alpha}_t}{\bar{\alpha}{t-1}}} \quad \text{for all } t$ and $s = t - 1$ at all timesteps $t \in [1, T]$, the processes resemble those of DDPM. Conversely, when $\sigma_t = 0$ for all $t$, the forward process (and hence the reverse process) becomes deterministic, leaving the only stochasticity to the initial diffusion latent. Other sampling schedules have been proposed (Karras et al. (2022), Liu et al. (2022), Song et al. (2023), Zhang et al. (2023)) to facilitate fast generation while maintaining high quality.

Diffusion Guidance

Guidance methods have been proposed to steer the diffusion generation towards catering the faithfulness of certain classes or conditions, specifically, the classifier guidance (Dhariwal and Nichol (2021)) and classifier-free guidance (Ho and Salimans (2022)). Classifier guidance aims at explicitly incorporating class information during sampling via the gradients of a noise-aware classifier $f_\phi$. Given the class information $y$, we have:

Where $\nabla \log q(y)$ is irrelevant to $z_t$ and thus can be ignored. In (Dhariwal and Nichol (2021)), a score-based conditioning trick from (Song and Ermon (2021)) that draws a connection between diffusion models and score matching (Song and Ermon (2020)) is used, giving:

Where $\epsilon_\theta$ is the noise-predictor model. Then, by substituting into the score function, we have:

Which leads to a new noise prediction that integrates the score from the classifier:

Where $s$ denotes the constant guidance scale, which modulates the trade-off between sample fidelity and diversity. A larger scale enhances the fidelity and faithfulness to the class but reduces the diversity of the generated samples.

On the other hand, classifier-free guidance (Ho and Salimans (2022)) adheres to a similar intuition but functions without an explicit classifier. During training, classifier-free guidance employs a scheme that randomly masks the conditioning information, thereby capturing both the conditional and unconditional distributions. This method enables extrapolation with a specified guidance scale during sampling to achieve comparable trade-offs. For more details, we refer readers to the original paper (Ho and Salimans (2022)) for a comprehensive discussion.

Latent Diffusion Models

Another notable advancement in diffusion models is the advent of Latent Diffusion Models (Rombach et al. (2022)). The core concept involves performing the diffusion process on smaller spatial latent representations of the original images, defined by pre-trained Variational Autoencoders (VAE). This approach significantly reduces the computational overhead associated with training on high-resolution pixel-level space, while the latent representations produced by VAE are potentially more semantically meaningful. Specifically, using a learned VAE encoder $\mathcal{E}$ and image data $\mathbf{x}$, we train and sample diffusion models on $\mathbf{z} = \mathcal{E}(\mathbf{x})$ while keeping $\mathcal{E}$ fixed. Finally, we apply a VAE decoder $\mathcal{D}$ on $\mathbf{z}$ to obtain the generated samples $\hat{\mathbf{x}} = \mathcal{D}(\mathbf{z})$.

References

[1] L. Weng, “What are diffusion models?,” lilianweng.github.io, Jul. 2021. [Online]. Available: https://lilianweng.github.io/posts/2021-07-11-diffusion-models/

[2] J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli, “Deep unsupervised learning using nonequilibrium thermodynamics,” in Proceedings of the 32nd International Conference on Machine Learning, vol. 37, pp. 2256–2265, PMLR, 2015.

[3] J. Ho, A. Jain, and P. Abbeel, “Denoising diffusion probabilistic models,” in Advances in Neural Information Processing Systems, vol. 33, pp. 6840–6851, Curran Associates, Inc., 2020.

[4] D. P. Kingma and M. Welling, “Auto-encoding variational bayes,” arXiv preprint arXiv:1312.6114, 2013.

[5] A. Q. Nichol and P. Dhariwal, “Improved denoising diffusion probabilistic models,” in International Conference on Machine Learning, pp. 8162–8171, 2021.

[6] J. Song, C. Meng, and S. Ermon, “Denoising diffusion implicit models,” in International Conference on Learning Representations, 2021.

[7] R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer, “High-resolution image synthesis with latent diffusion models,” in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 10684–10695, June 2022.

Thanks for reading! If you find any mistakes or have any questions, please feel free to contact me :D

Issues with the rendering of formulas?: you may download the markdown file from here and view it locally, my apologies for this temporary inconvenience!