[Paper Review] Continuous Hierarchical Representations with Poincaré Variational Auto-Encoders

Émile Mathieu, Charline Le Lan+3 more|arXiv (Cornell University)|Jan 17, 2019

Generative Adversarial Networks and Image SynthesisComputer Science44 references49 citations

TL;DR

The paper introduces a VAE with a Poincaré ball (hyperbolic) latent space, deriving two Gaussian generalisations (Riemannian normal and wrapped normal), a geometry-aware decoder, and shows improved generalisation and hierarchical representation recovery over Euclidean VAEs.

ABSTRACT

The variational auto-encoder (VAE) is a popular method for learning a generative model and embeddings of the data. Many real datasets are hierarchically structured. However, traditional VAEs map data in a Euclidean latent space which cannot efficiently embed tree-like structures. Hyperbolic spaces with negative curvature can. We therefore endow VAEs with a Poincaré ball model of hyperbolic geometry as a latent space and rigorously derive the necessary methods to work with two main Gaussian generalisations on that space. We empirically show better generalisation to unseen data than the Euclidean counterpart, and can qualitatively and quantitatively better recover hierarchical structures.

Motivation & Objective

• Motivate learning hierarchical representations by embedding data in hyperbolic space rather than Euclidean space.
• Develop a variational auto-encoder framework with a Poincaré ball latent space.
• Derive and implement two Gaussian generalisations on the Poincaré ball for priors and posteriors.
• Design a decoder that explicitly respects hyperbolic geometry to improve reconstruction and latent interpretability.
• Empirically demonstrate better generalisation and interpretable hierarchies on synthetic data, MNIST, and graph datasets.

Proposed method

• Adopt a Poincaré ball model as the latent space for the VAE.
• Define two Gaussian generalisations on the ball: the Riemannian normal and the wrapped normal, with respective densities and reparameterisable sampling schemes.
• Use a hyperbolic prior p(z) = N_B^d(0, σ0^2) and a variational family q(z|x) = N_B^d(μ, σ^2).
• Propose a decoder architecture (gyroplane layer) that encodes hyperbolic geometry in the output mapping.
• Train by maximizing an evidence lower bound (ELBO) with reparameterised Monte Carlo estimates on the Poincaré ball.
• Parameterise encoder outputs as a Fréchet mean via the exponential map and a positive distortion for the posterior.

Experimental results

Research questions

• RQ1 Can hyperbolic (Poincaré) latent spaces better capture hierarchical structure in data than Euclidean VAEs?
• RQ2 How do two Gaussian generalisations on the Poincaré ball (Riemannian normal and wrapped normal) compare for prior/posterior modelling in VAEs?
• RQ3 Does a geometry-aware decoder improve generalisation and the interpretability of latent hierarchies?
• RQ4 What empirical gains in generalisation and downstream task performance arise when applying Poincaré VAEs to synthetic hierarchically structured data, MNIST, and graph data?

Key findings

• Poincaré VAEs outperform Euclidean VAEs in generalisation to unseen data on synthetic branching data and MNIST, especially at lower latent dimensions.
• Wrapped and Riemannian normal generalisations enable tractable reparameterisation and density formulations on the Poincaré ball, with the Riemannian normal offering slight advantages in some settings.
• The geometry-aware decoder (gyroplane layer) is crucial for leveraging hyperbolic latent space; ablations show performance gains over baseline decoders.
• MNIST embeddings from Poincaré VAEs yield higher digit-classification accuracy on 2D latent spaces, indicating more discriminative hierarchical structure.
• Poincaré VAEs improve link prediction in graph datasets compared to Euclidean VAEs, illustrating better generalisation for hierarchical network data.