Mixed-curvature Variational Autoencoders

Euclidean geometry has historically been the typical “ workhorse ” for machine learning applications due to its power and simplicity . However , it has recently been shown that geometric spaces with constant non-zero curvature improve representations and performance on a variety of data types and downstream tasks . Consequently , generative models like Variational Autoencoders ( VAEs ) have been successfully generalized to elliptical and hyperbolic latent spaces . While these approaches work well on data with particular kinds of biases e.g . tree-like data for a hyperbolic VAE , there exists no generic approach unifying and leveraging all three models . We develop a Mixed-curvature Variational Autoencoder , an efficient way to train a VAE whose latent space is a product of constant curvature Riemannian manifolds , where the per-component curvature is fixed or learnable . This generalizes the Euclidean VAE to curved latent spaces and recovers it when curvatures of all latent space components go to 0 . 1 INTRODUCTION . Generative models , a growing area of unsupervised learning , aim to model the data distribution p ( x ) over data points x from a space X , which is usually a high-dimensional Euclidean space Rn . This has desirable benefits like a naturally definable inner-product , vector addition , or a closedform distance function . Yet , many types of data have a strongly non-Euclidean latent structure ( Bronstein et al. , 2017 ) , e.g . the set of human-interpretable images . They are usually thought to live on a “ natural image manifold ” ( Zhu et al. , 2016 ) , a continuous lower-dimensional subset of the space in which they are represented . By moving along the manifold , one can continuously change the content and appearance of interpretable images . As noted in Nickel & Kiela ( 2017 ) , changing the geometry of the underlying latent space enables better representations of specific data types compared to Euclidean spaces of any dimensions , e.g . tree structures and scale-free networks . Motivated by these observations , a range of recent methods learn representations in different spaces of constant curvatures : spherical or elliptical ( Batmanghelich et al. , 2016 ) , hyperbolic ( Nickel & Kiela , 2017 ; Sala et al. , 2018 ; Tifrea et al. , 2019 ) and even in products of these spaces ( Gu et al. , 2019 ; Bachmann et al. , 2020 ) . Using a combination of different constant curvature spaces , Gu et al . ( 2019 ) aim to match the underlying geometry of the data better . However , an open question remains : how to choose the dimensionality and curvatures of each of the partial spaces ? A popular approach to generative modeling is the Variational Autoencoder ( Kingma & Welling , 2014 ) . VAEs provide a way to sidestep the intractability of marginalizing a joint probability model of the input and latent space p ( x , z ) , while allowing for a prior p ( z ) on the latent space . Recently , variants of the VAE have been introduced for spherical ( Davidson et al. , 2018 ; Xu & Durrett , 2018 ) and hyperbolic ( Mathieu et al. , 2019 ; Nagano et al. , 2019 ) latent spaces . Our approach , the Mixed-curvature Variational Autoencoder , is a generalization of VAEs to products of constant curvature spaces.1 It has the advantage of a better reduction in dimensionality , while maintaining efficient optimization . The resulting latent space is a “ non-constantly ” curved manifold that is more flexible than a single constant curvature manifold . 1Code is available on GitHub at https : //github.com/oskopek/mvae . Our contributions are the following : ( i ) we develop a principled framework for manipulating representations and modeling probability distributions in products of constant curvature spaces that smoothly transitions across curvatures of different signs , ( ii ) we generalize Variational Autoencoders to learn latent representations on products of constant curvature spaces with generalized Gaussianlike priors , and ( iii ) empirically , our models outperform current benchmarks on a synthetic tree dataset ( Mathieu et al. , 2019 ) and on image reconstruction on the MNIST ( LeCun , 1998 ) , Omniglot ( Lake et al. , 2015 ) , and CIFAR ( Krizhevsky , 2009 ) datasets for some latent space dimensions . 2 GEOMETRY AND PROBABILITY IN RIEMANNIAN MANIFOLDS . To define constantly curved spaces , we first need to define the notion of sectional curvatureK ( τx ) of two linearly independent vectors in the tangent space at a point x ∈M spanning a two-dimensional plane τx ( Berger , 2012 ) . Since we deal with constant curvature spaces where all the sectional curvatures are equal , we denote a manifold ’ s curvature as K. Instead of curvature K , we sometimes use the generalized notion of a radius : R = 1/ √ |K| . There are three different types of manifoldsM we can define with respect to the sign of the curvature : a positively curved space , a “ flat ” space , and a negatively curved space . Common realizations of those manifolds are the hypersphere SK , the Euclidean space E , and the hyperboloid HK : M = SnK = { x ∈ Rn+1 : 〈x , x〉2 = 1/K } , for K > 0 En = Rn , for K = 0 HnK = { x ∈ Rn+1 : 〈x , x〉L = 1/K } , for K < 0 where 〈· , ·〉2 is the standard Euclidean inner product , and 〈· , ·〉L is the Lorentz inner product , 〈x , y〉L = −x1y1 + n+1∑ i=2 xiyi ∀x , y ∈ Rn+1 . We will need to define the exponential map , logarithmic map , and parallel transport in all spaces we consider . The exponential map in Euclidean space is defined as expx ( v ) = x + v , for all x ∈ En and v ∈ TxEn . Its inverse , the logarithmic map is logx ( y ) = y − x , for all x , y ∈ En . Parallel transport in Euclidean space is simply an identity PTx→y ( v ) = v , for all x , y ∈ En and v ∈ TxEn . An overview of these operations in the hyperboloid HnK and the hypersphere SnK can be found in Table 1 . For more details , refer to Petersen et al . ( 2006 ) , Cannon et al . ( 1997 ) , or Appendix A . 2.1 STEREOGRAPHICALLY PROJECTED SPACES . The above spaces are enough to cover any possible value of the curvature , and they define all the necessary operations we will need to train VAEs in them . However , both the hypersphere and the hyperboloid have an unsuitable property , namely the non-convergence of the norm of points as the curvature goes to 0 . Both spaces grow as K → 0 and become locally “ flatter ” , but to do that , their points have to go away from the origin of the coordinate space 0 to be able to satisfy the manifold ’ s definition . A good example of a point that diverges is the origin of the hyperboloid ( or a pole of the hypersphere ) µ0 = ( 1/ √ |K| , 0 , . . . , 0 ) T . In general , we can see that ‖µ0‖2 = 1/K K→0−−−→ ±∞ . Additionally , the distance and the metric tensors of these spaces do not converge to their Euclidean variants as K → 0 , hence the spaces themselves do not converge to Rd . This makes both of these spaces unsuitable for trying to learn sign-agnostic curvatures . Luckily , there exist well-defined non-Euclidean spaces that inherit most properties from the hyperboloid and the hypersphere , yet do not have these properties – namely , the Poincaré ball and the projected sphere , respectively . We obtain them by applying stereographic conformal projections , meaning that angles are preserved by this transformation . Since the distance function on the hyperboloid and hypersphere only depend on the radius and angles between points , they are isometric . We first need to define the projection function ρK . For a point ( ξ ; xT ) T ∈ Rn+1 and curvature K ∈ R , where ξ ∈ R , x , y ∈ Rn ρK ( ( ξ ; x T ) T ) = x 1 + √ |K|ξ , ρ−1K ( y ) = ( 1√ |K| 1−K ‖y‖22 1 +K ‖y‖22 ; 2yT 1 +K ‖y‖22 ) T . The formulas correspond to the classical stereographic projections defined for these models ( Lee , 1997 ) . Note that both of these projections map the point µ0 = ( 1/ √ |K| , 0 , . . . , 0 ) in the original space to µ0 = 0 in the projected space , and back . Since the stereographic projection is conformal , the metric tensors of both spaces will be conformal . In this case , the metric tensors of both spaces are the same , except for the sign of K : gDKx = g PK x = ( λKx ) 2gE , for all x in the respective manifold ( Ganea et al. , 2018a ) , and gEy = I for all y ∈ E. The conformal factor λKx is then defined as λ K x = 2/ ( 1 + K ‖x‖ 2 2 ) . Among other things , this form of the metric tensor has the consequence that we unfortunately can not define a single unified inner product in all tangent spaces at all points . The inner product at x ∈ M has the form of 〈u , v〉x = ( λKx ) 2 〈u , v〉2 for all u , v ∈ TxM . We can now define the two models corresponding to K > 0 and K < 0 . The curvature of the projected manifold is the same as the original manifold . An n-dimensional projected hypersphere ( K > 0 ) is defined as the set DnK = ρK ( SnK \ { −µ0 } ) = Rn , where µ0 = ( 1/ √ |K| , 0 , . . . , 0 ) T ∈ SnK , along with the induced distance function . The n-dimensional Poincaré ball PnK ( also called the Poincaré disk when n = 2 ) for a given curvature K < 0 is defined as PnK = ρK ( HnK ) = { x ∈ Rn : 〈x , x〉2 < − 1 K } , with the induced distance function . 2.2 GYROVECTOR SPACES . An important analogy to vector spaces ( vector addition and scalar multiplication ) in non-Euclidean geometry is the notion of gyrovector spaces ( Ungar , 2008 ) . Both of the above spaces DK and PK ( jointly denoted asMK ) share the same structure , hence they also share the following definition of addition . The Möbius addition ⊕K of x , y ∈MK ( for both signs of K ) is defined as x⊕K y = ( 1− 2K 〈x , y〉2 −K ‖y‖ 2 2 ) x+ ( 1 +K ‖x‖ 2 2 ) y 1− 2K 〈x , y〉2 +K2 ‖x‖ 2 2 ‖y‖ 2 2 . We can , therefore , define “ gyrospace distances ” for both of the above spaces , which are alternative curvature-aware distance functions dDgyr ( x , y ) = 2√ K tan−1 ( √ K ‖−x⊕K y‖2 ) , dPgyr ( x , y ) = 2√ −K tanh−1 ( √ −K ‖−x⊕K y‖2 ) . These two distances are equivalent to their non-gyrospace variants dM ( x , y ) = dMgyr ( x , y ) , as is shown in Theorem A.4 and its hypersphere equivalent . Additionally , Theorem A.5 shows that dMgyr ( x , y ) K→0−−−→ 2 ‖x− y‖2 , which means that the distance functions converge to the Euclidean distance function as K → 0 . We can notice that most statements and operations in constant curvature spaces have a dual statement or operation in the corresponding space with the opposite curvature sign . The notion of duality is one which comes up very often and in our case is based on Euler ’ s formula eix = cos ( x ) + i sin ( x ) and the notion of principal square roots √ −K = i √ K. This provides a connection between trigonometric , hyperbolic trigonometric , and exponential functions . Thus , we can convert all the hyperbolic formulas above to their spherical equivalents and vice-versa . Since Ganea et al . ( 2018a ) and Tifrea et al . ( 2019 ) used the same gyrovector spaces to define an exponential map , its inverse logarithmic map , and parallel transport in the Poincaré ball , we can reuse them for the projected hypersphere by applying the transformations above , as they share the same formalism . For parallel transport , we additionally need the notion of gyration ( Ungar , 2008 ) gyr [ x , y ] v = K ( x⊕K y ) ⊕K ( x⊕K ( y ⊕K v ) ) . Parallel transport in the both the projected hypersphere and the Poincaré ball then is PTKx→y ( v ) = ( λKx /λ K y ) gyr [ y , −x ] v , for all x , y ∈ MnK and v ∈ TxMnK . Using a curvature-aware definition scalar products ( 〈x , y〉K = 〈x , y〉2 if K ≥ 0 , 〈x , y〉L if K < 0 ) and of trigonometric functions sinK = { sin if K > 0 sinh if K < 0 cosK = { cos if K > 0 cosh if K < 0 tanK = { tan if K > 0 tanh if K < 0 we can summarize all the necessary operations in all manifolds compactly in Table 1 and Table 2 . Table 1 : Summary of operations in SK and HK . Distance d ( x , y ) = 1√ |K| cos−1K ( |K| 〈x , y〉K ) Exponential map expKx ( v ) = cosK ( √ |K| ‖v‖K ) x+ sinK ( √ |K| ‖v‖K ) v√ |K| ‖v‖K Logarithmic map logKx ( y ) = cos−1K ( K 〈x , y〉K ) sinK ( cos−1K ( K 〈x , y〉K ) ) ( y −K 〈x , y〉K x ) Parallel transport PTKx→y ( v ) = v − K 〈y , v〉K 1 +K 〈x , y〉K ( x+ y ) Table 2 : Summary of operations in projected spaces DK and PK . Distance d ( x , y ) = 1√ |K| cos−1K ( 1− 2K ‖x− y‖22 ( 1 +K ‖x‖22 ) ( 1 +K ‖y‖ 2 2 ) ) Gyrospace distance dgyr ( x , y ) = 2√ |K| tan−1K ( √ |K| ‖−x⊕K y‖2 ) Exponential map expKx ( v ) = x⊕K ( tanK ( √ |K| λKx ‖v‖2 2 ) v√ |K| ‖v‖2 ) Logarithmic map logKx ( y ) = 2√ |K|λKx tan−1K ( √ |K| ‖−x⊕K y‖2 ) −x⊕K y ‖−x⊕K y‖2 Parallel transport PTKx→y ( v ) = λKx λKy gyr [ y , −x ] v

Summary: This paper introduces a general formulation of the notion of a VAE with a latent space composed by a curved manifold. It follows the current trend of learning representations on curved spaces by proposing a formulation of the latent distributions of the VAE in a variety of fixed-curvature spaces, and introduces an approach to learn the curvature of the space itself. Extensive mathematical derivations are provided, as well as experiments illustrating the impact of various choices of latent manifolds on the performance of the VAE.