Beyond Vanilla Variational Autoencoders: Detecting Posterior Collapse in Conditional and Hierarchical Variational Autoencoders

Dear Reviewer 1CHE,

Thank you for your thoughtful review and valuable feedback. Below we address your concerns. We have also revised the current manuscript according to the suggestions of the reviewers. The changes are in blue color. Please kindly see the revised manuscript, as well as our General Response and Summary of Revision.

Q1: The reviewer appreciates little contribution in the paper. Yes, the authors generalize the results obtained for the linear VAE by Wang & Ziyin (2022) and Dai et al. 2020 to linear CVAEs and hierarchical linear VAES, and this contribution is acknowledged. However, I wonder if generalizing the analysis to these new models resulted in different conclusions or insights about avoiding posterior collapse. In other words, is anything in Table 1 significantly different from what we learned from Wang and Ziyin? At the end, the conclusion is that both $\beta$ and $\eta_{dec}$ should be small.

A1: We respectfully disagree with the reviewer that our contributions are little. Please allow us to clarify our contributions below.

1/ Implications of our theoretical results for CVAE and MHVAE.

For CVAE, our theoretical analysis is the first to prove the existence of posterior collapse in CVAE, an important and popularly used variant of VAE [1, 2, 3]. We rigorously study the effect of parameters to the rank of the models and the level of posterior collapse at optimality, including $\theta$’s (the singular values of matrix $Z = \mathbb{E}(x \tilde{x}^{\top})$, defined at Theorem 2), $\beta$ (scaling scalar of the KL term of ELBO at Equation (5) in our paper) and $\eta_{dec}$ (magnitude of the variance of the generating distribution). The results imply that a sufficiently small $\beta$ and/or $\eta_{dec}$ can mitigate posterior collapse. Especially, we prove that the correlation of the input condition and output of training data relates to the collapse level (see the paragraph after Theorem 2 in our manuscript). This insight obviously cannot be derived from the standard VAE setting which does not incorporate input conditions into the generative process.

For MHVAE, our theoretical analysis is again the first to prove the existence of posterior collapse in MHVAE, another important and popularly used variant of VAE [4, 5, 6, 7]. We also rigorously study the effect of parameters to the rank of the models and the level of posterior collapse at optimality, including $\theta$’s, $\beta_1$, $\beta_2$ and $\eta_{dec}$. Moreover, since the ELBO loss of MHVAE consists of multiple KL-regularized terms, we study the effect of each KL-regularized term and the trade-offs between latents via the magnitude of the hyperparameters $\beta_1$ and $\beta_2$. We find that decreasing $\beta_2$ can help to alleviate posterior collapse and increase the rank of the model, while decreasing $\beta_1$ will have the opposite effect. This insight cannot be derived from the standard VAE setting which has only one latent and one KL-regularized term in the ELBO. In addition, we observe that when any latent variable of MHVAE suffers complete collapse, its higher-level latents become useless. Thus, we suggest creating additional mapping between the input data and every latent to prevent this adverse situation.

Furthermore, we find that unlearnable encoder variance can prevent posterior collapse (see Section 3 in our paper) for VAE, CVAE and MHVAE. This insight is also novel.

2/ As per Reviewer vZoP comment, our paper “provides an important step forward for the understanding of the inner workings of VAEs and provides the basis and theoretical tool set for future research”. Specifically, current state-of-the-art VAE models are mostly CVAE and HVAE [1, 2, 3, 4, 5, 6, 7], not standard VAE with only one latent, so our settings are more practical compared with prior works and provide the basis and theoretical tool set for future research of CVAE and HVAE. Furthermore, as the discussion in Section 4.2 in our manuscript, diffusion models share many similarities with deep MHVAE model where the encoding process of diffusion models also consists of consecutive linear maps with injected noise and with unlearnable isotropic variance, and their training loss function is also ELBO. Therefore, our framework potentially sheds light on the understanding of diffusion models.

Dear Reviewer fUP2,

Thank you for your thoughtful review and valuable feedback. Below we address your concerns. We have also revised the current manuscript according to the suggestions of the reviewers. The changes are in blue color. Please kindly see the revised manuscript, as well as our General Response and Summary of Revision.

Q1: The primary weakness was in exposition. The paper tends to discuss intuitions after theorems, which makes the theorems a bit challenging on first read.

Q2: Ok the explanation after the theorem is helpful. It would be preferable to help readers see what is coming ahead of time, too.

A1 + A2: Thanks for your suggestions. We have added the description before every theorem of what the incoming theorems are about and what they try to accomplish to facilitate the reading.

Q3: Moreover, at several points the authors assume a lot of the reader. Specifically, consider the case of unlearnable Sigma, and comparisons between theorem 1 and the previous work (the result from which is not stated in the text). Taken together the paper was more difficult than need be to understand.

Q4: "We note that our results generalize Theorem 1 in (Wang & Ziyin, 2022) where σi’s are all equal to a constant." It would really be helpful to the reader to present the previous result and explain clearly why.

A3 + A4: Thank you for your valuable comment. We have revised the paper to clarify the previous results whenever we mention them. For the comparison between our Theorem 1 and previous results, we have updated the paper to describe explicitly the differences between our Theorem 1 and Proposition 2 in (Wang & Ziyin, 2022), please see the paragraphs before and after Theorem 1 in our manuscript.

Specifically, our Theorem 1 allows for arbitrary predefined values of $\\{ \sigma_i \\}\_{i=1}^{d_1}$, thus is more general than the Proposition 2 in (Wang & Ziyin, 2022) where $\sigma_i$'s are all equal to a constant. Under broader settings, there are two notable points from Theorem 1 that has not been captured in the previous result of (Wang & Ziyin, 2022): i) at optimality, the singular matrix $\mathbf{T}$ of the encoder map $\mathbf{V}$ sorts the set $\\{ \sigma_i \\}\_{i=1}^{d_1}$ in non-decreasing order, and ii) singular values $\omega\_{i}$ of the decoder map $\mathbf{U}$ and $\lambda\_{i}$ of the encoder map $\mathbf{V}$ are calculated via the $i$-th smallest value $\sigma^{\prime}\_{i}$ of the set $\\{ \sigma_i \\}\_{i=1}^{d_1}$, not necessarily the $i$-th element $\sigma_i$.

Q5: "Interestingly, we find that the correlation of the training input and training output is one of the factors that decides the collapse level" Why is this interesting? And some typos: "We study model with", "beside the eigenvalue"

A5: Thanks for your suggestion and for pointing out the typos. We used the word “Interestingly” because we found that the finding was novel. In our revision, we removed this word. We have also fixed the typos that the reviewer mentioned.

Q6: What is the meaning of Theorem 1? It isn't yet clear to me what we are proving or why?

Q7: “Thus, learnable latent variance is among the causes of posterior collapse, opposite to the results in Wang & Ziyin (2022) that it is not the cause of posterior collapse." This warrants some explanation. Why is the conclusion opposite?

A6 + A7: Our Theorem 1 characterizes the global minima of the ELBO training problem, which includes the parameters from the matrix $\mathbf{U}$ (the decoder), $\mathbf{W}$ (the encoder) as optimization variables. The encoder variance $\mathbf{\Sigma}$ is fixed/unlearnable in this setting. We prove that the optimal parameters $(\mathbf{U}^{\ast}, \mathbf{W}^{\ast})$ have their singular values can be calculated via the closed-form formula stated in Theorem 1 in our manuscript. Thus, the ranks of the encoder map and the decoder map depend on $\theta$’s (the singular values of the matrix $\mathbf{Z} := \mathbb{E}(x \tilde{x}^{\top})$ ), $\beta$ (scaling factor of the KL term in ELBO in Eqn. (3)) and $\eta_{dec}$ (the variance magnitude of the generating distribution $p(x | z)$). After Theorem 1, in Section 3, we discuss the conditions for posterior collapse to occur, including the rank of the encoder/decoder and learnability of the encoder variance.

Dear Reviewer VR5S,

Thank you for your thoughtful review and valuable feedback. Below we address your concerns. We have also revised the current manuscript according to the suggestions of the reviewers. The changes are in blue color. Please kindly see the revised manuscript, as well as our General Response and Summary of Revision.

Q1: Much of the empirical results are relegated to the appendix with the paper itself being somewhat light on experiments. This isn’t necessarily a bad thing but it would be nice to see more thorough experiments in the main paper, e.g. vanilla VAE, linear CVAE/MHVAE, learnable vs. unlearnable variance, just for comparison.

A1: Thank you for your comment. We have moved an additional experiment to the main paper, where we compare the level of posterior collapse for VAE with learnable and unlearnable encoder variance to validate our claim about the role of learnability encoder variance to posterior collapse (Section 3). Please kindly see Figure 2 in the revised manuscript. Due to the main text’s space limit of 9 pages and our work consists of many theoretical results that require detailed explanation, we currently include other experimental results in the Appendix.

Q2: The theoretical results are generally restricted to highly specific scenarios, such as linear and two latent for MHVAE, and the setting of both variances being learnable in MHVAE is not considered. It would be nice to see more theoretical motivation for why these results might generalize to more realistic training settings, but instead that is largely directed towards prior work.

A2: We agree with the reviewer that our theoretical results focus on the settings of linear networks. However, the theoretical analysis of nonlinear networks is very challenging and, in fact, there has been no rigorous theory for deep nonlinear networks yet to the best of our knowledge. Even for the linear case, a complete theoretical understanding of the VAE setting is still lacking. We aim to fill this gap and use the insights from our work for deep linear VAEs to empirically examine and improve deep nonlinear VAEs: we have experiments at Figure 2 + Figure 3 in our main paper and more experiments on nonlinear networks in the Appendix that support our theoretical insights.

We also want to emphasize that analyzing deep linear networks has been an important step in studying deep nonlinear networks. [1, 2, 3, 4] show that the optimization of deep linear models exhibits similar properties to those of the optimization of deep nonlinear models. In practice, deep linear networks can help improve the training and performance of deep nonlinear networks [5, 6, 7]. Specifically, [5] empirically proves that linear overparameterization in nonlinear networks improves generalization on classification tasks (see Section 4 in [5]). In particular, [5] expands each linear layer into a succession of multiple linear layers and does not include any non-linearities in between, which results in a considerable increase in performance. [6] applies a similar strategy for compact networks, and their experiments show that training such expanded networks yields better results than training the original compact networks. We discuss these importances of studying linear networks in the Related Work section in Appendix B in our paper, which is also our theoretical motivation to study this setting.

Regarding the motivation for the settings of CVAE and MHVAE, we note that current state-of-the-art VAE models are mostly deep HVAE and CVAE [8, 9, 10, 11, 12], not standard VAE with only one latent, so our settings are practical and we provide the basis and theoretical tool set for future research of CVAE and HVAE (as per Reviewer vZoP comment). Furthermore, as discussed in Section 4.2 in our paper, diffusion models share many similarities with MHVAE model where the encoding process of diffusion models also consists of consecutive linear maps with injected noise and with unlearnable isotropic variances, and their training loss function is also the ELBO. Therefore, our framework potentially sheds light on the understanding of diffusion models.

We note that MHVAE with two latents is cumbersome to study theoretically. The complexities added when increasing the number of latent variables from one (standard VAE) to two can be seen via the ELBO training problem of two models in our paper: $L_{VAE}$ in Eqn. (3) and $L_{HVAE}$ in page 6. Due to some technical challenges, we currently leave the setting of both learnable encoders as future work.

Dear Reviewer Bn1w,

Thank you for your thoughtful review and valuable feedback. Below we address your concerns. We have also revised the current manuscript according to the suggestions of the reviewers. The changes are in blue color. Please kindly see the revised manuscript, as well as our General Response and Summary of Revision.

Q1: The paper seems to claim a result that has already been proven. Theorem 1 seems to be exactly identical to proposition 2 in (Wang & Ziying, 2022). Maybe the functional form of the result is identical but the authors have more general assumptions on $\sigma_i$. In either case, the statement that this theorem is one of the novel contributions is false or at best misleading.

Q2: Can you clarify on the novelty claim for theorem 1?

A1 + A2: Thank you for your comment. We agree with the reviewer that the word “novel” may be inappropriate here. In our revision, we have changed the wording from “Our novel contributions” to “Our contributions” in Section 3. Please allow us to clarify the contributions of Theorem 1 in our manuscript and the differences with Proposition 2 in (Wang & Ziyin, 2022) as follows.

First, compared to Proposition 2 in (Wang & Ziyin, 2022), our Theorem 1 characterizes the global solutions of standard linear VAE for unlearnable $\mathbf{\Sigma}$ with arbitrary elements $\sigma_i$ on the diagonal while Proposition 2 and Theorem 1 in (Wang & Ziyin, 2022) studies the isotropic $\mathbf{\Sigma} = \sigma^2 \mathbf{I}$ case only. Another new and interesting result from Theorem 1 in our manuscript is that we prove the left singular matrix $\mathbf{T}$ of the encoder map $\mathbf{V}$ is the sorting matrix that sorts the values of $\mathbf{\Sigma}$ in non-decreasing order. This property of $\mathbf{T}$ is one of the factors that decide whether posterior collapse happens or not, which has not been considered in (Wang & Ziyin, 2022).

Second, after Theorem 1 in (Wang & Ziyin, 2022), they claim that “optimizable encoder variance is not essential to posterior collapse problem” (Section 4) and “a learnable (data-dependent or not) latent variance is not the cause of posterior collapse” (Section 4.5). Moreover, in Section 4.2 of (Wang & Ziyin, 2022), after Theorem 1, under the setting of unlearnable isotropic variance $\sigma_i = \eta_{enc}, \forall i$, they claim the existence of posterior collapse when the learned model is low-rank (i.e., some singular values $\lambda$ and $\theta$ are 0’s). In particular, they argue “...the learned model becomes low-rank. Namely, some of the dimensions collapse with the prior”. We observe that these claims are not totally correct, and we explain in our manuscript why posterior collapse may not happen in the above unlearnable setting (see the paragraph after our Theorem 1). For convenience, we repeat the explanation below.

We first let the SVD of the encoder matrix $\mathbf{V}$ (i.e, the encoder map) to be $\mathbf{V} = \mathbf{T} \Lambda \mathbf{S}^{\top}$ then at the global minimizer of the training problem, we prove in Theorem 1 that $\mathbf{T}$ is the matrix that sorts the diagonal values of $\Sigma$ in non-decreasing order. When $\mathbf{V}$ is low-rank, $\Lambda$ will have zero rows, then $\Lambda \mathbf{S}^{\top}$ will have zero rows. However, in the case of isotropic $\mathbf{\Sigma} = \eta_{enc}^{2} \mathbf{I}$ (as considered in Theorem 1 in (Wang & Ziyin, 2022)), $\mathbf{T}$ can be any orthonormal matrix (since $\sigma_i$’s are all equal and in non-decreasing order already), and thus, $\mathbf{V} = \mathbf{T} \Lambda \mathbf{S}^{\top}$ may have no zero rows. Then, for any data $x$, the mean $\mathbf{V} \tilde{x}$ of the approximate posterior $q(z|x) = \mathcal{N} (\mathbf{V} \tilde{x}, \Sigma)$ might have no zero component, and thus the posterior may not collapse to the prior $\mathcal{N}(0, \eta^2_{enc})$. On the other hand, if all values $\sigma_i$’s are distinct, Theorem 1 suggests that $\mathbf{T}$ is a permutation matrix and hence, $\mathbf{V}$ will have zero rows, which corresponds with the dimensions of latent $z$ that collapse.

For the case of learnable $\mathbf{\Sigma}$, low-rank condition leads to posterior collapse. This is because at optimality, the diagonality of $\mathbf{\Sigma}$ makes $\mathbf{U}^{\top} \mathbf{U}$ diagonal. Therefore, $\mathbf{U}$ has zero columns, and subsequently, $\mathbf{V}$ will have zero rows (see detailed explanation in the paragraph before Section 4 in our manuscript). Finally, the corresponding dimensions of $z = \mathbf{V} \tilde{x} + \epsilon$ collapse to its prior $\mathcal{N}(0, \eta^2_{\text{enc}})$. Thus, our findings suggest that the learnability of encoder variance plays an important role in the existence of posterior collapse. We have updated the manuscript to describe carefully the differences between our Theorem 1 and Proposition 2 in (Wang & Ziyin, 2022) in the paragraphs after Theorem 1 in our revision.

Dear Reviewer vZoP,

Thank you for your thoughtful review and valuable feedback. Below we address your concerns. We have also revised our manuscript according to the suggestions of the reviewers. The changes are in blue color. Please kindly see the revised manuscript, as well as our General Response and Summary of Revision.

Q1: The conditions provided in the paper are only for the linear (single layer encoder and decoder) case, which of course limits the practical applicability of the framework. Yet, I believe that the work provides an important step forward in the understanding of VAEs and do not see a big drawback in this.

A1: Thank you for your comment and appreciation of our work. We agree with the reviewer that our theoretical results focus on the settings of linear networks. However, the theoretical analysis of nonlinear networks is very challenging and, in fact, there has been no rigorous theory for deep nonlinear networks yet to the best of our knowledge. Even for the linear case, the complete theoretical understanding in the VAE setting is still lacking. We aim to fill this gap and use the insights from our work for deep linear VAEs to empirically examine and improve the deep nonlinear VAEs: we have experiments at Figure 2 + Figure 3 in our main paper and more experiments on nonlinear networks in the Appendix that support our theoretical insights. Throughout this process, our frameworks and theorems play an important role as a starting point for further theoretical and empirical study of the nonlinear conditional VAEs, hierarchical VAEs or even diffusion models, potentially.

We also want to note that analyzing deep linear networks has been an important step in studying deep nonlinear networks. For example, using only linear regression, [1] can recover several phenomena observed in large-scale deep nonlinear networks, including the double descent phenomenon [2]. [3, 4, 5, 6] show that the optimization of deep linear models exhibits similar properties to those of the optimization of deep nonlinear models. In practice, deep linear networks can help improve the training and performance of deep nonlinear networks [7, 8, 9]. Specifically, [7] empirically proves that linear overparameterization in nonlinear networks improves generalization on classification tasks (see Section 4 in [7]). In particular, [7] expands each linear layer into a succession of multiple linear layers and does not include any non-linearities in between, which results in a considerable increase in performance. [8] applies a similar strategy for compact networks, and their experiments show that training such expanded networks yields better results than training the original compact networks. [9] shows that linear overparameterization, i.e., the use of a deep linear network in place of a classic linear model, induces on gradient descent a particular preconditioning scheme that can accelerate optimization. The preconditioning scheme that deep linear layers introduce can be interpreted as using momentum and adaptive learning rate.

Q2: In Fig. 3 (b) the case for beta_1=1.0 and beta_2=1.0 has been left out. Just for completeness it might be interesting to visualize the result for this "standard" ELBO (beta_1=beta_2=1.0).

A2: Thanks for your suggestion. We already have this setting in the paper. In particular, we conducted the experiment for either $\beta_1$ or $\beta_2$ is 1 and included the results in Figure 4 on page 14 of the Appendix. We agree with the reviewer that this setting is standard, and we have moved it back to the main text in our revision.

References:

[1] Trevor Hastie et al. “Surprises in High-Dimensional Ridgeless Least Squares Interpolation”, Annals of statistics, 2022

[2] Preetum Nakkiran et al. “Deep Double Descent: Where Bigger Models and More Data Hurt”, Journal of Statistical Mechanics, 2021

[3] Andrew M. Saxe, James L. McClelland, Surya Ganguli. “Exact solutions to the nonlinear dynamics of learning in deep linear neural networks”, 2013

[4] Kenji Kawaguchi. “Deep Learning without Poor Local Minima”, NeurIPS 2016

[5] Moritz Hardt, Tengyu Ma. “Identity Matters in Deep Learning”, ICLR 2017

[6] Thomas Laurent, James von Brecht. “Deep linear neural networks with arbitrary loss: All local minima are global”, ICML 2018

[7] Minyoung Huh et al. “The Low-Rank Simplicity Bias in Deep Networks”, TMLR 2022

[8] Shuxuan Guo, Jose M. Alvarez, Mathieu Salzmann. “ExpandNets: Linear Over-parameterization to Train Compact Convolutional Networks”, NeurIPS 2020

[9] Sanjeev Arora, Nadav Cohen, Elad Hazan. “On the Optimization of Deep Networks: Implicit Acceleration by Overparameterization”, ICML 2018

A1 (Continue):

2/ Theoretical contributions and motivations

We acknowledge that our theoretical results study the simplistic settings of linear networks. However, the theoretical analysis of nonlinear networks is very challenging and, in fact, there has been no rigorous theory for deep nonlinear networks yet to the best of our knowledge. Even for the linear case, a complete theoretical understanding of the VAE setting is still lacking. We aim to fill this gap and use the insights from our work for deep linear VAEs to empirically examine and improve deep nonlinear VAEs.

We also want to emphasize that analyzing deep linear networks has been an important step in studying deep nonlinear networks. [8, 9, 10, 11] show that the optimization of deep linear models exhibits similar properties to those of the optimization of deep nonlinear models. In practice, deep linear networks can help improve the training and performance of deep nonlinear networks [12, 13, 14]. Specifically, [12] empirically proves that linear overparameterization in nonlinear networks improves generalization on classification tasks (see Section 4 in [12]). In particular, [12] expands each linear layer into a succession of multiple linear layers and does not include any non-linearities in between, which results in a considerable increase in performance.

3/ Our proofs are methodologically novel and are not just a generalization from linear VAE

For example, for MHVAE, both the architecture and the training problem are more complicated than the standard VAEs with only one latent variable. Adding one more level of latent requires two more mappings to map the first latent to the second latent in the encoding process and vice versa for the decoding process. The training problem of maximizing the ELBO must have an extra KL-regularizer term between the two latents, i.e., $D_{\text{KL}}(q(z_1 | x) || p(z_1 | z_{2}))$. This term significantly complicates the training problem and our proof uses new techniques to handle it. Specifically, we prove a property that $\mathbf{V}\_{1} \mathbf{V}\_{1}^{\top}$ and $\mathbf{U}\_{2} \mathbf{U}\_{2}^{\top}$ are simultaneously diagonalizable (see page 38), and thus, we are able to convert the zero gradient condition into relations of their singular values $\lambda$’s and $\omega$’s (see Eqn. (86) at page 38). Thanks to these relations between $\lambda$’s and $\omega$’s , the loss function now can be converted to a function of singular values. Again, the formulations for MHVAE are significantly more cumbersome and require different calculations compared to linear VAE cases, which can be seen via the gradient calculation from Eqn. (75)-(78) and how we leverage these critical points conditions to simplify the loss function in Eqn. (79)-(88).

Q2: The opposite conclusion with previous work (Wang & Ziyin, 2022) regarding the role of learnable latent variance.

A2: After the Theorem 1 in (Wang & Ziying, 2022), they claim that “optimizable encoder variance is not essential to posterior collapse problem” (Section 4), “a learnable (data-dependent or not) latent variance is not the cause of posterior collapse” (Section 4.5). Moreover, in Section 4.2 in (Wang & Ziying, 2022), after Theorem 1, under the setting of unlearnable isotropic variance $\sigma_i = \eta_{enc}, \forall i$, they claimed the existence of posterior collapse when the learned model is low-rank (i.e., some singular values $\lambda$ and $\theta$ are 0’s). In particular, they argue “... the learned model becomes low-rank. Namely, some of the dimensions collapse with the prior”. We observe that these claims are not totally correct and we explain in our manuscript why posterior collapse may not happen in the above unlearnable setting (see the first paragraph after our Theorem 1 in the paper). For convenience, we repeat the explanation below.