# Justification of the fixed variational distribution in diffusion models

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)$$?

• I have the same question. Have you figured it out? Commented Nov 26, 2023 at 1:22
• @Newton I'm not completely sure I fully grasp this yet, but my current understanding is that the ELBO can be made tight also by updating the parameters of the model (and not only through $q$): if the model is flexible enough, we can find $\theta$ such that $\mathrm{KL}(q(z|x)||p(z|x; \theta)) = 0$, so the bound will be tight. What I still don't get is why the bound is tight in the region of interest, that is, at the maximum of the ELBO (as I've tried to draw in this figure). Commented Nov 28, 2023 at 8:17
• In case it's helpful, here is what Jascha Sohl-Dickstein said (personal communication): "[T]e variational bound can be made tight by either changing p to match q, or by changing q to match p. In a VAE both are done at once. In a (vanilla) diffusion model, only the second is done -- but this is still sufficient to bring the variational bound to 0 if the generative model is flexible enough, or has a structure which is closely enough matched to that of the inference model." And here is Simon Prince's answer. Commented Nov 28, 2023 at 8:21
• Why would $D_\text{KL}[q(z|x) || p(z|x; \theta)]$ become close to zero just because it can? Unlike in a VAE where maximizing ELBO over $\phi$ implies minimizing KL, there is nothing in the diffusion loss function that encourages this. Commented Nov 29, 2023 at 9:50
• I agree with you, @Newton! Please let me know if you find something more about this. This question has been nagging me for a long time. And I thought I had been missing something obvious, given that every textbook takes this formulation for granted. But at least now we are two :-) Commented Nov 29, 2023 at 21:07

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.

• Thanks for the answer, Dario! I guess my question boils to the following: What guarantee do we have that by maximizing the ELBO with a fixed distribution $q$, we also improve the marginal log-likelihood $\log p_\theta(x)$? I'm pretty sure that we can come up with a distribution $q$ that will make the ELBO a poor lower bound and as such it will provide no information on the true marginal log-likelihood. So, what properties does $q$ have to satisfy to ensure a tight lower bound for all parameters $\theta$? Commented Jan 17, 2023 at 10:44
• The distribution $q$ is the true denoising process, which is dependent on data. I think it is the other way around: we are getting $p\theta$ close to $q$. Tell me if this answers your question! Commented Jan 17, 2023 at 12:17
• I sort of understand your point, but if the goal is to get $p_\theta$ close to (a fixed, true) $q$ then why do we care maximizing the log-likelihood? Note that $\log p_\theta(x)$ doesn't depend on $q$, so maximizing $\log p_\theta$ may yield in principle a completely different objective. I understand that we never get to optimize the true log likelihood (due to intractability), but then I feel that the story is somewhat backwards here. What do you think? Commented Jan 17, 2023 at 12:46
• The maximization of log-likelihood (or the minimization of negative log-likelihood in our case) is the way we estimate the parameters of a statistical model, in a way that it gives high probability to the data points of the dataset used for training. Notice that the ELBO does depend on the reverse process $q$, and in fact we are learning $p_\theta$ to be closer to $q$ by using the ELBO as a proxy. The ELBO also depends on $q$, but in this case it becomes tractable because during training we can condition also on the original image $x$. Commented Jan 17, 2023 at 13:53
• I have edited the original answer to include these comments, I hope it is more clear now Commented Jan 17, 2023 at 13:54