SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We hold  $70\%$  entries in each user's interactions as the training set,  $10\%$  entries as the validation set for model tuning, and the remaining  $20\%$  for model testing. All interactions (e.g., ratings) are binarized following the implicit feedback setting [33]. For all methods, the total dimension  $\sum_{i=1}^{N} b_i$  is set to 100 for a fair comparison to ensure the same model size. The curvatures for spherical and hyperbolic models are set to 1 and  $-1$ , respectively.  $N$  is set to 5.  $K$  is tuned among  $\{1,2,3,4\}$ . Regularization rate is chosen from  $\{0.1,0.01,0.001\}$ .  $m$  is fixed to 0.5. Adam is adopted as the optimizer. We measure the performance based on the widely adopted metrics in personalized ranking: mean average precision (MAP), precision, and recall at the top-ranked ten items (Precision@10 and Recall@10). For all experiments, we report the average over five runs.
Main Results. Figure 3 summarizes the performances. The proposed method achieves the best scores on all three metrics, outperforming the product space counterpart with the same signature and other single space models. Also, dot-product based approach remains competitive in this task. We estimate the global average curvature with the algorithm described in [22] and obtain a value of 0.695 for MovieLens 1M and 0.19 for MovieLens 100K, which suggests that they lean towards cyclical structures. However, from the performance comparison among  $\mathbb{E}^{100}$ ,  $\mathbb{D}^{100}$ , and  $\mathbb{P}^{100}$ , we find that the model built in Euclidean space offers the best performance, outperforming spherical and hyperbolic counterparts. This observation indicates that the global average curvature estimation algorithms proposed by Gu et al. [22] may give misleading information, which also ascertains the difficulty of data-space alignment.
5.2 Link Prediction for Relational Graphs
Relational graphs have emerged as an effective way to integrate disparate data sources and model underlying relationships. Encoding the nodes and relations of relational graphs into low-dimensional vector spaces is vital to downstream applications such as missing facts completion, question answering, information extraction, and logical reasoning [3].
Given a relational graph  $\mathcal{G}$  with a set of entities  $\mathcal{E}$  and a set of relations  $\mathcal{R}$ . Each triple, abbreviated as  $(h,r,t)$ , in  $\mathcal{G}$  is composed by two entities (i.e., head entity  $h\in \mathcal{E}$  and tail entity  $t\in \mathcal{E}$ ), and the relationship  $r\in \mathcal{R}$  between them. In the  $i^{\mathrm{th}}$  space, entities  $h,t$  are represented by vectors  $\mathbf{e}_h^{(i)},\mathbf{e}_t^{(i)}\in \mathbb{R}^b$  and relation  $r$  is represented by two translation vectors  $\alpha_r^{(i)},\beta_r^{(i)}\in \mathbb{R}^b$  and a rotation vector  $\gamma_r^{(i)}\in \mathbb{R}^b$ . Also, each head (tail) entity is associated with a bias term  $b_{h}(b_{t})\in \mathbb{R}$ .
Inspired by the method (RotE and RotH) proposed in [11], we propose a generic model, RotX, that subsumes the two methods and can be extended to spherical space. In RotX, the head entity is translated twice via Möbius addition and rotated once. Formally, the head entity is processed as follows:
$$
Q ^ {(i)} (h, r) = \operatorname {R O T A T E} \left(\exp_ {\mathbf {0}} ^ {c} \left(\mathbf {e} _ {h} ^ {(i)}\right) \oplus_ {c} \exp_ {\mathbf {0}} ^ {c} \left(\boldsymbol {\alpha} _ {r} ^ {(i)}\right), \gamma_ {r} ^ {(i)}\right) \oplus_ {c} \exp_ {\mathbf {0}} ^ {c} \left(\boldsymbol {\beta} _ {r} ^ {(i)}\right), \tag {18}
$$
where  $\exp_0^c$  is the exponential map over origin. ROTATE is a rotation function, and  $\gamma_r^{(i)}$  is the rotation matrix. The transformed head entity is then compared with the tail entities using squared distance.
Figure 4: Performance on relational graphs link prediction. We use the performance of TransE-AT [46] as the baseline (e.g.,  $\Delta$  MRR = MRR(Model) - MRR(TransE-AT)) for the lollipop charts. In our method, for WN18RR, the geometric space pool consists of  $4 \times \mathbb{D}^{100}$  and one  $\mathbb{E}^{100}$ ,  $K = 4$ ; for FB15K-237, the geometric space pool consists of  $3 \times \mathbb{P}^{100}$  and  $2 \times \mathbb{D}^{100}$ ,  $K = 4$ .
The final scoring function is as follows:
$$
s _ {\mathcal {M} ^ {(i)}} (h, r, t) = - d _ {\mathcal {M} ^ {(i)}} ^ {2} \left(Q ^ {(i)} (h, r), \exp_ {\mathbf {0}} ^ {c} \left(\mathbf {e} _ {t} ^ {(i)}\right)\right) + b _ {h} ^ {(i)} + b _ {t} ^ {(i)}. \tag {19}
$$
We make the curvature relation specific and trainable. Each relation has a curvature parameter and is trained simultaneously with the model. We adopt the following cross-entropy loss as the objective function:
$$
\ell_ {\text {t a s k}} = \sum_ {(h, r, t) \in \Omega} \log (1 + \exp (- Y _ {(h, r, t)} s (h, r, t))), \tag {20}
$$
where  $Y_{(h,r,t)} \in \{1, - 1\}$  is a binary label indicating whether a triple is factual (1) or not (-1).  $\Omega$  represents the training collection, including positive and negative triples.
We use two datasets including WN18RR [8, 16] and FB15K-237 [8, 16] for model evaluation. WN18RR is taken from WordNet, a lexical database of semantic relations between words. It has 40,943 entities, 11 relations, and 86,835/3,034/3,134 training/validation/test triples. FB15K-237 is a subset of the Freebase knowledge graph, a global resource of common and general information. It has 14,541 entities, 237 relations, and 272,115/17,535/20,466 training/validation/test triples. The performance is evaluated using two standard metrics, including mean reciprocal rank (MRR) and hit rate (HR) with a given cut-off value  $\{1,3,10\}$ .
The total dimension is fixed to 500 for a fair comparison. Learning rate is tuned amongst  $\{0.01, 0.005, 0.001\}$ . For all experiments, we report the average over 5 runs. We set the kernel size to 5 and stride to 3 for convolution operation in the gating network.  $N$  is set to 5 and  $K$  is tuned among  $\{1,2,3,4\}$ . The number of negative samples (uniformly sampled) per factual triple is set to 50. Optimizer Adam is used for model learning. We perform early stopping if the validation MRR stops increasing after 10 epochs.
We compare our method with several baselines, including Euclidean methods TransE [8], DistMult [45], ConvE [16], TuckER [6], RotE [11], BoxE [1]; complex number based methods ComplExN3, [27] and RotatE [38]; quaternion model QuatE [48]; spherical models MuRS [43]; hyperbolic methods MurP [5], RotH [11]; Cartesian product space model  $\mathbf{M}^2\mathbf{G}\mathbf{N}\mathbf{N}$  [43], and Transformer/pretrained model based models HittER [14] and CSProm-KG  $[13]^1$ . CSProm-KG uses the pretrained large language model (LLM), BERT [17], as the backbone.
Main Results. We report the performance comparison in Figure 4 where we use TransE-AT [46] as baseline $^2$  and report the differences. We make the following observations. Firstly, our method outperforms all compared baselines without pretrained large language model as external sources. Specifically, it outperforms Euclidean, hyperbolic, spherical, complex-valued, and Cartesian product spaces methods and obtains a clear performance gain over the second-best model. Secondly, the
$^{1}$ A concurrent work.
${}^{2}$  TransE-AT on WN18RR: MRR=0.479,HR@1=0.434,HR@3=0.495,HR@10=0.571;on FB15k-237: MRR=0.351,HR@1=0.257,HR@3=0.386,HR@10=0.538.
Figure 5: Left: distribution of the selected geometric space on WN18RR; Center: distribution of the selected geometric space on FB15K-237; Right: geometric space distribution of four relations from FB15K-237.
Figure 6: Distribution of the selected spaces shown in a sub-graph of the FB15K-237.
improvement over the second-best model on FB15K-237 is larger than that on WN18RR. One possible explanation is that FB15K-237 has a more diverse set of relations and a higher structure heterogeneity [5], so a better geometric space alignment could bring more value. It is worth noting that WN18RR is sampled from the WordNet lexical database and BERT is trained on a large corpus of lexical data, it is unsurprising that CSProm-KG can obtain very high score on WN18RR. However, its relatively modest scores on FB15K-237 suggest limited practicality for non-lexical knowledge graphs. As such, it is not fair to compare an LLM based approach with our approach as using LLM may lead to data leakage.
Distribution of the selected Geometric Space. Figure 5 (left and center) presents the distribution of the selected spaces. We observe that certain product spaces are preferred depending on the dataset. For example, the product space  $\mathbb{D}_2\mathbb{D}_3$  is the most selected space on WN18RR. While on FB15K-237, the product space  $\mathbb{P}_1\mathbb{P}_2\mathbb{P}_3\mathbb{D}_1$  is more preferable. To show more fine-grained examples, we randomly select four relations from FB15K-237 and visualize the distribution in Figure 5 (Right). We find that each relation has their own desirable product spaces. For example, relation "/location/location/time Zones" prefers space  $\mathbb{P}_1\mathbb{P}_2\mathbb{P}_3\mathbb{D}_1$  but relation "people/person/profession" prefers space  $\mathbb{P}_1\mathbb{P}_3\mathbb{D}_1\mathbb{D}_2$ . This reconfirms the specification capability of our method.
To show how the space distribution coincides with data structures, we randomly extract three connected graph components from the test set of FB15K-237 and visualize them in Figure 6. Each triple can be represented with an edge and its two end nodes. We render edges with different colors based on the type of product spaces used to model the corresponding triples. From the left two graphs, we see some clear clusters with a dominant color, which suggests that neighborhood triples are more likely to be processed with the same selected product space. From the right-most graph, we find that different relations usually
prefer different spaces, which indicates that the type of relations plays a critical role in the space selection process.
M	Signature	WN18RR	FB15K-237
MRR	HR@3	MRR	HR@3
Single	D500	0.492	0.514	0.293	0.322
Product	(P100)3(D100)2	0.484	0.498	0.311	0.344
(D100)4E100	0.479	0.497	0.312	0.344
(P100)4D100	0.468	0.488	0.321	0.356
(P100)2(D100)2E100	0.479	0.496	0.308	0.342
Ours	3 × P100, 2 × D100	0.500	0.521	0.530	0.545
4 × D100, E100	0.526	0.549	0.525	0.531
4 × P100, D100	0.504	0.526	0.515	0.527
2 × P100, 2 × D100, E100	0.522	0.546	0.526	0.535
Table 1: Comparison with Cartesian product space models [22] (a few signatures are reported due to length constraints).  $K = 2$  for WN18RR and  $K = 4$  for FB15K-237.
Comparison with Cartesian Product Space with Varying Signatures. Compared with  $\mathbf{M}^2\mathrm{GNN}$  (its signature is  $\mathbb{P}^{200}\mathbb{D}^{200}\mathbb{E}^{200}$ ) (in Figure 4). Our model achieves the best performance even with
fewer trainable parameters. Then, we compare our method with Cartesian product space models with varying signatures in Table 1. We observe that our method constantly outperforms pure product spaces with the same prior signatures. It is worth noting that our method requires merely two/four active spaces to outperform product space models with five active spaces. These observations reconfirm the effectiveness of the data-informed space alignment method.
Impact of N and K.
The left figure of 7 shows the impact of  $N$  by fixing  $K$ . We observe that (a) the model performance does not benefit much from increasing  $N$  on this dataset; (b) the inference time remains nearly constant when we increase  $N$ , confirming its efficiency. The effect of  $K$  is shown in Figure 7 (right). We find that  $K$  has a higher impact on the
Figure 7: Effect of  $N$  (left) and  $K$  (right) on WN18RR.
model performance and inference time. Increasing  $K$  will generally degrade the performance on this dataset and yield additional computational overhead.
In this paper, we propose a data-informed geometric space selection method. Our method enables an automatic alignment between data and geometric spaces by allowing each data point to automatically select its preferred geometric spaces via a sparsely-gated MOE network. It brings a greater extent of specification for geometric representations while remaining efficient. We show that the proposed approach can boost representation performance in real-world applications with a noticeable margin.
Broader Impact. This approach exhibits considerable potential in facilitating enhanced and versatile geometric representation learning, thereby engendering broader and more comprehensive applications across various domains.
[1] Ralph Abboud, Ismail Ceylan, Thomas Lukasiewicz, and Tommaso Salvatori. Boxe: A box embedding model for knowledge base completion. Advances in Neural Information Processing Systems, 33, 2020.
[2] Rahaf Aljundi, Punarjay Chakravarty, and Tinne Tuytelaars. Expert gate: Lifelong learning with a network of experts. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 3366-3375, 2017.
[3] Alfonso Amayuelas, Shuai Zhang, Xi Susie Rao, and Ce Zhang. Neural methods for logical reasoning over knowledge graphs. In International Conference on Learning Representations, 2022.
[4] Gregor Bachmann, Gary Bécigneul, and Octavian Ganea. Constant curvature graph convolutional networks. In International Conference on Machine Learning, pages 486-496. PMLR, 2020.
[5] Ivana Balazevic, Carl Allen, and Timothy Hospedales. Multi-relational poincaré graph embeddings. In Advances in Neural Information Processing Systems, pages 4463-4473, 2019.
[6] Ivana Balazevic, Carl Allen, and Timothy Hospedales. Tucker: Tensor factorization for knowledge graph completion. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 5188-5197, 2019.
[7] Silvere Bonnabel. Stochastic gradient descent on riemannian manifolds. IEEE Transactions on Automatic Control, 58(9):2217-2229, 2013.
[42] Lucas Vinh Tran, Yi Tay, Shuai Zhang, Gao Cong, and Xiaoli Li. Hyperml: A boosting metric learning approach in hyperbolic space for recommender systems. In Proceedings of the 13th International Conference on Web Search and Data Mining, pages 609-617, 2020.
[43] Shen Wang, Xiaokai Wei, Cícero Nogueira dos Santos, Zhiguo Wang, Ramesh Nallapati, Andrew Arnold, Bing Xiang, S Yu Philip, and Isabel F Cruz. Mixed-curvature multi-relational graph neural network for knowledge graph completion. In The Web Conference, 2021.
[44] Richard C Wilson, Edwin R Hancock, Elżbieta Pekalska, and Robert PW Duin. Spherical and hyperbolic embeddings of data. IEEE transactions on pattern analysis and machine intelligence, 36(11):2255-2269, 2014.
[45] Bishan Yang, Wen-tau Yih, Xiaodong He, Jianfeng Gao, and Li Deng. Embedding entities and relations for learning and inference in knowledge bases. ICLR, 2015.
[46] Jina Yang, Yongjie Shi, Xin Tong, Robin Wang, Taiyan Chen, and Xianghua Ying. Improving knowledge graph embedding using affine transformations of entities corresponding to each relation. In Findings of the Association for Computational Linguistics: EMNLP 2021, pages 508-517, Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics.
[47] Bangpeng Yao, Dirk Walther, Diane Beck, and Li Fei-Fei. Hierarchical mixture of classification experts uncovers interactions between brain regions. Advances in Neural Information Processing Systems, 22:2178–2186, 2009.
[48] Shuai Zhang, Yi Tay, Lina Yao, and Qi Liu. Quaternion knowledge graph embeddings. In Advances in Neural Information Processing Systems, pages 2735-2745, 2019.