Diffusion probabilistic models - Introduction

Author : Philippe Esling (esling@ircam.fr)

This second notebook continues the exploration of diffusion probabilistic models [ 1 ] in our four notebook series.

1. Score matching and Langevin dynamics.
2. Diffusion probabilistic models and denoising
3. Applications to waveforms with WaveGrad
4. Implicit models to accelerate inference

Here, we quickly recall the basics of score matching [ 3 ] and Langevin dynamics seen in the previous notebook. Then, we introduce the original formulation of diffusion probabilistic models based on thermodynamics [ 2 ] , and more recent formulations from denoising [ 1 ] .

Theoretical bases - quick recap

In this section we provide a quick recap on score matching from the previous notebook, still based on the swiss roll dataset.

Score matching

Score matching aims to learn the gradients (termed score) of $\log p(\mathbf{x})$ with respect to $\mathbf{x}$ instead of directly $\log p(\mathbf{x})$. Therefore, we seek a model to approximate $${\mathcal{F}}_{\theta }(\mathbf{x})\approx {\nabla }_{\mathbf{x}}\log p(\mathbf{x})$$

We have seen that optimizing this model with an MSE objective was equivalent to optimizing

$${\mathcal{L}}_{matching}=E_{\mathbf{x}\sim p(\mathbf{x})}\left[\text{ tr}\left({\nabla }_{\mathbf{x}}{\mathcal{F}}_{\theta }(\mathbf{x})\right)+\frac{1}{2}{\Vert {\mathcal{F}}_{\theta }(\mathbf{x})\Vert }_2^2\right],$$

where ${\nabla }_{\mathbf{x}}{\mathcal{F}}_{\theta }(\mathbf{x})$ denotes the Jacobian of ${\mathcal{F}}_{\theta }(\mathbf{x})$ with respect to $\mathbf{x}$. The problem with this formulation lies in the computation of this Jacobian, which does not scale well to high-dimensional data. This leads to a more efficient formulation of sliced score matching, which relies on random projections to approximate the computation of the Jacobian with

$$E_{\mathbf{v}\sim \mathcal{N}(0,1)}{E}_{\mathbf{x}\sim p(\mathbf{x})}\left[{\mathbf{v}}^T{\nabla }_{\mathbf{x}}{\mathcal{F}}_{\theta }(\mathbf{x})\mathbf{v}+\frac{1}{2}{\Vert {\mathbf{v}}^T{\mathcal{F}}_{\theta }(\mathbf{x})\Vert }_2^2\right],$$

where $\mathbf{v}\sim \mathcal{N}(0,1)$ are a set of Normal-distributed vectors. They show that this can be computed by using forward mode auto-differentiation, which is computationally efficient, as shown in the following implementation

Denoising score matching

Originally, denoising score matching was discussed by Vincent [ 3 ] in the context of denoising auto-encoders. In our case, we can completely remove the use of ${\nabla }_{\mathbf{x}}{\mathcal{F}}_{\theta }(\mathbf{x})$ in the computation of score matching, by corrupting the inputs through a distribution $q_{\sigma }(\tilde{\mathbf{x}}\mid \mathbf{x})$. It has been shown that the optimal network ${\mathcal{F}}_{\theta }(\mathbf{x})\approx {\nabla }_{\mathbf{x}}\log p(\mathbf{x})$ can be found by minimizing the following objective

$$E_{q_{\sigma }(\tilde{\mathbf{x}}\mid \mathbf{x})}{E}_{\mathbf{x}\sim p(\mathbf{x})}\left[{\Vert {\mathcal{F}}_{\theta }(\tilde{\mathbf{x}})-{\nabla }_{\tilde{\mathbf{x}}}\log q_{\sigma }(\tilde{\mathbf{x}}\mid \mathbf{x})\Vert }_2^2\right],$$

An important remark is that ${\mathcal{F}}_{\theta }(\mathbf{x})={\nabla }_{\mathbf{x}}\log q_{\sigma }(\mathbf{x})\approx {\nabla }_{\mathbf{x}}\log p(\mathbf{x})$ is only true when the noise is small enough $q_{\sigma }(\mathbf{x})\approx p(\mathbf{x})$. As it has been shown in [ 3 ] , [ 8 ] , if we choose the noise distribution to be $q_{\sigma }(\tilde{\mathbf{x}}\mid \mathbf{x})=\mathcal{N}(\tilde{\mathbf{x}}\mid \mathbf{x},{\sigma }^2\mathbf{I})$, then we have ${\nabla }_{\tilde{\mathbf{x}}}\log q_{\sigma }(\tilde{\mathbf{x}}\mid \mathbf{x})=\frac{\tilde{\mathbf{x}}-\mathbf{x}}{{\sigma }^2}$. Therefore, the denoising score matching loss simply becomes

$$l(\theta ;\sigma )=E_{q_{\sigma }(\tilde{\mathbf{x}}\mid \mathbf{x})}{E}_{\mathbf{x}\sim p(\mathbf{x})}\left[{\Vert {\mathcal{F}}_{\theta }(\tilde{\mathbf{x}})+\frac{\tilde{\mathbf{x}}-\mathbf{x}}{{\sigma }^2}\Vert }_2^2\right],$$

We can implement the denoising score matching loss as follows

Regarding optimization, we can perform a very simple implementation of this process, by define ${\mathcal{F}}_{\theta }(\mathbf{x})$ as being any type of neural network. We can perform the minimalistic implementation as follows

We can observe that our model has learned to represent ${\mathcal{F}}_{\theta }(\mathbf{x})\approx {\nabla }_x\log p(x)$ by plotting the output value across the input space

Langevin sampling

We have also seen that Langevin dynamics is a process from thermodynamics that can produce true samples from a density $p(\mathbf{x})$, by relying only on ${\nabla }_{\mathbf{x}}\log p(\mathbf{x})$

$${\mathbf{x}}_{t+1}={\mathbf{x}}_t+\frac{ϵ}{2}{\nabla }_{{\mathbf{x}}_t}logp({\mathbf{x}}_t)+\sqrt{ϵ}{\mathbf{z}}_t$$

where ${\mathbf{z}}_t\sim \mathcal{N}(0,\mathbf{I})$ and under $ϵ\to 0,t\to \inf$: ${\mathbf{x}}_t$ converges to an exact sample from $p(\mathbf{x})$. This is a key idea behind the score-based generative modeling approach.

Diffusion models

Diffusion probabilistic models were originally proposed by Sohl-Dickstein et al. [ 1 ] based on non-equilibrium thermodynamics. These models are based on two reciprocal processes that represent two Markov chains of random variables. One process $q({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1})$ that gradually adds noise to the input data (called the diffusion or forward process), destroying the signal up to full noise. In the opposite direction, the reverse process $p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)$ tries to learn how to invert this diffusion process (transform random noise into a high-quality waveform). This is examplified in the following figure, where we can see the whole model.

As we can see, the forward (and fixed) process $q({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1})$ gradually introduces noise at each step. Oppositely, the reverse (parametric) process $p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)$ must learn how to denoise local perturbations. Hence, learning involves estimating a large number of small perturbations, which is more tractable than trying to directly estimate the full distribution with a single potential function.

Both processes can be defined as parametrized Markov chains, but the diffusion process is usually simplified to inject pre-selected amounts of noise at each step. The reverse process is trained using variational inference, and can be modeled as conditional Gaussians, which allow for neural network parameterization and tractable estimation.

Formalization

Diffusion models are based on a series of latent variables ${\mathbf{x}}_1,\cdots ,{\mathbf{x}}_T$ that have the same dimensionality as a given input data, which is labeled as ${\mathbf{x}}_0\sim q({\mathbf{x}}_0)$. Then, we need to define the behavior of two process

$$\text{forward (diffusion)}:q({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1})$$ $$\text{reverse (parametric)}:p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)$$

Forward process

In the forward process, the data distribution $q({\mathbf{x}}_0)$ is gradually converted into an analytically tractable distribution $\pi (\mathbf{y})$, by repeated application of a Markov diffusion kernel $T_{\pi }(\mathbf{y}\mid {\mathbf{y}}^{\prime };\beta )$, with a given diffusion rate $\beta$.

$$q({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1})=T_{\pi }({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1};{\beta }_t)$$

This diffusion kernel can be set to gradually inject Gaussian noise, given a variance schedule ${\beta }_1,\cdots ,{\beta }_T$ such that

$$q({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1})=\mathcal{N}({\mathbf{x}}_t;\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I})$$

The complete distribution $q({\mathbf{x}}_{0:T})$ is called the diffusion process and is defined as

$$q({\mathbf{x}}_{0:T})=q({\mathbf{x}}_0)\prod _{t=1}^{T}q({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1})$$

Here, we show how to perform a naive implementation of the simple forward diffusion process with a constant variance schedule

We can define any type of variance schedules for ${\beta }_1,\cdots ,{\beta }_n$, as provided in the following function

Interestingly, the forward process admits sampling ${\mathbf{x}}_t$ at an arbitrary timestep $t$. Using notations ${\alpha }_t=1-{\beta }_t$ and ${\bar{\alpha }}_t={\prod }_{s=1}^t{\alpha }_s$, we have

$$q({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1})=\mathcal{N}({\mathbf{x}}_t;\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_{t-1},(1-{\bar{\alpha }}_t)\mathbf{I})$$

Therefore, we can update our diffusion sampling function to allow for this mecanism. Note that this depends on the given variance schedule of ${\beta }_1,\cdots ,{\beta }_T$ that we compute prior to the function.

This allows to perform a very efficient implementation of the forward process, where we can directly sample at any given timesteps, as shown in the following code.

Note that for training, we will also need to have access to the mean and variance of the posterior distribution of this process.

Reverse process

The generative distribution that we aim to learn will be trained to perform the reverse trajectory, starting from Gaussian noise to gradually remove local perturbations. Therefore the reverse process starts with our given tractable distribution $p({\mathbf{x}}_T)=\pi ({\mathbf{x}}_T)$ and is described as

$$p_{\theta }({\mathbf{x}}_{0:T})=p({\mathbf{x}}_T)\prod _{t=1}^T{p}_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)$$

Each of the transitions in this process can simply be defined as conditional Gaussians (note: which is reminiscent of the definition of VAEs). Therefore, during learning, only the mean and covariancce for a Gaussian diffusion kernel needs to be trained

$$p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)=\mathcal{N}({\mathbf{x}}_{t-1};{\mu }_{\theta }({\mathbf{x}}_t,t),{\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t))$$

The two functions defining the mean ${\mu }_{\theta }({\mathbf{x}}_t,t)$ and covariance ${\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t)$ can be parametrized by deep neural networks. Note also that these functions are parametrized by $t$, which means that a single model can be used for all time steps.

Here, we show a naive implementation of this process, where we have a given model to infer variance. Note that this model is shared across all time steps but conditionned on that said time step.

As we can see, the reverse process consists in inferring the values of the mean and log variance for a given timestep. Then, once we have learned the correponding model, we can perform the denoising of any given timestep, by providing both the sample ${\mathbf{x}}_t$ at a given time step, and that time step $t$ that we can use to condition the models for ${\mu }_{\theta }({\mathbf{x}}_t,t)$ and ${\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t)$.

Finally, obtaining samples from the model is given by running through the whole Markov chain in reverse, starting from the normal distribution to obtain samples from the target distribution. Note that this process can be very slow if we have a large number of steps, as we need to wait for a given ${\mathbf{x}}_t$ to infer the following ${\mathbf{x}}_{t-1}$

Model probability

The complete probability of the generative model is defined as

$$p_{\theta }({\mathbf{x}}_0)=\int p_{\theta }({\mathbf{x}}_{0:T})d{\mathbf{x}}_{1:T}$$

At first sight, this integral appears intractable. However, using a similar approach than variational inference, this integral can be rewritten as

$$\begin{array}{rl}{p}_{\theta }({\mathbf{x}}_0)& =\int p_{\theta }({\mathbf{x}}_{0:T})\frac{q({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)}{q({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)}d{\mathbf{x}}_{1:T}\\ & =\int q({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)\frac{p({\mathbf{x}}_{0:T})}{q({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)}d{\mathbf{x}}_{1:T}\end{array}$$

Training

By using Jensen's inequality on the previous expression, we can see that the training may be performed by optimizing the variational bound on negative log-likelihood

$$\begin{array}{rl}\mathbb{E}\left[-\log p_{\theta }({\mathbf{x}}_0)\right]& \le {\mathbb{E}}_q\left[-\log \frac{p_{\theta }({\mathbf{x}}_{0:T})}{q({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)}\right]\\ \mathcal{L}& ={\mathbb{E}}_q\left[-\log p({\mathbf{x}}_T)-\sum _{t\ge 1}\log \frac{p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)}{q({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1})}\right]\end{array}$$

Therefore, efficient training is allowed by optimizing random terms of $\mathcal{L}$ with gradient descent.

To optimize this loss, we will need several computational tools, notably the KL divergence between two gaussians, and the entropy of a Gaussian.

Training loss

In the original paper by Sohl-Dickstein et al. [ 1 ] , this loss is shown to be reducible to

$$\begin{array}{rl}K=-{\mathbb{E}}_q[& D_{KL}(q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0)\Vert p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t))\\ & +H_q({\mathbf{X}}_T|{\mathbf{X}}_0)-H_q({\mathbf{X}}_1|{\mathbf{X}}_0)-H_p({\mathbf{X}}_T)]\end{array}$$

Hence, all parts of this loss can be quite easily estimated, as we are dealing with Gaussian distributions in all cases

Training random time steps

The way that the model is trained is slightly counterintuitive, since we select a timestep at random to train for each of the batch input. The implementation taken from the DDIM repo provides a form of antithetic sampling, which allows to ensure that symmetrical points in the different chains are trained jointly. Therefore, the final procedure consists in first run the forward process on each input at a given (random) time steps (performing diffusion). Then we run the reverse process on this sample, and compute the loss.

We can very simply optimize this loss with the following training loop.

Denoising diffusion probabilistic models (DDPM)

In a very recent article, Ho et al. [ 1 ]  constructed over the diffusion models idea, by proposing several enhancements allowing to enhance the quality of the results. First, they proposed to rely on the following parameterization for the mean function $${\mu }_{\theta }({\mathbf{x}}_t,t)=\frac{1}{\sqrt{{\alpha }_t}}\left(({\mathbf{x}}_t-\frac{{\beta }_t}{{\sqrt{1-\bar{\alpha }}}_t}{ϵ}_{\theta }({\mathbf{x}}_t,t)\right)$$

Note that now, the model is trained at outputing directly a form of noise function, which is used in the sampling process. Furthermore, the authors suggest to rather use a fixed variance function

$${\mathbf{x}}_{t-1}=\frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t-\frac{1-{\alpha }_t}{\sqrt{1-\overline{{\alpha }_t}}}{ϵ}_{\theta }({\mathbf{x}}_t,t)\right)+{\sigma }_t\mathbf{z}$$

This leads to a new sampling procedure for the reverse process as follows (we also quickly redefine the model to output the correct dimensionality).

Notably, the forward process posteriors are tractable when conditioned on ${\mathbf{x}}_0$

$$\begin{array}{rl}q({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1})& =\mathcal{N}({\mathbf{x}}_{t-1};1-{\beta }_{\mathbf{t}}{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I})\end{array}$$

And we can obtain the corresponding mean ${\tilde{\mu }}_t({\mathbf{x}}_t,{\mathbf{x}}_0)$ and variance $\tilde{\beta }$ as

$$\begin{array}{rl}{\tilde{\mu }}_t({\mathbf{x}}_t,{\mathbf{x}}_0)& =\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{\mathbf{x}}_0+\frac{\sqrt{{\bar{\alpha }}_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{\mathbf{x}}_t\\ {\tilde{\beta }}_t& =\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t\end{array}$$

Training in DDPM

Further improvements come from variance reduction by rewriting $\mathcal{L}$ as a sum of KL divergences

$$\begin{array}{rl}\mathcal{L}& ={\mathbb{E}}_q\left[{\mathcal{L}}_T+\sum _{t>1}{\mathcal{L}}_{t-1}+{\mathcal{L}}_0\right]\\ {\mathcal{L}}_T& =D_{KL}(q({\mathbf{x}}_T\mid {\mathbf{x}}_0)\Vert p({\mathbf{x}}_T))\\ {\mathcal{L}}_{t-1}& =D_{KL}(q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0)\Vert p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t))\\ {\mathcal{L}}_0& =-\log p_{\theta }({\mathbf{x}}_0\mid {\mathbf{x}}_1)\end{array}$$

All the KL divergences defined in this equation compare Gaussians, which means that they have a closed-form solution.

This leads to a new loss function as implemented in the following (note that this objective does not provide large change to the optimization itself).

Simplifying loss to denoising score matching

The paper by Ho et al. [ 1 ] proposes a new parameterization for the mean of the reverse process $${\mathbf{x}}_{t-1}=\frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t-\frac{1-{\alpha }_t}{\sqrt{1-\overline{{\alpha }_t}}}{ϵ}_{\theta }({\mathbf{x}}_t,t)\right)+{\sigma }_t\mathbf{z}$$ Based on this parametrization, they show that the training objective can simplify to $${\mathcal{L}}_{t-1}-C={\mathbb{E}}_{{\mathbf{x}}_0,ϵ}\left[\frac{{\beta }_t^2}{2{\sigma }_t^2{\alpha }_t(1-{\bar{\alpha }}_t)}\Vert ϵ-{ϵ}_{\theta }(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,t){\Vert }^2\right]$$ which resembles denoising score matching over multiple noise scales indexed by $t$.

Further simplified training objective

The authors discuss the fact that it is beneficial to the sample quality to completely remove the complicated factor at the beginning of the loss. This further simplifies the objective as $${\mathcal{L}}_{\text{simple}}={\mathbb{E}}_{t,{\mathbf{x}}_0,ϵ}\left[\Vert ϵ-{ϵ}_{\theta }(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,t){\Vert }^2\right].$$ We can see that this objective now very closely ressemble the denoising score matching formulation. Furthermore, it provides an extremely simple implementation.

Stabilizing training with Exponential Moving Average (EMA)

This idea is found in most of the implementations, which allows to implement a form of model momentum. Instead of directly updating the weights of the model, we keep a copy of the previous values of the weights, and then update a weighted mean between the previous and new version of the weights. Here, we reuse the implementation proposed in the DDIM repository.

The training loop is finally obtained with the following code

Bibliography

[1] Ho, J., Jain, A., & Abbeel, P. (2020). Denoising diffusion probabilistic models. arXiv preprint arXiv:2006.11239.

[2] Sohl-Dickstein, J., Weiss, E. A., Maheswaranathan, N., & Ganguli, S. (2015). Deep unsupervised learning using nonequilibrium thermodynamics. arXiv preprint arXiv:1503.03585.

[3] Vincent, P. (2011). A connection between score matching and denoising autoencoders. Neural computation, 23(7), 1661-1674.

[4] Song, J., Meng, C., & Ermon, S. (2020). Denoising Diffusion Implicit Models. arXiv preprint arXiv:2010.02502.

[5] Chen, N., Zhang, Y., Zen, H., Weiss, R. J., Norouzi, M., & Chan, W. (2020). WaveGrad: Estimating gradients for waveform generation. arXiv preprint arXiv:2009.00713.

[6] Hyvärinen, A. (2005). Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(Apr), 695-709.

[7] Song, Y., Garg, S., Shi, J., & Ermon, S. (2020, August). Sliced score matching: A scalable approach to density and score estimation. In Uncertainty in Artificial Intelligence (pp. 574-584). PMLR.

[8] Song, Y., & Ermon, S. (2019). Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems (pp. 11918-11930).

Inspiration and sources

https://colab.research.google.com/github/google/jax/blob/master/docs/notebooks/score_matching.ipynb https://github.com/ermongroup/sliced_score_matching