SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Here we introduce padding-free topk operation to avoid the overhead in batched execution. The key idea is to convert batched execution of different small samples into execution of a single large sample, which can be paralleled by existing operations in deep learning frameworks. For example, the batched execution of topk([[1,3],[2,1,0]]) can be converted into a multi-key sort problem over
$[0,1],[0,3],[1,2],[1,1],[1,0]]$ , where the first key is the index of the sample in the batch and the second key is the original input. The multi-key sort is then implemented by composing stable single-key sort operations in deep learning frameworks. See App. C for details.
Algorithm 1 A*Net
Input: head entity  $u$ , query relation  $q$ , #iterations  $T$   
Output:  $p(v|u, q)$  for all  $v \in \mathcal{V}$
1: for  $v \in \mathcal{V}$  do
2:  $\pmb{h}_q^{(0)}(u,v) \gets \mathbb{1}_q(u = v)$
3: end for
4: for  $t \gets 1$  to  $T$  do
5:  $\mathcal{X}^{(t)}\gets \mathrm{TopK}(s_{uq}^{(t - 1)}(x)|x\in \mathcal{V})$
6:  $\mathcal{E}^{(t)}\gets \bigcup_{x\in \mathcal{X}^{(t)}}\mathcal{E}(x)$
7:  $\mathcal{E}^{(t)}\gets \mathrm{TopL}(s_{uq}^{(t - 1)}(v)|(x,r,v)\in \mathcal{E}^{(t)})$
8:  $\mathcal{V}^{(t)}\gets \bigcup_{(x,r,v)\in \mathcal{E}^{(t)}}\{v\}$
9: for  $v\in \mathcal{V}^{(t)}$  do
10: Compute  $h_{q}^{(t)}(u,v)$  with Eqn. 12.
11: Compute priority  $s_{\mu q}^{(t)}(v)$  with Eqn. 10, 11.
12: end for
13: end for
14:  $\triangleright$  Share weights between  $s_{uq}(v)$  and the predictor
15: return  $s_{uq}^{(T)}(v)$  as  $p(v|u,q)$  for all  $v \in \mathcal{V}$
We evaluate A*Net on standard transductive and inductive knowledge graph reasoning datasets, including a million-scale one ogbl-wikikg2. We conduct ablation studies to verify our design choices and visualize the important paths learned by the priority function in A*Net.
Datasets & Evaluation We evaluate A*Net on 4 standard knowledge graphs, FB15k-237 [40], WN18RR [16], YAGO3-10 [30] and ogbl-wikikg2 [25]. For the transductive setting, we use the standard splits from their original works [40, 16]. For the inductive setting, we use the splits provided by [39], which contains 4 different versions for each dataset. As for evaluation, we use the standard filtered ranking protocol [6] for knowledge graph reasoning. Each triplet  $(u,q,v)$  is ranked against all negative triplets  $(u,q,v^{\prime})$  or  $(u^{\prime},q,v)$  that are not present in the knowledge graph. We measure the performance with mean reciprocal rank (MRR) and HITS at K (H@K). Efficiency is measured by the average number of messages (#message) per step, wall time per epoch and memory cost. To plot the convergence curves for each model, we dump checkpoints during training with a high frequency, and evaluate the checkpoints later on the validation set. See more details in App. D.
Implementation Details Our work is developed based on the open-source codebase of path-based reasoning with Bellman-Ford algorithm<sup>4</sup>. For a fair comparison with existing path-based methods, we follow the implementation of NBFNet [58] and parameterize  $\bigoplus$  with principal neighborhood aggregation (PNA) [13] or sum aggregation, and parameterize  $\bigotimes$  with the relation operation from DistMult [49], i.e., vector multiplication. The indicator function (Eqn. 2)  $\mathbb{1}_q(u = v) = \mathbb{1}(u = v)\pmb{q}$  is parameterized with a query embedding  $\pmb{q}$  for all datasets except ogbl-wikikg2, where we augment the indicator function with learnable embeddings based on a soft distance from  $u$  to  $v$  (see App. E for more details). The edge representation (Eqn. 12)  $\pmb{w}_{q}(\boldsymbol{x},\boldsymbol{r},\boldsymbol{v}) = \pmb{W}_{r}\pmb{q} + \pmb{b}_{r}$  is parameterized as a linear function over the query relation  $q$  for all datasets except WN18RR, where we use a simple embedding  $\pmb{w}_{q}(\boldsymbol{x},\boldsymbol{r},\boldsymbol{v}) = \pmb{r}$ . We use the same preprocessing steps as in [58], including augmenting each triplet with a flipped triplet, and dropping out query edges during training.
For the neural priority function, we have two hyperparameters:  $K$  for the maximum number of nodes and  $L$  for the maximum number of edges. To make hyperparameter tuning easier, we define maximum node ratio  $\alpha = K / |\mathcal{V}|$  and maximum average degree ratio  $\beta = L|\mathcal{V}| / K|\mathcal{E}|$ , and tune the ratios for each dataset. The maximum edge ratio is determined by  $\alpha \beta$ . The other hyperparameters are kept the same as the values in [58]. We train A*Net with 4 Tesla A100 GPUs (40 GB), and select the best model based on validation performance. See App. E for more details.
Baselines We compare A*Net against embedding methods, GNNs and path-based methods. The embedding methods are TransE [6], ComplEx [42], RotatE [38], HAKE [55], RotH [7], PairRE [8], ComplEx+Relation Prediction [12] and ConE [3]. The GNNs are RGCN [36], CompGCN [43] and GraIL [39]. The path-based methods are MINERVA [14], Multi-Hop [29], CURL [52], NeuralLP [50], DRUM [35], NBFNet [58] and RED-GNN [54]. Note that path-finding methods [14, 29, 52] that use reinforcement learning and assume sparse answers can only be evaluated on tail prediction. Training time of all baselines are measured based on their official open-source implementations, except that we use a more recent implementation<sup>5</sup> of TransE and ComplEx.
Tab. 1 shows that A*Net outperforms all embedding methods and GNNs, and is on par with NBFNet on transductive knowledge graph reasoning. We also observe a similar trend of A*Net and NBFNet over path-finding methods on tail prediction (Tab. 2). Since path-finding methods select only one path with reinforcement learning, such results imply the advantage of aggregating multiple paths in A*Net. A*Net also converges faster than all the other methods (Fig. 5). Notably, unlike NBFNet that propagates through all nodes and edges, A*Net only propagates through  $10\%$  nodes and  $10\%$  edges on both datasets, which suggests that most nodes and edges are not important for path-based reasoning. Tab. 3 shows that A*Net reduces the number of messages by  $14.1\times$  and  $42.9\times$  compared to NBFNet on two datasets respectively. Note that the reduction in time and memory is less than the reduction in the number of messages, since A*Net operates on subgraphs with dynamic sizes and is harder to parallel than NBFNet on GPUs. We leave better parallel implementation as future work.
Tab. 4 shows the performance on ogbl-wikikg2, which has 2.5 million entities and 16 million triplets. While NBFNet faces out-of-memory (OOM) problem even for a batch size of 1, A*Net can perform reasoning by propagating through  $0.2\%$  nodes and  $0.2\%$  edges at each step. Surprisingly, even with
<sup>4</sup>https://github.com/DeepGraphLearning/NBFNet.MIT license.
<sup>5</sup>https://github.com/DeepGraphLearning/KnowledgeGraphEmbedding.MIT license.
Table 1: Performance on transductive knowledge graph reasoning. Results of embedding methods are from [3]. Results of GNNs and path-based methods are from [58]. Performance and efficiency on YAGO3-10 are in App. F.
Method	FB15k-237	WN18RR
MRR	H@1	H@3	H@10	MRR	H@1	H@3	H@10
TransE	0.294	-	-	0.465	0.226	-	0.403	0.532
RotatE	0.338	0.241	0.375	0.533	0.476	0.428	0.492	0.571
HAKE	0.341	0.243	0.378	0.535	0.496	0.451	0.513	0.582
RotH	0.344	0.246	0.380	0.535	0.495	0.449	0.514	0.586
ComplEx+RP	0.388	0.298	0.425	0.568	0.488	0.443	0.505	0.578
ConE	0.345	0.247	0.381	0.540	0.496	0.453	0.515	0.579
RGCN	0.273	0.182	0.303	0.456	0.402	0.345	0.437	0.494
CompGCN	0.355	0.264	0.390	0.535	0.479	0.443	0.494	0.546
NeuralLP	0.240	-	-	0.362	0.435	0.371	0.434	0.566
DRUM	0.343	0.255	0.378	0.516	0.486	0.425	0.513	0.586
NBFNet	0.415	0.321	0.454	0.599	0.551	0.497	0.573	0.666
RED-GNN	0.374	0.283	-	0.558	0.533	0.485	-	0.624
A*Net	0.411	0.321	0.453	0.586	0.549	0.495	0.573	0.659
Table 2: Tail prediction performance on transductive knowledge graphs. Results of compared methods are from [14, 29, 52].
Method	FB15k-237	WN18RR
MRR	H@1	H@3	H@10	MRR	H@1	H@3	H@10
MINERVA	0.293	0.217	0.329	0.456	0.448	0.413	0.456	0.513
Multi-Hop	0.393	0.329	-	0.544	0.472	0.437	-	0.542
CURL	0.306	0.224	0.341	0.470	0.460	0.429	0.471	0.523
NBFNet	0.509	0.411	0.562	0.697	0.557	0.503	0.579	0.669
A*Net	0.505	0.410	0.556	0.687	0.557	0.504	0.580	0.666
Table 3: Efficiency on transductive knowledge graph reasoning.
Method	FB15k-237	WN18RR
#message	time	memory	#message	time	memory
NBFNet	544,230	16.8 min	19.1 GiB	173,670	9.42 min	26.4 GiB
A*Net	38,610	8.07 min	11.1 GiB	4,049	1.39 min	5.04 GiB
Improvement	14.1×	2.1×	1.7×	42.9×	6.8×	5.2×
Figure 5: Validation MRR w.r.t. training time (1 A100 GPU).
Table 4: Performance on ogbl-wikikg2 (MRR). Results of compared methods are from [8, 12].
Method	ogbl-wikigk2
Test	Valid	#Params
TransE	0.4256	0.4272	1,251 M
ComplEx	0.4027	0.3759	1,251 M
RotatE	0.4332	0.4353	1,251 M
PairRE	0.5208	0.5423	500 M
ComplEx+RP	0.6392	0.6561	250 M
NBFNet	OOM	OOM	OOM
A*Net	0.6767	0.6851	6.83 M
such sparse propagation, A*Net outperforms embedding methods and achieves a new state-of-the-art result. Moreover, the validation curve in Fig. 1 shows that A*Net converges significantly faster than embedding methods. Since A*Net only learns parameters for relations but not entities, it only uses 6.83 million parameters, which is  $36.6 \times$  less than the best embedding method ComplEx+RP.
Tab. 5 shows the performance on inductive knowledge graph reasoning. A*Net is on par with NBFNet and significantly outperforms all the other methods. Note that embedding methods cannot deal with the inductive setting. Other metrics (H@1, H@10) and efficiency results are in App. F.
4.3 Ablation Studies
Priority Function To verify the effectiveness of neural priority function, we compare it against three handcrafted priority functions: personalized PageRank (PPR), Degree and Random. PPR selects nodes with higher PPR scores w.r.t. the query head entity  $u$ . Degree selects nodes with larger degrees, while Random selects nodes uniformly. Tab. 6 shows that the neural priority function outperforms all three handcrafted priority functions, suggesting the necessity of learning a neural priority function.
Sharing Weights As discussed in Sec. 3.2, we share the weights between the neural priority function and the reasoning predictor to help train the neural priority function. Tab. 6 compares A*Net trained with and without sharing weights. It can be observed that sharing weights is essential to training a good neural priority function in A*Net.
Trade-off between Performance and Efficiency While A*Net matches the performance of NBFNet in less training time, one may further trade off performance and efficiency in A*Net by adjusting the ratios  $\alpha$  and  $\beta$ . Fig. 6 plots curves of performance and speedup ratio w.r.t. different
Table 5: Performance on inductive knowledge graph reasoning (MRR). V1-v4 are 4 standard inductive splits. Results of compared methods are taken from [54].  $\alpha = 50\%$  and  $\beta = 100\%$  for FB15k237.  $\alpha = 5\%$  and  $\beta = 100\%$  for WN18RR. More metrics and efficiency results are in App. F.
Method	FB15k-237	WN18RR
v1	v2	v3	v4	v1	v2	v3	v4
GraIL	0.279	0.276	0.251	0.227	0.627	0.625	0.323	0.553
NeuralLP	0.325	0.389	0.400	0.396	0.649	0.635	0.361	0.628
DRUM	0.333	0.395	0.402	0.410	0.666	0.646	0.380	0.627
NBFNet	0.422	0.514	0.476	0.453	0.741	0.704	0.452	0.641
RED-GNN	0.369	0.469	0.445	0.442	0.701	0.690	0.427	0.651
A*Net	0.457	0.510	0.476	0.466	0.727	0.704	0.441	0.661
Table 6: Ablation studies of A*Net on transductive FB15k-237.
(a) Choices of priority function.
Priority Function	FB15k-237
MRR	H@1	H@3	H@10
PPR	0.266	0.212	0.296	0.371
Degree	0.347	0.268	0.383	0.501
Random	0.378	0.288	0.413	0.556
Neural	0.411	0.321	0.453	0.586
(b) W/ or w/o sharing weights.
Sharing
Weights	FB15k-237
MRR	H@1	H@3	H@10
No	0.374	0.282	0.413	0.557
Yes	0.411	0.321	0.453	0.586
Figure 6: Performance and efficiency trade-off w.r.t. node ratio  $\alpha$  and degree ratio  $\beta$ . Speedup ratio is relative to NBFNet.
(Bandai, industry, ?)
Figure 7: Visualization of important paths learned by the neural priority function in A*Net.
$\alpha$  and  $\beta$ . If we can accept a performance similar to embedding methods (e.g., ConE [3]), we can set either  $\alpha$  to  $1\%$  or  $\beta$  to  $10\%$ , resulting in  $8.7\times$  speedup compared to NBFNet.
4.4 Visualization of Learned Important Paths
We can extract the important paths from the neural priority function in A*Net for interpretation. For a given query  $(u, q, ?)$  and a predicted entity  $v$ , we can use the node priority  $s_{uq}^{(t)}(x)$  at each step to estimate the importance of a path. Empirically, the importance of a path  $s_q(P)$  is estimated by
$$
s _ {q} (P) = \frac {1}{| P |} \sum_ {t = 1, P ^ {(t)} = (x, r, y)} ^ {| P |} \frac {s _ {u q} ^ {(t - 1)} (x)}{S _ {u q} ^ {(t - 1)}} \tag {14}
$$
where  $S_{uq}^{(t-1)} = \max_{x \in \mathcal{V}^{(t-1)}} s_{uq}^{(t-1)}(x)$  is a normalizer to normalize the priority score for each step  $t$ . To extract the important paths with large  $s_q(P)$ , we perform beam search over the priority function  $s_{uq}^{(t-1)}(x)$  of each step. Fig. 7 shows the important paths learned by A*Net for a test sample in FB15k-237. Given the query (Bandai, industry, ?), we can see both paths Bandai  $\leftarrow \overset{\text{subsidiary}}{\underset{B\text{andai}}{\rightleftharpoons}}$  Namco  $\overset{\text{industry}}{\rightarrow}$  video game and Bandai  $\overset{\text{industry}}{\rightarrow}$  media  $\overset{\text{industry}}{\rightarrow}$  Pony Canyon  $\overset{\text{industry}}{\rightarrow}$  video game are consistent with human cognition. More visualization results can be found in App. G.
Path-based Reasoning Path-based methods use paths between entities for knowledge graph reasoning. Early methods like Path Ranking [28, 19] collect relational paths as symbolic features for classification. Path-RNN [31, 15] and PathCon [45] improve Path Ranking by learning the representations of paths with recurrent neural networks (RNN). However, these works operate on the full set of paths between two entities, which grows exponentially w.r.t. the path length. Typically, these methods can only be applied to paths with at most 3 edges.
To avoid the exhaustive search of paths, many methods learn to sample important paths for reasoning. DeepPath [47] and MINERVA [14] learn an agent to collect meaningful paths on the knowledge graph
through reinforcement learning. These methods are hard to train due to the extremely sparse rewards. Later works improve them by engineering the reward function [29] or the search strategy [37], using multiple agents for positive and negative paths [24] or for coarse- and fine-grained paths [52]. [11] and [33] use a variational formulation to learn a sparse prior for path sampling. Another category of methods utilizes the dynamic programming to search paths in a polynomial time. NeuralLP [50] and DRUM [35] use dynamic programming to learn linear combination of logic rules. All-Paths [41] adopts a Floyd-Warshall-like algorithm to learn path representations between all pairs of entities. Recently, NBFNet [58] and RED-GNN [54] leverage a Bellman-Ford-like algorithm to learn path representations from a single-source entity to all entities. While dynamic programming methods achieve state-of-the-art results among path-based methods, they need to perform message passing on the full knowledge graph. By comparison, our A*Net learns a priority function and only explores a subset of paths, which is more scalable than existing dynamic programming methods.
Efficient Graph Neural Networks Our work is also related to efficient graph neural networks, since both try to improve the scalability of graph neural networks (GNNs). Sampling methods [21, 9, 26, 51] reduce the cost of message passing by computing GNNs with a sampled subset of nodes and edges. Non-parametric GNNs [27, 46, 18, 10] decouple feature propagation from feature transformation, and reduce time complexity by preprocessing feature propagation. However, both sampling methods and non-parametric GNNs are designed for homogeneous graphs, and it is not straightforward to adapt them to knowledge graphs. On knowledge graphs, RS-GCN [17] learns to sample neighborhood with reinforcement learning. DPMPN [48] learns an attention to iteratively select nodes for message passing. SQUALER [1] first predicts important path types based on the query, and then applies GNNs on the subgraph extracted by the predicted paths. Our A*Net shares the same goal with these methods, but learns a neural priority function to iteratively select important paths.
6 Discussion and Conclusion
Limitation and Future Work One limitation for A*Net is that we focus on algorithm design rather than system design. As a result, the improvement in time and memory cost is much less than the improvement in the number of messages (Tab. 3 and App. F). In the future, we will co-design the algorithm and the system to further improve the efficiency.
Societal Impact This work proposes a scalable model for path-based reasoning. On the positive side, it reduces the training and test time of reasoning models, which helps control carbon emission. On the negative side, reasoning models might be used in malicious activities, such as discovering sensitive relationship in anonymized data, which could be augmented by a more scalable model.
Conclusion We propose A*Net, a scalable path-based method, to solve knowledge graph reasoning by searching for important paths, which is guided by a neural priority function. Experiments on both transductive and inductive knowledge graphs verify the performance and efficiency of A*Net. Meanwhile, A*Net is the first path-based method that scales to million-scale knowledge graphs.
This project is supported by Intel-MILA partnership program, the Natural Sciences and Engineering Research Council (NSERC) Discovery Grant, the Canada CIFAR AI Chair Program, collaboration grants between Microsoft Research and Mila, Samsung Electronics Co., Ltd., Amazon Faculty Research Award, Tencent AI Lab Rhino-Bird Gift Fund and a NRC Collaborative R&D Project (AI4D-CORE-06). This project was also partially funded by IVADO Fundamental Research Project grant PRF-2019-3583139727. The computation resource of this project is supported by Mila<sup>6</sup>, Calcul Quebec<sup>7</sup> and the Digital Research Alliance of Canada<sup>8</sup>.
We would like to thank Zuobai Zhang, Jiarui Lu and Minghao Xu for helpful discussions and comments. We also appreciate all anonymous reviewers for their constructive suggestions.
$^{6}$ https://mila.quebec/
<sup>7</sup>https://www.calculquebec.ca/
<sup>8</sup>https://alliancecan.ca/
[50] Fan Yang, Zhilin Yang, and William W Cohen. Differentiable learning of logical rules for knowledge base reasoning. In Advances in Neural Information Processing Systems, pages 2316-2325, 2017.
[51] Hanqing Zeng, Hongkuan Zhou, Ajitesh Srivastava, Rajgopal Kannan, and Viktor Prasanna. Graphsaint: Graph sampling based inductive learning method. In International Conference on Learning Representations, 2019.
[52] Denghui Zhang, Zixuan Yuan, Hao Liu, Hui Xiong, et al. Learning to walk with dual agents for knowledge graph reasoning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 5932-5941, 2022.
[53] Fuzheng Zhang, Nicholas Jing Yuan, Defu Lian, Xing Xie, and Wei-Ying Ma. Collaborative knowledge base embedding for recommender systems. In Proceedings of the 22nd ACM SIGKDD international conference on knowledge discovery and data mining, pages 353-362, 2016.
[54] Yongqi Zhang and Quanming Yao. Knowledge graph reasoning with relational digraph. In Proceedings of the ACM Web Conference 2022, pages 912-924, 2022.
[55] Zhanqiu Zhang, Jianyu Cai, Yongdong Zhang, and Jie Wang. Learning hierarchy-aware knowledge graph embeddings for link prediction. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 3065-3072, 2020.
[56] Da Zheng, Xiang Song, Chao Ma, Zeyuan Tan, Zihao Ye, Jin Dong, Hao Xiong, Zheng Zhang, and George Karypis. Dgl-ke: Training knowledge graph embeddings at scale. In Proceedings of the 43rd International ACM SIGIR Conference on Research and Development in Information Retrieval, pages 739–748, 2020.
[57] Zhaocheng Zhu, Shizhen Xu, Jian Tang, and Meng Qu. Graphvite: A high-performance cpu-gpu hybrid system for node embedding. In The World Wide Web Conference, pages 2494-2504, 2019.
[58] Zhaocheng Zhu, Zuobai Zhang, Louis-Pascal Xhonneux, and Jian Tang. Neural bellman-ford networks: A general graph neural network framework for link prediction. Advances in Neural Information Processing Systems, 34, 2021.