SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Haitz Sáez de Océariz Borde*  
Oxford Robotics Institute  
University of Oxford
Álvaro Arroyo* Oxford-Man Institute University of Oxford
Ismael Morales López  
Mathematical Institute  
University of Oxford
Ingmar Posner  
Oxford Robotics Institute  
University of Oxford
Xiaowen Dong  
Machine Learning Research Group  
University of Oxford
Recent research indicates that the performance of machine learning models can be improved by aligning the geometry of the latent space with the underlying data structure. Rather than relying solely on Euclidean space, researchers have proposed using hyperbolic and spherical spaces with constant curvature, or combinations thereof, to better model the latent space and enhance model performance. However, little attention has been given to the problem of automatically identifying the optimal latent geometry for the downstream task. We mathematically define this novel formulation and coin it as neural latent geometry search (NLGS). More specifically, we introduce an initial attempt to search for a latent geometry composed of a product of constant curvature model spaces with a small number of query evaluations, under some simplifying assumptions. To accomplish this, we propose a novel notion of distance between candidate latent geometries based on the Gromov-Hausdorff distance from metric geometry. In order to compute the Gromov-Hausdorff distance, we introduce a mapping function that enables the comparison of different manifolds by embedding them in a common high-dimensional ambient space. We then design a graph search space based on the notion of smoothness between latent geometries, and employ the calculated distances as an additional inductive bias. Finally, we use Bayesian optimization to search for the optimal latent geometry in a query-efficient manner. This is a general method which can be applied to search for the optimal latent geometry for a variety of models and downstream tasks. We perform experiments on synthetic and real-world datasets to identify the optimal latent geometry for multiple machine learning problems.
There has been a recent surge of research employing ideas from differential geometry and topology to improve the performance of learning algorithms [Bortoli et al., 2022, Hensel et al., 2021, Chamberlain et al., 2021, Huang et al., 2022, Barbero et al., 2022a,b]. Traditionally, Euclidean spaces have been the preferred choice to model the geometry of latent spaces in the ML community [Weber, 2019, Bronstein et al., 2021]. However, recent work has found that representing the latent space with a geometry that better matches the structure of the data can provide significant performance enhancements in
both reconstruction and other downstream tasks [Shukla et al., 2018]. In particular, most works have employed constant curvature model spaces such as the Poincaré ball model [Mathieu et al., 2019], the hyperboloid [Chami et al., 2019], or the hypersphere [Zhao et al., 2019], to encode latent representations of data in a relatively simple and computationally tractable way.
While individual model spaces have sometimes shown superior performance when compared to their Euclidean counterparts, more recent works [Gu et al., 2018, Skopek et al., 2019, Sáez de Ocáriz Borde et al., 2023b,a, Zhang et al., 2020, Fumero et al., 2021, Pfau et al., 2020] have leveraged the notion of product spaces (also known as product manifolds) to model the latent space. This idea allows to generate more complex representations of the latent space for improved performance by taking Cartesian products of model spaces, while retaining the computational tractability of mathematical objects such as exponential maps or geodesic distances, see Sáez de Ocáriz Borde et al. [2023b]. Despite the success of this methodology, there exists no principled way of obtaining the product manifold signature (i.e., the choice and number of manifold components used to generate the product manifold and their respective dimensionalities) to optimally represent the data for downstream task performance. This procedure is typically performed heuristically, often involving a random search over the discrete combinatorial space of all possible combinations of product manifold signatures, which is an exceedingly large search space that hampers computational efficiency and practical applicability. Some other work related to latent space geometry modelling can be found in Lubold et al. [2023], Hauberg et al. [2012], Arvanitidis et al. [2017].
Contributions. 1) In this paper, we consider a novel problem setting where we aim to search for an optimal latent geometry that best suits the model and downstream task. As a particular instance of this setting, we consider searching for the optimal product manifold signature. Due to the conceptual similarity with neural architecture search (NAS) strategies [Elsken et al., 2018, Zoph and Le, 2016, Pham et al., 2018], we coin this problem as neural latent geometry search (NLGS), which we hope will encourage additional work in the direction of optimal latent geometry inference. We test our framework on a variety of use cases, such as autoencoder reconstruction [Mathieu et al., 2019] and latent graph inference [Sáez de Ocariz Borde et al., 2023b], for which we create a set of custom datasets.
2) To search for the product manifold signature, we must be able to compare product manifolds. This is traditionally done by computing the Hausdorff distance between manifolds [Taha and Hanbury, 2015], which however requires manifolds to reside in the same metric space. To address this limitation, in this work we develop to our knowledge the first computational method to mathematically compare product manifolds. Our approach generalizes classical algorithms and allows the comparison of manifolds existing in different spaces. This is achieved by defining an isometric embedding that maps the manifolds to a common high-dimensional ambient space, which enables the computation of their Gromov-Hausdorff distances.
3) Leveraging the Gromov-Hausdorff distances between candidate latent space manifolds, we design a principled and query-efficient framework to search for an optimal latent geometry, in the sense that it yields the best performance with respect to a given machine learning model and downstream task.
Our approach consists of constructing a geometry-informed graph search space where each node in the graph represents a unique candidate product manifold, associated with the model performance using this manifold as its embedding space. The strength of edges in the graph are based on the inverse of the Gromov-Hausdorff distance, thereby encoding a notion of "closeness" between manifolds in the search space. We then perform efficient search over this space using Bayesian optimization (BO). We compare our proposed method with other search algorithms that lack the topological prior inherent in our model. Empirical results demonstrate that our method outperforms the baselines by a significant margin in finding the optimal product manifold.
Outline. In Section 2 we discuss manifold learning, and product manifolds of constant curvature model spaces such as the Euclidean plane, the hyperboloid, and the hypersphere. We also review relevant mathematical concepts, particularly the Hausdorff and Gromov-Hausdorff distances from metric geometry [Gopal et al., 2020]. Section 3 presents the problem formulation, the proposed methodology to compare product manifolds, as well as how the search space over which to perform geometry-informed Bayesian optimization is constructed. Finally, Section 4 explains how our custom
synthetic and real-world datasets were obtained, and the empirical results. Lastly, in Section 5 we conclude and discuss avenues for future work. $^2$
Manifold Learning Manifold learning is a sub-field of machine learning that uses tools from differential geometry to model high-dimensional datasets by mapping them to a low-dimensional latent space. This allows researchers to analyze the underlying structure of data and improve machine learning models by capturing the geometry of the data more accurately. Manifold learning is based on the assumption that most observed data can be encoded within a low-dimensional manifold (see Figure 1) embedded in a high-dimensional space [Fefferman et al., 2013]. This has seen applications in dimensionality reduction [Roweis and Saul, 2000, Tenenbaum et al., 2000], generative models [Goodfellow et al., 2014, Du et al., 2021, Bortoli et al., 2022], and graph structure learning for graph neural networks (GNNs) [Topping et al., 2021]. In all these application, the key is to find a topological representation as an abstract encoding that describes the data optimally for the downstream task.
Product Manifolds. In this work, we model the geometry using model space Riemannian manifolds (Appendix A.1) and Cartesian products of such manifolds. The three so-called model spaces with constant curvature are the Euclidean plane,  $\mathbb{E}^n = \mathbb{E}_{K_{\mathbb{E}}}^{d_{\mathbb{E}}} = \mathbb{R}^{d_{\mathbb{E}}}$ , where the curvature  $K_{\mathbb{E}} = 0$ ; the hyperboloid,  $\mathbb{H}^n = \mathbb{H}_{K_{\mathbb{H}}}^{d_{\mathbb{H}}} = \{\mathbf{x}_p \in \mathbb{R}^{d_{\mathbb{H}} + 1} : \langle \mathbf{x}_p, \mathbf{x}_p \rangle_{\mathcal{L}} = 1 / K_{\mathbb{H}}\}$ , where  $K_{\mathbb{H}} < 0$  and  $\langle \cdot, \cdot \rangle_{\mathcal{L}}$  is the Lorentz inner product; and the hypersphere,  $\mathbb{S}^n = \mathbb{S}_{K_{\mathbb{S}}}^{d_{\mathbb{S}}} = \{\mathbf{x}_p \in \mathbb{R}^{d_{\mathbb{S}} + 1} : \langle \mathbf{x}_p, \mathbf{x}_p \rangle_2 = 1 / K_{\mathbb{S}}\}$ , where  $K_{\mathbb{S}} > 0$  and  $\langle \cdot, \cdot \rangle_2$  is the standard Euclidean inner product. These have associated exponential maps and distance functions with closed form solutions, which can be found in Appendix A.2. A product manifold can be constructed using the Cartesian product  $\mathcal{P} = \times_{i=1}^{n_{\mathcal{P}}} \mathcal{M}_{K_i}^{d_i}$  of  $n_{\mathcal{P}}$  manifolds with curvature  $K_i$  and dimensionality  $d_i$ . Note that both  $n_{\mathcal{P}}$  and  $d_i$  are hyperparameters that define the product manifold  $\mathcal{P}$  and that must be set a priori. On the other hand, the curvature of each model space  $K_i$  can be learned via
gradient descent. One must note that the product manifold construction makes it possible to generate more complex embedding spaces than the original constant curvature model spaces, but it does not allow to generate any arbitrary manifold nor to control local curvature.
Figure 1: Schematic of a manifold  $\mathcal{M}$  and open subsets  $U_{i}$  and  $U_{j}$ . An open chart is a homeomorphism of an open subset of the manifold onto an open subset of the Euclidean hyperplane. Here,  $\psi_{ij}$  is a transition function.
Hausdorff and Gromov-Hausdorff Distances for Comparing Manifolds. The Hausdorff distance between two subsets of a metric space refers to the greatest distance between any point on the first set and its closest point on the second set [Jungeblut et al., 2021]. Given a metric space  $X$  with metric  $d_X$ , and two subsets  $A$  and  $B$ , we can define the Hausdorff distance between  $A$  and  $B$  in  $X$  by
$$
\mathrm {d} _ {\mathrm {H}} ^ {X} (A, B) = \max  \left(\sup  _ {a \in A} d _ {X} (a, B), \sup  _ {b \in B} d _ {X} (b, A)\right). \tag {1}
$$
A priori this quantity may be infinite. Hence we will restrict to compact subsets  $A$  and  $B$ . In this case, we can equivalently define  $\mathrm{d_H}(A,B)$  as the smallest real number  $c\geqslant 0$  such that for every  $a\in A$  and every  $b\in B$  there exist  $a^\prime \in A$  and  $b^{\prime}\in B$  such that both  $d_{X}(a,b^{\prime})$  and  $d_{X}(a^{\prime},b)$  are at most  $c$ .
We note that the previous definition does not require any differentiable structures on  $X$ ,  $A$  and  $B$ . They can be merely metric spaces. This generality allows us to distinguish the metric properties of Euclidean, hyperbolic and spherical geometries beyond analytic notions such as curvature. However, the definition in Equation 1 only allows to compare spaces  $A$  and  $B$  that are embedded in a certain metric space  $X$ . The notion of distance that we shall consider is the Gromov-Hausdorff distance, which we define below in Equation 3.
Given a metric space  $X$  and two isometric embeddings  $f:A\to X$  and  $g:B\to X$ , we define
$$
\mathrm {d} _ {\mathrm {H}} ^ {X, f, g} (A, B) = \mathrm {d} _ {\mathrm {H}} ^ {X} (f (A), g (B)). \tag {2}
$$
2For a high level description of the proposed framework, we recommend skipping to Section 3
Now, given two metric spaces  $A$  and  $B$ , we denote by  $\operatorname{ES}(A, B)$  (standing for "embedding spaces of  $A$  and  $B$ ) as the triple  $(X, f, g)$  where  $X$  is a metric space and  $f: A \to X$  and  $g: B \to X$  are isometric embeddings. We define the Gromov-Hausdorff distance between  $A$  and  $B$  as:
$$
\mathrm {d} _ {\mathrm {G H}} (A, B) = \inf  _ {(X, f, g) \in \operatorname {E S} (A, B)} \mathrm {d} _ {\mathrm {H}} ^ {X, f, g} (f (A), g (B)). \tag {3}
$$
We should note that, since we assume that both  $A$  and  $B$  are compact, there is a trivial upper bound for their Gromov-Hausdorff distance in terms of their diameters. The diameter of a metric space  $Y$  is defined to be  $\mathrm{diam}(Y) = \sup_{y,y' \in Y} d_Y(y,y')$ . Given  $a_0 \in A$  and  $b_0 \in B$ , we can define the isometric embeddings  $f: A \to A \times B$  and  $g: B \to A \times B$  given by  $f(a) = (a,b_0)$  and  $g(b) = (a_0,b)$ . It is easy to see that  $\mathrm{d_H^{A \times B,f,g}}(A,B) \leqslant \max(\mathrm{diam}(A), \mathrm{diam}(B))$ . Since the triple  $(A \times B,f,g)$  belongs to  $\operatorname{ES}(A,B)$ , we can estimate  $\mathrm{d_{GH}}(A,B) \leqslant \max(\mathrm{diam}(A), \mathrm{diam}(B))$ . To compare  $\mathbb{E}^n$ ,  $\mathbb{H}^n$  and  $\mathbb{S}^n$ , we propose taking closed balls of radius one in each space. Since balls of radius one in any of these spaces are homogeneous Riemannian manifolds, they are isometric to each other. By estimating or providing an upper bound for their Gromov-Hausdorff distance, we can compare the spaces. This notion of distance between candidate latent geometries will later be used to generate a search space for our framework. With exactly an analogous argument as before, we can notice that given two compact balls of radius one  $B$  and  $B'$ , of centres  $x_0$  and  $x_0'$ , we can embed  $B$  into  $B \times B'$  by the mapping  $f: b \mapsto (b,x_0')$ . From here, it is obvious to see that  $\mathrm{d_H^{B \times B',f,id}}(B,B \times B') = 1$ . In particular, this gives us the bound  $\mathrm{d_{GH}}(B,B') \leqslant 1$ . This is the estimation we will take as the Gromov-Hausdorff distances of product manifolds that simply differ in one coordinate (such as, say,  $\mathbb{E}^2$  and  $\mathbb{E}^2 \times \mathbb{H}^2$ ).
3 Neural Latent Geometry Search: Latent Product Manifold Inference
In this section, we leverage ideas discussed in Section 2 to introduce a principled way to find the optimal latent product manifold. First, we introduce the problem formulation of NLGS. Next, we outline the strategy used to compute the Gromov-Hausdorff distance between product manifolds, and we discuss how this notion of similarity can be used in practice to construct a graph search space of latent geometries. Lastly, we explain how the Gromov-Hausdorff-informed graph search space can be used to perform NLGS via BO.
The problem of NLGS can be formulated as follows. Given a search space  $\mathfrak{G}$  denoting the set of all possible latent geometries, and the objective function  $L_{T,A}(g)$  which evaluates the performance of a given geometry  $g$  on a downstream task  $T$  for a machine learning model architecture  $A$ , the objective is to find an optimal latent geometry  $g^{*}$ :
$$
g ^ {*} = \underset {g \in \mathfrak {G}} {\arg \min } L _ {T, A} (g). \tag {4}
$$
In our case we model the latent space geometry using product manifolds. Hence we effectively restrict Equation 4 to finding the optimal product manifold signature:
$$
n _ {\mathcal {P}} ^ {*}, \left\{d _ {i} \right\} _ {i \in n _ {\mathcal {P}} ^ {*}} ^ {*}, \left\{K _ {i} \right\} _ {i \in n _ {\mathcal {P}} ^ {*}} ^ {*} = \underset {n _ {\mathcal {P}} \in \mathbb {Z}, d _ {i} \in \mathbb {Z}, K _ {i} \in \mathbb {R}} {\arg \min } L _ {T, A} \left(n _ {\mathcal {P}}, \left\{d _ {i} \right\} _ {i \in n _ {\mathcal {P}}}, \left\{K _ {i} \right\} _ {i \in n _ {\mathcal {P}}}\right), \tag {5}
$$
where  $n_{\mathcal{P}}$  is the number of model spaces composing the product manifold  $\mathcal{P}$ ,  $\{d_i\}_{i \in n_{\mathcal{P}}}$  are the dimensions of each of the model spaces of constant curvature, and  $\{K_i\}_{i \in n_{\mathcal{P}}}$  their respective curvatures. We further simplify the problem by setting  $d_i = 2$ ,  $\forall i$ , and by restricting ourselves to  $K_i \in \{-1,0,1\}$ , in order to limit the size of the hypothesis space.
3.2 Quantifying the Difference Between Product Manifolds
Motivation. From a computational perspective, we can think of the Hausdorff distance as a measure of dissimilarity between two point sets, each representing a discretized version of the two underlying continuous manifolds we wish to compare. Taha and Hanbury [2015] proposed an efficient algorithm to compute the exact Hausdorff distance between two point sets with nearly-linear complexity leveraging early breaking and random sampling in place of scanning. However, the original algorithm assumes that both point sets live in the same space and have the same dimensionality, which is a
limiting requirement. Gromov-Hausdorff distances, as opposed to usual Hausdorff distances, allow us to measure the distance between two metric spaces that a priori are not embedded in a common bigger ambient space. However, this has the caveat that they are not computable. For instance, for our application we must calculate the distance between each pair of the following three spaces:  $\mathbb{E}^n$ ,  $\mathbb{S}^n$  and  $\mathbb{H}^n$ . However,  $\mathbb{S}^n$  does not isometrically embed in  $\mathbb{E}^n$  and hence in order to compare  $\mathbb{E}^n$  and  $\mathbb{S}^n$ , we must work in a higher dimensional ambient space such as  $\mathbb{E}^{n + 1}$ . In the case of  $\mathbb{H}^n$ , finding an embedding to a Euclidean space is more complicated and is described in Appendix B.2 ( $\mathbb{H}^n$  will be embedded isometrically into  $\mathbb{E}^{6n - 6}$ ). However, in this process there will be choices made about which embeddings to consider (in particular, we cannot exactly compute the infimum that appears in the definition of Gromov-Hausdorff distance in Equation 3). Likewise, using the original algorithm by Taha and Hanbury [2015] it is not possible to compute the Hausdorff distance between product manifolds based on an unequal number of model spaces, for instance, there is no way of computing the distance between  $\mathbb{E}^n$  and  $\mathbb{E}^n \times \mathbb{H}^n$ . In this section we give an upper bound for  $\mathrm{d}_{\mathrm{GH}}(\mathbb{E}^n,\mathbb{S}^n)$  and then describe an algorithm to give an upper bound for  $\mathrm{d}_{\mathrm{GH}}(\mathbb{E}^n,\mathbb{H}^n)$  and  $\mathrm{d}_{\mathrm{GH}}(\mathbb{S}^n,\mathbb{H}^n)$ .
Strategy. The spaces  $\mathbb{E}^n$  and  $\mathbb{S}^n$  isometrically embed into  $\mathbb{E}^{6n - 6}$  in many ways. This may seem redundant because both spaces already embed in  $\mathbb{E}^{n + 1}$ . However, the interest of considering this higher dimensional Euclidean space is that  $\mathbb{H}^n$  will also isometrically embed into it. Crucially, this will provide a common underlying space in which to compute Hausdorff distances between our geometries  $\mathbb{E}^n$ ,  $\mathbb{S}^n$  and  $\mathbb{H}^n$ , which will lead to an estimation of their mutual Gromov-Hausdorff distance. The embedding of  $\mathbb{H}^n$  into  $\mathbb{E}^{6n - 6}$  that we shall describe appears in Henke and Nettekoven [1987] and is made explicit in [Blanuša, 1955]. We also refer the reader to the exposition [Brander, 2003, Chapter 5], which puts this results in a broader context while also summarising related advances on the topic of isometrically embedding homogeneous spaces into higher dimensional ones. For  $n = 2$ , we name this embedding  $F:\mathbb{H}^2\to \mathbb{E}^6$  (Appendix B.2). For simplicity, in our experiments in Section 4, we will work with product manifolds generated based on constant curvature model spaces of dimension  $n = 2$ .
Now we can summarise our strategy to estimate  $\mathrm{d}_{\mathrm{GH}}(B_{\mathbb{E}^2}, B_{\mathbb{H}^2})$  as follows (for  $\mathrm{d}_{\mathrm{GH}}(B_{\mathbb{H}^2}, B_{\mathbb{S}^2})$  it will be entirely analogous). The first step consists of approximating our infinite smooth spaces by finite discrete ones. For this, we consider several collections of points  $\{P_i\}_{i \in I}$  in  $\mathbb{E}^2$  that are sufficiently well distributed. The exponential map can be applied to the collection of points  $\exp : T_0\mathbb{H}^2 \cong \mathbb{R}^2 \to B_{\mathbb{H}^2}$  to get several collections of points  $Q$  in  $B_{\mathbb{H}^2}$  (again, well distributed by construction). In addition, we will consider several isometric embeddings  $f_k : B_{\mathbb{E}^2} \to \mathbb{R}^6$ . Hence, we take
$$
\mathrm {d} _ {\mathrm {G H}} \left(B _ {\mathbb {E} ^ {2}}, B _ {\mathbb {H} ^ {2}}\right) \approx \min  _ {i, j, k} \mathrm {d} _ {\mathrm {H}} ^ {\mathbb {R} ^ {6}, f _ {k}, F} \left(P _ {i}, Q _ {j}\right) = \min  _ {i, j, k} \mathrm {d} _ {\mathrm {H}} ^ {\mathbb {R} ^ {6}} \left(f _ {k} \left(P _ {i}\right), F \left(Q _ {j}\right)\right). \tag {6}
$$
In Appendix B, we gradually unravel the previous formula and give explicit examples of the involved elements. In particular, in Appendix B.1 we explain how to generate points in the balls of radius one, and in Appendix B.2 how to describe the isometric embedding  $F: B_{\mathbb{H}^2} \to \mathbb{E}^6$ . The results obtained for the Gromov-Hausdorff distances between product manifolds are used to generate the graph search space introduced in the next section.
3.3 The Gromov-Hausdorff-Informed Graph Search Space
Gromov-Hausdorff Edge Weights. In Section 2, we used Cartesian products of constant curvature model spaces to generate candidate latent geometries. We now turn our attention to constructing a search space to find the optimal latent geometry. To do so, we first consider all possible combinations of product manifolds based on a given number of model spaces, represented by  $n_s$ . Furthermore, we denote the total number of products (model spaces) used to form the product manifold  $\mathcal{P}$  with  $n_p$ . Conceptually,  $n_s$  is the number of model space types used, while  $n_p$  refers to the overall number, or quantity, of model spaces that form the resulting product space. For instance, if only the Euclidean plane, the hyperboloid, and the hypersphere
Table 1: Estimated Gromov-Hausdorff distances (up to two decimal places) between model spaces and corresponding edge weights in the graph search space.
Comparison Pair	dGH(·)	w(·)
(E2, S2)	0.23	4.35
(E2, H2)	0.77	1.30
(S2, H2)	0.84	1.20
are taken into account, then  $n_s = 3$ . If all product mani
fold combinations are considered, the number of elements in the search space increases to  $\sum_{i=1}^{n_p} n_s^i$ . However, we assume commutativity for latent product manifolds, implying that the output of a trained neural network with a latent geometry  $\mathcal{M}_i \times \mathcal{M}_j$  should be the same as that with  $\mathcal{M}_j \times \mathcal{M}_i$ . We refer to this concept as the symmetry of parameterization, as neural networks can rearrange the weight matrices of their neurons to achieve optimal performance for two equivalent latent manifolds. This assumption reduces the search space from growing exponentially to  $\mathcal{O}(n_p^2)$ , assuming three constant curvature model spaces are considered (see Appendix B.5 for a more complete explanation).