Neural Bellman-Ford Networks: A General Graph Neural Network Framework for Link Prediction

We consider the link prediction problem on both knowledge graphs and homogeneous graphs. A knowledge graph is denoted by ${\mathcal{G}}=({\mathcal{V}},{\mathcal{E}},{\mathcal{R}})$, where ${\mathcal{V}}$ and ${\mathcal{E}}$ represent the set of entities (nodes) and relations (edges) respectively, and ${\mathcal{R}}$ is the set of relation types. We use ${\mathcal{N}}(u)$ to denote the set of nodes connected to $u$, and ${\mathcal{E}}(u)$ to denote the set of edges ending with node $u$. A homogeneous graph ${\mathcal{G}}=({\mathcal{V}},{\mathcal{E}})$ can be viewed as a special case of knowledge graphs, with only one relation type for all edges. Throughout this paper, we use bold terms, ${\bm{w}}_{q}(e)$ or ${\bm{h}}_{q}(u,v)$, to denote vector representations, and italic terms, $w_{e}$ or $w_{uv}$, to denote scalars like the weight of edge $(u,v)$ in homogeneous graphs or triplet $(u,r,v)$ in knowledge graphs. Without loss of generality, we derive our method based on knowledge graphs, while our method can also be applied to homogeneous graphs.

Path Formulation. Link prediction is aimed at predicting the existence of a query relation $q$ between a head entity $u$ and a tail entity $v$. From a representation learning perspective, this requires to learn a pair representation ${\bm{h}}_{q}(u,v)$, which captures the local subgraph structure between $u$ and $v$ w.r.t. the query relation $q$. In traditional methods, such a local structure is encoded by counting different types of random walks from $u$ to $v$ [35, 15]. Inspired by this construction, we formulate the pair representation as a generalized sum of path representations between $u$ and $v$ with a commutative summation operator $\oplus$. Each path representation ${\bm{h}}_{q}(P)$ is defined as a generalized product of the edge representations in the path with the multiplication operator $\otimes$.

where ${\mathcal{P}}_{uv}$ denotes the set of paths from $u$ to $v$ and ${\bm{w}}_{q}(e_{i})$ is the representation of edge $e_{i}$. Note the multiplication operator $\otimes$ is not required to be commutative (e.g., matrix multiplication), therefore we define $\bigotimes$ to compute the product following the exact order. Intuitively, the path formulation can be interpreted as a depth-first-search (DFS) algorithm, where one searches all possible paths from $u$ to $v$, computes their representations (Equation 2) and aggregates the results (Equation 1). Such a formulation is capable of modeling several traditional link prediction methods, as well as graph theory algorithms. Formally, Theorem 1-5 state the corresponding path formulations for 3 link prediction methods and 2 graph theory algorithms respectively. See Appendix A for proofs.

Generalized Bellman-Ford Algorithm. While the above formulation is able to model important heuristics for link prediction, it is computationally expensive since the number of paths grows exponentially with the path length. Previous works [40, 11, 62] that directly computes the exponential number of paths can only afford a maximal path length of 3. A more scalable solution is to use the generalized Bellman-Ford algorithm [4]. Specifically, assuming the operators $\langle\oplus,\otimes\rangle$ satisfy a semiring system [21] with summation identity $\text{\scriptsize{0}⃝}_{q}$ and multiplication identity $\text{\scriptsize{1}⃝}_{q}$, we have the following algorithm.

where $\text{1}_{q}(u=v)$ is the indicator function that outputs $\text{\scriptsize{1}⃝}_{q}$ if $u=v$ and $\text{\scriptsize{0}⃝}_{q}$ otherwise. ${\bm{w}}_{q}(x,r,v)$ is the representation for edge $e=(x,r,v)$ and $r$ is the relation type of the edge. Equation 3 is known as the boundary condition, while Equation 4 is known as the Bellman-Ford iteration. The high-level idea of the generalized Bellman-Ford algorithm is to compute the pair representation ${\bm{h}}_{q}(u,v)$ for a given entity $u$, a given query relation $q$ and all $v\in{\mathcal{V}}$ in parallel, and reduce the total computation by the distributive property of multiplication over summation. Since $u$ and $q$ are fixed in the generalized Bellman-Ford algorithm, we may abbreviate ${\bm{h}}_{q}^{(t)}(u,v)$ as ${\bm{h}}^{(t)}_{v}$ when the context is clear. When $\oplus=min$ and $\otimes=+$, it recovers the original Bellman-Ford algorithm for the shortest path problem [5]. See Appendix B for preliminaries and the proof of the above algorithm.

While the generalized Bellman-Ford algorithm can solve many classical methods (Theorem 6), these methods instantiate the path formulation with handcrafted operators (Table 2), and may not be optimal for link prediction. To improve the capacity of path formulation, we propose a general framework, Neural Bellman-Ford Networks (NBFNet), to learn the operators in the pair representations.

Neural Parameterization. We relax the semiring assumption and parameterize the generalized Bellman-Ford algorithm (Equation 3 and 4) with 3 neural functions, namely Indicator, Message and Aggregate functions. The Indicator function replaces the indicator function $\text{1}_{q}(u=v)$. The Message function replaces the binary multiplication operator $\otimes$. The Aggregate function is a permutation invariant function over sets that replaces the n-ary summation operator $\bigoplus$. Note that one may alternatively define Aggregate as the commutative binary operator $\oplus$ and apply it to a sequence of messages. However, this will make the parameterization more complicated.

Now consider the generalized Bellman-Ford algorithm for a given entity $u$ and relation $q$. In this context, we abbreviate ${\bm{h}}_{q}^{(t)}(u,v)$ as ${\bm{h}}^{(t)}_{v}$, i.e., a representation on entity $v$ in the $t$-th iteration. It should be stressed that ${\bm{h}}^{(t)}_{v}$ is still a pair representation, rather than a node representation. By substituting the neural functions into Equation 3 and 4, we get our Neural Bellman-Ford Networks.

NBFNet can be interpreted as a novel GNN framework for learning pair representations. Compared to common GNN frameworks [32, 48] that compute the pair representation as two independent node representations ${\bm{h}}_{q}(u)$ and ${\bm{h}}_{q}(v)$, NBFNet initializes a representation on the source node $u$, and readouts the pair representation on the target node $v$. Intuitively, our framework can be viewed as a source-specific message passing process, where every node learns a representation conditioned on the source node. The pseudo code of NBFNet is outlined in Algorithm 1.

Design Space. Now we discuss some principled designs for Message, Aggregate and Indicator functions by drawing insights from traditional methods. Note the potential design space for NBFNet is way larger than what is presented here, as one can always borrow Message and Aggregate from the arsenal of message-passing GNNs [19, 16, 60, 65].

For the Message function, traditional methods instantiate it as natural summation, natural multiplication or min over scalars. Therefore, we may use the vectorized version of summation or multiplication. Intuitively, summation of ${\bm{h}}^{(t-1)}_{x}$ and ${\bm{w}}_{q}(x,r,v)$ can be interpreted as a translation of ${\bm{h}}^{(t-1)}_{x}$ by ${\bm{w}}_{q}(x,r,v)$ in the pair representation space, while multiplication corresponds to scaling. Such transformations correspond to the relational operators [18, 45] in knowledge graph embeddings [6, 68, 58, 31, 52]. For example, translation and scaling are the relational operators used in TransE [6] and DistMult [68] respectively. We also consider the rotation operator in RotatE [52].

The Aggregate function is instantiated as natural summation, max or min in traditional methods, which are reminiscent of set aggregation functions [71, 65, 9] used in GNNs. Therefore, we specify the Aggregate function to be sum, mean, or max, followed by a linear transformation and a non-linear activation. We also consider the principal neighborhood aggregation (PNA) proposed in a recent work [9], which jointly learns the types and scales of the aggregation function.

The Indicator function is aimed at providing a non-trivial representation for the source node $u$ as the boundary condition. Therefore, we learn a query embedding ${\bm{q}}$ for $\text{\scriptsize{1}⃝}_{q}$ and define Indicator function as $\text{1}(u=v)*{\bm{q}}$. Note it is also possible to additionally learn an embedding for $\text{\scriptsize{0}⃝}_{q}$. However, we find a single query embedding works better in practice.

The edge representations are instantiated as transition probabilities or length in traditional methods. We notice that an edge may have different contribution in answering different query relations. Therefore, we parameterize the edge representations as a linear function over the query relation, i.e., ${\bm{w}}_{q}(x,r,v)={\bm{W}}_{r}{\bm{q}}+{\bm{b}}_{r}$. For homogeneous graphs or knowledge graphs with very few relations, we simplify the parameterization to ${\bm{w}}_{q}(x,r,v)={\bm{b}}_{r}$ to prevent overfitting. Note that one may also parameterize ${\bm{w}}_{q}(x,r,v)$ with learnable entity embeddings ${\bm{x}}$ and ${\bm{v}}$, but such a parameterization cannot solve the inductive setting. Similar to NeuralLP [69] & DRUM [46], we use different edge representations for different iterations, which is able to distinguish noncommutative edges in paths, e.g., father’s mother v.s. mother’s father.

Link Prediction. We now show how to apply the learned pair representations ${\bm{h}}_{q}(u,v)$ to the link prediction problem. We predict the conditional likelihood of the tail entity $v$ as $p(v|u,q)=\sigma(f({\bm{h}}_{q}(u,v)))$, where $\sigma(\cdot)$ is the sigmoid function and $f(\cdot)$ is a feed-forward neural network. The conditional likelihood of the head entity $u$ can be predicted by $p(u|v,q^{-1})=\sigma(f({\bm{h}}_{q^{-1}}(v,u)))$ with the same model. Following previous works [6, 52], we minimize the negative log-likelihood of positive and negative triplets (Equation 7). The negative samples are generated according to Partial Completeness Assumption (PCA) [14], which corrupts one of the entities in a positive triplet to create a negative sample. For undirected graphs, we symmetrize the representations and define $p_{q}(u,v)=\sigma(f({\bm{h}}_{q}(u,v)+{\bm{h}}_{q}(v,u)))$. Equation 8 shows the loss for homogeneous graphs.

where $n$ is the number of negative samples per positive sample and $(u_{i}^{\prime},q,v_{i}^{\prime})$ and $(u_{i}^{\prime},v_{i}^{\prime})$ are the $i$-th negative samples for knowledge graphs and homogeneous graphs, respectively.

Time Complexity. One advantage of NBFNet is that it has a relatively low time complexity during inference⁴⁴4Although the same analysis can be applied to training on a fixed number of samples, we note it is less instructive since one can trade-off samples for performance, and the trade-off varies from method to method.. Consider a scenario where a model is required to infer the conditional likelihood of all possible triplets $p(v|u,q)$. We group triplets with the same condition $u,q$ together, where each group contains $|{\mathcal{V}}|$ triplets. For each group, we only need to execute Algorithm 1 once to get their predictions. Since a small constant number of iterations $T$ is enough for NBFNet to converge (Table 6(b)), Algorithm 1 has a time complexity of $O(|{\mathcal{E}}|d+|{\mathcal{V}}|d^{2})$, where $d$ is the dimension of representations. Therefore, the amortized time complexity for a single triplet is $O\left(\frac{|{\mathcal{E}}|d}{|{\mathcal{V}}|}+d^{2}\right)$. For a detailed derivation of time complexity of other GNN frameworks, please refer to Appendix C.

We evaluate NBFNet in three settings, knowledge graph completion, homogeneous graph link prediction and inductive relation prediction. The former two are transductive settings, while the last is an inductive setting. For knowledge graphs, we use FB15k-237 [56] and WN18RR [13]. We use the standard transductive splits [56, 13] and inductive splits [55] of these datasets. For homogeneous graphs, we use Cora, Citeseer and PubMed [49]. Following previous works [32, 12], we split the edges into train/valid/test with a ratio of 85:5:10. Statistics of datasets can be found in Appendix E. Additional experiments of NBFNet on OGB [25] datasets can be found in Appendix G.

Implementation Details. Our implementation generally follows the open source codebases of knowledge graph completion⁵⁵5https://github.com/DeepGraphLearning/KnowledgeGraphEmbedding. MIT license. and homogeneous graph link prediction⁶⁶6https://github.com/tkipf/gae. MIT license.. For knowledge graphs, we follow [69, 46] and augment each triplet $\langle$u, q, v$\rangle$ with a flipped triplet $\langle$v, q^-1, u$\rangle$. For homogeneous graphs, we follow [33, 32] and augment each node $u$ with a self loop $\langle$u, u$\rangle$. We instantiate NBFNet with 6 layers, each with 32 hidden units. The feed-forward network $f(\cdot)$ is set to a 2-layer MLP with 64 hidden units. ReLU is used as the activation function for all hidden layers. We drop out edges that directly connect query node pairs during training to encourage the model to capture longer paths and prevent overfitting. Our model is trained on 4 Tesla V100 GPUs for 20 epochs. We select the models based on their performance on the validation set. See Appendix F for more details.

Evaluation. We follow the filtered ranking protocol [6] for knowledge graph completion. For a test triplet $\langle$u, q, v$\rangle$, we rank it against all negative triplets $\langle$u, q, v’$\rangle$ or $\langle$u’, q, v$\rangle$ that do not appear in the knowledge graph. We report mean rank (MR), mean reciprocal rank (MRR) and HITS at N (H@N) for knowledge graph completion. For inductive relation prediction, we follow [55] and draw 50 negative triplets for each positive triplet and use the above filtered ranking. We report HITS@10 for inductive relation prediction. For homogeneous graph link prediction, we follow [32] and compare the positive edges against the same number of negative edges. We report area under the receiver operating characteristic curve (AUROC) and average precision (AP) for homogeneous graphs.

Baselines. We compare NBFNet against path-based methods, embedding methods, and GNNs. These include 11 baselines for knowledge graph completion, 10 baselines for homogeneous graph link prediction and 4 baselines for inductive relation prediction. Note the inductive setting only includes path-based methods and GNNs, since existing embedding methods cannot handle this setting.

Table 3 summarizes the results on knowledge graph completion. NBFNet significantly outperforms existing methods on all metrics and both datasets. NBFNet achieves an average relative gain of 21% in HITS@1 compared to the best path-based method, DRUM [46], on two datasets. Since DRUM is a special instance of NBFNet with natural summation and multiplication operators, this indicates the importance of learning Message and Aggregate functions in NBFNet. NBFNet also outperforms the best embedding method, LowFER [1], with an average relative performance gain of 18% in HITS@1 on two datasets. Meanwhile, NBFNet requires much less parameters than embedding methods. NBFNet only uses 3M parameters on FB15k-237, while TransE needs 30M parameters. See Appendix D for details on the number of parameters.

Table 4 shows the results on homogeneous graph link prediction. NBFNet gets the best results on Cora and PubMed, meanwhile achieves competitive results on CiteSeer. Note CiteSeer is extremely sparse (Appendix E), which makes it hard to learn good representations with NBFNet. One thing to note here is that unlike other GNN methods, NBFNet does not use the node features provided by the datasets but is still able to outperform most other methods. We leave how to effectively combine node features and structural representations for link prediction as our future work.

Table 5 summarizes the results on inductive relation prediction. On all inductive splits of two datasets, NBFNet achieves the best result. NBFNet outperforms the previous best method, GraIL [55], with an average relative performance gain of 22% in HITS@10. Note that GraIL explicitly encodes the local subgraph surrounding each node pair and has a high time complexity (Appendix C). Usually, GraIL can at most encode a 2-hop subgraph, while our NBFNet can efficiently explore longer paths.

Message & Aggregate Functions. Table 6(a) shows the results of different Message and Aggregate functions. Generally, NBFNet benefits from advanced embedding methods (DistMult, RotatE > TransE) and aggregation functions (PNA > sum, mean, max). Among simple Aggregate functions (sum, mean, max), combinations of Message and Aggregate functions (TransE & max, DistMult & sum) that satisfy the semiring assumption⁷⁷7Here semiring is discussed under the assumption of linear activation functions. Rigorously, no combination satisfies a semiring if we consider non-linearity in the model. of the generalized Bellman-Ford algorithm, achieve locally optimal performance. PNA significantly improves over simple counterparts, which highlights the importance of learning more powerful Aggregate functions.

Number of GNN Layers. Table 6(b) compares the results of NBFNet with different number of layers. Although it has been reported that GNNs with deep layers often result in significant performance drop [36, 77], we observe NBFNet does not have this issue. The performance increases monotonically with more layers, hitting a saturation after 6 layers. We conjecture the reason is that longer paths have negligible contribution, and paths not longer than 6 are enough for link prediction.

Performance by Relation Category. We break down the performance of NBFNet by the categories of query relations: one-to-one, one-to-many, many-to-one and many-to-many⁸⁸8The categories are defined same as [63]. We compute the average number of tails per head and the average number of heads per tail. The category is one if the average number is smaller than 1.5 and many otherwise.. Table 6(c) shows the prediction results for each category. It is observed that NBFNet not only improves on easy one-to-one cases, but also on hard cases where there are multiple true answers for the query.

One advantage of NBFNet is that we can interpret its predictions through paths, which may be important for users to understand and debug the model. Intuitively, the interpretations should contain paths that contribute most to the prediction $p(u,q,v)$. Following local interpretation methods [3, 72], we approximate the local landscape of NBFNet with a linear model over the set of all paths, i.e., 1st-order Taylor polynomial. We define the importance of a path as its weight in the linear model, which can be computed by the partial derivative of the prediction w.r.t. the path. Formally, the top-k path interpretations for $p(u,q,v)$ are defined as

Note this formulation generalizes the definition of logical rules [69, 46] to non-linear models. While directly computing the importance of all paths is intractable, we approximate them with edge importance. Specifically, the importance of each path is approximated by the sum of the importance of edges in that path, where edge importance is obtained via auto differentiation. Then the top-k path interpretations are equivalent to the top-k longest paths on the edge importance graph, which can be solved by a Bellman-Ford-style beam search. Better approximation is left as a future work.

Table 7 visualizes path interpretations from FB15k-237 test set. While users may have different insights towards the visualization, here is our understanding. 1) In the first example, NBFNet learns soft logical entailment, such as $\emph{impersonate}^{-1}\land\emph{nationality}\implies\emph{nationality}$ and $\emph{ethnicity}^{-1}\land\emph{distribution}\implies\emph{nationality}$. 2) In second example, NBFNet performs analogical reasoning by leveraging the fact that Florence is similar to Rome. 3) In the last example, NBFNet extracts longer paths, since there is no obvious connection between Pearl Harbor (film) and Japanese language.