Variational diffusion

• “Understanding Diffusion Models: A Unified Perspective” + “Variational Diffusion” from Kingma

• The easiest way to think of a Variational Diffusion Model (VDM) is simply as a Markovian Hierarchical Variational Autoencoder with three key restrictions:

  1. The latent dimension is exactly equal to the data dimension
     1. This means that $q(z_{t+1}|z_t)=q(x_{t+1}|x_t)$. Everything stays in $x$ space.
  2. The structure of the latent encoder at each timestep is not learned
     • It is pre-defined as a linear Gaussian model. In other words, it is a Gaussian distribution centered around the output of the previous timestep
  3. The Gaussian parameters of the latent encoders vary over time in such a way that the distribution of the latent at final timestep $T$ is a standard Gaussian

• Diagram

  • 

ELBO derivation

The variational lower bound loss is derived from ${\mathbb{E}}_{q(x_{1:T}|x_0)}[log\frac{p(x_{0:T})}{q(x_{1:T}|x_0)}]$ where we take advantage of the fact that both $p(x_{0:T})$ and $q(x_{1:T}|x_0)$ are Markovian, and thus the log-product can be easily decomposed into incremental steps.

• We obtain $logp(x)\ge {\mathbb{E}}_{q(x_1|x_0)}[log{p}_{\theta }(x_0|x_1)]-{\sum }_{t=2}^T{\mathbb{E}}_{q(x_t|x_0)}[D_{KL}(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))]-D_{KL}(q(x_T|x_0)||p(x_T))$
• Reconstruction term
  • where the first term is the reconstruction term, like its analogue in the ELBO of vanilla VAE.
    • This term can be approximated and optimized using a Monte Carlo estimate
• Denoising matching terms
  • the $T-2$ terms are denoising matching terms.
    • We learn desired denoising transition step $p_{\theta }(x_{t-1}|x_t)$ as an approximation to tractable, ground truth denoising transition step $q(x_{t-1}|x_t,x_0)$ (which closed form is derived in The diffusion process)
    • Detail: when originally deriving the ELBO, you get “consistency terms”
      • i.e. $E_{q(x_{t-1},x_{t+1}|x_0)}[D_{KL}(q(x_t|x_{t-1})||p_{\theta }(x_t|x_{t+1}))]$, where a denoising step from a noisier image should match the corresponding noising step from a cleaner image.
      • However, actually optimizing the ELBO using the terms we just derived might be suboptimal; because the consistency term is computed as an expectation over two random variables $\{x_{t-1},x_{t+1}\}$ for every timestep, the variance of its Monte Carlo estimate could potentially be higher than a term that is estimated using only one random variable per timestep.
      • The trick to obtain denoising matching terms is to rewrite the encoder transition as $q(x_t|x_{t-1})=q(x_t|x_{t-1},x_0)$, where the extra conditioning is superfluous due to the Markov property and then use Bayes rule to rewrite each transition as $q(x_t|x_{t-1},x_0)=\frac{q(x_{t-1}|x_t,x_0)q(x_t|x_0)}{q(x_{t-1}|x_0)}$
• Prior matching term
  • The last form enforces isotropic Gaussian at the end of the diffusion, it’s never optimized.

Optimizing in practice

Maximizing ELBO, equivalent to denoising (when variance schedule is fixed)

• Remember that using Bayes theorem, one can calculate the posterior $q(x_{t-1}|x_t,x_0)$ in terms of ${\tilde{\beta }}_t$ and $\tilde{\mu }(x_t,x_0)$ which are defined as follows:

  • ${\tilde{\beta }}_t=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t$ (posterior variance schedule)
  • $\tilde{\mu }(x_t,x_0)=\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{x}_0+\frac{\sqrt{{\bar{\alpha }}_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{x}_t$
  • $q(x_{t-1}|x_t,x_0)=\mathcal{N}(x_{t-1};\tilde{\mu }(x_t,x_0),{\tilde{\beta }}_{t}I)$

• The KL-divergence between two Gaussian disitributions is composed of an MSE of the two means (divided by some variance term) + some terms about the variances

• In the case, where we fix the variance schedule ${\beta }_t$,

  • we can match exactly the two distributions variances, and thus minimizing the KL-divergence is exactly equivalent to minimize the MSE between the two means ⇒ denoising objective: $\underset{\theta }{argmin}{D}_{KL}(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))=\underset{\theta }{argmin}\frac{1}{2{\tilde{\beta }}_t}||{\mu }_{\theta }(x_t,t)-\tilde{\mu }(x_t,x_0)|{|}^2$
  • By setting ${\mu }_{\theta }(x_t,t)=\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{\hat{x}}_{\theta }(x_t,t)+\frac{\sqrt{{\bar{\alpha }}_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{x}_t$ (similar parametrization as $\tilde{\mu }(x_t,x_0)$), we recover: $\underset{\theta }{argmin}\frac{1}{2{\tilde{\beta }}_t}\frac{{\bar{\alpha }}_{t-1}(1-{\alpha }_t{)}^2}{(1-{\bar{\alpha }}_t{)}^2}||{\hat{x}}_{\theta }(x_t,t)-x_0|{|}^2$
  • Optimizing a VDM boils down to learning a neural network to predict the original ground truth image from an arbitrarily noisified version of it
    • Just do $\underset{\theta }{argmin}{\mathbb{E}}_{t\sim U\{2,T\}}[\text{denoising term at time }t]$

Learning Diffusion Noise Parameters (variance schedule is not fixed)

• Above, we wrote our per-timestep objective as $\frac{1}{2{\hat{\beta }}_t}\frac{{\bar{\alpha }}_{t-1}(1-{\alpha }_t{)}^2}{(1-{\bar{\alpha }}_t{)}^2}||{\hat{x}}_{\theta }(x_t,t)-x_0|{|}^2$
• By plugging in ${\tilde{\beta }}_t=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t=\frac{1-{\bar{\alpha }}_{t-1}(1-{\alpha }_t)}{1-{\bar{\alpha }}_t}$ in the above equation, we get $\frac{1}{2}(\frac{{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_{t-1}}-\frac{{\bar{\alpha }}_t}{1-{\bar{\alpha }}_t})||{\hat{x}}_{\theta }(x_t,t)-x_0|{|}^2=\frac{1}{2}(\mathrm{SNR}(t-1)-\mathrm{SNR}(t))||{\hat{x}}_{\theta }(x_t,t)-x_0|{|}^2$
  • Recall that $q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{{\bar{\alpha }}_t}{x}_0,(1-\overline{{\alpha }_t})I)$.
  • Then following the definition of the signal-to-noise ratio (SNR) as $SNR=\frac{{\mu }^2}{{\sigma }^2}$, we can write $\mathrm{SNR}(t)=\frac{{\bar{\alpha }}_t}{1-{\bar{\alpha }}_t}$
• . In a diffusion model, we require the SNR to monotonically decrease as timestep $t$ increases.

Parametrizing the SNR

• We can directly parameterize the SNR at each timestep using a neural network, and learn it jointly along with the diffusion model.
• As the SNR must monotonically decrease over time, we can represent it as:
  • ${\mathrm{SNR}}_{\nu }(t)=\exp (-w_{\nu }(t))$
  • where $w_{\nu }(t)$ is modeled as a monotonically increasing neural network with parameters $\nu$.
  • We can write elegant forms for the $\alpha$‘s
    • $\frac{{\bar{\alpha }}_t}{1-{\bar{\alpha }}_t}=\exp (-w_{\nu }(t))$
    • ${\bar{\alpha }}_t=\mathrm{sigmoid}(-w_{\nu }(t)$
    • $1-{\bar{\alpha }}_t=\mathrm{sigmoid}(w_{\nu }(t)$

Optimizing everything

• We now optimize $\underset{\theta ,\nu }{argmin}{\mathbb{E}}_{t\sim U\{2,T\}}[{\mathbb{E}}_{q(x_t|x_0)}[D_{KL}(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))]]$
• where we showed above that $D_{KL}(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))=\frac{1}{2}({\mathrm{SNR}}_{\nu }(t-1)-{\mathrm{SNR}}_{\nu }(t))||{\hat{x}}_{\theta }(x_t,t)-x_0|{|}^2$