Diffusion Primer

[Sneha Bangalore]

Diffusion, a class of generative models, boils down to an MSE between added noise and network predicted noise.
Why would learning to predict the noise help with image generation (which is what it is most used for)? How did we arrive at MSE? This post dives deep into the math to answer these questions.

Background

One way to interpret Diffusion is as a continuous VAE (Variational Auto Encoder).
A VAE computes a lower bound on the likelihood of generating real data samples ($\log p_{\theta }\left(x\right)$) by approximating the unknown posterior $p_{\theta }\left(z|x\right)$ with a learnable one $q_{\varphi }\left(z|x\right)$ [Fig. 1]:

Figure 1: VAE

$\begin{array}{ccc}-\log p_{\theta }\left(x\right)& \le -\log p_{\theta }\left(x\right)& \\ & \le -\log p_{\theta }\left(x\right)+D_{KL}(q_{\varphi }\left(z|x\right)\parallel p_{\theta }\left(z|x\right))& \text{[KL is always positive]}\\ & \le -\log p_{\theta }\left(x\right)+\int q_{\varphi }\left(z|x\right)\log \left(\frac{q_{\varphi }\left(z|x\right)}{p_{\theta }\left(z|x\right)}\right)dz& \text{[Definition of KL]}\\ & \le -\log p_{\theta }\left(x\right)+\int q_{\varphi }\left(z|x\right)\log \left(\frac{q_{\varphi }\left(z|x\right)p_{\theta }\left(x\right)}{p_{\theta }(z,x)}\right)dz& \text{[conditional to joint]}\\ & \le -\log p_{\theta }\left(x\right)+\int q_{\varphi }\left(z|x\right)(\log p_{\theta }\left(x\right)+\log \left(\frac{q_{\varphi }\left(z|x\right)}{p_{\theta }(z,x)}\right))dz& \\ & \le -\log p_{\theta }\left(x\right)+\log p_{\theta }\left(x\right)+\int q_{\varphi }\left(z|x\right)\left(\log \left(\frac{q_{\varphi }\left(z|x\right)}{p_{\theta }\left(x|z\right)p_{\theta }\left(z\right)}\right)\right)dz& \text{[}{p}_{\theta }\text{independent of}z\text{; joint to conditional]}\\ & =-E_{z_{\sim }{q}_{\varphi }\left(z|x\right)}[\log \left(\frac{q_{\varphi }\left(z|x\right)}{p_{\theta }\left(z\right)}\right)-\log p_{\theta }\left(x|z\right)]& \text{[Definition of}E\text{for continuous variable}z]\\ & =-E_{z_{\sim }{q}_{\varphi }\left(z|x\right)}\log p_{\theta }\left(x|z\right)+D_{KL}(q_{\varphi }\left(z|x\right)\parallel p_{\theta }\left(z\right))& \text{[Tractable]}\end{array}$

The process of encoding and decoding leads to regularized compression of images in the $z$ space (hence the name auto encoder; variational comes from approximating distributions). During inference, random images can be generated from the decoder, based on sampling of $z$.

Derivation

Figure 2: Markov chain of forward [$q$] (reverse [$p_{\theta }$]) diffusion process of generating a sample by slowly adding (removing) noise. (Image source: Ho et al. 2020)

Similar to VAE, diffusion models compute a lower bound on the likelihood of generating real data samples ($\log p_{\theta }\left(x_0\right)$) by approximating the unknown posterior $p_{\theta }\left(x_{t-1}|x_t\right)$ with a computable one $q_{\varphi }(x_{t-1}|x_t,x_0)$, for all $t$ [Fig. 2] (It follows the same math as VAE for all $t$, therefore could be thought of as its continuous version):

$\begin{array}{ccc}-\log p_{\theta }\left(x_0\right)& \le -\log p_{\theta }\left(x_0\right)& \\ & \le -\log p_{\theta }\left(x_0\right)+D_{KL}(q\left(x_{1:T}|x_0\right)\parallel p_{\theta }\left(x_{1:T}|x_0\right))& \text{[KL is always positive]}\\ & =-\log p_{\theta }\left(x_0\right)+E_{x_{1:T}\sim q\left(x_{1:T}|x_0\right)}\left[\log \left(\frac{q\left(x_{1:T}|x_0\right)}{p_{\theta }\left(x_{1:T}\right)|x_0}\right)\right]& \\ & =-\log p_{\theta }\left(x_0\right)+E_{x_{1:T}\sim q\left(x_{1:T}|x_0\right)}\left[\log \left(\frac{q\left(x_{1:T}|x_0\right)}{\left(\frac{p_{\theta }\left(x_0|x_{1:T}\right)p_{\theta }\left(x_{1:T}\right)}{p_{\theta }\left(x_0\right)}\right)}\right)\right]& \text{[Bayes' rule]}\\ & =-\log p_{\theta }\left(x_0\right)+E_{x_{1:T}\sim q\left(x_{1:T}|x_0\right)}\left[\log \left(\frac{q\left(x_{1:T}|x_0\right)}{\left(\frac{p_{\theta }(x_0,x_{1:T})}{p_{\theta }\left(x_0\right)}\right)}\right)\right]& \text{[conditional to joint]}\\ & =-\log p_{\theta }\left(x_0\right)+E_{x_{1:T}\sim q\left(x_{1:T}|x_0\right)}\left[\log \left(\frac{q\left(x_{1:T}|x_0\right)}{\left(\frac{p_{\theta }\left(x_{0:T}\right)}{p_{\theta }\left(x_0\right)}\right)}\right)\right]& \text{[merge joint]}\\ & =-\log p_{\theta }\left(x_0\right)+E_{x_{1:T}\sim q\left(x_{1:T}|x_0\right)}\left[\log \left(\frac{q\left(x_{1:T}|x_0\right)}{p_{\theta }\left(x_{0:T}\right)}\right)\right]+\log \left(p_{\theta }\left(x_0\right)\right)& \\ E_q[-\log \left(p_{\theta }\left(x_0\right)\right)]& \le E_q\left[\log \left(\frac{q\left(x_{1:T}|x_0\right)}{p_{\theta }\left(x_{0:T}\right)}\right)\right]& \end{array}$

Solving for R.H.S.:

$\begin{array}{cc}\log \left(\frac{q\left(x_{1:T}|x_0\right)}{p_{\theta }\left(x_{0:T}\right)}\right)& =\log \left(\frac{{\prod }_{t=1}^{T}q\left(x_t|x_{t-1}\right)}{p_{\theta }\left(x_T\right){\prod }_{t=1}^T{p}_{\theta }\left(x_{t-1}|x_t\right)}\right)\\ & =-\log p_{\theta }\left(x_T\right)+\sum _{t=1}^T\log \left(\frac{q\left(x_t|x_{t-1}\right)}{p_{\theta }\left(x_{t-1}|x_t\right)}\right)\\ & =-\log p_{\theta }\left(x_T\right)+\sum _{t=2}^T\log \left(\frac{q\left(x_t|x_{t-1}\right)}{p_{\theta }\left(x_{t-1}|x_t\right)}\right)+\log \left(\frac{q\left(x_1|x_0\right)}{p_{\theta }\left(x_0|x_1\right)}\right)\\ & =-\log p_{\theta }\left(x_T\right)+\sum _{t=2}^T\log \left(\frac{q\left(x_{t-1}|x_t\right)q\left(x_t\right)}{p_{\theta }\left(x_{t-1}|x_t\right)q\left(x_{t-1}\right)}\right)+\log \left(\frac{q\left(x_1|x_0\right)}{p_{\theta }\left(x_0|x_1\right)}\right)\\ & =-\log p_{\theta }\left(x_T\right)+\sum _{t=2}^T\log \left(\frac{q(x_{t-1}|x_t,x_0)q\left(x_t|x_0\right)}{p_{\theta }\left(x_{t-1}|x_t\right)q\left(x_{t-1}|x_0\right)}\right)+\log \left(\frac{q\left(x_1|x_0\right)}{p_{\theta }\left(x_0|x_1\right)}\right)\text{[conditioning on}{x}_0\text{for tractability]}\\ & =-\log p_{\theta }\left(x_T\right)+\sum _{t=2}^T\log \left(\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }\left(x_{t-1}|x_t\right)}\right)+\sum _{t=2}^T\log \left(\frac{q\left(x_t|x_0\right)}{q\left(x_{t-1}|x_0\right)}\right)+\log \left(\frac{q\left(x_1|x_0\right)}{p_{\theta }\left(x_0|x_1\right)}\right)\\ & =-\log p_{\theta }\left(x_T\right)+\sum _{t=2}^T\log \left(\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }\left(x_{t-1}|x_t\right)}\right)+\log \left(\frac{q\left(x_2|x_0\right)\dots q\left(x_{T-1}|x_0\right)q\left(x_T|x_0\right)}{q\left(x_1|x_0\right)q\left(x_2|x_0\right)\dots q\left(x_{T-1}|x_0\right)}\right)+\log \left(\frac{q\left(x_1|x_0\right)}{p_{\theta }\left(x_0|x_1\right)}\right)\\ & =-\log p_{\theta }\left(x_T\right)+\sum _{t=2}^T\log \left(\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }\left(x_{t-1}|x_t\right)}\right)+\log \left(\frac{q\left(x_T|x_0\right)}{q\left(x_1|x_0\right)}\right)+\log \left(\frac{q\left(x_1|x_0\right)}{p_{\theta }\left(x_0|x_1\right)}\right)\\ & =-\log p_{\theta }\left(x_T\right)+\sum _{t=2}^T\log \left(\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }\left(x_{t-1}|x_t\right)}\right)+\log \left(\frac{q\left(x_T|x_0\right)q\left(x_1|x_0\right)}{q\left(x_1|x_0\right)p_{\theta }\left(x_0|x_1\right)}\right)\\ & =\log q\left(x_T|x_0\right)-\log p_{\theta }\left(x_T\right)+\sum _{t=2}^T\log \left(\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }\left(x_{t-1}|x_t\right)}\right)-\log p_{\theta }\left(x_0|x_1\right)\\ & =\log \left(\frac{q\left(x_T|x_0\right)}{p_{\theta }\left(x_T\right)}\right)+\sum _{t=2}^T\log \left(\frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }\left(x_{t-1}|x_t\right)}\right)-\log p_{\theta }\left(x_0|x_1\right)\\ \log \left(\frac{q\left(x_{1:T}|x_0\right)}{p_{\theta }\left(x_{0:T}\right)}\right)& =\underset{L_T}{\underset{}{D_{KL}(q\left(x_T|x_0\right)\parallel p_{\theta }\left(x_T\right))}}+\sum _{t=2}^T\underset{L_{t-1}}{\underset{}{D_{KL}(q(x_{t-1}|x_t,x_0)\parallel p_{\theta }\left(x_{t-1}|x_t\right))}}\underset{L_0}{\underset{}{-\log p_{\theta }\left(x_0|x_1\right)}}\end{array}$

Taking each piece apart:

$L_T$: Can be ignored while training since $q$ has no learnable parameters, and $x_T$ is just Gaussian noise.

$L_{t-1}$: Match the unknown ($p_{\theta }$) to the tractable quantity ($q$). Define the reverse process as:

$p_{\theta }\left(x_{t-1}|x_t\right)=N(x_{t-1};{\mu }_{\theta }(x_t,t),{\Sigma }_{\theta }(x_t,t))$

Simplifying to known variance:

$p_{\theta }\left(x_{t-1}|x_t\right)=N(x_{t-1};{\mu }_{\theta }(x_t,t),{\sigma }_t^{2}I)$

Solving $q$ (zoom out if not fully visible):

	Scratch Work
$\begin{array}{cc}q(x_{t-1}|x_t,x_0)& =N(x_t;{\tilde{\mu }}_t(x_t,x_0),{\tilde{\beta }}_{t}I)\\ q(x_{t-1}|x_t,x_0)& =N(x_t;\sqrt{{\alpha }_t}{x}_{t-1},{\beta }_{t}I)\left(\frac{N(x_{t-1};\sqrt{{\overline{\alpha }}_{t-1}}{x}_0,1-{\overline{\alpha }}_{t-1}I)}{N(x_t;\sqrt{{\overline{\alpha }}_t}{x}_0,1-{\overline{\alpha }}_{t}I)}\right)\end{array}$	$\begin{array}{ccc}q(x_{t-1}|x_t,x_0)& =q(x_t|x_{t-1},x_0)\left(\frac{q\left(x_{t-1}|x_0\right)}{q\left(x_t|x_0\right)}\right)& \text{[Bayes' rule]}\\ q\left(x_{1:T}|x_0\right)& =\prod _{t=1}^{T}q\left(x_t|x_{t-1}\right)& \text{[Gaussian noising]}\\ {\beta }_t& =1-{\alpha }_t& \text{[Variance schedule]}\\ \therefore q\left(x_t|x_{t-1}\right)& =N(x_t;\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_{t}I)& \\ x_t& =\sqrt{{\alpha }_t}{x}_{t-1}+\sqrt{1-{\alpha }_t}{ϵ}_{t-1}& \text{[Reparameterization trick]}\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{\overline{ϵ}}_{t-2}& \text{[}{\overline{ϵ}}_{t-2}\text{merges two Gaussians]}\\ & =\cdots & \\ & =\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}ϵ& \text{[}{\overline{\alpha }}_t=\prod _{i=1}^t{\alpha }_i\text{;}ϵ\sim N(0,I)\text{]}\\ \therefore q\left(x_n|x_0\right)& =N(x_n;\sqrt{{\overline{\alpha }}_n}{x}_0,(1-{\overline{\alpha }}_n)I)& \\ \text{Also,}{x}_0& =\frac{1}{\sqrt{{\overline{\alpha }}_t}}(x_t-\sqrt{1-{\overline{\alpha }}_t}ϵ)& \text{[1]}\end{array}$
$\begin{array}{cc}q(x_{t-1}|x_t,x_0)& \propto \exp (-\frac{1}{2}[\frac{{(x_t-\sqrt{{\alpha }_t}{x}_{t-1})}^2}{{\beta }_t}+\frac{{(x_{t-1}-\sqrt{{\overline{\alpha }}_{t-1}}{x}_0)}^2}{1-{\overline{\alpha }}_{t-1}}-\frac{{(x_t-\sqrt{{\overline{\alpha }}_t}{x}_0)}^2}{1-{\overline{\alpha }}_t}])\\ & \propto \exp (-\frac{1}{2}[\underset{1/{\sigma }^2}{\underset{}{(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\overline{\alpha }}_{t-1}})}}{x}_{t-1}^2-\underset{2\mu /{\sigma }^2}{\underset{}{(\frac{2\sqrt{{\alpha }_t}}{{\beta }_t}{x}_t+\frac{2\sqrt{{\overline{\alpha }}_{t-1}}}{1-{\overline{\alpha }}_{t-1}}{x}_0)}}{x}_{t-1}+\underset{{\mu }^2/{\sigma }^2}{\underset{}{f(x_t,x_0)}}])\end{array}$	$\begin{array}{cc}N\left(x\right)& \propto \exp (-\frac{1}{2}\left[\frac{{(x-\mu )}^2}{{\sigma }^2}\right])\\ & =\exp (-\frac{1}{2}\left[\frac{x^2-2\mu x+{\mu }^2}{{\sigma }^2}\right])\\ & =\exp (-\frac{1}{2}[\frac{1}{{\sigma }^2}{x}^2-2\frac{\mu }{{\sigma }^2}x+\frac{{\mu }^2}{{\sigma }^2}])\end{array}$

Computing the parameters ($\sigma$ and $\mu$) of the new Gaussian:

$\begin{array}{ccc}{\sigma }^2={\tilde{\beta }}_t& =\frac{1}{(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\overline{\alpha }}_{t-1}})}=\left(\frac{1-{\overline{\alpha }}_{t-1}}{1-{\overline{\alpha }}_t}\right){\beta }_t& \\ {\tilde{\mu }}_t(x_t,x_0)& =(\frac{\sqrt{{\alpha }_t}}{{\beta }_t}{x}_t+\frac{\sqrt{{\overline{\alpha }}_{t-1}}}{1-{\overline{\alpha }}_{t-1}}{x}_0)\frac{1}{{\sigma }^2}& \\ & =(\frac{\sqrt{{\alpha }_t}}{{\beta }_t}{x}_t+\frac{\sqrt{{\overline{\alpha }}_{t-1}}}{1-{\overline{\alpha }}_{t-1}}{x}_0)\left(\frac{1-{\overline{\alpha }}_{t-1}}{1-{\overline{\alpha }}_t}\right){\beta }_t& \\ & =\frac{\sqrt{{\alpha }_t}(1-{\overline{\alpha }}_{t-1})}{1-{\overline{\alpha }}_t}{x}_t+\frac{\sqrt{{\overline{\alpha }}_{t-1}}{\beta }_t}{1-{\overline{\alpha }}_t}{x}_0& \\ {\tilde{\mu }}_t(x_t,x_0)& =\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\overline{\alpha }}_t}}{ϵ}_t)& \text{[From 1]}\end{array}$

${\mu }_{\theta }$ must match ${\tilde{\mu }}_t$, and since $x_t$ is given, learn to predict the noise at step $t$, ${\epsilon }_t$.

$\begin{array}{c}{\mu }_{\theta }(x_t,t)=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\overline{\alpha }}_t}}{ϵ}_{\theta }(x_t,t))\text{[2]}\end{array}$

Which means
$\begin{array}{c}{x}_{t-1}=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\overline{\alpha }}_t}}{ϵ}_{\theta }(x_t,t))+{\sigma }_{t}z\text{[3]}\end{array}$

The KL divergence between Gaussians with mean as the only parameter leads to the following loss:

$\begin{array}{cc}{L}_t& =D_{KL}(q(x_{t-1}|x_t,x_0)\parallel p_{\theta }\left(x_{t-1}|x_t\right))\\ & =\frac{1}{2{\sigma }_t^2}{\parallel {\tilde{\mu }}_t(x_t,x_0)-{\mu }_{\theta }(x_t,t)\parallel }^2\\ & =\frac{1}{2{\sigma }_t^2}{\Vert \frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\overline{\alpha }}_t}}{ϵ}_t)-\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\overline{\alpha }}_t}}{ϵ}_{\theta }(x_t,t))\Vert }^2\\ & =\frac{{(1-{\alpha }_t)}^2}{2{\alpha }_t(1-{\overline{\alpha }}_t){\sigma }_t^2}{\parallel {ϵ}_t-{ϵ}_{\theta }(x_t,t)\parallel }^2\end{array}$

In the DDPM paper^[1], it is found empirically that the training works better if the scaling factor is omitted, i.e., the loss can be simplified to

$\begin{array}{c}{L}_t={\parallel {ϵ}_t-{ϵ}_{\theta }(x_t,t)\parallel }^2\end{array}$

$L_0$: The image ($x_0$) is originally constructed from D $\text{discrete}$ pixels with values $[0,1,2,\dots ,255]$ scaled to the range $[-1,1]$. The DDPM paper^[1] suggests using an independent discrete decoder derived from the Gaussian (since $L_t$ doesn’t apply here):

$p_{\theta }\left(x_0|x_1\right)={\prod }_{i=1}^D{\int }_{{\delta }_{-}\left(x_0^i\right)}^{{\delta }_{+}\left(x_0^i\right)}N(x_1,{\mu }_{\theta }^i(x_1,1),{\sigma }_1^2)dx$

${\delta }_{+}\left(x\right)=\{\begin{array}{cc}\infty & \text{if}x=1\\ x+\frac{1}{255}& \text{if}x<1\end{array}{\delta }_{-}\left(x\right)=\{\begin{array}{cc}-\infty & \text{if}x=-1\\ x-\frac{1}{255}& \text{if}x>-1\end{array}$

Intuitively, the network predicts the mean of a pixel which is used to draw from a Gaussian distribution which is integrated^[2] over according to the real pixel value as in the original image, from the real value minus $1/255$ to the real value plus $1/255$. If the predicted mean is close to the original value of the pixel, the result of the integral will be high.

The integral is approximated by the Gaussian probability density function times the bin width, which makes the math workout similar to the $L_{t-1}$ case:$\begin{array}{cc}{p}_{\theta }\left(x_0|x_1\right)& \approx \frac{1}{\sqrt{2\pi }{\sigma }_1}\exp (-\frac{1}{2}\frac{{\parallel x_0-{\mu }_{\theta }(x_1,1)\parallel }^2}{{\sigma }_1^2})\left(\frac{2}{255}\right)\\ L_0=-\log \left(p_{\theta }\left(x_0|x_1\right)\right)& \approx \frac{1}{2{\sigma }_1^2}{\parallel x_0-{\mu }_{\theta }(x_1,1)\parallel }^2+C\\ & \approx \frac{1}{2{\sigma }_1^2}{\Vert \frac{1}{\sqrt{{\alpha }_1}}(x_1-\sqrt{1-{\alpha }_1}{ϵ}_1)-\frac{1}{\sqrt{{\alpha }_1}}(x_1-\frac{1-{\alpha }_1}{\sqrt{1-{\alpha }_1}}{ϵ}_{\theta }(x_1,1))\Vert }^2+C\text{[From 1, 2]}\\ & \approx {\parallel {ϵ}_1-{ϵ}_{\theta }(x_1,1)\parallel }^2\end{array}$

Putting the pieces together, the DDPM paper^[1] suggests using this loss function for training (followed by the algorithm):
$\begin{array}{cc}{L}_{\text{simple}}& :=E_{t\sim U[1,T],x_0,ϵ}\left[{\parallel ϵ-{ϵ}_{\theta }(x_t,t)\parallel }^2\right]\\ & =E_{t\sim U[1,T],x_0,ϵ}\left[{\Vert ϵ-{ϵ}_{\theta }(\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}ϵ,t)\Vert }^2\right]\end{array}$

Algorithm 1 DDPM Training
$\begin{array}{c}\text{repeat}\\ x_0\sim q\left(x_0\right)\\ t\sim \mathrm{Uniform}\left([1,\dots ,T]\right)\\ ϵ\sim N(0,I)\\ \text{Take gradient descent step on}\\ {\nabla }_{\theta }{\Vert ϵ-{ϵ}_{\theta }(\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}ϵ,t)\Vert }^2\\ \text{until}\text{converged}\end{array}$

We saw how we can go from noise to an image via the reverse process, but how do we go from random noise $\to$ random (aggregate of training set) image to random noise $\to$ image of interest? The next section talks about this.

Conditioning

Dhariwal & Nichol (2021)^[3] use gradients ${\nabla }_x\log f_{\varphi }\left(y|x_t\right)$ of a classifier $f_{\varphi }(y|x_t,t)$ to condition diffusion sampling on $y$ (e.g. a target class label) where $x_t$ is noisy image. The score function for the joint distribution $q(x_t,y)$ is computed as,

$\begin{array}{cc}{\nabla }_{x_t}\log q(x_t,y)& ={\nabla }_{x_t}\log q\left(x_t\right)+{\nabla }_{x_t}\log q\left(y|x_t\right)\\ & \approx -\frac{1}{\sqrt{1-{\overline{\alpha }}_t}}{ϵ}_{\theta }(x_t,t)+{\nabla }_{x_t}\log f_{\varphi }\left(y|x_t\right)\\ & =-\frac{1}{\sqrt{1-{\overline{\alpha }}_t}}({ϵ}_{\theta }(x_t,t)-\sqrt{1-{\overline{\alpha }}_t}{\nabla }_{x_t}\log f_{\varphi }\left(y|x_t\right))\end{array}$

Thus, a new classifier-guided predictor ${\overline{ϵ}}_{\theta }$ modifies the initial predictor ${ϵ}_{\theta }$ as follows,

${\overline{ϵ}}_{\theta }(x_t,t)={ϵ}_{\theta }(x_t,t)-\sqrt{1-{\overline{\alpha }}_t}{\nabla }_{x_t}\log f_{\varphi }\left(y|x_t\right)$

To control the strength of the classifier guidance, a weight $w$ is added to the delta part,
$\begin{array}{cc}{\overline{ϵ}}_{\theta }(x_t,t)={ϵ}_{\theta }(x_t,t)-\sqrt{1-{\overline{\alpha }}_t}w{\nabla }_{x_t}\log f_{\varphi }\left(y|x_t\right)& \text{[4]}\end{array}$

Ho & Salimans (2021)^[4] propose to run conditional diffusion without an independent classifier $f_{\varphi }$. Consider an unconditional denoising diffusion model $p_{\theta }\left(x\right)$ parameterized through a score estimator ${ϵ}_{\theta }(x_t,t)$ and a conditional model $p_{\theta }\left(x|y\right)$ parameterized through ${ϵ}_{\theta }(x_t,t,y)$. These two models can be learned via a single neural network. Specifically, a conditional diffusion model $p_{\theta }\left(x|y\right)$ is trained on paired data $(x,y)$, where the conditioning information $y$ gets discarded periodically at random such that the model knows how to generate images unconditionally as well, i.e. ${ϵ}_{\theta }(x_t,t)={ϵ}_{\theta }(x_t,t,y=\varnothing )$.

The gradient of an implicit classifier $p\left(y|x_t\right)$ is proportional to $p\left(x_t|y\right)p\left(x_t\right)$. Taking log on both sides,

$\begin{array}{ccc}{\nabla }_{x_t}\log p\left(y|x_t\right)& ={\nabla }_{x_t}\log p\left(x_t|y\right)-{\nabla }_{x_t}\log p\left(x_t\right)& \\ & =-\frac{1}{\sqrt{1-{\overline{\alpha }}_t}}({ϵ}_{\theta }(x_t,t,y)-{ϵ}_{\theta }(x_t,t))& \\ {\overline{ϵ}}_{\theta }(x_t,t,y)& ={ϵ}_{\theta }(x_t,t,y)-\sqrt{1-{\overline{\alpha }}_t}w{\nabla }_{x_t}\log p\left(y|x_t\right)& \text{[From 4]}\\ & ={ϵ}_{\theta }(x_t,t,y)+w({ϵ}_{\theta }(x_t,t,y)-{ϵ}_{\theta }(x_t,t))& \\ & =(w+1){ϵ}_{\theta }(x_t,t,y)-w{ϵ}_{\theta }(x_t,t)& \end{array}$

Below is a comparison of sampling algorithms across DDPM, DDIM and DDIM with classifier guidance.

Algorithm 2 DDPM $\begin{array}{c}\text{Input:}\text{diffusion model}{ϵ}_{\theta }(x_t,t)\\ x_T\sim N(0,I)\\ \text{for}t\text{from}T\text{to 1}\\ z\sim N(0,I)\text{if}t>1\text{, else}z=0\\ x_{t-1}\leftarrow \frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\overline{\alpha }}_t}}{ϵ}_{\theta }(x_t,t))+{\sigma }_{t}z\\ \text{end for}\\ \text{return}{x}_0\end{array}$

Algorithm 3 DDIM $\begin{array}{c}\text{Input:}\text{diffusion model}{ϵ}_{\theta }(x_t,t)\\ x_T\sim N(0,I)\\ \text{for}t\text{from}T\text{to 1}\\ x_{t-1}\leftarrow \sqrt{{\overline{\alpha }}_{t-1}}\left(\frac{x_t-\sqrt{1-{\overline{\alpha }}_t}{ϵ}_{\theta }(x_t,t)}{\sqrt{{\overline{\alpha }}_t}}\right)+\sqrt{1-{\overline{\alpha }}_{t-1}}{ϵ}_{\theta }(x_t,t)\\ \text{end for}\\ \text{return}{x}_0\end{array}$

Algorithm 4 DDIM with classifier guidance $\begin{array}{c}\text{Input:}\text{class label}y\text{, diffusion model}{ϵ}_{\theta }(x_t,t)\text{and classifier}{f}_{\varphi }\left(y|x_t\right)\\ x_T\sim N(0,I)\\ \text{for}t\text{from}T\text{to 1}\\ \widehat{ϵ}\leftarrow {ϵ}_{\theta }(x_t,t)-\sqrt{1-{\overline{\alpha }}_t}{\nabla }_{x_t}\log f_{\varphi }\left(y|x_t\right)\\ x_{t-1}\leftarrow \sqrt{{\overline{\alpha }}_{t-1}}\left(\frac{x_t-\sqrt{1-{\overline{\alpha }}_t}\widehat{ϵ}}{\sqrt{{\overline{\alpha }}_t}}\right)+\sqrt{1-{\overline{\alpha }}_{t-1}}\widehat{ϵ}\\ \text{end for}\\ \text{return}{x}_0\end{array}$

Optimization

So far we discussed training and sampling algorithms, but what is a good value for T? Can we make the sampling process efficient and reduce the number of steps somehow? We explore this part here.

Langevin dynamics can generate samples from a probability density $q\left(x\right)$ using only the gradients ${\nabla }_x\log q\left(x\right)$ in a Markov chain of updates:

$x_t=x_{t-1}+\frac{\delta }{2}{\nabla }_x\log q\left(x_{t-1}\right)+\sqrt{\delta }{ϵ}_t,\text{where}{ϵ}_t\sim N(0,I)$

where $\delta$ is the step size. When $T\to \infty ,ϵ\to 0,x_T$ equals the true probability density $q\left(x\right)$. Compared to SGD, Langevin dynamics injects Gaussian noise into the parameter updates to avoid collapses into local minima.

Score-based generative modeling trains a score network to estimate the above gradient i.e. $s_{\theta }\left(x\right)\approx {\nabla }_x\log q\left(x\right)$. To make the data points cover the whole space, the score network jointly estimates the scores of data perturbed at different noise levels (Noise Conditional Score Network) i.e. $s_{\theta }(x_t,t)\approx {\nabla }_{x_t}\log q\left(x_t\right)$. The schedule of increasing noise levels resembles the forward diffusion process.

Given a Gaussian distribution $x\sim N(\mu ,{\sigma }^{2}I)$, the derivative of the logarithm of its density function ${\nabla }_x\log p\left(x\right)$ is equal to ${\nabla }_x(-\frac{1}{2}{\left[\frac{x-\mu }{\sigma }\right]}^2)=-\left(\frac{x-\mu }{\sigma }\right)=-\frac{ϵ}{\sigma }$ where $ϵ\sim N(0,I)$.

Recall that $q\left(x_t|x_0\right)\sim N(\sqrt{{\overline{\alpha }}_t}{x}_0,(1-{\overline{\alpha }}_t)I)$. Therefore,
$s_{\theta }(x_t,t)\approx {\nabla }_{x_t}\log q\left(x_t\right)=E_{q\left(x_0\right)}\left[{\nabla }_{x_t}\log q\left(x_t|x_0\right)\right]=E_{q\left(x_0\right)}[-\frac{{ϵ}_{\theta }(x_t,t)}{\sqrt{1-{\overline{\alpha }}_t}}]=-\frac{{ϵ}_{\theta }(x_t,t)}{\sqrt{1-{\overline{\alpha }}_t}}$

Figure 3: Non Markov chain of forward [$q$] (reverse [$p_{\theta }$]) diffusion process of generating a sample by slowly adding (removing) noise. (Image source: Song et al. 2022)

To make diffusion more efficient, DDIM [Fig. 3] propose fewer sampling steps by modeling the forward process as a non Markovian chain: $q\left(x_{1:T}|x_0\right)=q\left(x_T|x_0\right){\prod }_{t=2}^{T}q(x_{t-1}|x_t,x_0)$ (DDPM omits the $x_0$ conditioning). Similar to [3], $x_{t-1}$ can be written as
$\begin{array}{cccc}& =\sqrt{{\overline{\alpha }}_{t-1}}{x}_0+\sqrt{(1-{\overline{\alpha }}_{t-1}-{\sigma }_t^2){ϵ}_{\theta }(x_t,t)}+{\sigma }_t{ϵ}_t& & \\ & =\underset{\text{"predicted}{x}_0\text{"}}{\underset{}{\sqrt{{\overline{\alpha }}_{t-1}}\left(\frac{x_t-\sqrt{1-{\overline{\alpha }}_t}{ϵ}_{\theta }(x_t,t)}{\sqrt{{\overline{\alpha }}_t}}\right)}}+\underset{\text{"direction pointing to}{x}_t\text{"}}{\underset{}{\sqrt{(1-{\overline{\alpha }}_{t-1}-{\sigma }_t^2){ϵ}_{\theta }(x_t,t)}}}+\underset{\text{random noise}}{\underset{}{{\sigma }_t{ϵ}_t}}& & \\ & =\sqrt{{\overline{\alpha }}_{t-1}}\left(\frac{x_t-\sqrt{1-{\overline{\alpha }}_t}{ϵ}_{\theta }(x_t,t)}{\sqrt{{\overline{\alpha }}_t}}\right)+\sqrt{(1-{\overline{\alpha }}_{t-1}){ϵ}_{\theta }(x_t,t)}& \text{[}\sigma \text{= 0]}& \end{array}$

If $\sigma$ is 0, the sampling process is deterministic and because $x_0$ is predicted (which is what we finally want + we can replace $t$ with $\tau$) each time, some steps [T..1] can be skipped.

Conclusion

We answered all the questions we asked—loss formulation, (required) image generation, sampling efficiency and background connecting to VAE. There are (many) additional papers that have come out since the initial set, but this post focusses on the fundamentals and therefore ends here.

REFERENCES

1. ^

   Denoising Diffusion Probabilistic Models

2. ^

   The integration of Gaussian distribution over the given range is necessary because $x_0$ is discrete. The alternate is to treat $x_0$ as continuous and compute the value of the Gaussian distribution at $x_0$. Both approaches lead to an MSE though, since the integral is approximated in DDPM.

3. ^

   Diffusion Models Beat GANs on Image Synthesis

4. ^

   DENOISING DIFFUSION IMPLICIT MODELS

   https://​​dzdata.medium.com/​​intro-to-diffusion-model-part-3-5d699e5f0714

   https://​​lilianweng.github.io/​​posts/​​2021-07-11-diffusion-models/​​#nice

   https://​​medium.com/​​better-programming/​​diffusion-models-ddpms-ddims-and-classifier-free-guidance-e07b297b2869

   https://​​mbernste.github.io/​​posts/​​vae/​​