Hyperbolic VAE via Latent Gaussian Distributions
This work addresses the challenge of learning stable and effective hyperbolic representations in VAEs, which is incremental but offers practical improvements for tasks like density estimation and reinforcement learning.
The authors tackled the problem of improving variational auto-encoders by using a latent space of Gaussian distributions, which forms a hyperbolic space, and demonstrated that their GM-VAE outperforms other variants on density estimation tasks and shows competitive performance in model-based reinforcement learning, with strong numerical stability.
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.