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: The paper studies the VAE approximation error, when the encoder and decoder distributions are from conditional exponential families (EFs). Theorem 1 characterizes the form of the joint probability (an EF-Harmonium) which is consistent, i.e. where q_{phi}(z|x) \equiv p_{theta}(z|x). Sec 4.2 shows a v interesting example of Gaussian VAEs for CelebA images, where it is shown that a VAE "unlearns" the ground truth solution (being pulled away from the likelihood optimizer towards the consistent subset). Sec 4.3 shows another interesting example, where a careful formulation of the model identifies that it should use word counts rather than frequencies, and this is demonstrated empirically (Table 2).