Mixed-curvature Variational Autoencoders

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.

• The paper introduces a novel framework that extends VAEs to operate in mixed-curvature latent spaces combining Euclidean, hyperbolic, and spherical geometries.
• It employs rigorous Riemannian geometry tools and Wrapped Normal distributions for efficient latent space parameterization to handle complex data structures.
• Experimental results on datasets like MNIST and CIFAR-10 demonstrate enhanced representation learning, especially when curvature parameters are learnable.

Mixed-curvature Variational Autoencoders

Introduction

The paper "Mixed-curvature Variational Autoencoders" (1911.08411) introduces a novel framework for extending Variational Autoencoders (VAEs) into spaces characterized by mixed curvature. Unlike traditional approaches that primarily utilize Euclidean geometries, this research leverages Riemannian manifolds with constant curvature. This innovative application allows VAEs to be generalized to combinations of elliptical, hyperbolic, and Euclidean spaces. The authors argue that distinct geometrical layouts can significantly enhance the capacity of VAEs to model complex and non-linear structures inherent in many data types, such as images and hierarchical data.

Theoretical Framework

The framework is built upon the notion that certain data types are more aptly represented in non-Euclidean spaces. For instance, hierarchical and tree-structured data lend themselves well to hyperbolic spaces due to their expansive volume growth, while spherical and elliptical spaces may suit data with cyclic or closed-loop structures. The paper proposes a product manifold approach wherein VAEs can learn latent space representations situated within products of constant curvature spaces. Each component space within this product can have a fixed or learnable curvature, covering Euclidean, hyperbolic, and spherical configurations.

Geometry and Probability in Riemannian Manifolds

The authors effectively define and manipulate geometric constructs in constant curvature spaces by using exponential and logarithmic maps, alongside parallel transport mechanisms. These constructs are central to developing VAEs that maintain coherent and differentiable operations across the manifold. Integration over these spaces permits the formulation of Wrapped Normal distributions, which stand in for the Gaussian distributions typically employed in Euclidean latent spaces. This choice is made due to the advantageous properties of Wrapped Normal distributions, which exhibit reparameterization capabilities and facilitate computationally efficient estimations of probability densities.

Experimental Results

The experimental evaluation covers diverse datasets, including dynamically-binarized MNIST and CIFAR-10, under varying dimensional and curvature configurations. The results delineate the strengths of mixed-curvature VAEs, particularly when the curvature of latent spaces is allowed to be learnable, which has empirically shown improved performance for low-dimensional latent spaces. Among these, the Riemannian Normal Poincaré ball VAE performed exceptionally well.

Figure 1: Learned curvature across epochs (with standard deviation) with latent space dimension of 6, diagonal covariance parametrization, on the MNIST dataset.

Implications and Future Work

The implications of this work are twofold: first, it significantly diversifies the scope of VAEs by enabling them to operate in spaces that naturally mirror the complex structures of various data forms. This adaptability promises better representation learning, particularly in fields involving hierarchical and structured data. Second, the study opens pathways to exploring other generative models, such as GANs, within these geometrical frameworks, potentially enhancing sampling qualities and feature representations.

Moving forward, the paper suggests several engaging directions, including optimizing VAEs tailored for graph data and exploring adversarial training methodologies. Another prospective avenue is the integration of normalizing flows to further enhance posterior approximation flexibility.

Conclusion

"Mixed-curvature Variational Autoencoders" (1911.08411) presents a critical enhancement in the field of generative models by transcending traditional Euclidean constraints. It proposes a versatile framework adaptable to various synthetic and real-world data scenarios, advocating a robust expansion of the VAE design space. This advancement underscores a pivotal step in aligning machine learning models more intrinsically with the geometric nature of data, fostering potential breakthroughs in representation learning and deep generative modeling.