PAPER KEY: HWSLN6S2
TITLE: Understanding Diffusion Models: A Unified Perspective
AUTHORS: Luo, Calvin

arXiv:2208.11970v1 [cs.LG] 25 Aug 2022

Understanding Diﬀusion Models: A Uniﬁed Perspective
Calvin Luo Google Research, Brain Team
calvinluo@google.com
August 26, 2022
Contents
Introduction: Generative Models . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 1 Background: ELBO, VAE, and Hierarchical VAE . . . . . . . . . . . . . . . . . . . . . . . . 2
Evidence Lower Bound . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 2 Variational Autoencoders . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4 Hierarchical Variational Autoencoders . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 5 Variational Diﬀusion Models . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 6 Learning Diﬀusion Noise Parameters . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 14 Three Equivalent Interpretations . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 15 Score-based Generative Models . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 17 Guidance . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 20 Classiﬁer Guidance . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 21 Classiﬁer-Free Guidance . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 21 Closing . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 22
Introduction: Generative Models
Given observed samples x from a distribution of interest, the goal of a generative model is to learn to model its true data distribution p(x). Once learned, we can generate new samples from our approximate model at will. Furthermore, under some formulations, we are able to use the learned model to evaluate the likelihood of observed or sampled data as well. There are several well-known directions in current literature, that we will only introduce brieﬂy at a high level. Generative Adversarial Networks (GANs) model the sampling procedure of a complex distribution, which is learned in an adversarial manner. Another class of generative models, termed "likelihood-based", seeks to learn a model that assigns a high likelihood to the observed data samples. This includes autoregressive models, normalizing ﬂows, and Variational Autoencoders (VAEs). Another similar approach is energy-based modeling, in which a distribution is learned as an arbitrarily ﬂexible energy function that is then normalized.
1

Score-based generative models are highly related; instead of learning to model the energy function itself, they learn the score of the energy-based model as a neural network. In this work we explore and review diﬀusion models, which as we will demonstrate, have both likelihood-based and score-based interpretations. We showcase the math behind such models in excruciating detail, with the aim that anyone can follow along and understand what diﬀusion models are and how they work.

Background: ELBO, VAE, and Hierarchical VAE
For many modalities, we can think of the data we observe as represented or generated by an associated unseen latent variable, which we can denote by random variable z. The best intuition for expressing this idea is through Plato’s Allegory of the Cave. In the allegory, a group of people are chained inside a cave their entire life and can only see the two-dimensional shadows projected onto a wall in front of them, which are generated by unseen three-dimensional objects passed before a ﬁre. To such people, everything they observe is actually determined by higher-dimensional abstract concepts that they can never behold.
Analogously, the objects that we encounter in the actual world may also be generated as a function of some higher-level representations; for example, such representations may encapsulate abstract properties such as color, size, shape, and more. Then, what we observe can be interpreted as a three-dimensional projection or instantiation of such abstract concepts, just as what the cave people observe is actually a two-dimensional projection of three-dimensional objects. Whereas the cave people can never see (or even fully comprehend) the hidden objects, they can still reason and draw inferences about them; in a similar way, we can approximate latent representations that describe the data we observe.
Whereas Plato’s Allegory illustrates the idea behind latent variables as potentially unobservable representations that determine observations, a caveat of this analogy is that in generative modeling, we generally seek to learn lower-dimensional latent representations rather than higher-dimensional ones. This is because trying to learn a representation of higher dimension than the observation is a fruitless endeavor without strong priors. On the other hand, learning lower-dimensional latents can also be seen as a form of compression, and can potentially uncover semantically meaningful structure describing observations.

Evidence Lower Bound

Mathematically, we can imagine the latent variables and the data we observe as modeled by a joint distribution p(x, z). Recall one approach of generative modeling, termed "likelihood-based", is to learn a model to maximize the likelihood p(x) of all observed x. There are two ways we can manipulate this joint distribution to recover the likelihood of purely our observed data p(x); we can explicitly marginalize out the latent variable z:

p(x) = p(x, z)dz

(1)

or, we could also appeal to the chain rule of probability:

p(x, z) p(x) =

(2)

p(z|x)

Directly computing and maximizing the likelihood p(x) is diﬃcult because it either involves integrating out all latent variables z in Equation 1, which is intractable for complex models, or it involves having access to a ground truth latent encoder p(z|x) in Equation 2. However, using these two equations, we can derive a term called the Evidence Lower Bound (ELBO), which as its name suggests, is a lower bound of the evidence. The evidence is quantiﬁed in this case as the log likelihood of the observed data. Then, maximizing the ELBO becomes a proxy objective with which to optimize a latent variable model; in the best case, when the ELBO is powerfully parameterized and perfectly optimized, it becomes exactly equivalent to the evidence. Formally, the equation of the ELBO is:

p(x, z)

Eqφ (z |x)

log qφ(z|x)

(3)

2

To make the relationship with the evidence explicit, we can mathematically write:

p(x, z)

log p(x) ≥ Eqφ(z|x)

log qφ(z|x)

(4)

Here, qφ(z|x) is a ﬂexible approximate variational distribution with parameters φ that we seek to optimize. Intuitively, it can be thought of as a parameterizable model that is learned to estimate the true distribution over latent variables for given observations x; in other words, it seeks to approximate true posterior p(z|x). As we will see when exploring the Variational Autoencoder, as we increase the lower bound by tuning the parameters φ to maximize the ELBO, we gain access to components that can be used to model the true data distribution and sample from it, thus learning a generative model. For now, let us try to dive deeper into why the ELBO is an objective we would like to maximize.

Let us begin by deriving the ELBO, using Equation 1:

log p(x) = log p(x, z)dz

(Apply Equation 1)

(5)

= log p(x, z)qφ(z|x) dz qφ(z|x)

p(x, z) = log Eqφ(z|x) qφ(z|x)

p(x, z)

≥ Eqφ(z|x)

log qφ(z|x)

(Multiply by 1 = qφ(z|x) )

(6)

qφ(z|x)

(Deﬁnition of Expectation)

(7)

(Apply Jensen’s Inequality)

(8)

In this derivation, we directly arrive at our lower bound by applying Jensen’s Inequality. However, this does not supply us much useful information about what is actually going on underneath the hood; crucially, this proof gives no intuition on exactly why the ELBO is actually a lower bound of the evidence, as Jensen’s Inequality handwaves it away. Furthermore, simply knowing that the ELBO is truly a lower bound of the data does not really tell us why we want to maximize it as an objective. To better understand the relationship between the evidence and the ELBO, let us perform another derivation, this time using Equation 2:

log p(x) = log p(x) qφ(z|x)dz

(Multiply by 1 = qφ(z|x)dz) (9)

= qφ(z|x)(log p(x))dz

(Bring evidence into integral) (10)

= Eqφ(z|x) [log p(x)]

p(x, z)

= Eqφ(z|x)

log p(z|x)

(Deﬁnition of Expectation)

(11)

(Apply Equation 2)

(12)

= Eqφ(z|x)

log p(x, z)qφ(z|x) p(z|x)qφ(z|x)

(Multiply by 1 = qφ(z|x) )

(13)

qφ(z|x)

= Eqφ(z|x)

p(x, z) log
qφ(z|x)

+ Eqφ(z|x)

log qφ(z|x) p(z|x)

(Split the Expectation)

(14)

p(x, z)

= Eqφ(z|x)

log qφ(z|x)

+ DKL(qφ(z|x)

p(z|x))

(Deﬁnition of KL Divergence)

(15)

p(x, z)

≥ Eqφ(z|x)

log qφ(z|x)

(KL Divergence always ≥ 0)

(16)

From this derivation, we clearly observe from Equation 15 that the evidence is equal to the ELBO plus the KL Divergence between the approximate posterior qφ(z|x) and the true posterior p(z|x). In fact, it was this KL Divergence term that was magically removed by Jensen’s Inequality in Equation 8 of the ﬁrst derivation. Understanding this term is the key to understanding not only the relationship between the ELBO and the evidence, but also the reason why optimizing the ELBO is an appropriate objective at all.

Firstly, we now know why the ELBO is indeed a lower bound: the diﬀerence between the evidence and the ELBO is a strictly non-negative KL term, thus the value of the ELBO can never exceed the evidence.

3

Figure 1: A Variational Autoencoder graphically represented. Here, encoder q(z|x) deﬁnes a distribution over latent variables z for observations x, and p(x|z) decodes latent variables into observations.

Secondly, we explore why we seek to maximize the ELBO. Having introduced latent variables z that we would like to model, our goal is to learn this underlying latent structure that describes our observed data. In other words, we want to optimize the parameters of our variational posterior qφ(z|x) to exactly match the true posterior distribution p(z|x), which is achieved by minimizing their KL Divergence (ideally to zero). Unfortunately, it is intractable to minimize this KL Divergence term directly, as we do not have access to the ground truth p(z|x) distribution. However, notice that on the left hand side of Equation 15, the likelihood of our data (and therefore our evidence term log p(x)) is always a constant with respect to φ, as it is computed by marginalizing out all latents z from the joint distribution p(x, z) and does not depend on φ whatsoever. Since the ELBO and KL Divergence terms sum up to a constant, any maximization of the ELBO term with respect to φ necessarily invokes an equal minimization of the KL Divergence term. Thus, the ELBO can be maximized as a proxy for learning how to perfectly model the true latent posterior distribution; the more we optimize the ELBO, the closer our approximate posterior gets to the true posterior. Additionally, once trained, the ELBO can be used to estimate the likelihood of observed or generated data as well, since it is learned to approximate the model evidence log p(x).

Variational Autoencoders
In the default formulation of the Variational Autoencoder (VAE) [1], we directly maximize the ELBO. This approach is variational, because we optimize for the best qφ(z|x) amongst a family of potential posterior distributions parameterized by φ. It is called an autoencoder because it is reminiscent of a traditional autoencoder model, where input data is trained to predict itself after undergoing an intermediate bottlenecking representation step. To make this connection explicit, let us dissect the ELBO term further:

p(x, z)

Eqφ (z |x)

log qφ(z|x)

= Eqφ(z|x)

log pθ(x|z)p(z) qφ(z|x)

= Eqφ(z|x) [log pθ(x|z)] + Eqφ(z|x)

p(z) log
qφ(z|x)

= Eqφ(z|x) [log pθ(x|z)] − DKL(qφ(z|x) p(z))

reconstruction term

prior matching term

(Chain Rule of Probability) (17)

(Split the Expectation)

(18)

(Deﬁnition of KL Divergence) (19)

In this case, we learn an intermediate bottlenecking distribution qφ(z|x) that can be treated as an encoder ; it transforms inputs into a distribution over possible latents. Simultaneously, we learn a deterministic function pθ(x|z) to convert a given latent vector z into an observation x, which can be interpreted as a decoder.
The two terms in Equation 19 each have intuitive descriptions: the ﬁrst term measures the reconstruction likelihood of the decoder from our variational distribution; this ensures that the learned distribution is modeling eﬀective latents that the original data can be regenerated from. The second term measures how similar the learned variational distribution is to a prior belief held over latent variables. Minimizing this term encourages the encoder to actually learn a distribution rather than collapse into a Dirac delta function. Maximizing the ELBO is thus equivalent to maximizing its ﬁrst term and minimizing its second term.

4

A deﬁning feature of the VAE is how the ELBO is optimized jointly over parameters φ and θ. The encoder of the VAE is commonly chosen to model a multivariate Gaussian with diagonal covariance, and the prior is often selected to be a standard multivariate Gaussian:

qφ(z|x) = N (z; µφ(x), σφ2 (x)I)

(20)

p(z) = N (z; 0, I)

(21)

Then, the KL divergence term of the ELBO can be computed analytically, and the reconstruction term can be approximated using a Monte Carlo estimate. Our objective can then be rewritten as:

L

arg max Eqφ(z|x) [log pθ(x|z)] − DKL(qφ(z|x) p(z)) ≈ arg max log pθ(x|z(l)) − DKL(qφ(z|x) p(z)) (22)

φ,θ

φ,θ l=1

where latents {z(l)}Ll=1 are sampled from qφ(z|x), for every observation x in the dataset. However, a problem arises in this default setup: each z(l) that our loss is computed on is generated by a stochastic sampling procedure, which is generally non-diﬀerentiable. Fortunately, this can be addressed via the reparameterization trick when qφ(z|x) is designed to model certain distributions, including the multivariate Gaussian.
The reparameterization trick rewrites a random variable as a deterministic function of a noise variable; this allows for the optimization of the non-stochastic terms through gradient descent. For example, samples from a normal distribution x ∼ N (x; µ, σ2) with arbitrary mean µ and variance σ2 can be rewritten as:

x = µ + σ with ∼ N ( ; 0, I)

In other words, arbitrary Gaussian distributions can be interpreted as standard Gaussians (of which is a sample) that have their mean shifted from zero to the target mean µ by addition, and their variance stretched by the target variance σ2. Therefore, by the reparameterization trick, sampling from an arbitrary Gaussian distribution can be performed by sampling from a standard Gaussian, scaling the result by the target standard deviation, and shifting it by the target mean.
In a VAE, each z is thus computed as a deterministic function of input x and auxiliary noise variable :

z = µφ(x) + σφ(x)

with ∼ N ( ; 0, I)

where represents an element-wise product. Under this reparameterized version of z, gradients can then be computed with respect to φ as desired, to optimize µφ and σφ. The VAE therefore utilizes the reparameterization trick and Monte Carlo estimates to optimize the ELBO jointly over φ and θ.
After training a VAE, generating new data can be performed by sampling directly from the latent space p(z) and then running it through the decoder. Variational Autoencoders are particularly interesting when the dimensionality of z is less than that of input x, as we might then be learning compact, useful representations. Furthermore, when a semantically meaningful latent space is learned, latent vectors can be edited before being passed to the decoder to more precisely control the data generated.

Hierarchical Variational Autoencoders
A Hierarchical Variational Autoencoder (HVAE) [2, 3] is a generalization of a VAE that extends to multiple hierarchies over latent variables. Under this formulation, latent variables themselves are interpreted as generated from other higher-level, more abstract latents. Intuitively, just as we treat our three-dimensional observed objects as generated from a higher-level abstract latent, the people in Plato’s cave treat threedimensional objects as latents that generate their two-dimensional observations. Therefore, from the perspective of Plato’s cave dwellers, their observations can be treated as modeled by a latent hierarchy of depth two (or more).
Whereas in the general HVAE with T hierarchical levels, each latent is allowed to condition on all previous latents, in this work we focus on a special case which we call a Markovian HVAE (MHVAE). In a MHVAE, the generative process is a Markov chain; that is, each transition down the hierarchy is Markovian, where

5

Figure 2: A Markovian Hierarchical Variational Autoencoder with T hierarchical latents. The generative process is modeled as a Markov chain, where each latent zt is generated only from the previous latent zt+1.

decoding each latent zt only conditions on previous latent zt+1. Intuitively, and visually, this can be seen as simply stacking VAEs on top of each other, as depicted in Figure 2; another appropriate term describing this model is a Recursive VAE. Mathematically, we represent the joint distribution and the posterior of a Markovian HVAE as:

T
p(x, z1:T ) = p(zT )pθ(x|z1) pθ(zt−1|zt)
t=2 T
qφ(z1:T |x) = qφ(z1|x) qφ(zt|zt−1)
t=2

(23) (24)

Then, we can easily extend the ELBO to be:

log p(x) = log p(x, z1:T )dz1:T

(Apply Equation 1)

(25)

= log

p(x,

z1:T )qφ qφ(z1:T

(z1:T |x)

|x)

dz1:T

= log Eqφ(z1:T |x)

p(x, z1:T ) qφ(z1:T |x)

≥ Eqφ(z1:T |x)

log p(x, z1:T ) qφ(z1:T |x)

(Multiply by 1 = qφ(z1:T |x) )

(26)

qφ(z1:T |x)

(Deﬁnition of Expectation)

(27)

(Apply Jensen’s Inequality)

(28)

We can then plug our joint distribution (Equation 23) and posterior (Equation 24) into Equation 28 to produce an alternate form:

Eqφ(z1:T |x)

log p(x, z1:T ) qφ(z1:T |x)

= Eqφ(z1:T |x)

log

p(zT )pθ(x|z1)

T t=2

pθ

(zt−1|zt)

qφ(z1|x)

T t=2

qφ(zt

|zt−1

)

(29)

As we will show below, when we investigate Variational Diﬀusion Models, this objective can be further decomposed into interpretable components.

Variational Diﬀusion Models
The easiest way to think of a Variational Diﬀusion Model (VDM) [4, 5, 6] is simply as a Markovian Hierarchical Variational Autoencoder with three key restrictions:
• The latent dimension is exactly equal to the data dimension
• The structure of the latent encoder at each timestep is not learned; it is pre-deﬁned as a linear Gaussian model. In other words, it is a Gaussian distribution centered around the output of the previous timestep
• The Gaussian parameters of the latent encoders vary over time in such a way that the distribution of the latent at ﬁnal timestep T is a standard Gaussian

6

Figure 3: A visual representation of a Variational Diﬀusion Model; x0 represents true data observations such as natural images, xT represents pure Gaussian noise, and xt is an intermediate noisy version of x0. Each q(xt|xt−1) is modeled as a Gaussian distribution that uses the output of the previous state as its mean.

Furthermore, we explicitly maintain the Markov property between hierarchical transitions from a standard Markovian Hierarchical Variational Autoencoder.

Let us expand on the implications of these assumptions. From the ﬁrst restriction, with some abuse of notation, we can now represent both true data samples and latent variables as xt, where t = 0 represents true data samples and t ∈ [1, T ] represents a corresponding latent with hierarchy indexed by t. The VDM posterior is the same as the MHVAE posterior (Equation 24), but can now be rewritten as:

T
q(x1:T |x0) = q(xt|xt−1)
t=1

(30)

From the second assumption, we know that the distribution of each latent variable in the encoder is a

Gaussian centered around its previous hierarchical latent. Unlike a Mark