1.1 VAE Setting
A long time ago, in a galaxy far far away, we want to build a generator that generates images from a latent code. The simplest (and perhaps one of the most classical) approach is to consider an encoder-decoder pair shown below. This is called a variational autoencoder (VAE) [1, 2, 3].
The autoencoder has an input variable x and a latent variable z. For the sake of understanding the subject, we treat x as a beautiful image and z as some kind of vector living in some high dimensional space.
The name “variational” comes from the factor that we use probability distributions to describe x and z. Instead of resorting to a deterministic procedure of converting x to z, we are more interested in ensuring that the distribution p(x) can be mapped to a desired distribution p(z), and go backwards to p(x). Because of the distributional setting, we need to consider a few distributions.
• p(x): The distribution of x. It is never known. If we knew it, we would have become a billionaire. The whole galaxy of diffusion models is to find ways to draw samples from p(x).
• p(z): The distribution of the latent variable. Because we are all lazy, let’s just make it a zero-mean unit-variance Gaussian p(z) = N(0, I).
• p(z|x): The conditional distribution associated with the encoder, which tells us the likelihood of z when given x. We have no access to it. p(z|x) itself is not the encoder, but the encoder has to do something so that it will behave consistently with p(z|x).
• p(x|z): The conditional distribution associated with the decoder, which tells us the posterior probability of getting x given z. Again, we have no access to it.
The four distributions above are not too mysterious. Here is a somewhat trivial but educational example that can illustrate the idea.
Example. Consider a random variable X distributed according to a Gaussian mixture model with a latent variable denoting the cluster identity such that
for k = 1, . . . , K. We assume
-th cluster only, the conditional distribution of X given Z is
Smart readers like you will certainly complain: “Your example is so trivially unreal.” No worries. We understand. Life is of course a lot harder than a Gaussian mixture model with known means and known variance. But one thing we realize is that if we want to find the magical encoder and decoder, we must have a way to find the two conditional distributions. However, they are both high-dimensional creatures. So, in order for us to say something more meaningful, we need to impose additional structures so that we can generalize the concept to harder problems.
In the literature of VAE, people come up with an idea to consider the following two proxy distributions:
• ): The proxy for p(z|x). We will make it a Gaussian. Why Gaussian? No particular good reason. Perhaps we are just ordinary (aka lazy) human beings.
• ): The proxy for p(x|z). Believe it or not, we will make it a Gaussian too. But the role of this Gaussian is slightly different from the Gaussian
). While we will need to estimate the mean and variance for the Gaussian
), we do not need to estimate anything for the Gaussian
). Instead, we will need a decoder neural network to turn z into x. The Gaussian
) will be used to inform us how good our generated image x is.
The relationship between the input x and the latent z, as well as the conditional distributions, are summarized in Figure 1. There are two nodes x and z. The “forward” relationship is specified by p(z|x) (and approximated by )), whereas the “reverse” relationship is specified by p(x|z) (and approximated by
)).
Figure 1: In a variational autoencoder, the variables x and z are connected by the conditional distributions p(x|z) and p(z|x). To make things work, we introduce two proxy distributions ) and
), respectively.
1.2 Evidence Lower Bound
How do we use these two proxy distributions to achieve our goal of determining the encoder and the decoder? If we treat and
as optimization variables, then we need an objective function (or the loss function) so that we can optimize
and
through training samples. To this end, we need to set up a loss function in terms of
and
. The loss function we use here is called the Evidence Lower BOund (ELBO) [1]:
You are certainly puzzled how on the Earth people can come up with this loss function!? Let’s see what ELBO means and how it is derived.
In a nutshell, ELBO is a lower bound for the prior distribution log p(x) because we can show that
where the inequality follows from the fact that the KL divergence is always non-negative. Therefore, ELBO is a valid lower bound for log p(x). Since we never have access to log p(x), if we somehow have access to ELBO and if ELBO is a good lower bound, then we can effectively maximize ELBO to achieve the goal of maximizing log p(x) which is the gold standard. Now, the question is how good the lower bound is. As you can see from the equation and also Figure 2, the inequality will become an equality when our proxy ) can match the true distribution p(z|x) exactly. So, part of the game is to ensure
) is close to p(z|x).
Figure 2: Visualization of log p(x) and ELBO. The gap between the two is determined by the KL divergence )).
We now have ELBO. But this ELBO is still not too useful because it involves p(x, z), something we have no access to. So, we need to do a little more things. Let’s take a closer look at ELBO
where we secretly replaced the inaccessible p(x|z) by its proxy ). This is a beautiful result. We just showed something very easy to understand.
There are two terms in Eqn (6):
• Reconstruction. The first term is about the decoder. We want the decoder to produce a good image x if we feed a latent z into the decoder (of course!!). So, we want to maximize log ). It is similar to maximum likelihood where we want to find the model parameter to maximize the likelihood of observing the image. The expectation here is taken with respect to the samples z (conditioned on x). This shouldn’t be a surprise because the samples z are used to assess the quality of the decoder. It cannot be an arbitrary noise vector but a meaningful latent vector. So, z needs to be sampled from
).
• Prior Matching. The second term is the KL divergence for the encoder. We want the encoder to turn x into a latent vector z such that the latent vector will follow our choice of (lazy) distribution N(0, I). To be slightly more general, we write p(z) as the target distribution. Because KL is a distance (which increases when the two distributions become more dissimilar), we need to put a negative sign in front so that it increases when the two distributions become more similar.
The reconstruction term and the prior matching terms are illustrated in Figure 3. In both cases, and during training, we assume that we have access to both z and x, where z needs to be sampled from ). Then for reconstruction, we estimate
to maximize
). For prior matching, we find
to minimize the KL divergence. The optimization can be challenging, because if you update
, the distribution
) will change.
Figure 3: Interpreting the reconstruction term and the prior matching term in ELBO for a variational autoencoder.
1.3 Training VAE
Now that we understand the meaning of ELBO, we can discuss how to train the VAE. To train a VAE, we need the ground truth pairs (x, z). We know how to get x; it is just the image from a dataset. But correspondingly what should z be?
Let’s talk about the encoder. We know that z is generated from the distribution ). We also know that
) is a Gaussian. Assume that this Gaussian has a mean
and a covariance matrix
(Ha! Our laziness again! We do not use a general covariance matrix but assume an equal variance).
The tricky part is how to determine and
from the input image x. Okay, if you run out of clue, do not worry. Welcome to the Dark Side of the Force. We construct a deep neural network(s) such that
Therefore, the samples (where
denotes the
-th training sample in the training set) can be sampled from the Gaussian distribution
The idea is summarized in Figure 4 where we use a neural network to estimate the Gaussian parameters, and from the Gaussian we draw samples. Note that ) and
) are functions of
. Thus, for a different
we will have a different Gaussian.
Figure 4: Implementation of a VAE encoder. We use a neural network to take the image x and estimate the mean and variance
of the Gaussian distribution.
Let’s talk about the decoder. The decoder is implemented through a neural network. For notation simplicity, let’s define it as decodewhere
denotes the network parameters. The job of the decoder network is to take a latent variable z and generates an image
:
Now let’s make one more (crazy) assumption that the error between the decoded image and the ground truth image x is Gaussian. (Wait, Gaussian again?!) We assume that
Then, it follows that the distribution ) is
where D is the dimension of x. This equation says that the maximization of the likelihood term in ELBO is literally just the loss between the decoded image and ground truth. The idea is shown in Figure 5.
Figure 5: Implementation of a VAE decoder. We use a neural network to take the latent vector z and generate an image . The log likelihood will give us a quadratic equation if we assume a Gaussian distribution.
1.4 Loss Function
Once you understand the structure of the encoder and the decoder, the loss function is easy to understand. We approximate the expectation by Monte-Carlo simulation:
where is the
-th sample in the training set, and
is sampled from
). The distribution
is
).
The z in the KL divergence term does not depend on because we are measuring the KL divergence between two distributions. The variable z here is a dummy.
One last thing we need to clarify is the KL divergence. Since ) and p(z) = N(0, I), we are essentially coming two Gaussian distributions. If you go to Wikipedia, you can see that the KL divergence for two d-dimensional Gaussian distributions
) and
) is
12
Substituting our distributions by considering , we can show that the KL divergence has an analytic expression
where d is the dimension of the vector z. Therefore, the overall loss function Eqn (12) is differentiable. So, we can train the encoder and the decoder end-to-end by backpropagating the gradients.
1.5 Inference with VAE
For inference, we can simply throw a latent vector z (which is sampled from p(z) = N(0, I)) into the decoder decodeand get an image x. That’s it; see Figure 6.
Figure 6: Using VAE to generate image is as simple as sending a latent noise code z through the decoder.
If you would like to read more, we highly recommend the tutorial by Kingma and Welling [1]. A shorter tutorial can be found at [2]. If you type VAE tutorial PyTorch in Google, you will be able to find hundreds if not thousands programming tutorials and videos.
In this section, we will discuss the DDPM by Ho et al. [4]. If you are confused by the thousands of tutorials online, rest assured that DDPM is not that complicated. All you need to understand is the following summary:
Why increment? It’s like turning the direction of a giant ship. You need to turn the ship slowly towards your desired direction or otherwise you will lose control. The same principle applies to your life, your company HR, your university administration, your spouse, your children, and anything around your life. “Bend one inch at a time!” (Credit: Sergio Goma who made this comment at Electronic Imaging 2023.)
The structure of the diffusion model is shown below. It is called the variational diffusion model [5]. The variational diffusion model has a sequence of states :
• : It is the original image, which is the same as x in VAE. •
: It is the latent variable, which is the same as z in VAE. Since we are all lazy, we want
). •
: They are the intermediate states. They are also the latent variables, but they are not white Gaussian.
The structure of the variational diffusion model is shown in Figure 7. The forward and the reverse paths are analogous to the paths of a single-step variational autoencoder. The difference is that the encoders and decoders have identical input-output dimensions. The assembly of all the forward building blocks will give us the encoder, and the assembly of all the reverse building blocks will give us the decoder.
Figure 7: Variational diffusion model. In this model, the input image is and the white noise is
. The intermediate variables (or states)
are latent variables. The transition from
to
is analogous to the forward step (encoder) in VAE, whereas the transition from
to
is analogous to the reverse step (decoder) in VAE. Note, however, that the input dimension and the output dimension of the encoders/decoders here are identical.
2.1 Building Blocks
Transition Block The t-th transition block consists of three states , and
. There are two possible paths to get to state
, as illustrated in Figure 8.
• The forward transition that goes from to
. The associated transition distribution is
). In plain words, if you tell us
, we can tell you
according to
). However, just like a VAE, the transition distribution
) is never accessible. But this is okay. Lazy people like us
will just approximate it by a Gaussian ). We will discuss the exact form
later, but it is just some Gaussian.
• The reverse transition goes from to
. Again, we never know
) but that’s okay. We just use another Gaussian
) to approximate the true distribution, but its mean needs to be estimated by a neural network.
Figure 8: The transition block of a variational diffusion model consists of three nodes. The transition distributions ) and
) are not accessible, but we can approximate them by Gaussians.
Initial Block The initial block of the variational diffusion model focuses on the state . Since all problems we study starts at
, there is only the reverse transition from
to
, and nothing that goes from
to
. Therefore, we only need to worry about
). But since
) is never accessible, we approximate it by a Gaussian
) where the mean is computed through a neural network. See Figure 9 for illustration.
Figure 9: The initial block of a variational diffusion model focuses on the node . Since there is no state before time
to
.
Final Block. The final block focuses on the state . Remember that
is supposed to be our final latent variable which is a white Gaussian noise vector. Because it is the final block, there is only a forward transition from
to
, and nothing such as
to
. The forward transition is approximated by
), which is a Gaussian. See Figure 10 for illustration.
Figure 10: The final block of a variational diffusion model focuses on the node . Since there is no state after time t = T, we only have a forward transition from
to
.
Understanding the Transition Distribution. Before we proceed further, we need to detour slightly to talk about the transition distribution ). We know that it is Gaussian. But we still need to know its formal definition, and the origin of this definition.
In other words, the mean is and the variance is 1
. The choice of the scaling factor
is to make sure that the variance magnitude is preserved so that it will not explode and vanish after many iterations.
Remark. For those who would like to understand how we derive the probability density of a mixture model in Eqn (16), we can show a simple derivation. Consider a mixture model
If we consider a new variable +
where
), then the distribution of y can be derived by using the law of total probability:
Since y|k is a linear combination of a Gaussian random variable x and another Gaussian random variable , the sum y will remain as a Gaussian. The mean is
2.2 The magical scalars √αt and 1 − αt
You may wonder how the genie (the authors of the denoising diffusion) come up with the magical scalars and (1
) for the above transition probability. To demystify this, let’s begin with two unrelated scalars
and
, and we define the transition distribution as
Here is the rule of thumb:
So, if we want limCov[
(so that the distribution of
will approach N(0, I), then b =
. Now, if we let
, then
. This will give us
Or equivalently, (1
). You can replace
by
, if you prefer a scheduler.
2.3 Distribution qϕ(xt|x0)
With the understanding of the magical scalars, we can talk about the distribution ). That is, we want to know how
will be distributed if we are given
.
The utility of the new distribution ) is its one-shot forward diffusion step compared to the chain
. In every step of the forward diffusion model, since we already know
and we assume that all subsequence transitions are Gaussian, we will know
immediately for any t. The situation can be understood from Figure 11.
Figure 11: The difference between ) and
).
If you are curious about how the probability distribution evolves over time t, we show in Figure 12 the trajectory of the distribution. You can see that when we are at t = 0, the initial distribution is a mixture of two Gaussians. As we progress by following the transition defined in Eqn (26), we can see that the distribution gradually becomes the single Gaussian N(0, 1).
Figure 12: Trajectory plot of the Gaussian mixture, as we progress to transit the probability distribution to N(0, 1).
In the same plot, we overlay and show a few instantaneous trajectories of the random samples as a function of time t. The equation we used to generate the samples is
As you can see, the trajectories of more or less follow the distribution
).
2.4 Evidence Lower Bound
Now that we understand the structure of the variational diffusion model, we can write down the ELBO and hence train the model. The ELBO for the variational diffusion model is
We can interpret the meaning of this ELBO. The ELBO here consists of three components:
• Reconstruction. The reconstruction term is based on the initial block. We use the log-likelihood ) to measure how good the neural network associated with
can recover the image
from the latent variable
. The expectation is taken with respect to the samples drawn from
) which is the distribution that generates
. If you are puzzled why we want to draw samples from
), just think about that where the samples
should come from. The samples
do not come from the sky. Since they are the intermediate latent variables, they are created by the forward transition
). So, we should generate samples from
).
• Prior Matching. The prior matching term is based on the final block. We use KL divergence to measure the difference between ) and
). The first distribution
) is the forward transition from
to
. This is how
is generated. The second distribution is
). Because of our laziness,
) is N(0, I). We want
) to be as close to N(0, I) as possible. The samples here are
which are drawn from
) because
) provides the forward sample generation process.
• Consistency. The consistency term is based on the transition blocks. There are two directions. The forward transition is determined by the distribution ) whereas the reverse transition is determined by the neural network
). The consistency term uses the KL divergence to measure the deviation. The expectation is taken with respect to samples (
) drawn from the joint distribution
). Oh, what is
)? No worries. We will get rid of it soon.
At this moment, we will skip the training and inference because this formulation is not ready for implementation. We will discuss one more tricks, and then we will talk about the implementation.
Proof of Eqn (27). Let’s define the following notation: means the collection of all state variables from t = 0 to t = T. We also recall that the prior distribution p(x) is the distribution for the image
. So it is equivalent to
). With these in mind, we can show that
where we notice that the conditional expectation can be simplified to samples and
only, because log
only depends on
and
. Finally, we look at the product term. We can show that
2.5 Rewrite the Consistency Term
The nightmare of the above variation diffusion model is that we need to draw samples () from a joint distribution
). We don’t know what
) is! Well, it is a Gaussian, of course, but still we need to use future samples
to draw the current sample
. This is odd, and it is not fun.
Inspecting the consistency term, we notice that ) and
) are moving along two opposite directions. Thus, it is unavoidable that we need to use
and
. The question we need to ask is: Can we come up with something so that we do not need to handle two opposite directions while we are able to check consistency?
So, here is the simple trick called Bayes theorem.
With this change of the conditioning order, we can switch ) to
) by adding one more condition variable
. The direction
) is now parallel to
) as shown in Figure 13. So, if we want to rewrite the consistency term, a natural option is to calculate the KL divergence between
) and
).
Figure 13: If we consider the Bayes theorem in Eqn (31), we can define a distribution ) that has a direction parallel to
).
If we manage to go through a few (boring) algebraic derivations, we can show that the ELBO is now:
Let’s quickly make three interpretations:
• Reconstruction. The new reconstruction term is the same as before. We are still maximizing the log-likelihood.
• Prior Matching. The new prior matching is simplified to the KL divergence between ) and
). The change is due to the fact that we now condition upon
. Thus, there is no need to draw samples from
) and take expectation.
• Consistency. The new consistency term is different from the previous one in two ways. Firstly, the running index t starts at t = 2 and ends at t = T. Previously it was from t = 1 to t = 1. Accompanied with this is the distribution matching, which is now between
) and
). So, instead of asking a forward transition to match with a reverse transition, we use
to construct a reverse transition and use it to match with
.
2.6 Derivation of qϕ(xt−1|xt, x0)
Now that we know the new ELBO for the variational diffusion model, we should spend some time discussing its core component which is ). In a nutshell, what we want to show is that
• ) is not as crazy as you think. It is still a Gaussian. • Since it is a Gaussian, it is fully characterized by the mean and covariance. It turns out that
The interesting part of Eqn (35) is that ) is completely characterized by
and
. There is no neural network required to estimate the mean and variance! (You can compare this with VAE where a network is needed.) Since a network is not needed, there is really nothing to “learn”. The distribution
) is automatically determined if we know
and
. No guessing, no estimation, nothing.
The realization here is important. If we look at the consistency term, it is a summation of many KL divergence terms where the t-th term is
As we just said, there is nothing to do with ). But we need to do something to
) so that we can calculate the KL divergence.
So, what should we do? We know that ) is Gaussian. If we want to quickly calculate the KL divergence, then obviously we need to assume that
) is also a Gaussian. Yeah, no kidding. We have no justification why it is a Gaussian. But since
is a distribution we can choose, we should of course choose something easier. To this end, we pick
where we assume that the mean vector can be determined using a neural network. As for the variance, we choose the variance to be ). This is identical to Eqn (37)! Thus, if we put Eqn (35) side by side with
), we notice a parallel relation between the two:
Therefore, the KL divergence is simplified to
where we used the fact that the KL divergence between two identical-variance Gaussians is just the Euclidean distance square between the two mean vectors. If we go back to the definition of ELBO in Eqn (32), we can rewrite it as
A few observations are interesting:
• We dropped all the subscripts because q is completely described as long as we know
. We are just adding (different level of) white noise to each of
. This will give us an ELBO that only requires us to optimize over
.
• The parameter is realized through the network
). It is the network weight for
). • The sampling from
) is done according to Eqn (21) which states that
(1
).
• Given ), we can calculate log
), which is just log
(1)I). So, as soon as we know
, we can send it to a network
) to return us a mean estimate. The mean estimate will then be used to compute the likelihood.
Before we go further down, let’s complete the story by discussing how Eqn (35) was determined.
2.7 Training and Inference
The ELBO in Eqn (43) suggests that we need to find a network that can somehow minimize this loss:
Since is our design, there is no reason why we cannot define it as something more convenient. So here is an option:
Substituting Eqn (51) and Eqn (52) into Eqn (50) will give us
Therefore ELBO can be simplified into
The first term is
Substituting Eqn (54) into Eqn (53) will simplify ELBO as
Therefore, the training of the neural network boils down to a simple loss function:
The loss function defined in Eqn (55) is very intuitive. Ignoring the constants and expectations, the main subject of interest, for a particular , is
This is nothing but a denoising problem because we need to find a network such that the denoised image
) will be close to the ground truth
. What makes it not a typical denoiser is that
• : We are not trying to denoise any random noisy image. Instead, we are carefully choosing the noisy image to be
Figure 14: Forward sampling process. The forward sampling process is originally a chain of operations. However, if we assume Gaussian, then we can simplify the sampling process as a one-step data generation.
• : We do not weight the denoising loss equally for all steps. Instead, there is a scheduler to controls the relative emphasis on each denoising loss. However, for simplicity, we can drop these. It has minor impacts.
• : The summation can be replaced by a uniform distribution
Uniform[1, T].
Once the denoiser is trained, we can apply it to do the inference. The inference is about sampling images from the distributions
) over the sequence of states
. Since it is the reverse
Figure 15: Training of a denoising diffusion probabilistic model. For the same neural network , we send noisy inputs
to the network. The gradient of the loss is back-propagated to update the network. Note that the noisy images are not arbitrary. They are generated according to the forward sampling process.
diffusion process, we need to do it recursively via:
This leads to the following inferencing algorithm.
Figure 16: Inference of a denoising diffusion probabilistic model.
2.8 Derivation based on Noise Vector
If you are familiar with the denoising literature, you probably know the residue-type of algorithm that predicts the noise instead of the signal. The same spirit applies to denoising diffusion, where we can learn to predict the noise. To see why this is the case, we consider Eqn (24). If we re-arrange the terms we will obtain
Substituting this into ), we can show that
So, if we can design our mean estimator , we can freely choose it to match for the form:
Substituting Eqn (57) and Eqn (58) into Eqn (50) will give us a new ELBO
Therefore, if you give us , we will return you a predicted noise
). This will give us an alternative training scheme
Consequently, the inference step can be derived through
Summarizing it here, we have
2.9 Inversion by Direct Denoising (InDI)
If we look at the DDPM equation, we will see that the update Eqn (56) takes the following form:
In other words, the (1)-th estimate is a linear combination of three terms: the current estimate
, the denoised version denoise(
) and a noise term. The current estimate and the noise term are easy to understand. But what is “denoise”? An interesting paper by Delbracio and Milanfar [6] looked at the generative diffusion models from a pure denoising perspective. As it turns out, this surprisingly simple perspective is consistent with the other more advanced diffusion models in some good ways.
)? Denoising is a generic procedure that removes noise from a noisy image. In the good old days of statistical signal processing, a standard textbook problem is to derive the optimal denoiser for white noise. Given the observation model
can you construct an estimator ) such that the mean squared error is minimized? We shall skip the derivation of the solution to this classical problem because you can find it in any probability textbook, e.g., [7, Chapter 8]. The solution is
So, going back to our problem: If we assume that
then clearly the denoiser is the conditional expectation of the posterior distribution:
Thus, if we are given the distribution ), then the optimal denoiser is just the conditional expectation of this distribution. Such a denoiser is called the minimum mean squared error (MMSE) denoiser. MMSE denoiser is not the “best” denoiser; It is only the optimal denoiser with respect to the mean squared error. Since mean squared error is never a good metric for image quality, minimizing the MSE will not necessarily give us a better image. Nevertheless, people like MMSE denoisers because they are easy to derive.
Incremental Denoising Steps. If you understand that an MMSE denoiser is equivalent to the conditional expectation of the posterior distribution, you will appreciate the incremental denoising. Here is how it works. Suppose that we have a clean image and a noise image y. Our goal is to form a linear combination of
and y via a simple equation
Now, consider a small step previous to time t. The following result, showed by [6], provides some useful utilities:
If we define as the left hand side, replace
by
, and write
] as denoise(
), then the above equation will become
where is a small step in time. Eqn (64) gives us an inference step. If you tell us the denoiser and suppose that you start with a noisy image y, then we can iteratively apply Eqn (64) to retrieve the images
, . . . ,
.
Training. The training of the iterative scheme requires a denoiser that generates denoise(). To this end, we can train a neural network denoise
(where
denotes the network weight):
Here, the distribution “uniform” specifies that the time step t is drawn uniformly from a given distribution. Therefore, we are training one denoiser for all time steps t. The expectation (x, y) is generally fulfilled when you use a pair of noisy and clean images from the training dataset. After training, we can perform the incremental update via Eqn (64).
Connection with Denoising Score-Matching. Although we have not yet discussed score-matching (which will be presented in the next Section), an interesting fact about the above iterative denoising procedure is that it is related to denoising score-matching. At the high level, we can rewrite the iteration as
This is an ordinary differential equation (ODE). If we let so that the noise level in
is
, then we can use several results in the literature to show that
Therefore, the incremental denoising iteration is equivalent to the denoising score-matching, at least in the limiting case determined by the ODE.
Adding Stochastic Steps. The above incremental denoising iteration can be equipped with stochastic perturbation. For the inference step, we can define a sequence of noise levels , and define
As or training, one can train a denoiser via
where .
The literature of DDPM is quickly exploding. The original paper by Sohl-Dickstein et al. [10] and Ho et al. [4] are the must-reads to understand the topic. For a more “user-friendly” version, we found that the tutorial by Luo very useful [11]. Some follow up works are highly cited, including the denoising diffusion implicit models by Song et al. [12]. In terms of application, people have been using DDPM for various image synthesis applications, e.g., [13, 14].
Score-based generative models [8] are alternative approaches to generate data from a desired distribution. There are several core ingredients: the Langevin dynamics, the (Stein) score function, and the score-matching loss. In this section, we will look at these topics one by one.
3.1 Langevin Dynamics
An interesting starting point of our discussion is the Langevin dynamics. It is a very physics topic that will appear to have nothing to do with generative models. But please don’t worry. They are related, in fact, in a good way.
Instead of telling you the physics right a way, let’s talk about how Langevin dynamics can be used to draw samples from a distribution. Imagine that we are given a distribution p(x) and suppose that we want to draw samples from p(x). Langevin dynamics is an iterative procedure that allows us to draw samples according to the following equation.
You may wonder, what on the Earth is this mysterious equation about?! Here is the short and quick answer. If you ignore the noise termat the end, the Langevin dynamics equation in Eqn (68) is literally gradient descent. The descent direction
log p(x) is carefully chosen that
will converge to the distribution p(x). If you watch any YouTube videos mumbling Langevin dynamics equations for 10 minutes without explaining what it is, you can gently tell them the following:
Consider a distribution p(x). The shape of this distribution is defined and is fixed as soon as the model parameters are defined. For example, if you choose a Gaussian, the shape and location of the Gaussian is fixed once you specify the mean and the variance. The value p(x) is nothing but the probability density evaluated at a data point x. Therefore, going from one x to another , we are just moving from one value p(x) to a different value
). The underlying shape of the Gaussian is not changed.
Suppose that we start with some arbitrary location in . We want to move it to (one of) the peak(s) of the distribution. The peak is a special place because it is where the probability is the highest. So, if we say that a sample x is drawn from a distribution p(x), certainly the “optimal” location for x is where p(x) is maximized. If p(x) has multiple local minima, any one of them would be fine. So, naturally, the goal of sampling is equivalent to solving the optimization
We emphasize again that this is not maximum likelihood estimation. In maximum likelihood, the data point x is fixed but the model parameters are changing. Here, the model parameters are fixed but the data point is changing. The table below summarizes the difference.
The optimization can be solved in many ways. The cheapest way is, of course, gradient descent. For log p(x), we see that the gradient descent step is
where log
) denotes the gradient of log p(x) evaluated at
, and
is the step size. Here we use “+” instead of the typical “
” because we are solving a maximization problem.
Example. Consider a Gaussian mixture ) +
). We can numerically calculate
log p(x). For demonstration, we choose
2,
2. We initialize
05. We run the above gradient descent iteration for
) for t = 1, . . . , T. As we can see in the figure below, the sequence
simply follows the shape of the Gaussian and climb to one of the peaks.
What is more interesting is when we add the noise term. Instead of landing at the peak, the sequence move around the peak and finishes somewhere near the peak. The closer we are to the peak, the higher probability we will stop there.
Figure 17 shows an interesting description of the sample trajectory. Starting with an arbitrary location, the data point will do a random walk according to the Langevin dynamics equation. The direction of the random walk is not completely arbitrary. There is a certain amount of pre-defined drift while at every step there is some level of randomness. The drift is determined by
log p(x) where as the randomness comes from z.
Figure 17: Trajectory of sample evolutions using the Langevin dynamics. We colored the two modes of the Gaussian mixture in different colors for better visualization. The setting here is identical to the example above, except that the step size is 001.
As we can see from the example above, the addition of the noise term actually changes the gradient descent to stochastic gradient descent. Instead of shooting for the deterministic optimum, the stochastic gradient descent climbs up the hill randomly. Since we use a constant step size, the final solution will just oscillate around the peak. So, we can summarize Langevin dynamics equation as
But why do we want to do stochastic gradient descent instead of gradient descent? The key is that we are not interested in solving the optimization problem. Instead, we are more interested in sampling from a distribution. By introducing the random noise to the gradient descent step, we randomly pick a sample that is following the objective function’s trajectory while not staying at where it is. If we are closer to the peak, we will move left and right slightly. If we are far from the peak, the gradient direction will pull us towards the peak. If the curvature around the peak is sharp, we will concentrate most of the steady state points there. If the curvature around the peak is flat, we will spread around. Therefore, by repeatedly initialize the stochastic gradient descent algorithm at a uniformly distributed location, we will eventually collect samples that will follow the distribution we designate.
Example. Consider a Gaussian mixture ) +
). We can numerically calculate
log p(x). For demonstration, we choose
2,
2. Suppose we initialize
Uniform[
3]. We run Langevin updates for t = 100 steps. The histograms of generated samples are shown in the figures below.
3.2 (Stein’s) Score Function
The second component of the Langevin dynamics equation is the gradient log p(x). It has a formal name known as the
, denoted by
We should be careful not to confuse Stein’s score function with the ordinary score function which is defined as
The ordinary score function is the gradient (wrt ) of the log-likelihood. In contrast, the Stein’s score function is the gradient wrt the data point x. Maximum likelihood estimation uses the ordinary score function, whereas Langevin dynamics uses Stein’s score function. However, since most people in the diffusion literature calls Stein’s score function as the score function, we follow this culture.
The “score function” in Langevin dynamics is more accurately known as the Stein’s score function.
The way to understand the score function is to remember that it is the gradient with respect to the data x. For any high-dimensional distribution p(x), the gradient will give us vector field
Let’s consider two examples.
The probability density function and the corresponding score function of the above two examples are shown in Figure 18.
Figure 18: Examples of score functions
• The magnitude of the vectors are the strongest at places where the change of log p(x) is the biggest. Therefore, in regions where log p(x) is close to the peak will be mostly very weak gradient.
• The vector field indicates how a data point should travel in the contour map. In Figure 19 we show the contour map of a Gaussian mixture (with two Gaussians). We draw arrows to indicate the vector field. Now if we consider a data point living the space, the Langevin dynamics equation will basically move the data points along the direction pointed by the vector field towards the basin.
• In physics, the score function is equivalent to the “drift”. This name suggests how the diffusion particles should flow to the lowest energy state.
Figure 19: The contour map of the score function, and the corresponding trajectory of two samples.
3.3 Score Matching Techniques
The most difficult question in Langevin dynamics is how to obtain ) because we have no access to p(x). Let’s recall the definition of the (Stein’s) score function
where we put a subscript to denote that
will be implemented via a network. Since the right hand side of the above equation is not known, we need some cheap and dirty ways to approximate it. In this section, we briefly discuss two approximation.
Explicit Score-Matching. Suppose that we are given a dataset . The solution people came up with is to consider the classical kernel density estimation by defining a distribution
where h is just some hyperparameter for the kernel function ), and
is the m-th sample in the training set. Figure 20 illustrates the idea of kernel density estimation. In the cartoon figure shown on the left, we show multiple kernels
) centered at different data points
. The sum of all these individual kernels gives us the overall kernel density estimate q(x). On the right hand side we show a real histogram and the corresponding kernel density estimate. We remark that q(x) is at best an approximation to the true data distribution p(x) which is never known.
Figure 20: Illustration of kernel density estimation.
Since q(x) is an approximation to p(x) which is never accessible, we can learn ) based on q(x). This leads to the following definition of the a loss function which can be used to train a network.
By substituting the kernel density estimation, we can show that the loss is
So, we have derived a loss function that can be used to train the network. Once we train the network , we can replace it in the Langevin dynamics equation to obtain the recursion:
The issue of explicit score matching is that the kernel density estimation is a fairly poor non-parameter estimation of the true distribution. Especially when we have limited number of samples and the samples live in a high dimensional space, the kernel density estimation performance can be poor.
Denoising Score Matching. Given the potential drawbacks of explicit score matching, we now introduce a more popular score matching known as the denoising score matching (DSM). In DSM, the loss function is defined as follows.
The key difference here is that we replace the distribution q(x) by a conditional distribution ). The former requires an approximation, e.g., via kernel density estimation, whereas the latter does not. Here is an example.
In the special case where ), we can let
+
. This will give us
As a result, the loss function of the denoising score matching becomes
If we replace the dummy variable by x, and we note that sampling from q(x) can be replaced by sampling from p(x) when we are given a training dataset, we can conclude the following.
The beauty about Eqn (85) is that it is highly interpretable. The quantity is effectively adding noise
to a clean image x. The score function
is supposed to take this noisy image and predict the noise
. Predicting noise is equivalent to denoising, because any denoised image plus the predicted noise will give us the noisy observation. Therefore, Eqn (85) is a denoising step. Figure 21 illustrates the training procedure of the score function
).
Figure 21: Training of for denoising score matching. The network
is trained to estimate the noise.
The training step can simply described as follows: You give us a training dataset , we train a network
with the goal to
The bigger question here is why Eqn (84) would even make sense in the first place. This needs to be answered through the equivalence between the explicit score matching loss and the denoising score matching loss.
The equivalence between the explicit score matching and the denoising score matching is a major discovery. The proof below is based on the original work of Vincent 2011.
For inference, we assume that we have already trained the score estimator . To generate an image, we perform the following procedure for t = 1, . . . , T:
Congratulations! We are done. This is all about Score-based Generative Models.
Additional readings about score-matching should start with Vincent’s technical report [9]. A very popular paper in the recent literature is Song and Ermon [15], their follow up work [16], and [8]. In practice, training a score function requires a noise schedule by considering a sequence of noise levels. We will briefly discuss this when we explain the variance exploding SDE in the next section.
Thus far we have derived the diffusion iterations via the DDPM and the SMLD perspective. In this section, we will introduce a third perspective through the lens of differential equation. It may not be obvious why our iterative schemes suddenly become the complicated differential equations. So, before we derive any equation, we should briefly discuss how differential equations can be relevant to us.
4.1 Motivating Examples
What we observe in this motivating example are a two interesting facts:
• The discrete-time iterative scheme can be written as a continuous-time ordinary differential equation.
• For simple ODEs, we can write down the analytic solution in closed form. More complicated ODE would be hard to write an analytic solution. But we can still use ODE tools to analyze the behavior of the solution. We can also derive the limiting solution 0.
Forward and Backward Updates.
Let’s use the gradient descent example to illustrate one more aspect of the ODE. Going back to Eqn (92), we recognize that the recursion can be written equivalently as (assuming ):
where the continuous equation holds when we set ∆0 and ∆
0. The interesting point about this equality is that it gives us a summary of the update ∆x by writing it in terms of dt. It says that if we move the along the time axis by dt, then the solution x will be updated by dx.
Eqn (95) defines the relationship between changes. If we consider a sequence of iterates i = 1, 2, . . . , N, and if we are told that the progression of the iterates follows Eqn (95), then we can write
We call this as the forward equation because we update x by x + ∆x assuming that + ∆t. Now, consider a sequence of iterates
1. If we are told that the progression of the iterates follows Eqn (95), then the time-reversal iterates will be
Note the change in sign when reversing the progression direction. We call this the reverse equation.
4.2 Forward and Backward Iterations in SDE
The concept of differential equation for diffusion is not too far from the above gradient descent algorithm. If we introduce a noise term ) to the gradient descent algorithm, then the ODE will become a stochastic differential equation (SDE). To see this, we just follow the same discretization scheme by defining x(t) as a continuous function for 0
1. Suppose that there are N steps in the interval so that the interval [0, 1] can be divided into a sequence
. The discretization will give us
), and
(
). The interval step is ∆
, and the set of all t’s is
. Using these definitions, we can write
Now, let’s define a random process w(t) such that z(t) = w(t + ∆∆t for a very small ∆t. In computation, we can generate such a w(t) by integrating z(t) (which is a Wiener process). With w(t) defined, we can write
The equation above reveals a generic form of the SDE. We summarize it as follows.
The two terms f(x, t) and g(t) carry physical meaning. The draft coefficient is a vector-valued function f(x, t) defining how molecules in a closed system would move in the absence of random effects. For the gradient descent algorithm, the drift is defined by the negative gradient of the objective function. That is, we want the solution trajectory to follow the gradient of the objective.
The diffusion coefficient g(t) is a scalar function describing how the molecules would randomly walk from one position to another. The function g(t) determines how strong the random movement is.
We can plot the function as below. The initial point
process is moving forward in time.
Remark. As you can see, the differential is defined as the Wiener process which is a white Gaussian vector. The individual
is not a Gaussian, but the difference
is a Gaussian.
The reverse direction of the diffusion equation is to move backward in time. The reverse-time SDE, according to Anderson [17], is given as follows.
where ) is the probability distribution of x at time t, and w is the Wiener process when time flows backward.
4.3 Stochastic Differential Equation for DDPM
In order to draw the connection between DDPM and SDE, we consider the discrete-time DDPM iteration. For i = 1, 2, . . . , N:
We can show that this equation can be derived from the forward SDE equation below.
To see why this is the case, we define a step size ∆, and consider an auxiliary noise level
where
. Then
where we assume that in the ) which is a continuous time function for 0
1. Similarly, we define
Hence, we have
Thus, as ∆0, we have
Therefore, we showed that the DDPM forward update iteration can be equivalently written as an SDE.
Being able to write the DDPM forward update iteration as an SDE means that the DDPM estimates can be determined by solving the SDE. In other words, for an appropriately defined SDE solver, we can throw the SDE into the solver. The solution returned by an appropriately chosen solver will be the DDPM estimate. Of course, we are not required to use the SDE solver because the DDPM iteration itself is solving the SDE. It may not be the best SDE solver because the DDPM iteration is only a first order method. Nevertheless, if we are not interested in using the SDE solver, we can still use the DDPM iteration to obtain a solution. Here is an example.
Example. Consider the DDPM forward equation with 05 for all
1. We initialize the sample
by drawing it from a Gaussian mixture such that
The reverse diffusion equation follows from Eqn (97) by substituting the appropriate quantities: f(x, t) = and
). This will give us
which will give us the following equation:
The iterative update scheme can be written by considering ), and
). Then, letting
, we can show that
By grouping the terms, and assuming that 1, we recognize that
Then, following the discretization scheme by letting , ∆
, and
, we can show that
where ) is the probability density function of x at time i. For practical implementation, we can replace
log
) by the estimated score function
).
So, we have recovered the DDPM iteration that is consistent with the one defined by Song and Ermon in [8]. This is an interesting result, because it allows us to connect DDPM’s iteration using the score function. Song and Ermon [8] called the SDE an variance preserving (VP) SDE.
4.4 Stochastic Differential Equation for SMLD
The score-matching Langevin Dynamics model can also be described by an SDE. To start with, we notice that in the SMLD setting, there isn’t really a “forward diffusion step”. However, we can roughly argue that if we divide the noise scale in the SMLD training into N levels, then the recursion should follow a Markov chain
This is not too hard to see. If we assume that the variance of is
, then we can show that
Therefore, given a sequence of noise levels, Eqn (104) will indeed generate estimates such that the noise statistics will satisfy the desired property.
If we agree Eqn (104), it is easy to derive the SDE associated with Eqn (104). Assuming that in the limit becomes the continuous time
) for 0
1, and
becomes x(t) where
) if we let
. Then we have
At the limit when ∆0, the equation converges to
We summarize our result as follows.
Mapping this to Eqn (96), we recognize that
As a result, if we write the reverse equation Eqn (97), we should have
This will give us the following reverse equation:
For the discrete-time iterations, we first define . Then, using the same set of discretization setups as the DDPM case, we can show that
which is identical to the SMLD reverse update equation. Song and Ermon [8] called the SDE an variance exploding (VE) SDE.
4.5 Solving SDE
In this subsection, we briefly discuss how the differential equations are solved numerically. To make our discussion slightly easier, we shall focus on the ODE. Consider the following ODE
If the ODE is a scalar ODE, then the ODE is ).
Euler Method. Euler method is a first order numerical method for solving the ODE. Given f(x(t), t), and
, Euler method solves the problem via an iterative scheme for
1 such that
where is the step size. Let’s consider a simple example.
Runge-Kutta (RK) Method. Another popularly used ODE solver is the Runge-Kutta (RK) method. The classical RK-4 algorithm solves the ODE via the iteration
where the quantities and
are defined as
For more details, you can consult numerical methods textbooks such as [18].
Predictor-Corrector Algorithm. Since different numerical solvers have different behavior in terms of the error of approximation, throwing the ODE (or SDE) into an off-the-shelf numerical solver will result in various degrees of error [19]. However, if we are specifically trying to solve the reverse diffusion equation, it is possible to use techniques other than numerical ODE/SDE solvers to make the appropriate corrections, as illustrated in Figure 22.
Figure 22: Prediction and correction algorithm.
Let’s use DDPM as an example. In DDPM, the reverse diffusion equation is given by
We can consider it as an Euler method for the reverse diffusion. However, if we have already trained the score function ), we can run the score-matching equation, i.e.,
for M times to make the correction. Algorithm 4.5 summarizes the idea. (Note that we have replaced the score function by the estimate.) For the SMLD algorithm, the two equations are:
We can pair them up as in the case of DDPM’s prediction-correction algorithm by repeating the correction iteration a few times.
Accelerate the SDE Solver. While generic ODE solvers can be used to solve the ODE, the forward and reverse diffusion equations we encounter are quite special. In fact, they take the form of
for some choice of functions a(t) and b(t), with the initial condition . This is not a complicated ODE. It is just a first order ODE. In [20], Lu et al. observed that because of the special structure of the ODE (they called the semi-linear structure), it is possible to separately handle a(t)x(t) and b(t). To understand how things work, we use a textbook result shown below.
We can further simplify the second term above by noticing that
Of particular interest as presented in [20] was the reverse diffusion equation derived from [8]:
where , and
2
. Using the Variation of Constants Theorem, we can solve the ODE exactly at time t by the formula
Then, by defining ), and with additional simplifications outlined in [20], this equation can be simplified to
To evaluate this equation, one just needs to run a numerical integrator for the integration shown on the right hand side. Of course, there are other numerical acceleration methods for solving ODEs which we shall skip for brevity.
Some of you may wonder: Why do we want to map the iterative schemes to differential equations? There are several reasons, some are legitimate whereas some are speculative.
• By unifying multiple diffusion models to the same SDE framework, one can compare the algorithms. In some cases, one can improve the numerical scheme by borrowing ideas from the SDE literature as well as the probabilistic sampling literature. For example, the predictor-corrector scheme in [8] was a hybrid SDE solver coupled with a Markov Chain Monte Carlo.
• Mapping the diffusion iterations to SDE, according to some papers such as [22], offers more design flexibility.
• Outside the context diffusion algorithms, in general stochastic gradient descent algorithms have corresponding SDE such as the Fokker-Planck equations. People have demonstrated how to theoretically analyze the limiting distribution of the estimates, in exact closed-form. This alleviates the difficulty of analyzing the random algorithm by means of analyzing a well-defined limiting distribution.
This tutorial covers a few basic concepts underpinning the development of diffusion-based generative models in the recent literature. Given the sheer volume (and quickly expanding) of literature, we find it particularly important to describe the fundamental ideas instead of recycling the Python demos. A few lessons we learned from writing this tutorial are:
• The same diffusion idea can be derived independently from multiple perspectives, namely VAE, DDPM, SMLD, and SDE. There is no particular reason why one is more superior/inferior than the other although some may argue differently.
• The main reason why denoising diffusion works is because of its small increment which was not realized in the age of GAN and VAE.
• Although iterative denoising is the current state-of-the-art, the approach itself does not appear to be the ultimate solution. Humans do not generate images from pure noise. Moreover, because of the small increment nature of diffusion models, speed will continue to be a major hurdle although some efforts in knowledge distillation have been made to improve the situation.
• Some questions regarding generating noise from non-Gaussian might require justification. If the whole reason of introducing Gaussian distribution is to make the derivations easier, why should we switch to another type of noise by making our lives harder?
• Application of diffusion models to inverse problem is readily available. For any existing inverse solvers such as the Plug-and-Play ADMM algorithm, we can replace the denoiser by an explicit diffusion sampler. People have demonstrated improved image restoration results based on this approach.
[1] D. P. Kingma and M. Welling, “An introduction to variational autoencoders,” Foundations and Trends in Machine Learning, vol. 12, no. 4, pp. 307–392, 2019. https://arxiv.org/abs/1906.02691.
[2] C. Doersch, “Tutorial on variational autoencoders,” 2016. https://arxiv.org/abs/1606.05908.
[3] D. P. Kingma and M. Welling, “Auto-encoding variational Bayes,” in ICLR, 2014. https:// openreview.net/forum?id=33X9fd2-9FyZd.
[4] J. Ho, A. Jain, and P. Abbeel, “Denoising diffusion probabilistic models,” in NeurIPS, 2020. https: //arxiv.org/abs/2006.11239.
[5] D. P. Kingma, T. Salimans, B. Poole, and J. Ho, “Variational diffusion models,” in NeurIPS, 2021. https://arxiv.org/abs/2107.00630.
[6] M. Delbracio and P. Milanfar, “Inversion by direct iteration: An alternative to denoising diffusion for image restoration,” Transactions on Machine Learning Research, 2023. https://openreview.net/ forum?id=VmyFF5lL3F.
[7] S. H. Chan, Introduction to Probability for Data Science. Michigan Publishing, 2021. https: //probability4datascience.com/.
[8] Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole, “Score-based generative modeling through stochastic differential equations,” in ICLR, 2021. https://openreview.net/forum? id=PxTIG12RRHS.
[9] P. Vincent, “A connection between score matching and denoising autoencoders,” Neural Computation, vol. 23, no. 7, pp. 1661–1674, 2011. https://www.iro.umontreal.ca/~vincentp/Publications/ smdae_techreport.pdf.
[10] J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli, “Deep unsupervised learning using nonequilibrium thermodynamics,” in ICML, vol. 37, pp. 2256–2265, 2015. https://arxiv.org/abs/ 1503.03585.
[11] C. Luo, “Understanding diffusion models: A unified perspective,” 2022. https://arxiv.org/abs/ 2208.11970.
[12] J. Song, C. Meng, and S. Ermon, “Denoising diffusion implicit models,” in ICLR, 2023. https:// openreview.net/forum?id=St1giarCHLP.
[13] R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer, “High-resolution image synthesis with latent diffusion models,” in CVPR, pp. 10684–10695, 2022. https://arxiv.org/abs/2112.10752.
[14] C. Saharia, W. Chan, S. Saxena, L. Li, J. Whang, E. L. Denton, K. Ghasemipour, R. Gontijo Lopes, B. Karagol Ayan, T. Salimans, J. Ho, D. J. Fleet, and M. Norouzi, “Photorealistic text-to-image diffusion models with deep language understanding,” in NeurIPS, vol. 35, pp. 36479–36494, 2022. https: //arxiv.org/abs/2205.11487.
[15] Y. Song and S. Ermon, “Generative modeling by estimating gradients of the data distribution,” in NeurIPS, 2019. https://arxiv.org/abs/1907.05600.
[16] Y. Song and S. Ermon, “Improved techniques for training score-based generative models,” in NeurIPS, 2020. https://arxiv.org/abs/2006.09011.
[17] B. Anderson, “Reverse-time diffusion equation models,” Stochastic Process. Appl., vol. 12, pp. 313–326, May 1982. https://www.sciencedirect.com/science/article/pii/0304414982900515.
[18] K. Atkinson, W. Han, and D. Stewart, Numerical solution of ordinary differential equations. Wiley, 2009. https://homepage.math.uiowa.edu/~atkinson/papers/NAODE_Book.pdf.
[19] T. Karras, M. Aittala, T. Aila, and S. Laine, “Elucidating the design space of diffusion-based generative models,” in NeurIPS, 2022. https://arxiv.org/abs/2206.00364.
[20] C. Lu, Y. Zhou, F. Bao, J. Chen, C. Li, and J. Zhu, “DPM-Solver: A fast ODE solver for diffusion probabilistic model sampling in around 10 steps,” in NeurIPS, 2022. https://arxiv.org/abs/2206. 00927.
[21] G. Nagy, “MTH 235 differential equations,” 2024. https://users.math.msu.edu/users/gnagy/ teaching/ade.pdf.
[22] M. S. Albergo, N. M. Boffi, and E. Vanden-Eijnden, “Stochastic interpolants: A unifying framework for flows and diffusions.” https://arxiv.org/abs/2303.08797.