Forward diffusion

To present the model, we need some notation.

• ${\displaystyle \beta _{1},...,\beta _{T}\in (0,1)}$ are fixed constants.
• ${\displaystyle \alpha _{t}:=1-\beta _{t}}$
• ${\displaystyle {\bar {\alpha }}_{t}:=\alpha _{1}\cdots \alpha _{t}}$
• ${\displaystyle {\tilde {\beta }}_{t}:={\frac {1-{\bar {\alpha }}_{t-1}}{1-{\bar {\alpha }}_{t}}}\beta _{t}}$
• ${\displaystyle {\tilde {\mu }}_{t}(x_{t},x_{0}):={\frac {{\sqrt {\alpha _{t}}}(1-{\bar {\alpha }}_{t-1})x_{t}+{\sqrt {{\bar {\alpha }}_{t-1}}}(1-\alpha _{t})x_{0}}{1-{\bar {\alpha }}_{t}}}}$
• ${\displaystyle N(\mu ,\Sigma )}$ is the normal distribution with mean ${\displaystyle \mu }$ and variance ${\displaystyle \Sigma }$, and ${\displaystyle N(x|\mu ,\Sigma )}$ is the probability density at ${\displaystyle x}$.
• A vertical bar denotes conditioning.

A forward diffusion process starts at some starting point ${\displaystyle x_{0}\sim q}$, where ${\displaystyle q}$ is the probability distribution to be learned, then repeatedly adds noise to it by

$${\displaystyle x_{t}={\sqrt {1-\beta _{t}}}x_{t-1}+{\sqrt {\beta _{t}}}z_{t}}$$

where ${\displaystyle z_{1},...,z_{T}}$ are IID samples from ${\displaystyle N(0,I)}$. This is designed so that for any starting distribution of ${\displaystyle x_{0}}$, we have ${\displaystyle \lim _{t}x_{t}|x_{0}}$ converging to ${\displaystyle N(0,I)}$. The entire diffusion process then satisfies

$${\displaystyle q(x_{0:T})=q(x_{0})q(x_{1}|x_{0})\cdots q(x_{T}|x_{T-1})=q(x_{0})N(x_{1}|{\sqrt {\alpha _{1}}}x_{0},\beta _{1}I)\cdots N(x_{T}|{\sqrt {\alpha _{T}}}x_{T-1},\beta _{T}I)}$$

or

$${\displaystyle \ln q(x_{0:T})=\ln q(x_{0})-\sum _{t=1}^{T}{\frac {1}{2\beta _{t}}}\|x_{t}-{\sqrt {1-\beta _{t}}}x_{t-1}\|^{2}+C}$$

where ${\displaystyle C}$ is a normalization constant and often omitted. In particular, we note that ${\displaystyle x_{1:T}|x_{0}}$ is a gaussian process, which affords us considerable freedom in reparameterization. For example, by standard manipulation with gaussian process,

$${\displaystyle x_{t}|x_{0}\sim N\left({\sqrt {{\bar {\alpha }}_{t}}}x_{0},(1-{\bar {\alpha }}_{t})I\right)}$$

$${\displaystyle x_{t-1}|x_{t},x_{0}\sim N({\tilde {\mu }}_{t}(x_{t},x_{0}),{\tilde {\beta }}_{t}I)}$$

In particular, notice that for large ${\displaystyle t}$, the variable ${\displaystyle x_{t}|x_{0}\sim N\left({\sqrt {{\bar {\alpha }}_{t}}}x_{0},(1-{\bar {\alpha }}_{t})I\right)}$ converges to ${\displaystyle N(0,I)}$. That is, after a long enough diffusion process, we end up with some ${\displaystyle x_{T}}$ that is very close to ${\displaystyle N(0,I)}$, with all traces of the original ${\displaystyle x_{0}\sim q}$ gone. For example, since

$${\displaystyle x_{t}|x_{0}\sim N\left({\sqrt {{\bar {\alpha }}_{t}}}x_{0},(1-{\bar {\alpha }}_{t})I\right)}$$

we can sample ${\displaystyle x_{t}|x_{0}}$ directly "in one step", instead of going through all the intermediate steps ${\displaystyle x_{1},x_{2},...,x_{t-1}}$.

Derivation by reparameterization We know ${\textstyle x_{t-1}|x_{0}}$ is a gaussian, and ${\textstyle x_{t}|x_{t-1}}$ is another gaussian. We also know that these are independent. Thus we can perform a reparameterization:

$${\displaystyle x_{t-1}={\sqrt {{\bar {\alpha }}_{t-1}}}x_{0}+{\sqrt {1-{\bar {\alpha }}_{t-1}}}z}$$

$${\displaystyle x_{t}={\sqrt {\alpha _{t}}}x_{t-1}+{\sqrt {1-\alpha _{t}}}z'}$$

where ${\textstyle z,z'}$ are IID gaussians.

There are 5 variables ${\textstyle x_{0},x_{t-1},x_{t},z,z'}$ and two linear equations. The two sources of randomness are ${\textstyle z,z'}$, which can be reparameterized by rotation, since the IID gaussian distribution is rotationally symmetric.

By plugging in the equations, we can solve for the first reparameterization:

$${\displaystyle x_{t}={\sqrt {{\bar {\alpha }}_{t}}}x_{0}+\underbrace {{\sqrt {\alpha _{t}-{\bar {\alpha }}_{t}}}z+{\sqrt {1-\alpha _{t}}}z'} _{={\sqrt {1-{\bar {\alpha }}_{t}}}z''}}$$

where ${\textstyle z''}$ is a gaussian with mean zero and variance one. To find the second one, we complete the rotational matrix:

$${\displaystyle {\begin{bmatrix}z''\\z'''\end{bmatrix}}={\begin{bmatrix}{\frac {\sqrt {\alpha _{t}-{\bar {\alpha }}_{t}}}{\sqrt {1-{\bar {\alpha }}_{t}}}}&{\frac {\sqrt {\beta _{t}}}{\sqrt {1-{\bar {\alpha }}_{t}}}}\\?&?\end{bmatrix}}{\begin{bmatrix}z\\z'\end{bmatrix}}}$$

Since rotational matrices are all of the form ${\textstyle {\begin{bmatrix}\cos \theta &\sin \theta \\-\sin \theta &\cos \theta \end{bmatrix}}}$, we know the matrix must be

$${\displaystyle {\begin{bmatrix}z''\\z'''\end{bmatrix}}={\begin{bmatrix}{\frac {\sqrt {\alpha _{t}-{\bar {\alpha }}_{t}}}{\sqrt {1-{\bar {\alpha }}_{t}}}}&{\frac {\sqrt {\beta _{t}}}{\sqrt {1-{\bar {\alpha }}_{t}}}}\\-{\frac {\sqrt {\beta _{t}}}{\sqrt {1-{\bar {\alpha }}_{t}}}}&{\frac {\sqrt {\alpha _{t}-{\bar {\alpha }}_{t}}}{\sqrt {1-{\bar {\alpha }}_{t}}}}\end{bmatrix}}{\begin{bmatrix}z\\z'\end{bmatrix}}}$$

and since the inverse of rotational matrix is its transpose,

$${\displaystyle {\begin{bmatrix}z\\z'\end{bmatrix}}={\begin{bmatrix}{\frac {\sqrt {\alpha _{t}-{\bar {\alpha }}_{t}}}{\sqrt {1-{\bar {\alpha }}_{t}}}}&-{\frac {\sqrt {\beta _{t}}}{\sqrt {1-{\bar {\alpha }}_{t}}}}\\{\frac {\sqrt {\beta _{t}}}{\sqrt {1-{\bar {\alpha }}_{t}}}}&{\frac {\sqrt {\alpha _{t}-{\bar {\alpha }}_{t}}}{\sqrt {1-{\bar {\alpha }}_{t}}}}\end{bmatrix}}{\begin{bmatrix}z''\\z'''\end{bmatrix}}}$$

Plugging back, and simplifying, we have

$${\displaystyle x_{t}={\sqrt {{\bar {\alpha }}_{t}}}x_{0}+{\sqrt {1-{\bar {\alpha }}_{t}}}z''}$$

$${\displaystyle x_{t-1}={\tilde {\mu }}_{t}(x_{t},x_{0})-{\sqrt {{\tilde {\beta }}_{t}}}z'''}$$

Backward diffusion

The key idea of DDPM is to use a neural network parametrized by ${\displaystyle \theta }$. The network takes in two arguments ${\displaystyle x_{t},t}$, and outputs a vector ${\displaystyle \mu _{\theta }(x_{t},t)}$ and a matrix ${\displaystyle \Sigma _{\theta }(x_{t},t)}$, such that each step in the forward diffusion process can be approximately undone by ${\displaystyle x_{t-1}\sim N(\mu _{\theta }(x_{t},t),\Sigma _{\theta }(x_{t},t))}$. This then gives us a backward diffusion process ${\displaystyle p_{\theta }}$ defined by

$${\displaystyle p_{\theta }(x_{T})=N(x_{T}|0,I)}$$

$${\displaystyle p_{\theta }(x_{t-1}|x_{t})=N(x_{t-1}|\mu _{\theta }(x_{t},t),\Sigma _{\theta }(x_{t},t))}$$

The goal now is to learn the parameters such that ${\displaystyle p_{\theta }(x_{0})}$ is as close to ${\displaystyle q(x_{0})}$ as possible. To do that, we use maximum likelihood estimation with variational inference.

Variational inference

The ELBO inequality states that ${\displaystyle \ln p_{\theta }(x_{0})\geq E_{x_{1:T}\sim q(\cdot |x_{0})}[\ln p_{\theta }(x_{0:T})-\ln q(x_{1:T}|x_{0})]}$, and taking one more expectation, we get

$${\displaystyle E_{x_{0}\sim q}[\ln p_{\theta }(x_{0})]\geq E_{x_{0:T}\sim q}[\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 function

$${\displaystyle L(\theta ):=-E_{x_{0:T}\sim q}[\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^[8]

$${\displaystyle L(\theta )=\sum _{t=1}^{T}E_{x_{t-1},x_{t}\sim q}[-\ln p_{\theta }(x_{t-1}|x_{t})]+E_{x_{0}\sim q}[D_{KL}(q(x_{T}|x_{0})\|p_{\theta }(x_{T}))]+C}$$

where ${\displaystyle C}$ does not depend on the parameter, and thus can be ignored. Since ${\displaystyle p_{\theta }(x_{T})=N(x_{T}|0,I)}$ also does not depend on the parameter, the term ${\displaystyle E_{x_{0}\sim q}[D_{KL}(q(x_{T}|x_{0})\|p_{\theta }(x_{T}))]}$ can also be ignored. This leaves just ${\displaystyle L(\theta )=\sum _{t=1}^{T}L_{t}}$ with ${\displaystyle L_{t}=E_{x_{t-1},x_{t}\sim q}[-\ln p_{\theta }(x_{t-1}|x_{t})]}$ to be minimized.

Noise prediction network

Since ${\displaystyle x_{t-1}|x_{t},x_{0}\sim N({\tilde {\mu }}_{t}(x_{t},x_{0}),{\tilde {\beta }}_{t}I)}$, this suggests that we should use ${\displaystyle \mu _{\theta }(x_{t},t)={\tilde {\mu }}_{t}(x_{t},x_{0})}$; however, the network does not have access to ${\displaystyle x_{0}}$, and so it has to estimate it instead. Now, since ${\displaystyle x_{t}|x_{0}\sim N\left({\sqrt {{\bar {\alpha }}_{t}}}x_{0},(1-{\bar {\alpha }}_{t})I\right)}$, we may write ${\displaystyle x_{t}={\sqrt {{\bar {\alpha }}_{t}}}x_{0}+{\sqrt {1-{\bar {\alpha }}_{t}}}z}$, where ${\displaystyle z}$ is some unknown gaussian noise. Now we see that estimating ${\displaystyle x_{0}}$ is equivalent to estimating ${\displaystyle z}$.

Therefore, let the network output a noise vector ${\displaystyle \epsilon _{\theta }(x_{t},t)}$, and let it predict

$${\displaystyle \mu _{\theta }(x_{t},t)={\tilde {\mu }}_{t}\left(x_{t},{\frac {x_{t}-{\sqrt {1-{\bar {\alpha }}_{t}}}\epsilon _{\theta }(x_{t},t)}{\sqrt {{\bar {\alpha }}_{t}}}}\right)={\frac {x_{t}-\epsilon _{\theta }(x_{t},t)\beta _{t}/{\sqrt {1-{\bar {\alpha }}_{t}}}}{\sqrt {\alpha _{t}}}}}$$

It remains to design ${\displaystyle \Sigma _{\theta }(x_{t},t)}$. The DDPM paper suggested not learning it (since it resulted in "unstable training and poorer sample quality"), but fixing it at some value ${\displaystyle \Sigma _{\theta }(x_{t},t)=\sigma _{t}^{2}I}$, where either ${\displaystyle \sigma _{t}^{2}=\beta _{t}{\text{ or }}{\tilde {\beta }}_{t}}$ yielded similar performance. With this, the loss simplifies to

$${\displaystyle L_{t}={\frac {\beta _{t}^{2}}{2\alpha _{t}(1-{\bar {\alpha }}_{t})\sigma _{t}^{2}}}E_{x_{0}\sim q;z\sim N(0,I)}\left[\left\|\epsilon _{\theta }(x_{t},t)-z\right\|^{2}\right]+C}$$

which may be minimized by stochastic gradient descent. The paper noted empirically that an even simpler loss function

$${\displaystyle L_{simple,t}=E_{x_{0}\sim q;z\sim N(0,I)}\left[\left\|\epsilon _{\theta }(x_{t},t)-z\right\|^{2}\right]}$$

resulted in better models.

The idea of score functions

Consider the problem of image generation. Let ${\displaystyle x}$ represent an image, and let ${\displaystyle q(x)}$ be the probability distribution over all possible images. If we have ${\displaystyle 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 ${\displaystyle q(x)}$ itself, but rather, ${\displaystyle \nabla _{x}\ln q(x)}$. This has two major effects:

• One, we no longer need to normalize ${\displaystyle q(x)}$, but can use any ${\displaystyle {\tilde {q}}(x)=Cq(x)}$, where ${\displaystyle C=\int {\tilde {q}}(x)dx>0}$ is any unknown constant that is of no concern to us.
• Two, we are comparing ${\displaystyle q(x)}$ neighbors ${\displaystyle q(x+dx)}$, by ${\displaystyle {\frac {q(x)}{q(x+dx)}}=e^{-\langle \nabla _{x}\ln p,dx\rangle }}$

Let the score function be ${\displaystyle s(x):=\nabla _{x}\ln q(x)}$; then consider what we can do with ${\displaystyle s(x)}$.

As it turns out, ${\displaystyle s(x)}$ allows us to sample from ${\displaystyle q(x)}$ using thermodynamics. Specifically, if we have a potential energy function ${\displaystyle U(x)=-\ln q(x)}$, and a lot of particles in the potential well, then the distribution at thermodynamic equilibrium is the Boltzmann distribution ${\displaystyle q_{U}(x)\propto e^{-U(x)/k_{B}T}=q(x)^{1/k_{B}T}}$. At temperature ${\displaystyle k_{B}T=1}$, the Boltzmann distribution is exactly ${\displaystyle q(x)}$.

Therefore, to model ${\displaystyle 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 equation

$${\displaystyle dx_{t}=-\nabla _{x_{t}}U(x_{t})dt+dW_{t}}$$

and the Boltzmann distribution is, by Fokker-Planck equation, the unique thermodynamic equilibrium. So no matter what distribution ${\displaystyle x_{0}}$ has, the distribution of ${\displaystyle x_{t}}$ converges in distribution to ${\displaystyle q}$ as ${\displaystyle t\to \infty }$.

Learning the score function

Given a density ${\displaystyle q}$, we wish to learn a score function approximation ${\displaystyle f_{\theta }\approx \nabla \ln q}$. This is score matching.^[11] Typically, score matching is formalized as minimizing Fisher divergence function ${\displaystyle E_{q}[\|f_{\theta }(x)-\nabla \ln q(x)\|^{2}]}$. By expanding the integral, and performing an integration by parts,

$${\displaystyle E_{q}[\|f_{\theta }(x)-\nabla \ln q(x)\|^{2}]=E_{q}[\|f_{\theta }\|^{2}+2\nabla ^{2}\cdot f_{\theta }]+C}$$

giving 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 ${\displaystyle x_{0}\sim N(0,I)}$, a white-noise image. Now, most white-noise images do not look like real images, so ${\displaystyle q(x_{0})\approx 0}$ for large swaths of ${\displaystyle x_{0}\sim N(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 ${\displaystyle \nabla _{x_{t}}\ln q(x_{t})}$ at that point, then we cannot impose the time-evolution equation on a particle:

$${\displaystyle dx_{t}=\nabla _{x_{t}}\ln q(x_{t})dt+dW_{t}}$$

To deal with this problem, we perform annealing. If ${\displaystyle 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.

Forward diffusion process

Consider again the forward diffusion process, but this time in continuous time:

$${\displaystyle x_{t}={\sqrt {1-\beta _{t}}}x_{t-1}+{\sqrt {\beta _{t}}}z_{t}}$$

By taking the ${\displaystyle \beta _{t}\to \beta (t)dt,{\sqrt {dt}}z_{t}\to dW_{t}}$ limit, we obtain a continuous diffusion process, in the form of a stochastic differential equation:

$${\displaystyle dx_{t}=-{\frac {1}{2}}\beta (t)x_{t}dt+{\sqrt {\beta (t)}}dW_{t}}$$

where ${\displaystyle W_{t}}$ is a Wiener process (multidimensional Brownian motion). Now, the equation is exactly a special case of the overdamped Langevin equation

$${\displaystyle dx_{t}=-{\frac {D}{k_{B}T}}(\nabla _{x}U)dt+{\sqrt {2D}}dW_{t}}$$

where ${\displaystyle D}$ is diffusion tensor, ${\displaystyle T}$ is temperature, and ${\displaystyle U}$ is potential energy field. If we substitute in ${\displaystyle D={\frac {1}{2}}\beta (t)I,k_{B}T=1,U={\frac {1}{2}}\|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 ${\displaystyle q}$ at time ${\displaystyle t=0}$, then after a long time, the cloud of particles would settle into the stable distribution of ${\displaystyle N(0,I)}$. Let ${\displaystyle \rho _{t}}$ be the density of the cloud of particles at time ${\displaystyle t}$, then we have

$${\displaystyle \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

$${\displaystyle \partial _{t}\ln \rho _{t}={\frac {1}{2}}\beta (t)\left(n+(x+\nabla \ln \rho _{t})\cdot \nabla \ln \rho _{t}+\Delta \ln \rho _{t}\right)}$$

where ${\displaystyle n}$ is the dimension of space, and ${\displaystyle \Delta }$ is the Laplace operator.

Backward diffusion process

If we have solved ${\displaystyle \rho _{t}}$ for time ${\displaystyle t\in [0,T]}$, then we can exactly reverse the evolution of the cloud. Suppose we start with another cloud of particles with density ${\displaystyle \nu _{0}=\rho _{T}}$, and let the particles in the cloud evolve according to

$${\displaystyle dy_{t}={\frac {1}{2}}\beta (T-t)y_{t}dt+\beta (T-t)\underbrace {\nabla _{y_{t}}\ln \rho _{T-t}\left(y_{t}\right)} _{\text{score function }}dt+{\sqrt {\beta (T-t)}}dW_{t}}$$

then by plugging into the Fokker-Planck equation, we find that ${\displaystyle \partial _{t}\rho _{T-t}=\partial _{t}\nu _{t}}$. Thus this cloud of points is the original cloud, evolving backwards.^[12]