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

1 Introduction . Predicting the interactions between nodes ( a.k.a . link prediction ) is a fundamental task in the field of graph machine learning . Given the ubiquitous existence of graphs , such a task has many applications , such as recommender system [ 34 ] , knowledge graph completion [ 41 ] and drug repurposing [ 27 ] . Traditional methods of link prediction usually define different heuristic metrics over the paths between a pair of nodes . For example , Katz index [ 30 ] is defined as a weighted count of paths between two nodes . Personalized PageRank [ 42 ] measures the similarity of two nodes as the random walk probability from one to the other . Graph distance [ 37 ] uses the length of the shortest path between two nodes to predict their association . These methods can be directly applied to new graphs , i.e. , inductive setting , enjoy good interpretability and scale up to large graphs . However , they are designed based on handcrafted metrics and may not be optimal for link prediction on real-world graphs . 1Unless stated otherwise , we use summation and multiplication to refer the generalized operators in the path formulation , rather than the basic operations of arithmetic . 2Code is available at https : //github.com/DeepGraphLearning/NBFNet 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . To address these limitations , some link prediction methods adopt graph neural networks ( GNNs ) [ 32 , 48 , 59 ] to automatically extract important features from local neighborhoods for link prediction . Thanks to the high expressiveness of GNNs , these methods have shown state-of-the-art performance . However , these methods can only be applied to predict new links on the training graph , i.e . transductive setting , and lack interpretability . While some recent methods [ 73 , 55 ] extract features from local subgraphs with GNNs and support inductive setting , the scalability of these methods is compromised . Therefore , we wonder if there exists an approach that enjoys the advantages of both traditional path-based methods and recent approaches based on graph neural networks , i.e. , generalization in the inductive setting , interpretability , high model capacity and scalability . In this paper , we propose such a solution . Inspired by traditional path-based methods , our goal is to develop a general and flexible representation learning framework for link prediction based on the paths between two nodes . Specifically , we define the representation of a pair of nodes as the generalized sum of all the path representations between them , where each path representation is defined as the generalized product of the edge representations in the path . Many link prediction methods , such as Katz index [ 30 ] , personalized PageRank [ 42 ] , graph distance [ 37 ] , as well as graph theory algorithms like widest path [ 4 ] and most reliable path [ 4 ] , are special instances of this path formulation with different summation and multiplication operators . Motivated by the polynomial-time algorithm for the shortest path problem [ 5 ] , we show that such a formulation can be efficiently solved via the generalized Bellman-Ford algorithm [ 4 ] under mild conditions and scale up to large graphs . The operators in the generalized Bellman-Ford algorithm—summation and multiplication—are handcrafted , which have limited flexibility . Therefore , we further propose the Neural Bellman-Ford Networks ( NBFNet ) , a graph neural network framework that solves the above path formulation with learned operators in the generalized Bellman-Ford algorithm . Specifically , NBFNet parameterizes the generalized Bellman-Ford algorithm with three neural components , namely INDICATOR , MESSAGE and AGGREGATE functions . The INDICATOR function initializes a representation on each node , which is taken as the boundary condition of the generalized Bellman-Ford algorithm . The MESSAGE and the AGGREGATE functions learn the multiplication and summation operators respectively . We show that the MESSAGE function can be defined according to the relational operators in knowledge graph embeddings [ 6 , 68 , 58 , 31 , 52 ] , e.g. , as a translation in Euclidean space induced by the relational operators of TransE [ 6 ] . The AGGREGATE function can be defined as learnable set aggregation functions [ 71 , 65 , 9 ] . With such parameterization , NBFNet can generalize to the inductive setting , meanwhile achieve one of the lowest time complexity among inductive GNN methods . A comparison of NBFNet and other GNN frameworks for link prediction is showed in Table 1 . With other instantiations of MESSAGE and AGGREGATE functions , our framework can also recover some existing works on learning logic rules [ 69 , 46 ] for link prediction on knowledge graphs ( Table 2 ) . Our NBFNet framework can be applied to several link prediction variants , covering not only singlerelational graphs ( e.g. , homogeneous graphs ) but also multi-relational graphs ( e.g. , knowledge graphs ) . We empirically evaluate the proposed NBFNet for link prediction on homogeneous graphs and knowledge graphs in both transductive and inductive settings . Experimental results show that the proposed NBFNet outperforms existing state-of-the-art methods by a large margin in all settings , with an average relative performance gain of 18 % on knowledge graph completion ( HITS @ 1 ) and 22 % on inductive relation prediction ( HITS @ 10 ) . We also show that the proposed NBFNet is indeed interpretable by visualizing the top-k relevant paths for link prediction on knowledge graphs . 2 Related Work . Existing work on link prediction can be generally classified into 3 main paradigms : path-based methods , embedding methods , and graph neural networks . Path-based Methods . Early methods on homogeneous graphs compute the similarity between two nodes based on the weighted count of paths ( Katz index [ 30 ] ) , random walk probability ( personalized PageRank [ 42 ] ) or the length of the shortest path ( graph distance [ 37 ] ) . SimRank [ 28 ] uses advanced metrics such as the expected meeting distance on homogeneous graphs , which is extended by PathSim [ 51 ] to heterogeneous graphs . On knowledge graphs , Path Ranking [ 35 , 15 ] directly uses relational paths as symbolic features for prediction . Rule mining methods , such as NeuralLP [ 69 ] and DRUM [ 46 ] , learn probabilistic logical rules to weight different paths . Path representation methods , such as Path-RNN [ 40 ] and its successors [ 11 , 62 ] , encode each path with recurrent neural networks ( RNNs ) , and aggregate paths for prediction . However , these methods need to traverse an exponential number of paths and are limited to very short paths , e.g. , ≤ 3 edges . To scale up path-based methods , All-Paths [ 57 ] proposes to efficiently aggregate all paths with dynamic programming . However , All-Paths is restricted to bilinear models and has limited model capacity . Another stream of works [ 64 , 10 , 22 ] learns an agent to collect useful paths for link prediction . While these methods can produce interpretable paths , they suffer from extremely sparse rewards and require careful engineering of the reward function [ 38 ] or the search strategy [ 50 ] . Some other works [ 8 , 44 ] adopt variational inference to learn a path finder and a path reasoner for link prediction . Embedding Methods . Embedding methods learn a distributed representation for each node and edge by preserving the edge structure of the graph . Representative methods include DeepWalk [ 43 ] and LINE [ 53 ] on homogeneous graphs , and TransE [ 6 ] , DistMult [ 68 ] and RotatE [ 52 ] on knowledge graphs . Later works improve embedding methods with new score functions [ 58 , 13 , 31 , 52 , 54 , 76 ] that capture common semantic patterns of the relations , or search the score function in a general design space [ 75 ] . Embedding methods achieve promising results on link prediction , and can be scaled to very large graphs using multiple GPUs [ 78 ] . However , embedding methods do not explicitly encode local subgraphs between node pairs and can not be applied to the inductive setting . Graph Neural Networks . Graph neural networks ( GNNs ) [ 47 , 33 , 60 , 65 ] are a family of representation learning models that encode topological structures of graphs . For link prediction , the prevalent frameworks [ 32 , 48 , 12 , 59 ] adopt an auto-encoder formulation , which uses GNNs to encode node representations , and decodes edges as a function over node pairs . Such frameworks are potentially inductive if the dataset provides node features , but are transductive only when node features are unavailable . Another stream of frameworks , such as SEAL [ 73 ] and GraIL [ 55 ] , explicitly encodes the subgraph around each node pair for link prediction . While these frameworks are proved to be more powerful than the auto-encoder formulation [ 74 ] and can solve the inductive setting , they require to materialize a subgraph for each link , which is not scalable to large graphs . By contrast , our NBFNet explicitly captures the paths between two nodes for link prediction , meanwhile achieves a relatively low time complexity ( Table 1 ) . ID-GNN [ 70 ] formalizes link prediction as a conditional node classification task , and augments GNNs with the identity of the source node . While the architecture of NBFNet shares some spirits with ID-GNN , our model is motivated by the generalized Bellman-Ford algorithm and has theoretical connections with traditional path-based methods . There are also some works trying to scale up GNNs for link prediction by dynamically pruning the set of nodes in message passing [ 66 , 20 ] . These methods are complementary to NBFNet , and may be incorporated into our method to further improve scalability . 3 Methodology . In this section , we first define a path formulation for link prediction . Our path formulation generalizes several traditional methods , and can be efficiently solved by the generalized Bellman-Ford algorithm . Then we propose Neural Bellman-Ford Networks to learn the path formulation with neural functions . 3.1 Path Formulation for Link Prediction . We consider the link prediction problem on both knowledge graphs and homogeneous graphs . A 3We consider the inductive setting where a model can generalize to entirely new graphs without node features . knowledge graph is denoted by G = ( V , E , R ) , where V and E represent the set of entities ( nodes ) and relations ( edges ) respectively , andR is the set of relation types . We useN ( u ) to denote the set of nodes connected to u , and E ( u ) to denote the set of edges ending with node u . A homogeneous graph G = ( V , 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 , wq ( e ) or hq ( u , v ) , to denote vector representations , and italic terms , we or wuv , 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 hq ( 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 ⊕ . Each path representation hq ( P ) is defined as a generalized product of the edge representations in the path with the multiplication operator ⊗ . hq ( u , v ) = hq ( P1 ) ⊕ hq ( P2 ) ⊕ ... ⊕ hq ( P|Puv| ) |Pi∈Puv , ⊕ P∈Puv hq ( P ) ( 1 ) hq ( P = ( e1 , e2 , ... , e|P | ) ) = wq ( e1 ) ⊗wq ( e2 ) ⊗ ... ⊗wq ( e|P | ) , |P |⊗ i=1 wq ( ei ) ( 2 ) where Puv denotes the set of paths from u to v and wq ( ei ) is the representation of edge ei . Note the multiplication operator ⊗ is not required to be commutative ( e.g. , matrix multiplication ) , therefore we define ⊗ 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 . Theorem 1 Katz index is a path formulation with ⊕ = + , ⊗ = × and wq ( e ) = βwe . Theorem 2 Personalized PageRank is a path formulation with ⊕ = + , ⊗ = × and wq ( e ) = αwuv/ ∑ v′∈N ( u ) wuv′ . Theorem 3 Graph distance is a path formulation with ⊕ = min , ⊗ = + and wq ( e ) = we . Theorem 4 Widest path is a path formulation with ⊕ = max , ⊗ = min and wq ( e ) = we . Theorem 5 Most reliable path is a path formulation with ⊕ = max , ⊗ = × and wq ( e ) = we . 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 〈⊕ , ⊗〉 satisfy a semiring system [ 21 ] with summation identity 0©q and multiplication identity 1©q , we have the following algorithm . h ( 0 ) q ( u , v ) ← 1q ( u = v ) ( 3 ) h ( t ) q ( u , v ) ← ⊕ ( x , r , v ) ∈E ( v ) h ( t−1 ) q ( u , x ) ⊗wq ( x , r , v ) ⊕ h ( 0 ) q ( u , v ) ( 4 ) where 1q ( u = v ) is the indicator function that outputs 1©q if u = v and 0©q otherwise . wq ( 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 hq ( u , v ) for a given entity u , a given query relation q and all v ∈ 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 h ( t ) q ( u , v ) as h ( t ) v when the context is clear . When ⊕ = min and ⊗ = + , 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 . Theorem 6 Katz index , personalized PageRank , graph distance , widest path and most reliable path can be solved via the generalized Bellman-Ford algorithm .

Summary: Summary: In this study the authors propose a general framework to learn neural INDICATOR, MESSAGE and AGGREGATE functions for the generalized Bellman-Ford (BF) algorithm. They apply the neural BF to link prediction problems in multirelational knowledge graphs in the in-sample setting as well as in the out-of-sample setting, and in homogeneous graphs. On a practical level the approach trades compute with memory and shows good results on benchmark datasets.