Variational Diffusion Autoencoders with Random Walk Sampling

Variational inference ( VI ) methods and especially variational autoencoders ( VAEs ) specify scalable generative models that enjoy an intuitive connection to manifold learning — with many default priors the posterior/likelihood pair q ( z|x ) /p ( x|z ) can be viewed as an approximate homeomorphism ( and its inverse ) between the data manifold and a latent Euclidean space . However , these approximations are well-documented to become degenerate in training . Unless the subjective prior is carefully chosen , the topologies of the prior and data distributions often will not match . Conversely , diffusion maps ( DM ) automatically infer the data topology and enjoy a rigorous connection to manifold learning , but do not scale easily or provide the inverse homeomorphism . In this paper , we propose a ) a principled measure for recognizing the mismatch between data and latent distributions and b ) a method that combines the advantages of variational inference and diffusion maps to learn a homeomorphic generative model . The measure , the locally bi-Lipschitz property , is a sufficient condition for a homeomorphism and easy to compute and interpret . The method , the variational diffusion autoencoder ( VDAE ) , is a novel generative algorithm that first infers the topology of the data distribution , then models a diffusion random walk over the data . To achieve efficient computation in VDAEs , we use stochastic versions of both variational inference and manifold learning optimization . We prove approximation theoretic results for the dimension dependence of VDAEs , and that locally isotropic sampling in the latent space results in a random walk over the reconstructed manifold . Finally , we demonstrate our method on various real and synthetic datasets , and show that it exhibits performance superior to other generative models . 1 INTRODUCTION . Recent developments in generative models such as variational auto-encoders ( VAEs , Kingma & Welling ( 2013 ) ) and generative adversarial networks ( GANs , Goodfellow et al . ( 2014 ) ) have made it possible to sample remarkably realistic points from complex high dimensional distributions at low computational cost . While their methods are very different — one is derived from variational inference and the other from game theory — their ends both involve learning smooth mappings from a user-defined prior distribution to the modeled distribution . These maps are closely tied to manifold learning when the prior is supported over a Euclidean space ( e.g . Gaussian or uniform priors ) and the data lie on a manifold ( also known as the Manifold Hypothesis , see Narayanan & Mitter ( 2010 ) ; Fefferman et al . ( 2016 ) ) . This is because manifolds themselves are defined by sets that have homeomorphisms to such spaces . Learning such maps is beneficial to any machine learning task , and may shed light on the success of VAEs and GANs in modeling complex distributions . Furthermore , the connection to manifold learning may explain why these generative models fail when they do . Known as posterior collapse in VAEs ( Alemi et al. , 2017 ; Zhao et al. , 2017 ; He et al. , 2019 ; Razavi et al. , 2019 ) and mode collapse in GANs ( Goodfellow , 2017 ) , both describe cases where the forward/reverse mapping to/from Euclidean space collapses large parts of the input to a single output . This violates the bijective requirement of the homeomorphic mapping . It also results in degenerate latent spaces and poor generative performance . A major cause of such failings is when the geometries of the prior and target data do not agree . We explore this issue of prior mismatch and previous treatments of it in Section 3 . Given their connection to manifold learning , it is natural to look to classical approaches in the field for ways to improve VAEs . One of the most principled methods is spectral learning ( Schölkopf et al. , 1998 ; Roweis & Saul , 2000 ; Belkin & Niyogi , 2002 ) which involves describing data from a manifold X ⊂ MX by the eigenfunctions of a kernel on MX . We focus specifically on DMs , where Coifman & Lafon ( 2006 ) show that normalizations of the kernel approximate a very specific diffusion process , the heat kernel over MX . A crucial property of the heat kernel is that , like its physical analogue , it defines a diffusion process that has a uniform stationary distribution — in other words , drawing from this stationary distribution draws uniformly from the data manifold . Moreover , Jones et al . ( 2008 ) established another crucial property of DMs , namely that distances in local neighborhoods in the eigenfunction space are nearly isometric to corresponding geodesic distances on the manifold . However , despite its strong theoretical guarantees , DMs are poorly equipped for large scale generative modeling as they are not easily scalable and do not provide an inverse mapping from the intrinsic feature space . In this paper we address issues in variational inference and manifold learning by combining ideas from both . Theory in manifold learning allows us to better recognize prior mismatch , whereas variational inference provides a method to learn the difficult to approximate inverse diffusion map . Our contributions : 1 ) We introduce the locally bi-Lipschitz property , a sufficient condition for a homeomorphism , for measuring the stability of a mapping between latent and data distributions . 2 ) We introduce VDAEs , a class of variational autoencoders whose encoder-decoder feedforward pass approximates the diffusion process on the data manifold with respect to a user-defined kernel k. 3 ) We show that deep neural networks are capable of learning such diffusion processes , and 4 ) that networks approximating this process produce random walks that have certain desirable properties , including well defined transition and stationary distributions . 5 ) Finally , we demonstrate the utility of the VDAE framework on a set of real and synthetic datasets , and show that they have superior performance and satisfy the locally bi-Lipschitz property where GANs and VAEs do not . 2 BACKGROUND . Variational inference ( VI , Jordan et al . ( 1999 ) ; Wainwright et al . ( 2008 ) ) is a machine learning method that combines Bayesian statistics and latent variable models to approximate some probability density p ( x ) . VI assumes and exploits a latent variable structure in the assumed data generation process , that the observations x ∼ p ( x ) are conditionally distributed given unobserved latent vari- ables z . By modeling the conditional distribution , then marginalizing over z , as in pθ ( x ) = ∫ z pθ ( x|z ) p ( z ) dz , ( 1 ) we obtain the model evidence , or likelihood that x could have instead been drawn from pθ ( x ) . Maximizing Eq . 1 leads to an algorithm for finding likely approximations of p ( x ) . As the cost of computing this integral scales exponentially with the dimension of z , we instead maximize the evidence lower bound ( ELBO ) : log pθ ( x ) ≥ −DKL ( q ( z|x ) ||p ( z ) ) + Ez∼q ( z|x ) [ log pθ ( x|z ) ] , ( 2 ) where q ( z|x ) is usually an approximation of pθ ( z|x ) . Optimizing the ELBO is sped up by taking stochastic gradients ( Hoffman et al. , 2013 ) , and further accelerated by learning a global function approximator qφ in an autoencoding structure ( Kingma & Welling , 2013 ) . Diffusion maps ( DMs , Coifman & Lafon ( 2006 ) ) on the other hand , are a class of kernel methods that perform non-linear dimensionality reduction on a set of observations X ⊆ MX , whereMX is the data manifold . Given a symmetric and positive kernel k , DM considers the induced random walk on the graph of X , where given x , y ∈ X , the transition probabilities p ( y|x ) = p ( x , y ) are row normalized versions of k ( x , y ) . Moreover , the diffusion map ψ embeds the data X ∈ Rm into the Euclidean space RD so that the diffusion distance is approximated by Euclidean distance . This is a powerful property , as it allows the arbitrarily complex random walk induced by k on MX to become an isotropic Gaussian random walk on ψ ( MX ) . SpectralNet is an algorithm introduced by algorithm in Shaham et al . ( 2018b ) to speed up the diffusion map . Until recently , the method ψk could only be computed via the eigendecomposition of K. As a result , DMs were only be tractable for small datasets , or on larger datasets by combining landmark-based estimates and Nystrom approximation techniques . However , Shaham et al . ( 2018b ) propose approximations of the function ψ itself in the case that the kernel k is symmetric . In particular , we will leverage SpectralNet to enforce our diffusion embedding prior . Locally bi-lipschitz coordinates by kernel eigenfunctions . ( Jones et al . ( 2008 ) ) analyzed the construction of local coordinates of Riemannian manifolds by Laplacian eigenfunctions and diffusion map coordinates . They establish , for all x ∈ X , the existence of some neighborhood U ( x ) and d spectral coordinates given U ( x ) that define a bi-Lipschitz mapping from U ( x ) to Rd . With a smooth compact Riemannian manifold , U ( x ) can be chosen to be a geodesic ball with radius a constant multiple of the inradius ( the radius of the largest possible ball around x without intersecting with the manifold boundary ) , where the constant is uniform for all x , but the indices of the d spectral coordinates as well as the local bi-Lipschitz constants may depend on x . Specifically , the Lipschitz constants involve inverse of the inradius at x multiplied again by some global constants . For completeness we give a simplified statement of the Jones et al . ( 2008 ) result in the supplementary material . Using the compactness of the manifold , one can always cover the manifold with m many neighborhoods ( geodesic balls ) on which the bi-Lipschitz property in Jones et al . ( 2008 ) holds . As a result , there are a total of D spectral coordinates , D ≤ md ( in practice D is much smaller than md , since the selected spectral coordinates in the proof of Jones et al . ( 2008 ) tend to be low-frequency ones , and thus the selection on different neighborhoods tend to overlap ) , such that on each of them neighborhoods , there exists a subset of d spectral coordinates out of the D ones which are bi-Lipschitz on the neighborhood , and the Lipschitz constants can be bounded uniformly from below and above . 3 MOTIVATION AND RELATED WORK . Our proposed measure and model is motivated by degenerate latent spaces and poor generative performance in a variational inference framework arising from prior mismatch : when the topologies of the data and prior distributions do not agree . In real world data , this is usually due to two factors : first , when the dimensionalities of the distributions do not match , and second , when the geometries do not match . It is easy to see that homeomorphisms between the distributions will not exist in either case : pointwise correspondences can not be established , thus the bijective condition can not be met . As a result , the model has poor generative performance — for each point not captured in the pointwise correspondence , the latent or generated distribution loses expressivity . Though the default choice of Gaussian distribution for p ( z ) is mathematically elegant and computationally expedient , there are many datasets , real and synthetic , for which this distribution is ill-suited . It is well known that spherical distributions are superior for modeling directional data ( Fisher et al. , 1993 ; Mardia , 2014 ) , which can be found in fields as diverse as bioinformatics ( Hamelryck et al. , 2006 ) , geology ( Peel et al. , 2001 ) , material science ( Krieger Lassen et al. , 1994 ) , natural image processing ( Bahlmann , 2006 ) , and simply preprocessed datasets1 . Additionally observe that no homeomorphism exists between Rk and S1 for any k. For data distributed on more complex manifolds , the literature is sparse due to the difficult nature of such study . However , the manifold hypothesis is well-known and studied ( Narayanan & Mitter , 2010 ; Fefferman et al. , 2016 ) . Previous research on alleviating prior mismatch exists . Davidson et al . ( 2018 ) ; Xu & Durrett ( 2018 ) consider VAEs with the von-Mises Fisher prior , a geometrically hyperspherical prior . ( Rey et al. , 2019 ) further model arbitrarily complex manifolds as priors , but require explicit knowledge of the manifold ( i.e . its projection map , scalar curvature , and volume ) . Finally , Tomczak & Welling ( 2017 ) consider mixtures of any pre-existing priors . But while these methods increase the expressivity of the priors available , they do not prescribe a method for choosing the prior itself . That responsibility still lies with the user . Convserly , our method chooses the best prior automatically . To our knowledge , ours is the first to take a data-driven approach to prior selection . By using some data to inform the prior , we not only guarantee the existence of a homeomorphism between data and prior distributions , we explicitly define it by the learned diffusion map ψ̃ .

Summary: The paper studies the problem of density estimation and learning accurate generative models. The authors start from the observation that this problem has been approaches either using variational inference models, that scale very well but whose approximations may lead to degenerate results in practice, and diffusion maps, that scale poorly but are very effective in capturing the underlying data manifold. From here, the authors propose integrating the notion of random walk from diffusion maps into VAEs to avoid degenerate conditions. The proposed method is first defined in its generality, a practical implementation is presented, theoretical guarantees are provided, and empirical evidence of its effectiveness is reported.