The Usual Suspects? Reassessing Blame for VAE Posterior Collapse

1 INTRODUCTION . The variational autoencoder ( VAE ) ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) represents a powerful generative model of data points that are assumed to possess some complex yet unknown latent structure . This assumption is instantiated via the marginalized distribution pθ ( x ) = ∫ pθ ( x|z ) p ( z ) dz , ( 1 ) which forms the basis of prevailing VAE models . Here z ∈ Rκ is a collection of unobservable latent factors of variation that , when drawn from the prior p ( z ) , are colloquially said to generate an observed data point x ∈ Rd through the conditional distribution pθ ( x|z ) . The latter is controlled by parameters θ that can , at least conceptually speaking , be optimized by maximum likelihood over pθ ( x ) given available training examples . In particular , assuming n training points X = [ x ( 1 ) , . . . , x ( n ) ] , maximum likelihood estimation is tantamount to minimizing the negative log-likelihood expression 1n ∑ i− log [ pθ ( x ( i ) ) ] . Proceeding further , because the marginalization over z in ( 1 ) is often intractable , the VAE instead minimizes a convenient variational upper bound given by L ( θ , φ ) , 1 n n∑ i=1 { −Eqφ ( z|x ( i ) ) [ log pθ ( x ( i ) |z ) ] + KL [ qφ ( z|x ( i ) ||p ( z ) ] } ≥ 1n n∑ i=1 − log [ pθ ( x ( i ) ) ] , ( 2 ) with equality iff qφ ( z|x ( i ) ) = pθ ( z|x ( i ) ) for all i . The additional parameters φ govern the shape of the variational distribution qφ ( z|x ) that is designed to approximate the true but often intractable latent posterior pθ ( z|x ) . The VAE energy from ( 2 ) is composed of two terms , a data-fitting loss that borrows the basic structure of an autoencoder ( AE ) , and a KL-divergence-based regularization factor . The former incentivizes assigning high probability to latent codes z that facilitate accurate reconstructions of each x ( i ) . In fact , if qφ ( z|x ) is a Dirac delta function , this term is exactly equivalent to a deterministic AE with data reconstruction loss defined by− log pθ ( x|z ) . Overall , it is because of this association that qφ ( z|x ) is generally referred to as the encoder distribution , while pθ ( x|z ) denotes the decoder distribution . Additionally , the KL regularizer KL [ qφ ( z|x ) ||p ( z ) ] pushes the encoder distribution towards the prior without violating the variational bound . For continuous data , which will be our primary focus herein , it is typical to assume that p ( z ) = N ( z|0 , I ) , pθ ( x|z ) = N ( x|µx , γI ) , and qφ ( z|x ) = N ( z|µz , Σz ) , ( 3 ) where γ > 0 is a scalar variance parameter , while the Gaussian moments µx ≡ µx ( z ; θ ) , µz ≡ µz ( x ; φ ) , and Σz ≡ diag [ σz ( x ; φ ) ] 2 are computed via feedforward neural network layers . The encoder network parameterized by φ takes x as an input and outputs µz and Σz . Similarly the decoder network parameterized by θ converts a latent code z into µx . Given these assumptions , the generic VAE objective from ( 2 ) can be refined to L ( θ , φ ) = 1n n∑ i=1 { Eqφ ( z|x ( i ) ) [ 1 γ ‖x ( i ) − µx ( z ; θ ) ‖22 ] ( 4 ) + d log γ + ∥∥∥σz ( x ( i ) ; φ ) ∥∥∥2 2 − log ∣∣∣∣diag [ σz ( x ( i ) ; φ ) ] 2∣∣∣∣+ ∥∥∥µz ( x ( i ) ; φ ) ∥∥∥2 2 } , excluding an inconsequential factor of 1/2 . This expression can be optimized over using SGD and a simple reparameterization strategy ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) to produce parameter estimates { θ∗ , φ∗ } . Among other things , new samples approximating the training data can then be generated via the ancestral process znew ∼ N ( z|0 , I ) and xnew ∼ pθ∗ ( x|znew ) . Although it has been argued that global minima of ( 4 ) may correspond with the optimal recovery of ground truth distributions in certain asymptotic settings ( Dai & Wipf , 2019 ) , it is well known that in practice , VAE models are at risk of converging to degenerate solutions where , for example , it may be that qφ ( z|x ) = p ( z ) . This phenomena , commonly referred to as VAE posterior collapse ( He et al. , 2019 ; Razavi et al. , 2019 ) , has been acknowledged and analyzed from a variety of different perspectives as we detail in Section 2 . That being said , we would argue that there remains lingering ambiguity regarding the different types and respective causes of posterior collapse . Consequently , Section 3 provides a useful taxonomy that will serve to contextualize our main technical contributions . These include the following : • Building upon existing analysis of affine VAE decoder models , in Section 4 we prove that even arbitrarily small nonlinear activations can introduce suboptimal local minima exhibiting posterior collapse . • We demonstrate in Section 5 that if the encoder/decoder networks are incapable of sufficiently reducing the VAE reconstruction errors , even in a deterministic setting with no KL-divergence regularizer , there will exist an implicit lower bound on the optimal value of γ . Moreover , we prove that if this γ is sufficiently large , the VAE will behave like an aggressive thresholding operator , enforcing exact posterior collapse , i.e. , qφ ( z|x ) = p ( z ) . • Based on these observations , we present experiments in Section 6 establishing that as network depth/capacity is increased , even for deterministic AE models with no regularization , reconstruction errors become worse . This bounds the effective VAE trade-off parameter γ such that posterior collapse is essentially inevitable . Collectively then , we provide convincing evidence that posterior collapse is , at least in certain settings , the fault of deep AE local minima , and need not be exclusively a consequence of usual suspects such as the KL-divergence term . We conclude in Section 7 with practical take-home messages , and motivate the search for improved AE architectures and training regimes that might be leveraged by analogous VAE models . 2 RECENT WORK AND THE USUAL SUSPECTS FOR INSTIGATING COLLAPSE . Posterior collapse under various guises is one of the most frequently addressed topics related to VAE performance . Depending on the context , arguably the most common and seemingly transparent suspect for causing collapse is the KL regularization factor that is obviously minimized by qφ ( z|x ) = p ( z ) . This perception has inspired various countermeasures , including heuristic annealing of the KL penalty or KL warm-start ( Bowman et al. , 2015 ; Huang et al. , 2018 ; Sønderby et al. , 2016 ) , tighter bounds on the log-likelihood ( Burda et al. , 2015 ; Rezende & Mohamed , 2015 ) , more complex priors ( Bauer & Mnih , 2018 ; Tomczak & Welling , 2018 ) , modified decoder architectures ( Cai et al. , 2017 ; Dieng et al. , 2018 ; Yeung et al. , 2017 ) , or efforts to explicitly disallow the prior from ever equaling the variational distribution ( Razavi et al. , 2019 ) . Thus far though , most published results do not indicate success generating high-resolution images , and in the majority of cases , evaluations are limited to small images and/or relatively shallow networks . This suggests that there may be more nuance involved in pinpointing the causes and potential remedies of posterior collapse . One notable exception though is the BIVA model from ( Maaløe et al. , 2019 ) , which employs a bidirectional hierarchy of latent variables , in part to combat posterior collapse . While improvements in NLL scores have been demonstrated with BIVA using relatively deep encoder/decoders , this model is significantly more complex and difficult to analyze . On the analysis side , there have been various efforts to explicitly characterize posterior collapse in restricted settings . For example , Lucas et al . ( 2019 ) demonstrate that if γ is fixed to a sufficiently large value , then a VAE energy function with an affine decoder mean will have minima that overprune latent dimensions . A related linearized approximation to the VAE objective is analyzed in ( Rolinek et al. , 2019 ) ; however , collapsed latent dimensions are excluded and it remains somewhat unclear how the surrogate objective relates to the original . Posterior collapse has also been associated with data-dependent decoder covariance networks Σx ( z ; θ ) 6= γI ( Mattei & Frellsen , 2018 ) , which allows for degenerate solutions assigning infinite density to a single data point and a diffuse , collapsed density everywhere else . Finally , from the perspective of training dynamics , ( He et al. , 2019 ) argue that a lagging inference network can also lead to posterior collapse . 3 TAXONOMY OF POSTERIOR COLLAPSE . Although there is now a vast literature on the various potential causes of posterior collapse , there remains ambiguity as to exactly what this phenomena is referring to . In this regard , we believe that it is critical to differentiate five subtle yet quite distinct scenarios that could reasonably fall under the generic rubric of posterior collapse : ( i ) Latent dimensions of z that are not needed for providing good reconstructions of the training data are set to the prior , meaning qφ ( zj |x ) ≈ p ( zj ) = N ( 0 , 1 ) at any superfluous dimension j . Along other dimensions σ2z will be near zero and µz will provide a usable predictive signal leading to accurate reconstructions of the training data . This case can actually be viewed as a desirable form of selective posterior collapse that , as argued in ( Dai & Wipf , 2019 ) , is a necessary ( albeit not sufficient ) condition for generating good samples . ( ii ) The decoder variance γ is not learned but fixed to a large value1 such that the KL term from ( 2 ) is overly dominant , forcing most or all dimensions of z to follow the prior N ( 0 , 1 ) . In this scenario , the actual global optimum of the VAE energy ( conditioned on γ being fixed ) will lead to deleterious posterior collapse and the model reconstructions of the training data will be poor . In fact , even the original marginal log-likelihood can potentially default to a trivial/useless solution if γ is fixed too large , assigning a small marginal likelihood to the training data , provably so in the affine case ( Lucas et al. , 2019 ) . ( iii ) As mentioned previously , if the Gaussian decoder covariance is learned as a separate network structure ( instead of simply Σx ( z ; θ ) = γI ) , there can exist degenerate solutions that assign infinite density to a single data point and a diffuse , isotropic Gaussian elsewhere ( Mattei & Frellsen , 2018 ) . This implies that ( 4 ) can be unbounded from below at what amounts to a posterior collapsed solution and bad reconstructions almost everywhere . ( iv ) When powerful non-Gaussian decoders are used , and in particular those that can parameterize complex distributions regardless of the value of z ( e.g. , PixelCNN-based ( Van den Oord et al. , 2016 ) ) , it is possible for the VAE to assign high-probability to the training data even if qφ ( z|x ) = p ( z ) ( Alemi et al. , 2017 ; Bowman et al. , 2015 ; Chen et al. , 2016 ) . This category of posterior collapse is quite distinct from categories ( ii ) and ( iii ) above in that , although the reconstructions are similarly poor , the associated NLL scores can still be good . ( v ) The previous four categories of posterior collapse can all be directly associated with emergent properties of the VAE global minimum under various modeling conditions . In contrast , a fifth type of collapse exists that is the explicit progeny of bad VAE local minima . More 1Or equivalently , a KL scaling parameter such as used by the β-VAE ( Higgins et al. , 2017 ) is set too large . specifically , as we will argue shortly , when deeper encoder/decoder networks are used , the risk of converging to bad , overregularized solutions increases . The remainder of this paper will primarily focus on category ( v ) , with brief mention of the other types for comparison purposes where appropriate . Our rationale for this selection bias is that , unlike the others , category ( i ) collapse is actually advantageous and hence need not be mitigated . In contrast , while category ( ii ) is undesirable , it be can be avoided by learning γ . As for category ( iii ) , this represents an unavoidable consequence of models with flexible decoder covariances capable of detecting outliers ( Dai et al. , 2019 ) . In fact , even simpler inlier/outlier decomposition models such as robust PCA are inevitably at risk for this phenomena ( Candès et al. , 2011 ) . Regardless , when Σz ( x ; θ ) = γI this problem goes away . And finally , we do not address category ( iv ) in depth simply because it is unrelated to the canonical Gaussian VAE models of continuous data that we have chosen to examine herein . Regardless , it is still worthwhile to explicitly differentiate these five types and bare them in mind when considering attempts to both explain and improve VAE models .

Summary: The paper theoretically investigates the role of “local optima” of the variational objective in ignoring latent variables (leading to posterior collapse) in variational autoencoders. The paper first discusses various potential causes for posterior collapse before diving deeper into a particular cause: local optima. The paper considers a class of near-affine decoders and characterise the relationship between the variance (gamma) in the likelihood and local optima. The paper then extends this discussion for deeper architecture and vanilla autoencoders and illustrate how this can arise when the reconstruction cost is high. The paper considers several experiments to illustrate this issue.