Diffusion model explained

In machine learning, diffusion models, also known as diffusion probabilistic models or score-based generative models, are a class of latent variable generative models. A diffusion model consists of three major components: the forward process, the reverse process, and the sampling procedure.[1] The goal of diffusion models is to learn a diffusion process for a given dataset, such that the process can generate new elements that are distributed similarly as the original dataset. A diffusion model models data as generated by a diffusion process, whereby a new datum performs a random walk with drift through the space of all possible data. A trained diffusion model can be sampled in many ways, with different efficiency and quality.

There are various equivalent formalisms, include Markov chains, denoising diffusion probabilistic models, noise conditioned score networks, and stochastic differential equations.[2] They are typically trained using variational inference. The model responsible for denoising is typically called its "backbone". The backbone may be of any kind, but they are typically U-nets or transformers.

, diffusion models are mainly used for computer vision tasks, including image denoising, inpainting, super-resolution, image generation, and video generation. These typically involves training a neural network to sequentially denoise images blurred with Gaussian noise.[3] [4] The model is trained to reverse the process of adding noise to an image. After training to convergence, it can be used for image generation by starting with an image composed of random noise, and applying the network iteratively to denoise the image.

Diffusion-based image generators have seen widespread commercial interest, such as Stable Diffusion and DALL-E. These models typically combine diffusion models with other models, such as text-encoders and cross-attention modules to allow text-conditioned generation.

Other than computer vision, diffusion models have also found applications in natural language processing[5] such as text generation[6] [7] and summarization[8], sound generation[9], and reinforcement learning.

Denoising diffusion model

Non-equilibrium thermodynamics

Diffusion models were introduced in 2015 as a method to learn a model that can sample from a highly complex probability distribution. They used techniques from non-equilibrium thermodynamics, especially diffusion.[10]

N(0,I)

. A model that can approximately undo the diffusion can then be used to sample from the original distribution. This is studied in "non-equilibrium" thermodynamics, as the starting distribution is not in equilibrium, unlike the final distribution.

The equilibrium distribution is the Gaussian distribution

N(0,I)

, with pdf

\rho(x)\propto

-122
\|x\|
e
. This is just the Maxwell–Boltzmann distribution of particles in a potential well

V(x)=

12
\|x\|

2

at temperature 1. The initial distribution, being very much out of equilibrium, would diffuse towards the equilibrium distribution, making biased random steps that are a sum of pure randomness (like a Brownian walker) and gradient descent down the potential well. The randomness is necessary: if the particles were to undergo only gradient descent, then they will all fall to the origin, collapsing the distribution.

Denoising Diffusion Probabilistic Model (DDPM)

The 2020 paper proposed the Denoising Diffusion Probabilistic Model (DDPM), which improves upon the previous method by variational inference.[11]

Forward diffusion

To present the model, we need some notation.

\beta1,...,\betaT\in(0,1)

are fixed constants.

\alphat:=1-\betat

\bar\alphat:=\alpha1\alphat

\tilde\betat:=

1-\bar\alphat-1
1-\bar\alphat

\betat

\tilde\mut(xt,x0):=

\sqrt{\alphat
(1-\bar

\alphat-1)xt+\sqrt{\bar\alphat-1

}(1-\alpha_)x_0}

N(\mu,\Sigma)

is the normal distribution with mean

\mu

and variance

\Sigma

, and

N(x|\mu,\Sigma)

is the probability density at

x

.

A forward diffusion process starts at some starting point

x0\simq

, where

q

is the probability distribution to be learned, then repeatedly adds noise to it byx_t = \sqrt x_ + \sqrt z_twhere

z1,...,zT

are IID samples from

N(0,I)

. This is designed so that for any starting distribution of

x0

, we have

\limtxt|x0

converging to

N(0,I)

.

The entire diffusion process then satisfiesq(x_) = q(x_0)q(x_1|x_0) \cdots q(x_T|x_) = q(x_0) N(x_1 | \sqrt x_0, \beta_1 I) \cdots N(x_T | \sqrt x_, \beta_T I)or\ln q(x_) = \ln q(x_0) - \sum_^T \frac \| x_t - \sqrtx_\|^2 + Cwhere

C

is a normalization constant and often omitted. In particular, we note that

x1:T|x0

is a gaussian process, which affords us considerable freedom in reparameterization. For example, by standard manipulation with gaussian process, x_|x_0 \sim N\left(\sqrt x_, (1-\bar\alpha_t) I \right)x_ | x_t, x_0 \sim N(\tilde\mu_t(x_t, x_0), \tilde \beta_t I)In particular, notice that for large

t

, the variable

xt|x0\simN\left(\sqrt{\bar\alphat}x0,(1-\bar\alphat)I\right)

converges to

N(0,I)

. That is, after a long enough diffusion process, we end up with some

xT

that is very close to

N(0,I)

, with all traces of the original

x0\simq

gone.

For example, sincex_|x_0 \sim N\left(\sqrt x_, (1-\bar\alpha_t) I \right)we can sample

xt|x0

directly "in one step", instead of going through all the intermediate steps

x1,x2,...,xt-1

.

Backward diffusion

The key idea of DDPM is to use a neural network parametrized by

\theta

. The network takes in two arguments

xt,t

, and outputs a vector

\mu\theta(xt,t)

and a matrix

\Sigma\theta(xt,t)

, such that each step in the forward diffusion process can be approximately undone by

xt-1\simN(\mu\theta(xt,t),\Sigma\theta(xt,t))

. This then gives us a backward diffusion process

p\theta

defined byp_\theta(x_T) = N(x_T | 0, I)p_\theta(x_ | x_t) = N(x_ | \mu_\theta(x_t, t), \Sigma_\theta(x_t, t))The goal now is to learn the parameters such that

p\theta(x0)

is as close to

q(x0)

as possible. To do that, we use maximum likelihood estimation with variational inference.

Variational inference

The ELBO inequality states that

lnp\theta(x0)\geq

E
x1:T\simq(|x0)

[lnp\theta(x0:T)-lnq(x1:T|x0)]

, and taking one more expectation, we getE_[\ln p_\theta(x_0)] \geq E_[\ln p_\theta(x_{0:T}) - \ln q(x_{1:T}|x_0)] We see that maximizing the quantity on the right would give us a lower bound on the likelihood of observed data. This allows us to perform variational inference.

Define the loss functionL(\theta) := -E_[\ln p_\theta(x_{0:T}) - \ln q(x_{1:T}|x_0)]and now the goal is to minimize the loss by stochastic gradient descent. The expression may be simplified to[12] L(\theta) = \sum_^T E_[-\ln p_\theta(x_{t-1} | x_t)] + E_[D_{KL}(q(x_T|x_0) \| p_\theta(x_T))] + Cwhere

C

does not depend on the parameter, and thus can be ignored. Since

p\theta(xT)=N(xT|0,I)

also does not depend on the parameter, the term
E
x0\simq

[DKL(q(xT|x0)\|p\theta(xT))]

can also be ignored. This leaves just

L(\theta)=

T
\sum
t=1

Lt

with

Lt=

E
xt-1,xt\simq

[-lnp\theta(xt-1|xt)]

to be minimized.

Noise prediction network

Since

xt-1|xt,x0\simN(\tilde\mut(xt,x0),\tilde\betatI)

, this suggests that we should use

\mu\theta(xt,t)=\tilde\mut(xt,x0)

; however, the network does not have access to

x0

, and so it has to estimate it instead. Now, since

xt|x0\simN\left(\sqrt{\bar\alphat}x0,(1-\bar\alphat)I\right)

, we may write

xt=\sqrt{\bar\alphat}x0+\sqrt{1-\bar\alphat}z

, where

z

is some unknown gaussian noise. Now we see that estimating

x0

is equivalent to estimating

z

.

Therefore, let the network output a noise vector

\epsilon\theta(xt,t)

, and let it predict\mu_\theta(x_t, t) =\tilde\mu_t\left(x_t, \frac\right) = \fracIt remains to design

\Sigma\theta(xt,t)

. The DDPM paper suggested not learning it (since it resulted in "unstable training and poorer sample quality"), but fixing it at some value

\Sigma\theta(xt,t)=

2
\sigma
t

I

, where either
2
\sigma
t

=\betator\tilde\betat

yielded similar performance.

With this, the loss simplifies to L_t = \frac E_\left[\left\| \epsilon_\theta(x_t, t) - z \right\|^2\right] + Cwhich may be minimized by stochastic gradient descent. The paper noted empirically that an even simpler loss functionL_ = E_\left[\left\| \epsilon_\theta(x_t, t) - z \right\|^2\right]resulted in better models.

Score-based generative model

Score-based generative model is another formulation of diffusion modelling. They are also called noise conditional score network (NCSN) or score-matching with Langevin dynamics (SMLD).[13] [14]

Score matching

The idea of score functions

Consider the problem of image generation. Let

x

represent an image, and let

q(x)

be the probability distribution over all possible images. If we have

q(x)

itself, then we can say for certain how likely a certain image is. However, this is intractable in general.

Most often, we are uninterested in knowing the absolute probability of a certain image. Instead, we are usually only interested in knowing how likely a certain image is compared to its immediate neighbors — e.g. how much more likely is an image of cat compared to some small variants of it? Is it more likely if the image contains two whiskers, or three, or with some Gaussian noise added?

Consequently, we are actually quite uninterested in

q(x)

itself, but rather,

\nablaxlnq(x)

. This has two major effects:

q(x)

, but can use any

\tildeq(x)=Cq(x)

, where

C=\int\tildeq(x)dx>0

is any unknown constant that is of no concern to us.

q(x)

neighbors

q(x+dx)

, by
q(x)
q(x+dx)
-\langle\nablaxlnq,dx\rangle
=e

Let the score function be

s(x):=\nablaxlnq(x)

; then consider what we can do with

s(x)

.

As it turns out,

s(x)

allows us to sample from

q(x)

using thermodynamics. Specifically, if we have a potential energy function

U(x)=-lnq(x)

, and a lot of particles in the potential well, then the distribution at thermodynamic equilibrium is the Boltzmann distribution

qU(x)\propto

-U(x)/kBT
e

=

1/kBT
q(x)
. At temperature

kBT=1

, the Boltzmann distribution is exactly

q(x)

.

Therefore, to model

q(x)

, we may start with a particle sampled at any convenient distribution (such as the standard gaussian distribution), then simulate the motion of the particle forwards according to the Langevin equationdx_= -\nabla_U(x_t) d t+d W_tand the Boltzmann distribution is, by Fokker-Planck equation, the unique thermodynamic equilibrium. So no matter what distribution

x0

has, the distribution of

xt

converges in distribution to

q

as

t\toinfty

.

Learning the score function

Given a density

q

, we wish to learn a score function approximation

f\theta\nablalnq

. This is score matching.[15] Typically, score matching is formalized as minimizing Fisher divergence function

Eq[\|f\theta(x)-\nablalnq(x)\|2]

. By expanding the integral, and performing an integration by parts, E_q[\|f_\theta(x) - \nabla \ln q(x)\|^2] = E_q[\|f_\theta\|^2 + 2\nabla^2\cdot f_\theta] + Cgiving us a loss function, also known as the Hyvärinen scoring rule, that can be minimized by stochastic gradient descent.

Annealing the score function

Suppose we need to model the distribution of images, and we want

x0\simN(0,I)

, a white-noise image. Now, most white-noise images do not look like real images, so

q(x0)0

for large swaths of

x0\simN(0,I)

. This presents a problem for learning the score function, because if there are no samples around a certain point, then we can't learn the score function at that point. If we do not know the score function
\nabla
xt

lnq(xt)

at that point, then we cannot impose the time-evolution equation on a particle:dx_= \nabla_\ln q(x_t) d t+d W_tTo deal with this problem, we perform annealing. If

q

is too different from a white-noise distribution, then progressively add noise until it is indistinguishable from one. That is, we perform a forward diffusion, then learn the score function, then use the score function to perform a backward diffusion.

Continuous diffusion processes

Forward diffusion process

Consider again the forward diffusion process, but this time in continuous time:x_t = \sqrt x_ + \sqrt z_tBy taking the

\betat\to\beta(t)dt,\sqrt{dt}zt\todWt

limit, we obtain a continuous diffusion process, in the form of a stochastic differential equation:dx_t = -\frac 12 \beta(t) x_t dt + \sqrt dW_twhere

Wt

is a Wiener process (multidimensional Brownian motion).

Now, the equation is exactly a special case of the overdamped Langevin equationdx_t = -\frac (\nabla_x U)dt + \sqrtdW_twhere

D

is diffusion tensor,

T

is temperature, and

U

is potential energy field. If we substitute in

D=

12
\beta(t)I,

kBT=1,U=

12
\|x\|

2

, we recover the above equation. This explains why the phrase "Langevin dynamics" is sometimes used in diffusion models.

Now the above equation is for the stochastic motion of a single particle. Suppose we have a cloud of particles distributed according to

q

at time

t=0

, then after a long time, the cloud of particles would settle into the stable distribution of

N(0,I)

. Let

\rhot

be the density of the cloud of particles at time

t

, then we have\rho_0 = q; \quad \rho_T \approx N(0, I)and the goal is to somehow reverse the process, so that we can start at the end and diffuse back to the beginning.

By Fokker-Planck equation, the density of the cloud evolves according to\partial_t \ln \rho_t = \frac 12 \beta(t) \left(n + (x+ \nabla\ln\rho_t) \cdot \nabla \ln\rho_t + \Delta\ln\rho_t\right)where

n

is the dimension of space, and

\Delta

is the Laplace operator.

Backward diffusion process

If we have solved

\rhot

for time

t\in[0,T]

, then we can exactly reverse the evolution of the cloud. Suppose we start with another cloud of particles with density

\nu0=\rhoT

, and let the particles in the cloud evolve according tody_t = \frac \beta(T-t) y_ d t + \beta(T-t) \underbrace_ d t+\sqrt d W_tthen by plugging into the Fokker-Planck equation, we find that

\partialt\rhoT-t=\partialt\nut

. Thus this cloud of points is the original cloud, evolving backwards.[16]

Noise conditional score network (NCSN)

At the continuous limit, \bar \alpha_t = (1-\beta_1) \cdots (1-\beta_t) = e^ \to e^ and so x_|x_0 \sim N\left(e^ x_, \left(1- e^\right) I \right)In particular, we see that we can directly sample from any point in the continuous diffusion process without going through the intermediate steps, by first sampling

x0\simq,z\simN(0,I)

, then get

xt=

-12\intt\beta(t)dt
0
e

x0+\left(1-

t
-\int\beta(t)dt
0
e

\right)z

. That is, we can quickly sample

xt\sim\rhot

for any

t\geq0

.

Now, define a certain probability distribution

\gamma

over

[0,infty)

, then the score-matching loss function is defined as the expected Fisher divergence:L(\theta) = E_[\|f_\theta(x_t, t)\|^2 + 2\nabla\cdot f_\theta(x_t, t)]After training,

f\theta(xt,t)\nablaln\rhot

, so we can perform the backwards diffusion process by first sampling

xT\simN(0,I)

, then integrating the SDE from

t=T

to

t=0

:x_=x_t + \frac \beta(t) x_ d t + \beta(t) f_\theta(x_t, t) d t+\sqrt d W_tThis may be done by any SDE integration method, such as Euler–Maruyama method.

The name "noise conditional score network" is explained thus:

f\theta

is implemented as a neural network.

\nablaln\rhot

.

\rhot

is equal to

\rho0

blurred by an added gaussian noise that increases with time, and so the score function depends on the amount of noise added.

Their equivalence

DDPM and score-based generative models are equivalent.[17] This means that a network trained using DDPM can be used as a NCSN, and vice versa.

We know that

xt|x0\simN\left(\sqrt{\bar\alphat}x0,(1-\bar\alphat)I\right)

, so by Tweedie's formula, we have\nabla_\ln q(x_t) = \frac(-x_t + \sqrt E_q[x_0|x_t])As described previously, the DDPM loss function is

\sumtLsimple,

withL_ = E_\left[\left\| \epsilon_\theta(x_t, t) - z \right\|^2\right]where

xt=\sqrt{\bar\alphat}x0+\sqrt{1-\bar\alphat}z

. By a change of variables,L_ = E_\left[\left\| \epsilon_\theta(x_t, t) - \frac{x_t -\sqrt{\bar\alpha_t} x_{0}}{\sqrt{1-\bar\alpha_t}} \right\|^2\right] = E_\left[\left\| \epsilon_\theta(x_t, t) - \frac{x_t -\sqrt{\bar\alpha_t} x_{0}}{\sqrt{1-\bar\alpha_t}} \right\|^2\right]and the term inside becomes a least squares regression, so if the network actually reaches the global minimum of loss, then we have

\epsilon\theta(xt,t)=

xt-\sqrt{\bar\alphat
E

q[x0|xt]}{\sqrt{1-\bar\alphat}}=-\sqrt{1-\bar\alphat}\nabla

xt

lnq(xt)

.

Now, the continuous limit

xt-1=xt-dt,\betat=\beta(t)dt,zt\sqrt{dt}=dWt

of the backward equationx_ = \frac- \frac \epsilon_\theta(x_t, t) + \sqrt z_t; \quad z_t \sim N(0, I)gives us precisely the same equation as score-based diffusion:x_ = x_t(1+\beta(t)dt / 2) + \beta(t) \nabla_\ln q(x_t) dt + \sqrtdW_t

Main variants

Denoising Diffusion Implicit Model (DDIM)

The original DDPM method for generating images is slow, since the forward diffusion process usually takes

T\sim1000

to make the distribution of

xT

to appear close to gaussian. However this means the backward diffusion process also take 1000 steps. Unlike the forward diffusion process, which can skip steps as

xt|x0

is gaussian for all

t\geq1

, the backward diffusion process does not allow skipping steps. For example, to sample

xt-2|xt-1\simN(\mu\theta(xt-1,t-1),\Sigma\theta(xt-1,t-1))

requires the model to first sample

xt-1

. Attempting to directly sample

xt-2|xt

would require us to marginalize out

xt-1

, which is generally intractable.

DDIM[18] is a method to take any model trained on DDPM loss, and use it to sample with some steps skipped, sacrificing an adjustable amount of quality. If we generate the Markovian chain case in DDPM to non-Markovian case, DDIM corresponds to the case that the reverse process has variance equals to 0. In other words, the reverse process (and also the forward process) is deterministic. When using fewer sampling steps, DDIM outperforms DDPM.

Latent diffusion model (LDM)

Since the diffusion model is a general method for modelling probability distributions, if one wants to model a distribution over images, one can first encode the images into a lower-dimensional space by an encoder, then use a diffusion model to model the distribution over encoded images. Then to generate an image, one can sample from the diffusion model, then use a decoder to decode it into an image.[19]

The encoder-decoder pair is most often a variational autoencoder (VAE).

Classifier guidance

Suppose we wish to sample not from the entire distribution of images, but conditional on the image description. We don't want to sample a generic image, but an image that fits the description "black cat with red eyes". Generally, we want to sample from the distribution

p(x|y)

, where

x

ranges over images, and

y

ranges over classes of images (a description "black cat with red eyes" is just a very detailed class, and a class "cat" is just a very vague description).

Taking the perspective of the noisy channel model, we can understand the process as follows: To generate an image

x

conditional on description

y

, we imagine that the requester really had in mind an image

x

, but the image is passed through a noisy channel and came out garbled, as

y

. Image generation is then nothing but inferring which

x

the requester had in mind.

In other words, conditional image generation is simply "translating from a textual language into a pictorial language". Then, as in noisy-channel model, we use Bayes theorem to getp(x|y) \propto p(y|x)p(x) in other words, if we have a good model of the space of all images, and a good image-to-class translator, we get a class-to-image translator "for free". In the equation for backward diffusion, the score

\nablalnp(x)

can be replaced by\nabla_x \ln p(x|y) = \nabla_x \ln p(y|x) + \nabla_x \ln p(x) where

\nablaxlnp(x)

is the score function, trained as previously described, and

\nablaxlnp(y|x)

is found by using a differentiable image classifier.

With temperature

The classifier-guided diffusion model samples from

p(x|y)

, which is concentrated around the maximum a posteriori estimate

\argmaxxp(x|y)

. If we want to force the model to move towards the maximum likelihood estimate

\argmaxxp(y|x)

, we can use p_\beta(x|y) \propto p(y|x)^\beta p(x) where

\beta>0

is interpretable as inverse temperature. In the context of diffusion models, it is usually called the guidance scale. A high

\beta

would force the model to sample from a distribution concentrated around

\argmaxxp(y|x)

. This often improves quality of generated images.[20]

This can be done simply by SGLD with\nabla_x \ln p_\beta(x|y) = \beta\nabla_x \ln p(y|x) + \nabla_x \ln p(x)

Classifier-free guidance (CFG)

If we do not have a classifier

p(y|x)

, we could still extract one out of the image model itself:[21] \nabla_x \ln p_\beta(x|y) = (1-\beta) \nabla_x \ln p(x) + \beta \nabla_x \ln p(x|y) Such a model is usually trained by presenting it with both

(x,y)

and

(x,{\rmNone})

, allowing it to model both

\nablaxlnp(x|y)

and

\nablaxlnp(x)

.

Samplers

Given a diffusion model, one may regard it either as a continuous process, and sample from it by integrating a SDE, or one can regard it as a discrete process, and sample from it by iterating the discrete steps. The choice of the "noise schedule"

\betat

can also affect the quality of samples. In the DDPM perspective, one can use the DDPM itself (with noise), or DDIM (with adjustable amount of noise). The case where one adds noise is sometimes called ancestral sampling.[22] One can interpolate between noise and no noise. The amount of noise is denoted

η

("eta value") in the DDIM paper, with

η=0

denoting no noise (as in deterministic DDIM), and

η=1

denoting full noise (as in DDPM).

In the perspective of SDE, one can use any of the numerical integration methods, such as Euler–Maruyama method, Heun's method, linear multistep methods, etc. Just as in the discrete case, one can add an adjustable amount of noise during the integration.

A survey and comparison of samplers in the context of image generation is in.[23]

Other examples

Notable variants include[24] Poisson flow generative model[25], consistency model[26], critically-damped Langevin diffusion[27], GenPhys, cold diffusion[28], etc.

Flow-based diffusion model

Abstractly speaking, the idea of diffusion model is to take an unknown probability distribution (the distribution of natural-looking images), then progressively convert it to a known probability distribution (standard gaussian distribution), by building an absolutely continuous probability path connecting them. The probability path is in fact defined implicitly by the score function

\nablalnpt

.

In denoising diffusion models, the forward process adds noise, and the backward process removes noise. Both the forward and backward processes are SDEs, though the forward process is integrable in closed-form, so it can be done at no computational cost. The backward process is not integrable in closed-form, so it must be integrated step-by-step by standard SDE solvers, which can be very expensive. The probability path in diffusions model is defined through an Itô process and one can retrieve the deterministic process by using the Probability ODE flow formulation.

In flow-based diffusion models, the forward process is a both deterministic flow along a time-dependent vector field, and the backward process is the same vector field, but going backwards. Both processes are solutions to ODEs. If the vector field is well-behaved, the ODE will also be well-behaved.

Given two distributions

\pi0

and

\pi1

, a flow-based model is a time-dependent velocity field

vt(x)

in

[0,1] x Rd

, such that if we start by sampling a point

x\sim\pi0

, and let it move according to the velocity field:\frac \phi_t(x) = v_t(\phi_t(x)) \quad t \in [0,1], \quad \text\phi_0(x) = xwe end up with a point

x1\sim\pi1

. The solution

\phit

of the above ODE define a probability path

pt=[\phit]\#\pi0

by the Pushforward measure operator. Espectially, one has

[\phi1]\#\pi0=\pi1

.

The probability path and the velocity field also satisfy the Continuity equation, in the sense of probability distribution:\partial_t p_t + \mathrm(v_t \cdot p_t) = 0To construct a probability path, we start by construct a conditional probability path

pt(x\vertz)

and the corresponding conditional velocity field

vt(x\vertz)

on some conditional distribution

q(z)

. A natural choice is the Gaussian conditional probability path:p_t(x \vert z) = \mathcal \left(m_t(z), \sigma_t^2 I \right) The conditional velocity field which corresponds to the geodesic path between conditional Gaussian path is v_t(x \vert z) = \frac (x - m_t(z)) + m_t'(z)The probability path and velocity field are then computed by marginalizing

pt(x)=\intpt(x\vertz)q(z)dz    and    vt(x)=Eq(z)\left[

vt(x\vertz)pt(x\vertz)
pt(x)

\right]

Optimal Transport Flow

The idea of optimal transport flow [29] is to construct a probability path minimizing the Wasserstein metric. The distribution on which we condition is the optimal transport plan between

\pi0

and

\pi1

:

z=(x0,x1)

and

q(z)=\Gamma(\pi0,\pi1)

, where

\Gamma

is the optimal transport plan, which can be approximated by mini-batch optimal transport.

Rectified flow

The idea of rectified flow[30] [31] is to learn a flow model such that the velocity is nearly constant along each flow path. This is beneficial, because we can integrate along such a vector field with very few steps. For example, if an ODE

\phit

(x)=vt(\phit(x))

follows perfectly straight paths, it simplifies to

\phit(x)=x0+tv0(x0)

, allowing for exact solutions in one step. In practice, we cannot reach such perfection, but when the flow field is nearly so, we can take a few large steps instead of many little steps. The general idea is to start with two distributions

\pi0

and

\pi1

, then construct a flow field

\phi0=\{\phit:t\in[0,1]\}

from it, then repeatedly apply a "reflow" operation to obtain successive flow fields

\phi1,\phi2,...

, each straighter than the previous one. When the flow field is straight enough for the application, we stop.

Generally, for any time-differentiable process

\phit

,

vt

can be estimated by solving:\min_ \int_0^1 \mathbb_\left [\lVert{v_t(x, \theta) - v_t(x)}\rVert^2\right] \,\mathrmt.

In rectified flow, by injecting strong priors that intermediate trajectories are straight, it can achieve both theoretical relevance for optimal transport and computational efficiency, as ODEs with straight paths can be simulated precisely without time discretization.Specifically, rectified flow seeks to match an ODE with the marginal distributions of the linear interpolation between points from distributions

\pi0

and

\pi1

. Given observations

x0\sim\pi0

and

x1\sim\pi1

, the canonical linear interpolation

xt=tx1+(1-t)x0,t\in[0,1]

yields a trivial case
x

t=x1-x0

, which cannot be causally simulated without

x1

. To address this,

xt

is "projected" into a space of causally simulatable ODEs, by minimizing the least squares loss with respect to the direction

x1-x0

:\min_ \int_0^1 \mathbb_\left [\lVert{(x_1-x_0) - v_t(x_t)}\rVert^2\right] \,\mathrmt.

The data pair

(x0,x1)

can be any coupling of

\pi0

and

\pi1

, typically independent (i.e.,

(x0,x1)\sim\pi0 x \pi1

) obtained by randomly combining observations from

\pi0

and

\pi1

. This process ensures that the trajectories closely mirror the density map of

xt

trajectories but reroute at intersections to ensure causality. This rectifying process is also known as Flow Matching, Stochastic Interpolation, and Alpha-Blending.

A distinctive aspect of rectified flow is its capability for "reflow", which straightens the trajectory of ODE paths. Denote the rectified flow

\phi0=\{\phit:t\in[0,1]\}

induced from

(x0,x1)

as

\phi0=Rectflow((x0,x1))

. Recursively applying this

Rectflow()

operator generates a series of rectified flows

\phik+1=

k(x
Rectflow((\phi
0),
k(x
\phi
1)))
. This "reflow" process not only reduces transport costs but also straightens the paths of rectified flows, making

\phik

paths straighter with increasing

k

.

Rectified flow includes a nonlinear extension where linear interpolation

xt

is replaced with any time-differentiable curve that connects

x0

and

x1

, given by

xt=\alphatx1+\betatx0

. This framework encompasses DDIM and probability flow ODEs as special cases, with particular choices of

\alphat

and

\betat

. However, in the case where the path of

xt

is not straight, the reflow process no longer ensures a reduction in convex transport costs, and also no longer straighten the paths of

\phit

.

See [32] for a tutorial on flow matching, with animations.

Choice of architecture

Diffusion model

For generating images by DDPM, we need a neural network that takes a time

t

and a noisy image

xt

, and predicts a noise

\epsilon\theta(xt,t)

from it. Since predicting the noise is the same as predicting the denoised image, then subtracting it from

xt

, denoising architectures tend to work well. For example, the U-Net, which was found to be good for denoising images, is often used for denoising diffusion models that generate images.[33]

For DDPM, the underlying architecture ("backbone") does not have to be a U-Net. It just has to predict the noise somehow. For example, the diffusion transformer (DiT) uses a Transformer to predict the mean and diagonal covariance of the noise, given the textual conditioning and the partially denoised image. It is the same as standard U-Net-based denoising diffusion model, with a Transformer replacing the U-Net.[34] Mixture of experts-Transformer can also be applied.

DDPM can be used to model general data distributions, not just natural-looking images. For example, Human Motion Diffusion[35] models human motion trajectory by DDPM. Each human motion trajectory is a sequence of poses, represented by either joint rotations or positions. It uses a Transformer network to generate a less noisy trajectory out of a noisy one.

Conditioning

The base diffusion model can only generate unconditionally from the whole distribution. For example, a diffusion model learned on ImageNet would generate images that look like a random image from ImageNet. To generate images from just one category, one would need to impose the condition. Whatever condition one wants to impose, one needs to first convert the conditioning into a vector of floating point numbers, then feed it into the underlying diffusion model neural network. However, one has freedom in choosing how to convert the conditioning into a vector.

Stable Diffusion, for example, imposes conditioning in the form of cross-attention mechanism, where the query is an intermediate representation of the image in the U-Net, and both key and value are the conditioning vectors. The conditioning can be selectively applied to only parts of an image, and new kinds of conditionings can be finetuned upon the base model, as used in ControlNet.[36]

As a particularly simple example, consider image inpainting. The conditions are

\tildex

, the reference image, and

m

, the inpainting mask. The conditioning is imposed at each step of the backward diffusion process, by first sampling

\tildext\simN\left(\sqrt{\bar\alphat}\tildex,(1-\bar\alphat)I\right)

, a noisy version of

\tildex

, then replacing

xt

with

(1-m)\odotxt+m\odot\tildext

, where

\odot

means elementwise multiplication.[37]

Conditioning is not limited to just generating images from a specific category, or according to a specific caption (as in text-to-image). For example, demonstrated generating human motion, conditioned on an audio clip of human walking (allowing syncing motion to a soundtrack), or video of human running, or a text description of human motion, etc.

Upscaling

As generating an image takes a long time, one can try to generate a small image by a base diffusion model, then upscale it by other models. Upscaling can be done by GAN,[38] Transformer,[39] or signal processing methods like Lanczos resampling.

Diffusion models themselves can be used to perform upscaling. Cascading diffusion model stacks multiple diffusion models one after another, in the style of Progressive GAN. The lowest level is a standard diffusion model that generate 32x32 image, then the image would be upscaled by a diffusion model specifically trained for upscaling, and the process repeats.

In more detail, the diffusion upscaler is trained as follows:

(x0,z0,c)

, where

x0

is the high-resolution image,

z0

is the same image but scaled down to a low-resolution, and

c

is the conditioning, which can be the caption of the image, the class of the image, etc.

\epsilonx,\epsilonz

, two time-steps

tx,tz

. Compute the noisy versions of the high-resolution and low-resolution images:
\begin{cases} x
tx

&=

\sqrt{\bar\alpha
tx
} x_0 + \sqrt \epsilon_x\\z_ &= \sqrt z_0 + \sqrt \epsilon_z\end.

\epsilonx

given
x
tx

,

z
tz

,tx,tz,c

. That is, apply gradient descent on

\theta

on the L2 loss

\|\epsilon\theta(x

tx

,

z
tz

,tx,tz,c)-\epsilonx

2
\|
2
.

Examples

This section collects some notable diffusion models, and briefly describes their architecture.

OpenAI

See main article: DALL-E and Sora (text-to-video model). The DALL-E series by OpenAI are text-conditional diffusion models of images.

The first version of DALL-E (2021) is not actually a diffusion model. Instead, it uses a Transformer architecture that generates a sequence of tokens, which is then converted to an image by the decoder of a discrete VAE. Released with DALL-E was the CLIP classifier, which was used by DALL-E to rank generated images according to how close the image fits the text.

GLIDE (2022-03)[40] is a 3.5-billion diffusion model, and a small version was released publicly. Soon after, DALL-E 2 was released (2022-04).[41] DALL-E 2 is a 3.5-billion cascaded diffusion model that generates images from text by "inverting the CLIP image encoder", the technique which they termed "unCLIP".

Sora (2024-02) is a diffusion Transformer model (DiT).

Stability AI

See main article: Stable Diffusion.

Stable Diffusion, released by Stability AI, consists of a denoising latent diffusion model (860 million parameters), a VAE, and a text encoder. The denoising network is a U-Net, with cross-attention blocks to allow for conditional image generation.[42]

Stable Diffusion 3 changed the latent diffusion model from the UNet to a Transformer model, and so it is a DiT. It uses rectified flow.

Stable Video 4D is a latent diffusion model for videos of 3D objects.

Google

Imagen[43] [44] uses a T5 language model to encode the input text into an embedding vector. It is a cascaded diffusion model with three steps. The first step denoises a white noise to a 64×64 image, conditional on the embedding vector of the text. The second step upscales the image by 64×64→256×256, conditional on embedding. The third step is similar, upscaling by 256×256→1024×1024. The three denoising networks are all U-Nets.

Imagen 2 is also diffusion-based. It can generate images based on a prompt that mixes images and text. No further information available.[45]

Veo generates videos by latent diffusion. The diffusion is conditioned on a vector that encodes both a text prompt and an image prompt.[46]

See also

Further reading

Notes and References

  1. Chang . Ziyi . Koulieris . George Alex . Shum . Hubert P. H. . On the Design Fundamentals of Diffusion Models: A Survey . 2023 . 2306.04542 . cs.LG.
  2. Croitoru . Florinel-Alin . Hondru . Vlad . Ionescu . Radu Tudor . Shah . Mubarak . 2023 . Diffusion Models in Vision: A Survey . IEEE Transactions on Pattern Analysis and Machine Intelligence . 45 . 9 . 10850–10869 . 2209.04747 . 10.1109/TPAMI.2023.3261988 . 37030794 . 252199918.
  3. Song . Yang . Sohl-Dickstein . Jascha . Kingma . Diederik P. . Kumar . Abhishek . Ermon . Stefano . Poole . Ben . 2021-02-10 . Score-Based Generative Modeling through Stochastic Differential Equations . cs.LG . 2011.13456 .
  4. Gu . Shuyang . Chen . Dong . Bao . Jianmin . Wen . Fang . Zhang . Bo . Chen . Dongdong . Yuan . Lu . Guo . Baining . Vector Quantized Diffusion Model for Text-to-Image Synthesis . 2021 . cs.CV . 2111.14822.
  5. Book: Li . Yifan . Zhou . Kun . Zhao . Wayne Xin . Wen . Ji-Rong . Diffusion Models for Non-autoregressive Text Generation: A Survey . August 2023 . 6692–6701 . Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence . http://dx.doi.org/10.24963/ijcai.2023/750 . California . International Joint Conferences on Artificial Intelligence Organization . 10.24963/ijcai.2023/750. 2303.06574 . 978-1-956792-03-4 .
  6. Han . Xiaochuang . Kumar . Sachin . Tsvetkov . Yulia . 2023 . SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control . Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) . 11575–11596 . Stroudsburg, PA, USA . Association for Computational Linguistics . 10.18653/v1/2023.acl-long.647. 2210.17432 .
  7. Xu . Weijie . Hu . Wenxiang . Wu . Fanyou . Sengamedu . Srinivasan . 2023 . DeTiME: Diffusion-Enhanced Topic Modeling using Encoder-decoder based LLM . Findings of the Association for Computational Linguistics: EMNLP 2023 . 9040–9057 . Stroudsburg, PA, USA . Association for Computational Linguistics . 10.18653/v1/2023.findings-emnlp.606. 2310.15296 .
  8. Zhang . Haopeng . Liu . Xiao . Zhang . Jiawei . 2023 . DiffuSum: Generation Enhanced Extractive Summarization with Diffusion . Findings of the Association for Computational Linguistics: ACL 2023 . 13089–13100 . Stroudsburg, PA, USA . Association for Computational Linguistics . 10.18653/v1/2023.findings-acl.828. 2305.01735 .
  9. Yang . Dongchao . Yu . Jianwei . Wang . Helin . Wang . Wen . Weng . Chao . Zou . Yuexian . Yu . Dong . 2023 . Diffsound: Discrete Diffusion Model for Text-to-Sound Generation . IEEE/ACM Transactions on Audio, Speech, and Language Processing . 31 . 1720–1733 . 10.1109/taslp.2023.3268730 . 2329-9290.
  10. Sohl-Dickstein . Jascha . Weiss . Eric . Maheswaranathan . Niru . Ganguli . Surya . 2015-06-01 . Deep Unsupervised Learning using Nonequilibrium Thermodynamics . Proceedings of the 32nd International Conference on Machine Learning . en . PMLR . 37 . 2256–2265.
  11. Ho . Jonathan . Jain . Ajay . Abbeel . Pieter . 2020 . Denoising Diffusion Probabilistic Models . Advances in Neural Information Processing Systems . Curran Associates, Inc. . 33 . 6840–6851.
  12. Web site: Weng . Lilian . 2021-07-11 . What are Diffusion Models? . 2023-09-24 . lilianweng.github.io . en.
  13. Web site: Generative Modeling by Estimating Gradients of the Data Distribution Yang Song . 2023-09-24 . yang-song.net.
  14. 2011.13456 . cs.LG . Yang . Song . Jascha . Sohl-Dickstein . Score-Based Generative Modeling through Stochastic Differential Equations . 2021-02-10 . Kingma . Diederik P. . Kumar . Abhishek . Ermon . Stefano . Poole . Ben.
  15. Web site: Sliced Score Matching: A Scalable Approach to Density and Score Estimation Yang Song . 2023-09-24 . yang-song.net.
  16. Anderson . Brian D.O. . May 1982 . Reverse-time diffusion equation models . Stochastic Processes and Their Applications . 12 . 3 . 313–326 . 10.1016/0304-4149(82)90051-5 . 0304-4149.
  17. 2208.11970v1 . Luo . Calvin . 2022 . Understanding Diffusion Models: A Unified Perspective. cs.LG .
  18. Song . Jiaming . Meng . Chenlin . Ermon . Stefano . 3 Oct 2023 . Denoising Diffusion Implicit Models . cs.LG . 2010.02502.
  19. Rombach . Robin . Blattmann . Andreas . Lorenz . Dominik . Esser . Patrick . Ommer . Björn . 13 April 2022 . High-Resolution Image Synthesis With Latent Diffusion Models . cs.CV . 2112.10752 .
  20. Dhariwal . Prafulla . Nichol . Alex . 2021-06-01 . Diffusion Models Beat GANs on Image Synthesis . cs.LG . 2105.05233 .
  21. Ho . Jonathan . Salimans . Tim . 2022-07-25 . Classifier-Free Diffusion Guidance . cs.LG . 2207.12598 .
  22. 2206.00364 . Yang . Ling . Zhang . Zhilong . Song . Yang . Hong . Shenda . Xu . Runsheng . Zhao . Yue . Zhang . Wentao . Cui . Bin . Yang . Ming-Hsuan . 2022 . Diffusion Models: A Comprehensive Survey of Methods and Applications. cs.CV .
  23. 2206.00364v2 . Karras . Tero . Aittala . Miika . Aila . Timo . Laine . Samuli . 2022 . Elucidating the Design Space of Diffusion-Based Generative Models. cs.CV .
  24. Cao . Hanqun . Tan . Cheng . Gao . Zhangyang . Xu . Yilun . Chen . Guangyong . Heng . Pheng-Ann . Li . Stan Z. . July 2024 . A Survey on Generative Diffusion Models . IEEE Transactions on Knowledge and Data Engineering . 36 . 7 . 2814–2830 . 10.1109/TKDE.2024.3361474 . 1041-4347.
  25. Xu . Yilun . Liu . Ziming . Tian . Yonglong . Tong . Shangyuan . Tegmark . Max . Jaakkola . Tommi . 2023-07-03 . PFGM++: Unlocking the Potential of Physics-Inspired Generative Models . Proceedings of the 40th International Conference on Machine Learning . en . PMLR . 38566–38591.
  26. Song . Yang . Dhariwal . Prafulla . Chen . Mark . Sutskever . Ilya . 2023-07-03 . Consistency Models . Proceedings of the 40th International Conference on Machine Learning . en . PMLR . 32211–32252.
  27. Dockhorn . Tim . Vahdat . Arash . Kreis . Karsten . 2021-10-06 . Score-Based Generative Modeling with Critically-Damped Langevin Diffusion . en.
  28. Bansal . Arpit . Borgnia . Eitan . Chu . Hong-Min . Li . Jie . Kazemi . Hamid . Huang . Furong . Goldblum . Micah . Geiping . Jonas . Goldstein . Tom . 2023-12-15 . Cold Diffusion: Inverting Arbitrary Image Transforms Without Noise . Advances in Neural Information Processing Systems . en . 36 . 41259–41282.
  29. Tong . Alexander . Fatras . Kilian . Malkin . Nikolay . Huguet . Guillaume . Zhang . Yanlei . Rector-Brooks . Jarrid . Wolf . Guy . Bengio . Yoshua . 2023-11-08 . Improving and generalizing flow-based generative models with minibatch optimal transport . Transactions on Machine Learning Research . en . 2835-8856.
  30. Liu . Xingchao . Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow . 2022-09-07 . 2209.03003 . Gong . Chengyue . Liu . Qiang. cs.LG .
  31. Liu . Qiang . Rectified Flow: A Marginal Preserving Approach to Optimal Transport . 2022-09-29 . stat.ML . 2209.14577.
  32. Web site: An introduction to Flow Matching · Cambridge MLG Blog . 2024-08-20 . mlg.eng.cam.ac.uk.
  33. Ho . Jonathan . Saharia . Chitwan . Chan . William . Fleet . David J. . Norouzi . Mohammad . Salimans . Tim . 2022-01-01 . Cascaded diffusion models for high fidelity image generation . The Journal of Machine Learning Research . 23 . 1 . 47:2249–47:2281 . 2106.15282 . 1532-4435.
  34. 2212.09748v2 . Peebles . William . Xie . Saining . March 2023 . Scalable Diffusion Models with Transformers . cs.CV . en.
  35. 2209.14916 . Tevet . Guy . Raab . Sigal . Gordon . Brian . Shafir . Yonatan . Cohen-Or . Daniel . Bermano . Amit H. . 2022 . Human Motion Diffusion Model. cs.CV .
  36. Zhang . Lvmin . Rao . Anyi . Agrawala . Maneesh . 2023 . Adding Conditional Control to Text-to-Image Diffusion Models . cs.CV . 2302.05543.
  37. 2201.09865v4 . Lugmayr . Andreas . Danelljan . Martin . Romero . Andres . Yu . Fisher . Timofte . Radu . Van Gool . Luc . 2022 . RePaint: Inpainting Using Denoising Diffusion Probabilistic Models . cs.CV . en.
  38. Wang . Xintao . Xie . Liangbin . Dong . Chao . Shan . Ying . 2021 . Real-ESRGAN: Training Real-World Blind Super-Resolution With Pure Synthetic Data . International Conference on Computer Vision . Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV) Workshops, 2021 . en . 1905–1914. 2107.10833 .
  39. Liang . Jingyun . Cao . Jiezhang . Sun . Guolei . Zhang . Kai . Van Gool . Luc . Timofte . Radu . 2021 . SwinIR: Image Restoration Using Swin Transformer . Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV) Workshops . International Conference on Computer Vision, 2021 . en . 1833–1844. 2108.10257v1 .
  40. 2112.10741 . cs.CV . Alex . Nichol . Prafulla . Dhariwal . GLIDE: Towards Photorealistic Image Generation and Editing with Text-Guided Diffusion Models . 2022-03-08 . Ramesh . Aditya . Shyam . Pranav . Mishkin . Pamela . McGrew . Bob . Sutskever . Ilya . Chen . Mark.
  41. 2204.06125 . cs.CV . Aditya . Ramesh . Prafulla . Dhariwal . Hierarchical Text-Conditional Image Generation with CLIP Latents . 2022-04-12 . Nichol . Alex . Chu . Casey . Chen . Mark.
  42. Web site: Alammar . Jay . The Illustrated Stable Diffusion . 2022-10-31 . jalammar.github.io.
  43. Web site: Imagen: Text-to-Image Diffusion Models . 2024-04-04 . imagen.research.google.
  44. Saharia . Chitwan . Chan . William . Saxena . Saurabh . Li . Lala . Whang . Jay . Denton . Emily L. . Ghasemipour . Kamyar . Gontijo Lopes . Raphael . Karagol Ayan . Burcu . Salimans . Tim . Ho . Jonathan . Fleet . David J. . Norouzi . Mohammad . 2022-12-06 . Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding . Advances in Neural Information Processing Systems . en . 35 . 36479–36494. 2205.11487 .
  45. Web site: Imagen 2 - our most advanced text-to-image technology . 2024-04-04 . Google DeepMind . en.
  46. Web site: 2024-05-14 . Veo . 2024-05-17 . Google DeepMind . en.