VAE Approximation Error: ELBO and Exponential Families

1 INTRODUCTION . Variational autoencoders proposed by Kingma & Welling ( 2014 ) strive at learning complex data distributions pd ( x ) , x ∈ X in a generative way . They introduce latent variables z ∈ Z and model the joint distribution as pθ ( x | z ) p ( z ) , where p ( z ) is a simple distribution that is usually assumed to be known . The conditional distribution pθ ( x | z ) , called decoder , is modeled in terms of a deep network parametrized by θ ∈ Θ . Models defined in this way allow to sample from pθ ( x ) = Ep ( z ) pθ ( x | z ) easily , however at the price that computing the posterior pθ ( z |x ) = pθ ( x | z ) p ( z ) /pθ ( x ) is usually intractable . To handle this problem , VAE approximates the posterior pθ ( z |x ) by an amortized inference encoder model qφ ( z |x ) parametrized by φ ∈ Φ . Given the empirical data distribution pd ( x ) , the model is learned by maximizing the evidence lower bound ( ELBO ) of the data loglikelihood L ( θ ) = Epd log pθ ( x ) . It can be expressed in the following two equivalent forms : LB ( θ , φ ) = Epd [ Eqφ log pθ ( x | z ) −DKL ( qφ ( z |x ) ‖ p ( z ) ) ] ( 1a ) = L ( θ ) − Epd [ DKL ( qφ ( z |x ) ‖ pθ ( z |x ) ) ] . ( 1b ) The first form allows for stochastic optimization of ELBO while the second form shows that the gap between log-likelihood and ELBO is exactly the mismatch between the encoder and the posterior . VAEs constitute a powerful deep learning extension of the expectation maximization approach to handle latent variables . They are useful not only as generative models but also , e.g. , in semisupervised learning ( Kingma et al. , 2014 ; Mattei & Frellsen , 2019 ) . Furthermore the encoder part constructs an efficient embedding of the data in the latent space , useful in many applications . The outreach of the VAE approach requires therefore a careful empirical and theoretical analysis of the problems and trade offs involved . The most important ones are ( i ) posterior collapse ( He et al. , 2019 ; Lucas et al. , 2019 ; Dai et al. , 2018 ; Dai & Wipf , 2019 ) and ( ii ) approximation errors caused by an inappropriate choice of the encoder family . The VAE approximation error has been studied ( e.g. , Cremer et al . 2018 ; Hjelm et al . 2016 ; Kim et al . 2018 ) so far mainly empirically . The problem also occurs and is well-recognized in the context of variational inference and variational Bayesian inference , where the target posterior distribution is expected to be complex . It is commonly understood , that the mean field approximation of pθ ( z |x ) by qφ ( z |x ) in ( 1b ) significantly limits variational Bayesian inference . In contrast , in VAEs , the decoder may adopt to compensate for the chosen encoder family . The effect of this coupling , we believe , is not fully understood . The phenomenon of decoder adopting to the posterior was experimentally observed , e.g. , by Cremer et al . ( 2018 , Section 5.4 ) , noting that the approximation error is often dominated by the amortization error . Turner & Sahani ( 2011 , Sec . 1.4 ) analytically show for linear state space models that simpler variational approximations ( such a mean-field ) can lead to less bias in parameter estimation than more complicated structured approximations . Similarly , Shu et al . ( 2018 ) view the VAE objective as providing a regularization and show that making the amortized inference model smoother , while increasing the amortization gap , leads to a better generalization . The common ( empirical ) understanding of the importance of the gap between the approximate and the true posterior has led to many generalizations of standard VAEs , which achieve impressive practical results , notably , tighter bounds using importance weighting ( Burda et al. , 2016 ; Nowozin , 2018 ) , encoders employing normalizing flows ( Rezende & Mohamed , 2015 ; Kingma et al. , 2016 ) , hierarchical and autoregressive encoders ( Vahdat & Kautz , 2020 ; Sønderby et al. , 2016 ; Ranganath et al. , 2016 ) , MRF encoders ( Vahdat et al. , 2020 ) and more . While these extensions mitigate the posterior mismatch problem , they often come at a price of a more difficult training and more expensive inference . Furthermore , simpler encoders may be of practical interest . Burda et al . ( 2016 , Appendix C ) illustrates that IWAE approximate posteriors are less regular and more spread out . In contrast , factorized encoders provide simple embeddings useful for downstream tasks such as semantic hashing ( Chaidaroon & Fang , 2017 ) . The aim of this paper is to study the approximation error of VAEs and its impact on the learned decoder . We consider a setting that generalizes many common VAEs , in particular popular models where encoder and decoder are conditionally independent Bernoulli or Gaussian distributions : we assume that both decoder and encoder are conditional exponential families . We identify the subclass of generative models where the encoder can model the posterior exactly , referred to as consistent VAEs . We give a characterization of consistent VAEs revealing that this set in fact does not depend on the complexity of the involved neural networks . We further show that the ELBO optimizer is pulled towards this set away from the likelihood optimizer . Specializing the characterization to several common VAE models , we show that the respective consistent models turn out to be RBMlike in many cases . We experimentally investigate the detrimental effect in one case and show that a simpler but more consistent VAE can perform better in the other . 2 PROBLEM STATEMENT . We adopt the following notion of approximation error . Consider a generative model class PΘ = { pθ ( x , z ) | θ∈Θ } , the encoder class QΦ = { qφ ( z |x ) | φ∈Φ } and the data distribution pd ( x ) . The maximum likelihood generative model is given by θML ∈ argmaxθ∈Θ Epd ( x ) log pθ ( x ) . For a decoder with parameters θ we define its approximation error as the likelihood difference L ( θML ) − L ( θ ) . Respectively , the VAE approximation error is defined for a given θ as : L ( θML ) −maxφ LB ( θ , φ ) ≥ L ( θML ) − L ( θ ) . ( 2 ) In order for this error to become zero , two conditions are necessary and sufficient : • Parameters ( θ , φ ) must be optimal for the ELBO objective . • ELBO must be tight at ( θ , φ ) , i.e. , LB ( θ , φ ) = L ( θ ) . Assuming that the optimality can be achieved , we study the non-tightness gap L ( θ ) − LB ( θ , φ ) . From ( 1b ) it expresses as Epd [ DKL ( qφ ( z |x ) ‖ pθ ( z |x ) ) ] . It follows that ELBO is tight at ( θ , φ ) iff qφ ( z |x ) ≡ pθ ( z |x ) . Hence , we define the consistent set ΘΦ ⊆ Θ as the subset of distributions pθ ( x , z ) whose posteriors are in QΦ , i.e. , ΘΦ = { θ ∈ Θ ∣∣ ∃φ ∈ Φ : qφ ( z |x ) ≡ pθ ( z |x ) } . ( 3 ) The KL-divergence in the ELBO objective ( 1b ) can vanish only if θ ∈ ΘΦ . If the likelihood maximizer θML is not contained in ΘΦ , then this KL-divergence pulls the optimizer towards ΘΦ and away from θML as illustrated in Fig . 1 . We characterize the consistent set ΘΦ , on which the bound is tight , and show that this set is quite narrow and does not depend on the complexity of the encoder and decoder networks beyond simple 1-layer linear mappings of sufficient statistics . 3 THEORETICAL ANALYSIS . We consider a general class of VAEs , where both encoder and decoder are defined as exponential families . This class includes many common models , in particular Gaussian VAEs and Bernoulli VAEs with conditional independence assumptions , but also more complex ones , e.g. , where the encoder is a conditional random field ( Vahdat et al. , 2020 ) 1 . Assumption 1 ( Exponential family VAE ) . LetX andZ be sets of observations and latent variables , respectively . We consider VAE models defined by pθ ( x | z ) = h ( x ) exp [ 〈ν ( x ) , fθ ( z ) 〉 −A ( fθ ( z ) ) ] ( 4a ) qφ ( z |x ) = h′ ( z ) exp [ 〈ψ ( z ) , gφ ( x ) 〉 −B ( gφ ( x ) ) ] , ( 4b ) where ν : X → Rn and ψ : Z → Rm are fixed sufficient statistics of dimensionality n and m ; fθ : Z → Rn and gφ : X → Rm are the the decoder , resp. , encoder networks , depending on learnable parameters θ , resp . φ ; h : X → R+ , h′ : Z → R+ are strictly positive base measures and A , B denote the respective log-partition functions . Notice that this assumption imposes no restrictions on the nature of random variables x and z . They can be discrete or continuous , univariate or multivariate . Similarly , it imposes no restrictions on the complexity of the decoder and encoder networks fθ ( z ) and gφ ( x ) . Characterization of the consistent set . In the first step of our analysis , we investigate the conditions under which the approximation error of an exponential family VAE can be made exactly zero . As discussed above , a tight VAE ( θ , φ ) must satisfy ∀ ( x , z ) qφ ( z |x ) = pθ ( z |x ) , which leads to the following theorem . Theorem 1 . The consistent set ΘΦ of an exponential family VAE is given by decoders of the form p ( x | z ) = h ( x ) exp [ 〈ν ( x ) , Wψ ( z ) 〉+ 〈ν ( x ) , u〉 −A ( z ) ] , ( 5 ) where W is a n×m matrix and u ∈ Rn . Moreover , the corresponding encoders have the form q ( z |x ) = h′ ( z ) exp [ 〈 ψ ( z ) , WT ν ( x ) 〉 + 〈ψ ( z ) , v〉 −B ( x ) ] , ( 6 ) where v ∈ Rm . This is a direct consequence of a theorem by Arnold & Strauss ( 1991 ) ( see Appendix A.1 for more details ) . For a tight VAE , Theorem 1 states that the decoder and encoder take the form ( 5 ) and ( 6 ) with the interaction between x and z parametrized by a single matrix W instead of the complex neural networks with parameters θ , φ . Thus they turn out to be generalized linear models ( GLMs ) . The corresponding joint probability distribution is an EF Harmonium ( Welling et al. , 2005 ) : p ( x , z ) = h ( x ) h′ ( z ) exp ( 〈ν ( x ) , Wψ ( z ) 〉+ 〈ν ( x ) , u〉+ 〈ψ ( z ) , v〉 −A ) . ( 7 ) 1Notice , however , that this class does not include VAEs with advanced encoder families like normalizing flows , hierarchical and autoregressive encoders . Corollary 1 . The subset ΘΦ of consistent models can not be enlarged by considering more complex encoder networks g ( x ) , provided that the affine family WTν ( x ) can already be represented . Corollary 2 . Let the decoder network family be affine in ψ ( z ) , i.e. , f ( z ) = Wψ ( z ) + a and let the encoder network family g ( x ) include at least all affine maps V ν ( x ) + b . Then any global optimum of ELBO attains a zero approximation error . VAEs can escape consistency when they degenerate to an invertible flow . In practice , VAE models are almost never tight . It is therefore natural to ask , whether a small VAE posterior mismatch error implies closeness of the optimal decoder to some decoder in the consistent set . Definition 1 . A VAE ( pθ , qφ ) is ε-tight for some ε > 0 if Epd ( x ) [ DKL ( qφ ( z |x ) ‖ pθ ( z |x ) ) ] ≤ ε . It turns out that the the KL divergence has a leak , which allows a VAE to approach tightness while not approaching consistency . In the continuous case an example satisfying ε-tightness with non-linear decoder follows from Dai & Wipf ( 2019 , Theorem 2 ) . They show for a class of Gaussian VAEs with general neural networks fθ , gφ , that it is possible to build a sequence of network parameters θt , φt with the following properties : i ) the target distribution is approximated arbitrary well , ii ) the posterior mismatch DKL ( qφt ( z |x ) ‖ pθt ( z |x ) ) approaches zero and iii ) both the encoder and decoder approach deterministic mappings . The VAE thus approaches a flow model ( or invertible neural network ) between the data manifold and a subspace of the latent space ( Dai & Wipf , 2019 ) . Clearly , in a general case the flow must be non-linear . A similar case can be made for discrete variables , see Example A.1 . Non-deterministic nearly-tight VAEs approach consistency . We would however argue that the mode where the decoder and encoder are nearly-deterministic is not a natural VAE solution . By making additional assumptions , excluding such deterministic solutions , and restricting ourselves to the finite space in order to simplify the analysis , we can show that the decoder of an ε-tight VAE does indeed approach a GLM . Theorem 2 . Let ( pθ , qφ ) be an exponential family VAE ( Assumption 1 ) on a discrete space X × Z with strictly positive encoder qφ ( z |x ) and decoder posterior pθ ( z |x ) , both bounded from below by α > 0 . If the VAE is ε-tight , then there exists W ∈ Rn , m such that the decoder can be approximated by a GLM of the form p̃ ( x | z ) = h̃ ( x ) exp [ 〈ν ( x ) , Wψ ( z ) 〉 − Ã ( z ) ] with the error Epd ( x ) [ ( log pθ ( x | z ) − log p̃ ( x | z ) ) 2 ] ≤ ε2α2 + o ( ε ) ∀z ∈ Z . ( 8 ) The proof is given in Appendix A.3 .

Summary: Summary. This paper presents an analysis of the approximation error of VAE models when the encoder and decoder are from exponential families. They show that when the model is consistent (i.e., the encoder is able to match the posterior), the encoder and decoder distributions are exponential family distributions that are linearly parameterized. The paper also shows that barring pathological cases, when the model is tight the decoder distribution is close to a linearly parameterized exponential family distribution. In particular, this shows that in this case additional depth is not useful.  Despite the restriction to exponential family distributions, the paper shows a number of real instances of models that lie in this class including Gaussian VAEs and VAEs with RBM encoders. For Bernoulli VAEs for semantic hashing, the paper proposes and verifies improvements to a model presented in the literature.