Posterior Collapse of a Linear Latent Variable Model
Anonymous Author(s)
AffiliationAddress
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 latentvariable model that includes linear variational autoencoders as a special case, weprecisely identify the nature of posterior collapse to be the competition betweenthe likelihood and the regularization of the mean due to the prior.
Our result alsosuggests that posterior collapse may be a general problem of learning for deeperarchitectures and deepens our understanding of Bayesian deep learning.
Introduction
Bayesian approaches to deep learning have attracted much attention because they allow for a moreprincipled 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 unresolvedproblem for the Bayesian deep learning practice is the problem of posterior collapse, where theposterior distribution of the learned latent variables partially 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 or the cause of the posterior collapseproblem is 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 ishard to understand in general; (2) the use of approximate loss functions such as the evidence lowerbound (ELBO) complicates the problem.
Consider a problem where one wants to model the data distribution p ( x,y ) through a latent variablez , 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 lossis 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 priordistribution of the latent variable z and is often a low-complexity distribution such as a zero-meanunit-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 complexityterm (the second term ℓ KL ) that encourages a simpler solution.
Learning under this loss functionproceeds by balancing the prediction error and the model simplicity.
Moreover, learning under thisloss 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 providesan in-depth study of the posterior collapse problem of Eq.
(1), when the decoder q ( y ∣ z ) and encoderq ( 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 linearVAE as a special case under Objective (1);• we find the precise condition when posterior collapse occurs, where the global minimum isthe origin;• we pinpoint the cause of the posterior collapse to be the excessively strong regularizationeffect 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 collapseproblem.
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 resultsand analyzes them in detail.
Section 5 presents numerical examples.
The last section concludes thiswork and points to the remaining open problems.
The Appendix D treats the case of a learnabledecoder variance, Appendix B investigates the effect of the bias term, and Appendix C details theeffect of a data-dependent encoder variance.
Related Works
Approximate Bayesian Deep Learning .
Bayesian deep learning in general and VAE training, inparticular, rely heavily on approximate methods such as the ELBO objective because the exactprobabilities are intractable.
The connection of approximate Bayesian learning and probabilistic
PCA (pPCA) is well known (Nakajima and Sugiyama, 2010; Nakajima et al., 2013, 2015; Lucaset al., 2019).
Causes of Posterior Collapse .
Earlier touches on the problem tend to attribute the cause of posteriorcollapse 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 capacityof modern neural networks (Alemi et al., 2018; Ziyin et al., 2022b).
However, Lucas et al.
(2019)showed that for a simplified linear model, the ELBO is not the cause of posterior collapse becausethe posterior collapse exists even in the exact posterior.
It also implies that the posterior collapse isnot 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 andthat a nonlearnable decoder variance is the cause of posterior.
However, our results show that thisis not the case: for both learnable and nonlearnable decoder variance, there exist situations where acollapse happens or does not happen, which implies that the learnability of the decoder variance doesnot 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 fordeep-learning-based latent-variable models (not just unique to autoencoding models) and that theyshare a common cause.
Meanwhile, (L ¨ ucke et al., 2020) shows that posterior collapse can happendue to the tradeoff between the decoding performance and the decoding entropy.
(Shekhovtsovet al., 2022) demonstrated the relationship between model consistency and posterior collapse andsuggested 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 it to understand the dynamics of learning of nonlinear networks.
The linear nets are the same as a linear regression model in terms of expressivity.
However, theloss landscape is highly complicated due to depth.
(Kawaguchi, 2016; Hardt and Ma, 2016; Laurentand Brecht, 2018; Ziyin et al., 2022a).
Our work essentially studies the loss landscape of linearnetworks.
While each encoder and decoder we use consists of a single linear layer, they effectivelyconstitute a two-layer linear network when trained together.
Problem Setting
We consider (generalized) variational autoencoders 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 isidentical to the input x , it reduces to the standard VAE.
The VAE formalism assumes that there isan intermediate “latent variable” z that captures the data generation process.
In the main text, the encoder and decoder are considered as linear transformations without bias terms, and the learnablebias is treated in Appendix B, which shows that the effect of the bias terms is equivalent to centeringboth the input and target to be zero-mean ( x → x − E [ x ] , y → y − E
[ y ] ).
Incorporating the bias termsthus 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 variancematrix Σ = diag ( σ 21 ,...,σ 2 d 1 ) is assumed to be diagonal and independent from x .
The decoderparametrizes the distribution p ( y ∣ z ) =
N( Uz,η 2dec I ) , where the variance η 2dec
I is to be isotropic andinput-independent.
In alignment with the standard practice, we also assume the prior distribution oflatent variable p ( z ) =
N( 0 ,η 2enc I ) is an isotropic normal distribution, and the encoding variancesmatrix Σ is learned from the data distribution while η 2dec is not learnable.
Lastly, we weigh the KLterm by a coefficient β , which is a common practice in VAE training (Higgins et al., 2016).
Hence,the objective of linear VAE reads,
[Equation]
η 2enc Tr
[Equation]
·„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„‚„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„¶ ℓ mean ⎤ ⎥ ⎥⎥⎥⎥⎥⎥⎥⎦ + d∑ i =
β 2 ( σ 2 i η 2enc − 1 − log σ 2 i η 2enc ) ·„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„‚„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„„¶ ℓ var , (5)
where A ∶= E x
[ xx ⊺ ] is the second moment of the input data.
Note that a crucial feature of the KLterm is that it decomposes into two terms, one that regularizes the variance of z ( ℓ var ) and anotherthat regularizes the mean of z ( ℓ mean ).
We will see that it is precisely the ℓ mean term that causes theposterior collapse.
Eq.
(5) has ignored the partition function of the decoder because we treat η dec asa 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 becauseour result also considers the effect of β and allows for the case
d 1 ≥ d 2 .
It is also worth commentingon the difference between this setting and that of the pPCA setting (Nakajima et al., 2015): (1) theeffect of β is included in the VAE loss, (2) the prior of VAE is over the latent variable,
whereaspPCA 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 targetvariable, 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 eigenvaluedecomposition of A is A =
P A Ψ P ⊺ A , where Ψ = diag ( ϕ 1 , ⋯ ,ϕ d 0 ) is the diagonal matrix for all 115 d 0 ≤ D 0 positive eigenvalues and P A =
[ p 1 , ⋯ ,p d 0 ] are matrices by concatenating d 0 eigenvectorsp
i ∈ R d 0 .
W and U are learnable linear transformation matrice 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 ofdecoded samples.
A frequently used quantity is a whitened and rotated x : ˜ x ∶= Φ − 12 P ⊺ A x .
Note thatthis 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 decompositionof 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 oftena 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 thiscomparison 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 VAEtraining objective and will be useful for us to find the global minimum.
We first show that the VAEobjective 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]
∗ ,
[Equation]
∗
[Equation]
∶= arg min
[Equation]
Given a fixed Σ , the minimizer of L VAE ( U,W, Σ ) is ( U ∗ ,W ∗ ) , where W ∗ 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 ) canbe compared with the regularized singular value decomposition problem (Zheng et al., 2018).
Wesee that the first term is the standard matrix factorization objective, while the second and third is aunique regularization effect 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 withstandard 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 anarbitrary orthogonal matrix in R d 1 × d 1 , and Λ ∈ R d 2 × d 1 and Θ ∈ R d 0 × d 1 are rectangular diagonalmatrices 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, theobjective 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 afree 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 iscrucial, and can encourage the parameters U and V to be low-rank.
Recall that σ i is the eigenvaluevalue of ZZ T
=
E
[ y ˜ x ⊺ ] E [ y ˜ x ⊺ ] ⊺ , one can roughly identify ζ 2 i as the the strength of the alignmentbetween the input x and the target y .
To see this, consider a simplified scenario where the targety = γMx is a linear function of the input, where γ is the overall strength of the signal and ∣∣ M ∣∣
=
1is a normalized orientation matrix, then
[Equation]
which is a positive semidefinite matrix.
We see that there are two distinctive sources of contributionto 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 inputfeatures, 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 simplyplugging σ 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 onthe sign of ζ
i − √ βη dec for all i .
When the sign is negative for some i , the learned model becomeslow-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 havea distribution identical to the prior.
A complete posterior collapse happens if and only if max i ζ i − 179 √ βη 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 impliesthat having a sufficiently small β will always prevent posterior collapse.
The second observation isthat the effect of β is identical to that of η dec because √ β and η dec always appear together, and soone alternative way to fix posterior collapses is to use a sufficiently small η dec .
From a Bayesianperspective, the latter method of tuning η dec is better because η comes directly from the (assumed)likelihood p ( x ∣ η ) .
In contrast, the β parameter is only an implementation technique that has obscuremeaning in the Bayesian framework.
Therefore, using a small η dec can be a fix to the problem that isjustified by the Bayesian principle.
The third observation is that the condition for posterior collapseis completely independent of the parameter η enc , which is the desired variance according to the priorp ( z ) .
This means that under a Gaussian assumption, the prior does not affect the posterior collapseat all.
Lastly, one also notices a potential problem.
The eigenvalue of the second layer U increases with 192 √ βη dec , while the first layer decreases with √ βη dec , and so having a too-small β or η dec causesthe model to have a very large norm, which can cause a significant problem for both empiricaloptimization and generalization.
This problem is well-known in the studies about the use of L 2regularization in deep learning: suppose we apply weight decay to two different layers of a ReLUnet, and decrease the weight decay strength of one layer to zero, then the norm of this layer will tendto infinity, and the norm of the other layer will tend to zero (Mehta et al., 2018).
However, in thenext 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 onthe input x .
We make the simplification that Σ is just a data-independent optimizable diagonalmatrix, 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 remain unchanged.
Thefollowing Corollary gives the optimal training objective as a function Σ and is a direct consequenceof 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 lessthan d ∗ =
min ( d 0 ,d 2 ) .
This is the common situation for VAE application.
It indicates that themodel learns the large eigenvalues and ignores the small eigenvalues.
This means that to find the optimal σ i of Eq.
(5), one only have to find the global minimum of a reduced objective:
[Equation]
∶= l i ( σ i ) ⎤⎥⎥⎥⎥⎥⎥⎥⎥⎥⎦ + constant.
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 thresholdvalue for β .
If β is sufficiently large, σ i will be identical to the prior value η enc , in agreement withour 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 forposterior collapse to happen in a learnable Σ setting.
Thus, the learned variance being identical tothe prior variance is also a signature of posterior collapse.
The following theorem gives the preciseform 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]
We note that for σ i = 0 when i > min ( d 0 ,d 2 ) .
The optimal Σ ∗ = diag ( σ ∗ 21 ,...,σ ∗ 2 d 1 ) such that
[Equation]
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 completeor partial posterior collapse remain unchanged, which implies that a learnable latent variance isneither qualitatively nor quantitatively relevant for the posterior collapse problem even though thefunctional form of the eigenvalues changed; (b) the magnitude of each of the two layers no longerscales with √ βη dec , and so using a small β or η will not directly cause the model to diverge in termsof norm, which suggests using that making Σ learnable can have the unexpected practical advantageof stabilizing the training.
Additionally, one also notices that βη 2dec has the effect of keeping the learned model low-rank byremoving all the eigenvalues of the learned model below it.
This can be directly compared withthe 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 theremaining 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 hardthresholding estimation in statistics (Wasserman, 2013).
This suggests that partial posterior collapses are not necessarily undesirable because, during a partial posterior collapse, the latent variablemodels automatically perform a degree of sparse learning, which is theoretically understood to helpdenoising the signal and lead to better generalization (Markovsky, 2012).
That being said, completeposterior 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 partitionfunction of the decoder, proportional to log η 2dec , that has been ignored in Eq.
See 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, weconclude that η dec does not have a causal relation with posterior collapse.
Our analysis also suggestsa 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 thatcomplete 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 anyproblem, whereas for an overparametrized model, a small β causes the decoder variance to convergetowards 0 .
4.5 Implications
Our main results have implications for both the problem of posterior collapse and the practice oflatent variable models in general.
The cause of posterior collapse .
One important implication is the identification of the cause of theposterior 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 causallyrelated to the posterior collapse problem;
• the values of η dec and β are crucial for the posterior collapse;
• making η dec learnable is not sufficient to fix complete 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 causeof 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 ofposterior collapse.
Ignoring the ℓ var Eq.
5, one sees that the posterior collapse is caused by thecompetition between the likelihood and ℓ mean , which is precisely the regularization effect on themean of z .
There is an interesting alternative perspective on the nature of the posterior from the viewpoint ofthe architecture geometry.
The following theorem states that the origin (where all parameters arezero) is either a saddle or the global minimum for this problem.
Since we have shown that σ i doesnot affect the collapse, we simply let σ i =
η enc as in Section 4.2.
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 theorigin is a local minimum but not global.
Therefore, the origin is, in fact, a very special pointin the landscape of a latent variable model, in the sense that a key global property of the landscape
L rec KL
L VAE     0.
i
L rec KL
L VAE    0.
i
(namely, the global minimum) is determined by the local geometry of the model at the origin.
Notingthat our model can be seen as a direct generalization of the Bayesian linear regression to a deeperarchitecture, it also becomes reasonable to suspect that the posterior collapse problem is a uniqueproblem of deep learning because the standard Bayesian linear regression does not suffer posteriorcollapse 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 implicationof our work is broader.
The origin is not only a special point for the autoencoding model families, butcan actually be a special point for a very broad of model classes (namely, the model class of generallatent variable models).
The problem of posterior collapse is thus not limited to autoencoders, butcan also be relevant to the common regression and classification tasks.
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 (includingVAE) in practice when there is no complete posterior collapse.
Specifically, our results suggest that• 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 ofmodel parameters;• when η dec is not learned, the effect of increasing β is identical to the effect of decreasingη dec ;
• when it is learned, one needs to pay special care to choosing a suitable β .
Numerical Examples
This section empirically examines our theoretical claims for linear models and demonstrates that ourkey 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, wesample 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 naturaldata, we also experiment on the standard MNIST data.
Following common practices, we choose
[Figure]
η 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 acitvation functions and with hiddendimension d h .
For synthetic dataset d h = 8 , and d h = 2048 for real-world data.
In contrast toour 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 thecomparison can justify the correctness of our theory.
The model is optimized by Adam with learningrate 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, and so we onlypresent 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 withempirical 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 anonlinear encoder and decoder.
See Figure 2.
For illustration, we also present the generated MNISTimages by non-linear β -VAE trained with different choices of β in Figure 3.
The latent dimensionis 5 as described before.
When there are 5 non-collapsed modes, the generated images are bothsharp and contain meaningful variations.
As the number of remaining non-collapsed modes reducesto zero, we see that the generated images become increasingly blurred and the variation betweenthe data also diminishes.
When the model completely collapses, the model outputs a constant, asthe 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 ] .
Wesee that the theoretical thresholds provide good predictive power for the behavior of mode collapsequalitatively.
Outlook
In this work, we have tackled the problem of posterior collapse from a theoretical point of view.
Our work also contributes to the fundamental theory of deep learning.
The linear VAE architecturecan be seen as a deep linear model with two layers, whose loss landscape is highly nontrivial.
Inthis perspective, our results advance those results in Ziyin et al.
(2022a), where the dimension of theoutput space is limited to 1d.
The limitation of our work is obvious: our theory only deals with linearmodels, and it is unclear to extent the result carries over to the nonlinear case.
Does a nonlinear VAEalways experience posterior collapse when β is large enough?
If not, what is the precise conditionthat posterior collapse can happen?
Moreover, can we design a Bayesian-principled method foravoiding posterior collapse that does not involve hand-tuning β or η dec ?
These are the remainingopen questions whose solutions can help us design principled methods that prevent posterior collapseonce and for all and deepen our understanding of how deep learning works.
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]
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 exper-
imental results (either in the supplemental material or as a URL)?
[No] The experi-
ments 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., typeof GPUs, internal cluster, or cloud provider)?
[Yes] They are done on a single 3080Ti
GPU.4.
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 identifi-
able information or offensive content?
[N/A]
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
[Table]
N
[Table]