[Paper Review] Explorations in Homeomorphic Variational Auto-Encoding

Luca Falorsi, Pim de Haan+5 more|arXiv (Cornell University)|Jul 12, 2018

Neural Networks and Applications10 references59 citations

TL;DR

The paper extends VAEs to manifold-valued latent variables on Lie groups (notably SO(3)) by a reparameterization trick, showing topology-matching latent spaces preserve structure and improve learning.

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.

Motivation & Objective

• Motivate the use of manifold-valued latent variables when data lie on non-trivial topologies to preserve topological structure in latent space.
• Develop a reparameterization trick for distributions on compact connected Lie groups, with a focus on SO(3).
• Design an encoder that learns a homeomorphism from the data manifold to the latent Lie group and a decoder that leverages group actions to respect structure.
• Empirically show that matching latent topology to data manifold improves continuity and reconstruction and that group-action decoders outperform standard decoders.

Proposed method

• Reparameterization: Sample from a scale-reparameterizable distribution on R^3, map via the exponential map to so(3), then left-multiply by a group element to center the distribution on SO(3).
• Density: Prove the induced pushforward density on SO(3) is absolutely continuous w.r.t. the Haar measure and derive the density via the exponential map with Rodrigues formula.
• Encoder design: Split into enc^mu (maps to a group element R_mu) and enc^sigma (maps to a Euclidean scale); enc^mu relies on a fixed surjective pi: Y -> SO(3) and an embedding of SO(3) into Y to achieve a homeomorphism.
• Decoder: Implement a group-action decoder that rotates latent representations using SO(3) and projects back to the data space; propose a Fourier-analytic decoder using Wigner-D matrices for pose handling.
• Experimentation: Compare SO(3) latent variables with Gaussian and hyperspherical baselines on synthetic SO(3) embeddings and on images of rotated colored cubes; evaluate continuity and likelihood via NLL/ELBO and reconstruction.
• Exploration: Analyze mean parameterizations (e.g., q, alg, s2s1, s2s2) and show topology-aware parameterizations yield continuous embeddings.

Experimental results

Research questions

• RQ1 Can latent variables living on a Lie group (e.g., SO(3)) be reparameterized for VAEs to enable end-to-end learning with topological guarantees?
• RQ2 Does aligning the latent space topology with the data manifold improve continuity, reconstruction quality, and log-likelihood compared to standard Gaussian VAEs?
• RQ3 Does a group-action decoder better preserve and utilize latent structure than a plain MLP decoder?
• RQ4 How do different mean parameterizations for SO(3) (e.g., quaternions, Lie algebra, sphere products) affect continuity and representation of the manifold?
• RQ5 Can the proposed approach generalize to other compact, connected Lie groups beyond SO(3)?

Key findings

• A reparameterization trick for SO(3) distributions is derived, enabling VAE training with SO(3) latent variables.
• The encoder learns a homeomorphism between the SO(3) data manifold and the SO(3) latent manifold, preserving topological structure.
• A group-action decoder improves performance and encourages latent space structure matching the group, outperforming MLP decoders.
• Topology-matching latent variables (especially S^2 × S^2 mean parameterization) yield more continuous and faithful latent embeddings than flat Gaussian latents.
• Gaussian and Lie-algebra-based mean parametrizations can introduce discontinuities, whereas topology-aware parametrizations maintain continuity in latent trajectories.
• Experiments on SO(3)-embedded data and rotated cube images demonstrate that only proper manifold-valued latents achieve continuous embeddings and higher log-likelihoods.