Diagnosing and Enhancing VAE Models

1 INTRODUCTION . Our starting point is the desire to learn a probabilistic generative model of observable variables x ∈ χ , where χ is a r-dimensional manifold embedded in Rd . Note that if r = d , then this assumption places no restriction on the distribution of x ∈ Rd whatsoever ; however , the added formalism is introduced to handle the frequently encountered case where x possesses low-dimensional structure relative to a high-dimensional ambient space , i.e. , r d. In fact , the very utility of generative models of continuous data , and their attendant low-dimensional representations , often hinges on this assumption ( Bengio et al. , 2013 ) . It therefore behooves us to explicitly account for this situation . Beyond this , we assume that χ is a simple Riemannian manifold , which means there exists a diffeomorphism ϕ between χ and Rr , or more explicitly , the mapping ϕ : χ 7→ Rr is invertible and differentiable . Denote a ground-truth probability measure on χ as µgt such that the probability mass of an infinitesimal dx on the manifold is µgt ( dx ) and ∫ χ µgt ( dx ) = 1 . The variational autoencoder ( VAE ) ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) attempts to approximate this ground-truth measure using a parameterized density pθ ( x ) defined across all of Rd since any underlying generative manifold is unknown in advance . This density is further assumed to admit the latent decomposition pθ ( x ) = ∫ pθ ( x|z ) p ( z ) dz , where z ∈ Rκ serves as a lowdimensional representation , with κ ≈ r and prior p ( z ) = N ( z|0 , I ) . Ideally we might like to minimize the negative log-likelihood − log pθ ( x ) averaged across the ground-truth measure µgt , i.e. , solve minθ ∫ χ− log pθ ( x ) µgt ( dx ) . Unfortunately though , the required marginalization over z is generally infeasible . Instead the VAE model relies on tractable encoder qφ ( z|x ) and decoder pθ ( x|z ) distributions , where φ represents additional trainable parameters . The canonical VAE cost is a bound on the average negative log-likelihood given by L ( θ , φ ) , ∫ χ { − log pθ ( x ) + KL [ qφ ( z|x ) ||pθ ( z|x ) ] } µgt ( dx ) ≥ ∫ χ− log pθ ( x ) µgt ( dx ) , ( 1 ) where the inequality follows directly from the non-negativity of the KL-divergence . Here φ can be viewed as tuning the tightness of bound , while θ dictates the actual estimation of µgt . Using a few standard manipulations , this bound can also be expressed as L ( θ , φ ) = ∫ χ { −Eqφ ( z|x ) [ log pθ ( x|z ) ] + KL [ qφ ( z|x ) ||p ( z ) ] } µgt ( dx ) , ( 2 ) which explicitly involves the encoder/decoder distributions and is conveniently amenable to SGD optimization of { θ , φ } via a reparameterization trick ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) . The first term in ( 2 ) can be viewed as a reconstruction cost ( or a stochastic analog of a traditional autoencoder ) , while the second penalizes posterior deviations from the prior p ( z ) . Additionally , for any realizable implementation via SGD , the integration over χ must be approximated via a finite sum across training samples { x ( i ) } ni=1 drawn from µgt . Nonetheless , examining the true objective L ( θ , φ ) can lead to important , practically-relevant insights . At least in principle , qφ ( z|x ) and pθ ( x|z ) can be arbitrary distributions , in which case we could simply enforce qφ ( z|x ) = pθ ( z|x ) ∝ pθ ( x|z ) p ( z ) such that the bound from ( 1 ) is tight . Unfortunately though , this is essentially always an intractable undertaking . Consequently , largely to facilitate practical implementation , a commonly adopted distributional assumption for continuous data is that both qφ ( z|x ) and pθ ( x|z ) are Gaussian . This design choice has previously been cited as a key limitation of VAEs ( Burda et al. , 2015 ; Kingma et al. , 2016 ) , and existing quantitative tests of generative modeling quality thus far dramatically favor contemporary alternatives such as generative adversarial networks ( GAN ) ( Goodfellow et al. , 2014 ) . Regardless , because the VAE possesses certain desirable properties relative to GAN models ( e.g. , stable training ( Tolstikhin et al. , 2018 ) , interpretable encoder/inference network ( Brock et al. , 2016 ) , outlier-robustness ( Dai et al. , 2018 ) , etc . ) , it remains a highly influential paradigm worthy of examination and enhancement . In Section 2 we closely investigate the implications of VAE Gaussian assumptions leading to a number of interesting diagnostic conclusions . In particular , we differentiate the situation where r = d , in which case we prove that recovering the ground-truth distribution is actually possible iff the VAE global optimum is reached , and r < d , in which case the VAE global optimum can be reached by solutions that reflect the ground-truth distribution almost everywhere , but not necessarily uniquely so . In other words , there could exist alternative solutions that both reach the global optimum and yet do not assign the same probability measure as µgt . Section 3 then further probes this non-uniqueness issue by inspecting necessary conditions of global optima when r < d. This analysis reveals that an optimal VAE parameterization will provide an encoder/decoder pair capable of perfectly reconstructing all x ∈ χ using any z drawn from qφ ( z|x ) . Moreover , we demonstrate that the VAE accomplishes this using a degenerate latent code whereby only r dimensions are effectively active . Collectively , these results indicate that the VAE global optimum can in fact uniquely learn a mapping to the correct ground-truth manifold when r < d , but not necessarily the correct probability measure within this manifold , a critical distinction . Next we leverage these analytical results in Section 4 to motivate an almost trivially-simple , twostage VAE enhancement for addressing typical regimes when r < d. In brief , the first stage just learns the manifold per the allowances from Section 3 , and in doing so , provides a mapping to a lower dimensional intermediate representation with no degenerate dimensions that mirrors the r = d regime . The second ( much smaller ) stage then only needs to learn the correct probability measure on this intermediate representation , which is possible per the analysis from Section 2 . Experiments from Sections 5 and 6 empirically corroborate motivational theory and reveal that the proposed two-stage procedure can generate high-quality samples , reducing the blurriness often attributed to VAE models in the past ( Dosovitskiy & Brox , 2016 ; Larsen et al. , 2015 ) . And to the best of our knowledge , this is the first demonstration of a VAE pipeline that can produce stable FID scores , an influential recent metric for evaluating generated sample quality ( Heusel et al. , 2017 ) , that are comparable to GAN models under neutral testing conditions . Moreover , this is accomplished without additional penalties , cost function modifications , or sensitive tuning parameters . Finally , an extended version of this work can be found in ( Dai & Wipf , 2019 ) . There we include additional results , consideration of disentangled representations , as well as a comparative discussion of broader VAE modeling paradigms such as those involving normalizing flows or parameterized families for p ( z ) . 2 HIGH-LEVEL IMPACT OF VAE GAUSSIAN ASSUMPTIONS . Conventional wisdom suggests that VAE Gaussian assumptions will introduce a gap between L ( θ , φ ) and the ideal negative log-likelihood ∫ χ− log pθ ( x ) µgt ( dx ) , compromising efforts to learn the ground-truth measure . However , we will now argue that this pessimism is in some sense premature . In fact , we will demonstrate that , even with the stated Gaussian distributions , there exist parameters φ and θ that can simultaneously : ( i ) Globally optimize the VAE objective and , ( ii ) Recover the ground-truth probability measure in a certain sense described below . This is possible because , at least for some coordinated values of φ and θ , qφ ( z|x ) and pθ ( z|x ) can indeed become arbitrarily close . Before presenting the details , we first formalize a κ-simple VAE , which is merely a VAE model with explicit Gaussian assumptions and parameterizations : Definition 1 A κ-simple VAE is defined as a VAE model with dim [ z ] = κ latent dimensions , the Gaussian encoder qφ ( z|x ) = N ( z|µz , Σz ) , and the Gaussian decoder pθ ( x|z ) = N ( x|µx , Σx ) . Moreover , the encoder moments are defined as µz = fµz ( x ; φ ) and Σz = SzS > z with Sz = fSz ( x ; φ ) . Likewise , the decoder moments are µx = fµx ( z ; θ ) and Σx = γI . Here γ > 0 is a tunable scalar , while fµz , fSz and fµx specify parameterized differentiable functional forms that can be arbitrarily complex , e.g. , a deep neural network . Equipped with these definitions , we will now demonstrate that a κ-simple VAE , with κ ≥ r , can achieve the optimality criteria ( i ) and ( ii ) from above . In doing so , we first consider the simpler case where r = d , followed by the extended scenario with r < d. The distinction between these two cases turns out to be significant , with practical implications to be explored in Section 4 . 2.1 MANIFOLD DIMENSION EQUAL TO AMBIENT SPACE DIMENSION ( r = d ) We first analyze the specialized situation where r = d. Assuming pgt ( x ) , µgt ( dx ) /dx exists everywhere in Rd , then pgt ( x ) represents the ground-truth probability density with respect to the standard Lebesgue measure in Euclidean space . Given these considerations , the minimal possible value of ( 1 ) will necessarily occur if KL [ qφ ( z|x ) ||pθ ( z|x ) ] = 0 and pθ ( x ) = pgt ( x ) almost everywhere . ( 3 ) This follows because by VAE design it must be that L ( θ , φ ) ≥ − ∫ pgt ( x ) log pgt ( x ) dx , and in the present context , this lower bound is achievable iff the conditions from ( 3 ) hold . Collectively , this implies that the approximate posterior produced by the encoder qφ ( z|x ) is in fact perfectly matched to the actual posterior pθ ( z|x ) , while the corresponding marginalized data distribution pθ ( x ) is perfectly matched the ground-truth density pgt ( x ) as desired . Perhaps surprisingly , a κ-simple VAE can actually achieve such a solution : Theorem 1 Suppose that r = d and there exists a density pgt ( x ) associated with the ground-truth measure µgt that is nonzero everywhere on Rd.1 . Then for any κ ≥ r , there is a sequence of κ-simple VAE model parameters { θ∗t , φ∗t } such that lim t→∞ KL [ qφ∗t ( z|x ) ||pθ∗t ( z|x ) ] = 0 and lim t→∞ pθ∗t ( x ) = pgt ( x ) almost everywhere . ( 4 ) All the proofs can be found in ( Dai & Wipf , 2019 ) . So at least when r = d , the VAE Gaussian assumptions need not actually prevent the optimal ground-truth probability measure from being recovered , as long as the latent dimension is sufficiently large ( i.e. , κ ≥ r ) . And contrary to popular notions , a richer class of distributions is not required to achieve this . Of course Theorem 1 only applies to a restricted case that excludes d > r ; however , later we will demonstrate that a key consequence of this result can nonetheless be leveraged to dramatically enhance VAE performance . 2.2 MANIFOLD DIMENSION LESS THAN AMBIENT SPACE DIMENSION ( r < d ) When r < d , additional subtleties are introduced that will be unpacked both here and in the sequel . To begin , if both qφ ( z|x ) and pθ ( x|z ) are arbitrary/unconstrained ( i.e. , not necessarily Gaussian ) , then infφ , θ L ( θ , φ ) = −∞ . To achieve this global optimum , we need only choose φ such that qφ ( z|x ) = pθ ( z|x ) ( minimizing the KL term from ( 1 ) ) while selecting θ such that all probability mass collapses to the correct manifold χ . In this scenario the density pθ ( x ) will become unbounded on χ and zero elsewhere , such that ∫ χ− log pθ ( x ) µgt ( dx ) will approach negative infinity . But of course the stated Gaussian assumptions from the κ-simple VAE model could ostensibly prevent this from occurring by causing the KL term to blow up , counteracting the negative loglikelihood factor . We will now analyze this case to demonstrate that this need not happen . Before 1This nonzero assumption can be replaced with a much looser condition . Specifically , if there exists a diffeomorphism between the set { x|pgt ( x ) 6= 0 } and Rd , then it can be shown that Theorem 1 still holds even if pgt ( x ) = 0 for some x ∈ Rd . proceeding to this result , we first define a manifold density p̃gt ( x ) as the probability density ( assuming it exists ) of µgt with respect to the volume measure of the manifold χ . If d = r then this volume measure reduces to the standard Lebesgue measure in Rd and p̃gt ( x ) = pgt ( x ) ; however , when d > r a density pgt ( x ) defined in Rd will not technically exist , while p̃gt ( x ) is still perfectly well-defined . We then have the following : Theorem 2 Assume r < d and that there exists a manifold density p̃gt ( x ) associated with the ground-truth measure µgt that is nonzero everywhere on χ . Then for any κ ≥ r , there is a sequence of κ-simple VAE model parameters { θ∗t , φ∗t } such that ( i ) lim t→∞ KL [ qφ∗t ( z|x ) ||pθ∗t ( z|x ) ] = 0 and lim t→∞ ∫ χ− log pθ∗t ( x ) µgt ( dx ) = −∞ , ( 5 ) ( ii ) lim t→∞ ∫ x∈A pθ∗t ( x ) dx = µgt ( A ∩ χ ) ( 6 ) for all measurable sets A ⊆ Rd with µgt ( ∂A ∩ χ ) = 0 , where ∂A is the boundary of A . Technical details notwithstanding , Theorem 2 admits a very intuitive interpretation . First , ( 5 ) directly implies that the VAE Gaussian assumptions do not prevent minimization of L ( θ , φ ) from converging to minus infinity , which can be trivially viewed as a globally optimum solution . Furthermore , based on ( 6 ) , this solution can be achieved with a limiting density estimate that will assign a probability mass to most all measurable subsets of Rd that is indistinguishable from the groundtruth measure ( which confines all mass to χ ) . Hence this solution is more-or-less an arbitrarily-good approximation to µgt for all practical purposes.2 Regardless , there is an absolutely crucial distinction between Theorem 2 and the simpler case quantified by Theorem 1 . Although both describe conditions whereby the κ-simple VAE can achieve the minimal possible objective , in the r = d case achieving the lower bound ( whether the specific parameterization for doing so is unique or not ) necessitates that the ground-truth probability measure has been recovered almost everywhere . But the r < d situation is quite different because we have not ruled out the possibility that a different set of parameters { θ , φ } could push L ( θ , φ ) to −∞ and yet not achieve ( 6 ) . In other words , the VAE could reach the lower bound but fail to closely approximate µgt . And we stress that this uniqueness issue is not a consequence of the VAE Gaussian assumptions per se ; even if qφ ( z|x ) were unconstrained the same lack of uniqueness can persist . Rather , the intrinsic difficulty is that , because the VAE model does not have access to the groundtruth low-dimensional manifold , it must implicitly rely on a density pθ ( x ) defined across all of Rd as mentioned previously . Moreover , if this density converges towards infinity on the manifold during training without increasing the KL term at the same rate , the VAE cost can be unbounded from below , even in cases where ( 6 ) is not satisfied , meaning incorrect assignment of probability mass . To conclude , the key take-home message from this section is that , at least in principle , VAE Gaussian assumptions need not actually be the root cause of any failure to recover ground-truth distributions . Instead we expose a structural deficiency that lies elsewhere , namely , the non-uniqueness of solutions that can optimize the VAE objective without necessarily learning a close approximation to µgt . But to probe this issue further and motivate possible workarounds , it is critical to further disambiguate these optimal solutions and their relationship with ground-truth manifolds . This will be the task of Section 3 , where we will explicitly differentiate the problem of locating the correct groundtruth manifold , from the task of learning the correct probability measure within the manifold . Note that the only comparable prior work we are aware of related to the results in this section comes from Doersch ( 2016 ) , where the implications of adopting Gaussian encoder/decoder pairs in the specialized case of r = d = 1 are briefly considered . Moreover , the analysis there requires additional much stronger assumptions than ours , namely , that pgt ( x ) should be nonzero and infinitely differentiable everywhere in the requisite 1D ambient space . These requirements of course exclude essentially all practical usage regimes where d = r > 1 or d > r , or when ground-truth densities are not sufficiently smooth . 2Note that ( 6 ) is only framed in this technical way to accommodate the difficulty of comparing a measure µgt restricted to χ with the VAE density pθ ( x ) defined everywhere in Rd . See ( Dai & Wipf , 2019 ) for details .

Summary: The paper provides a number of novel interesting theoretical results on "vanilla" Gaussian Variational Auto-Encoders (VAEs) (sections 1, 2, and 3), which are then used to build a new algorithm called "2 stage VAEs" (Section 4). The resulting algorithm is as stable as VAEs to train (it is free of any sort of adversarial training, it comes with a little overhead in terms of extra parameters), while achieving a quality of samples which is *very impressive* for an Auto-Encoder (AE) based generative modeling techniques (Section 5). In particular, the method achieves FID score 24 on the CelebA dataset which is on par with the best GAN-based models as reported in [1], thus sufficiently reducing the gap between the generative quality of the GAN-based and AE-based models reported in the literature.