[Paper Review] Poincaré Wasserstein Autoencoder
This paper introduces the Poincaré Wasserstein Autoencoder (PWA), a variational autoencoder that models the latent space using the Poincaré ball model of hyperbolic geometry to capture hierarchical data structures. By leveraging the intrinsic curvature of hyperbolic space, PWA achieves improved performance on hierarchical data, demonstrating competitive results in link prediction and effective latent space organization on synthetic and visual datasets.
This work presents a reformulation of the recently proposed Wasserstein autoencoder framework on a non-Euclidean manifold, the Poincaré ball model of the hyperbolic space. By assuming the latent space to be hyperbolic, we can use its intrinsic hierarchy to impose structure on the learned latent space representations. We demonstrate the model in the visual domain to analyze some of its properties and show competitive results on a graph link prediction task.
Motivation & Objective
- To address the limitation of standard VAEs in modeling hierarchical data structures due to their Euclidean latent space.
- To explore the use of hyperbolic geometry—specifically the Poincaré ball model—for learning structured, hierarchical representations in latent spaces.
- To develop a Wasserstein autoencoder framework on a Riemannian manifold with constant negative curvature to better capture hierarchical dependencies.
- To demonstrate that hyperbolic latent spaces can outperform Euclidean spaces in tasks with intrinsic tree-like or partially ordered structures.
- To derive closed-form expressions for Gaussian distributions and sampling procedures in hyperbolic space for use in variational inference.
Proposed method
- Reformulates the Wasserstein autoencoder framework on the Poincaré ball model of hyperbolic space $\mathbb{H}^n$ to enable amortized variational inference on non-Euclidean manifolds.
- Defines a Gaussian distribution in hyperbolic space using the Riemannian exponential map and Fisher information metric, enabling proper prior and posterior parametrization.
- Uses Maximum Mean Discrepancy (MMD) to match the prior and posterior distributions in the latent space, enabling end-to-end training.
- Employs parallel transport and exponential maps for sampling from the prior and posterior distributions in hyperbolic space.
- Applies the model to both visual generation (MNIST) and graph link prediction (Cora, Citeseer, Pubmed) to evaluate performance and structural fidelity.
- Uses distortion measures to quantitatively compare the quality of Euclidean vs. hyperbolic latent space representations on hierarchical synthetic data.
Experimental results
Research questions
- RQ1Can modeling the latent space in hyperbolic geometry improve the representation of hierarchical data structures compared to Euclidean VAEs?
- RQ2How does the Poincaré Wasserstein Autoencoder perform in capturing hierarchical structure in low-dimensional latent spaces?
- RQ3What are the quantitative advantages of using hyperbolic geometry in autoencoders for graph link prediction tasks?
- RQ4How does the PWA model compare to existing models like Euclidean VAEs and hyperspherical VAEs in terms of reconstruction quality and downstream performance?
- RQ5What challenges arise when applying MMD-based objectives in hyperbolic space, and how can they be mitigated?
Key findings
- On synthetic hierarchical datasets, the PWA achieved a mean distortion measure of 0.49, significantly outperforming T-SNE (0.73) and the Euclidean VAE (0.82).
- On the MNIST dataset, the PWA produced latent representations that were pushed toward the boundary of the Poincaré disk, reflecting the exponential volume growth of hyperbolic space.
- The PWA achieved a test AUC of 93.9 ± 0.2 and AP of 93.2 ± 0.4 on the Cora citation network, outperforming the Euclidean VAE (AUC: 92.7 ± 0.2) and matching the hyperspherical VAE.
- On Citeseer, the PWA achieved an AUC of 92.2 ± 0.2 and AP of 91.8 ± 0.2, outperforming the Euclidean VAE (AUC: 90.3 ± 0.5) and approaching the performance of the hyperspherical VAE.
- On Pubmed, the PWA achieved an AUC of 95.9 ± 0.2 and AP of 96.3 ± 0.2, slightly underperforming the Euclidean VAE (AUC: 97.1 ± 0.0) but still showing strong performance.
- The model demonstrated that hyperbolic latent spaces can effectively encode hierarchical information with lower-dimensional representations, consistent with findings in prior work on hyperbolic word embeddings.
Better researchstarts right now
From reading papers to final review, dramatically reduce your research time.
No credit card · Free plan available
This review was created by AI and reviewed by human editors.