SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shuai Zhang
cheungshuai@outlook.com
Wenqi Jiang
wenqi.jiang@inf.ethz.ch
Geometric representation learning (e.g., hyperbolic and spherical geometry) has proven to be efficacious in solving many intricate machine learning tasks. The fundamental challenge of geometric representation learning lies in aligning the inherent geometric bias with the underlying structure of the data, which is a rarely explored topic in the literature. Existing methods heavily rely on heuristic assumptions on the data structure to decide the type of geometry to be adopted, which often leads to suboptimal performance. This work aims to automate the alignment process via a data-informed strategy such that we optimize model performance with minimal overhead. Specifically, a sparse gating mechanism is employed to enable each input data point  $p$  to select  $K$  geometric spaces from a given candidate geometric space pool with  $N$  ( $K < N$ ) spaces of different geometry. The selected  $K$  spaces are then tightly integrated to formulate a Cartesian product space, which is leveraged to process this input data  $p$ . In doing so, each input data is processed by the spaces it selected with maximum specialization. We empirically show that this method can effectively align data and spaces without human interventions and further boost performance on real-world tasks, demonstrating its potential in eliciting the expressive power of geometric representations and practical usability.
Representing data in low-dimensional geometric spaces (e.g., Euclidean, hyperbolic, and spherical) has been an invaluable technique in a variety of machine learning tasks. Euclidean space has been the most widely adopted and predominant geometric space across many tasks [8, 31]. Recent research has shown that non-Euclidean geometry, such as hyperbolic and spherical geometry, can be more effective in representing data with specific/intricate inherent structures. Typically, hyperbolic space, a type of manifold with constant negative (sectional) curvature, provides a greater ability to model tree-like structures and capture hierarchical patterns [32, 11, 21]. Conversely, spherical space, with positive curvature, is well-suited for fitting data with cyclical structures [44, 30]. Different geometry endows models with different inductive bias that prioritizes learning specific patterns in the data.
Geometric representation learning faces a significant challenge in effectively aligning underlying data structures with appropriate geometric spaces. Earlier approaches tend to map all data to a single geometric space, operating under the implicit assumption that the majority, if not all, of the data conforms to the same geometric pattern. This assumption of a uniform geometric pattern across all data points is unrealistic. As such, there has been growing interest in hybridizing different geometric spaces [22, 37, 4] via Cartesian product to form a mixed curvature space. By leveraging the benefits of multiple geometric spaces, this approach can offer improved performance compared to methods that rely on a single geometric space.
Real-world data is usually characterized by intricate patterns that a single geometric space cannot adequately capture. Thus, employing a one-space-fits-all approach is unlikely to produce optimal solutions. In addition, the Cartesian product space representation learning technique does not discriminate between input examples based on their underlying geometric structures. For example,
if a Cartesian product space is constructed by combining a spherical space with a hyperbolic space, both spaces will be used to fit all input data examples, irrespective of whether a data point only aligns with one of the geometry. Ideally, each space should calibrate its internal parameters using the input examples that conform to the corresponding manifold to ensure that the modeling capacity of all geometric spaces is maximized.
Accurately aligning input data with appropriate geometric spaces is a crucial challenge in geometric representation learning. However, neither the single geometric space method nor the Cartesian product approach fully address this issue. Our primary goal is to tackle this challenge by developing an automated approach that seamlessly aligns each data point with its appropriate space. Instead of relying on oversimplified heuristics or manual interventions, we aim to learn the data-space alignment component together with the task-specific model using an end-to-end framework.
To this end, we propose a data-informed method for geometric space selection. Given  $N$  candidate geometric spaces with different geometry, we let each data point  $p$  choose  $K(1 \leq K < N)$  spaces from the pool. The chosen  $K$  spaces are then tightly combined via Cartesian product to form a Cartesian product space, which is used to process  $p$ . The selection process is parameterized by a sparsely gated neural network that can adaptively activate a subset of the given spaces based on the given input data. This module can be easily plugged into any geometric representation learning task to enable end-to-end optimization. Theoretically, we can construct  $_{N}C_{K} = \frac{N!}{(N - K)!K!}$  possible Cartesian product spaces in the end, which brings a greater degree of freedom and personalization for representation learning. Moreover, the proposed method will only incur a negligible computational overhead since only the  $K$  selected geometric spaces participate in the model inference no matter how large the space pool is.
The contributions of this work are summarized as follows:
- We propose a data-informed approach for geometric space selection that offers the advantage of automatically aligning each input data point with its preferred geometric space, thereby increasing specificity while being computationally efficient.
- We empirically demonstrate the practical effectiveness of this method with real-world datasets across two downstream tasks. We examine the proposed method's inner workings (e.g., distribution of the selected geometric spaces) via visualization, ablation, and case studies, reaffirming the benefits of the proposed method.
We briefly review two areas of active research that intersect with the present work.
2.1 Non-Euclidean Geometric Representation Learning
Two popular non-Euclidean geometric spaces with non-zero curvature, including hyperbolic and spherical spaces, have gained increasing attention and can be found in a broad spectrum of representation learning applications. Hyperbolic space has negative curvature and is reminiscent of a continuous version of trees, which makes it excel in capturing hierarchical structures. A bunch of work has demonstrated its effectiveness in tasks such as graph modeling and knowledge representation [32, 42, 39, 5, 12, 35, 28, 10, 15, 20, 9, 26, 40]. On the other hand, spherical space with positive curvature is more suitable for directional similarity modeling and cyclical-structured data. We can find applications of spherical spaces in text embeddings [30], texture mapping [44], time-warping functions embedding [44]. In these papers, the choice of geometric space is primarily based on expert heuristics about the data characteristics.
To combine the best of different worlds, Gu et al. [22, 37, 4] proposed a method to construct a mixed curvature space via Cartesian product of different spaces. The component spaces in Cartesian product spaces are selected from Euclidean, hyperbolic, and spherical spaces. Each component of Cartesian product space has constant curvature, while the resulting mixed space has a non-constant curvature, which enables it to capture a wide range of patterns with a single model. Empirically, Cartesian product spaces demonstrate its efficacy in graph reconstruction, word embedding with low dimensions [22], node classification in graphs [4], and image reconstruction [37]. Product space is designed to handle data with mixed patterns, but the alignment process between data and geometric
space is still oversimplified and not customized for each data point. Consequently, the challenge of aligning appropriate spaces with input data points persists, and a solution to this problem remains elusive.
2.2 Sparsely-gated Mixture of Experts
Another related research area is sparsely-gated neural networks, which are a type of mixture-of-experts (MOE) architecture [24, 25]. Over the past few decades, many MOE architectures have been proposed, such as hierarchical structures [47], sequential experts [2], deep MOE [18], and sparsely-gated MOE [36]. MOE is based on the divide-and-conquer principle, which divides the problem into homogeneous regions and assigns an expert to each region. The final prediction is made by combining the predictions of all experts via a gating network. The gating network controls the contribution of each expert using a probabilistic gating function. If the gating function is sparse, i.e., restricted to assigning non-zero weights to only a subset of the experts, unused experts need not be computed. MOE has been widely studied in various tasks, such as multi-task learning [29], large language models [36, 19], and scaling up vision transformers [34].
We draw inspiration from sparsely gated MOE to automate the alignment between data and its appropriate geometric space. Our approach allows each input data point to control the gating function and select suitable geometric spaces from a candidate geometric space pool, thus enabling customized assignment of data to its suitable spaces. It brings higher specialization, a greater degree of freedom, and superior expressiveness.
3.1 Riemannian Geometric Space
Stereographic Projection model. In general, there are three types of constant curvature spaces with respect to the sign of the curvature (Fig. 1). Common realizations are Euclidean space  $\mathbb{E}$  (flat), hypersphere  $\mathbb{S}$  (positively curved) and hyperboloid  $\mathbb{H}$  (negatively curved). For the latter two, we prefer their stereographic projection model: projected sphere  $\mathbb{D}$  and Poincare ball  $\mathbb{P}$ . These models are easier to optimize and avoid the problem of nonconvergence of norm of points with a curvature close to 0, and the projection is conformal, i.e., does not affect the angles between points [32, 37].
Figure 1: Left: Hyperbolic space  $\mathbb{H}$ ; Center: Euclidean plane  $\mathbb{E}$ ; Right: Spherical space  $\mathbb{S}$ .
An alternative to vector space in non-Euclidean geometry is gyrovector space [41], which defines operations such as vector addition and multiplication.
Definition 1 For  $\mathbb{D}$  and  $\mathbb{P}$  (jointly denoted as  $\mathcal{M}_c$ , where  $c$  denotes curvature), the addition between two points  $\mathbf{x}, \mathbf{y} \in \mathcal{M}_c$ , also known as Möbius addition  $\oplus_c$  (for both signs of  $c$ ), is defined as:
$$
\mathbf {x} \oplus_ {c} \mathbf {y} = \frac {\left(1 - 2 c \langle \mathbf {x} , \mathbf {y} \rangle - c \| \mathbf {y} \| _ {2} ^ {2}\right) \mathbf {x} + \left(1 + c \| \mathbf {x} \| _ {2} ^ {2}\right) \mathbf {y}}{1 - 2 c \langle \mathbf {x} , \mathbf {y} \rangle + c ^ {2} \| \mathbf {x} \| _ {2} ^ {2} \| \mathbf {y} \| _ {2} ^ {2}}, \tag {1}
$$
Where  $\langle ,\rangle$  is Euclidean inner product. The distance between points in the gyrovector space is defined as:
$$
d _ {\mathcal {M} _ {c}} (\mathbf {x}, \mathbf {y}) = \frac {2}{\sqrt {| c |}} \tan_ {c} ^ {- 1} \left(\sqrt {| c |} \| - \mathbf {x} \oplus_ {c} \mathbf {y} \| _ {2}\right), \tag {2}
$$
where  $\tan_c$  stands for  $\tan$  if  $c > 0$  and  $\tanh$  if  $c < 0$ . In both spaces, we have Euclidean geometry when  $c \to 0$ . It is easy to prove that:
$$
d _ {\mathcal {M} _ {c}} (\mathbf {x}, \mathbf {y}) \xrightarrow {c \rightarrow 0} 2 \| \mathbf {x} - \mathbf {y} \| _ {2}, \tag {3}
$$
which means the gyroscope distance converges to Euclidean distance when limiting  $c$  to zero.
Definition 2 Let  $\mathcal{T}_{\mathbf{x}}\mathcal{M}_c$  be the tangent space to the point  $\mathbf{x} \in \mathcal{M}_c$ . Mapping between (Euclidean) tangent space and hyperbolic/spherical space is performed with exponential map:  $\mathcal{T}_{\mathbf{x}}\mathcal{M}_c \to \mathcal{M}_c$  and logarithmic map:  $\mathcal{M}_c \to \mathcal{T}_{\mathbf{x}}\mathcal{M}_c$ , which are defined as:
$$
\log_ {\mathbf {x}} ^ {c} (\mathbf {y}) = \frac {2}{\sqrt {| c | \lambda_ {\mathbf {x}} ^ {c}}} \tan_ {c} ^ {- 1} \left(\sqrt {| c |} \| - \mathbf {x} \oplus_ {c} \mathbf {y} \| _ {2}\right) \frac {- \mathbf {x} \oplus_ {c} \mathbf {y}}{\| - \mathbf {x} \oplus_ {c} \mathbf {y} \| _ {2}},
$$
$$
\exp_ {\mathbf {x}} ^ {c} (\mathbf {v}) = \mathbf {x} \oplus_ {c} \left(\tan_ {c} \left(\sqrt {| c |} \frac {\lambda_ {\mathbf {x}} ^ {c} \| \mathbf {v} \| _ {2}}{2}\right) \frac {\mathbf {v}}{\sqrt {| c |} \| \mathbf {v} \| _ {2}}\right),
$$
where  $\lambda_{\mathbf{x}}^{c}$  is a conformal factor, defined as  $\lambda_{\mathbf{x}}^{c} = 2 / (1 + c\| \mathbf{x}\|_{2}^{2})$  , which is used to transform the metric tensors between (Euclidean) tangent space and non-Euclidean space.
3.2 Cartesian Product of Spaces
Definition 3 Product space is defined as the Cartesian product of multiple spaces with varying dimensionality and curvature. Let  $\mathcal{P}$  denote a product space composed by  $N$  independent component spaces  $\mathcal{M}^{(1)},\mathcal{M}^{(2)},\dots,\mathcal{M}^{(N)}$ . The mixed space  $\mathcal{P}$  has the form (operator  $\times$  can be omitted):
$$
\mathcal {P} = \bigotimes_ {i = 1} ^ {N} \mathcal {M} ^ {(i)} = \mathcal {M} ^ {(1)} \times \mathcal {M} ^ {(2)} \times \dots \times \mathcal {M} ^ {(N)}. \tag {5}
$$
Definition 4 The product space  $\mathcal{P}$  also has distance functions. The squared distance between points  $\mathbf{x},\mathbf{y}\in \mathcal{P}$  is defined as:
$$
d _ {\mathcal {P}} ^ {2} (\mathbf {x}, \mathbf {y}) = \sum_ {i = 1} ^ {N} d _ {\mathcal {M} ^ {(i)}} ^ {2} \left(\mathbf {x} _ {\mathcal {M} ^ {(i)}}, \mathbf {y} _ {\mathcal {M} ^ {(i)}}\right), \tag {6}
$$
where  $\mathbf{x}_{\mathcal{M}^{(i)}}$  and  $\mathbf{y}_{\mathcal{M}^{(i)}}$  denote the corresponding vectors on the component space  $\mathcal{M}^{(i)}$ .
Other operations such as exponential map and logarithmic map are element-wise, meaning that we can decompose the points into component spaces, apply operations on each component space, and then compose them back (e.g., concatenation) to the product space. A product space's signature (i.e., parametrization) refers to the types of space, the number of spaces, the dimensionality, and the curvature of each space. For instance,  $(\mathbb{P}^{100})^3 (\mathbb{D}^{50})^2$  denotes a Cartesian product space of three hyperbolic spaces of hidden dimension 100 and two spherical spaces of hidden dimension 50.
4 Data-Informed Space Selection
4.1 Geometric Space Pool
Suppose we have a space pool with  $N$  geometric spaces,  $\mathcal{M}_c^{(i)}\in \{\mathbb{E},\mathbb{D},\mathbb{P}\} ,i = \{1,\dots,N\}$ . For simplicity, we assume all spaces have the same dimensionality  $b$ . The goal is to select  $K(1\leq K < N)$  spaces from the given  $N$  spaces for each data point  $p$ . If  $K > 1$ , the selected  $K$  spaces will form a product space. Theoretically, there will be  $N C_K = \frac{N!}{(N - K)!K!}$  possible space combinations. Figure 2 illustrates the proposed data-informed space selection pipeline. Here,  $\mathbb{P}$  and  $\mathbb{E}$  are selected by data  $p$  and the Cartesian product  $\mathbb{P}\times \mathbb{E}$  is used to handle  $p$ .
Figure 2: Data-informed geometric space selection.  $\mathbb{P}$  and  $\mathbb{E}$  are selected to form  $\mathbb{P} \times \mathbb{E}$  for input data  $p$ .
4.2 Embeddings
For each geometric space  $\mathcal{M}_c^{(i)}$ , we define an embedding dictionary  $\mathbf{e}^{(i)}$ . In the  $i^{\mathrm{th}}$  space, data point  $p$  is represented by a vector  $\mathbf{e}_p^{(i)} \in \mathbb{R}^b$ , where  $b$  is the hidden dimension of this embedding dictionary.
We initialize all the embedding vectors in the tangent space, and the exponential map will be applied to recover them into hyperbolic or spherical space when necessary. In doing so, standard Euclidean optimization algorithms can be directly applied to enable an end-to-end learning framework to avoid the cumbersome Riemannian optimization [7].
4.3 Geometric Space Selector via Sparsely-gated MOE
We use a sparsely-gate MOE network to select a suitable space subset from the space pool. The selector is a differentiable network that takes as input the embedding vectors  $\mathbf{e}_p^{(i)}$  and outputs the probability of each space to be selected. Formally, the input is defined as:
$$
\mathbf {x} _ {p} = \left[ \mathbf {e} _ {p} ^ {(1)}, \mathbf {e} _ {p} ^ {(2)}, \dots , \mathbf {e} _ {p} ^ {(N)} \right] ^ {\top}, \tag {7}
$$
where  $\mathbf{x}_p$  is a two-dimensional matrix of shape  $N\times b$ . We use a convolutional neural network to transform this matrix into a vector of dimension  $N$ . Other types of neural networks are also viable, we adopt CNN because it does not introduce a lot of additional parameters. Then, the last hidden layer of the gating network has the following form:
$$
f \left(\mathbf {x} _ {p}\right) = f _ {1} \left(\mathbf {x} _ {p}\right) + \operatorname {r a n d n} () \cdot \ln (1 + \exp \left(f _ {2} \left(\mathbf {x} _ {p}\right)\right)), \tag {8}
$$
where  $f_{1}$  and  $f_{2}$  represent CNN layers and  $f_{*}(\mathbf{x}_{p})\in \mathbb{R}^{N}$ ; function "randn") is used to product Gaussian noise in the training stage to improve loading balancing, i.e., balance the number of samples accepted by each space.
To impose sparsity, we employ a TopK function to obtain the largest  $K$  elements from  $f(\mathbf{x}_p)$ :
$$
f \left(\mathbf {x} _ {p}\right) \leftarrow \operatorname {T o p K} \left(f \left(\mathbf {x} _ {p}\right)\right), \tag {9}
$$
where TopK returns the original value if the element is in the top  $K$  list; otherwise, it returns  $-\infty$ . Afterward,  $N$  softmax gating functions  $g_{1}, g_{2},.., g_{N}$  are leveraged to control which  $K$  space(s) to be selected. For  $-\infty$ , it will output a zero.
$$
g _ {i} \left(\mathbf {x} _ {p}\right) = \frac {\exp \left(f \left(\mathbf {x} _ {p}\right) _ {i}\right)}{\sum_ {j = 1} ^ {N} \exp \left(f \left(\mathbf {x} _ {p}\right) _ {j}\right)}, i = \{1, \dots , N \}. \tag {10}
$$
Later, spaces with nonzero gating values will be activated while others remain idle. The selected  $K$  space will form a Cartesian product space.
4.4 Task Specific Prediction Function
The prediction function is specific to the type of task. For instance, in each component space  $\mathcal{M}_c^{(i)}$ , we can adopt squared distance to produce the model prediction. Given two points  $p$  and  $q$ , the prediction function can be defined as  $s_{\mathcal{M}_c^{(i)}}(\mathbf{e}_p^{(i)},\mathbf{e}_q^{(i)})$ . If space  $i$  is not selected,  $s_{\mathcal{M}_c^{(i)}}$  is set to 0. More details will be provided in the experiment section for each task. If  $K > 1$ , the final prediction score is computed as:
$$
s (p, q) = \log \left(\sum_ {i = 1} ^ {N} g _ {i} \left(\mathbf {x} _ {p}\right) \exp \left(s _ {\mathcal {M} ^ {(i)}} \left(\mathbf {e} _ {p} ^ {(i)}, \mathbf {e} _ {q} ^ {(i)}\right)\right)\right). \tag {11}
$$
Here, a log-sum-exp technique is leveraged to improve the numerical stability to avoid problems such as underflow and overflow. Multiplying the prediction by the gating probability is optional in the framework.
If  $s_{\mathcal{M}^{(i)}}$  is calculated with squared distance and  $K = N$ , we can recover the squared distance function defined for product space by removing  $g_i(\mathbf{x})$  and the log-sum-exp technique. Our prediction function cannot be viewed as a standard distance metric as it does not satisfy the triangle inequality property. Also, it is flexible to use other prediction functions to replace squared distances.
4.5 Load Balancing Regularization
We use a load balancing regularization term to improve numerical stability and avoid some spaces from being over-selected. In specific, two additional regularization terms are used following [36].
For a batch of  $X$  inputs,
$$
\ell_ {1} = \psi \left(\sum_ {p \in X} g \left(\mathbf {x} _ {p}\right)\right) ^ {2}, \tag {12}
$$
where  $\psi (v) = \frac{\text{variance}(v)}{\text{mean}(v)}$  is the coefficient of variation.
$$
\ell_ {2} = \psi (\kappa (X)) ^ {2}, \tag {13}
$$
$$
\kappa (X) _ {i} = \sum_ {p \in X} \Phi \left(\frac {f _ {1} \left(\mathbf {x} _ {p}\right) - \xi \left(f \left(\mathbf {x} _ {p}\right) , k , i\right)}{\ln \left(1 + \exp \left(f _ {2} \left(\mathbf {x} _ {p}\right)\right)\right)}\right), \tag {14}
$$
where  $\xi (v,k,i)$  is the  $k^{\mathrm{th}}$  highest component of  $v$  excluding component  $i$ .  $\Phi$  is the cumulative distribution function of standard normal distribution.
The two regularization terms are added to the task-specific loss for model optimization:
$$
\ell = \ell_ {\text {t a s k}} + \mu_ {1} \ell_ {1} + \mu_ {2} \ell_ {2}, \tag {15}
$$
where  $\mu_{1}$  and  $\mu_{2}$  are scaling factors (set to 0.01 by default).
4.6 Limitation
Compared with Cartesian product space [22], the proposed method introduces one additional hyperparameter,  $K$ . As we will show later, a relatively small  $K$  (e.g.,  $K <= 5$ ) can introduce a significant performance boost. Just like determining the signature of Cartesian product space, we need to determine what geometric spaces to be included in the space pool. This overhead is inherited from Cartesian product space and can be potentially mitigated via hyper-parameters optimization algorithms [23].
In this section, we evaluate the proposed method on real-world datasets to demonstrate its capability in dealing with practical tasks, including personalized ranking and link prediction for relational graphs.
5.1 Personalized Ranking
The task of personalized ranking is to provide a user with a tailored, ranked list of recommended items [33]. Given a partially-observed interaction matrix between users and items, where each entry of the matrix represents the rating/preference a user  $u$  gives to an item  $v$  if  $u$  has interacted with the item (e.g., watched the movie) and is otherwise missing. We conduct experiments on two well-known datasets, MovieLens 100K and MovieLens 1M. The first dataset consists of 100K ratings between 943 users and 1,682 movies, and the second dataset consists of one million ratings between 6,040 users and 3,706 movies. Over  $90\%$  of the entries of these interaction matrices are missing. Completing this matrix enables us to recommend to customers what to watch next.
The personalized ranking model takes a user vector  $\mathbf{u}_u^{(i)}\in \mathbb{R}^b$  and an item vector  $\mathbf{v}_v^{(i)}\in \mathbb{R}^b$  as input, and approximate the rating with squared distance [42]. That is, the rating of user  $u$  gives to movie  $v$  is estimated by:
$$
s _ {\mathcal {M} ^ {(i)}} (u, v) = - d _ {\mathcal {M} ^ {(i)}} ^ {2} \left(\exp_ {\mathbf {0}} ^ {c} \left(\mathbf {u} _ {u} ^ {(i)}\right), \exp_ {\mathbf {0}} ^ {c} \left(\mathbf {v} _ {v} ^ {(i)}\right)\right). \tag {16}
$$
Alternatively, we can use the dot product,  $\mathbf{u}_u^{(i)}\cdot \mathbf{u}_v^{(i)}$ , to approximate each entry of the matrix [33].
The input of the gating network is the concatenation of the user and item embeddings. Specifically, a vector of size  $\mathbb{R}^{2Nb}$  is used as the input, and a linear layer is used as the sparsely-gated network. Let  $s(u,v)$  denote the final prediction. We treat the task as a learning-to-rank problem and optimize it by minimizing the following contrastive loss:
$$
\ell_ {\text {t a s k}} = \sum_ {(u, v _ {1}, v _ {2}) \in \Omega} \max  (0, s (u, v _ {1}) + m - s (u, v _ {2})), \tag {17}
$$
Figure 3: Personalized ranking performance on MovieLens 1M and MovieLens 100K.
where  $\Omega$  is the collection of training samples,  $v_{1}$  is the movie that  $u$  watched, and  $v_{2}$  represents an unobserved movie for the user;  $m$  is a hyper-parameter.