[Paper Review] Hyperbolic VAE via Latent Gaussian Distributions

Seunghyuk Cho, Juyong Lee+1 more|arXiv (Cornell University)|Sep 30, 2022

Model Reduction and Neural Networks4 citations

TL;DR

This paper proposes Gaussian Manifold VAE (GM-VAE), a variational autoencoder that models the latent space as a Riemannian manifold of univariate Gaussian distributions using the Fisher information metric, forming a hyperbolic geometry. By introducing a pseudo-Gaussian manifold normal distribution based on KL divergence, the method enables stable, efficient sampling and closed-form KL computation, outperforming Euclidean and hyperbolic baselines in density estimation and showing competitive performance in model-based reinforcement learning with improved numerical stability.

ABSTRACT

We propose a Gaussian manifold variational auto-encoder (GM-VAE) whose latent space consists of a set of Gaussian distributions. It is known that the set of the univariate Gaussian distributions with the Fisher information metric form a hyperbolic space, which we call a Gaussian manifold. To learn the VAE endowed with the Gaussian manifolds, we propose a pseudo-Gaussian manifold normal distribution based on the Kullback-Leibler divergence, a local approximation of the squared Fisher-Rao distance, to define a density over the latent space. In experiments, we demonstrate the efficacy of GM-VAE on two different tasks: density estimation of image datasets and environment modeling in model-based reinforcement learning. GM-VAE outperforms the other variants of hyperbolic- and Euclidean-VAEs on density estimation tasks and shows competitive performance in model-based reinforcement learning. We observe that our model provides strong numerical stability, addressing a common limitation reported in previous hyperbolic-VAEs.

Motivation & Objective

• To address numerical instability and sampling inefficiency in existing hyperbolic VAEs based on Poincaré or wrapped normal distributions.
• To leverage the intrinsic hyperbolic geometry of the statistical manifold formed by univariate Gaussians under the Fisher information metric.
• To develop a new prior and variational distribution with closed-form KL divergence and stable sampling for training VAEs on this manifold.
• To empirically validate the model on density estimation and model-based reinforcement learning, demonstrating improved stability and performance.

Proposed method

• The latent space is defined as the Gaussian manifold, where each point corresponds to a univariate Gaussian distribution (μ, σ) with μ ∈ ℝ and σ > 0.
• The Fisher information metric on this manifold induces a hyperbolic geometry isomorphic to the Poincaré half-plane model.
• A pseudo-Gaussian manifold normal distribution is proposed, using the Kullback-Leibler divergence as a local approximation of the squared Fisher-Rao distance.
• The KL divergence between two such distributions decomposes into a product of Gaussian and Gamma distributions, enabling closed-form computation.
• The model uses the log-space transformation for σ to simplify manifold constraints, avoiding complex projections or clipping.
• The method is applied to VAEs for density estimation and world model learning in model-based RL, with stable training and efficient sampling.

Experimental results

Research questions

• RQ1 Can a statistical manifold of univariate Gaussians with the Fisher information metric serve as a stable and effective hyperbolic latent space for VAEs?
• RQ2 Does replacing the Poincaré normal distribution with a KL-based pseudo-Gaussian manifold normal distribution improve numerical stability and sampling efficiency in hyperbolic VAEs?
• RQ3 Can the proposed GM-VAE achieve better generalization in density estimation on image datasets compared to Euclidean and hyperbolic baselines?
• RQ4 Does the hyperbolic latent space in GM-VAE better preserve hierarchical structure in reinforcement learning environments?
• RQ5 Is the proposed method scalable and practical for real-world generative modeling tasks?

Key findings

• GM-VAE achieves state-of-the-art performance on density estimation across multiple image datasets, outperforming both Euclidean and hyperbolic VAE baselines.
• The model demonstrates superior numerical stability, avoiding the common issues of gradient explosion and unstable KL computation seen in prior hyperbolic VAEs.
• The pseudo-Gaussian manifold normal distribution enables closed-form KL divergence and efficient sampling, reducing computational overhead.
• In model-based RL, GM-VAE learns a world model with competitive performance to state-of-the-art baselines, capturing hierarchical structure in the latent space.
• The latent space analysis shows a Pearson correlation of 0.46 between state norms and cumulative rewards, outperforming the Euclidean baseline (0.40), indicating better hierarchy preservation.
• Training time for GM-VAE is comparable to Euclidean VAEs and significantly faster than baseline hyperbolic VAEs, due to simpler manifold constraints and no need for projection or clipping.