It has been shown that using geometric spaces with non-zero curvature instead of plain Euclidean spaces with zero curvature improves performance on a range of Machine Learning tasks for learning representations. Recent work has leveraged these geometries to learn latent variable models like Variational Autoencoders (VAEs) in spherical and hyperbolic spaces with constant curvature. While these approaches work well on particular kinds of data that they were designed for e.g.~tree-like data for a hyperbolic VAE, there exists no generic approach unifying 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 can be learned. This generalizes the Euclidean VAE to curved latent spaces, as the model essentially reduces to the Euclidean VAE if curvatures of all latent space components go to 0. Generative models are a growing area of unsupervised learning, that aim to model the data distribution p(x) over data points x in a high-dimensional space X (Doersch, 2016) , usually a subset of a high-dimensional Euclidean space R n , with all the associated benefits: a naturally definable scalar product, vector addition, and others. Yet, many types of data have a strongly non-Euclidean latent structure (Bronstein et al., 2017) , like the set of human-interpretable images. They are usually thought to live on a "natural image manifold" (Zhu et al., 2016) , a lower-dimensional subset of the space in which they are represented. On this continuous manifold, one finds all the images that humans can interpret using their visual system. By moving along the manifold, we can continuously change the content and appearance of interpretable images. As mentioned in Nickel & Kiela (2017) , changing the geometry of the underlying latent space enables us to represent some data better than is possible in the equivalent Euclidean space. Motivated by these observations, a range of methods to learn representations in different spaces of constant curvature have recently been introduced: learning embeddings in spherical spaces (Batmanghelich et al., 2016) , hyperbolic spaces (Nickel & Kiela, 2017; Tifrea et al., 2019; Sala et al., 2018) , and even in products of these spaces (Gu et al., 2019; Anonymous, 2020) . By using a combination of different constant curvature spaces, it aims to match the underlying geometry of the data even closer than the others. However, an open question that remains, is how to choose the dimensionality of partial spaces and their curvatures. A popular approach to generative modeling is the Variational Autoencoder (Kingma & Welling, 2014, VAE) . VAEs provide us with 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 is a generalization of the VAE to products of constant curvature spaces, which have the advantage that we can obtain a better reduction in dimensionality while not making optimization of the model significantly more complex. The resulting latent space is then a "non-constantly" curved manifold in an ambient Euclidean space. Modeling the latent space as a single constant curvature manifold limits the flexibility of the space to assume a shape similar to that of the hypothetical intrinsic manifold. 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 different curvature signs, (ii) we show how to generalize Variational Au-toencoders to learn latent representations on products of constant curvature spaces with generalized Gaussian-like priors, and (iii) our approaches outperform current benchmarks on a synthetic tree dataset (Mathieu et al., 2019) and image reconstruction on MNIST (LeCun, 1998) , Omniglot (Lake et al., 2015) , and CIFAR (Krizhevsky, 2009) for some latent space dimensions. By transforming the latent space and associated prior distributions onto Riemannian manifolds of constant curvature, it has previously been shown that we can learn representations on curved space. Generalizing on the above ideas, we have extended the theory of learning VAEs to products of constant curvature spaces. To do that, we have derived the necessary operations in several models of constant curvature spaces, extended existing probability distribution families to these manifolds, and generalized VAEs to latent spaces that are products of smaller "component" spaces, with learnable curvature. On various datasets, we show that our approach is competitive and additionally has the property that it generalizes the Euclidean variational autoencoder -if the curvatures of all components go to 0, we recover the VAE of Kingma & Welling (2014 An elementary notion in Riemannian geometry is that of a real, smooth manifold M ⊆ R n , which is a collection of real vectors x that is locally similar to a linear space, and lives in the ambient space R n . At each point of the manifold x ∈ M a real vector space of the same dimensionality as M is defined, called the tangent space at point x: T x M. Intuitively, the tangent space contains all the directions and speeds at which one can pass through x. Given a matrix representation G(x) ∈ R n×n of the Riemannian metric tensor g(x), we can define a scalar product on the tangent space: A Riemannian manifold is then the tuple (M, g). The scalar product induces a norm on the tangent space T x M: ||a|| x = a, a x ∀a ∈ T x M (Petersen et al., 2006). Although it seems like the manifold only defines a local geometry, it induces global quantities by integrating the local contributions. The metric tensor induces a local infinitesimal volume element on each tangent space T x M and hence a measure is induced as well dM(x) = |G(x)|dx where dx is the Lebesgue measure. The length of a curve γ : Straight lines are generalized to constant speed curves giving the shortest path between pairs of points x, y ∈ M, so called geodesics, for which it holds that γ * = arg min γ L(γ), such that γ(0) = x, γ(1) = y, and . Using this metric, we can go on to define a metric space (M, d M ). Moving from a point x ∈ M in a given direction v ∈ T x M with constant velocity is formalized by the exponential map: exp x : T x M → M. There exists a unique unit speed geodesic γ such that γ(0) = x and The corresponding exponential map is then defined as exp x (v) = γ(1). The logarithmic map is the inverse log x = exp −1 x : M → T x M. For geodesically complete manifolds, i.e. manifolds in which there exists a length-minimizing geodesic between every x, y ∈ M, such as the Lorentz model, hypersphere, and many others, exp x is well-defined on the full tangent space T x M. To connect vectors in tangent spaces, we use parallel transport PT x→y : T x M → T y M, which is an isomorphism between the two tangent spaces, so that the transported vectors stay parallel to the connection. It corresponds to moving tangent vectors along geodesics and defines a canonical way to connect tangent spaces.
Variational Autoencoders with latent spaces modeled as products of constant curvature Riemannian manifolds improve on image reconstruction over single-manifold variants. This paper introduces a general formulation of the notion of a VAE with a latent space composed by a curved manifold. This paper is about developing VAEs in non-Euclidean spaces.