Posterior Collapse and Latent Variable Non-identifiability

1 Introduction . Variational autoencoders ( VAE ) are powerful generative models for high-dimensional data [ 28 , 46 ] . Their key idea is to combine the inference principles of probabilistic modeling with the flexibility of neural networks . In a VAE , each datapoint is independently generated by a low-dimensional latent variable drawn from a prior , then mapped to a flexible distribution parametrized by a neural network . Unfortunately , VAE often suffer from posterior collapse , an important and widely studied phenomenon where the posterior of the latent variables is equal to prior [ 6 , 8 , 38 , 62 ] . This phenomenon is also known as latent variable collapse , KL vanishing , and over-pruning . Posterior collapse renders the VAE useless to produce meaningful representations , in so much as its per-datapoint latent variables all have the exact same posterior . Posterior collapse is commonly observed in the VAE whose generative model is highly flexible , leading to the common speculation that posterior collapse occurs because VAE involve flexible neural networks in the generative model [ 11 ] , or because it uses variational inference [ 59 ] . Based on these hypotheses , many of the proposed strategies for mitigating posterior collapse thus focus on modifying the variational inference objective ( e.g . [ 44 ] ) , designing special optimization schemes for variational inference in VAE ( e.g . [ 18 , 25 , 32 ] ) , or limiting the capacity of the generative model ( e.g . [ 6 , 16 , 60 ] . ) 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In this paper , we consider posterior collapse as a problem of latent variable non-identifiability . We prove that posterior collapse occurs if and only if the latent variable is non-identifiable in the generative model , which loosely means the likelihood function does not depend on the latent variable [ 40 , 42 , 56 ] . Below , we formally establish this equivalence by appealing to recent results in Bayesian non-identifiability [ 40 , 42 , 43 , 49 , 58 ] . More broadly , the relationship between posterior collapse and latent variable non-identifiability implies that posterior collapse is not a phenomenon specific to the use of neural networks or variational inference . Rather , it can also occur in classical probabilistic models fitted with exact inference methods , such as Gaussian mixture models and probabilistic principal component analysis ( PPCA ) . This relationship also leads to a new perspective on existing methods for avoiding posterior collapse , such as the delta-VAE [ 44 ] or the Ø-VAE [ 19 ] . These methods heuristically adjust the approximate inference procedure embedded in the optimization of the model parameters . Though originally motivated by the goal of patching the variational objective , the results here suggest that these adjustments are useful because they help avoid parameters at which the latent variable is nonidentifiable and , consequently , avoid posterior collapse . The relationship between posterior collapse and non-identifiability points to a direct solution to the problem : we must make the latent variable identifiable . To this end , we propose latent-identifiable VAE , a class of VAE that is as flexible as classical VAE while also being identifiable . Latentidentifiable VAE resolves the latent variable non-identifiability by leveraging Brenier maps [ 36 , 39 ] and parameterizing them with input-convex neural networks [ 2 , 35 ] . Inference on identifiable VAE uses the standard variational inference objective , without special modifications or optimization tricks . Across synthetic and real datasets , we show that identifiable VAE mitigates posterior collapse without sacrificing fidelity to the data . Related work . Existing approaches to avoiding posterior collapse often modify the variational inference objective , design new initialization or optimization schemes for VAE , or add neural network links between each data point and their latent variables [ 1 , 3 , 6 , 8 , 12 , 15 , 16 , 17 , 18 , 21 , 25 , 27 , 32 , 34 , 38 , 44 , 50 , 51 , 52 , 55 , 61 , 62 , 63 ] . Several recent papers also attempt to provide explanations for posterior collapse . Chen et al . [ 8 ] explains how the inexact variational approximation can lead to inefficiency of coding in VAE , which could lead to posterior collapse due to a form of information preference . Dai et al . [ 11 ] argues that posterior collapse can be partially attributed to the local optima in training VAE with deep neural networks . Lucas et al . [ 33 ] shows that posterior collapse is not specific to the variational inference training objective ; absent a variational approximation , the log marginal likelihood of PPCA has bad local optima that can lead to posterior collapse . Yacoby et al . [ 59 ] discusses how variational approximation can select an undesirable generative model when the generative model parameters are non-identifiable . In contrast to these works , we consider posterior collapse solely as a problem of latent variable non-identifiability , and not of optimization , variational approximations , or neural networks per se . We use this result to propose the identifiable VAE as a way to directly avoid posterior collapse . Outside VAE , latent variable identifiability in probabilistic models has long been studied in the statistics literature [ 40 , 42 , 42 , 43 , 49 , 56 , 58 ] . More recently , Betancourt [ 5 ] studies the effect of latent variable identifiability on Bayesian computation for Gaussian mixtures . Khemakhem et al . [ 23 , 24 ] propose to resolve the non-identifiability in deep generative models by appealing to auxiliary data . Kumar & Poole [ 29 ] study how the variational family can help resolve the nonidentifiability of VAE . These works address the identifiability issue for a different goal : they develop identifiability conditions for different subsets of VAE , aiming for recovering true causal factors of the data and improving disentanglement or out-of-distribution generalization . Related to these papers , we demonstrate posterior collapse as an additional way that the concept of identifiability , though classical , can be instrumental in modern probabilistic modeling . Considering identifiability leads to new solutions to posterior collapse . Contributions . We prove that posterior collapse occurs if and only if the latent variable in the generative model is non-identifiable . We then propose latent-identifiable VAE , a class of VAE that are as flexible as classical VAE but have latent variables that are provably identifiable . Across synthetic and real datasets , we demonstrate that latent-identifiable VAE mitigates posterior collapse without modifying VAE objectives or applying special optimization tricks . 2 Posterior collapse and latent variable non-identifiability . Consider a dataset x = ( x1 , . . . , xn ) ; each datapoint is m-dimensional . Positing n latent variables z = ( z1 , . . . , zn ) , a variational autoencoder ( VAE ) assumes that each datapoint xi is generated by a K-dimensional latent variable zi : zi ª p ( zi ) , xi | zi ª p ( xi | zi ; µ ) =EF ( xi | fµ ( zi ) ) , ( 1 ) where xi follows an exponential family distribution with parameters fµ ( zi ) ; fµ parameterizes the conditional likelihood . In a deep generative model fµ is a parameterized neural network . Classical probabilistic models like Gaussian mixture model [ 45 ] and probabilistic PCA [ 10 , 47 , 48 , 54 ] are also special cases of Eq . 1 . To fit the model , VAE optimizes the parameters µ by maximizing a variational approximation of the log marginal likelihood . After finding an optimal µ̂ , we can form a representation of the data using the approximate posterior q¡̂ ( z |x ) with variational parameters ¡̂ or its expectation Eq¡̂ ( z |x ) [ z |x ] . Note that here we abstract away computational considerations and consider the ideal case where the variational approximation is exact . This choice is sensible : if the exact posterior suffers from posterior collapse then so will the approximate posterior ( a variational approximation can not “ uncollapse ” a collapsed posterior ) . That said we also note that there exist in practice situations where variational inference alone can lead to posterior collapse . A notable example is when the variational approximating family is overly restrictive : it is then possible to have non-collapsing exact posteriors but collapsing approximate posteriors . 2.1 Posterior collapse , Latent variable non-identifiability . We first define posterior collapse and latent variable non-identifiability , then proving their connection . Definition 1 ( Posterior collapse [ 6 , 8 , 38 , 62 ] ) . Given a probability model p ( x , z ; µ ) , a parameter value µ = µ̂ , and a dataset x= ( x1 , . . . , xn ) , the posterior of the latent variables z collapses if p ( z |x ; µ̂ ) = p ( z ) . ( 2 ) The posterior collapse phenomenon can occur in a variety of probabilistic models and with different latent variables . When the probability model is a VAE , it only has local latent variables z = ( z1 , . . . , zn ) , and Eq . 2 is equivalent to the common definition of posterior collapse p ( zi |xi ; µ̂ ) = p ( zi ) for all i [ 12 , 17 , 33 , 44 ] . Posterior collapse has also been observed in Gaussian mixture models [ 5 ] ; the posterior of the latent mixture weights resembles their prior when the number of mixture components in the model is larger than that of the data generating process . Regardless of the model , when posterior collapse occurs , it prevents the latent variable from providing meaningful summary of the dataset . Definition 2 ( Latent variable non-identifiability [ 42 , 56 ] ) . Given a likelihood function p ( x | z ; µ ) , a parameter value µ = µ̂ , and a dataset x= ( x1 , . . . , xn ) , the latent variable z is non-identifiable if p ( x | z = z̃0 ; µ̂ ) = p ( x | z = z̃ ; µ̂ ) 8z̃0 , z̃ 2Z , ( 3 ) where Z denotes the domain of z , and z̃0 , z̃ refer to two arbitrary values the latent variable z can take . As a consequence , for any prior p ( z ) on z , we have the conditional likelihood equal to the marginal p ( x | z = z̃ ; µ̂ ) = R p ( x | z ; µ̂ ) p ( z ) dz = p ( x ; µ̂ ) 8z̃ 2Z . Definition 2 says a latent variable z is non-identifiable when the likelihood of the dataset x does not depend on z . It is also known as practical non-identifiability [ 42 , 56 ] and is closely related to the definition of z being conditionally non-identifiable ( or conditionally uninformative ) given µ̂ [ 40 , 42 , 43 , 49 , 58 ] . To enforce latent variable identifiability , it is sufficient to ensure that the likelihood p ( x | z , µ ) is an injective ( a.k.a . one-to-one ) function of z for all µ . If this condition holds then z̃0 6= z̃ ) p ( x | z = z̃0 ; µ̂ ) 6= p ( x | z = z̃ ; µ̂ ) . ( 4 ) Note that latent variable non-identifiability only requires Eq . 3 be true for a given dataset x and parameter value µ̂ . Thus a latent variable may be identifiable in a model given one dataset but not another , and at one µ but not another . See examples in Appendix A . Latent variable identifiability ( Definition 2 ) [ 42 , 56 ] differs from model identifiability [ 41 ] , a related notion that has also been cited as a contributing factor to posterior collapse [ 59 ] . Latent variable identifiability is a weaker requirement : it only requires the latent variable z be identifiable at a particular parameter value µ = µ̂ , while model identifiability requires both z and µ be identifiable . We now establish the equivalence between posterior collapse and latent variable non-identifiability . Theorem 1 ( Latent variable non-identifiability , Posterior collapse ) . Consider a probability model p ( x , z ; µ ) , a dataset x , and a parameter value µ = µ̂ . The local latent variables z are non-identifiable at µ̂ if and only if the posterior of the latent variable z collapses , p ( z |x ) = p ( z ) . Proof . To prove that non-identifiability implies posterior collapse , note that , by Bayes rule , p ( z |x ; µ̂ ) / p ( z ) p ( x | z ; µ̂ ) = p ( z ) p ( x ; µ̂ ) / p ( z ) , ( 5 ) where the middle equality is due to the definition of latent variable non-identifiability . It implies p ( z |x ; µ̂ ) = p ( z ) as both are densities . To prove that posterior collapse implies latent variable nonidentifiability , we again invoke Bayes rule . Posterior collapse implies that p ( z ) = p ( z |x ; µ̂ ) / p ( z ) · p ( x | z ; µ̂ ) , which further implies that p ( x | z ; µ̂ ) is constant in z . If p ( x | z ; µ̂ ) nontrivially depends on z , then p ( z ) must be different from p ( z ) p ( x | z ; µ̂ ) as a function of z . The proof of Theorem 1 is straightforward , but Theorem 1 has an important implication . It shows that the problem of posterior collapse mainly arises from the model and the data , rather than from inference or optimization . If the maximum likelihood parameters µ̂ of the VAE renders the latent variable z non-identifiable , then we will observe posterior collapse . Theorem 1 also clarifies why posteriors may change from non-collapsed to collapsed ( and back ) while fitting a VAE . When fitting a VAE , Some parameter iterates may lead to posterior collapse ; others may not . Theorem 1 points to why existing approaches can help mitigate posterior collapse . Consider the ØVAE [ 19 ] , the VAE lagging encoder [ 18 ] , and the semi-amortized VAE [ 25 ] . Though motivated by other perspectives , these methods modify the optimization objectives or algorithms of VAE to avoid parameter values µ at which the latent variable is non-identifiable . The resulting posterior may not collapse , though the optimal parameters for these algorithms no longer approximates the maximum likelihood estimate . Theorem 1 can also help us understand posterior collapse observed in practice , which manifests as the phenomenon that the posterior is approximately ( as opposed to exactly ) equal to the prior , p ( z |x ; µ̂ ) º p ( z ) . In several empirical studies of VAE ( e.g . [ 12 , 18 , 25 ] ) , we observe that the Kullback-Leibler ( KL ) divergence between the prior and posterior is close to zero but not exactly zero , a property that stems from the likelihood p ( x | z ) being nearly constant in the latents z . In these cases , Theorem 1 provides the intuition that the latent variable is nearly non-identifiable , p ( x | z̃0 ) º p ( x | z̃ ) ,8z̃ , z̃0 and so Eq . 2 holds approximately .

Summary: The paper proposes an explanation for the phenomenon of posterior collapse in variational autoencoders and proves that an exact notion of posterior collapse happens if and only if the latent variables are non-identifiable. Then, it proposes a new variational autoencoder model based on Brenier maps which ensures that the latent variables are identifiable. The new model can simply be “plugged into” existing variational autoencoder models. Experiments demonstrate that the new model mitigates posterior collapse.