Posterior Collapse of a Linear Latent Variable Model
Abstract
This work identifies the existence and cause of a type of posterior collapse that frequently occurs in the Bayesian deep learning practice.
For a general linear latent variable model that includes linear variational autoencoders as a special case, we precisely identify the nature of posterior collapse to be the competition between the likelihood and the regularization of the mean due to the prior.
Our result suggests that posterior collapse may be related to neural collapse and dimensional collapse and could be a subclass of a general problem of learning for deeper architectures.
Introduction
Bayesian approaches to deep learning have attracted much attention because they allow for a more principled treatment of inference and uncertainty estimation (Mackay, 1992; Neal, 2012; Wang and
Yeung, 2020; Jiang and Ahn, 2020; Zhao et al., 2021; Liu, 2021).
One long-standing and unresolved problem for the Bayesian deep learning practice is the problem of posterior collapse, where the posterior distribution of the learned latent variables partially or completely collapses with the prior (Bowman et al., 2015; Huang et al., 2018;
Lucas et al., 2019; Razavi et al., 2019; Kingma et al., 2016; Wang et al., 2021).
Up to now, the study of the nature of the cause of the posterior collapse problem has been limited.
There are two main challenges that prevent our understanding of the problem: (1) posterior collapses mainly occur in deep learning, and the landscape of deep neural networks is hard to understand in general; (2) the use of approximate loss functions such as the evidence lower bound (ELBO) complicates the problem.
Consider a problem where one wants to model the data distribution p ( x,y ) through a latent variable z .
The evidence lower bound (ELBO) loss function reads
[Equation]
[Equation]
where q is the approximate distribution, we rely on to approximate the true distribution p .
This loss is more general than the standard ELBO for variational autoencoders (VAE) (Kingma and Welling, 2013).
Meanwhile, it can be seen as the simplest type of loss for a conditional VAE (Sohn et al., 2015), where one aims to model a conditional distribution p ( y ∣ x ) .
The distribution p ( z ) is the prior distribution of the latent variable z and is often a low-complexity distribution such as a zero-mean unit-variance Gaussian.
This loss function thus has a clean interpretation as the sum of a prediction accuracy term (the first term ℓ rec ) that encourages better prediction accuracy and a complexity term (the second term ℓ KL ) that encourages a simpler solution.
Learning under this loss function proceeds by balancing the prediction error and the model simplicity.
Moreover, learning under this loss function has also been used as one of the primary theoretical models in neuroscience (Friston, 2009), and its understanding may also help advance theoretical neuroscience.
This work provides an in-depth study of the posterior collapse problem of Eq.
(1), when the decoder q ( y ∣ z ) and encoder q ( z ∣ x ) are each parametrized by a linear model.
Specifically, our contributions include:
• we find the global minima of a general linear latent variable model that includes the linear VAE as a special case under the Objective (1); • we find the precise condition when posterior collapse occurs, where the global minimum is the origin; • we pinpoint the cause of the posterior collapse to be the excessively strong regularization effect on the mean of the latent variables due to the prior.
To the best of our knowledge, our work is the first to pinpoint the cause of the posterior collapse problem.
This work is organized as follows.
The next section discusses the previous literature.
Section 3 describes the theoretical problem setting.
Section 4 presents our main technical results and analyzes them in detail.
Section 5 presents numerical examples.
The last section concludes this work and points to the remaining open problems.
The Appendix B investigates the effect of the bias term, Appendix C details the effect of a data-dependent encoder variance, and Appendix D treats the case of a learnable decoder variance.
Related Works
Approximate Bayesian Deep Learning .
Bayesian deep learning in general and VAE training, in particular, rely heavily on approximate methods such as the ELBO objective because the exact probabilities are intractable.
The connection of approximate Bayesian learning and probabilistic
PCA (pPCA) has been extensively studied (Nakajima and Sugiyama, 2010; Nakajima et al., 2013, 2015; Lucas et al., 2019).
Causes of Posterior Collapse .
Earlier touches on the problem tend to attribute the cause of posterior collapse to the use of approximate methods, namely, to the use of the ELBO (Bowman et al., 2015;
Huang et al., 2018; Razavi et al., 2019).
Another line of work attributes the cause to the high capacity of modern neural networks (Alemi et al., 2018; Ziyin et al., 2022c).
However, Lucas et al.
(2019) showed that for a simplified linear model, the ELBO is not the cause of posterior collapse because the posterior collapse exists even in the exact posterior.
It also implies that the posterior collapse is not due to the high capacity of the models because linear models have a limited capacity.
Lucas et al.
(2019) then suggested that making the decoder variance learnable can fix the collapse problem and that an unlearnable decoder variance is the cause of the posterior collapse.
However, our results show that this is not the case: for both learnable and unlearnable decoder variance, there exist situations where a collapse happens or does not happen, which implies that the learnability of the decoder variance does not have a causal relation with posterior collapse, nor is it sufficient to fix the problem (Section D).
In terms of the problem setting, ours is also more general than Lucas et al.
(2019) because our result (1) applies to general latent variable models (one example being the conditional VAE (Sohn et al., 2015)) and (2) considers the case of β -VAE with a general β when the decoder variance is learned.
An important implication of our work is that posterior collapses can be a ubiquitous problem for deep-learning-based latent-variable models (not just unique to autoencoding models) and that they share a common cause.
Meanwhile, (L ¨ ucke et al., 2020) shows that posterior collapse can happen due to the tradeoff between the decoding performance and the decoding entropy.
(Shekhovtsov et al., 2022) demonstrated the relationship between model consistency and posterior collapse and suggested that a proper choice of data processing or architecture may alleviate collapse.
Linear Networks .
Deep linear nets have been extensively used to understand the landscape of nonlinear networks.
For example, linear regressors are shown to be relevant for understanding the generalization behavior of modern overparametrized networks (Hastie et al., 2019).
Saxe et al.
(2013) used a two-layer linear network to understand the dynamics of learning nonlinear networks.
The linear nets are the same as a linear regression model in terms of expressivity.
However, the loss landscape is highly complicated due to depth.
(Kawaguchi, 2016; Hardt and Ma, 2016; Laurent and Brecht, 2018; Ziyin et al., 2022a).
Our work essentially studies the loss landscape of linear networks.
While each encoder and decoder we use consists of a single linear layer, they effectively constitute a two-layer linear network when trained together.
Problem Setting
We consider a general linear latent variable model with input space x ∈ R D 0 , latent space z ∈ R d 1 , and target space
y ∈ R
d 2 .
In general, y = f ( x ) is an arbitrary function of x .
When the target y is identical to the input x , it reduces to the standard VAE.
The VAE formalism assumes that there is an intermediate “latent variable” z that captures the data generation process.
In the main text, the encoder and decoder are linear transformations without bias terms, and the learnable bias is treated in Appendix B, which shows that the effect of the bias terms is equivalent to centering both the input and target to be zero-mean ( x → x − E [ x ] , y → y − E
[ y ] ).
Incorporating the bias terms thus does not affect the main results.
Specifically, the encoder is defined as z = W ⊺ x + ϵ , where ϵ ∼ N( 0 , Σ ) is the noise distribution introduced by the reparameterization trick where the variance matrix
Σ = diag ( σ 21 ,...,σ 2 d 1 ) is assumed to be diagonal and independent from x .
The decoder parametrizes the distribution p ( y ∣ z ) =
N( Uz,η 2dec I ) , where the variance η 2dec
I is to be isotropic and input-independent.
In alignment with the standard practice, we also assume the prior distribution of latent variable p ( z ) =
N( 0 ,η 2enc I ) is an isotropic normal distribution, and the encoding variances matrix Σ is learned from the data distribution while η 2dec is not learnable.
Lastly, we weigh the KL term by a coefficient β , which is a common practice in VAE training (Higgins et al., 2016).
Hence, the objective of such a linear model reads,
[Equation]
ℓ mean ⎤⎥⎥⎥⎥⎥⎥⎥⎥⎦ + d
[Equation]
=
[Equation]
2 i
[Equation]
2 i
[Equation]
) ·„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„‚„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„¶ ℓ var , (5) where A ∶= E x [ xx ⊺ ] is the second moment of the input data.
Note that a crucial feature of the KL term is that it decomposes into two terms, one that regularizes the variance of z ( ℓ var ) and another that regularizes the mean of z ( ℓ mean ).
We will see that it is precisely the ℓ mean term that causes the posterior collapse.
Eq.
(5) has ignored the partition function of the decoder because we treat η dec as a constant.
We study the case of a learnable η dec in Section D.
In comparison to the previous works (Lucas et al., 2019) that have treated the case of a learnable η dec , our result is more general because our result also considers the effect of β and allows for the case
d 1 ≥ d 2 .
It is also worth commenting on the difference between this setting and that of the pPCA setting (Nakajima et al., 2015): (1) the effect of β is included in the VAE loss, (2) the prior of VAE is over the latent variable, whereas pPCA has it over the model parameters, and (3) the model can be overparametrized ( d 1 ≥ d 0 ).
Notation .
To summarize, we use x,y , and z to denote the input variable, latent variable, and target variable, respectively.
E x denotes the expectation over the training set.
A ∶= E x [ xx T ] is the second moment matrix of the input x .
A is thus positive semidefinite by definition.
The eigenvalue decomposition of A is A = P A Φ P ⊺ A , where Φ = diag ( ϕ 1 , ⋯ ,ϕ d 0 ) is the diagonal matrix for all d 0 ≤ D 0 positive eigenvalues and P A
=
[ p 1 , ⋯ ,p d 0 ] are matrices by concatenating d 0 eigenvectors p i
∈ R D 0 .
W and U are learnable linear transformation matrices for the linear encoding and decoding processes.
Σ is the learnable diagonal latent variance matrix for encoder with diagonal entries σ i .
η enc is the standard deviation of the prior distribution p ( z ) .
η dec is the standard deviation of decoded samples.
A frequently used quantity is a whitened and rotated x : ˜ x ∶= Φ − 12 P ⊺ A x .
Note that this transformation can be inverted: x = P A
Φ 12 ˜ x .
We see that E x ˜ x ˜ x ⊺ =
I .
Furthermore, we define
Z ∶= E ˜ x
[ y ˜ x ⊺ ] =
E x [ yx ⊺ P A
Φ − 12 ] ∈ R d 2 × d 0 .
Let Z = F Σ Z G ⊺ be the singular value decomposition of Z , where F ∈ R d 2 × d 2 and G ∈ R d 0 × d 0 are two orthogonal matrices.
Σ Z ∈ R d 2 × d 0 is a rectangular diagonal matrix with d ∗ =
min ( d 0 ,d 2 ) singular values of Z in the non-increasing order, i.e., ζ 1 ≥ ζ 2 ≥ ⋯ ≥ ζ d ∗ ≥ 0 .
Main Results
This section discusses the main results, whose proofs are presented in Appendix E. While Σ is often a learnable parameter, we first assume that the KL term is sufficiently strong such that σ 1 = ⋯ =
σ d 1 ≈ η enc is close to the prior value.
We then compare with the case when it is learnable, and this comparison reveals that an optimizable σ i is not essential to the posterior collapse problem.
4.1 General Result
In this section, we prove two results that will be useful for understanding the nature of the VAE training objective and will be useful for us to find the global minimum.
We first show that the VAE objective is equivalent to a matrix factorization problem with a special type of regularization.
Proposition 1.
Let ˜ x ∶=
Φ − 12 P ⊺ A x , Z ∶= E ˜ x
[ y ˜ x ⊺ ] , and
[Equation]
Given a fixed Σ , the minimizer of L VAE ( U,W, Σ ) is ( U ∗ ,W ∗ ) , where W ∗ is any solution of Φ 12 P ⊺ A W = V ∗ .
Proof sketch .
The term ℓ var is irrelevant to finding the optimal U ∗ and V ∗ when Σ is fixed.
Thus, the relevant objective can be obtained with the change of variables ˜ x
= Φ − 12 P ⊺ A x . ◻
The condition Φ 12 P ⊺
A W = V ∗ shows that when the data is low-rank, each solution ( U ∗ ,V ∗ ) corresponds to a manifold of solutions in the original parameter space.
The effective loss L ( U,V ) can be compared with the regularized singular value decomposition problem (Zheng et al., 2018).
We see that the first term is the standard matrix factorization objective, while the second and third are unique regularization effects due to the VAE structure and the ELBO objective.
In addition, the term Σ in the second is the strength of the regularization for the norm of U , and a crucial difference with standard regularized matrix factorization is that Σ is also a learnable matrix.
The next proposition finds, for any fixed Σ , the global minima ( U ∗ ,V ∗ ) of Eq.
In particular, the learning is characterized by the learning of the singular values of U and V .
Proposition 2.
The optimal solution ( U ∗ ,V ∗ ) of min U,V L ( U,V ) is given by
[Equation]
where F ∈ R d 2 × d 2 and G ∈ R d 0 × d 0 are orthogonal matrices derived by the SVD of Z , P is an arbitrary orthogonal matrix in R d 1 × d 1 , and Λ ∈ R d 2 × d 1 and Θ ∈ R d 0 × d 1 are rectangular diagonal matrices with the diagonal elements
[Equation]
For convention, we let ζ i =
when i > d ∗ = min ( d 0 ,d 2 ) .
Proof sketch .
The optimal V ∗ is a function of U under the zero gradient condition.
Thus, the objective reduces to single-variate with respect to U .
The optimal U ∗ is constructed by its SVD
U = Q Λ P , where the optimal Q ∗ and Λ ∗ can be determined given the SVD of Z , and P is left as a free orthogonal matrix.
V ∗ is determined once U ∗ is obtained.
◻
The readers are recommended to examine the form of the solutions closely.
There are a few interesting features of the global minimum.
One note that the sign of the term ζ i − √ βσ i η dec / η enc is crucial, and can encourage the parameters U and V to be low-rank.
Recall that σ i is the eigenvalue value of ZZ T
=
E
[ y ˜ x ⊺ ] E [ y ˜ x ⊺ ] ⊺ , one can roughly identify ζ 2 i as the the strength of the alignment between the input x and the target y .
To see this, consider a simplified scenario where the target y =
γMx is a linear function of the input, where γ is the overall strength of the signal and ∣∣ M ∣∣ =is a normalized orientation matrix, then
[Equation]
which is a positive semidefinite matrix.
We see that there are two distinctive sources of contribution to the magnitude of the eigenvalues of ZZ T .
Its eigenvalues are large if either the overall strength γ is large or if the orientation matrix M aligns well with the covariance of the input feature A .
Additionally, in the case of VAE, γM =
I , and ZZ T
=
A is nothing but the covariance of input features, and ζ 2 i are the eigenvalues of A .
4.2 Linear VAE without Learnable Σ
We first consider the case where σ i is a constant that is completely determined by the prior: σ i =
η enc .
This allows us to find a simplified form for the global minimum.
The proof follows by plugging σ i =
η enc into Proposition 2.
Theorem 1.
Let σ i =
η enc for all i .
Then, the global minimum has
[Equation]
There are three interesting observations of the global minimum.
First of all, it depends crucially on the sign of ζ
i − √ βη dec for all i .
When the sign is negative for some i , the learned model becomes low-rank.
Namely, some of the dimensions collapse with the prior.
When the signs are all negative, we have a complete posterior collapse: both U and V are identically zero, so the latent variables have a distribution identical to the prior.
A complete posterior collapse happens if and only if max i ζ
i − √ βη dec ≤ 0 .
A partial posterior collapse happens if there exists i such that ζ 2 i − √ βη dec ≤ 0 .
These two conditions give the precise conditions of posterior collapse in this scenario.
This implies that having a sufficiently small β will always prevent posterior collapse.
The second observation is that the effect of β is identical to that of η dec because √ β and η dec always appear together, and so one alternative way to fix posterior collapses is to use a sufficiently small η dec .
From a Bayesian perspective, the latter method of tuning η dec is better because η dec comes directly from the (assumed) likelihood p ( x ∣ η dec ) .
In contrast, the β parameter is only an implementation technique that has obscure meaning in the Bayesian framework.
Therefore, using a small η dec can be a fix to the problem that is justified by the Bayesian principle.
The third observation is that the condition for posterior collapse is completely independent of the parameter η enc , which is the desired variance according to the prior p ( z ) .
This means that under a Gaussian assumption, the prior does not affect the posterior collapse at all.
Lastly, one also notices a potential problem.
The eigenvalue of the second layer U increases with √ βη dec , while the first layer decreases with √ βη dec , and so having a too-small β or η dec causes the model to have a very large norm, which can cause a significant problem for both empirical optimization and generalization.
This problem is well-known in the studies about the use of Lregularization in deep learning: suppose we apply weight decay to two different layers of a ReLU net, and decrease the weight decay strength of one layer to zero, then the norm of this layer will tend to infinity, and the norm of the other layer will tend to zero (Mehta et al., 2018).
However, in the next section, we will see that this problem is miraculously solved for VAE when σ i is learnable.
4.3 Linear VAE with Learnable Σ
Now, we consider the more general case of a learnable Σ .
In practice, Σ is often dependent on the input x .
We make the simplification that Σ is just a data-independent optimizable diagonal matrix, which is the common assumption in the related works (Lucas et al., 2019).
In Section C, we consider the case when Σ is data-dependent and show that our result remains unchanged.
The following Corollary gives the optimal training objective as a function Σ and is a direct consequence of proposition 2.
Corollary 1.
[Equation]
where the indicator 1 f > 0 = 1 when the corresponding inequality condition f > 0 is true, and 1 f > 0 = 0 otherwise.
The constant term ∑ d ∗ i =
d 1 + 1 ζ 2 i in Equation (11) only appears when the latent dimension d 1 is less than d ∗ =
min ( d 0 ,d 2 ) .
This is the common situation for VAE applications.
It indicates that the model learns the large eigenvalues and ignores the small eigenvalues.
This means that to find the optimal σ i of Eq.
(5), one only has to find the global minimum of a reduced objective:
[Equation]
·„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„‚„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„¶ ∶= l i ( σ i ) ⎤⎥⎥⎥⎥⎥⎥⎥⎥⎥⎦ + constant
[Equation]
The optimal σ ∗ i can thus be obtained by minimizing each l i independently: σ ∗ i = arg min σ >
l i ( σ ) .
Proposition 3.
The optimal σ ∗ i of l i ( σ ) is
[Equation]
This proposition gives an explicit expression for σ ∗ i .
On the one hand, we see that there is a threshold value for β .
If β is sufficiently large, σ i will be identical to the prior value η enc , in agreement with our assumption in the previous section.
On the other hand, the learned variance σ ∗ i is a function of β if β is below a threshold.
We will see that this threshold is the necessary and sufficient condition for posterior collapse to happen in a learnable Σ setting.
Thus, the learned variance being identical to the prior variance is also a signature of posterior collapse.
The following theorem gives the precise form of the global minimum.
Theorem 2.
The global minimum of L VAE ( U,W, Σ ) is given by
[Equation]
W ∗ is the solution of
[Equation]
16) where F and G are derived by the SVD of Z
, P is an arbitrary orthogonal matrix in R d 1 × d 1 , and Λ =
diag ( λ 1 ,...,λ d 1 ) and Θ = diag ( θ 1 ,...,
θ d 1 ) are diagonal matrices such that
[Equation]
when ζ 2 i >
[Equation]
otherwise .
[Equation]
The optimal Σ ∗ = diag ( σ ∗ 21 ,...,σ ∗ 2 d 1 ) such that
[Equation] [Equation]
for i ≤ min ( d 0 ,d 2 ) .
For i > min ( d 0 ,d 2 ) , σ ∗ i = 0 .
Proof.
The optimal solution U ∗ ,W ∗ , Σ ∗ are obtained by combining proposition 1, 2, and 3.
Comparing with the solution in section 4.2, one notices two things: (a) the conditions for complete or partial posterior collapse remain unchanged, which implies that a learnable latent variance is neither qualitatively nor quantitatively relevant for the posterior collapse problem even though the functional form of the eigenvalues changed; (b) the magnitude of each of the two layers no longer scales with √ βη dec , and so using a small β or η will not directly cause the model to diverge in the norm, which suggests using that making Σ learnable can have the unexpected practical advantage of stabilizing the training.
Additionally, one also notices that βη 2dec has the effect of keeping the learned model low-rank by removing all the eigenvalues of the learned model below it.
This can be directly compared with the effect of using a latent dimension smaller than the input dimension:
d 1 < d 0 .
In the latter case, the smallest d 1 − d 0 singular values are also pruned.
There is a difference between the two types of low-rankness: using a large βη 2dec both removes all the singular values below it and shrinks the remaining ones while using a small latent dimension only removes the smaller singular values without affecting the rest.
This is similar to the difference between soft thresholding estimation and hard thresholding estimation in statistics (Wasserman, 2013).
This suggests that partial posterior collapses are not necessarily undesirable because, during a partial posterior collapse, the latent variable models automatically perform a degree of sparse learning, which is theoretically understood to help denoising the signal and lead to better generalization (Markovsky, 2012).
That being said, complete posterior collapse should always be avoided.
4.4 Learnable η dec
Our result can also be extended to the case when η dec is learnable, which has been suggested by
Lucas et al.
as a remedy for the posterior collapse.
To do this, we need to include the partition function of the decoder, proportional to log η 2dec , that has been ignored in Eq.
We present the detailed analysis in Section D.
Our analysis shows that even if η dec is learnable, posterior collapse can happen for some datasets.
In addition to the fact that it is also possible for collapses not to happen when η dec is not learned, we conclude that η dec does not have a causal relation with posterior collapse.
Our analysis also suggests a way to fix posterior collapse for VAE: make η dec learnable and set β < d 2 / d ∗ .
Note that the condition β < d 2 / d ∗ is tight in the sense that if it does not hold, then there exists a data distribution such that complete collapse can happen.
This condition also highlights that it is important to introduce the β coefficient for VAEs because, for VAE, d 2 / d ∗ = 1 , and this condition translates to β < 1 ; namely, vanilla VAE cannot avoid complete collapse.
Our result also implies that β has a highly nonlinear effect on learning depending on the architecture.
For example, when the model is underparameterized ( d 1 < d ∗ ), using a small β does not cause any problem, whereas for an overparametrized model, a small β causes the decoder variance to converge towards 0 .
4.5 Implications
Our main results have implications for both the problem of posterior collapse and the practice of latent variable models in general.
The cause of posterior collapse .
One important implication is the identification of the cause of the posterior collapse problem and the potential ways to fix it.
Our results suggest that
• a learnable (data-dependent or not) latent variance is not the cause of posterior collapse; • changing the variance of the prior cannot fix or influence the posterior collapse problem; • comparing with the results in Lucas et al.
(2019), η dec being learnable or not is causally related to the posterior collapse problem; • the values of η dec and β are crucial for the posterior collapse; • choosing appropriate β is still needed: a sufficiently small β can avoid posterior collapse.
Note that the effect of a small β (large η dec ) weakens the prior (reconstruction) term, and so the cause of the posterior collapse must be the competition between the prior term, which regularizes the complexity of the model, and the likelihood term, which encourages accurate recognition/reconstruction.
Our results suggest that one can ignore the effect of the ℓ var term in studying the mechanism of posterior collapse.
Ignoring the ℓ var Eq.
5, one sees that the posterior collapse is caused by the competition between the likelihood and ℓ mean , which is precisely the regularization effect on the mean of z .
There is an interesting alternative perspective on the nature of the posterior from the viewpoint of the loss landscape geometry.
The following theorem states that the origin (where all parameters are zero) is either a saddle or the global minimum for this problem.
Since we have shown that σ i does not affect the collapse, we simply let σ i =
η enc as in Section 4.2.
L rec KL
L VAE     0.
i
L rec KL
L VAE    0.
i
Theorem 3.
The Hessian of Eq.
5 at 0 is positive semidefinite if and only if it is the global minimum.
The surprising aspect is that for the latent variable model, there is no intermediate case where the origin is a local minimum but not global.
Therefore, the origin is, in fact, a very special point in the landscape of a latent variable model, in the sense that a key global property of the landscape (namely, the global minimum) is determined by the local geometry of the model at the origin.
Noting that our model can be seen as a direct generalization of the Bayesian linear regression to a deeper architecture, it also becomes reasonable to suspect that the posterior collapse problem is a unique problem of deep learning because the standard Bayesian linear regression does not suffer posterior collapse because the origin can never be a local maximum of the posterior (Bishop and Nasrabadi, 2006).
Dai et al.
(2020) also finds the origin to be a very special point in a general deep nonlinear
VAE structure and that it can be a local minimum under various settings.
However, the implication of our work is broader.
The origin is not only a special point for the autoencoding model families but can actually be a special point for a very broad of model classes (namely, the model class of general latent variable models).
The problem of posterior collapse is thus not limited to autoencoders but can also be relevant to common regression and classification tasks.
Connection to other types of collapses .
Our result suggests that there are some interesting connections between the posterior collapse phenomenon and the neural collapse phenomenon in supervised learning (Papyan et al., 2020) and dimensional collapse phenomenon in self-supervised learning (SSL)
(Jing et al., 2021).
Ziyin et al.
(2022a) and Ziyin and Ueda (2022) shows that the neural collapse phenomenon for a two-layer model can be understood through the change of the stability at the origin, which is determined by the competition between the signal strength ( E [ xy ] ) of the data distribution and the regularization strength of weight decay.
For SSL, Ziyin et al.
(2022b) also shows that the stability of the origin is important and that it is decided by the competition between the level of data variation and the data augmentation strength.
Our result suggests that the posterior collapse problem can also be understood through the stability at the origin.
This might imply that there could be some universal cause of all these collapses that have been discovered independently in different subfields of deep learning, and one important future direction would be to study these phenomena from a unified perspective.
Insights for latent variable model practices .
While we have primarily focused on discussing the phenomena of posterior collapse, our results also shed light on latent variable models (including
VAE) in practice when there is no complete posterior collapse.
Specifically, our results suggest that
[Figure]
• latent variable models perform sparse learning through soft thresholding or hard thresholding or both; • thus, partial posterior collapse may actually be desirable; • making the latent variance learnable can help stabilize training and avoid divergence of model parameters; • when η dec is not learned, the effect of increasing β is identical to the effect of decreasing η dec ; • when η dec is learned, one needs to pay special care to choose a suitable β .
Numerical Examples
This section empirically examines our theoretical claims for linear models and demonstrates that our key theoretical insights generalize well to nonlinear models and natural data.
Setting.
We illustrate our results on both synthetic data and natural data.
For synthetic data, we sample input data x from multivariate normal distribution N( 0 ,A ) , and target data y =
Mx is obtained by a linear transformation.
Specifically, we choose d 0 = d 2 = 5 .
As an example of natural data, we also experiment with the standard MNIST data.
Following common practices, we choose η dec = η enc = 1 .
For non-linear VAE models, we consider two-layer fully connected neural networks for the encoder and decoder with both ReLU and Tanh activation functions and with hidden dimension d h .
For synthetic dataset d h = 8 , and d h = 2048 for real-world data.
In contrast to our assumption that the variances Σ of encoded z are independent from the input x , we parameterize the variance of each encoded z by a linear transformation or a two-layer neural network, i.e., Σ ( x ) =
[ Linear / MLP ]( x ) .
This data-dependent modeling is closer to the common practice, and the comparison can justify the correctness of our theory.
The model is optimized by Adam with a learning rate of 10 − 3 .
The results are reported after the convergence.
For MNIST, the learning rate is 10 − 4 .
Results .
Linear models are found to agree precisely with the theoretical results, so we only present the results in the appendix.
We focus on exploring the nonlinear models in the main text.
We first consider a simple regression task with MLP encoder and decoders with the ReLU activation (Figure 1).
Here, we see that the theoretical prediction of loss function L VAE agrees well with empirical observation.
Moreover, the threshold of complete posterior collapse is also perfectly predicted.
For completeness, we also present the case when (1) the activation is Tanh in Appendix A.1.
We note that the results are similar.
The observation is similar to the standard MNIST dataset with a nonlinear encoder and decoder.
See Figure 2.
For illustration, we also present the generated MNIST images by non-linear β -VAE trained with different choices of β in Figure 3.
The latent dimension is five as described before.
When there are 5 non-collapsed modes, the generated images are both sharp and contain meaningful variations.
As the number of remaining non-collapsed modes reduces to zero, we see that the generated images become increasingly blurred, and the variation between the data also diminishes.
When the model completely collapses, the model outputs a constant, as the theory suggests.
Moreover, we note that the values of β are chosen according to the theoretical thresholds for each mode to collapse, i e, the top-5 ζ i are [ 5 . 12 , 3 . 74 , 3 . 25 , 2 . 84 , 2 . 57 ] .
We see that the theoretical thresholds provide good predictive power for the behavior of mode collapse qualitatively.
Outlook
In this work we have tackled the problem of posterior collapse from a loss landscape point of view.
Our work also contributes to the fundamental theory of deep learning.
The linear VAE architecture can be seen as a deep linear model with two layers, whose loss landscape is highly nontrivial.
In this perspective, our results advance those results in Ziyin et al.
(2022a), where the dimension of the output space is limited to 1d.
The limitation of our work is obvious: our theory only deals with the landscape, and it is unclear how the dynamics of gradient-based methods could contribute to the collapse problem.
In fact, there is strong evidence that stochastic gradient descent can bias the model towards low-rank or sparse solutions (Arora et al., 2019; Ziyin et al., 2021), and, in the context of posterior collapse, these are precisely the collapsed solutions.
One important future direction is thus to study the role of dynamics in influencing posterior collapse.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix. 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[No] The experiments are only for demonstration and are straightforward to reproduce following the theory.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[No] The fluctuations are visually negligible.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] They are done on a single 3080Ti GPU.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[N/A] (b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[-⋃ ˆ d 1 p = 
-]
[-G ′ ( s ˆ d ) = 
-]
[N/A]
[-(1) t ( 1 ) i > 0 : ∂J Q ∗ ∂t i > 0 when t i > t ( 1 ) i and ∂J Q ∗ ∂t i < 
-]