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.

What problem does this paper attempt to address?

The main problem that this paper attempts to solve is to improve the performance of variational auto - encoders (VAE) when dealing with data with hierarchical structures. Specifically, the author proposes a new model named Gaussian Manifold - Variational Auto - Encoder (GM - VAE), whose latent space consists of a set of Gaussian distributions that form a hyperbolic space, called the Gaussian manifold, under the Riemannian metric. In this way, GM - VAE aims to overcome the problems of numerical instability and high computational complexity existing in the existing hyperbolic VAEs, while improving the density estimation performance on image datasets and demonstrating competitiveness in model - based reinforcement learning tasks. ### Main Contributions 1. **Introduction of a new latent space**: The author defines a statistical manifold, which is composed of univariate Gaussian distributions and is called the Gaussian manifold. This manifold has an arbitrary curvature, thus providing greater flexibility. 2. **Proposal of a new distribution**: The author proposes a Pseudo - Gaussian Manifold Normal Distribution (PGM normal distribution), which is easy to sample and has a closed - form KL - divergence, and is suitable for training VAE on the Gaussian manifold. 3. **Experimental verification**: Through experiments in density estimation tasks on multiple image datasets and model - based reinforcement learning tasks, the stability and superior performance of GM - VAE are verified. In particular, GM - VAE significantly outperforms the baseline models in the density estimation tasks of image datasets and shows performance comparable to the state - of - the - art methods in model - based reinforcement learning tasks. ### Key Technologies - **Gaussian manifold**: By re - parameterizing the univariate Gaussian distribution, the statistical manifold can have an arbitrary curvature. - **Pseudo - Gaussian Manifold Normal Distribution**: Using the probability density function defined by KL - divergence, combining the advantages of Gaussian distribution and Gamma distribution, making the distribution easy to sample and having a closed - form KL - divergence. - **Numerical stability**: By avoiding complex Riemannian operations and numerically unstable calculations, the numerical stability of the model is improved. ### Experimental Results - **Density estimation**: On multiple image datasets, GM - VAE performs excellently in density estimation tasks. Especially on the CUB and Oxford102 datasets, regardless of the curvature value, GM - VAE outperforms the baseline models. - **Model - based reinforcement learning**: In the Atari2600 game environment, GM - VAE shows performance comparable to the baseline models with Euclidean and discrete latent spaces in the world model learning task, while the models using common hyperbolic distributions cannot be successfully trained due to numerical stability problems. ### Conclusion This paper successfully solves the problems of numerical instability and high computational complexity existing in the existing hyperbolic VAEs by introducing the Gaussian manifold and the Pseudo - Gaussian Manifold Normal Distribution, and at the same time shows superior performance in multiple tasks. This provides a new and effective method for dealing with data with hierarchical structures.

Hyperbolic VAE via Latent Gaussian Distributions

Variational autoencoders with latent high-dimensional steady geometric flows for dynamics

VTAE: Variational Transformer Autoencoder with Manifolds Learning

Manifold Learning by Mixture Models of VAEs for Inverse Problems

Variational Inference with Gaussian Mixture Model and Householder Flow

Hidden Talents of the Variational Autoencoder

Diffusion Variational Autoencoders

Variational Autoencoder With Optimizing Gaussian Mixture Model Priors.

Learning Manifold Dimensions with Conditional Variational Autoencoders

Deep Unsupervised Clustering with Gaussian Mixture Variational Autoencoders

A Non-negative VAE:the Generalized Gamma Belief Network

Adaptive Learning of the Latent Space of Wasserstein Generative Adversarial Networks

HGMVAE: hierarchical disentanglement in Gaussian mixture variational autoencoder

Improving Variational Autoencoders with Density Gap-based Regularization

Implicit Gaussian process representation of vector fields over arbitrary latent manifolds

LDC-VAE: A Latent Distribution Consistency Approach to Variational AutoEncoders

Neighborhood Geometric Structure-Preserving Variational Autoencoder for Smooth and Bounded Data Sources

Gaussian Copula Variational Autoencoders for Mixed Data

GD-VAEs: Geometric Dynamic Variational Autoencoders for Learning Nonlinear Dynamics and Dimension Reductions

Variational Diffusion Auto-encoder: Latent Space Extraction from Pre-trained Diffusion Models

Physics-enhanced Gaussian Process Variational Autoencoder