Hyperbolic VAE via Latent Gaussian Distributions
Authors: Seunghyuk Cho
Juyong Lee
Dongwoo Kim
Venue: NeurIPS 2023 poster (2023)
Keywords: Hyperbolic space
VAE
Distribution on hyperbolic space
Hierarchical representation learning
Reinforcement Learning
TL;DR: We propose a hyperbolic VAE whose latent space consists of Gaussian distributions and provides strong numerical stability, addressing a common limitation reported in previous hyperbolic-VAEs.

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
Seunghyuk Cho POSTECH GSAI
Juyong Lee
Dongwoo Kim POSTECH GSAI & CSED
KAIST Al
shhj1198@postech. ac..cch......................... ac.kr kr
agi wss.kiistts.cc... ac.
dongwoo kim@postech. ac kr
Abstract
We propose Gaussian manifold variational auto-encoder (GM-VAE) whose latent space consists of a set of Gaussian distributions. is known that the set the uni- variate Gaussian distributions with the Fisher information metric form hyperbolic space, which we call a Gaussian manifold. To learn the VAE endowed with the Gaussian manifolds, we propose Speudo-Gaussian manifold normal distribution based on the Kullback-Leible divergence, a local approximation Fisher-Rao distance, define 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 common limitation reported in previous hyperbolic-VAN The implementation is available at https: /tt://gttttt github.... com/ml- postech/tt. VAE..
the squared
Introduction
The geometry of latent space in generative models, such variational auto-encoders (VAE) (Kingma & Welling, 2013), reflects the structure of the data representations. Mathieu et al. (2019); Nagano al. (2019); Cho 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 (201 ); 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 distri- butions. Unlike the Gaussian distribution Euclidean space, however, these distributions suffer from several shortcomings, including the absence of closed-form Kullback-Lein (KL) divergence, numerical instability (Mathieu et al., 2019; Skopek al., 2019), and high computational cost in sampling (Mathieu et al., 2019).
form Riemannian manifold from the of univariate Gaussian distributions by equipping the Fisher information metric (FIM). It is known that the FIM of univariate Gaussian distributions is akin 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, 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 new distribution with better analytic properties.
Meanwhile,
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 hyperbolic distribution by
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

p(a)
(01.01)
(o2) m2))
(Ng. as)
(a) Gaussian manifold
(b) Univariate Gaussians
Figure 1: (a) The visualization of the Gaussian manifold consisting of set of Gaussian distributions.
Each point of the Gaussian manifold pair of two parameters univariate Gaussian distribution: (u, .) E IR x R The dashed lines are the geodesies, which are either the ellipses with eccentricity 1/v2 with the origin placed on the jı-axis or straight lines parallel to the o-axís. (b) Three univariate Gaussian distributions correspond three points the Gaussian manifold (a).
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 hyperbolic VAF, named Gaussian manifold VAE (GM-VAE), whose latent space is 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 reinforeement learning (RL) verifies the feasibility of using hyperbolic space another domain of task.
We summarize our contributions follows:
We introduce a variant of VAE whose latent space defined on a statistical manifold formed by univariate Gaussian distributions, namely Gaussian manifold.
We propose new distribution, called a pseudo Gaussian manifold normal distribution, which easy to sample and has closed-form KL divergence, to train VAE the Gaussian manifold.
We empirically verify that the newly proposed VAE performs stable training without numer- ieal issues 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 RI. Specifically, we replace the latent space of the world model with hyperbolic space for learning environments, showing competitive results with a state-of-the-ar baseline.
Preliminaries
In this section, 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-dimeensionl Riemannian manifold consists of manifold M and metric tensor M paxx, which is smooth map from each point x M to symmetric

---

positive definite matrix. The metric tensor g(x) defines the inner product of two tangent vectors for each point of the manifold )x TxM T&M R. where TxM the tangent space of x.
The metric tensor induces basic Riemannian operations, such as a geodesic, exponential map, log map, and parallel transport. Given points x, E M, geodesic 7 [0, 1] + M is unit speed curve on MM being the shortest path between (0) x and 7(1) y. This curve can be interpreted as generalized path of a straight line Euclidean space. The exponential map expx : Tx + .M is defined as expx (v) = 71) = y given is geodesic starting from (() = x and > (0) = v E TxM The log map logx M Tx M the inverse of the exponential map, i.e., logx (expx (v)) v. The
parallel transport PT T + M moves the tangent vector v along the geodesic between and y. The geodesic distance d. y) can be induced by the metric tensor as follows:
(x J.' V /A((0), (()))))
Hyperbolic space. One method of classifying Riemannian manifolds, a basic question of differen- tial geometry, is based on curvature. Among different types of curvatures, one popular curvature is the sectional curvature Kg. which is generalization of the Gaussian curvature in classical surface geometry. Given two linearly independent vector fields X, Y E ((M), the sectional curvature Kig can be computed with Riemannian curvature tensor R: X(M) x X(M) x X(M) - X(M) as below:
(R(X, X), Kig XX,X) - (X,
where R(X, Y returns tensor field assigning a tensor to each point of the Riemannian manifold M. The hyperbolic space is Riemannian manifold that has the sectional curvature value of constant negative (Nickel Kiela, 2018). The hyperbolic space known to be able embed tree-structured data with arbitrarily low distortion (Sarkar, 2011 1).
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 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 (Mathicu 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 an alteration of the Poincaré disk model (Nickel & Kiela, 2018; Nagano et al., 2019; Bose al., 2020; Cho al., 2022). The Lorentz model uses half hyperboloid manifold, where the closed form of the Riemannian operations exists, so the numerical stability issue of employing the Poincaré disk model relieved.
The Poincaré half-plane model is another well-known model of hyperbolic space with an open half-plane manifold. The metric tensor of point of the two-dimensional Poincaré half-plane model (z, y) is 2diag(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 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 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é
3

---

disk model and the log mapping of the Lorentz model are unstable Riemannian operations, which are written as:
xx y|1 x (1 [lyl2)
cosh (a) Va²
dp(x, cosh -1 1
logu(v):
ou);
where = Vo -i=1 U;V; is the Lorentzian inner product of For the Poincaré disk model, when the points x, y are near the boundary the unit disk, the floating point representation of the norm values |xx|? ||yll becomes onc. The denominator of the geodesic distance then becomes zero. For the Lorentz model, = v and u contains large values the coordinates, a becomes less than one which results in NaN because the domain cosh - (x) is x 2
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 pe X + R. where 9 E .M, equipped with the Fisher information metric (FIM) forms a Riemannian manifold (Rao, 1992). The FIM is defined as:
96g 955(0) = /
b log (x) à log po( (2) dr. 00,
00;
In the parameter space univariate Gaussian distributions {(u, o)) μ E R, e Rz0), the FIM can be simplified as two-dimensional diagonal matrix o diagg 2) (Costa et al., 2015).
Connection to the Poincaré half-plane model. The diagonal form of the FIM implies that the Riemannian manifold with {(u o) μ E R, a E |> has the same set points as the manifold of the Poincaré half-plane, but with different curvature value of -055.
The parameter space of the re-dimensional diagonal Gaussian distributions becomes the product of 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 X For example, an exponential map applied on point (p: )= ey E M, for each i {I, n),
i= MM, are defined manifold-wist
Mi, with tangent vector be represented (expp. (v:))
Distance in the statistical manifold. In the statistical manifold, distance functions measure the difference between two distributions 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 DKL g]
Je p(x) log q(æ)
p(x) dæ for distributions p, q in the same statistical manifold. One notable property of KI. divergence is that it can locally approximate the squared Fisher-Rao distance (Tifrea et al., 2018):
DkL(p(:0 + d0) p(:0)) = Σ (i + (()
8
3 Method
In this section, first present the concept of the Gaussian manifold, which can have an arbitrary curvature by reparameterizing univariate Gaussian distribution. We then propose pseudo Gaussian manifold normal distribution. Finally, we suggest a new variant the VAE defined over the Gaussian manifold with PGM normal prior. We denote the density function of the Gaussian distribution as N(ip,, (/(V)T exp (-(μ- z)² /(20)))

---

3.1 Gaussian manifold with arbitrary curvature
Previous studies on hyperbolic space emphasize the importance of having an arbitrary curva- ture (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, shown in Section 2.2, the univariate Gaussian distributions form 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 M((22c,, o²) be the reparameterized univariate Gaussian distribution with additional parameter reparameterization leads to the FIM of o 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.
0. The
We call the statistical manifold with the reparameterized univariate Gaussian distributions and the extended FIM as the Gaussian manifold and denote as Go, where -c is the curvature of the Gaussian manifold.
We then verify that the KL divergence between the distributions of the Gaussian manifold approxi- mates the geodesic distance, even in the presence of arbitrary curvature in the Gaussian manifold. Let (u, a)) E Ge an arbitrary point of the Gaussian manifold. The KL divergence between (μ, σr) and its neighbor (μ + du, da) can be computed as:
DRt. + dp), (a + do)))) o²))
du 2 do
-
0 oo?
du do
+ ((do)
2c
(1)
where the first term is the squared Riemannian norm of the tangent vector (du, da) approximating the squared Fisher-Rao distance. The detailed derivation of the KI. 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 (μ, E Ge be 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:
(- DKi u, 0.) || (V(2....,,3)) a, (2,3)) (
Kc(,,o;a,3,1)) B(c,B,n)
x exp
(2)
2c 20
where (a, 3) E Go. and y E lR50 are the parameters of the distribution. The distribution is centered at (a, 3) with additional scale parameter 7, We verify the convergence of the PGM normal and compute the normalizing constant (c,, B, 7) over the probability measure of the Gaussian manifold at Appendix C.I. As shown in Equation 1, the KL divergence of the Gaussian manifold approximates the Fisher-Rao Rao distance between (v2c 2c a, (8) and N μ., o²). 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:
Kolls, a; 3, (del(g) = Nu;; a, 82g2). 2o Gamma (0); 4c
4c,32n22
where Gamma(z; a, b) [(a) =a-1exxp exp (-bz) and is the FIM of the Gaussian manifold. The /dtt g term enables us sample and compute the KL divergence in Euclidean manner. Thanks to the properties of Gaussian and gamma distributions, the PGM normal is easy to sample and has closed-form KL divergence. The detailed derivation 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.
5

---

We note that the PGM normal can be easily extended for the diagonal Gaussian manifold, 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 {(1s, E) la E IR, Σ 1R as:
IE (xx det(g)
[logpo(x 14, E) DKL (ns(1, Σ x) Vdel(g) II p(1,,) E) vdel(g))
where (x μ. Σ) is the decoder network, qa (e. E x) is the encoder network and p(p,, Σ) is the prior. The variational distribution is set (p. Σ x) (x) &&
(x), %x))), where ag (x)
I" and (x). Tx E R55. and the prior is set to p( p. Σ) = (0, I, in our experiments given curvature c. The pseudo-algorithm for the decoder of GM-VAE present 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 investi- gated 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 VAE 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 VAE decoder modeled as probability distributions. Arvanitidis et al. (2021) suggest method of using the pull-back metric defined with arbitrary decoders on the latent space. Our work focuses more 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 al., 2019; Nagano et al., 2019; Cho 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 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 novel viewpoint 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 exchange for high computational complexity.
Although these distributions are widely adopted in many applications (Skopek et al.. 2019; Mathieu & Nickel, 2020; Cho al., 2022), one can barely adopt the full covariance matrix due to the difficulties
Algorithm Decoder
Input Parameter (a, 3) E Go, Decoding layers Dec(-) Output Reconstruction x'
1: Sample μ (2,,3222))
2: Sample log 2 ~ Gamma Jen + 1, J0
3: x'= =Dee([[, =Dec([u,logoo")
4: return x'

---

Table I: Density estimation real-world datasets. denotes the latent dimension. We report the negative test log-likelihoods of average 10 runs for Breakout, CUB, Food101, and Oxford102 with 95% confidence interval. N/A in the log- likeliheeoodd indicates that the results not available due to the failure of all runs, and N/A in the standard deviation indicates the results are not available due to failures of some runs. The best results are bolded.
GM-VAE (c 1)
GM-VAE (c-
GM-VAE (c 3/2)
E-VAE
L-VAE
P-VAE
2
124.74
122.58NAA
270.05
121.5211.00
122.64.1.13
122.47
Breakout
4
66.39 10.76
66. 66.7010.32
271.73
66.83_0.49 0.44
66.39
65.8020.49
8
10.97
45.2510.27
81.55164.61
45.14
45.31
45.3610044
992. 192.05+1.55
993.034...c
990.4942.22
985.46=3.82
086.27+3.81
979. 179.74
CUB
969.99 13.13
968. 39.70
964.0223..5
958.00 13.22
960.88.13.
956.77
70
949. 1419113
948.88
4.440
039.088a..1
942.3443.44
937. 937.15
50
1297. 81 14.451
1298. 46.6632
1293. 1293.26
1286.30+6.19
1209. 1209.58
1290.51
18.0. 18.23
Food101
1224. 03..31
1227. 16
1218.09 33.88
1213.31 13...113.88 =3.88
1216.63
1207 23.
1164.95 11.4.9543.80
1165.3915.54
1165.91 4.9.1
1152.80
1160.97 1:.11
1149.5613.41
50
1297.
41
1296. 1.56
1294.12
1292.90=3.6
1289.43
43.244
1289.99
Oxford102
60
1253.80
12.55
1256. 12.99
1251.77
1245,49+2.18
1248. 1148.7211 1..11.2
1247. 12.551
70
1231 1231.52440 1.241.2
1229.38
1219.75+1.72
1215.07=2.02
1218.5443.85
1214.85
in Monte-Cari based KL approximation. The number of samples to approximate the KL divergence increases exponentially when the full covariance matrix is used (Cho al., 2022), so 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.
RL with hy 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 VAEs. 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 RI. 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- MM-VAE shown in the experiments with additional analyses.
5.1 Density estimation on image datasets
We conduct density estimation on image datasets 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 (0xford102) (Nilsback & Zisserman, 2008), Food101 (Bossard et al., 2014), and Caltech-UCSD Birds- 200-2011 (CUB) (Wah et al., 2011). The datasets chosen with the four lowest d-hyperbolicity (6-H), a metric that measures how the given images are well-embed in hyperbolic space. Low 6-H implies that the dataset is likely to embed in hyperbolic space. The details about ô-H are available in Appendix D. The values of 5-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 implicit hierarchical structure (Nagano et al., 2019; Li et al., 2019; Bossard et al., 2014; Kerdels & Peters, 2015).
We compare GM-VAR with the three baseline models: VAE with Fuclidean latent space (E- (ABV,)), and hyperbolic VAB equipped with HWN (C-VAE) and Poincaré normal (P-VAE). We use the product
7

---

latent space for both C- and P- P-E, and set the curvature value to -1. The other details on the implementation and experimental setups are described in Appendix E.
The results are reported at Table 1. GM-VAE outperforms the baselines in all the settings, except one case of Breakout. Especially in CUB and Oxford102, GM-VAR outperforms the baselines regardless of the curvature value. In Breakout, P-VAE shows inferior performance due to unstable training, and C-VAE fails in some of the runs with small latent dimension. The results of P-VAE and C-VAE with non-product latent space, 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-VAE on various tasks. Specifically, we apply GM-VAE to world model, which aims to learn the representation of the environments (Ha & Schmidhuber, 2018; Hafner et al., b, 2020). We use Dreamer (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 Pa (zz, (4, 22), the reward ( [/,, 21), and the discount factor P [b2, 2,), where x is the observation which the format is the image, Tr is the reward, "ee is the discounting factor, h& is the deterministic recurrent state, and 2 is the stochastic state. The model is trained by maximizing the likelihood of p(x, r,y a) given observations x, rewards r, and discount factors Y 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 go (z (4,, mex) with the following objective [(0, 0) as:
T
Σ log Po To, Tolle, =)) + 3 DEL (=0(=e))e) Poleelhall)
C(0, 0) =
where B is KL. loss scaling factor, T is the length of input sequence, Pe is the prior, and gs is the approximated posterior. GM-VAE is employed by replacing the space of 2, with the Gaussian manifold and two components in RSSM, the representation model ( (e) and transition predictor Pa (ze (hg), 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 5-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 F.2 with the 5-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, GME 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 C-VAE and P-VAE. During the experiments, L- C-V 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 sampling-based based KL divergence computation can be unstable.
We first show that sampling from PGM normal can be stabilized via simple reparameterization trick. To train GM-VAE, one needs to obtain sample μ and from PGN normal Ca B, 7). Sampling μ can be done from (aa, 32n²) without numerical issues. Sampling a can be done from Gamma(a. (a, b) where a = 1/4e72
1/4c:2222 as shown in Appendix C.2. However, due to machine precision error, often, 3 violates the manifold constraints, i.e., 3 O. Eventually, direct sampling of a can cause the numerical instability. To avoid 3 being zero, we use the output of the VAE encoder as log 32 whose value ranges over the entire real numbers and is more stable even when 3 is close to
8

---

Latent space
Euc.
Disc.
Hyp.
8-H H
Breakout
329.0
256.8
319.3
0.12
Alien
3412.5
3120.0
3485.0
0.14
Zaxxon
34275
38825
38950
0.14
Ice Hockey
25.50
11.80
20.75
0.14
Freeway
32.8
33.0
33.0
0.38
Krull
53290
36135
66185
0.38
(a) Latent space analysis
(b) Results
Figure 2: The results of model-based RL experiment. (a) The dots from yellow to purple represent the latent states from the world model in the Atari2600 Breakout with decreasing rewards. Along the red geodesic dashed line passing, we sample for images to visualize the learned representations. As the sample shows, we can observe hierarchical structure at different stages of the game along the geodesic. (b) We compare the methods of using Euclidean, discrete, and hyperbolic latent space. We report averaged rewards over four runs and bold the best reward.
zero. With log 3², instead of sampling o². we sample log from Gamma a, 1), through the reparameterization of the Gamma distribution, where log can be directly computed from log 3
o² = log + log b, where e is sampled
We can show that the KL divergence between arbitrary PGM normal and prior distribution (O, I, 1) has a closed-form solution without any sampling. The KL divergence of PGM normal is the sum of the KL divergences between tions, as shown in Appendix C.3. First, the KL. divergence between a univariate Gaussian distribution N(,, d²) and the prior distribution can be obtained with log o² shown in Equation 4. Second, the KL divergence between the two Gamma distributions, Gamma(as (01, b1) and Gamma(a2, b2), written as:
Gaussian distributions and between two Gamma distribu-
b
I(a1)
+ +(a - a2))((a))) (1 (1- b1
b2
DKL (oam@aaaa(a (1. b1) || Gamma (22 b2)) = log
In
b2
I(a2)
where &i is the digamma function, can be computed using log
Training time comparison. Common bottle- necks of the mode hyperbolic VAEs arise from the complex manifold constraints and the dif- ficulty of sampling from the hyperbolic distri- butions. For example, the Poincaré disk model of P-VAE and the Lorentz model of C-VAB re- quires the samples to be inside of a unit disk hyperholoid with constraint (x E [R=+1 -xy + Σ Such manifolds need complex transformations, e.g., clipping, projection, or geometric transfor- mations using the Riemannian operations, to 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 a.
Table 2: The training time of the VAEs in density estimation of the Breakout image dataset. We re- port the training time of the VAEs in seconds per epoch. GM-VAE 1.93x faster than P-VA PAVAE and 1.41x faster than L- C-VAE in the experiments held on a single A100 40GB PCI GPU.
and to be
1= x. respectively.
E-VAE
C-VAE
P-VAE
GM-VAE
24.5
35.9
49.2
25.5
We report the time consumptions of the VAEs with the latent dimension of per epoch in the density estimation of Breakout at Table 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 Figure for qualitative analysis. We take the world model with GM-VA 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 well captured in the latent space.
Note that in the Poincaré disk model, the depth of the hierarchy is expected 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, 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 Appendix G.2.
Conclusion & Future Work
In this work, we propose novel method of representation learning with GM-VAE, utilizing the Gaussian manifold for the latent space. With the newly-proposee 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 model-based RL compared to baselines. We explain the behavior of GM-VAE 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- hong-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 (ITP) grant funded by the Korea govermment (MSIT) (No.2019-0-01906 Artificial In- telligence Graduate School Program (POSTECH) and National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (No. RS-2023-00217288 and National Research Foun- dation of Korea (NRF) grant funded by the Korea government (MSIT) (NRF-20218) (C1C1011375)
10

---

References
Arvanitidis, G., González-Duque, M., Pouplin, A., Kalatzis, D., and Hauberg, S. Pulling back information geometry. arXiv preprint arXiv:2106.05330) 05367, 2021.
Bay, A. and Sengupta, B. 2ee: Information geometric sequence-to-sequenee networks. arXiv preprint arXiv: 17X0000/1 2017.
Bose, J., Smofsky, A., Liao, R., Panangaden, P., and Hamilton, W. Latent variable modelling with hyperbolic normalizing flows. In International Conference 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. Hyperbolic deep reinforcement learning. arXiv preprint arXiv:2210.0154: 2022.
Chen, W., Han, X., Lin, Y., Zhao, H., Liu, Z., Li, P., Sun, M., and Zhou, Fully hyperbolic neural networks. arXiv preprint arXiv:2105. J4686, 2021.
Child, R. Very deep VAEs generalize autoregressive models and outperform them 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, 69, 2015.
Davidson, T. R., Falorsi, L., De Cao, N., Kipf. T., and Tomezak, 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., 17666))576-779 jun 2015. ISSN 0020-0190. doi: 10.1016fj.p1..015.02. 002.
Gogianu, F., Berariu, T., Busoniu, L., and Burceanu, E. Atari agents, 2022. URL http:::::///iithub com/flor hwtp/fooringgiannn www.niinnggnna/aoooiiooiaooaanaa
Gomes, E. D. C., Alberge, F., Duhamel, P., and Piantanida, P. Igeood: An information geometry
approach out-of-distribution detection. arXiv preprint w0Xiv:2203.07798, 2022.
Ha, D. and Schmidhuber, J. World models. arXiv preprint arXiv: 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. International conference on machine learning, pp.
PMLR, 2019b.
Hafner, D., Lillicrap. T., Norouzi, M., and Ba, Mastering atari with discrete world models. arXiv preprint 2rXiv:2010.021111 2020.
Han, T., Zhang, J., and Wu, Y. From em-projections variational auto-encoder. 2020.
Karakida, R., Akaho, S., and Amari, S.-i. Universal statistics of fisher information in deep neural networks: Mean field approach. The 22nd International Conference on Artificial Intelligence and Statistics, pp. 1032-1041. PMLR, 2019.
Kerdels, and Peters, 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. and Welling, M. Auto-encoding variational bayes. arXiv preprint arXiv:1312.614 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., Ried- miller, M., Fidjeland, A. K.. Ostrovski, G., et al. Human-level control through deep reinforcement learning, nature, 51807540 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. 4693- 4702. 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 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 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): 2006.
Radford, A., Metz, L., and Chintala, S. Unsupervised representation learning with deep convolutional generative ad adversarial networks. arXiv preprint arXiv:1511.06444 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 1-611--9919-- doi: 10.1007/ 978- 778-1-4612-0919-5166
Sarkar, R. Low distortion delaunay embedding of trees hyperbolic plane, In Proceedings of the 19th International Conference Graph Drawing, GD'11, pp. 355-366, Berlin, Heidelberg, 2011. Springer-Verttee ISBN 9783642258770. doi:
Shi, L. X., Lim, J. J., and Lee, Y. Skill-based model-based reinforcement learning. In Conference Robot Learning, 2022.
Skopek, O., Ganea, O.-E., and Bécigneul, G. Mixed-curvateee variational autoencoders. arXiv 2019.
preprint
Tifrea, A., Bécigneul, G., and Ganea, O.-E. Poincar 'e glove: Hyperbolic word embeddings. arXiv preprint arxiv: 181000661, 2018.
Wah, C., Branson, S., Welinder, P., Perona, P., and Belongie, S. dataset. Technical Report CNS-TR-2011-00 Californi Institute of Technology, 2011.
The caltech-uesd
Xu, J. and Durrett, G. Spherical latent spaces for stable variational autoencoders. arXiv preprint 18Xiv: 10805, 2018.
Yu, T. and Sa, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Numerically accurate hyperbolic embeddings using tiling tiling-based models. In
Yu, T. and Sa, C. Representing hyperbolic space accurately using multi-component floats. In
Advances in Neural Information Processing Systems, 1.2
12

---

A Machine Precision Error Analysis on Hyperbolic Space
10t0
10th
The Poincare half- plane model
Tao and Sa (2021), M-2
10
The Lorentz model The Poincare disk model
105
100
10"5
80
origin
Figure 3: The machine precision of the hyperbolic models. For the Poincare half-plane model, we report the upper bound of the machine precision enor derived by Yu & Sa (2021). For the Lorentz model, we report the eтrоr between the Lorentzian product of the point and -1, which is the manifold constraint of the Lorentz model. For the Poincare disk model, we report the threshold that the norm of the point becomes one. The precision is to 32 bits. The analysis reveals that all the three models suffer from numerical stability issue with points which the distance between the origin is farther than 20.
B
Gaussian Manifold
B.1 Curvature of the Gaussian manifold
We construct a Riemannian manifold {(H, o) E R, a E R>o} with positive constant and the metric tensor 1dia) which we will name Gaussian manifold. We need to show the value of the curvature.
First, need to compute the Christoeffel symbols of the Gaussian manifold defined as:
1ggt 0ga
Ogi
2
0g: 0g;
0g.
where 95z is the (i,j) element of the metric tensor and 9Dو is the (i,j) element of the inverse of the metric tensor.
The Christoeffel symbols of the Gaussian manifold are:
r
Then, the sectional curvature of the space Kg with given tangent vectors du, du computed as: (R(du, do doz, dp)
Ki
det.
@m 22
01
12
DP PM
1
Im 2p
del
91m
0p
1p
13

---

Note that the sectional curvature of two-dimensional Riemannian manifold is same as the Gaussian curvature where (dp, du) (da, do) (dd, do) det g.
B.2
Gaussian manifold with KL--ivergence
Between two univariate Gaussian distributions (N,,, and (12,, o3), compute the KI. divergence as:
DKL II (V2,, 03))
o of (y) (22 log of
-1 -1)
(4)
o
We extend the KL divergence for an arbitrary curvature of the Gaussian manifold as:
D Dii (11,,, (12, ( || (VV2V222, o5))
2c
Now, we show that the extended KI. divergence still approximates the Riemannian distance of manifold as:
a²
(a + do)² + 2c(d1)2
Di ((p + di, a + do), (d,)(u,))) (1,
log
+
-11
2 2c
(a + do)2
1+ do
20da + (do)
+ 20(du)2
-2 log
+
2 2c
σ²
o²
do
(do)²
20do (do)²
+ 2c(du)²
+ ((do)))
2 2c
20
0²
o²
1 du
T
dp
((do))3)
2
do
do
oa?
C Pseudo Gaussian Manifold Normal Distribution
In this section, we derive the normalizing constant 2(c,8,7) and the factorization of the PGM normal which the density function is defined as:
D& ((4, σ), (a, B))
K,
7(c,B,7)
exp
72
C.1 Normalizing Constant
The given probability density function needs to satisfy the following condition:
[ 8. d(p,o) 1,
(5)
where Vdet g d(u,o) )) is the probability measure over the Gaussian manifold induced from the Lebesgue measure d(p, o). The normalizing factor (cc, 3, 7) can be computed using the condition
14

---

Equation 5 as:
100
00
D Kki (),
2(c,8,7)
exp
du
-00
72
Voo2
1
40y? 1o0 a 0.20))
un (
(1)
(- o²
Vc
2072 exp
exp
4c32g3
do
(11 a)²
x
exp
25222 df
1
sey
2 3d
exp
r
2/0
4072
1072
4072
(%
x
Gamma
0²;
40y2
1,
4c82222
do²)
(6)
A0y
2/c
exp
4
4cy2
4
Finally, the logarithm of the normalizing factor is computed as:
1
log 2(c, 8,2)
log (2g(2))+3log 3 log coggcoo222 2+ log 2o +log-
(1 ++og(Ae777)
2
2
4coy "oo7"
C.2 Sampling
Suppose that o) = Kelu, a; V. g. Sampling μ and from the probability distribution p(u. o) can be done by sampling from the marginal distribution p(μ) and then sampling a from the conditional distribution p(oll)). The marginal distribution p(μ) can be derived from Equation 6
p(p)
o(H,o)do
0
Ko det da
-1% Gamma (0?
1 4og2
do² a,
40B222
a $2n2).
Equation 6 also implies that μ and are independent in the aspect of p(u, o)) so the conditional distribution p(a (c) is identical the marginal distribution p(o). The marginal distribution p(a) is computed as:
p(o)
plu, o) du
a a; B,>) / du
4c8273) )) ( ((2),,
20 Gamma
- 4072
02: 4072
1
20 Gamma
+ 1, 4c3222
Here, sampling from p(a) can be easily replaced by the procedure of sampling o² from p(o²).
which is identical to p(a)/(2a) Gamma the sample 2 to get
(g2: ăoyz
+ 1, 1033.² and then applying square root to
15

---

C.3 KL Divergence
Suppose that p(p, σ), g(u, σ) are two different PGM normal multplied with /dt g. As shown in Appendix C.2, and a are independent so the KL divergence between p.g is same as (pL(p())ll(u)) + (pp())))l(()))) The first term is the KL divergence between normal distribution. The second term is same with DRL (p(a")llg((²) due to the change-of-variable for- mula, so it is the KI. divergence between two gamma distribution.
16

---

D 8-Hyperbolicity
In this section, we explain about &-hyperbolicity (6-H) and show the 6-H values of the candidate datasets of our experiments.
D.I d-hyperbolicity
Given metric space (X, d), the metric space is said be o-hyperbolic if, for any geodesic trian- gle, i.e., a triangle where each side is geodesic curve, any point on the side of the geodesic tri- angle 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 2, y/ E X is given as:
A
6 = 0
D
2), (d(x. v) + d(x,2) - d(y, 2)), 2
(a) 6-H 0
(b) 8-H =0
then the following inequality holds for any four points z, y, 2, E X, we call the metric space 5-hyperbolic:
Figure 4: The illustration of 5-H for given geodesic triangles. The blue lines denote the geodesic curves between the points. (a) When the side AB contained in the union of the two 6-neighbor regions of the side AC and BC. we say that the geodesic triangle is 5-hyperbolic. 5-H of the geodesic triangle is then defined as the minimum possible value of (b) Any tree-structured trian- gle, where the geodesic curves correspond to the tree edges, is O-H. For example, given geodesic triangle ABC, one side AB is already occupied by the two sides AB and BC, as the geodesic be- tween the points A and B is the union of edge AD and BD being occupied by other counterparts.
z) 2 min((2, w) ()=)) -
(7)
5-hyperbolicity (6-H) is defined to be the mini- mum value of 5 satisfying Equation 7 and is used measurement quantifying how given met- ric space well embeds in hyperbolic Khrulkov al. (2020); Cetin et al. (2022). Figure illus- trates the concept of ô-H. We note that the lower 5-H the metric space has, the less deviation from the exact hyperbolic space is.
Computation. We measure the 5-H values of 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ı min:y (Dx, Dus)) - D.
To reduce the scale difference between the datasets, 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 5-H of randomly sampled 1,000 images from X. We repeat this 10 times and report the average 6-H.
report the value 28(X) / /diam (X) where
D.2 Image datasets
We measure the 5-H of the image datasets using an ImageNet pre-trained Inception V3 as a feature extractor.
Breakout
CUB
Food101
Oxford102
CIFAR-10
SVHN
CelebA
0.124
0.223
0.233
0.238
0.248
0.283
0.287
D.3 Atari2600 environments
We collect the Atari2600 images using the pre-trained agents of Gogianu et al. (2022) and measure the 5-H using the image encoder from the agents. For each environment, we report the 8-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.
17

---

Game
6-H
Reward
Game
6-H
Reward
Breakout
0.12
340
AirRaid
0.25
11729
Alien
0.14
6855
Frostbite
0.25
8653
Zaxxon
0.14
13100
Pooyan
0.25
7973
IceHockey
0.14
3
ChopperCommand
0.26
10530
Gravitar
0.17
1004
KungFuMaster
0.26
27733
Carnival
0.17
5605
YarsRevenge
0.26
56405
RoadRunner
0.18
60050
Phoenix
0.27
30208
Pong
0.18
21
JocaeyEscape
0.27
-707
Tutankham
0.18
236
Spacelnvaders
0.27
15623
Boxing
0.19
99
Hero
0.28
28592
Solaris
0.19
1727
Enduro
0.28
2089
WizardO/Wor
0.20
17217
Assault
0.28
3217
Seaquest
0.20
29745
Venture
0.29
1529
Gopher
0.20
24173
StarGunner
0.29
68417
Centipede
0.20
4663
Atlantis
0.29
919750
PrivateEye
0.21
195
DoubleDunk
0.30
18
Pitfall
0.21
0
TimePilot
0.30
10914
ElevatorAction
0.21
64917
BeamRider
0.30
7069
Robotank
0.21
75
BattleZZne
0.30
40333
MsPacman
0.21
5455
NameThisGame
0.31
12925
DemonAttack
0.22
23876
Skinng
0.31
-10021
Asterix
0.22
26792
Amidar
0.32
2826
Qbert
0.22
17364
Asteroids
0.34
1405
Jamesbond
0.22
886
MontezumaRevenge
0.34
VideoPinball
0.23
632563
Bowling
0.35
49
UpNDown
0.23
27030
BankHeist
0.36
1563
FishingDerby
0.23
54
Kangaroo
0.37
14283
Tennis
0,23
22
CrazyClimber
0.37
132517
Berzerk
0.24
904
Freeway
0.38
34
Riverraid
0.24
15103
Krull
0.38
9016
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 al., 2019), Ox- ford102 (Nilsback & Zisserman, 2008), and CUB (Wah et al., 2). The images of Breakout collected from plays with ore-trained Deep Q-Network (Mnih et al., 2015). The size of images of all the datasets is resized to 64 x 64, while Breakout 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.
Split
Breakout
CUB
Food101
Oxford102
Train
80,000
4,795
6,000
5,120
Validation
9,503
1,199
1,000
1,228
Test
9,934
5,794
1,000
1,025
18

---

We design the encoder and decoder similar to the generator and discriminator of DCGAN (Radford al., 2015). The details of the architecture are Table 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- aggliiklihood 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.
Table 3: The architectures of encoder and decoder used in the density estimation experiments. n. is number of channels of the image, n. the latent dimension. nt, is a coefficient that depends on the VAE, i.e., aa is 2, 2, 1.5, 1.5 for E-VAE, L-VAE, P-VAE, and GM-VAE, respectively.
Encoder
Decoder
Layer
Size
Layer
Size
Input
64 x 64 x nc
Input
1 x nd TranspssdCConvolttion22
Convolution2D
32 x 32 x 32
4. x 4 256
LeakyReLU
ReLU
Convolution2D
16 16 x 64
TranspsedCConvolution22
8 X 8 128
LeakyReLU
ReLU TranspssedConvolutionnD
Convolution2D
8 x 8 X 128
16 x 16 x 64
LeakyReLU
ReLU
Convolution2D
4 4 4 256
TransoscaCConvolttion22
32 32 32
LeakyReLU
ReLU
Linear
na 1ld
TransposeCConvvoutioo22
64 64 me
E.2 Model-based RL
We use the official TensorFlow implementation from Dreamerv2 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 zt with the Gaussian manifold and two components in RSSM, the representation model (z 25g and transition predictor (2 [he), with PGM normal. The byperparameters are set to be the same 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.
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 reveals that the non-product hyperbolic VAEs fail in most of the settings and the product hyperbolic space makes the hyperbolic VAEs much more stable.
G Model-Based RL
G.I Learning curves
G.2 Latent space analysis
We conduct analysis of the latent space of the agents learned 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.
19

---

Table 4: Density estimation results of non con-product hyperbolic VAEs. d denotes the latent dimension. N/A in the logeliicl indicates that the results are available due the failure of all runs.
d
L- C-VAE
P-VAE
2
124.24-1.66
266.8646.01
Breakout
4
66. 66.20=0.11
N/A
8
44.76 -0.48
N/A
50
N/A
N/A
CUB
60
N/A
N/A
70
N/A
N/A
50
N/A
N/A
Oxford102
60
N/A
N/A
70
N/A
N/A
Allen
Breakout
Freeway
10000
30
7500
20-
5000
2500
1.0
0
0.00
1:02
0.00
1.02
Hockey
Krull
Zaxxon
150000
100000
40000
20000
0.00
1.02
0.00
1e8
hyperbalic
euclidean
discrete
We first propose an isometry between the Gaussian manifold and the Poincaré disk model Tp.
Pa Ua:
-2y
1 1 1- 222 Dc
(x, (Teex - 1)2 3 y²c²
(Ver - 1)22 y²c
and the inverse is:
Tg (2.1) ( /Vox² + - cx²
-2r
cx2 + (y 1)3
The translation of the Poincaré disk model can be derived using complex numbers. Let = x+ E C and (xr,)) E P and 20 C be the pivot point. Then the isometry that moves zo the origin is defined as T(z) Pi P1
2-20 Note that the translation of Euclidean space T-2カ2
20-
After transforming the latents on the Gaussian manifold to the Poincaré disk model and using the translation,
measure how the latents well-captures the hierarchical structure of data. We first pick 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 correlation coefficient 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.
20


Reviews:
[{"review_id": "xoJQ0bXA3P", "reviewer": "Reviewer_1SSZ", "summary": "Some 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.", "questions": "Please check other questions and suggestions pointed out with the weakness points. ", "limitations": "The 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. ", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "+ 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", "weaknesses": "- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank for the sincere feedback. We address the issues about i) the modeling capability with local approximation, ii) visualization of our method and iii) analysis on numerical stability.\n\n---\n\n**Q1)  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\nOne of the important properties of hyperbolic space is that the pairwise distances of a tree can be perfectly embedded with the hyperbolic distance function. This enables deep neural networks to learn hierarchical representations by incorporating hyperbolic space as the latent space. \n\nIn previous work, it is shown that KL divergence, the approximation of the hyperbolic distance function, can embed tree-structured data similar to the hyperbolic distance function [1, 2]. Based on the results, we believe that the approximation would not harm the capability of modeling the hierarchical structure.\n\n---\n\n**Q2) The authors may visualize the manifold of the proposed method and compare that with other methods, such as [27].  The sampling results of the VAE models should be visualized and compared with previous methods, such as [27] and [25].**\n\nWe add the density plots of the PGM normal at Figure 3. We run an experiment on tree-structured data to emphasize the visual benefits of hyperbolic VAEs. We train the models on a synthetic binary tree [2] and compare the performance in a density estimation task. Figure 2 describes the visualization of the learned latents. \n\n---\n\n**Q3) 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\nWe provide further explanation on the numerical instability of hyperbolic distributions and quantitative analysis.\n\nWe first analyze the causes of numerical instability of hyperbolic space with different models. The numerical issues of hyperbolic distributions occur when the mean or the sampled point does not belong to the manifold due to the machine precision error. For example, the denominator of the distance of the Poincare disk model becomes zero when the point is placed on the border of the disk. In the Lorentz model, the input of the arccosh function of the log mapping can be less than one when the mean of the distribution is far from the origin and the sampled point is near the mean point, which is out of the domain. Theoretically, these issues should not happen, but we have observed the problems during the experiments due to the limited machine precision. \n\nThese operations are essential in the sampling and KL divergence computation of the distributions. Inspired by the suggestion of reviewer bHMn, we conduct a quantitative analysis on numerical stability by following [1].  We compare the machine precision errors of the three models, the Poincare half-plane model, the Lorentz model, and the Poincare disk model which are used in PGM normal, HWN, and the Poincare normal with respectively. Figure 1 shows that the Poincare half-plane model, which is used in PGM normal, is more stable than the other two models in a wider regime. \n\nFurthermore, the sampling and KL divergence computation of PGM normal are free from the manifold of the Poincare half-plane model because we can express the operations with the log-space of the half-plane manifold, which is identical to Euclidean space. For more details, see the code of density_estimation/distributions/PGMNormal/distribution.py and Appendix F of the supplementary.\n\nWe also empirically show that PGM normal is the only hyperbolic distribution that runs with 16 bit precision in the RL task while the other distributions fail and output NaN.\n\n---\n\n**References**\n\n[1] Nagano et al., A Wrapped Normal Distribution on Hyperbolic Space for Gradient-Based Learning, ICML 2019\n\n[2] Cho et al., A Rotated Hyperbolic Wrapped Normal Distribution for Hierarchical Representation Learning, NeurIPS 2022"}, {"review_id": "0DdSodilHm", "reviewer": "Reviewer_ACLv", "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.", "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).\n\n2. Minor comments:\n\nTab. 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.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The proposed GM-VAE model is a non-trivial improvement of the VAE architecture with a solid theoretical background.\n\n2. 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).\n\n3. The proposed model has superior computational performance concerning the other hyperbolic VAEs considered.\n\n4. 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.\n\n2. Rather small-scale image datasets were used in experiments. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We acknowledge with sincere help for improving our work. We provide the answers for the questions about i) the comparable results in model-based RL and ii) large-scale datasets.\n\n---\n\n**Q1) The experimental results of the GM-VAE in the model-based RL task are only comparable to those of competitors.**\n\nWe agree that the results are comparable. However, to the best of our knowledge, it is the first time in model-based RL to successfully run and achieve comparable results to the previous baseline with the hyperbolic distribution. We attempt to provide a new-axis to the RL community by applying hyperbolic space for model-based RL algorithm, where previous work is focusing on model-free algorithms [1].\n\n---\n\n\n**Q2) Rather small-scale image datasets were used in experiments. 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).**\n\nDue to the short rebuttal period, we could not conduct an experiment with datasets with larger images. However, we believe that the results from the RL environment, which we train with sequences of 84x84 images of length 50, implies that our method can benefit on larger-scale datasets.\n\n---\n\n**Q3) Tab. 2 caption: N/A appears only in the standard deviation (or perhaps confidence interval?) indices.**\n\nSorry for the confusion. The N/A appears only in the confidence interval as mentioned in the question. Instead, in the experiment results of non-product space in Appendix E, we can observe the N/A also at the mean values.\n\n---\n\n**References**\n\n[1] Cetin et al., Hyperbolic Deep Reinforcement Learning, ICLR 2023"}, {"review_id": "xFWsceJoip", "reviewer": "Reviewer_bHMn", "summary": "This 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.", "questions": "1. 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?", "limitations": "Yes.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This 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.", "weaknesses": "There 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the careful comments for the novelty of our method. We address the questions about the i) novelty on numerical stability, ii) drawbacks of HWN, and other minors.\n\n---\n\n**Q1) PGM normal is another novel way to stabilize hyperbolic distribution.**\n\nWe thank for introducing a work related to the stability of hyperbolic space. Although [1] tackles the stability issue using the Poincare half-plane model similar to our approach, adopting the Poincare half-plane model to the previous hyperbolic distributions is difficult, and PGM normal resolves the stability issue in a different direction.\n\nAs suggested, we provide a machine error analysis on the models of hyperbolic space at section A of the general response. The results reveal that the Poincare half-plane model with MCF can be a better choice for HWN and Riemannian normal as suggested. However, it is hard to directly apply the proposed methods from [1] to the baseline distributions. Firstly, the computation of the normalizing constant and density function is necessary, and another new rejection sampling algorithm is also required for Riemannian normal. Achieving such technical contributions requires effort as done in previous work [2, 4] and is another interesting future research direction but out of the scope of our work. Secondly, we need to propose a bespoke architecture to apply MCF to VAE. In general, encoders in hyperbolic VAEs consist of Euclidean layers followed by a transformation layer to hyperbolic space. To apply MCF, we need to partition the encoder output to the components of MCF. This requires a carefully designed encoder layer and a partitioning algorithm because if the output of the encoder is too large, the output can be numerically unstable once transported to the Poincare half-plane. \n\nWe finally emphasize that PGM normal relieves numerical instability in a different manner, not because of the adoptation of the Poincare half-plane model. Remind that the instability comes when the Riemannian operation is applied over the sampled points that do not belong to the manifold. The sampling and KL divergence computation of PGM normal would not suffer from instability since we can sample $\\sigma$ in the log space, i.e., sample $\\log \\sigma$, which is much more stable than sampling $\\sigma$ directly. Consequently, the KL computation can be done with $\\log\\sigma$. Furthermore, the log-space of the Poincare half-plane model is the same as the Euclidean space, so we do not need any exponential and log mappings. We also empirically show that PGM normal successfully runs with 16 bit precision in the RL task because of the simplified sampling and KL computation without any floating point trick. For more details, see the code density_estimation/distributions/PGMNormal/distribution.py and Appendix F.\n\n---\n\n**Q2) The drawbacks of hyperbolic wrapped normal distribution.**\n\nWe respectfully disagree that the mentioned drawbacks only appear in the Poincare normal.\n\nFirstly, HWN suffers from numerical instability due to the log map of the Lorentz model. As shown in Figure 1, the points of the Lorentz model escape from the manifold earlier than the Poincare half-plane model. Hence, the input value of the arccosh of the log map can be less than one, which is out of the domain and causes numerical instability.\n\nSecondly, the sampling speed is slow. As shown in Table 2, the sampling speed of HWN is 8.39x slower than the PGM normal due to additional mappings between the Euclidean space and the Lorentz model. This also affects the speed of the KL divergence computation which is approximated with Monte Carlo sampling. This eventually slows the overall training time of HWN as shown in Table 3 of the main paper.\n\nLastly, the absence of KL divergence makes it hard to utilize complex covariance structure. Table 11 of Appendix D.3 of [3] shows that when using the KL divergence of non-diagonal covariance HWNs as the metric for word embedding, an extensive number of samples are required for stable training. This is not an issue when exploiting simple priors such as zero mean and identity covariance used in previous VAE. However, in the circumstances which require complex prior, e.g., conditional generation and world model of RL, it becomes problematic.\n\n---\n\n**Q3)  Product space vs non-product space.**\n\nOne of the advantages of non-product space is that we can adopt complex covariance structures for the distribution. While Poincare normal and PGM normal only allow isotropic covariance, HWN can take advantage of complex covariance structure. However, as mentioned in Q2, the absence of KL divergence bothers the usage of complex covariance. So all three distributions can take the advantage in limited cases. Also, the instability increases as seen in Appendix E.\n\nWith the product space, by setting the curvature for each space differently, we can express a complicated manifold structure in a different way.  Although we cannot model the complex covariance of product space, we have a more flexible manifold structure.\n\n---\n\n**Q4) Experiment on tree-like dataset.**\n\nWe appreciate suggestions for new dataset. We train the VAEs on synthetic binary tree [2] and compare the density estimation performance. Table 1 and Figure 2 show that the hyperbolic VAEs outperform the Euclidean baseline in both the performance and the visualization.\n\n---\n\n**Q5) Learnable curvature.**\n\nPlease see section B of the general response.\n\n---\n\n**Q6) Geometry of HWN.**\n\nPlease check section C of the general response.\n\n---\n\n**References**\n\n[1] Yu and De Sa, Representing Hyperbolic Space Accurately using Multi-Component Floats, NeurIPS 2021\n\n[2] Nagano et al., A Wrapped Normal Distribution on Hyperbolic Space for Gradient-Based Learning, ICML 2019\n\n[3] Cho et al., A Rotated Hyperbolic Wrapped Normal Distribution for Hierarchical Representation Learning, NeurIPS 2022\n\n[4] Mathieu et al., Continuous Hierarchical Representations with Poincaré Variational Auto-Encoders, NeurIPS 2022"}, {"review_id": "CCdar5Vker", "reviewer": "Reviewer_Darj", "summary": "The 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.", "questions": "1. 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", "limitations": "The paper does not provide much of a discussion of the limitations of the proposed method.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. 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).", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank for providing constructive feedback. We answer the questions about i) analysis on numerical stability, ii) the parameter $c$ and $\\beta$ of VAE, and iii) the approximation performance of KL divergence.\n\n---\n\n**Q1) Quantitative analysis on the numerical stability of hyperbolic distributions.**\n\nWe thank for the suggestion for strengthening our statement about numerical stability. Due to the time limit, instead of running more experiments, we conduct a machine precision error analysis on the hyperbolic distributions, which is the main cause of the numerical instability.\n\nWe first analyze the causes of numerical instability of hyperbolic space with different models. The numerical issues of hyperbolic distributions occur when the mean or the sampled point does not belong to the manifold due to the machine precision error. For example, the denominator of the distance of the Poincare disk model becomes zero when the point is placed on the border of the disk. In the Lorentz model, the input of the arccosh function of the log mapping can be less than one when the mean of the distribution is far from the origin and the sampled point is near the mean point, which is out of the domain. Theoretically, these issues should not happen, but we have observed the problems during the experiments due to the limited machine precision. \n\nThese operations are essential in the sampling and KL divergence computation of the distributions. Inspired by the suggestion of reviewer bHMn, we conduct a quantitative analysis on numerical stability by following [1].  We compare the machine precision errors of the three models, the Poincare half-plane model, the Lorentz model, and the Poincare disk model which are used in PGM normal, HWN, and the Poincare normal with respectively. Figure 1 shows that the Poincare half-plane model, which is used in PGM normal, is more stable than the other two models in a wider regime. \n\nFurthermore, the sampling and KL divergence computation of PGM normal are free from the manifold of the Poincare half-plane model because we can express the operations with the log-space of the half-plane manifold, which is identical to Euclidean space. For more details, see the code of density_estimation/distributions/PGMNormal/distribution.py and Appendix F of the supplementary.\n\nWe also empirically show that PGM normal is the only hyperbolic distribution that runs with 16 bit precision in the RL task while the other distributions fail and output NaN.\n\n---\n\n**Q2) The curvature $c$ of hyperbolic VAE.**\n\nThe curvature of the hyperbolic space is mostly set to minus one, the best value in previous work, or learned together [6]. Some work shows that the gradient of the curvature can be computed through the Riemannian operations and carefully learning it together can enhance the performance [6]. We also conducted an experiment on learning the curvature together. Please check the “Learning curvature experiment” section of the general response for details. \n\n---\n\n**Q3) The $\\beta$ parameter of VAE.**\n\nBy controlling $\\beta$, we can make a tradeoff between reconstruction and KL divergence [5]. If $\\beta$ is large, latents will better match with the prior, but the reconstruction performance will degrade. If $\\beta$ is low, the reconstruction performance will be enhanced, but the data would be overfit at the training data. The tradeoff is an important factor in the VAE literature and frequently studied. It is usually set to one, slightly increased as the training proceeds, or follows the previous work. In the RL setting, we set $\\beta$ to 0.1 which is used in the baseline model DreamerV2.\n\n---\n\n**Q4) Tree embedding with KL divergence.**\n\nIn previous work, it is empirically shown that KL divergence, the approximation of the hyperbolic distance function, can embed tree-structured data similar to the hyperbolic distance function [2, 3]. Based on the results, we believe that the approximation would not harm the capability of modeling the hierarchical structure.\n\n---\n\n**Q5) The approximation performance between KL divergence and wrapped normal.**\n\nThe geometry column of Table 1 from the main paper means that the distribution represents the density of a sample point with the hyperbolic distance. Although the HWN utilizes operations with Riemannian metric in sampling, we thought that the resulting distance is dependent on the Euclidean distance, not the hyperbolic distance due to the normal distribution on the tangent space. However, we find that the density of HWN can be seen as a local approximation of the hyperbolic distance [4], so we will remove the table in future revision. \n\nWe conduct a study on the approximation performance of KL divergence and wrapped normal operation, which are used in PGM normal and HWN with respectively. The results are at Figure 4. While the KL divergence shows better approximation in points with small squared distance, the wrapped normal better approximate the squared distance in far points.\n\n---\n\n**References**\n\n[1] Yu and De Sa, Representing Hyperbolic Space Accurately using Multi-Component Floats, NeurIPS 2021\n\n[2] Nagano et al., A Wrapped Normal Distribution on Hyperbolic Space for Gradient-Based Learning, ICML 2019\n\n[3] Cho et al., A Rotated Hyperbolic Wrapped Normal Distribution for Hierarchical Representation Learning, NeurIPS 2022\n\n[4] De Bortoli et al., Riemannian score-based generative modelling, NeurIPS 2022\n\n[5] Alemi et al., Fixing a Broken ELBO, ICML 2018\n\n[6] Skopek et al., Mixed-curvature variational autoencoders, ICLR 2020"}]

Author rebuttal:
We appreciate all reviewers for valuable comments which can improve our work. We reply to each reviewer with independent comments. We also provide supplementary material to complement the response to reviewers. The supplementary material consists of a PDF with figures and tables, and answers to frequently asked questions by the reviewers.

We organize the PDF with 4 figures and 3 tables. The figures are about the quantitative analysis on the numerical stability of the hyperbolic models [1SSZ, bHMn, Darj], the latent space of the hyperbolic VAEs [1SSZ, bHMn], the density plots of PGM normal [1SSZ], and the approximation performance between KL divergence and wrapped normal operation [Darj]. For the tables, we report the results of density estimation on synthetic binary tree [bHMn], detailed training speed analysis [bHMn], and the result of density estimation with learnable curvature GM-VAE  [bHMn, Darj]. 

In the next section, we answer the questions which are frequently asked by the reviewers.

---

**A. Quantitative analysis on the numerical instability [1SSZ, bHMn, Darj]**

We first analyze the causes of numerical instability of hyperbolic space with different models. The numerical issues of hyperbolic distributions occur when the mean or the sampled point does not belong to the manifold due to the machine precision error. For example, the denominator of the distance of the Poincare disk model becomes zero when the point is placed on the border of the disk. In the Lorentz model, the input of the arccosh function of the log mapping can be less than one, which is out of the domain. Theoretically, these issues should not happen, but we have observed the problems during the experiments due to the limited machine precision. 

These operations are essential in the sampling and KL divergence computation of the distributions. Inspired by the suggestion of reviewer bHMn, we conduct a quantitative analysis on numerical stability by following [1].  We compare the machine precision errors of the three models, the Poincare half-plane model, the Lorentz model, and the Poincare disk model which are used in PGM normal, HWN, and the Poincare normal with respectively. Figure 1 shows that the Poincare half-plane model with multi-component floats (MCF) is more stable than the other models in a wider regime. 

---

**B. Learnable curvature experiment [bHMn, Darj]**

PGM normal can leverage the learnable curvature. We test the PGM normal with learnable curvature on the density estimation task. We let all the curvatures learnable, and set the learning rate same as the VAE learning rate. We initialize the curvature to $-1$. Table 3 shows that there is no enhancement in the density estimation performance. Due to the short rebuttal period, we haven’t tested all possible configurations of learnable curvature, but we conjecture that carefully chosen learning rate and the initial value could improve the performances.


---

**C. Meaning of “Geometry” column [bHMn, Darj]**

The geometry column of Table 1 from the main paper means that the distribution represents the density of a sample point with the hyperbolic distance. Although the HWN utilizes operations with Riemannian metric in sampling, we thought that the resulting distance is dependent on the Euclidean distance, not the hyperbolic distance due to the normal distribution on the tangent space. However, we find that the density of HWN can be seen as a local approximation of the hyperbolic distance [2], so we will remove the table in future revision. 

We conduct a study on the approximation performance of KL divergence and wrapped normal operation, which are used in PGM normal and HWN with respectively. The results are at Figure 4. While the KL divergence shows better approximation in points with small squared distance, the wrapped normal better approximate the squared distance in far points.

---

**References**

[1] Yu and De Sa, Representing Hyperbolic Space Accurately using Multi-Component Floats, NeurIPS 2021

[2] De Bortoli et al., Riemannian score-based generative modelling, NeurIPS 2022


Meta-review: The paper introduces a novel method for VAEs with a latent space constructed using a hyperbolic distribution, modeled through the KL divergence to locally approximate the hyperbolic geodesic distance. The motivation is to improve numerical stability and make sampling easier. The experiments demonstrate its effectiveness on image datasets and model-based Reinforcement Learning (RL) benchmarks.

**Strengths:**
- Novel construction of a hyperbolic distribution, providing a unique way to form a hyperbolic latent space in a VAE setting.
- The method allows for straightforward analytic operations and is grounded in a clear motivation.
- The paper is well-written, with several reviewers noting its clarity and ease of understanding.
- Experimental results on multiple datasets showcase the potential superiority or comparability of this approach over its competitors.
- The factorization of the proposed distribution into known distributions is seen as an elegant feature.

**Weaknesses:**
- Several reviewers raised concerns about the clarity and thoroughness of the paper's exploration into the paper’s main benefit of numerical stability.
- Experimental results in the model-based RL task are only comparable to competitors, and some evaluations were performed on relatively small-scale image datasets.
- Some reviewers pointed out missing citations to relevant works.
- Presentation aspects, such as the introduction to Hyperbolic geometry and definition of certain models, were seen as lacking or non-standard.
- Reviewers noted that the authors could have delved deeper into potential limitations or challenges.

Given the aggregate of the reviews, the paper provides a technically solid contribution with a novel perspective on VAEs and hyperbolic spaces. However, there are concerns about the comprehensive demonstration of its benefits, especially numerical stability, and certain presentation aspects. A revision addressing the aforementioned weaknesses could bolster the paper's impact.
Decision: Accept (poster)