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 proposes a new generative model for unsupervised learning, based on a diffusion random walk principle inspired by the manifold learning literature. The basic idea is to (probabilistically) map points to a latent space, perform a random walk in that space, and then map back to the original space again. Learning of the suitable maps is achieved by casting the problem in a variational inference framework.