Mixed-curvature Variational Autoencoders

TL;DR

This paper introduces a novel Variational Autoencoder with a latent space composed of multiple constant curvature Riemannian manifolds, enabling flexible and unified geometric representations for diverse data types.

Contribution

It develops a mixed-curvature VAE that unifies Euclidean, hyperbolic, and elliptical latent spaces, allowing for learnable or fixed curvatures in each component.

Findings

Generalizes Euclidean VAEs to mixed-curvature spaces

Supports learnable and fixed curvature components

Recovers Euclidean VAE when all curvatures are zero

Abstract

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.