Title: Hyperbolic VAE via Latent Gaussian Distributions
Venue: NeurIPS 2023
Accepted: yes
Keywords: Hyperbolic space

VAE

Distribution on hyperbolic space

Hierarchical representation learning

Reinforcement Learning
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. We demonstrate the efficacy of GM-VAE on two different tasks: density estimation of image datasets and state representation learning for 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. The implementation is available at https://github.com/ml-postech/GM-VAE.

Hyperbolic VAE via Latent Gaussian Distributions

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. We demonstrate the efficacy of GM-VAE on two different tasks: density estimation of image datasets and state representation learning for model-based reinforcement learning. GMVAE 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. The implementation is available at https://github.com/ml-postech/GM-VAE .

1 Introduction

The geometry of latent space in generative models, such as variational auto-encoders (V AE) (Kingma &Welling, 2013), reflects the structure of the data representations. Mathieu et al. (2019); Nagano et al. (2019); Cho et al. (2022) show that employing hyperbolic space as the latent space improves the preservation of the hierarchical structure within the data. The theoretical background for adopting hyperbolic space lies in the analysis of Sarkar (2011); the tree-structured data can be embedded with arbitrary low distortion in hyperbolic space, while Euclidean space requires extensive dimensions.

Previously proposed hyperbolic VAEs rely on Poincaré normal distribution (Mathieu et al., 2019) or hyperbolic wrapped normal distribution Nagano et al. (2019) for the prior and variational distributions. Unlike the Gaussian distribution in Euclidean space, however, these distributions suffer from several shortcomings, including the absence of closed-form Kullback-Leibler (KL) divergence, numerical instability (Mathieu et al., 2019; Skopek et al., 2019), and high computational cost in sampling (Mathieu et al., 2019).

Meanwhile, we can form a Riemannian manifold from the set of univariate Gaussian distributions by equipping the Fisher information metric (FIM). It is known that the FIM of univariate Gaussian distributions is akin to that of the metric tensor of the Poincaré half-plane model (Costa et al., 2015), providing a perspective of viewing the points in hyperbolic space as univariate Gaussian distributions. In other words, a Gaussian distribution can be mapped to a single point in the open half-plane manifold as shown in Figure 1, where the FIM forms the shortest geodesic distance between two Gaussian distributions. Noting that the numerical issue of Poincaré normal arises from the geodesic distance of hyperbolic space, we question whether this perspective can lead us to define a new distribution with better analytic properties.

In this work, inspired by the fact that KL divergence itself is a statistical distance that locally approximates the geodesic distance (Tifrea et al., 2018), we propose a hyperbolic distribution by

▲

◆

●

(

µ

1

, σ

1

)

(

µ

2

, σ

2

)

(

µ

3

, σ

3

)

µ

σ

x

p

(

x

)

(b) Univariate Gaussians

substituting the geodesic distance of Poincaré normal with the KL divergence between the univariate Gaussian distributions. We then verify that this simple yet powerful alteration results in several practical analytic properties; the proposed distribution reduces into the product of two well-known distributions, i.e., the Gaussian and gamma distributions, which are easy to sample, with a closed-form KL divergence between the proposed distributions. By adopting the proposed hyperbolic distribution, we introduce a new variant of hyperbolic VAE, named Gaussian manifold VAE (GM-VAE), whose latent space is a set of Gaussian distributions.

During the experiments, we observe that the proposed distribution is robust in terms of sampling and KL divergence computation compared to the commonly-used hyperbolic distributions; we briefly explain the reason why others are numerically unstable. Experimental results on the density estimation task with image datasets show that GM-VAE can achieve outperforming generalization performances to unseen data against baselines of Euclidean and hyperbolic VAEs. Application of GM-VAE on model-based reinforcement learning (RL) verifies the feasibility of using hyperbolic space on another domain of task.

We summarize our contributions as follows:

We introduce a variant of V AE whose latent space is defined on a statistical manifold formed by univariate Gaussian distributions, namely Gaussian manifold.

We propose a new distribution, called a pseudo Gaussian manifold normal distribution, which is easy to sample and has closed-form KL divergence, to train VAE on the Gaussian manifold.

We empirically verify that the newly proposed V AE performs stable training without numerical issues on the density estimation task with several image datasets. The proposed model outperforms the baseline Euclidean VAE and other hyperbolic variants.

We show that our method can be used for model-based RL. Specifically, we replace the latent space of the world model with hyperbolic space for learning environments, showing competitive results with a state-of-the-art baseline.

2 Preliminaries

In this section, we first review the fundamental concepts of hyperbolic space and commonly used hyperbolic models. We then explain the Riemannian geometry between statistical objects, showing the connection between the statistical manifold and hyperbolic space.

2.1 Review on hyperbolic space

Riemannian manifold. A n -dimensional Riemannian manifold consists of a manifold M and a metric tensor g : M→ R n × n , which is a smooth map from each point x ∈ M to a symmetric

positive definite matrix. The metric tensor g ( x ) defines the inner product of two tangent vectors for each point of the manifold ⟨· , ·⟩ x : T x M×T x M→ R , where T x M is the tangent space of x .

The metric tensor induces basic Riemannian operations, such as a geodesic, exponential map, log map, and parallel transport. Given two points x , y ∈ M , geodesic γ x : [0 , 1] →M is a unit speed curve on M being the shortest path between γ (0) = x and γ (1) = y . This curve can be interpreted as a generalized path of a straight line in Euclidean space. The exponential map exp x : T x M→M is defined as exp x ( v ) = γ (1) = y given γ is a geodesic starting from γ (0) = x and γ ′ (0) = v ∈ T x M . The log map log x : M→T x M is the inverse of the exponential map, i.e., log x (exp x ( v )) = v . The parallel transport PT x → y : T x M→T y M moves the tangent vector v along the geodesic between x and y . The geodesic distance d M ( x , y ) can be induced by the metric tensor as follows:

Hyperbolic space. One method of classifying Riemannian manifolds, a basic question of differential geometry, is based on curvature. Among different types of curvatures, one popular curvature is the sectional curvature κ g , which is a generalization of the Gaussian curvature in classical surface geometry. Given two linearly independent vector fields X,Y ∈ X ( M ) , the sectional curvature κ g can be computed with Riemannian curvature tensor R : X ( M ) × X ( M ) × X ( M ) → X ( M ) as below:

where R ( X,Y ) Y returns a tensor field assigning a tensor to each point of the Riemannian manifold M . The hyperbolic space is a Riemannian manifold that has the sectional curvature value of constant negative (Nickel & Kiela, 2018). The hyperbolic space is known to be able to embed tree-structured data with arbitrarily low distortion (Sarkar, 2011).

Hyperbolic models. We utilize three famous models of hyperbolic space: the Poincaré disk model, the Lorentz model, and the Poincaré half-plane model.

The Poincaré disk model is a hyperbolic space with an open disk manifold. Earlier hyperbolic machine learning work uses the Poincaré disk model because it has a simple closed-form of the operations, such as exponential and log maps (Mathieu et al., 2019). However, the Poincaré disk model suffers from numerical stability issues when the points exist near the boundary of the manifold.

The Lorentz model is often used as an alteration of the Poincaré disk model (Nickel & Kiela, 2018; Nagano et al., 2019; Bose et al., 2020; Cho et al., 2022). The Lorentz model uses a half hyperboloid manifold, where the closed form of the Riemannian operations exists, so the numerical stability issue of employing the Poincaré disk model is relieved.

The Poincaré half-plane model is another well-known model of hyperbolic space with an open half-plane manifold. The metric tensor of a point of the two-dimensional Poincaré half-plane model ( x, y ) is y -2 diag (1 , 1) .

Numerical stability issues of the hyperbolic models. Hyperbolic space suffers from numerical stability when applied to machine learning algorithms (Yu & Sa, 2021; Skopek et al., 2019; Mathieu et al., 2019). The numerical stability mainly occurs for two reasons: machine precision error and unstable Riemannian operations.

First, due to the machine precision error, the hyperbolic points represented with floating point differ from the real value (Yu & Sa, 2021, 2019). In contrast to Euclidean space, the points of hyperbolic space need to satisfy manifold constraints, e.g., the Poincaré disk model allows points whose Euclidean norm is less than one. A point can be placed near the boundary during the optimization or inference processes. Although the point does not violate the manifold constraint in theory, it can be located on or out of the boundary when represented with a floating point due to machine precision error. We empirically observe that the manifold constraint violation occurs frequently when we need to embed many data points in hyperbolic space. Figure 3 demonstrates the machine precision error of each hyperbolic model.

Second, the Riemannian operations of hyperbolic space can result in a not-a-number (NaN) value when the input value is not in the manifold. For example, the geodesic distance from the Poincaré

disk model and the log mapping of the Lorentz model are unstable Riemannian operations, which are written as:

where α = u 0 v 0 -∑ n i =1 u i v i is the Lorentzian inner product of u , v . For the Poincaré disk model, when the points x , y are near the boundary of the unit disk, the floating point representation of the norm values ∥ x ∥ 2 , ∥ y ∥ 2 becomes one. The denominator of the geodesic distance then becomes zero. For the Lorentz model, if u = v and u contains large values in the coordinates, α becomes less than one which results in NaN because the domain of cosh -1 ( x ) is x ≥ 1 .

2.2 Statistical manifold of univariate Gaussians

A particular case of the Riemannian manifold is a statistical manifold, where each point in the manifold corresponds to a probability distribution. Specifically, the parameter manifold M of the probability distributions p θ : X → R , where θ ∈ M , equipped with the Fisher information metric (FIM) forms a Riemannian manifold (Rao, 1992). The FIM is defined as:

In the parameter space of univariate Gaussian distributions { ( µ, σ ) | µ ∈ R , σ ∈ R > 0 } , the FIM can be simplified as two-dimensional diagonal matrix σ -2 diag (1 , 2) (Costa et al., 2015).

Connection to the Poincaré half-plane model. The diagonal form of the FIM implies that the Riemannian manifold with { ( µ, σ ) | µ ∈ R , σ ∈ R > 0 } has the same set of points as the manifold of the Poincaré half-plane, but with different curvature value of -0 . 5 .

The parameter space of the n -dimensional diagonal Gaussian distributions becomes the product of n manifolds of the parameter space of univariate Gaussian distributions. In turn, the statistical manifold of n -dimensional diagonal Gaussian distributions can be viewed as the product of n hyperbolic spaces. The operations on the product of the Riemannian manifolds ⊗ n i =1 M i are defined manifold-wise. For example, an exponential map applied on a point ( p i ) n i =1 ∈ ⊗ n i =1 M i , with tangent vector v i ∈ T p i M i for each i ∈ { 1 , · · · , n } , can be represented as (exp p i ( v i )) n i =1 .

Distance in the statistical manifold. In the statistical manifold, distance functions measure the difference between two distributions on the statistical manifold. One example is the geodesic distance derived from the FIM, which is called the Fisher-Rao distance. The Fisher-Rao distance of the statistical manifold of univariate Gaussian distributions is the same as the geodesic distance of the Poincaré half-plane model with constant negative curvature -0 . 5 .

KL divergence is another widely-used statistical distance for distributions, defined as D KL ( p ∥ q ) := ∫ x p ( x ) log p ( x ) q ( x ) dx for two distributions p, q in the same statistical manifold. One notable property of KL divergence is that it can locally approximate the squared Fisher-Rao distance (Tifrea et al., 2018):

3 Method

In this section, we first present the concept of the Gaussian manifold, which can have an arbitrary curvature by reparameterizing univariate Gaussian distribution. We then propose a pseudo Gaussian manifold normal distribution. Finally, we suggest a new variant of the V AE defined over the Gaussian manifold with PGM normal as prior. We denote the density function of the Gaussian distribution as N ( x ; µ, σ 2 ) = 1 / ( √ 2 πσ 2 ) exp ( -( µ -x ) 2 / (2 σ 2 ) ) .

3.1 Gaussian manifold with arbitrary curvature

Previous studies on hyperbolic space emphasize the importance of having an arbitrary curvature (Skopek et al., 2019; Mathieu et al., 2019). These works empirically show that the generalization performances of hyperbolic VAEs can be improved with varying curvatures. However, as shown in Section 2.2, the univariate Gaussian distributions form a manifold with curvature value of -0 . 5 , limiting the flexibility of the manifold.

We show that the statistical manifold of univariate Gaussian distributions can have an arbitrary curvature by reparameterizing the univariate Gaussian distribution properly. Let N ( √ 2 cµ, σ 2 ) be the reparameterized univariate Gaussian distribution with additional parameter c > 0 . The reparameterization leads to the FIM of σ -2 diag (1 , 1 /c ) showing that the curvature of the statistical manifold is -c . The computation of the sectional curvature of the extended FIM is described in Appendix B.1.

We call the statistical manifold with the reparameterized univariate Gaussian distributions and the extended FIM as the Gaussian manifold and denote it as G c , where -c is the curvature of the Gaussian manifold.

We then verify that the KL divergence between the distributions of the Gaussian manifold approximates the geodesic distance, even in the presence of arbitrary curvature in the Gaussian manifold. Let ( µ, σ ) ∈ G c be an arbitrary point of the Gaussian manifold. The KL divergence between ( µ, σ ) and its neighbor ( µ + dµ, σ + dσ ) can be computed as:

where the first term is the squared Riemannian norm of the tangent vector ( dµ, dσ ) approximating the squared Fisher-Rao distance. The detailed derivation of the KL divergence of the Gaussian manifold is described in Appendix B.2.

3.2 Pseudo Gaussian manifold normal distribution

We propose a pseudo Gaussian manifold (PGM) normal distribution defined over the Gaussian manifold. Let ( µ, σ ) ∈ G c be a point in the Gaussian manifold. Inspired by the Riemannian normal distribution (Pennec, 2006), we define the probability density function of PGM normal with the KL divergence as:

where ( α, β ) ∈ G c , and γ ∈ R > 0 are the parameters of the distribution. The distribution is centered at ( α, β ) with additional scale parameter γ . We verify the convergence of the PGM normal and compute the normalizing constant Z ( c, β, γ ) over the probability measure of the Gaussian manifold at Appendix C.1. As shown in Equation 1, the KL divergence of the Gaussian manifold approximates the Fisher-Rao distance between N ( √ 2 c · α, β 2 ) and N ( √ 2 c · µ, σ 2 ) . Therefore, the PGM normal accounts for the geometric structure of the univariate Gaussian distributions.

The factorization of the probability density function in Equation 2 multiplied with the square root of the determinant of the FIM shows the advantages of the PGM normal, which can be written as:

where Gamma( z ; a, b ) = b a Γ( a ) z a -1 exp( -bz ) and g is the FIM of the Gaussian manifold. The √ det g term enables us to sample and compute the KL divergence in an Euclidean manner. Thanks to the properties of Gaussian and gamma distributions, the PGM normal is easy to sample and has a closed-form KL divergence. The detailed derivation is available in Appendix C.2 and Appendix C.3. The factorization has the same form as the well-known conjugate prior to the Gaussian distribution. In that sense, the PGM normal explicitly incorporates the geometric structure between Gaussians into the known prior distribution.

We note that the PGM normal can be easily extended for the diagonal Gaussian manifold, a manifold formed by diagonal Gaussian distributions since the diagonal Gaussian manifold is the product of the Gaussian manifolds.

3.3 Gaussian manifold VAE

We introduce a Gaussian manifold VAE (GM-VAE) whose latent space is defined over the diagonal Gaussian manifold with the help of the PGM normal. We use the PGM normal for variational and prior distributions. To be specific, with the PGM normal, the evidence lower bound (ELBO) of the GM-VAE can be formalized with the diagonal Gaussian manifold { ( µ , Σ) | µ ∈ R n , Σ ∈ R n > 0 } as:

where p θ ( x | µ , Σ) is the decoder network, q ϕ ( µ , Σ | x ) is the encoder network and p ( µ , Σ) is the prior. The variational distribution is set to q ϕ ( µ , Σ | x ) = K c ( α ϕ ( x ) , β ϕ ( x ) , γ ϕ ( x )) , where α θ ( x ) ∈ R n and β ϕ ( x ) , γ ϕ ( x ) ∈ R n > 0 , and the prior is set to p ( µ , Σ) = K c ( 0 , I, I ) in our experiments given curvature -c . The pseudo-algorithm for the decoder of GM-VAE is present at Algorithm 1.

4 Related Work

Information geometry on VAE. Focusing on the bridge between probability theory and differential geometry, the adaptation of information geometry to the deep learning framework has been investigated in various aspects (Karakida et al., 2019; Bay & Sengupta, 2017; Gomes et al., 2022). Having said that, Han et al. (2020) show that the training process of V AE can be seen as minimizing the distance between the two statistical manifolds: manifolds with the parameters of the decoder and the encoder. Not only can the parameters but the outputs from the V AE decoder be modeled as probability distributions. Arvanitidis et al. (2021) suggest a method of using the pull-back metric defined with arbitrary decoders on the latent space. Our work focuses more on the statistical manifolds lying on the outputs of the encoder with the benefits from the information geometry.

VAE with Riemannian manifold latent space. The latent space of VAE reflects the geometrical property of the representations of the data. The efficacy of setting the latent space to be hyperbolic space (Mathieu et al., 2019; Nagano et al., 2019; Cho et al., 2022) or elliptic space (Xu & Durrett, 2018; Davidson et al., 2018) has been verified for various datasets. Skopek et al. (2019) further extend the approach to enable the latent space to be the product of Riemannian manifolds with different learnable curvatures. On top of these, we explore the method of setting the latent space to be the diagonal Gaussian manifold, which can be viewed as the product of hyperbolic spaces, and provide a novel viewpoint on prior work with information geometry.

Distributions over hyperbolic space. Defining a tractable distribution over hyperbolic space is challenging. Nagano et al. (2019) suggest hyperbolic wrapped normal distribution (HWN) from the observation that the tangent space is Euclidean space. Leveraging operations defined on the tangent spaces, e.g., parallel transport, enables an easy sampling algorithm. Mathieu et al. (2019) propose a rejection sampling method for the Riemannian normal distribution defined on the Poincaré disk model, namely Poincaré normal distribution. This method rejects the pathological samples and enables accurate sampling from the distribution in exchange for high computational complexity.

Although these distributions are widely adopted in many applications (Skopek et al., 2019; Mathieu & Nickel, 2020; Cho et al., 2022), one can barely adopt the full covariance matrix due to the difficulties

Algorithm 1 Decoder

Input Parameter ( α, β ) ∈ G c , γ , Decoding layers Dec ( · ) Output Reconstruction x ′

1: Sample µ ∼ N ( α, β 2 γ 2 )

2: Sample log σ 2 ∼ Gamma ( 1 4 cγ 2 +1 , 1 4 cβ 2 γ 2 )

3: x ′ = Dec ([ µ, log σ 2 ])

4: return x ′

in Monte-Carlo based KL approximation. The number of samples to approximate the KL divergence increases exponentially when the full covariance matrix is used (Cho et al., 2022), so it is common to use isotropic or diagonal covariance instead. Especially in the Poincaré normal, the computation of KL divergence is slow due to the expensive rejection sampling.

RLwith hyperbolic space. The hierarchical relationship between the states lying on the trajectories earned from RL agents has been gaining attention recently. Nagano et al. (2019) have studied that the hierarchical structure of Atari2600 Breakout game states can be well-captured with hyperbolic V AEs. We compare the same task, where GM-VAE outperforms the previous work. Cetin et al. (2022) suggest using hyperbolic space as the geometric prior for representation learning in model-free RL agent, showing improvements in generalization performances. Here, we focus on model-based RL, especially the method of using the world model (Ha & Schmidhuber, 2018; Hafner et al., 2020), and open the possibility of applying hyperbolic space to broader domains of RL by solving the bottleneck of the numerical stability.

5 Experiments

In this section, we demonstrate the performances of GM-VAE on two tasks: density estimation of image datasets and model-based RL. We remark on the practical properties of GM-VAE shown in the experiments with additional analyses.

5.1 Density estimation on image datasets

We conduct density estimation on image datasets to measure the effectiveness of hyperbolic latent space against Euclidean space with the proposed GM-VAE. We use three datasets: the images from Atari2600 Breakout with binarization (Breakout) (Nagano et al., 2019), Oxford 102 Flower (Oxford102) (Nilsback & Zisserman, 2008), Food101 (Bossard et al., 2014), and Caltech-UCSD Birds200-2011 (CUB) (Wah et al., 2011). The datasets are chosen with the four lowest δ -hyperbolicity ( δ -H), a metric that measures how the given images are well-embed in hyperbolic space. Low δ -H implies that the dataset is likely to embed in hyperbolic space. The details about δ -H are available in Appendix D. The values of δ -H for the four datasets and other candidate datasets are in Appendix D.2. Several studies show that the images from the chosen datasets have an implicit hierarchical structure (Nagano et al., 2019; Li et al., 2019; Bossard et al., 2014; Kerdels & Peters, 2015).

We compare GM-VAE with the three baseline models: VAE with Euclidean latent space ( E -VAE), and hyperbolic VAE equipped with HWN ( L -VAE) and Poincaré normal ( P -VAE). We use the product

latent space for both L -VAE and P -VAE, and set the curvature value to -1 . The other details on the implementation and experimental setups are described in Appendix E.1.

The results are reported at Table 1. GM-V AE outperforms the baselines in all the settings, except one case of Breakout. Especially in CUB and Oxford102, GM-VAE outperforms the baselines regardless of the curvature value. In Breakout, P -VAE shows inferior performance due to unstable training, and L -VAE fails in some of the runs with small latent dimension. The results of P -VAE and L -VAE with non-product latent space, a common choice in previous work, are also present in Appendix F.

5.2 State representation learning in model-based RL

We focus on the model-based RL task to verify the utility of GM-V AE on various tasks. Specifically, we apply GM-VAE to a world model, which aims to learn the representation of the environments (Ha &Schmidhuber, 2018; Hafner et al., 2019a,b, 2020). We use DreamerV2 (Hafner et al., 2020) as the baseline model to evaluate the performance of GM-VAE in modeling environments. DreamerV2 is composed of a recurrent state space model (RSSM) (Hafner et al., 2019b) and three predictors for the image p ϕ ( x t | h t , z t ) , the reward p ϕ ( r t | h t , z t ) , and the discount factor p ϕ ( γ t | h t , z t ) , where x t is the observation which the format is the image, r t is the reward, γ t is the discounting factor, h t is the deterministic recurrent state, and z t is the stochastic state. The model is trained by maximizing the likelihood of p ( x , r , γ | a ) given observations x , rewards r , and discount factors γ earned from the sequence of actions a of an agent. By deriving the evidence lower bound of p ( x , r , γ | a ) , the world model is learned to optimize the likelihood with the variational distribution q θ ( z t | h t , x t ) with the following objective L ( ϕ, θ ) as:

where β is KL loss scaling factor, T is the length of input sequence, p ϕ is the prior, and q ϕ is the approximated posterior. GM-VAE is employed by replacing the space of z t with the Gaussian manifold and two components in RSSM, the representation model q θ ( z t | h t , x t ) and transition predictor p ϕ ( z t | h t ) , with PGM normal.

We compare evaluation scores between different types of latent space on world model learning over the Atari2600 environments. The agents are trained with 100M environment steps. We select games having the δ -H values of the four lowest and the two highest among 60 popular Atari2600 games. The other details on the implementation and experimental setups are described in Appendix E.2 with the δ -H for all 60 games in Appendix D.3.

With a commonly-used hyperbolic distribution, i.e., HWN, we observe that training the world model fails due to the numerical stability issue. On the other hand, GM-V AE shows competitive results with the baselines in Euclidean and discrete latent space in all the games we test. The results are reported in Figure 2b. We note that the reproduced Euclidean baseline results by using the official code are better than those reported in Hafner et al. (2020).

5.3 Remark on GM-VAE

Numerical stability. One notable property of GM-VAE is the numerical stability during training compared to L -VAE and P -VAE. During the experiments, L -VAE and P -VAE fail to run in some of the Breakout image density estimations and all the seeds of model-based RL due to the numerical instability. Similar observations are also reported in several previous works (Mathieu et al., 2019; Chen et al., 2021; Skopek et al., 2019). The sampling from a hyperbolic distribution is a major cause of the numerical instability. Consequently, the sampling-based KL divergence computation can be unstable.

We first show that sampling from PGM normal can be stabilized via a simple reparameterization trick. To train GM-VAE, one needs to obtain sample µ and σ from PGN normal K c ( α, β, γ ) . Sampling µ can be done from N ( α, β 2 γ 2 ) without numerical issues. Sampling σ can be done from Gamma ( a, b ) where a = 1 / 4 cγ 2 +1 , b = 1 / 4 cβ 2 γ 2 as shown in Appendix C.2. However, due to machine precision error, often, β violates the manifold constraints, i.e., β = 0 . Eventually, direct sampling of σ can cause the numerical instability. To avoid β being zero, we use the output of the VAE encoder as log β 2 whose value ranges over the entire real numbers and is more stable even when β is close to

(b) Results

zero. With log β 2 , instead of sampling σ 2 , we sample log σ 2 = log ϵ +log b , where ϵ is sampled from Gamma ( a, 1) , through the reparameterization of the Gamma distribution, where log b can be directly computed from log β 2 .

We can show that the KL divergence between an arbitrary PGM normal and prior distribution K c ( 0 , I, I ) has a closed-form solution without any sampling. The KL divergence of PGM normal is the sum of the KL divergences between two Gaussian distributions and between two Gamma distributions, as shown in Appendix C.3. First, the KL divergence between a univariate Gaussian distribution N ( µ, σ 2 ) and the prior distribution can be obtained with log σ 2 as shown in Equation 4. Second, the KL divergence between the two Gamma distributions, Gamma ( a 1 , b 1 ) and Gamma ( a 2 , b 2 ) , written as:

where ψ is the digamma function, can be computed using log b .

Training time comparison. Common bottlenecks of the mode hyperbolic VAEs arise from the complex manifold constraints and the difficulty of sampling from the hyperbolic distributions. For example, the Poincaré disk model of P -VAE and the Lorentz model of L -VAE requires the samples to be inside of a unit disk and to be on a hyperboloid with constraint { x ∈ R n +1 | -x 2 0 + ∑ n i =1 x 2 i } , respectively. Such manifolds need complex transformations, e.g., clipping, projection, or geometric transformations using the Riemannian operations, to

Table 2: The training time of the V AEs in density estimation of the Breakout image dataset. We report the training time of the V AEs in seconds per epoch. GM-VAE is 1.93x faster than P -VAE and 1.41x faster than L -VAE in the experiments held on a single A100 40GB PCI GPU.

match the manifold constraint so making the training of the hyperbolic VAEs slower. The Gaussian manifold, on the other hand, has a much simple manifold constraint and even does not require any transformations if we utilize the log space of σ .

We report the time consumptions of the VAEs with the latent dimension of 8 per epoch in the density estimation of Breakout at Table 2. The results demonstrate that the algorithmic distinctions enable GM-VAE to be trained much faster than the baseline hyperbolic VAEs and even similar to E -VAE.

Latent space analysis. We present a plot of the learned representation in the hyperbolic space at Figure 2a for qualitative analysis. We take the world model with GM-VAE trained for Breakout and illustrate the geodesic starting from the origin in the figure with four generated samples along with the geodesic. We also provide the scatter plot of game states with their cumulative rewards represented in different colors. The brighter the color, the higher the cumulative reward. The scatter

plot reveals that the states with high cumulative rewards are distributed near the origin. Together with the samples from the geodesic, we can observe that the hierarchical structure of Breakout is well captured in the latent space.

Note that in the Poincaré disk model, the depth of the hierarchy is expected to be shown as the distance from the origin (Nickel & Kiela, 2017). When the root node is placed near the origin, the leaf nodes are likely to be placed near the boundary of the open disk. The geodesic lines starting from the origin to the boundary of the Poincaré disk model are identical to the geodesics of the Gaussian manifold starting from (0 , 1) , i.e., the origin of the Gaussian manifold. The connection implies that the data hierarchy should be aligned along the geodesic curves if the hierarchy is well captured.

To quantitatively measure the correlation between the cumulative rewards and the states, we measure the Pearson correlation between the cumulative reward and the norm of the states. We obtain a correlation coefficient of 0.46 from the hyperbolic latent space, whereas the correlation coefficient of the Euclidean latent space is 0.40, showing the hyperbolic space better captures the hierarchy along the increasing norm. More experimental details are explained in Appendix G.2.

6 Conclusion & Future Work

In this work, we propose a novel method of representation learning with GM-VAE, utilizing the Gaussian manifold for the latent space. With the newly-proposed PGM normal defined over the Gaussian manifold, which shows better stability and ease of sampling compared to the commonly-used ones, we verify the efficacy of our method on several tasks. Our method achieves outperforming results on density estimation with image datasets and competitive results on model-based RL compared to the baselines. We explain the behavior of GM-VAE in terms of solving the frequent numerical issue of commonly-used hyperbolic VAEs. The analysis of latent space exhibits that the hierarchy lying in the dataset can be preserved by using GM-VAE.

We suggest that the numerical stability of our method can be helpful for scaling the generative models, e.g., very deep VAE (Child, 2021), endowed with hyperbolic geometrical priors. As GM-VAE is beneficial for capturing hierarchy with promising results in modeling RL environment, another potential future work can be extending the use of hyperbolic space, such as learning a skill tree for solving complex long-horizon tasks (Shi et al., 2022). We believe that the connection between the statistical manifold and hyperbolic space provides new insight to the research community and hope to see more interesting connections and analyses in the future.

Acknowledgments and Disclosure of Funding

This work was partly supported by Institute of Information & communications Technology Planning &Evaluation (IITP) grant funded by the Korea government (MSIT) (No.2019-0-01906, Artificial Intelligence Graduate School Program (POSTECH)) and National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (No. RS-2023-00217286) and National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (NRF-2021R1C1C1011375)

References

Arvanitidis, G., González-Duque, M., Pouplin, A., Kalatzis, D., and Hauberg, S. Pulling back information geometry. arXiv preprint arXiv:2106.05367 , 2021.

Bay, A. and Sengupta, B. Geoseq2seq: Information geometric sequence-to-sequence networks. arXiv preprint arXiv:1710.09363 , 2017.

Bose, J., Smofsky, A., Liao, R., Panangaden, P., and Hamilton, W. Latent variable modelling with hyperbolic normalizing flows. In International Conference on Machine Learning , pp. 1045-1055. PMLR, 2020.

Bossard, L., Guillaumin, M., and Van Gool, L. Food-101 - mining discriminative components with random forests. In European Conference on Computer Vision , 2014.

Cetin, E., Chamberlain, B., Bronstein, M., and Hunt, J. J. Hyperbolic deep reinforcement learning. arXiv preprint arXiv:2210.01542 , 2022.

Chen, W., Han, X., Lin, Y., Zhao, H., Liu, Z., Li, P., Sun, M., and Zhou, J. Fully hyperbolic neural networks. arXiv preprint arXiv:2105.14686 , 2021.

Child, R. Very deep VAEs generalize autoregressive models and can outperform them on images. In International Conference on Learning Representations , 2021.

Cho, S., Lee, J., Park, J., and Kim, D. A rotated hyperbolic wrapped normal distribution for hierarchical representation learning. arXiv preprint arXiv:2205.13371 , 2022.

Costa, S. I., Santos, S. A., and Strapasson, J. E. Fisher information distance: A geometrical reading. Discrete Applied Mathematics , 197:59-69, 2015.

Davidson, T. R., Falorsi, L., De Cao, N., Kipf, T., and Tomczak, J. M. Hyperspherical variational auto-encoders. arXiv preprint arXiv:1804.00891 , 2018.

Fournier, H., Ismail, A., and Vigneron, A. Computing the gromov hyperbolicity of a discrete metric space. Inf. Process. Lett. , 115(6):576-579, jun 2015. ISSN 0020-0190. doi: 10.1016/j.ipl.2015.02. 002.

Gogianu, F., Berariu, T., Bus , oniu, L., and Burceanu, E. Atari agents, 2022. URL https://github. com/floringogianu/atari-agents .

Gomes, E. D. C., Alberge, F., Duhamel, P., and Piantanida, P. Igeood: An information geometry approach to out-of-distribution detection. arXiv preprint arXiv:2203.07798 , 2022.

Ha, D. and Schmidhuber, J. World models. arXiv preprint arXiv:1803.10122 , 2018.

Hafner, D., Lillicrap, T., Ba, J., and Norouzi, M. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.01603 , 2019a.

Hafner, D., Lillicrap, T., Fischer, I., Villegas, R., Ha, D., Lee, H., and Davidson, J. Learning latent dynamics for planning from pixels. In International conference on machine learning , pp. 2555-2565. PMLR, 2019b.

Hafner, D., Lillicrap, T., Norouzi, M., and Ba, J. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02193 , 2020.

Han, T., Zhang, J., and Wu, Y. N. From em-projections to variational auto-encoder. 2020.

Karakida, R., Akaho, S., and Amari, S.-i. Universal statistics of fisher information in deep neural networks: Mean field approach. In The 22nd International Conference on Artificial Intelligence and Statistics , pp. 1032-1041. PMLR, 2019.

Kerdels, J. and Peters, G. Analysis of high-dimensional data using local input space histograms. Neurocomputing , 2015.

Khrulkov, V., Mirvakhabova, L., Ustinova, E., Oseledets, I., and Lempitsky, V. Hyperbolic image embeddings. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , June 2020.

Kingma, D. P. and Welling, M. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013.

Li, A.-X., Zhang, K.-X., and Wang, L.-W. Zero-shot fine-grained classification by deep feature learning with semantics. International Journal of Automation and Computing , 2019.

Mathieu, E. and Nickel, M. Riemannian continuous normalizing flows. Advances in Neural Information Processing Systems , 33:2503-2515, 2020.

Mathieu, E., Le Lan, C., Maddison, C. J., Tomioka, R., and Teh, Y. W. Continuous hierarchical representations with poincaré variational auto-encoders. Advances in neural information processing systems , 32, 2019.

Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidjeland, A. K., Ostrovski, G., et al. Human-level control through deep reinforcement learning. nature , 518(7540):529-533, 2015.

Nagano, Y., Yamaguchi, S., Fujita, Y., and Koyama, M. A wrapped normal distribution on hyperbolic space for gradient-based learning. In International Conference on Machine Learning , pp. 46934702. PMLR, 2019.

Nickel, M. and Kiela, D. Poincaré embeddings for learning hierarchical representations. In Advances in Neural Information Processing Systems , volume 30. Curran Associates, Inc., 2017.

Nickel, M. and Kiela, D. Learning continuous hierarchies in the lorentz model of hyperbolic geometry. In International Conference on Machine Learning , pp. 3779-3788. PMLR, 2018.

Nilsback, M.-E. and Zisserman, A. Automated flower classification over a large number of classes. In 2008 Sixth Indian Conference on Computer Vision, Graphics & Image Processing . IEEE, 2008.

Pennec, X. Intrinsic statistics on riemannian manifolds: Basic tools for geometric measurements. Journal of Mathematical Imaging and Vision , 25(1):127-154, 2006.

Radford, A., Metz, L., and Chintala, S. Unsupervised representation learning with deep convolutional generative adversarial networks. arXiv preprint arXiv:1511.06434 , 2015.

Rao, C. R. Information and the Accuracy Attainable in the Estimation of Statistical Parameters , pp. 235-247. Springer New York, New York, NY, 1992. ISBN 978-1-4612-0919-5. doi: 10.1007/ 978-1-4612-0919-5_16.

Sarkar, R. Low distortion delaunay embedding of trees in hyperbolic plane. In Proceedings of the 19th International Conference on Graph Drawing , GD'11, pp. 355-366, Berlin, Heidelberg, 2011. Springer-Verlag. ISBN 9783642258770. doi: 10.1007/978-3-642-25878-7_34.

Shi, L. X., Lim, J. J., and Lee, Y . Skill-based model-based reinforcement learning. In Conference on Robot Learning , 2022.

Skopek, O., Ganea, O.-E., and Bécigneul, G. Mixed-curvature variational autoencoders. arXiv preprint arXiv:1911.08411 , 2019.

Tifrea, A., Bécigneul, G., and Ganea, O.-E. Poincar \ 'e glove: Hyperbolic word embeddings. arXiv preprint arXiv:1810.06546 , 2018.

Wah, C., Branson, S., Welinder, P., Perona, P., and Belongie, S. The caltech-ucsd birds-200-2011 dataset. Technical Report CNS-TR-2011-001, California Institute of Technology, 2011.

Xu, J. and Durrett, G. Spherical latent spaces for stable variational autoencoders. arXiv preprint arXiv:1808.10805 , 2018.

Yu, T. and Sa, C. D. Numerically accurate hyperbolic embeddings using tiling-based models. In Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019.

Yu, T. and Sa, C. D. Representing hyperbolic space accurately using multi-component floats. In Advances in Neural Information Processing Systems , 2021.

A Machine Precision Error Analysis on Hyperbolic Space

1030

1025

1020

error

1015

The Poincare half-plane model

numerical

Tao and De Sa (2021), M=2

The Lorentzmodel

1010

The Poincare disk model

105

100

10-5

0

10

20

30

40

50

60

70

80

distancefromthe origin

B Gaussian Manifold

B.1 Curvature of the Gaussian manifold

We construct a Riemannian manifold { ( µ, σ ) | µ ∈ R , σ ∈ R > 0 } with a positive constant c and the metric tensor σ -2 diag (1 , 1 /c ) , which we will name Gaussian manifold. We need to show the value of the curvature.

First, we need to compute the Christoeffel symbols of the Gaussian manifold defined as:

where g ij is the ( i, j ) element of the metric tensor and g ij is the ( i, j ) element of the inverse of the metric tensor.

The Christoeffel symbols of the Gaussian manifold are:

Then, the sectional curvature of the space κ g with given tangent vectors du, dv is computed as:

Note that the sectional curvature of two-dimensional Riemannian manifold is same as the Gaussian curvature where ⟨ dµ, dµ ⟩⟨ dσ, dσ ⟩ - ⟨ dµ, dσ ⟩ 2 = det g .

B.2 Gaussian manifold with KL-divergence

Between two univariate Gaussian distributions N ( µ 1 , σ 2 1 ) and N ( µ 2 , σ 2 2 ) , we can compute the KL divergence as:

We extend the KL divergence for an arbitrary curvature of the Gaussian manifold as:

Now, we show that the extended KL divergence still approximates the Riemannian distance of the manifold as:

C Pseudo Gaussian Manifold Normal Distribution

In this section, we derive the normalizing constant Z ( c, β, γ ) and the factorization of the PGM normal which the density function is defined as:

C.1 Normalizing Constant

The given probability density function needs to satisfy the following condition:

where √ det g · d ( µ, σ ) is the probability measure over the Gaussian manifold induced from the Lebesgue measure d ( µ, σ ) . The normalizing factor Z ( c, β, γ ) can be computed using the condition

Equation 5 as:

Finally, the logarithm of the normalizing factor is computed as:

C.2 Sampling

Suppose that p ( µ, σ ) = K c ( µ, σ ; α, β, γ ) √ det g . Sampling µ and σ from the probability distribution p ( µ, σ ) can be done by sampling µ from the marginal distribution p ( µ ) and then sampling σ from the conditional distribution p ( σ | µ ) . The marginal distribution p ( µ ) can be derived from Equation 6 as:

Equation 6 also implies that µ and σ are independent in the aspect of p ( µ, σ ) so the conditional distribution p ( σ | µ ) is identical to the marginal distribution p ( σ ) . The marginal distribution p ( σ ) is computed as:

Here, sampling σ from p ( σ ) can be easily replaced by the procedure of sampling σ 2 from p ( σ 2 ) , which is identical to p ( σ ) / (2 σ ) = Gamma ( σ 2 ; 1 4 cγ 2 +1 , 1 4 cβ 2 γ 2 ) , and then applying square root to the sample σ 2 to get σ .

C.3 KL Divergence

Suppose that p ( µ, σ ) , q ( µ, σ ) are two different PGM normal multplied with det g . As shown in Appendix C.2, µ and σ are independent so the KL divergence between p, q is same as D KL ( p ( µ ) ∥ q ( µ )) + D KL ( p ( σ ) ∥ q ( σ )) . The first term is the KL divergence between two normal distribution. The second term is same with D KL ( p ( σ 2 ) ∥ q ( σ 2 )) due to the change-of-variable formula, so it is the KL divergence between two gamma distribution.

√

D δ -Hyperbolicity

In this section, we explain about δ -hyperbolicity ( δ -H) and show the δ -H values of the candidate datasets of our experiments.

D.1 δ -hyperbolicity

Given a metric space ( X,d ) , the metric space is said to be δ -hyperbolic if, for any geodesic triangle, i.e., a triangle where each side is a geodesic curve, any point on the side of the geodesic triangle is within the distance of less than or equal to δ of the other two sides; when such δ exists, ( X,d ) is said to be hyperbolic. To be specific, when the Gromov product between any three points x, y, z ∈ X is given as:

if then the following inequality holds for any four points x, y, z, w ∈ X , we call the metric space is δ -hyperbolic:

δ -hyperbolicity ( δ -H) is defined to be the minimumvalue of δ satisfying Equation 7 and is used as a measurement quantifying how a given metric space well embeds in hyperbolic Khrulkov et al. (2020); Cetin et al. (2022). Figure 4 illustrates the concept of δ -H. We note that the lower δ -H the metric space has, the less deviation from the exact hyperbolic space is.

Computation. We measure the δ -H values of

A

B

C

δ >

0

(a)

δ

-H

>

0

A

B

C

D

δ

= 0

(b)

δ

-H

= 0

the images X from the datasets by following the procedure from Fournier et al. (2015). We first extract the embeddings of the images using a pre-trained feature extractor to construct a metric space of the images. We then randomly sample a fixed point w and calculate the pairwise Gromov product of the embeddings D with w as Equation 7. We finally determine the δ -H of the images X by finding the largest coefficient of (max k min ij ( D ik , D kj )) -D .

To reduce the scale difference between the datasets, we report the value 2 δ ( X ) / diam ( X ) where diam ( X ) denotes the maximum pairwise distance of X . Because computing the matrix D among all the images X is computationally expensive, we compute the δ -H of randomly sampled 1,000 images from X . We repeat this 10 times and report the average δ -H.

D.2 Image datasets

We measure the δ -H of the image datasets using an ImageNet pre-trained Inception V3 as a feature extractor.

D.3 Atari2600 environments

We collect the Atari2600 images using the pre-trained agents of Gogianu et al. (2022) and measure the δ -H using the image encoder from the agents. For each environment, we report the δ -H of the images which are collected by the agent with the highest reward. We also report the corresponding reward. We run the agents for at least 6 episodes.

E Implementation Details

In this section, we introduce the implementation details of the experiments.

E.1 Density estimation on image datasets

We estimate the density of the images from Atari2600 Breakout (Nagano et al., 2019), Oxford102 (Nilsback & Zisserman, 2008), and CUB (Wah et al., 2011). The images of Breakout are collected from plays with a pre-trained Deep Q-Network (Mnih et al., 2015). The size of images of all the datasets is resized to 64 × 64 , while Breakout is binarized with a threshold value of 0.1; the threshold for Breakout is determined to visualize the components clearly.

We split the datasets into train, validation, and test. For Breakout and CUB, we split the original train set into train and validation sets. For Oxford102, because the original train set is too small, we merge the original train and test set and then split it into three splits. For Food101, we randomly sample the train set and validation set from the original train set, and also randomly sample the test set from the original test set.

We design the encoder and decoder similar to the generator and discriminator of DCGAN (Radford et al., 2015). The details of the architecture are at Table 3. We use learning rate 1e-3, batch size 100, and Adam optimizer for training. We use Bernoulli loss as the reconstruction loss for Breakout experiments and negative log-likelihood loss as the reconstruction loss for CUB, Food101, and Oxford102 experiments. We use the validation set for early stopping and report the negative log-likelihood on the test set with 50 importance weighted samples.

E.2 Model-based RL

We use the official TensorFlow implementation from Dreamerv2 1 to reproduce the baseline results, i.e., with Euclidean and discrete latent space. For the hyperbolic latent space results, we apply GM-VAE by replacing the latent space of z t with the Gaussian manifold and two components in RSSM, the representation model q θ ( z t | h t , x t ) and transition predictor p ϕ ( z t | h t ) , with PGM normal. The hyperparameters are set to be the same as suggested in the original paper, except for the training environment steps being 50M for Freeway and 100M for the others as we observe converging scores.

F Results of Non-Product Latent Space

Previous hyperbolic VAEs are implemented with a hyperbolic space, not the product of the hyperbolic spaces. We run the non-product hyperbolic VAEs in the density estimation of image datasets. Table 4 reveals that the non-product hyperbolic V AEs fail in most of the settings and the product hyperbolic space makes the hyperbolic VAEs much more stable.

G Model-Based RL

G.1 Learning curves

G.2 Latent space analysis

We conduct an analysis of the latent space of the agents learned to play Atari2600 Breakout. The purpose of the analysis is to measure how the latent spaces well-preserve the implicit hierarchy in the trajectory of the agents. To analyze the hyperbolic latent space, we need two isometries: the isometry between the Gaussian manifold and the Poincaré disk model and the translation of the Poincaré disk model.

0.00

1.02

0

2500

5000

7500

10000

Alien

0.00

1.02

0

200

400

600

Breakout

0

5

0

10

20

30

Freeway

0.00

1.02

1e8

20

0

20

Hockey

0.00

1.02

1e8

0

50000

100000

150000

Krull

0.00

1.02

1e8

0

20000

40000

60000

Zaxxon

hyperbolic

euclidean

discrete

We first propose an isometry between the Gaussian manifold and the Poincaré disk model T P c →U c : P c →U c :

and the inverse is:

The translation of the Poincaré disk model can be derived using complex numbers. Let z = x + yi ∈ C and ( x, y ) ∈ P 1 and z 0 ∈ C be the pivot point. Then the isometry that moves z 0 to the origin is defined as T ( z ) : P 1 →P 1 := z -z 0 1 -¯ z 0 z . Note that the translation of Euclidean space is z -z 0 .

After transforming the latents on the Gaussian manifold to the Poincaré disk model and using the translation, we can measure how the latents well-captures the hierarchical structure of data. We first pick a latent and then translate all the latents by setting the selected latent as the pivot point. We then measure the Pearson correlation between the cumulative reward of the latents and the norm.

We repeat this process for all the latents and compute the maximum of the correlations. We use the latents obtained from the agents which recorded at least 250 for long enough trajectories. We obtain a correlation coefficient of 0.46 from the hyperbolic latent space, whereas the correlation coefficient of the Euclidean latent space is 0.40, showing the hyperbolic space better captures the hierarchy along the increasing norm.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a pseudo-Gaussian distribution over hyperbolic spaces, which is suitable for constructing VAEs with hyperbolic latent spaces. This pseudo-Gaussian distribution arrives by replacing the usual Euclidean distance with a KL-divergence (this locally approximates the hyperbolic geodesic distance). The required properties of this distribution (sampling, normalization constant, KL-divergences) are derived and shown to be easily computable. Empirical results show promise.\n\nSTRENGTHS:\n1. The paper is clearly written and easy to follow.\n2. The paper provides a novel (to the best of my knowledge) construction of a hyperbolic distribution that seemingly solves a practical problem (numerical instability) in hyperbolic VAEs.\n3. I find it elegant how the proposed distribution factorizes into a product of known distributions (this was not obvious to me from the KL construction).\n\nWEAKNESSES:\n1. A claimed benefit of the proposed method is that it is numerically more stable than previous methods. While I do not doubt the correctness of this statement, I had wished for an empirical demonstration of the claim. As far as I can tell, numerical stability is the main practical gain over existing methods (note, this is an important gain), so I think it should be studied more extensively experimentally. In short: show me numbers that measure stability.\n2. The paper mentions tuning the curvature of the hyperbolic metric, but I did not find it clear how and when such tuning takes place. Is the $c$ parameter tuned by hand?\n3. Similarly, the paper mentions using a $\\beta$ parameter in the VAE, but it is not clear how such a parameter is chosen in practice. Is it manually tuned? If so, does that not suggest overfitting could be a problem?\n4. Since the geodesic distance is replaced by a local approximation (in the form of the KL divergence), the method is only somewhat geometric. It seems to be good enough, but I suppose that none of the established theoretical results on tree embeddings hold under the approximated metric. If so, it would be good to explicate.\n\n*Minor comment:*\nLine 126: The KL divergence is really not a \"distance\" and should not be described as such. Especially in a geometric paper, such confusions should be avoided.\n\nQUESTIONS:\n1. The KL is a local approximation to the hyperbolic metric. This is used as a distance measure in a pseudo-Gaussian distribution. I wonder how similar this construction is to a tangent space normal distribution. Both approaches rely on local approximations of the metric. Do you have comments on this link?\n2. In Table 1, what does $\\circ$ and $\\times$ mean? I'm struggling to make sense of those symbols.\n\nLIMITATIONS:\nThe paper does not provide much of a discussion of the limitations of the proposed method.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes to utilize a set of Gaussian distributions as the latent space of a VAE model. One benefit of doing so is that this forms a hyperbolic space under the Fisher information metric. A key insight is that this Gaussian manifold corresponds to the Poincare half-plane model of hyperbolic geometry. Using this construction the authors construct a Gaussian manifold VAE which locally approximates the metric and allows for fast calculation of the KL divergence and fast sampling. Empirical results demonstrate the efficacy of the proposed approach on image datasets as well as RL benchmarks.\n\nSTRENGTHS:\nThis paper proposes an interesting way to enforce a hyperbolic latent space in a VAE setting. The connection of the FIM and Poincare half-plane using a set of univariate Gaussian distributions is interesting. The definition of the PGM distribution is certainly novel and I especially enjoyed the construction as it has the benefits of exact KL computation while enabling fast sampling. Finally, the overall flow and presentation of the paper at a high level are clear, but certain aspects could be improved to increase readability.\n\nWEAKNESSES:\nThere are a few weaknesses in the paper which I outline below:\n\n**Novelty:**\n\n1. While the use of the PGM distribution is novel in a VAE setting the use of the Poincare Half plane model due to its enhanced numerical stability properties is not. In fact, this paper does not cite a key work in this space \"Representing Hyperbolic Space Accurately using\nMulti-Component Floats\" Tao and De Sa 2021. I believe their usage of geometry shares all the same benefits of numerical stability (and perhaps even more) than the one presented here. I think this paper should attempt to implement a VAE with the Poincare Half-plane latent space as proposed in this paper for a fair comparison between approaches.\n\n2. The authors comment that current hyperbolic normal distributions suffer from a variety of drawbacks such as numerical instability, lack of a closed form KL, and high computational cost in sampling. However, these criticisms seem mostly towards the Poincare and Riemannian Normal distribution. The wrapped normal does not have high sampling costs for instance, so the statement should be refined a bit. Furthermore, other hyperbolic distributions can be defined on higher dimensional hyperbolic spaces without resorting to product spaces. This is a limitation of the Poincare half-plane model itself, so I feel there should be more commentary about this and how it may affect experiments.\n\n**Presentation weaknesses:**\n\n1. The preliminaries section could be significantly improved. The introduction to Hyperbolic geometry is non-standard and the authors could lean on prior work to improve the exposition.\n2. The definition of the Poincare half-plane model is insufficient here. You need to be much more concrete in mathematical terms on how it is defined. Things that one could add is what is the exact form of standard operations using this geometry, e.g. distance, exp map, log map, etc ... Furthermore, a visualization could improve readability.\n\n**Experimental weaknesses:**\n\n1. The experiments section has a few weaknesses. First, it seems that the bolded numbers in the tables have std's that overlap with other table entries. In this case, both table entries should be bolded.\n2. I believe it is important to test this approach on actual tree-like datasets rather than image datasets. These include the Branching Diffusion Process dataset found in Nagano et. al 2019, as well as the lobster graphs dataset in Bose et. al 2020. A hyperbolic VAE should show more visual benefits in these settings.\n3. Please include numerical error analysis in the same vein as Tao and De Sa 2021 Fig 1 and Fig 2 to substantiate the claims that the proposed approach has less numerical instability.\n\n\n\n\n\n**Minor:**\n\n1. The authors say \"the tree-structured data can be embedded with arbitrary low distortion in hyperbolic space\". This fact is only true for 2D hyperbolic spaces. The cited paper does not extend this to higher dimensions\n\nQUESTIONS:\n1. Past hyperbolic VAE's have made the curvature constant learnable which leads to improvements. Can this also be leveraged here?\n2. In Table 1. the authors mark the HWN as not geometry aware. Can the authors explain why this is the case? The HWN explicitly uses the metric in all of its computations.\n3. Similarly, the HWN has been used in VAE settings, in which the case KL divergence is needed and is not difficult to calculate. Can the authors explain why they feel this is not the case?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-30"}

comments: SUMMARY:
In this paper, the authors propose a novel Gaussian manifold generative autoencoder (GM-VAE) whose latent space is a statistical manifold formed by univariate Gaussian distributions. They also introduce a pseudo-Gaussian manifold normal distribution on the latent, which is easy to sample from and improves training stability due to a closed form of KL divergence. The experiments conducted show that GM-VAE outperforms the considered state-of-the-art baselines in the density estimation task on a few image datasets, and gives comparable results to DreamerV2 in the model-based RL task.

STRENGTHS:
1. The proposed GM-VAE model is a non-trivial improvement of the VAE architecture with a solid theoretical background.

2. Experimental results on the CUB, Food101, Oxford 102, and (partially) Atari2600 Brekout datasets demonstrate the superiority of GM-VAE over its competitors ($\mathcal{E}$-VAE, $\mathcal{L}$-VAE, and $\mathcal{P}$-VAE).

3. The proposed model has superior computational performance concerning the other hyperbolic VAEs considered.

4. The paper is well written.

WEAKNESSES:
1. The experimental results of the GM-VAE in the model-based RL task are only comparable to those of competitors.

2. Rather small-scale image datasets were used in experiments.

QUESTIONS:
1. It would be interesting (and thus influence my opinion) to see the results of experiments performed on large-scale datasets (like CelebA). Also, the FID score would be reported (besides NLL).

2. Minor comments:

Tab. 2 caption: N/A appears only in the standard deviation (or perhaps confidence interval?) indices.

LIMITATIONS:
Generally yes, but for clarity I would suggest summarizing the limitations in a separate paragraph.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nSome VAE methods have been proposed to model the latent space as a hyperbolic space for modeling the structured representation of data. This paper proposes a hyperbolic distribution by using KL divergence to measure the geodesic distribution of Poincare normal, by using KL divergence as the local approximation of the geodesic distance. It makes the proposed method easy to sample and improve the numerical stability. Experiments on image density estimation and model-based RL are conducted.\n\nSTRENGTHS:\n+ The proposed method is with a clear motivation and relies on a simple assumption (in Eq. (2)) to make the analytic operations, which is straightforward.\n+ Experiments are conducted on two different applications. \n+ As shown in the reported results, the proposed method can perform better than or comparable with the related works.\n\nWEAKNESSES:\n- It is not clearly studied and discussed how the proposed pseudo Gaussian manifold normal distribution can model the hyperbolic distribution. Specifically, it is unclear how the introduced approximation may influence the modeling of the manifold. Although tractable KL is favorable, how the approximation introduced by the proposed method may influence the modeling capability? \n- The authors may visualize the manifold of the proposed method and compare that with other methods, such as [27]. \n- The sampling results of the VAE models should be visualized and compared with previous methods, such as [27] and [25], especially considering that numerical improvements shown in tables are insignificant. \n- As the claimed main benefit of the proposed method is numerical stability, the analysis and validation of the stability should be strengthened. How and why the previous works may introduce instability? How can the stability of the methods be measured and validated with empirical studies, numerical evaluation, and visualization?\n\nQUESTIONS:\nPlease check other questions and suggestions pointed out with the weakness points.\n\nLIMITATIONS:\nThe limitations are not explicitly discussed. The authors may check the weak points and consider discussing the potential limitation of the proposed method, such as the potential influence on the modeling capability that may be caused by the approximation.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}