Source-linked AI summary
Continuous Hierarchical Representations with Poincaré Variational Auto-Encoders
Emile Mathieu, Charline Le Lan, Chris J. Maddison, Ryota Tomioka, Yee Whye Teh
TL;DR
VAEs commonly use Euclidean latent spaces, although many datasets have hierarchical structure that Euclidean geometry does not efficiently represent. This paper equips VAEs with a Poincaré-ball latent space, derives trainable hyperbolic Gaussian machinery and a geometry-aware decoder, and reports better generalisation and more interpretable hierarchical representations.
Problem
Traditional VAEs use Euclidean latent spaces, while many datasets contain hierarchical structure that motivates more suitable representations.
Method
The paper builds Poincaré VAEs with two hyperbolic Gaussian generalisations, reparametrisable sampling, probability densities, and a decoder that accounts for hyperbolic geometry.
Results
Poincaré-ball VAEs can improve model generalisation and produce more interpretable representations for hierarchical structure.
Takeaways & Limitations
Hyperbolic latent geometry is beneficial for representing and discovering hierarchies when paired with a decoder that explicitly accounts for that geometry.
Takeaways & Limitations
It remains unclear which hyperbolic models should be preferred and which of their properties matter, and principled hierarchy assessment for datasets would be useful.
Abstract
from arXiv · showhide
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.
1 Introduction
VAEs learn low-dimensional representations, but hierarchical data motivates replacing Euclidean latent geometry with hyperbolic structure. The paper develops Poincaré VAEs and reports benefits for hierarchical representation, generalisation, and interpretability.
- VAEs encode high-dimensional observations into stochastic low-dimensional latent representations and decode them back into observations.These representations can support downstream tasks or be interpretable in their own right.
- Hierarchical organisation appears across domains including cognition, biology, and probabilistic modelling.The paper cites object categories and evolutionary relationships as examples of hierarchical structure.
- Hyperbolic spaces provide a continuous alternative to trees for learning hierarchical representations, with trees embeddable in the Poincaré disc at arbitrarily low error.The Poincaré surface area grows exponentially, supporting tree-like structure.
- The paper replaces Euclidean VAE latent components with hyperbolic generalisations to represent and discover hierarchies.Its goals are interpretable hierarchical relationships and more efficient representations that generalise better to hierarchically structured unseen data.
- The authors derive efficient reparametrisable sampling and densities for maximum-entropy and wrapped normal distributions on the Poincaré ball.These distributions provide ingredients required to train the proposed VAEs.
- The model introduces a geometry-aware decoder and empirically demonstrates improved generalisation and more interpretable representations.The decoder explicitly accounts for hyperbolic geometry, which the authors report as crucial.
2 The Poincaré Ball model of hyperbolic geometry
The Poincaré ball is a negatively curved Riemannian model in which distances, geodesics, and exponential maps are defined through its metric. Möbius addition and closed-form maps provide the geometric operations used by the model.
- 2.1 Review of Riemannian geometry: A Riemannian manifold equips each tangent space with a metric tensor that defines local inner products, lengths, angles, areas, and volumes.Global distances and related quantities arise by integrating these local geometric contributions.
- 2.1 Review of Riemannian geometry: Geodesics generalise straight lines as constant-speed shortest paths, while the exponential map moves from a point along a geodesic using an initial tangent vector.For the geodesically complete Poincaré ball, the exponential map is defined on the full tangent space.
- 2.2 The Poincaré ball model of hyperbolic geometry: Hyperbolic space is a complete, simply connected Riemannian manifold with constant negative curvature and multiple equivalent geometric models.The paper focuses on the Poincaré ball model.
- 2.2 The Poincaré ball model of hyperbolic geometry: The Poincaré ball defines hyperbolic geometry through a position-dependent metric tensor and its induced distance.The metric is expressed relative to the Euclidean metric tensor.
- 2.2 The Poincaré ball model of hyperbolic geometry: Möbius addition combines points in the Poincaré ball and converges to ordinary Euclidean vector addition as curvature c tends to zero.This operation supports closed-form exponential and logarithm maps for the hyperbolic model.
3 The Poincaré VAE
The Poincaré VAE replaces Euclidean latent-space components with hyperbolic distributions, geometry-aware encoder and decoder maps, and trainable reparameterisations. It uses maximum-entropy and wrapped normal distributions and optimises the ELBO with unbiased Monte Carlo estimation.
- The Pc-VAE maps observations into a Poincaré ball latent space and learns geometry-aware encoder and decoder maps by maximising the ELBO.
- 3.1 Prior and variational posterior distributions: The model considers maximum-entropy Riemannian normal and wrapped normal distributions as hyperbolic generalisations of Gaussian distributions.Both generalisations recover the usual normal distribution as curvature c approaches zero.
- 3.1 Prior and variational posterior distributions: The prior is a zero-mean hyperbolic normal distribution, while the variational family uses a hyperbolic mean and dispersion parameters.
- 3.2 Encoder and decoder: The decoder’s gyroplane layer represents hyperbolic decision boundaries as semi-hyperspheres orthogonal to the Poincaré ball boundary.These gyroplane operators are concatenated and then composed with a standard feed-forward neural network.
- 3.3 Training: Reparameterisable sampling enables an unbiased Monte Carlo estimator for the ELBO, with Riemannian-normal radii sampled efficiently by rejection sampling.The Riemannian normal radius is one-dimensional, so its sampling scheme is not directly affected by dimensionality.
- 3.3 Training: Model parameters in the Poincaré ball are represented through the exponential map so usual optimisation schemes can be applied.The Riemannian normal additionally uses implicit reparameterisation for gradients with respect to dispersion.
4 Related work
Related work includes discrete hierarchical models, hyperbolic representations for graphs and word embeddings, manifold-valued distributions, and auto-encoders with non-Euclidean latent spaces.
- Hierarchical models: Bayesian nonparametric models have long explicitly represented hierarchical structure, while continuous hyperbolic optimisation offers an alternative to discrete trees.
- Hierarchical models: Hyperbolic spaces have been empirically and theoretically useful for graph and word embeddings.
- Distributions on manifold: Manifold-valued probability measures model uncertainty for data living intrinsically or assumptively on non-Euclidean spaces.
- VAEs with Riemannian latent manifold: Prior VAE work used hyperspherical geometry or an SO(3) group structure, while concurrent work considered hyperbolic latent spaces for auto-encoders.
5 Experiments
Experiments evaluate hyperbolic VAEs on synthetic hierarchies, MNIST, and network link prediction. Across these settings, Poincaré-ball models improve generalisation and produce more interpretable hierarchical embeddings.
- 5.1 Branching diffusion process: On the synthetic branching-diffusion dataset, Pc-VAE outperforms the Euclidean N-VAE in test marginal likelihood.As curvature c tends to zero, N-VAE performance is recovered.
- 5.1 Branching diffusion process: Increasing prior distortion σ0 places embeddings closer to the Poincaré-ball border and improves generalisation.The learned Pc-VAE hierarchy is also the least distorted among the compared representations.
- 5.1 Branching diffusion process: Figure 5 compares P1-VAE, N-VAE, PCA, and GPLVM embeddings with posterior samples against the true hierarchy.Embeddings are black crosses, posterior samples are coloured dots, and true hierarchical links are blue lines.
- 5.2 Mnist digits: On MNIST, the hyperbolic model outperforms its Euclidean counterpart, especially at low latent dimension.The gain decreases as latent dimensionality increases, and the Riemannian normal distribution slightly outperforms the wrapped normal.
- 5.2 Mnist digits: The gyroplane decoder layer additionally improves MNIST generalisation over an MLP baseline, although gains decrease with latent dimensionality.Pre-composing the MLP with the logarithm map provides a smaller improvement.
- 5.2 Mnist digits: P-VAE embeddings yield a clearer partitioning of MNIST digits and an average 2% classifier-accuracy increase over the digits.The reported groupings include {4, 7, 9}, {0, 6}, {2, 3, 5, 8}, and {1}.
- 5.3 Graph embeddings: For network link prediction, the P-VAE generalises better than its Euclidean counterpart to unseen edges.Evaluation reports AUC and average precision on held-out edges and randomly sampled unconnected node pairs.
6 Conclusion
The paper develops VAEs with Poincaré-ball latent spaces, covering manifold probability distributions and the ingredients needed for training. It concludes that hyperbolic latents can improve generalisation and interpretability for hierarchically structured data, while leaving model choice and hierarchy detection as open directions.
- Poincaré-ball VAEs can improve model generalisation and yield more interpretable representations when data has hierarchical structure.
- The paper derives efficient, reparametrisable sampling schemes and probability density functions for wrapped and maximum-entropy normal distributions.
- The study introduces a decoder architecture that explicitly accounts for hyperbolic geometry and empirically finds it crucial for making the hyperbolic latent space useful.
- Future work should clarify which hyperbolic models and properties matter and develop principled methods for assessing whether datasets possess underlying hierarchical structure.
- Wrapped and maximum-entropy normal generalisations have different theoretical and computational advantages, including straightforward sampling for wrapped distributions and challenging sampling or normalisation for maximum-entropy distributions.
B.3 Wrapped hyperbolic normal distribution on Bd
The wrapped hyperbolic normal is constructed by mapping a Gaussian from the tangent space onto the Poincaré ball through the exponential map. Hyperbolic polar coordinates support its sampling and density calculations, including anisotropic settings.
- The wrapped normal pushes a tangent-space Gaussian forward along the exponential map centered at the Fréchet mean.
- Its probability density follows from the change-of-variable formula when the exponential-map Jacobian or its inverse is computable.
- Anisotropic wrapped normal measures are parameterised by a Fréchet mean, concentration matrix Σ = diag(σ), and curvature c = 1.
- In the isotropic case, the hyperbolic radius follows a χ-distribution-like radial construction, yielding the wrapped normal density.
B.4 Maximum entropy hyperbolic normal distribution on Bd
The maximum-entropy, or Riemannian, hyperbolic normal is defined by maximising entropy under mean and covariance constraints. Its density requires a normalisation constant, and its relation to the wrapped normal depends on curvature and dispersion.
- The Riemannian normal is the distribution that maximises entropy given an expectation and covariance matrix.
- Its density is formulated using the Riemannian distance and the metric-induced measure on the manifold.
- The density can be evaluated pointwise once the normalisation constant Z_R is known.
- As curvature and dispersion become smaller, the Riemannian normal approaches the wrapped normal; as they become larger, the two distributions diverge.
- Expanding the hyperbolic sinh term with the binomial formula produces a developed expression for the radial density.
B.4.2 Sampling
Sampling the maximum-entropy hyperbolic normal reduces to sampling a one-dimensional radial distribution and a direction. Because the radial density lacks an analytically invertible CDF, the paper uses rejection-based schemes, especially adaptive rejection sampling.
- The hyperbolic normal sampler uses a reparametrisation for the mean and an implicit reparametrisation for the dispersion gradients.
- Hyperbolic sampling is harder than Euclidean sampling because dimensional factorisation and simple rescaling do not extend to the hyperbolic geometry.
- The hyperbolic radius has an unfamiliar distribution with a seemingly non-invertible cumulative density function, motivating rejection sampling.
- Adaptive rejection sampling exploits the log-concavity of the radial density with a piecewise exponential proposal that adapts to σ, c, and d.
- Non-adaptive truncated-normal and Gamma proposals are derived, but their rejection rates scale poorly with dimensionality d and distortion σ.
- The normalisation constant is derived by integrating the expanded radial density and its cumulative distribution function.
- For c = 1 and d = 2, the derived expression recovers the formula reported by Said et al. (2014).
C Experimental details
This section provides additional details about the datasets, architecture designs, and optimisation schemes used for the experimental results in Section 5.
- The section expands the experimental description beyond the results presented in Section 5.
- It covers the datasets used in the experiments.
- It also describes model architectures and optimisation schemes.
C.1 Synthetic Branching Diffusion Process
The experiments use synthetic branching data, MNIST, and graph datasets with Euclidean and hyperbolic VAE architectures. Training details specify model parameterisations, latent dimensions, likelihoods, priors, and optimisation procedures.
- Synthetic Branching Diffusion Process: Synthetic branching diffusion data are generated as a hierarchy with depth 6, branching factor 2, observation dimension n = 50, and five noisy observations per node.
- Synthetic Branching Diffusion Process: N-VAE and Pc-VAE use unit-variance Gaussian decoders, while their encoders use isotropic normal and isotropic hyperbolic normal distributions, respectively.
- Synthetic Branching Diffusion Process: The synthetic models use two-dimensional latent spaces, with Pc-VAE incorporating gyroplane units in the decoder’s first layer.
- Synthetic Branching Diffusion Process: Synthetic data are centred, normalised to unit variance, and randomly split into training and testing sets with proportion 0.7.
- Optimisation: Training uses Adam with learning rate 1e−3 for synthetic models, 5e−4 for MNIST, and 1e−2 with full-batch updates for graph models.
- MNIST: MNIST contains 60,000 training and 10,000 test images of ten handwritten digits, with latent dimensions d = 2, 5, 10 and 20 evaluated.
- Graph datasets: Graph experiments cover advisor-advisee, phylogenetic, and disease-relationship networks using a two-layer GCN and metric-based edge probabilities.
- Graph datasets: The graph models use five-dimensional latent spaces, Wrapped Gaussian priors and posteriors, and manifold projection for Poincaré-ball encoder outputs.
D More experimental qualitative results
The additional figures visualise hyperbolic latent representations across curvatures and models, alongside MNIST classifier confusion matrices.
- Branching diffusion representations: Figures 9 and 10 show Pc-VAE branching-process representations for curvatures c = 1.2, 0.3, 0.1 from left to right.
- Model comparison: Figure 11 compares branching-process latent representations from P1-VAE and N-VAE, including a heatmap of log distance to the hyperplane.
- MNIST classification: Figure 12 presents average MNIST confusion matrices for classifiers trained on P1.4-VAE and N-VAE embeddings.