[论文解读] Hyperbolic VAE via Latent Gaussian Distributions

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

Model Reduction and Neural Networks被引用 4

一句话总结

本文提出高斯流形变分自编码器（GM-VAE），一种将潜在空间建模为单变量高斯分布的黎曼流形的变分自编码器，采用费舍尔信息度量，形成双曲几何。通过引入基于KL散度的伪高斯流形正态分布，该方法实现了稳定、高效的采样以及闭式KL计算，在密度估计任务中优于欧氏和双曲基线模型，并在基于模型的强化学习中表现出色，同时具备更高的数值稳定性。

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.

研究动机与目标

• 解决现有基于庞加莱或包裹正态分布的双曲VAE中存在的数值不稳定与采样低效问题。
• 利用单变量高斯分布在费舍尔信息度量下形成的统计流形的内在双曲几何结构。
• 为该流形上的VAE训练开发一种新的先验分布与变分分布，实现闭式KL散度计算与稳定采样。
• 在密度估计与基于模型的强化学习任务上对模型进行实证验证，展示其在数值稳定性与性能上的改进。

提出的方法

• 潜在空间被定义为高斯流形，其中每个点对应一个单变量高斯分布（μ, σ），其中 μ ∈ ℝ 且 σ > 0。
• 该流形上的费舍尔信息度量诱导出与庞加莱上半平面模型同构的双曲几何。
• 提出一种伪高斯流形正态分布，使用Kullback-Leibler散度作为平方费舍尔-罗距离的局部近似。
• 两个此类分布之间的KL散度可分解为高斯与伽马分布的乘积，从而实现闭式计算。
• 采用对数空间变换处理σ，以简化流形约束，避免复杂的投影或裁剪操作。
• 该方法被应用于VAE，用于密度估计与基于模型强化学习中的世界模型学习，实现稳定训练与高效采样。

实验结果

研究问题

• RQ1 以费舍尔信息度量为度量的单变量高斯分布统计流形，能否作为VAE的稳定且高效的双曲潜在空间？
• RQ2 将庞加莱正态分布替换为基于KL的伪高斯流形正态分布，是否能提升双曲VAE的数值稳定性与采样效率？
• RQ3 所提出的GM-VAE在图像数据集上的密度估计任务中，是否能实现优于欧氏与双曲基线模型的泛化性能？
• RQ4 GM-VAE中的双曲潜在空间是否能更好地保留强化学习环境中层次结构？
• RQ5 所提出的方法在真实世界生成建模任务中是否具备可扩展性与实用性？

主要发现

• GM-VAE在多个图像数据集的密度估计任务中达到最先进性能，优于欧氏与双曲VAE基线模型。
• 该模型展现出卓越的数值稳定性，避免了先前双曲VAE中常见的梯度爆炸与KL计算不稳定问题。
• 伪高斯流形正态分布实现了闭式KL散度计算与高效采样，显著降低了计算开销。
• 在基于模型的强化学习中，GM-VAE学习到的世界模型性能与最先进基线相当，且在潜在空间中成功捕捉了层次结构。
• 潜在空间分析显示，状态范数与累积奖励之间的皮尔逊相关系数为0.46，优于欧氏基线（0.40），表明其具有更好的层次结构保持能力。
• GM-VAE的训练时间与欧氏VAE相当，且显著快于基线双曲VAE，原因在于更简单的流形约束，且无需投影或裁剪操作。