This blog posts introduce derivation of diffusion models.

Diffusion Model

Diffusion models are inspired by non-equilibrium thermodynamics. They define a Markov chain of diffusion steps to gradually add random noise to the data and then learn to reverse the diffusion process to construct desired data samples from the noise. Unlike VAE or flow models, diffusion models are learned with a fixed $p(z \mid x)$ and the latent variable has high dimensionality (same as the original data).

Forward Process

Forward process adds noise to the sampled data graduately and finally turns it into noise equivalent to one sampled from an isotropic Gaussian distribution. Formally, given data point $x_0 \sim q(x)$, the forward process produce a sequence of noisy samples $x_1, \cdots, x_T$ in the form of markov chain. $x_t$ is generated by sampling from a Gaussian with mean of $\sqrt{1-\beta_t} x_{t-1}$ and variance $\sigma^2=\beta_t I$. $\beta_t$ is the step size controlling how fast it will turn into a zero-mean Gaussian noise.

The joint distribution of all the noisy samples given $x_0$ is $q ( x_1, \cdots, q_T \mid x_0 ) $ (denoted as $q(x_{1:T} \mid x_0)$):

The direct and simple solution to obtain the $t^{th}$ noisy sample $x_t$ is to iteratively sample according to $q(x_t\mid x_{t-1})$. A nice property of Gaussian is that this process can be written in a closed-form. That’s to say, we can sample $x_t$ at any arbitrary time step $t$ using reparameterization trick. Let $\alpha_t = 1 - \beta_t$, we have

Here, $\epsilon_t \sim N(0,1)$ is sampled from a standard Gaussian. We can further re-write $x_{t-1}$ in the same form:

The second and the third term is the sum of two Gaussian with mean $\mu=0$ and variance $\sigma_t^2 = 1-\alpha_t$ and $\sigma_{t-1}^2=\alpha_t(1 - \alpha_{t-1})$. The sum of two independent Gaussian $N(\mu_1, \sigma^2_1)$, $N(\mu_2, \sigma^2_2)$ is also an Gaussian $N(\mu_1+\mu_2, \sigma^2_1 + \sigma^2_2)$.

Thus, the sum of the second and the third term can be viewed as a new noise sampled from $N(0, 1-\alpha_t \alpha_{t-1})$. Using the reparameterization trick, we write it as $\sqrt{1 - \alpha_t \alpha_{t-1}}\hat \epsilon_{t-2}$. We can further expand the $x_t$ to $x_0$:

Here all $\epsilon_t$, $\hat\epsilon_{t-2}$ and $\epsilon$ are noise sampled form standard Gaussian $N(0,1)$. Thus, we can sample $x_t$ at any arbitrary time step t in a closed form:

where $\alpha_t = 1-\beta_t$ and $\bar\alpha_t=\alpha_t \alpha_{t-1} \cdots \alpha_1=\prod_{i=1}^t \alpha_i$.

Reverse Process

If we can reverse the process of each diffusion step $q(x_{t+1}\mid x_t)$ with $q(x_t\mid x_{t+1})$. Bayes’ Rule can be used to get $q(x_{t+1}\mid x_t)$, but $q(x_t)$ is intractable since $x_t$ is usually in high dimensional space (same dimension as input image).

Diffusion model thus use a $p(x_t\mid x_{t+1})$ to approximate $q(x_t\mid x_{t+1})$. $p(x_T)$ equals to $q(x_T)$ being the final noise distribution, which is a standard Gaussian. The joint disrtibution $p(x_0, x_1, \cdots, x_T) = p(x_{0:T})$ can be formulated by probability of the reverse trajectory:

The goal of optimization is to maximize the (log-) probability of $p(x_0)$

Training amounts to maximizing the expectation of log likelihood $E_{x_0 \sim q(x_0)}[\log p(x_0)]$. From the above, we can see there is intergral under logarithm in $p(x_0)$, which is problematic to optimize. We use Jensen’s inequality to put the logarithm under intergral:

We split the product under logarithm into sum of logarithm and peel off the contribution from $p(x_T)$:

The first term is essentially entropy of $x_T$ since $p(x_T)=q(x_T)\approx N(0,1)$.

Then we split the term of $t=1$ out for edge effect. We show later why this is neccessary.

Here, we show why it is neccessary to split the term with $t=1$. If the summation starts from $t=1$, then we will have to perform Bayes’ Rule on the first term which ended up with something like the following with a loop ($q(x_1\mid x_0)$ exists on both sides).

We take the second term out and split the product under logarithm into sum of logarithm. We find out that the second and third term is essentially conditional entropies and many of them are canceling each other out.

Now, we show how the first term can be converted into a KL divergence:

Now we have all the components to compute the optimization objective where we replace $-H(x_T)$ with intergral form ($-H(x_T)=\int q(x_{0:T})dx_{0:T} \log p(x_T)$):

Then we first merge the last two term canceling the $q(x_1\mid x_0)$ and then merge the denominator $q(x_T\mid x_0)$ with the first term:

Then we re-write the first term into KL Divergence:

Now take a look at each term. The first term

is essentially a small constant since $q(x_T\mid x_0)$ converges to $p(x_T)$ which is a standard Gaussian. The third term is tractable since $p(x_0\mid x_1) = N(\mu_\theta(x_1, 1), \beta_1)$ is in the form of Gaussian distribution:

A easier and more computational simple way is to get rid of the noise at $q(x_1\mid x_0)$. The second term contributes to the main loss of the training.

Both $q(x_{t-1}\mid x_0, x_t)=N(\tilde \mu_t(x_t, x_0), \tilde \beta_t I)$ and $p(x_{t-1}\mid x_t)=N(\mu_\theta(x_t, t), \Sigma_\theta(x_t, t) I)$ are Gaussian. We recall at the beginning that $q(x_{t-1}\mid x_t)$ is intractable, but if it’s conditional on $x_0$, then it’s tractable and we can give a closed form solution via Bayes’ rule:

All $q(x_{t}\mid x_{t-1}, x_0)$, $q(x_{t-1}\mid x_0)$ and $q(x_{t}\mid x_0)$ has analytical forms (desity function of Gaussian distribution). $q(x_{t+1}\mid x_t, x_0) = q(x_{t+1}\mid x_t)$ since $q$ is a markov process.

Then, we have $q(x_{t-1}\mid x_{t}, x_0)$. Essentially, we want to get $\tilde \mu_t(x_t, x_0)$ and $\tilde \beta_t$, which is the mean and variance of Gaussian distribution $q(x_{t-1} \vert x_t, x_0)$. $q(x_{t-1} \vert x_t, x_0)$ shares the similar density function as any Gaussian:

By taking the $q(x_{t}\mid x_{t-1}, x_0)$, $q(x_{t-1}\mid x_0)$ and $q(x_{t}\mid x_0)$ to the Bayes’ rule, we re-write it as following:

We can see,

We can replace $x_0$ with $x_t = \sqrt{\bar \alpha_t} x_0 + \sqrt{1 - \bar \alpha_t}\epsilon$,

Now, we derive the final loss of the second term and since we use a fixed $\beta_t$ meaning that $\Sigma_\theta(x_t, t)$ is the same as $\beta_t$.

Since we have $x_t$ during training, we re-write $\mu_\theta(x_t, t)$ in the same form as $\tilde{\mu}_t (x_t, x_0)$ and the objective now becomes predicting the noise $\epsilon$:

DDPM gets rid of the constant as it is empirically better.