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: This paper considers a link prediction framework called Neural Bellman-Ford Networks. The idea is to encode the relational structure between two nodes by using GNN message passing. The key here is to put the identity of the source node in the input feature. The provided framework demonstrates SoTA performance on the link prediction and KG completion datasets.