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.
翻译:我们提出了一种高斯流形变分自动编码器(GM-VAE),其潜在空间由一组高斯分布构成。已知具有Fisher信息度量的单变量高斯分布集合构成双曲空间,我们称之为高斯流形。为学习基于高斯流形的VAE,我们基于Kullback-Leibler散度(即Fisher-Rao距离平方的局部近似)提出了一种伪高斯流形正态分布,用于在潜在空间上定义密度。在实验中,我们通过两项不同任务验证了GM-VAE的有效性:图像数据集密度估计和基于模型的强化学习环境建模。在密度估计任务中,GM-VAE的表现优于其他双曲和欧几里得VAE变体,并在基于模型的强化学习中展现出具有竞争力的性能。我们观察到该模型具有强数值稳定性,解决了此前双曲VAE中普遍存在的局限性问题。