Why don’t diffusion models suffer posterior collapse?

Question

In VAEs, posterior collapse occurs when the approximated posterior $q_\theta(z|x)$ becomes the standard Gaussian prior $p(z)$ after training (Lucas et al. 2019). The forward process of diffusion models transforms the data $x$ into a standard Gaussian at time step $T$, which is similar to approximating the posterior $p(z|x)$ in hierarchical VAEs (see this blog post for more details). Therefore, I am wondering whether diffusion models suffer from the issue of posterior collapse, if that makes sense at all.

Answer 1

Your own cited blog explains the nice property to avoid posterior collapse in DDGMs simply due to the design of diffusion models.

If we recall the discussion about a potential issue of the posterior collapse in hierarchical VAEs, this should not be a problem anymore. Why? Because we should get a standard Gaussian distribution in the last layer by design... The key point is how we define these distributions. Before, we used normal distributions parameterized by DNNs but now we formulate them as the following Gaussian diffusion process (Sohl-Dickstein et al., 2015)... In flows, we know the inverse but we pay the price of calculating the Jacobian determinant while DDGMs require flexible parameterizations of the reverse diffusion but there are no extra strings attached.

Therefore in summary the specific convenience of DDGM is that the forward diffusion (i.e., the variational posteriors) are fixed as exemplified by your reference's equation $q_ϕ(\mathbf{z_i}|\mathbf{z_{i−1}})=N(\mathbf{z_i}|\sqrt{1−β_i}\mathbf{z_{i−1}},β_i\mathbf{I})$ where $\mathbf{z_0}=\mathbf{x}$ and usually in practice Ho et al suggests to change $β_i$ linearly from $β_1=10^{−4}$ to $β_T=0.02$, thus apparently $q_ϕ(\mathbf{z_T}|\mathbf{z_{T−1}})$ by design won't collapse to the standard Gaussian using the symbol $p_\theta(\mathbf{z_T)}$ in your reference.

Answer 2

A answer to help provide clarification on posterior collapse, why it happens in the training of VAEs and how these ideas relate to diffusion models.

As a first step in understanding posterior collapse, let's review the training objective for the $\beta$-VAE algorithm

\begin{equation} \begin{split} \mathcal{L}_{elbo} \: &= \: \underbrace{ \mathbb{E}_{z \sim q_{\phi}(z \vert x), \: x \sim p^{*}(x)} \: \Big\lbrack \: \log p_{\theta}(x \vert z) \: \Big\rbrack }_{\text{log-likelihood}} \: - \: \beta \cdot \underbrace{ \mathbb{E}_{x \sim p^{*}(x)} \: \Big\lbrack \: D_{KL} \big( \: q_{\phi}(z \vert x) \: \big\Vert \: p_{\theta}(z) \: \big) \: \Big\rbrack }_{\text{KL-divergence}} \end{split} \end{equation}

given by a known data distribution $p^{*}(x)$. The objective is a balance of both terms. The log-likelihood term improves the quality of the generated data point. And the KL-divergence serves to restrict the shape of the latent distribution $q_{\phi}(z \vert x)$ by keeping it in proximity to a Gaussian prior distribution, usually set as a standard Gaussian distribution $p(z)= p_{\theta}( z \vert 0,1)$. Properly tuning the $\beta$ parameter insures the latent distribution does not collapse onto the Gaussian prior, leading to what is hypothesized as limiting the number of latent variables exploited for generation of new samples. This phenomenon reduces the capacity of the VAE to produce samples from the entirety of the posterior, as well as reducing the quality of generation. This is what is meant by posterior collapse.

The forward diffusion process, as mentioned in DDPM (one type of manifestation of the algorithm), is by design constructed as a fixed linear Gaussian model for a discrete number of time-steps, each with known distributional parameters ($\mu_{t}, \sigma_{t})$ at each time-step

\begin{equation} \begin{split} q(z_{t} \vert z_{t-1}, x) \: = \: q(z_{t} \vert z_{t-1}) \: = \: \mathcal{N} \big( \sqrt{\smash[b]{\alpha_t}} \cdot z_{t-1} , \: (1 \: - \: \alpha_{t}) \cdot \textbf{I} \: \big) \end{split} \end{equation}

given by $\alpha_{t}= 1-\beta_{t}$. The goal is to train a neural network model to learn to predict these parameters so the reverse sequential (inference) process can proceed. Moreover, we define the diffusion objective to better understand how training occurs

\begin{equation} \begin{split} \mathcal{L}_{simple} \: &= \: \mathbb{E}_{\epsilon_{t} \sim \mathcal{N}(0, \textbf{I}), \: t \sim \mathcal{U}(1,T), \: x \sim q(0)} \Big\lbrack \: \big\Vert \: \epsilon_{t} \: - \: \epsilon_{\theta}(z_{t},t) \: \big\Vert_{2}^{2} \: \Big\rbrack \end{split} \end{equation}

The equation describes the procedure to minimize the difference between a randomly sampled noise and the model noise estimate for the conditional distribution $p_{\theta}(z_{t-1} \vert z_{t})$. The noise value is used to approximate the mean of the aforementioned conditional distribution enabling reverse sampling and generation.

What this discussion has tried to highlight is that VAEs are trained to learn a latent distribution (and therefore collapse of this distribution is possible). However, diffusion models directly learn the data distribution through said denoising process, avoiding the need for an explicit latent variable representation that can collapse.

The following excerpts are taken from my book on variational inference and generative ai. Learn more on the topic by visiting https://www.thevariationalbook.com/