Explorations in Homeomorphic Variational Auto-Encoding

TL;DR

This paper explores the use of manifold-valued latent variables, specifically Lie groups like SO(3), in VAEs to better capture data topology and improve latent space structure.

Contribution

It introduces a method to incorporate Lie group-valued latent variables into VAEs, extending the reparameterization trick for compact connected Lie groups.

Findings

Matching latent variable topology to data manifold preserves structure

Manifold-valued VAEs outperform Gaussian VAEs in topological tasks

Extending reparameterization trick enables Lie group latent variables

Abstract

The manifold hypothesis states that many kinds of high-dimensional data are concentrated near a low-dimensional manifold. If the topology of this data manifold is non-trivial, a continuous encoder network cannot embed it in a one-to-one manner without creating holes of low density in the latent space. This is at odds with the Gaussian prior assumption typically made in Variational Auto-Encoders (VAEs), because the density of a Gaussian concentrates near a blob-like manifold. In this paper we investigate the use of manifold-valued latent variables. Specifically, we focus on the important case of continuously differentiable symmetry groups (Lie groups), such as the group of 3D rotations $\operatorname{SO}(3)$. We show how a VAE with $\operatorname{SO}(3)$-valued latent variables can be constructed, by extending the reparameterization trick to compact connected Lie groups. Our experiments show that choosing manifold-valued latent variables that match the topology of the latent data manifold, is crucial to preserve the topological structure and learn a well-behaved latent space.