SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Note that the sectional curvature of two-dimensional Riemannian manifold is same as the Gaussian curvature where  $\langle d\mu, d\mu \rangle \langle d\sigma, d\sigma \rangle - \langle d\mu, d\sigma \rangle^2 = \det g$ .
B.2 Gaussian manifold with KL-divergence
Between two univariate Gaussian distributions  $\mathcal{N}(\mu_1,\sigma_1^2)$  and  $\mathcal{N}(\mu_2,\sigma_2^2)$ , we can compute the KL divergence as:
$$
D _ {\mathrm {K L}} \left(\mathcal {N} \left(\mu_ {1}, \sigma_ {1} ^ {2}\right) \| \mathcal {N} \left(\mu_ {2}, \sigma_ {2} ^ {2}\right)\right) = \frac {1}{2} \left(\log \frac {\sigma_ {2} ^ {2}}{\sigma_ {1} ^ {2}} + \frac {\sigma_ {1} ^ {2} + \left(\mu_ {1} - \mu_ {2}\right) ^ {2}}{\sigma_ {2} ^ {2}} - 1\right). \tag {4}
$$
We extend the KL divergence for an arbitrary curvature of the Gaussian manifold as:
$$
D _ {\mathrm {K L}} ^ {\mathcal {G} _ {c}} \left(\left(\mu_ {1}, \sigma_ {1}\right), \left(\mu_ {2}, \sigma_ {2}\right)\right) := \frac {D _ {\mathrm {K L}} \left(\mathcal {N} \left(\sqrt {2 c} \mu_ {1} , \sigma_ {1} ^ {2}\right) \parallel \mathcal {N} \left(\sqrt {2 c} \mu_ {2} , \sigma_ {2} ^ {2}\right)\right)}{2 c}.
$$
Now, we show that the extended KL divergence still approximates the Riemannian distance of the manifold as:
$$
\begin{array}{l} D _ {\mathrm {K L}} ^ {\mathcal {G} _ {c}} ((\mu + d \mu , \sigma + d \sigma), (\mu , \sigma)) = \frac {1}{2 \cdot 2 c} \left(\log \frac {\sigma^ {2}}{(\sigma + d \sigma) ^ {2}} + \frac {(\sigma + d \sigma) ^ {2} + 2 c (d \mu) ^ {2}}{\sigma^ {2}} - 1\right) \\ = \frac {1}{2 \cdot 2 c} \left(- 2 \log \left(1 + \frac {d \sigma}{\sigma}\right) + \frac {2 \sigma d \sigma + (d \sigma) ^ {2}}{\sigma^ {2}} + \frac {2 c (d \mu) ^ {2}}{\sigma^ {2}}\right) \\ = \frac {1}{2 \cdot 2 c} \left(- 2 \left(\frac {d \sigma}{\sigma} - \frac {(d \sigma) ^ {2}}{2 \sigma^ {2}}\right) + \frac {2 \sigma d \sigma + (d \sigma) ^ {2}}{\sigma^ {2}} + \frac {2 c (d \mu) ^ {2}}{\sigma^ {2}} + \mathcal {O} ((d \sigma) ^ {3})\right) \\ = \frac {1}{2} \left( \begin{array}{c} d \mu \\ d \sigma \end{array} \right) ^ {T} \left( \begin{array}{c c} \frac {1}{\sigma} & 0 \\ 0 & \frac {1}{c \sigma^ {2}} \end{array} \right) \left( \begin{array}{c} d \mu \\ d \sigma \end{array} \right) + \mathcal {O} ((d \sigma) ^ {3}). \\ \end{array}
$$
C Pseudo Gaussian Manifold Normal Distribution
In this section, we derive the normalizing constant  $Z(c,\beta ,\gamma)$  and the factorization of the PGM normal which the density function is defined as:
$$
\mathcal {K} _ {c} (\mu , \sigma ; \alpha , \beta , \gamma) = \frac {\sigma^ {3}}{Z (c , \beta , \gamma)} \exp \left(- \frac {D _ {\mathrm {K L}} ^ {\mathcal {G} _ {c}} ((\mu , \sigma) , (\alpha , \beta))}{\gamma^ {2}}\right).
$$
C.1 Normalizing Constant
The given probability density function needs to satisfy the following condition:
$$
\int_ {\mathcal {G} _ {c}} \mathcal {K} _ {c} (\mu , \sigma ; \alpha , \beta , \gamma) \sqrt {\det  g} \cdot d (\mu , \sigma) = 1, \tag {5}
$$
where  $\sqrt{\operatorname*{det}g}\cdot d(\mu ,\sigma)$  is the probability measure over the Gaussian manifold induced from the Lebesgue measure  $d(\mu ,\sigma)$ . The normalizing factor  $Z(c,\beta ,\gamma)$  can be computed using the condition
Equation 5 as:
$$
Z (c, \beta , \gamma) = \int_ {0} ^ {\infty} \int_ {- \infty} ^ {\infty} \sigma^ {3} \cdot \exp \left(- \frac {D _ {\mathrm {K L}} ^ {\mathcal {G} _ {c}} ((\mu , \sigma) , (\alpha , \beta))}{\gamma^ {2}}\right) \frac {1}{\sqrt {c} \sigma^ {2}} d \mu d \sigma
$$
$$
\begin{array}{l} = \frac {1}{\sqrt {c}} \left(\beta^ {- \frac {1}{- 2 c \gamma^ {2}}} \exp \left(\frac {1}{4 c \gamma^ {2}}\right) \int_ {0} ^ {\infty} \sigma \cdot (\sigma^ {2}) ^ {\left(\frac {1}{4 c \gamma^ {2}} + 1\right) - 1} \exp \left(- \frac {\sigma^ {2}}{4 c \beta^ {2} \gamma^ {2}}\right) d \sigma\right) \\ \times \left(\int_ {- \infty} ^ {\infty} \exp \left(- \frac {(\mu - \alpha) ^ {2}}{2 \beta^ {2} \gamma^ {2}}\right) d \mu\right) \\ \end{array}
$$
$$
\begin{array}{l} = \frac {1}{2 \sqrt {c}} \sqrt {2 \pi} \beta^ {3} \gamma \exp \left(\frac {1}{4 c \gamma^ {2}}\right) \Gamma \left(\frac {1}{4 c \gamma^ {2}}\right) \left(\frac {1}{4 c \gamma^ {2}}\right) ^ {- \frac {1}{4 c \gamma^ {2}}} \\ \times \left(\int_ {0} ^ {\infty} \operatorname {G a m m a} \left(\sigma^ {2}; \frac {1}{4 c \gamma^ {2}} + 1, \frac {1}{4 c \beta^ {2} \gamma^ {2}}\right) d \sigma^ {2}\right) \left(\int_ {- \infty} ^ {\infty} \mathcal {N} (\mu ; \alpha , \beta \gamma) d \mu\right) \tag {6} \\ \end{array}
$$
$$
= \frac {\sqrt {2 \pi} \beta^ {3}}{2 \sqrt {c}} \gamma \exp \left(\frac {1}{4 c \gamma^ {2}}\right) \Gamma \left(\frac {1}{4 c \gamma^ {2}}\right) \left(\frac {1}{4 c \gamma^ {2}}\right) ^ {- \frac {1}{4 c \gamma^ {2}}}.
$$
Finally, the logarithm of the normalizing factor is computed as:
$$
\log Z (c, \beta , \gamma) = \frac {1}{2} \log (2 \pi) + 3 \log \beta - \frac {1}{2} \log c - \log 2 + \frac {1}{2} \log \gamma^ {2} + \log \Gamma \left(\frac {1}{4 c \gamma^ {2}}\right) + \frac {1}{4 c \gamma^ {2}} (1 + \log (4 c \gamma^ {2})).
$$
C.2 Sampling
Suppose that  $p(\mu, \sigma) = \mathcal{K}_c(\mu, \sigma; \alpha, \beta, \gamma) \sqrt{\operatorname*{det} g}$ . Sampling  $\mu$  and  $\sigma$  from the probability distribution  $p(\mu, \sigma)$  can be done by sampling  $\mu$  from the marginal distribution  $p(\mu)$  and then sampling  $\sigma$  from the conditional distribution  $p(\sigma | \mu)$ . The marginal distribution  $p(\mu)$  can be derived from Equation 6 as:
$$
\begin{array}{l} p (\mu) = \int_ {0} ^ {\infty} p (\mu , \sigma) d \sigma \\ = \int_ {0} ^ {\infty} \mathcal {K} _ {c} (\mu , \sigma ; \alpha , \beta , \gamma) \sqrt {\det g} d \sigma \\ = \left(\int_ {0} ^ {\infty} \operatorname {G a m m a} \left(\sigma^ {2}; \frac {1}{4 c \gamma^ {2}} + 1, \frac {1}{4 c \beta^ {2} \gamma^ {2}}\right) d \sigma^ {2}\right) \mathcal {N} (\mu ; \alpha , \beta^ {2} \gamma^ {2}) \\ = \mathcal {N} (\mu ; \alpha , \beta^ {2} \gamma^ {2}). \\ \end{array}
$$
Equation 6 also implies that  $\mu$  and  $\sigma$  are independent in the aspect of  $p(\mu, \sigma)$  so the conditional distribution  $p(\sigma | \mu)$  is identical to the marginal distribution  $p(\sigma)$ . The marginal distribution  $p(\sigma)$  is computed as:
$$
\begin{array}{l} p (\sigma) = \int_ {- \infty} ^ {\infty} p (\mu , \sigma) d \mu \\ = \int_ {- \infty} ^ {\infty} \mathcal {K} _ {c} (\mu , \sigma ; \alpha , \beta , \gamma) \sqrt {\det g} d \mu \\ = \left(2 \sigma \cdot \operatorname {G a m m a} \left(\sigma^ {2}; \frac {1}{4 c \gamma^ {2}} + 1, \frac {1}{4 c \beta^ {2} \gamma^ {2}}\right)\right) \left(\int_ {- \infty} ^ {\infty} \mathcal {N} (\mu ; \alpha , \beta^ {2} \gamma^ {2}) d \mu\right) \\ = 2 \sigma \cdot \operatorname {G a m m a} \left(\sigma^ {2}; \frac {1}{4 c \gamma^ {2}} + 1, \frac {1}{4 c \beta^ {2} \gamma^ {2}}\right). \\ \end{array}
$$
Here, sampling  $\sigma$  from  $p(\sigma)$  can be easily replaced by the procedure of sampling  $\sigma^2$  from  $p(\sigma^2)$ , which is identical to  $p(\sigma) / (2\sigma) = \mathrm{Gamma}\left(\sigma^2; \frac{1}{4c\gamma^2} + 1, \frac{1}{4c\beta^2\gamma^2}\right)$ , and then applying square root to the sample  $\sigma^2$  to get  $\sigma$ .
C.3 KL Divergence
Suppose that  $p(\mu, \sigma), q(\mu, \sigma)$  are two different PGM normal multiplied with  $\sqrt{\operatorname*{det}g}$ . As shown in Appendix C.2,  $\mu$  and  $\sigma$  are independent so the KL divergence between  $p, q$  is same as  $D_{\mathrm{KL}}(p(\mu)\|q(\mu)) + D_{\mathrm{KL}}(p(\sigma)\|q(\sigma))$ . The first term is the KL divergence between two normal distribution. The second term is same with  $D_{\mathrm{KL}}(p(\sigma^2)\|q(\sigma^2))$  due to the change-of-variable formula, so it is the KL divergence between two gamma distribution.
D  $\delta$ -Hyperbolicity
In this section, we explain about  $\delta$ -hyperbolicity ( $\delta$ -H) and show the  $\delta$ -H values of the candidate datasets of our experiments.
D.1  $\delta$ -hyperbolicity
Given a metric space  $(X,d)$ , the metric space is said to be  $\delta$ -hyperbolic if, for any geodesic triangle, i.e., a triangle where each side is a geodesic curve, any point on the side of the geodesic triangle is within the distance of less than or equal to  $\delta$  of the other two sides; when such  $\delta$  exists,  $(X,d)$  is said to be hyperbolic. To be specific, when the Gromov product between any three points  $x,y,z\in X$  is given as:
$$
(y, z) _ {x} = \frac {1}{2} (d (x, y) + d (x, z) - d (y, z)),
$$
if then the following inequality holds for any four points  $x,y,z,w\in X$  , we call the metric space is  $\delta$  -hyperbolic:
$$
(x, z) _ {w} \geq \min  ((x, y) _ {w}, (y, z) _ {w})) - \delta . \tag {7}
$$
$\delta$ -hyperbolicity ( $\delta$ -H) is defined to be the minimum value of  $\delta$  satisfying Equation 7 and is used as a measurement quantifying how a given metric space well embeds in hyperbolic Khrulkov et al. (2020); Cetin et al. (2022). Figure 4 illustrates the concept of  $\delta$ -H. We note that the lower  $\delta$ -H the metric space has, the less deviation from the exact hyperbolic space is.
(a)  $\delta -\mathbf{H} > 0$
Figure 4: The illustration of  $\delta$ -H for given geodesic triangles. The blue lines denote the geodesic curves between the points. (a) When the side  $AB$  is contained in the union of the two  $\delta$ -neighbor regions of the side  $AC$  and  $BC$ , we say that the geodesic triangle is  $\delta$ -hyperbolic.  $\delta$ -H of the geodesic triangle is then defined as the minimum possible value of  $\delta$ . (b) Any tree-structured triangle, where the geodesic curves correspond to the tree edges, is  $0$ -H. For example, given geodesic triangle  $ABC$ , one side  $AB$  is already occupied by the two sides  $AB$  and  $BC$ , as the geodesic between the points  $A$  and  $B$  is the union of edge  $AD$  and  $BD$  being occupied by other counterparts.
(b)  $\delta \text{-H} = 0$
Computation. We measure the  $\delta$ -H values of  $BD$  and  $BD$  being occupied by other counterparts. 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  $\delta$ -H of the images  $X$  by finding the largest coefficient of  $(\max_k \min_{ij} (D_{ik}, D_{kj})) - D$ .
To reduce the scale difference between the datasets, we report the value  $2\delta(X)/\mathrm{diam}(X)$  where  $\mathrm{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  $\delta$ -H of randomly sampled 1,000 images from  $X$ . We repeat this 10 times and report the average  $\delta$ -H.
D.2 Image datasets
We measure the  $\delta$ -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  $\delta$ -H using the image encoder from the agents. For each environment, we report the  $\delta$ -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.
Game	δ-H	Reward
Breakout	0.12	340
Alien	0.14	6855
Zaxxon	0.14	13100
IceHockey	0.14	3
Gravitar	0.17	1004
Carnival	0.17	5605
RoadRunner	0.18	60050
Pong	0.18	21
Tutankham	0.18	236
Boxing	0.19	99
Solaris	0.19	1727
WizardOfWor	0.20	17217
Seaquest	0.20	29745
Gopher	0.20	24173
Centipede	0.20	4663
PrivateEye	0.21	195
Pitfall	0.21	0
ElevatorAction	0.21	64917
Robotank	0.21	75
MsPacman	0.21	5455
DemonAttack	0.22	23876
Asterix	0.22	26792
Qbert	0.22	17364
Jamesbond	0.22	886
VideoPinball	0.23	632563
UpNDown	0.23	27030
FishingDerby	0.23	54
Tennis	0.23	22
Berzerk	0.24	904
Riverraid	0.24	15103
Game	δ-H	Reward
AirRaid	0.25	11729
Frostbite	0.25	8653
Pooyan	0.25	7973
ChopperCommand	0.26	10530
KungFuMaster	0.26	27733
YarsRevenge	0.26	56405
Phoenix	0.27	30208
JourneyEscape	0.27	-707
SpaceInvaders	0.27	15623
Hero	0.28	28592
Enduro	0.28	2089
Assault	0.28	3217
Venture	0.29	1529
StarGunner	0.29	68417
Atlantis	0.29	919750
DoubleDunk	0.30	18
TimePilot	0.30	10914
BeamRider	0.30	7069
BattleZone	0.30	40333
NameThisGame	0.31	12925
Skiing	0.31	-10021
Amidar	0.32	2826
Asteroids	0.34	1405
MontezumaRevenge	0.34	0
Bowling	0.35	49
BankHeist	0.36	1563
Kangaroo	0.37	14283
CrazyClimber	0.37	132517
Freeway	0.38	34
Krull	0.38	9016
In this section, we introduce the implementation details of the experiments.
E.1 Density estimation on image datasets
We estimate the density of the images from Atari2600 Breakout (Nagano et al., 2019), Oxford102 (Nilsback & Zisserman, 2008), and CUB (Wah et al., 2011). The images of Breakout are collected from plays with a pre-trained Deep Q-Network (Mnih et al., 2015). The size of images of all the datasets is resized to  $64 \times 64$ , while Breakout is binarized with a threshold value of 0.1; the threshold for Breakout is determined to visualize the components clearly.
We split the datasets into train, validation, and test. For Breakout and CUB, we split the original train set into train and validation sets. For Oxford102, because the original train set is too small, we merge the original train and test set and then split it into three splits. For Food101, we randomly sample the train set and validation set from the original train set, and also randomly sample the test set from the original test set.
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
We design the encoder and decoder similar to the generator and discriminator of DCGAN (Radford et al., 2015). The details of the architecture are at Table 3. We use learning rate 1e-3, batch size 100, and Adam optimizer for training. We use Bernoulli loss as the reconstruction loss for Breakout experiments and negative log-likelihood loss as the reconstruction loss for CUB, Food101, and Oxford102 experiments. We use the validation set for early stopping and report the negative log-likelihood on the test set with 50 importance weighted samples.
Table 3: The architectures of encoder and decoder used in the density estimation experiments.  $n_c$  is the number of channels of the image,  $n_d$  is the latent dimension.  $n_a$  is a coefficient that depends on the VAE, i.e.,  $n_a$  is 2, 2, 1.5, 1.5 for  $\mathcal{E}$ -VAE,  $\mathcal{L}$ -VAE,  $\mathcal{P}$ -VAE, and GM-VAE, respectively.
Encoder	Decoder
Layer	Size	Layer	Size
Input	64 × 64 × nc	Input	1 × 1 × nd
Convolution2D	32 × 32 × 32	TransposedConvolution2D	4 × 4 × 256
LeakyReLU	ReLU
Convolution2D	16 × 16 × 64	TransposedConvolution2D	8 × 8 × 128
LeakyReLU	ReLU
Convolution2D	8 × 8 × 128	TransposedConvolution2D	16 × 16 × 64
LeakyReLU	ReLU
Convolution2D	4 × 4 × 256	TransposedConvolution2D	32 × 32 × 32
LeakyReLU	ReLU
Linear	na · nd	TransposedConvolution2D	64 × 64 × nc
E.2 Model-based RL
We use the official TensorFlow implementation from Dreamerv2 $^1$  to reproduce the baseline results, i.e., with Euclidean and discrete latent space. For the hyperbolic latent space results, we apply GM-VAE by replacing the latent space of  $z_{t}$  with the Gaussian manifold and two components in RSSM, the representation model  $q_{\theta}(z_t|h_t,x_t)$  and transition predictor  $p_{\phi}(z_t|h_t)$ , with PGM normal. The hyperparameters are set to be the same as suggested in the original paper, except for the training environment steps being 50M for Freeway and 100M for the others as we observe converging scores.
F Results of Non-Product Latent Space
Previous hyperbolic VAEs are implemented with a hyperbolic space, not the product of the hyperbolic spaces. We run the non-product hyperbolic VAEs in the density estimation of image datasets. Table 4 reveals that the non-product hyperbolic VAEs fail in most of the settings and the product hyperbolic space makes the hyperbolic VAEs much more stable.
G Model-Based RL
G.1 Learning curves
G.2 Latent space analysis
We conduct an analysis of the latent space of the agents learned to play Atari2600 Breakout. The purpose of the analysis is to measure how the latent spaces well-preserve the implicit hierarchy in the trajectory of the agents. To analyze the hyperbolic latent space, we need two isometries: the isometry between the Gaussian manifold and the Poincaré disk model and the translation of the Poincaré disk model.
$^{1}$ https://github.com/danijar/dreamerv2
Table 4: Density estimation results of non-product hyperbolic VAEs.  $d$  denotes the latent dimension. N/A in the log-likelihood indicates that the results are not available due to the failure of all runs.
d	L-VAE	P-VAE
Breakout	2	124.24±1.66	266.86±6.01
4	66.20±0.14	N/A
8	44.76±0.48	N/A
CUB	50	N/A	N/A
60	N/A	N/A
70	N/A	N/A
Oxford102	50	N/A	N/A
60	N/A	N/A
70	N/A	N/A
We first propose an isometry between the Gaussian manifold and the Poincaré disk model  $T_{\mathcal{P}_c \to \mathcal{U}_c}$ :  $\mathcal{P}_c \to \mathcal{U}_c$ :
$$
T _ {\mathcal {P} _ {c} \to \mathcal {G} _ {c}} (x, y) = \left(\frac {- 2 y}{(\sqrt {c} x - 1) ^ {2} + y ^ {2} c}, \frac {1 - (x ^ {2} + y ^ {2}) c}{(\sqrt {c} x - 1) ^ {2} + y ^ {2} c}\right),
$$
and the inverse is:
$$
T _ {\mathcal {P} _ {c} \to \mathcal {G} _ {c}} ^ {- 1} (x, y) \left(\frac {\sqrt {c} x ^ {2} + (y ^ {2} - 1) / \sqrt {c}}{c x ^ {2} + (y + 1) ^ {2}}, \frac {- 2 x}{c x ^ {2} + (y + 1) ^ {2}}\right).
$$
The translation of the Poincaré disk model can be derived using complex numbers. Let  $z = x + yi \in \mathbb{C}$  and  $(x,y) \in \mathcal{P}_1$  and  $z_0 \in \mathbb{C}$  be the pivot point. Then the isometry that moves  $z_0$  to the origin is defined as  $T(z): \mathcal{P}_1 \to \mathcal{P}_1 \coloneqq \frac{z - z_0}{1 - \bar{z}_0z}$ . Note that the translation of Euclidean space is  $z - z_0$ .
After transforming the latents on the Gaussian manifold to the Poincaré disk model and using the translation, we can measure how the latents well-captures the hierarchical structure of data. We first pick a latent and then translate all the latents by setting the selected latent as the pivot point. We then measure the Pearson correlation between the cumulative reward of the latents and the norm.
We repeat this process for all the latents and compute the maximum of the correlations. We use the latents obtained from the agents which recorded at least 250 for long enough trajectories. We obtain a correlation coefficient of 0.46 from the hyperbolic latent space, whereas the correlation coefficient of the Euclidean latent space is 0.40, showing the hyperbolic space better captures the hierarchy along the increasing norm.