A*Net: A Scalable Path-based Reasoning Approach for Knowledge Graphs
Authors: Zhaocheng Zhu
Xinyu Yuan
Mikhail Galkin
Sophie Xhonneux
Ming Zhang
Maxime Gazeau
Jian Tang
Venue: NeurIPS 2023 poster (2023)
Keywords: Knowledge Graph Reasoning
Path-based Methods
Scalability
A* Algorithm
TL;DR: A scalable path-based method for reasoning over knowledge graphs with millions of entities and triplets.

Abstract: Reasoning on large-scale knowledge graphs has been long dominated by embedding methods. While path-based methods possess the inductive capacity that embeddings lack, their scalability is limited by the exponential number of paths. Here we present A\*Net, a scalable path-based method for knowledge graph reasoning. Inspired by the A\* algorithm for shortest path problems, our A\*Net learns a priority function to select important nodes and edges at each iteration, to reduce time and memory footprint for both training and inference. The ratio of selected nodes and edges can be specified to trade off between performance and efficiency. Experiments on both transductive and inductive knowledge graph reasoning benchmarks show that A\*Net achieves competitive performance with existing state-of-the-art path-based methods, while merely visiting 10% nodes and 10% edges at each iteration. On a million-scale dataset ogbl-wikikg2, A\*Net not only achieves a new state-of-the-art result, but also converges faster than embedding methods. A\*Net is the first path-based method for knowledge graph reasoning at such scale.

A Neet: A Scalable Path-based Reasoning Approach for Knowledge Graphs
Zhaocheng Zhu Xinyu Yuan Mikkaal Galkin Sophie Xhonneux!2
Ming Zhang', Maxime Gazeau', Jian Tang
Milla Québec Al Institute, Uniieeriity of Montréal
Intel Al Lab, *Peking University, *LG Electronics Al Lab 'HEC Montréal, "CIFAR Al Chair
Abstract
Reasoning on parge-scale knowledge graphs has been long dominated by embedding methods. While path-based methods possess the inductive capacity that embed- dings lack, their scalability limited by the exponential number of paths. Here present A Net, scalable path-based method for knowledge graph reasoning. Inspired by the A algorithm for shortest path problems, our * learns a priority function to select important nodes and edges at each iteration, to reduce time and memory footprint for both training and inference. The ratio of selected nodes and edges can be specified to trade off between performance and efficiency. Experi- ments on both transductive and inductive knowledge graph reasoning benchmarks show that A Net achieves competitive performance with existing state-of-the-art path-based methods, while merely visiting 10% nodes and 10% edges at each iteration. On million-scale dataset wgbi-w2 A Net not only achieves a new state-of-the-ar result, but also converges faster than embedding methods. * is the first path-based method for knowledge graph reasoning at such scale.
Introduction
Reasoning, the ability to apply logic to draw new con- clusions from existing facts, has been long pursued as a goal of artificial intelligence [32 20). Knowledge graphs encapsulate facts in relational edges between entities, and serve foundation for reasoning. Reasoning over knowledge graphs is usually studied in the form of knowl- edge graph completion, where model is asked to predict missing triplets based on observed triplets in the knowl- edge graph. Such a task can be used existing knowledge graphs, but also improve downstream applications like multi-hop logical reasoning 34
0.5
MIR
0.4
Iransk:
CamplEs
0.2
Rotath:
ComplExx RP
only populate
0.0
10
15
Training
Time (hrs)
ques-
tion answering (5) and recommender systems
53
Figure I: Validation MRR w.r.t train- ing time on oibl-wi (1 A100 GPU). A Net achieves oftt--of-the-rt perfor- mance and the fastest convergence.
One challenge central to knowledge graph reasoning is the scalability of reasoning methods, as many real-world knowledge graphs 2 4 contain millions of entities and triplets. Typically, large-scale knowledge graph reasoning solved by embedding methods 6 42 38
which learn embedding for each entity and relation to reconstruct the structure of the knowledge
"Equal contribution. Code available abtpperr//shaaa..oooooooppppppps
Work done while at Mila Québee Al Institute.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

(
Fadeer
Friend
Brother
Molher
Motter
Mother
Wife
Friend
(a) Path-based reasoning
(b) Exhaustive search
(c) Bellman-Ford
(a)**
Figure 2: (a) Given a query (a, Mother, ?), only a few important paths (showed in red) are necessary for reasoning. Note that paths can go in the reverse direction of relations. (b) Exhaustive search algorithm (c.g., Path-RNN, PathCon) enumerates all paths in exponential time. (c) Bellman-Fort algorithm (e.g.. NeuralLP, DRUM, NBENet, RED-GNN) computes all paths in polynomial time, but needs to propagate through all nodes and edges. (d) ^ learns priority function to select subset of nodes and edges at each iteration, and avoids exploring all nodes and edges.
graph. Due to its simplicity, embedding methods have become the de facto standard for knowledge graphs with millions of entities and triplets. With the help of multi-GPU embedding systems [57] [56), they can further scale to knowledge graphs with billions of triplets.
58, predicts the relation between pair of entities based on the paths between them. Take the knowledge graph Fig. 2ia) as an example,
Another stream of works, path-based methods [28
we can prove that Mother(a. f) holds, because there are two paths a Funher, Wge f and BBothee
Mother As the semanties of paths purely determined by relations rather than entities, path-based methods naturally generalize to unseen entities (i.e., inductive setting). which cannot be handled by embedding methods. However, the number of paths grows exponentially w.r.t. the path length, which hinders the application of path-based methods on large-scale knowledge graphs.
Here we propose A*te to tackle the scalability issue of path-based methods. The key idea of our method to search for important paths rather than use all possible paths for reasoning, thereby reducing time and memory in training and inference. Inspired by the A algorithm [22] for shortest path problems, given a head entity u and a query relation q, we compute priority score for each entity to guide the search towards more important paths. At each iteration, we select K nodes and L edges according to their priority, and use message passing to update nodes in their neighborhood. Due to the complex semantics of knowledge graphs, it is hard to use handcrafted priority function like the A* algorithm without a significant performance drop (Tab, 6a). Instead, we design a neural priority function based on the node representations at the current iteration, which be end-to-end trained by the objective function of the reasoning task without any additional supervision.
We verify our method on 4 transductive and 2 inductive knowledge graph reasoning datasets. Ex- periments show that A *Net achieves competitive performance against state-of-the-art path-based methods on FB 153-277, WNI88 WNI8RR and YAGO3-10, even with only 10% of nodes and 10% edges at each iteration (Sec. 4.2). To verify the scalability of our method, we also evaluate Net on wibl-2ik million-scale knowledge graph that is magnitudes larger than datasets solved by previous path-based based methods. Surprisingly, with only 0.2% nodes and 0.2% edges, our method outperforms existing embedding methods and establishes new state-of-ther oftthh--rttttttt results (Sec. 4.2) the first non-embedding method on oibl-2i By adjusting the ratios of selected nodes and edges, can trade off between performance and efficiency (Sec. (43).. A A Net also converges significantly faster than embedding methods (Fig. []. which makes it a promising model for deployment on large-scale knowledge graphs. Additionally, A "ett offers interpretability that embeddings do not possess. Visualization shows that A Net captures important paths for reasoning (Sec. 4.4).
2 Preliminary
Knowledge Graph Reasoning knowledge graph G = EV,E, R.) consists of sets of entities (nodes) V7, facts (edges) & and relation types R. Each fact is triplet (æ, r, y) E V R x V. which indicates a relation from entity x to entity y. The task of knowledge graph reasoning aims at answering queries like (u, (u,q,?) ?) (?,g, u). Without loss of generality, assume the query is
2

---

(u, 9, ?), since g? u) equals to ? with q'' being the inverse of q. Given a query (u, ?), we need to predict the answer set ((...?)) such that Vu E (i.g.)) the triplet (u, 9, v) should be true.
58 solve knowledge graph reasoning by looking at the paths between a pair of entities in knowledge graph. For example, path a Fatter, - f may be used to predict Mother(a, f) in Fig. a). From representation learning perspective, path-based methods aim to learn representation h u4 u) to predict the triplet (u, 9. v) based on all paths Pa--y from entity u to entity v. Following the notation in [58] hg (u, v) is defined as
Path-based Methods Path-based methods 28 31
(u v)
+ PEPawe
(P)
(2,
(1)
PEPame (a.ra))P
where θ is permutation-invariaa aggregation function over paths (e.g., sum max), X is an aggregation function over edges that may be permutation- senssiive (e.g., matrix multiplication) and (,, r, 3) is the representation of triplet (z,r,y) conditioned on the query relation q. X is computed before ① Typically, (2 r, y) is designed to be independent of the entities æ and y. which enables path-based methods to generalize to the inductive setting. However, it is intractable to compute Eqn. since the number of paths usually grows exponentially w.r.t. the path length.
Path-based Reasoning with Bellman-Ford algorithm To reduce the time complexity of path- based methods, recent works 550
borrow the Bellman-Fore Ford algorithm from shortest path problems to solve path-based methods. Instead of enumerating each possible path, the Bellman- Ford algorithm iteratively propagates the representations of hops to compute the representations of hops, which achieves polynomial time complexity. Formally, let hg (u, v) be the representation of hops. The Bellman-Ford algorithm can be written as
35
58 54
(0 (u, v) < 1g(u v)
(2)
h (4) (u. t n (u, v) +
10 f1111) (u, & (2,rr,,
(3)
(0) where 1, is a learnable indicator function that defines the representations of 0 hops h4 (u, v). also known as the boundary condition the Bellman-Ford algorithm. {(V) the neighborhood of node v. Despite the polynomial time complexity achieved by the Bellman-Ford algorithm, Eqn. still needs to visit [V] nodes and |E| edges to compute h. (u, v) for all e each iteration, which is not feasible for lcage-scale knowledge graphs.
A Algorithm A algorithm 22 is extension of the Bellman- algorithm for shortest path problems. Unlike the Bellman-Fort algorithm that propagates through every node uniformly, the A algorithm prioritizes propagation through nodes with higher priority according to heuristic function specified by the user. With an appropriate heuristic function, A algorithm can reduce the search space of paths. Formally, with the notation from Eqn.
[ the priority function for node is
s(x) = d(u, x) ((
where d(u, æ) the length of current shortest path from u to z, and g(z, v) is a heuristic function estimating the cost from to the target node v. For instance, for grid-world shortest path problem (Fig. 4()), g(x, v) is usually defined as the L, distance from x to v, & is the addition operator, and s(x) is a lower bound for the shortest path length from u to e through x. During each iteration, the A algorithm prioritizes propagation through nodes with smaller sxx).
3 Proposed Method
We propose Net to scale up path-based methods with the A algorithm. We show that the * algorithm be derived from the observation that only a small of paths important for reasoning (Sec. 3.1 Since it is hard to handcraft a good priority function for knowledge graph reasoning (Tab. 6ab. design neural priority function, and train it end-to-end for reasoning (Sec. 3.2)
Bgp and @ are binary operations (akin x), while Q and X n-ary operations (akin I.ID.

---

3.1 Path-based Reasoning with A Algorithm
As discussed in Sec. 2 the Bellman-Foot Ford algorithm visits all [ nodes and || edges. However, in real-world knowledge graphs, only small portion of paths is related to the query. Based on this observation, we introduce the concept of important paths. We then show that the representations of important paths can be iteratively computed with the A* algorithm under mild assumptions.
Important Paths for Reasoning Given a query relation and a pair of entities, only some of the
paths between the entities are important for answering the query. Consider the example in Fig. 2 a), Friond Mother Friend, f cannot determine whether f an answer to Mother(a, ?) due to
the path a the use of the Friend relation in the path. On the other hand, kinship paths like Finherr b Brocher Mother, f are able to predict that Mother(a, f) is true. Formally, we define Pu Pu--ol S Pune to be the set of paths from u to u that is important to the query relation q. Mathematicallly we have
wife f or
(u, v)
(p) R
(P)
PEPami
PEP
In other words, any path P E Pu-mo PP-mell has negligible contribution to h, v). In real-world knowledge graphs, the number important paths [P /Pu-wolgl may be several orders of magnitudes smaller than the number of paths [Pu-oo [1. If we compute the representation hy(u,, v) using only the important paths, we can scale up path-based reasoning to large-scale knowledge graphs.
Iterative Computation of Important Paths Given a query (u, q, ?), we need to discover the set of im- portant paths P Pu-nol for all V. However, challenging to extract important paths from P Pu- since the size of Pu-ov is exponentially large. Our solution is to explore the structure of important paths and compute them iteratively. We first show that we can cover important paths with iterative path selec- tion (Eqn. 6 and 77. Then we approximate iterative path selection with iterative node selection (Eqn. 8).
$ a (a]
Pe 1x E (b,c,d)
p E {a
(b) (b) (h (b)
Figure 3: The colored paths are important paths Pu-elgs while the solid paths are the superset Pu Punvelo used in Eqn.
Notice that paths in Pa-- form tree structure (Fig. B. On the tree, a path is not important if any prefix of this path not important for the query. For example, in Fig. 2 a),
Friend
Mother
Friend, f is not important, its prefix Frioud, d is not important for the query Mother. Therefore, we assume there exists path selection function ma 2" + 2 that selects important paths from a set of paths given the query relation y. 2 is the set of all subsets of P. With my we construct the
following of paths p(t)
u--olg iteratively
B(0)
u--olg [(u, self loop, u)} if elsse ø
(6)
p(t)
U ZEV (x)),))E
{(+((r,)))) PE my (P }
(7)
where P+ ((r,,,)) concatenates the path P and the edge (r,r,u), The paths p(0)) uu-ull computed by
the above iteration is a superset of the important paths p(t) u--olo of length t (see Thm. A.I App. A
Due to the tree structure of paths, the above iterative path selection still requires exponential time.
Hence we further approximate iterative path selection with iterative node selection, by assuming
paths with the same length and the same stop node can be merged. The iterative node selection
replacing Eqn. 7 is (see Prop.
A.3 in App. A
p(O) anolog
U
P++ ((,,T,))) E p(t-1) )
(8)
(ER((-D(()) (1.7)))EE(v)
where nug 2 + 2 selects ending nodes of important paths of length from set of nodes.

---

Reasoning with A Algorithm Eqn. 8 iteratively computes the set of important paths P Pu--rl- In order to perform reasoning, we need to compute the representation (u, v) based on the important paths, which can be achieved by an iterative process similar Eqn. 8 (see Thm. A.44 App. A)
l (u, + (u
h/g (t ( (u (X,Tr,
(En(t-D(()) (2,15))EE())
Eqn. is the A iteration (Fig. 2()) for path- based reasoning. Note the A iteration uses the same boundary condition as Eqn. Inspired by the classical A algorithm, we parameterize ( (V) with
node priority function Bug V + [0, 1] and select top-K nodes based on their priority. However, there does not exist an oracle for the priority function ( (x:). We will discuss how learn
priority function Sug (z) in the following sections.
3.2 Path-based Reasoning with NeNe
Both the performance and the efficiency of the A algorithm heavily rely on the heuristic function. While it is straightforward to use L1 distance the heuristic function for grid-world shortest path problems. it is not clear what good priority func- tion for knowledge graph reasoning is due to the complex relation semantics knowledge graphs. Indeed, our experiments suggest that handcrafted priority functions largely hurt the performance of path-based methods (Tab. (Tab. In this section, we discuss a neural priority function, which can end-to-end trained by the reasoning task.
Fricad
Father
6+6
Mother
G+d
(a) A+ algorithm in grid world
(b) A Net on a knowledge graph
Figure (a) A* algorithm computes the current distance d(u, x) (blue), estimates the remaining distance g (z. v) (orange), and prioritizes shorter paths. (b) A Nett computes the current represen- tations h (u, (blue), estimates the remain-
Neural Priority Function To design the neural priority function Sug (r), we draw inspiration from the priority function in the A algorithm for short- est path problems (Eqn. 4 The priority function has terms d(u, and g(r, v), where d( d(u,,)) x:) is the current distance from node u to x, and g(x, v) estimates the remaining distance from node x to
ing representations g (() (u,x). q]) (orange) based the query q (green), and prioritizes paths more relevant to the query.
From a representation learning perspective, we need to learn a representation sug to predict the priority score Sug for each node Inspired by Eqn. 4. use the current representation h! (u,z) represent (t (u, .). However, is challenging to find representation for ((*) (Ir, (zr,)), since do not know the answer entity beforehand. Noticing that in the A* algorithm, the target node can be expressed by the source node plus displacement (Fig. (a)), we reparameterize the answer entity with the head entity and the query relation q in A Net. By replacing (9) v) with another function g ( q), the representation Sug (x) is parameterized as
h (u, w)
((h (u,
(10)
where g(.) is feed-forward network that outputs vector representation and [.,]] concatenates two representations. Intuitively, the learned representation q captures semantic of query relation g.
which serves the goal for answering query (u, g. ?). The function ((h
(t) (u, ql) compares the
current representation h (u,x) x)) with the goal q to estimate the remaining representation (Fig b)).
If h (u,x) is close to q, the remaining representation will be close to 0, and x is likely to be close the correct answer. The final priority score predicted by
( (x) a (s
(11)
where f(.) is ford-fr network and is the sigmoid function that maps the output to |0, 1].
5

---

Learning To learn the neural priority function, we incorporate it as weight for each message in the A iteration. For simplicity, let x Bugg
(nlg-)(())) (V) be the nodes we try to propagate through at
1-th iteration. We modify Eqn. 9 to be
h (u, h (u,
uu (x)
(h) (hi-") mg(0;
(12)
ZEX (1
Eqn. encourages the model learn larger weights Sug (x) for nodes that are important for reasoning. In practice, as some nodes may have very large degrees, we further select top-L edges miy (V²) (see App. By. A pseudo code of A Net is illustrated in Alg. Note the top-K and top-I. I functions are not differentiable.
(0)
from the neighborhood of fug
Nevertheless, it is still too challenging to train the neural priority function, since we do not know the ground truth for important paths, and there is no direct supervision for the priority function. Our solution is to share the weights between the priority function and the predictor for the reasoning task. The intuition is that the reasoning task can be viewed as a weak supervision for the priority function.
Recall that the goal of Suy (x)) is to determine whether there exists an important path from u to (Eqn. [). In the reasoning task, any positive answer entity must be present on least one important path, while negative answer entities are less likely to be important paths. Our ablation experiment demonstrates that sharing weights improve the performance of neural priority function (Tab, 6b). Following (38), A Net is trained to minimize the binary cross entropy loss over triplets
(t)
L log pogppuuggo0 q.u)
log(1 p(u, q, v())
(13)
where (u, q, u) is a positive sample and {(44, z, of) Ji= are negative samples. Each negative sample (u45 I, 14) is generated by corrupting the head or the tail in positive sample.
Efficient Implementation with Padding-Frre Free Operations Modern neural networks heavily rely on batched exccution to unleash the parallel capacity of GPUs. While Alg. iis easy to imple- ment for single sample (u, q, ?), it is trivial batch A Net for multiple samples. The challenge
Algorithm 1 *ett
Input: head entity u, query relation fiterations T Output: p(rlu, u, q) for all v E V
I: for v E V do
2: h4 (u, v) + 1,(u= =0)
3: end for
that different samples may have very different
4: for t
1 to T do
sizes for nodes 10t)) and edges (()) A common approach is to pad the set of nodes or edges to a predefined constant, which would severely coun- teract the acceleration brought by A Net..
x² + TopK "o (x)lar E V)
6:
g0o UsEx()) E(x)
g0o + Topl. (o (x)
8:
DIO Uen.weeete (v}
Here introduce padding-free topk operation avoid the overhead batched execution. The key idea is convert batched execution of different
9:
for E p do
10:
Compute h (u, v) with Eqn. 12
11:
Compute priority and (v) with Eqn. 10 11 12: end for
small samples into execution of single large sam-
ple, which can be paralleled by existing operations
13: end for
14: Share weights between suy (v) and the predictor
in deep learning frameworks. For example, the batched execution of topk([[1, 3). [2. O]] can be converted into a multi-key sort problem over [[o, 1], [0, 3), [1, 2), 1], [1, []] where the first key 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
15: return (TT) (v) as p(ulu, q) for E V
single-key sort operations in deep learning frameworks. See App. for details.
4 Experiments
We evaluate A*Net on standard transductive and inductive knowledge graph reasoning datasets, including a million-scal one ogbl-wikikg2. We conduct ablation studies to verify our design choices and visualize the important paths learned by the priority function in A Nett.

---

4.1 Experiment Setup
Datasets & Evaluation We evaluate A on 4 standard knowledge graphs, FB1 40)
WN18RR [16], YAGO3-10 [30] and ogbl-wikikg2 [25]. For the transductive setting, we use the standard splits from their original works [40 [6] 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 (v, q, v') I' or (u', q, u) 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 high frequency, and evaluate the checkpoints later 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-Fore algorithmf For fair comparison with existing path-based methods, we follow the implementation of NBFNet (58] and parameterize + with principal neighborhood aggregation (PNA) [3] sum aggregation, and parameterize X with the relation operation from
DistMult 49 i.e., vector multiplication. The indicator function (Eqn. 2 1, = v) = 1u u)q is parameterized with a query embedding q for all datasets except ogbi-wikikg2, where the indicator function with learnable embeddings based on a soft distance from to v (see App. E for more details). The edge representation (Eqn. T2) w. (2;, r, v) W.q + b, is parameterized as a linear function over the query relation for all datasets except WN18RR, where we use a simple embedding w,(2; r, v) = 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.
augment
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 K/[V| and maximum average degree ratio 3 L and tune the ratios for each dataset. The maximum edge ratio is determined by a.3. The other hyperparameters are kept the same the values in [58] We train A *ett with 4 Tesla A 100 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 bashd methods. The embedding methods are TransE (6) ComplEx [42] RotatE (38). HAKE (55), RotH 77), PairRE (8),
ComplEx+Relation Prediction 12 and ConE B]. The GNNs are RGCN 336. CompGCN
45 and
GralL 39
The path-based methods are MINERVA (14, Multi-Hop 29 CURL [2 NeuralLP
50
DRUM 35
NBFNet (58) and RED-GNN 55 Note that path-finding methods 14
29 52 that use
reinforcement learning and assume sparse answers can only be evaluated tail prediction. Training time of all basclines are measured based on their official open- implementations, except that we use a more recent implementation of TransE and ComplEx.
4.2 Main Results
Tab. T shows that *eet outperforms all embedding methods and GNNs, and is on par with NBFNet on transductive knowledge graph reasoning. We also observe similar trend of **tet 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 Nett A Nett also converges faster than all the other methods (Fig. 5). Notably, unlike NBFNet that propagates through all nodes and edges, *eNet 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 that "Net reduces the number of messages by 14.1 x and 42.9x compared to NBFNet on two datasets respectively. Note that the reduction in time and memory is less than the reduction 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. shows the performance oiki-22 which has 2.5 million entities and 16 million triplets. While NBFNet faces out-of out-of-memory (OOM) problem even for a batch size of 1, *ett can perform reasoning by propagating through 0.2% nodes and 0.2% edges at each step. Surprisingly, even with
httppr//gghhhhh.co///p/cpppaaaap MIT license.
http:://sshhhcccoooeeeeeel/ll/es MIT license.
7

---

Table 1: Performance on transductive knowledge graph reason-
FBB55k2237
ing. Results of embedding methods are from []. Results of
GNNs and path- methods are from 58
Performance and
efficiency on YAGO3-10 10 are in App.
MRX
Method
MRR
231
1663
H@@1
WNISRR
H@9
Hd
0.2
TransE
0.294
0.465
0.226
0.403
0.532
RotadE
0.338
0.241
0.575
0.533
0.476
0.428
0.492
0.571
0.1
RosatE
HAKE
RotH
0.341
0,344
0.243
0.246
0.378
0,580
0.535
0.535
0.496
0.495
0.451
0,449
0.514
0.513
0.582
0.586
Training Time (hrs)
ConF.
CompGCN
ComplEx+RP
0,388
0.298
0.425
0,568
0.488
0.443
0,505
0.578
WNISRR
NBENet
ConE
0.345
0.247
0.381
0.540
0.515
0.579
0,55
ANNa
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.566
MRR
0,50
DRUM
0,343
0.255
0.378
0.516
0.486
0.425
0.513
0.586
NBENet
0.415
0.321
0.454
0.599
0.666
0,45
RED-GNN
0.374
0.283
0.558
0.533
0.485
0.624
Net
0.411
0.321
0.453
0.586
0.549
0.495
0.573
0.689
Training Time (hrs)
Table 2: Tail prediction performance on transductive knowledge graphs. Results of compared methods are from [14 29
Figure 5: Validation MRR w.r.t. training time (1 A100 GPU).
52
FB15k-237
WN1SRR
Method
MRR
H@! 1
H@3
H@ 10
MRR
H@1
H@3
H@ 10
Table 4: Performance ogbl-
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
wikikg2 (MRR). Results of com-
CURL
0.306
0.224
0.341
0.470
0.460
0.429
0.471
0.523
pared methods are from 8 12.
NBFNet
0.509
0.411
0.562
0.697
0.587
0.503
0.579
0.669
A Net
0.505
0.410
0,556
0.687
0.557
0.504
0,580
0.666
ogbl-wikkkk2 Valid
Method
Test
#Params
Table 3: Efficiency on transductive knowledge graph reasoning.
Trand:
0.4256
0.4272
1,251 M
FB15K-237
WNISRR
ComplEx
0.4027
0.3759
1,251 M
Method
immsssage
time
memory
Armessage
time
memory
PaiRE
RotatE
0.5208
0,4332
0,5423
1,251
500 M
NBENet
544.230
16.8 min
19.1 GiB
173,670
9.42 min
26.4 GiB
ComplEx-RP
0.6392
0.6561
250 M
Net
38,610
8.07 mia
11.1 GiB
4,049
1.39 min
5.04 GiB
NBENet
OOM
0OM
0OM
Improvement
14.1x
2
1,7x
42.9x
688x
5.2x
N
0.6767
6,83
such sparse propagation, Aett outperforms embedding methods and achieves a new statc-of-the-art result. Moreover, the validation curve in Fig. [ shows that A *eNet converges significantly faster than embedding methods. Since A only learns parameters for relations but not entities, it only uses 6.83 million parameters, which 36.6x less than the best embedding method ComplEx+RP.
Tab. shows the performance on inductive knowledge graph reasoning. * 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.
4.3 Ablation Studies
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 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 neural priority function.
Priority Function
Sharing Weights As discussed in Sec. 3.2, share the weights between the neural priority function and the reasoning predictor to help train the neural priority function. Tab, ccompares A *ett trained with and without sharing weights. It can be observed that sharing weights essential to training a good neural priority function in A Net.
Trade-off between Performance and Efficiency While *eNet matches the performance of NBFNet in less training time, one may further trade off performance and efficiency in * by adjusting the ratios and B. Fig. plots curves of performance and speedup ratio w.r.t. different
8

---

Table 5: Performance on inductive knowledge graph reasoning (MRR). V1-v4 are 4 standard inductive splits. Results of com- pared methods are taken from [54) a = 50% and / 100% for FB15k237. a = 5% and [3 100% for WN18RR. More metries and efficiency results are in App.
Table 6: Ablation studies of A *Net on transductive FB 133--2377
(a) Choices of priority function.
Priority
FB15k-237
Function
MRR
H@l
H@3
H@ 10
Method
FB15k-237
WNI8RR
PPR
0.266
0.212
0.296
0.371
v1
v2
v1
v2
v3
Degree
0.347
0.268
0.383
0.501
Grall.
0.279
0.276
0.251
0.227
0.627
0.625
0.323
0.553
Neural
Random
0.378
0.411
0.288
0.321
0.413
0.453
0.556
0.586
NeuralLP
0.325
0.389
0.400
0.396
0.649
0.635
0.361
0.628
0.333
0.395
0.666
0.646
0,380
(b) W/ w/o sharing weights.
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
Sharing Weights
FB15K-237
Neet
0.457
0.510
0.476
0.466
0.727
0.704
0.441
0.661
MRR
H@l
H@3
H@ 10
0.282
0.413
0.557
Yes
0.411
0.321
0.453
0.586
Aitit
MRRR
0.44
Speedup Ratio OK)
Speakic Apcolleg Rlatinn OK)
(aandei, inchstry, ?)
300K
MILK
0,50
Bandai Namoo
subsidiary
industry
Banki
video gamne
0.3
10
20
100
32
100
Industry
Node
8-10055
(103), 0-10%%
incustry
media
Pony Canyon
Figure Performance and efficiency trade-off w.r.t. node ratio and degree ratio B. Speedup ratio is relative to NBFNet.
Figure 7: Visualization of important paths
learned by the neural priority function in Net.
a and 3. If we can accept a performance similar to embedding methods (e.g., ConE 3 ) we can set either to 1% 3 to 10%, resulting in 8.7x speedup compared to NBFNet.
4.4 Visualization of Learned Important Paths
We can extract the important paths from the neural priority function in Net for interpretation. For a given query (u. q. ?) and predicted entity v, we can use the node priority (t) (x) at each step to
estimate the importance of a path. Empirically, the importance of a path 8g (P) estimated by
sp(P)
1 P
Σ
Sug (x)
(t-1)
(14)
-(t-1)
P(0
where s(t-1) max Sug (x) is a normalizer normalize the priority score for each step t. To extract the important paths with large sy (P), we perform beam search over the priority function
(t-1)
(t-1) (x) of each step. Fig. 7 shows the important paths learned by Net for a test sample in Sug
FB15k-237. Given the query (Bandai, industry, ?), we can see both paths Bandai subsidiary
Bandai
Namco induztry video game and Bandai industry, media industry Pony Canyon industty video game are consistent with human cognition. More visualization results can be found in App. G
5 Related Work
Path-based Reasoning Path-based methods use paths between entities for knowledge graph rea- soning. Early methods like Path Ranking [28] [9] collect relational paths symbolic features for classification. Path-RNN [5 and PathCon [45] improve Path Ranking by learning the represen- tations of paths with recurrett neural networks (RNN). However, these works operate 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 edges.
To avoid the exhaustive search of paths, many methods learn to sample important paths for reasoning. DeepPath [47 and MINERVA [4] 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-graines paths 52
and [33 use a variational formulation to learn sparse prior for path sampling. Another category
of methods utillizes 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 [4]]
adopts Woyd-Warshall-Bie
algorithm to learn path representations between all pairs of entities.
Recently, NBFNet [58] and RED-GNN [54] leverage a Hollman-Ford-iike algorithm to learn path
representations from single-source entity to all entities. While dynamic programming methods achieve state- tfte-oof-the-ar results among path- based methods, they need to perform message passing on the full knowledge graph. By comparison, our A Neet learns priority function and only explores 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 2] [T reduce the cost of message passing by computing GNNs with sampled subset of nodes and edges. Non-parametris GNNs [27 46
18 IO] decouple feature propagation from feature transformation, and reduce time complexity by preprocessing feature propagation. However, both sampling methods and non-pas 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 attention to iteratively select nodes for message passing. SQALER [ first predicts important path types based on the query, and then applies GNNs on the subgraph extracted by the predicted paths. Our *eNet shares the same goal with these methods, but learns a neural priority function to iteratively select important paths.
Discussion and Conclusion
Limitation and Future Work One limitation for NeN 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. 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 more scalable model.
Conclusion We propose A Nett 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 *eNt. Meanwhile, A Net is the first path-bases method that scales to million-scale knowledge graphs.
Acknowledgement
This project is supported by Intel-MILA partnership program, the Natural Sciences and Engineering
Research Council (NSERC) Discovery Grant, the Canada CIFAR Al Chair Program, collaboration grants between Microsoft Research and Mila, Samsung Electronics Co., Ltd., Amazon Faculty Research Award, Tencent AI Lab Rhino-Birt -ird Gift Fund and NRC Collaborative R&D Project (A14D-CORE CORE-06). This project was also partially funded by IVADO Fundamental Research Project grant PRF- 2019-2019-20 0101999553319999 Québed and the Digital Research Alliance of C
The computation resource of this project is supported by Mile Calcul
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.
https://mila. quebec/
http://www.c.culuuueebe....c wttp://w.cll......mmmc cal ca/
https://alllanceeaa.
10

---

References
[1] Mattia Atzeni, Jasmina Bogojeska, and Andreas Loukas. Sqaler: Scaling question answering by decoupling multi-hop and logical reasoning. Advances in Neural Information Processing Systems, 34, 2021.
[2] Sören Auer, Christian Bizer, Georgi Kobilarov, Jens Lehmann, Richard Cyganiak, and Zachary Ives. Dbpedia: A nucleus for a web of open data. In The semantic web, pages 722-735. Springer, 2007.
[3] Yushi Bai, Zhitao Ying, Hongyu Ren, and Jure Leskovec. Modeling heterogeneous hierarchies with relation- relation-specific hyperbolic cones. Advances in Neural Information Processing Systems, 34, 2021.
[4] Richard Bellman. On a routing problem. Quarterly of applied mathematics, 16(1):87-900 1958.
[5] Jonathan Berant, Andrew Chou, Roy Frostig, and Percy Liang. Semantic parsing on freebase answer pairs. Proceedings of the 2013 conference on empirical methods in natural language processing, pages 1533-1544, 2013.
from question -eeestion-answer
[6] Antoine Bordes, Nicolas Usunier, Alberto Garcia-Dural Jason Weston, and Oksana Yakhnenko. Translating embeddings for modeling multi-relational data. Advances in neural information processing systems, 26, 2013.
[7] Ines Chami, Adva Wolf, Da- Da-Cheng Juan, Frederic Sala, Sujith Ravi, and Christopher Ré. Low- dimensional hyperbolic knowledge graph embeddings. Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 6901-6914, 2020.
[8] Linlin Chao, Jianshan He, Taifeng Wang, and Wei Chu. Pairre: Knowledge graph embeddings via paired relation vectors. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pages 4360-4369, 2021.
[9] Jie Chen, Tengfei Ma, and Cao Xiao. Fastgen: Fast learning with graph convolutional networks via importance sampling. In International Conference Learning Representations, 2018.
[10] Ming Chen, Zhewei Wei, Bolin Ding, Yaliang Li, Ye Yuan, Xiaoyong Du, and Ji- Ji-Rong Wen. Scalable graph neural networks via bidirectional propagation. Advances in neural information processing systems, 33: 15556-15666, 2020.
[11] Wenhu Chen, Wenhan Xiong, Xifeng Yan, and William Yang Wang. Variational knowledge graph reasoning. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume I (Long Papers), pages 1823-1832, 2018.
[12] Yihong Chen, Pasquale Minervini, Sebastian Riedel, and Pontus Stenetorp. Relation prediction as an auxiliary training objective for improving multi-relational graph representations. In 3rd Conference on Automated Knowledge Base Construction, 2021.
[13] Gabriele Corso, Luca Cavalleri, Dominique Beaini, Pietro Liò, and Petar Veličković. Principal neighbourhood aggregation for graph nets. Advances in Neural Information Processing Systems, 33:13260-13221 2020.
[14] Rajarshi Das, Shehzaad Dhuliawala, Manzil Zaheer, Luke Vilnis Ishan Durugkar, Akshay Krishnamurthy, Alex Smola, and Andrew McCallum. Go for walk and arrive at the answer: Reasoning over paths in knowledge bases using reinforcement learning. International Conference Learning Representations, 2018.
[15] Rajarshi Das, Arvind Neelakantan, David Belanger, and Andrew McCallum. Chains of reasoning over entities, relations, and text using recurrent neural networks. In Proceedings of the 15th Conference of the European Chapter of the Association for Computational Linguistics: Volume 1, Long Papers, pages 132-141, 2017.
11

---

[16] Tim Dettmers, Pasquale Minervini. Pontus Stenetorp, and Sebastian Riedel. Convolutional 2d knowledge graph embeddings. In Proceedings of the AAAI conference on artificial intelligence. volume 32, 2018.
[17] Arthur Feeney, Rishabh Cupta, Veronika Thost, Rico Angell, Gayathri Chandu, Yash Adhikari, and Tengfei Ma. Relation matters in sampling: A scalable multi-relational graph neural network for drug-drup drug interaction prediction. arXiv preprint arXiv:2105.13975 2021.
[18] Fabrizio Frasca, Emanuele Rossi, Davide Eynard, Ben Chamberlain, Michael Bronstein, and Federico Monti. Sign: Scalable inception graph neural networks. arXiv preprint arXiv:2004.11190 2020.
[19] Matt Gardner and Tom Mitchell. Efficient and expressive knowledge base completion using subgraph feature extraction. In Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing, pages 1488- 1498, 2015.
[20] Ben Goertzel and Cassio Pennachin. Artificial general intelligence, volume 2. Springer, 2007.
[21] Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs. Advances in neural information processing systems, 30, 2017.
[22] Peter E Hart, Nils J Nilsson, and Bertram Raphael. A formal basis for the heuristic determination of minimum cost paths. IEEE transactions on Systems Science and Cybernetics, 4(2):100-107, 1968.
[23] Udo Hebisch and Hanns Joachim Weinert. Semirings: algebraic theory and applications in computer science, volume 5. World Scientific, 1998.
[24] Marcel Hildebrandt, Jorge Andres Quintero Serna, Yunpu Ma, Martin Ringsquandl, Mitchell Joblin, and Volker Tresp. Reasoning on knowledge graphs with debate dynamics. In Proceedings of the AAAI Conference Artificial Intelligence, volume 34, pages 4123 223-4131, 2020.
[25] Weihua Hu, Matthias Fey, Hongyu Ren, Maho Nakata, Yuxiao Dong, and Jure Leskovec. Ogb- Isc: A large-scale challenge for machine learning on graphs. arXiv preprint arXiv:2103.094
2021.
[26] Wenbing Huang, Tong Zhang, Yu Rong, and Junzhou Huang. Adaptive sampling towards fast graph representation learning. Advances in neural information processing systems, 31, 2018.
[27] Johannes Klicpera, Aleksandar Bojchevski, and Stephan Günnemann. Predict then propagate: Graph neural networks meet personalized pagerank. In International Conference on Learning Representations, 2018.
[28] Ni Lao and William W Cohen. Relational retrieval using a combination of path-constrained random walks. Machine learning, 81(1):33-67, 2010.
[29] Xi Victoria Lin, Richard Socher, and Caiming Xiong. Multi-hop knowledge graph reasoning with reward shaping. In EMNLP. 2018
[30] Farzaneh Mahdisoltani, Joanna Biega, and Fabian Suchanek. Yago3: knowledge base from multilingual wikipedias. In 7th biennial conference on innovative data systems research. CIDR Conference, 2014.
[31] Arvind Neelakantan, Benjamin Roth, and Andrew McCallum. Compositional vector space models for knowledge base completion. In Proceedings of the 53rd Annual Meeting of the Association for Computational Linguistics and the 7th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pages 156-166, 2015.
[32] Judea Pearl. Probabilistic reasoning in intelligent systems: networks of plausible inference. Morgan kaufmann, 1988.
[33] Meng Qu, Junkun Chen, Louis-Pascal Xhonneux, Yoshua Bengio, and Jian Tang. Rnnlogic: Learning logic rules for reasoning on knowledge graphs. In International Conference on Learning Representations, 2021.
12

---

[34] Hongyu Ren, Mikhail Galkin, Michael Cochez, Zhaocheng Zhu, and Jure Leskovec. Neural graph reasoning: Complex logical query answering meets graph databases. arXiv preprint arXiv:2303. /4617, 2023.
[35] Ali Sadeghian, Mohammadreza Armandpour, Patrick Ding, and Daisy Zhe Wang. Drum: End-to-end differentiable rule mining on knowledge graphs. Advances in Neural Information Processing Systems, 32, 2019.
[36] Michael Schlichtkrull, Thomas Kipf. Peter Bloem, Rianne van den Berg, Ivan Titov, and Max Welling. Modeling relational data with graph convolutional networks. In European semantic web conference, pages 593-607. Springer, 2018.
[37] Yelong Shen, Jianshu Chen, Po-Sen Huang, Yuqing Guo, and Jianfeng Gao. M-walk: Learning to walk over graphs using monte carlo tree search. Advances in Neural Information Processing Systems, 31, 2018.
[38] Zhiqing Sun, Zhi-Hong Deng, Jian- Yun Nie, and Jian Tang. Rotate: Knowledge graph em- bedding by relational rotation in complex space, In International Conference on Learning Representations, 2019.
[39] Komal Teru, Etienne Denis, and Will Hamilton. Inductive relation prediction by subgraph reasoning. In International Conference on Machine Learning, pages 1458-9557. PMLR, 2020.
[40] Kristina Toutanova and Danqi Chen. Observed versus latent features for knowledge base and text inference. In Proceedings of the 3rd workshop on continuous vector space models and their compositionality, pages 37-66, 2015.
[41] Kristina Toutanova, Xi Victoria Lin, Wen-tau Yih, Hoifung Poon, and Chris Quirk. Composi- tional learning of embeddings for relation paths in knowledge base and text. In Proceedings of the 54th Anmual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1434-1444 1444, 2016.
[42] Théo Trouillon, Johannes Welbl, Sebastian Riedel, Éric Gaussier, and Guillaume Bouchard. Complex embeddings for simple link prediction. In International Conference on Machine Learning, pages 2071-2080. PMLR, 2016.
[43] Shikhar Vashishth, Soumya Sanyal, Vikram Nitin, and Partha Talukdar. Composition- multi-relational graph convolutional networks, In International Conference Learning Repre- sentations, 2020.
[44] Denny Vrandežić and Markus Krötzsch. Wikidata: a free collaborative knowledgebase. Com- munications of the ACM, 57(10):78-85, 2014.
[45] Hongwei Wang, Hongyu Ren, and Jure Leskovec. Relational message passing for knowledge graph completion. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining, pages 1697-1707, 2021.
[46] Felix Wu, Amauri Souza, Tianyi Zhang, Christopher Fifty, Tao Yu, and Kilian Weinberger. Simplifying graph convolutional networks. In International conference on machine learning. pages 6861-6871. PMLR, 2019.
[47] Wenhan Xiong, Thien Hoang, and William Yang Wang. Deeppath: A reinforcement learning method for knowledge graph reasoning. Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing (EMNLP 2017), Copenhagen, Denmark, September 2017. ACL.
[48] Xiaoran Xu, Wei Feng, Yunsheng Jiang, Xiaohui Xie, Zhiqing Sun, and Zhi-Hong Deng. Dynamically pruned message passing networks for large-scale knowledge graph reasoning. In International Conference on Learning Representations, 2019.
[49] Bishan Yang, Wen-tau Yih, Xiaodong He, Jianfeng Gao, and Li Deng. Embedding entities and relations for learning and inference in knowledge bases. International Conference on Learning Representations, 2015.
13

---

[50] Fan Yang, Zhilin Yang, and William W Cohen. Differentiable learning of logical rules for knowledge base reasoning. 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 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 Quamming 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 hicrarchy-award knowledge graph embeddings for link prediction. In Proceedings the AAAI Conference 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 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-performancee epu-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.
14


Reviews:
[{"review_id": "lBDoiWE9rF", "reviewer": "Reviewer_rgXX", "summary": "This paper proposes a scalable path-based knowledge graph reasoning approach. The idea is to extend only important paths from the exponentially growing set of all possible paths. A heuristic priority function is parametrized by a feed-forward network and is trained to predict the priority of nodes to expand. Experiments show that the proposed approach can significantly improve time and memory efficiency and also achieve good results.", "questions": "Please see Weaknesses", "limitations": "properly addressed", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- Scalability is an important issue for path-based reasoning approaches. The idea of selecting only important paths is interesting and sounds reasonable\n- The proposed approach is effective and supported by extensive experiments. Time and memory efficiency has been significantly improved. Benchmark results are also good.", "weaknesses": "My concern is mainly about the design of the priority function Eq (10)\n\nIn Eq (10), the first part $h_q^{(t)}(u, x)$ is already conditioned on q, u, and x, so in principle the second part $g([h_q^{(t)}(u, x), q])$ doesn't provide any additional information. Therefore, the priority function is purely based on the current path from the start and contains no information about the goal. In other words, the prediction of the priority function would be the same even if the goal changes. This is different from the design of the A* algorithm and may lose theoretical guarantees. \n\nIt is not appropriate to present the approach in the manner of A* algorithm", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your comments. We notice that you raised only one concern for the rating of 3. Feel free to bring up other concerns for discussion. Here is our response to your concern.\n\n**W1: The first part $h^{(t)}_q(u, x)$ is already conditioned on q, u, and x, so in principle the second part $g([h^{(t)}_q(u, x), q])$ doesn't provide any additional information. The priority function is purely based on the current path from the start and contains no information about the goal. The prediction of the priority function would be the same even if the goal changes. This is different from the design of the A\\* algorithm and may lose theoretical guarantees.**\n\nA1: It is a misunderstanding that the priority function contains no information about the goal. For a given query $(u,q,?)$, the goal is the set of answer nodes $\\mathcal{V}_{ans} = \\\\{ v | (u,q,v) \\in \\mathcal{E} \\\\}$.\n\nSince the set $\\mathcal{V}_{ans}$ is a function of $u$, $q$ and $\\mathcal{E}$, we can represent the goal with source node $u$ and the query relation $q$. Note that we exclude the ground truth triplets $\\mathcal{E}$ because they are unknown during inference. A more intuitive figure can be found in the PDF file attached to the global rebuttal response.\n\nThe vector representation $\\mathbf{s}_{uq}^{(t)}(x)$ for the priority function in A\\*Net (Eqn. 10) is designed to match the priority function $s(x)$ in the A\\* algorithm (Eqn. 4).\n\nFor a given $u$, $x$ and $t$, $\\mathbf{s}_{uq}^{(t)}(x)$ will be different if we change the goal $q$. We will make the correspondence between A\\*Net and the A\\* algorithm more explicit in the paper.\n\nHere we illustrate the correspondence between the terms in A\\*Net and the A\\* algorithm. The priority function in the A\\* algorithm (Eqn. 4) is\n\n$$s(x) = d(u,x) \\otimes g(x,v)$$\n\nwhere $d(u, x)$ is the length of the current shortest path from $u$ to $x$. $g(x, v)$ is a heuristic function that estimates the remaining length from $x$ to the target node $v$ (i.e. the goal). Typically, $g(x,v)$ is defined as the $L_1$ distance from $x$ to $v$.\n\nThe vector representation $\\mathbf{s}_{uq}^{(t)}(x)$ in A\\*Net (Eqn. 10) is\n\n$$\\mathbf{s}_{uq}^{(t)}(x) = \\mathbf{h}_q^{(t)}(u,x) \\otimes \\mathbf{g}([\\mathbf{h}_q^{(t)}(u,x), \\mathbf{q}])$$\n\nwhere $\\mathbf{h}_q^{(t)}(u, x)$ corresponds to $d(u, x)$ and represents the aggregation of paths from $u$ to $x$ within $t$ hops. The optional subscript $q$ means that the representation $\\mathbf{h}_q^{(t)}(u, x)$ is conditioned on the goal $q$. In other words, $\\mathbf{h}_q^{(t)}(u, x)$ mostly aggregates paths that are relevant to the goal $q$. One may also opt for an unconditioned version $\\mathbf{h}^{(t)}(u, x)$ which aggregates all the paths uniformly. Here we follow NBFNet[1] and use this conditioned parameterization.\n\nWe agree with you that the second part $\\mathbf{g}([\\mathbf{h}_q^{(t)}(u,x), \\mathbf{q}])$ is conditioned on the same set of variables $u$, $q$ and $x$ as the first part. However, two architectures that model the same information don’t suggest they have the same inductive bias and generalization ability. For example, CNNs are better than MLPs on images because they have the inductive bias for translation equivariance, despite that they take the same information as input[2]. In graph machine learning, several papers[3, 4, 5] also suggest algorithmic alignment is the key to the success of GNNs on graphs.\n\nHere we show why $\\mathbf{g}([\\mathbf{h}_q^{(t)}(u,x), \\mathbf{q}])$ is aligned with $g(x,v)$ in the A\\* algorithm. The learned representation $\\mathbf{q}$ captures the relative position from the source node $u$ to the target node $v$. For example, if $q$ is the mother relation, then the aggregation of paths between the source node $u$ and an answer node $v$ should roughly match the representation of mother $\\mathbf{q}$, i.e. $\\mathbf{q} \\approx \\mathbf{h}^{(T)}_q(u, v)$ for any $(u,q,v)\\in\\mathcal{E}$. The reason why $\\mathbf{q}$ can be independent of $u$ and $v$ is that the definition of a relation is independent of its triplet instances. $\\mathbf{h}_q^{(t)}(u,x)$ is the current aggregation of paths from $u$ to $x$ and represents the relative position from $u$ to $x$, conditioned on the query relation $q$. By transforming $\\mathbf{h}_q^{(t)}(u,x)$ and $\\mathbf{q}$ with a function $g(\\cdot)$ (e.g. a vector substraction function), we can obtain an estimate of the relative position from $x$ to $v$, which corresponds to $g(x,v)$.\n\n[1] Zhu et al. Neural Bellman-Ford Networks: A General Graph Neural Network Framework for Link Prediction. NeurIPS 2021.\n\n[2] Bronstein et al. Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges. arXiv 2021.\n\n[3] Xu et al. What Can Neural Networks Reason About? ICLR 2020.\n\n[4] Xu et al. How Neural Networks Extrapolate: From Feedforward to Graph Neural Networks. ICLR 2021.\n\n[5] Dudzik and Velickovic. Graph Neural Networks are Dynamic Programmers. NeurIPS 2022."}, {"review_id": "QZnCikyfWs", "reviewer": "Reviewer_7LW7", "summary": "The main contribution of this paper is presenting a scalable path-based method A*Net, for link prediction on large-scale knowledge graphs. A*Net is inspired by the A* algorithm for solving shortest path problems, where it learns a priority function to select important nodes and edges at each iteration. This allows for the time and memory reducing for both training and inference. From an efficiency perspective, this could be considered as a path-pruning method to progressively reduce the subgraph based on the learned priority function. The empirical results also demonstrate efficiency improvement.", "questions": "1. As addressed in Weaknesses 3, could you elaborate on how weight sharing aids in the selection of important paths?\n\n2. I observe that two handcrafted priority functions, PPR and Degree, are employed in the ablation studies. Given that high connectivity doesn't necessarily denote the importance of paths, what about the effectiveness and efficiency of a random pruning strategy, particularly with respect to the obgl_wikikg2 dataset?\n\n3. In the Visualization section, only the results of the proposed method are displayed without any comparison. Could you clarify what distinct paths the Neural function selects compared to the two handcrafted ones? Furthermore, does the Neural-based path selection align more closely with knowledge semantics?\n", "limitations": "Yes. The authors stated the limitation, future work and social impact.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The efficiency problem caused by the explosively increasing entities in deeper propagation layers is indeed serious in the recent GNN-based inductive methods. And the proposed method makes sense and technically sound. \n\n2. The experimental results are impressive. The paper demonstrates the practical applications of A*Net in various settings and datasets, with the efficiency improvement compared with several recent baselines. Furthermore, the paper sets a new state-of-the-art on the million-scale dataset ogbl-wikikg2 and converges faster than embedding methods. \n\n3. The paper's organization is well-executed and the content is easily comprehensible.\n", "weaknesses": "1. The paper's comparison to the A* algorithm seems somewhat overstated. As a derivative work of NBFNet, this paper draws an analogy to another shortest path algorithm, A*. Contrary to the Bellman-Ford algorithm that resolves the shortest path problem from the source to all other points, the A* algorithm typically addresses the shortest path problem from the source to a specific target point. However, in the context of knowledge graph (KG) reasoning, the target point is unknown, rendering the core principle of A*, assessing the estimated remaining cost to the target point, unfeasible. In fact, the A* algorithm's priority rule, involving the distance to the target node, is not pertinent to the priority function in the proposed model. The A* algorithm appears to function primarily as a promotional point, rather than as a guiding principle.\n\n2. Perhaps due to the overemphasis on the A* analogy, the paper's true technical contributions remain unclear. Comparing the core function of NBFNet in Eq. 3 and that of A*Net in Eq. 12, the only discernible difference lies in introducing the priority score, calculated based on the embeddings of the query and the current node. Stripping away the A* algorithm framework, it essentially seems to be a path-pruning technique reliant on an attention mechanism to select the top K nodes and top L edges in each layer for efficiency's sake.\n\n3. The paper lacks insightful contributions regarding important paths beyond a weighted version of the NBENet method. The theoretical appendix focuses solely on integrating path selection into the NBFNet framework, premised on the assumption that a certain function can distinguish important nodes. However, how to ensure that important paths are chosen is not clear. In response to this, the authors propose weight sharing between the priority function and the predictor, asserting that the reasoning task can be seen as a weak supervision for the priority function. However, this appears counterintuitive, given that the priority score is dependent on a specific query. A high predictor score, indicating that the node x answers the query (u, r_1), should not contribute to the priority score of x for a different query (u, r_2).\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your constructive comments. Here is our response to your concerns.\n\n**W1: In KG reasoning, the target point is unknown, rendering the core principle of A\\* unfeasible. In fact, the A\\* algorithm's priority rule, involving the distance to the target node, is not pertinent to the priority function in the proposed model.**\n\nA1: We agree that estimating the remaining distance is the core principle of the A\\* algorithm. Our A\\*Net follows this principle. We design Eqn. 10 in A\\*Net to match Eqn. 4 in the A\\* algorithm. $\\mathbf{h}_q^{(t)}(u, x)$ corresponds to the current length $d(u, x)$, while $\\mathbf{g}([\\mathbf{h}_q^{(t)}(u,x), \\mathbf{q}])$ corresponds to the remaining distance $g(x, v)$. A figure illustration is in the attached PDF file.\n\nYour misunderstanding might come from the fact that we don’t have $v$ in Eqn. 10. In fact, the learned representation $\\mathbf{q}$ encodes the relative position from $u$ to $v$. e.g. If $q$ is the mother relation, the aggregation of paths between $u$ and $v$ in a positive sample should roughly match the representation of mother $\\mathbf{q}$, i.e. $\\mathbf{q} \\approx \\mathbf{h}^{(T)}_q(u, v)$ for any $(u,q,v)\\in\\mathcal{E}$. The reason why $\\mathbf{q}$ can be independent of $u$ and $v$ is that the definition of a relation is independent of its triplet instances. The representation of the remaining cost is (due to TeX bugs, $g$ and $h$ should be $\\mathbf{g}$ and $\\mathbf{h}$ resp.)\n\n$g^{(t)}(x,v)\\approx g([h_{q}^{(t)}(u,x),h_{q}^{(T)}(u,v)])\\approx g([h_{q}^{(t)}(u,x),\\mathbf{q}])$\n\nwhere the first approximation says that the remaining cost between $x$ and $v$ can be estimated by the current aggregation of paths $\\mathbf{h}_{q}^{(t)}(u,x)$ between $u$ and $x$\n\nand the final path representations $\\mathbf{h}_{q}^{(T)}(u,v)$ between $u$ and $v$.\n\nThe second approximation replaces $\\mathbf{h}_{q}^{(T)}(u,v)$ with $\\mathbf{q}$. So $\\mathbf{g}([\\mathbf{h}_q^{(t)}(u,x), \\mathbf{q}])$ matches the remaining distance $g(x, v)$ in the A\\* algorithm.\n\n**W2: Due to the overemphasis on the A\\* analogy, the paper's true technical contributions remain unclear.**\n\nA2: We clarified the analogy between A\\*Net and the A\\* algorithm in A1 and they align at each term in the priority function. Our major contribution is the first path-based method that scales to ogbl-wikikg2. Note that ogbl-wikikg2 has been previously dominant by embedding methods, and none of them are inductive or interpretable like A\\*Net. Hence we think this contribution is very important to the community and may potentially change future research directions. The theoretical insights from the A\\* algorithm are a minor contribution compared to our empirical achievements.\n\n**W3 & Q1: How to ensure that important paths are chosen is not clear. The weight sharing between the priority function and the predictor appears counterintuitive, given that the priority score is dependent on a specific query.**\n\nA3: Since there isn’t any annotation of the important paths, we can’t verify whether important paths are chosen or not. However, the observation that A\\*Net matches the performance of NBFNet suggests that A\\*Net captures most important paths. Note that the performance of A\\*Net is upper bounded by NBFNet, as A\\*Net visits much fewer paths.\n\nWe agree that the priority score should be dependent on a specific query, and that’s what we designed A\\*Net to be. Both the priority function and the predictor take a **query-dependent** representation $\\mathbf{s}_{uq}^{(t)}(x)$ as input, but the parameters in both functions (i.e. $f(\\cdot)$ in Eqn. 11) are **query independent**. This is consistent with previous works[1, 2, 3]. We share the **query-independent** parameters in these functions.\n\nFor two queries $(u, q_1, ?)$ and $(u, q_2, ?)$, they will have different representations $\\mathbf{s}_{uq_1}^{(t)}(x)$\n\nand $\\mathbf{s}_{uq_2}^{(t)}(x)$, and hence different priority scores.\n\n**Q2: What about the effectiveness and efficiency of a random pruning strategy?**\n\nA4: We run experiments for the random pruning strategy on FB15k-237 and ogbl-wikikg2. We set the random pruning strategy to have the same node and edge ratios as A\\*Net.\n\nA\\*Net outperforms the random pruning strategy on both datasets, and is slightly better than the random pruning strategy in time and memory. We conjecture the reason is that A\\*Net tends to revisit nodes more often than the random pruning strategy, resulting in a more cache-friendly behavior.\n\n|FB15k-237|MRR|H@1|H@3|H@10|#message|time|memory|\n|---|---|---|---|---|---|---|---|\n|Random|0.378|0.288|0.413|0.556|39,017|9.20min|16.9GiB|\n|A\\*Net|**0.411**|**0.321**|**0.453**|**0.586**|38,610|8.07min|11.1GiB|\n\n|ogbl-wikikg2|TestMRR|ValidMRR|#Params|#message|time|memory|\n|---|---|---|---|---|---|---|\n|Random|0.5815|0.5827|**6.83M**|51,458|1.74hr|26.5GiB|\n|A\\*Net|**0.6767**|**0.6851**|**6.83M**|52,371|1.30hr|24.1GiB|\n\n**Q3: Only the results of the proposed method are visualized without any comparison. What distinct paths the neural function selects compared to the two handcrafted ones? Does the neural-based path selection align more closely with knowledge semantics?**\n\nA5: Due to space limits, we put the visualization in the global response. We observe that the paths captured by A\\*Net are concise. By contrast, NBFNet and PPR find longer and more redundant paths. The degree priority function results in the closest behavior as A\\*Net, but the entities it visits are less relevant than those in A\\*Net.\n\nIt is hard to confirm that A\\*Net aligns better with knowledge semantics, as we don’t have any ground truth for important paths. We don’t intend to claim the quality of the visualization of A\\*Net, but just show its interpretability.\n\n[1] Zhu et al. A General Graph Neural Network Framework for Link Prediction. NeurIPS 2021.\n\n[2] Yang et al. Differentiable Learning of Logical Rules for Knowledge Base Reasoning. NIPS 2017.\n\n[3] Teru et al. Inductive Relation Prediction by Subgraph Reasoning. ICML 2020."}, {"review_id": "x3rtwZsf9i", "reviewer": "Reviewer_2kB4", "summary": "This paper presents a scalable path-based method for knowledge graph reasoning, which is inspired by the A* algorithm for shortest path problems.", "questions": "1. In the context of KG reasoning, a crucial question is, how many steps are typically required for a query? According to the vanilla path reasoning in Equation 1, the number of paths increases exponentially with respect to path length. However, if the path length is typically small, this might not pose a significant problem? Moreover, when dealing with a large-scale KG, the BF algorithm would need to visit $|\\mathcal{V}|$ nodes and $|\\mathcal{E}|$ edges for each step, which can be quite computationally intensive. Given these considerations, it leads to the question: If the path length is usually small, could vanilla path reasoning be a more efficient choice compared to BF?\n\n2. Another question is, can we simply leverage the idea of beam search into vanilla path reasoning? For example, we keep top-K ranked paths for each step, which may also avoid the exponential growth of the number of paths.", "limitations": "Yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The intriguing approach of applying the A$^*$ algorithm's principle to path reasoning in KG is proposed in this paper, along with the introduction of novel methods for crafting the priority function.\n\n2. The paper achieves state-of-the-art results on the large-scale KG reasoning dataset, ogbl-wikikg2.\n\n3. There's a substantial enhancement in efficiency, considering both time and memory usage, as opposed to the top-performing baseline, NBFNet.", "weaknesses": "The proposed method performs slightly worse than NBFnet as shown in Table 1, and no results of NBFnet are reported on tail prediction in Table 2.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your recognition and constructive comments. Here is our response to your concerns.\n\n**W1: The proposed method performs slightly worse than NBFNet as shown in Table 1, and no results of NBFNet are reported on tail prediction in Table 2.**\n\nA1: By design, A\\*Net cannot be better than NBFNet in performance, and our goal is to show that the gap between A\\*Net and NBFNet is very small (e.g. only 1% absolute difference in H@10 and less than 1% absolute difference in other metrics in Tab. 1).\n\nThe reason is that A\\*Net explores strictly less paths (e.g. 10% nodes and 10% edges on FB15k-237) than NBFNet in both training and inference. In a fair comparison, the performance of A\\*Net should be roughly upper bounded by the performance of NBFNet. While it is possible to improve A\\*Net by tweaking the neural parameterization and hyperparameters, it will result in an unfair comparison and blur the contribution of this paper.\n\nHere are the results including NBFNet on tail prediction. We can see that A\\*Net achieves competitive performance compared to NBFNet on both datasets. We will update these results in the paper.\n\n|FB15k-237|MRR|H@1|H@3|H@10|\n|---|---|---|---|---|\n|MINERVA|0.293|0.217|0.329|0.456|\n|Multi-Hop|0.393|0.329|-|0.544|\n|CURL|0.306|0.224|0.341|0.470|\n|NBFNet|**0.509**|**0.411**|**0.562**|**0.697**|\n|A\\*Net|**0.505**|**0.410**|**0.556**|0.687|\n\n|WN18RR|MRR|H@1|H@3|H@10|\n|---|---|---|---|---|\n|MINERVA|0.448|0.413|0.456|0.513|\n|Multi-Hop|0.472|0.437|-|0.542|\n|CURL|0.460|0.429|0.471|0.523|\n|NBFNet|**0.557**|**0.503**|**0.579**|**0.669**|\n|A\\*Net|**0.557**|**0.504**|**0.580**|**0.666**|\n\n**Q1: In KG reasoning, how many steps are typically required for a query? If the path length is usually small, could vanilla path reasoning be a more efficient choice compared to Bellman-Ford?**\n\nA2: For path-based methods that use exhaustive search[1, 2, 3], they compute typically up to 3 steps due to poor scalability. For path-based methods that use the Bellman-Ford algorithm[4, 5], their ablation studies suggest that the performance keeps increasing with more steps until 6 steps.\n\nTo give an intuition of how infeasible exhaustive search is, we compute the number of paths for different lengths, averaged over all positive triplets in each dataset. Here is the statistics. The number of paths grow exponentially w.r.t. the length of the paths. We note that exhaustive search is more efficient than the Bellman-Ford algorithm only when the number of paths is less than the number of nodes $|\\mathcal{V}|$. This only holds for paths with length ≤ 1 on FB15k-237, length ≤ 4 on WN18RR and length ≤ 2 on ogbl-wikikg2 respectively. Therefore, vanilla path reasoning is usually not a good choice.\n\n||\\|V\\||Length=1|Length=2|Length=3|Length=4|Length=5|Length=6|\n|---|---|---|---|---|---|---|---|\n|FB15k-237|14,541|367.0|31943|9.014e6|1.023e9|2.411e11|3.199e13|\n|WN18RR|40,943|19.38|138.9|4705.5|35504|1.624e6|1.214e7|\n|ogbl-wikikg2|2,500,604|135698|2.337e6|3.792e11|7.027e12|1.077e18|2.131e19|\n\n**Q2: Can we simply leverage the idea of beam search into vanilla path reasoning? For example, we keep top-K ranked paths for each step, which may also avoid the exponential growth of the number of paths.**\n\nA3: Beam search can be applied to path-finding methods[6, 7, 8] that use reinforcement learning to find paths. However, these methods assume a sparse set of answers (typically ≤100 answers) and can only be evaluated on tail prediction. Their performance is much worse than path-based methods that operate on a dense set of paths (e.g. NBFNet, A\\*Net), as showed in Tab. 2.\n\nIt is not trivial to apply beam search on path-based methods like NBFNet. First, beam search requires a score function to prune intermediate steps, which isn’t present in the original NBFNet. Second, even if we have such a score function, beam search is not differentiable w.r.t. the scores at the intermediate steps. This means we can’t learn the score function but can only use a handcrafted score function, which is likely to be suboptimal.\n\n[1] Lao and Cohen. Relational Retrieval Using a Combination of Path-Constrained Random Walks. ML 2010.\n\n[2] Neelakantan et al. Compositional Vector Space Models for Knowledge Base Completion. IJCNLP  2015.\n\n[3] Wang et al. Relational Message Passing for Knowledge Graph Completion. KDD 2021.\n\n[4] Zhu et al. Neural Bellman-Ford Networks: A General Graph Neural Network Framework for Link Prediction. NeurIPS 2021.\n\n[5] Zhang and Yao. Knowledge Graph Reasoning with Relational Digraph. WWW 2022.\n\n[6] Xiong et al. DeepPath: A Reinforcement Learning Method for Knowledge Graph Reasoning. ACL 2017.\n\n[7] Das et al. Go for a Walk and Arrive at the Answer: Reasoning Over Paths in Knowledge Bases using Reinforcement Learning. EACL 2017.\n\n[8] Lin et al. Multi-Hop Knowledge Graph Reasoning with Reward Shaping. EMNLP 2018."}, {"review_id": "MMitGOytXm", "reviewer": "Reviewer_ALvG", "summary": "The paper introduces ANet, a scalable path-based approach for reasoning on extensive knowledge graphs (KGs). In contrast to embedding techniques, path-based methods exhibit inductive capabilities but encounter challenges in terms of scalability due to the exponential growth of paths. ANet addresses this issue by incorporating a priority function, inspired by the A\\* algorithm for shortest-path problems, which enables the selection of crucial nodes and edges during each iteration. This novel approach effectively reduces the time and memory requirements for both training and inference processes.", "questions": "Q1: The priority function in A\\*Net is similar to the attention used in RED-GNN except that A\\*Net selects the nodes and edges according to the attention score. In the case where memory allows, how does the performance of A\\* Net change when Top operation is disabled in Algorithm 1 (line 5 & line 7)?\n\nQ2: If some nodes and edges are discarded in the early phase of model training, it may introduce incorrect inductive biases and prevent the model from training effectively. How do you address this issue to avoid such problems or why is this not an issue in A\\*Net?", "limitations": "See **Weaknesses** and **Questions**.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "S1: This paper proposes an efficient GNN called A\\*Net for link prediction with good scalability.\n\nS2: A\\*Net shows impressive results on various KGs.", "weaknesses": "W1: Although the method proposed in this article has better scalability, the contributions from theoretical perspectives are incremental compared to NBFNet.\n\nW2: The introduction of the parameter sharing between the priority function and predictor is somewhat unclear, and the reason why the reasoning task can be regarded as weak supervision for the priority function is not well explained.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your comments. Here is our response to your concerns.\n\n**W1: Although the method proposed in this article has better scalability, the contributions from theoretical perspectives are incremental compared to NBFNet.**\n\nA1: The major contribution of this paper is scalability, not the theoretical insights. Previous path-based methods like NBFNet can only deal with graphs of tens of thousands of nodes. Our A\\*Net scales path-based methods to ogbl-wikikg2, a dataset containing 2.5 million entities and 16 million triplets, which is **2 magnitudes larger than what NBFNet can solve**. Note that ogbl-wikikg2 has been previously dominant by embedding methods, and none of the embedding methods are inductive or interpretable like A\\*Net. Therefore, we believe this contribution is very important to the knowledge graph community and may potentially change future research directions as new abilities may emerge with scalability.\n\nThe theoretical design insights from the A\\* algorithm are a minor contribution compared to our empirical achievements. However, we should emphasize that the design insights led us to develop A\\*Net without significant performance loss. As showed in Tab. 6(a), naive solutions like personalized PageRank or node degree sacrifice the performance when they prune the paths.\n\n**W2: The introduction of the parameter sharing between the priority function and predictor is somewhat unclear, and the reason why the reasoning task can be regarded as weak supervision for the priority function is not well explained.**\n\nA2: The priority function is $s_{uq}^{(t)}(x) = \\sigma(f(\\mathbf{s}_{uq}^{(t)}(x))$\n\nwhere $\\mathbf{s}_{uq}^{(t)}(x)$ is the representation computed by Eqn. 10 and $f(\\cdot)$ is a feed-forward network.\n\nThe predictor function is $p(v|u,q) = \\sigma(f'(\\mathbf{s}_{uq}^{(T)}(v)))$\n\nwhere $\\mathbf{s}_{uq}^{(T)}(v)$ is the representation from the last layer and $f'(\\cdot)$ is a feed-forward network. Since two functions have similar formulas, we share the parameters between $f(\\cdot)$ and $f'(\\cdot)$.\n\nFor the priority function, an ideal supervision should assign high scores to nodes on the important paths, and low scores to nodes that are not on the important paths. We notice that the reasoning task assigns labels of 1 for true answer nodes, and labels of 0 for false answer nodes. Since true answer nodes are always on the important paths, and false answer nodes are less likely to be on the important paths, the supervision from the reasoning task is correlated to the supervision we need for the priority function. Therefore, we share the weights between these two functions, and hope that a well-trained predictor can help the priority function to converge when the whole model is trained end-to-end. We illustrated this idea in Line 188-191 in the paper.\n\n**Q1: In the case where memory allows, how does the performance of A\\* Net change when the topk operation is disabled in Algorithm 1?**\n\nA3: Our ablation study (Fig. 6) suggests that the performance of A\\*Net hits a plateau when we select more nodes or edges than a certain threshold. We run experiments with the topk operations disabled, which makes A\\*Net almost identical to NBFNet. Here are the results. Generally, the performance doesn’t increase and is upper bounded by NBFNet.\n\n|FB15k-237|MRR|H@1|H@3|H@10|\n|---|---|---|---|---|\n|A\\*Net (10% nodes, 10% edges)|0.411|0.321|0.453|0.586|\n|A\\*Net (no pruning)|0.407|0.314|0.445|0.590|\n|NBFNet|0.415|0.321|0.454|0.599|\n\n**Q2: If some nodes and edges are discarded in the early phase of model training, it may introduce incorrect inductive biases and prevent the model from training effectively. How do you address this issue to avoid such problems or why is this not an issue in A\\*Net?**\n\nA4: That’s a valid concern. Empirically, this is not an issue for A\\*Net, as A\\*Net achieves competitive performance compared to its unpruned version, NBFNet. We conjecture the reason is that A\\*Net converges in a curriculum learning fashion with the help of the weight sharing trick. Since every edge in the training graph is a training sample, we automatically have a curriculum for training path-based methods. That is, the training samples cover consecutive distances from 1 to some finite value. As long as we start from some reasonable ratios of nodes and edges, the predictor function can converge on answers of distance 1. Through the weight sharing trick, this helps the priority function to find important nodes of distance 1, and reach answers of distance 2. Then the predictor function converges on answers of distance 2. Eventually, this procedure enables A\\*Net to converge on answers of all distances."}]

Author rebuttal:
# Summary of Responses

We would like to thank all reviewers for your time and patience on our submission. Here is a summary of reviewers’ points and our responses. **We attached a PDF file to illustrate the correspondance between A\*Net and the A\* algorithm.**

**Contributions**

- **Scalability is an important problem for path-based methods (all reviewers).**
- **A\*Net achieves impressive results in performance, time and memory, especially on a million-scale dataset ogbl-wikikg2 (all reviewers).**
- **The technical contribution of A\*Net is incremental (Reviewer 7LW7, ALvG):** The technical contribution involves mathematical derivation of A\*Net and the design of neural priority function inspired by the A\* algorithm. However, we emphasize the main contribution of this paper is to scale path-based methods to ogbl-wikikg2, which is **2 magnitudes larger** than datasets solved by previous path-based methods. This contribution is recognized by all reviewers. Considering that **ogbl-wikikg2 has been previously dominated by embedding methods**, this is an important breakthrough and may potentially change future research directions.

**Writing**

- **The neural priority function contains no information about the goal and is different from the A\* algorithm (Reviewer rgXX, 7LW7):** We design the representation for the neural priority function $\mathbf{s}_{uq}^{(t)}(x)$ in A\*Net (Eqn. 10) to match the priority function in the A\* algorithm (Eqn. 4). The current aggregation of paths $\mathbf{h}_q^{(t)}(u, x)$ corresponds to the current length $d(u, x)$ in the A\* algorithm, while the term $\mathbf{g}([\mathbf{h}_q^{(t)}(u,x), \mathbf{q}])$ corresponds to the remaining distance $g(x, v)$. While the A\* algorithm uses the target node $v$ as **the absolute goal**, A\*Net uses the learned vector $\mathbf{q}$ to represent **the relative goal** from the source node $u$ to the target node $v$. A more intuitive figure can be found in the attached PDF file.
- **The weight sharing between the priority function and the predictor is not clear. (Reviewer 7LW7, ALvG):** The priority function is $s_{uq}^{(t)}(x) = \sigma(f(\mathbf{s}_{uq}^{(t)}(x))$,

where $\mathbf{s}_{uq}^{(t)}(x)$ is the representation computed by Eqn. 10 and $f(\cdot)$ is a feed-forward network.

The predictor function is $p(v|u,q) = \sigma(f'(\mathbf{s}_{uq}^{(T)}(v)))$,

where $\mathbf{s}_{uq}^{(T)}(v)$ is the representation from the last layer and $f'(\cdot)$ is a feed-forward network. We share the parameters between $f(\cdot)$ and $f'(\cdot)$.

Note that $f(\cdot)$ is a query-independent function, but the node priority $s_{uq}^{(t)}(x)$ can be query dependent since the input representation $\mathbf{s}_{uq}^{(t)}(x)$ is dependent on the query relation $q$.
- We will carefully improve our writing to address these concerns in the revised version. Since these concerns are related to our technical contribution, please let us know if you have further questions.

**Experiments**

- **Comparison with a random pruning strategy (Reviewer 7LW7).**
- **Comparison on path visualization of A\*Net, NBFNet, and handcrafted priority functions (Reviewer 7LW7).**
- **Performance of NBFNet on tail prediction (Reviewer 2kB4).**
- **Performance of A\*Net when pruning is disabled (Reviewer ALvG).**
- We provide results to all the experiments required by the reviewers. All these experiments are consistent with the observations and claims in the paper. We will include them in the revised version to provide a better context to the readers.

**Questions**

- **If the path length is small, can vanilla path-based methods be more efficient than Bellman-Ford algorithm (Reviewer 2kB4):** We compute the number of paths for different lengths, averaged over all positive triplets in each dataset. We observe an exponential growth in the number of paths w.r.t. the length of the paths. Vanilla path-based methods are only efficient for length ≤ 1 on FB15k-237, length ≤ 4 on WN18RR and length ≤ 2 on ogbl-wikikg2 respectively. This is shorter than the optimal length of 6 reported by NBFNet and RED-GNN.

---

# Path Visualization of Different Methods

1. $(\text{Bandai}, \text{industry}, \text{?})$
- NBFNet
    - $\text{Bandai} \xleftarrow{\text{subsidiary}} \text{Bandai Namco Holdings} \xrightarrow{\text{webpage}} \text{official website}\xleftarrow{\text{webpage}}\text{Bandai Namco Entertainment}\xrightarrow{\text{industry}}\text{video game}$
    - $\text{Bandai}\xrightarrow{\text{webpage}}\text{official website}\xleftarrow{\text{webpage}} \text{Santa Clara}\xleftarrow{\text{location}} \text{Bandai Namco Entertainment}\xrightarrow{\text{industry}} \text{video game}$
- A\*Net (neural)
    - $\text{Bandai} \xleftarrow{\text{subsidiary}}\text{Bandai Namco Entertainment} \xrightarrow{\text{industry}}\text{video game}$
    - $\text{Bandai} \xrightarrow{\text{industry}}\text{media}\xleftarrow{\text{industry}}\text{Pony Canyon}\xrightarrow{\text{industry}}\text{video game}$
- A\*Net (PPR)
    - $\text{Bandai}\xrightarrow{\text{webpage}}\text{official website}\xleftarrow{\text{webpage}}\text{Bandai}\xrightarrow{\text{webpage}}\text{official website}\xleftarrow{\text{webpage}}\text{Bandai}\xleftarrow{\text{subsidiary}}\text{Bandai Namco Entertainment}\xrightarrow{\text{industry}}\text{video game}$
    - $\text{Bandai}\xrightarrow{\text{webpage}}\text{official website}\xleftarrow{\text{webpage}}\text{Bandai}\xleftarrow{\text{subsidiary}}\text{Bandai Namco Entertainment}\xrightarrow{\text{industry}}\text{video game}$
- A\*Net (degree)
    - $\text{Bandai}\xrightarrow{\text{webpage}}\text{official website}\xleftarrow{\text{webpage}}\text{Def Jam Recordings}\xrightarrow{\text{industry}}\text{video game}$
    - $\text{Bandai}\xleftarrow{\text{lead}}\text{CEO}\xrightarrow{\text{company}}\text{Microsoft}\xrightarrow{\text{industry}}\text{video game}$

Meta-review: This paper presents A*Net, a scalable path-based method for knowledge graph reasoning. A*Net is inspired by the A* algorithm for shortest path problems and learns a priority function to select important nodes and edges at each iteration, which aims to reduce time and memory consumption for both training and inference. The reviewers and authors had a few rounds of discussion during the author response period, and some reviewers also followed up with the AC afterwards. The AC also read the paper, reviews and discussions. Most concerns were resolved. The reviewers and AC applauded the paper for studying the scalability issue of path-based reasoning approaches as well as the proposed approach being supported by extensive experiments. The remaining concerns are mainly about the analogy between A* and A*Net, where the comparison to A* seems to make the paper's technical contributions/novelty less clear.  In addition, reviewers are skeptical of the response that A*Net can be faster than traditional KGE models at inference time as well as Figure 1 (lack of the implementation details to fairly evaluate the training time). The authors are strongly encouraged to include the newly added experiments and clarify the aspects discussed in the response period in the revised version.
Decision: Accept (poster)