Justification of the fixed variational distribution in diffusion models

Question

Diffusion models can be regarded as latent variable models (Ho et al., 2020; Section 2), with the latents being an hierarchical chain of random variables $z_T → \dots → z_t → z_{t-1} → \dots → z_1$ (finally, from $z_1$ we sample the observation $x$). Under this view the diffusion process (the noise-adding steps, $\dots ← z_t ← z_{t-1} ← \cdots$) defines an approximate posterior distribution $q(z_{1:T}|x)$. However, this distribution is fixed (in contrast to typical latent variable models, such as variational autoencoders, which attempt to learn the approximate posterior). So I was wondering why does this approach work?

In my mind, a learnable and flexible variational distribution is essential in ensuring that the evidence lower bound (ELBO) is a tight lower-bound of the marginal log-likelihood $\log p(x)$ (for example in the EM algorithm, the E step closes the gap by setting $q(z|x)$ to $p(z|x)$). Otherwise, if the ELBO is loose do we have any guarantee that we maximize the true marginal log-likelihood $\log p(x)$?

Answer 1

Under this view the diffusion process (the noise-adding steps, $z_T\dots ←z_t←z_{t−1}←\dots z_1$) defines an approximate posterior distribution $q(z_{1:T}|x)$

This is the forward trajectory, but the generative process we are interested in is in the reverse trajectory $q(z_{t-1}|z_t)$ (the denoising trajectory), since we want to provide the model a random noise and the model will generate an image.

Since the true denoising distribution $q(z_{t-1}|z_t)$ is intractable, we want to learn the parameterization for $p_\theta(z_{t-1}|z_t)$, assuming it Gaussian for a small $\beta$. $$p(z_{t-1}|z_t) = \mathcal{N}(z_{t-1};\mu_\theta(z_t),\Sigma_\theta(z_t))$$

It is shown empirically in Ho et al., 2020 that setting $\Sigma_\theta(z_t) = \sigma^2\mathbb{I} = \beta_t \mathbb{I}$ works well, so we can train a neural network to predict just the mean.

To do so we want to minimize the log-likelihood $$\mathbb{E}[-\log p_\theta(x_0)]$$

But since we cannot access $p_\theta(x_0)$ as we do for VAEs, we can obtain a simpler objective using the Jensen Inequality, finding an ELBO that is dependent on $q(z_{t-1}|z_t)$.

$$ \begin{split} \mathbb{E}[-\log p_\theta(x_0)] \leq {} \mathbb{E}_q & \bigl[ D_{KL}(q(z_T|z_0) || p(z_T)) \\ & + \sum_{t\geq 1} D_{KL}(q(z_{t-1}|z_t, x) || p_\theta(z_{t-1}|z_t)) \\ & - \log p_\theta(x | z_1) \bigr] \end{split} $$

Where KL is the Kullback-Leibler Divergence between the two distributions.

Notice that the reverse process becomes tractable when also conditioned on the real image $x$, which does not allow sampling starting from noise (our final objective)

Since we are minimizing this ELBO, by optimizing the parameters $\theta$ we are actually closing the gap between the two distributions $q(z_{t-1}|z_t)$ and our approximation with a NN $p_\theta(z_{t-1}|z_t)$.

• The first term does not depend on $\theta$ and can be ignored for optimization (We could learn $\beta_t$ though)
• The central term is the most important, where we actually close the gap
• The last term could improve the last diffusion step

All of this and the following computations are explained well in this blog, and further mathematical explanation is also available here.