Neural Link Prediction with Walk Pooling

1 Introduction . Graphs are a natural model for relational data such as coauthorship networks or the human protein interactome . Since real-world graphs are often only partially observed , a central problem across all scientific domains is to predict missing links ( Liben-Nowell & Kleinberg , 2007 ) . Link prediction finds applications in predicting protein interactions ( Qi et al. , 2006 ) , drug responses ( Stanfield et al. , 2017 ) , or completing the knowledge graph ( Nickel et al. , 2015 ) . It underpins recommender systems in social networks ( Adamic & Adar , 2003 ) and online marketplaces ( Lü et al. , 2012 ) . Successful link prediction demands an understanding of the principles behind graph formation . Observed graph Candidate links Subgraph Prefer triangles with large weights Prefer quadrangles with large weights Completed graph see Figure 1 for an illustration . But organizing rules are not universal across graphs and 1Mathematical topology studies ( global ) properties of shapes that are preserved under homeomorphisms . Our use of “ topology ” to refer to local patterns is common in the network literature . they need to be learned . The centerpiece of our algorithm is a new random-walk-based pooling mechanism called WalkPool which may be interpreted as a learnable version of topological heuristics . Our algorithm is universal in the sense of Chien et al . ( 2020 ) in that it models both heterophilic ( e.g. , 2D-grid in Figure 1 ) and homophilic ( e.g. , Triangle lattice in Figure 1 ) graphs unlike hardcoded heuristics or standard GNNs which often work better on homophilic graphs . WalkPool can learn the topological organizing rules in real graphs ( for further discussion of lattices from Figure 1 see also Section 4.4 ) , and achieves state-ofthe-art link prediction results on all common benchmark datasets , sometimes by a significant margin , even on datasets like USAir ( Batagelj & Mrvar , 2006 ) where the previous state of the art is as high as 97 % . 1.1 Related work . Early studies on link prediction use heuristics from social network analysis ( Lü & Zhou , 2011 ) . The homophily mechanism for example ( McPherson et al. , 2001 ) asserts that “ similar ” nodes connect . Most heuristics are based on connectivity : the common-neighbor index ( CN ) scores a pair of nodes by the number of shared neighbors , yielding predicted graphs with many triangles ; Adamić–Adar index ( AA ) assumes that a highly-connected neighbor contributes less to the score of a focal link ( Adamic & Adar , 2003 ) . CN and AA rely on paths of length two . Others heuristics use longer paths , explicitly accounting for long-range correlations ( Lü & Zhou , 2011 ) . The Katz index ( Katz , 1953 ) scores a pair of nodes by a weighted sum of the number of walks of all possible lengths . PageRank scores a link by the probability that a walker starting at one end of the link reaches the other in the stationary state under a random walk model with restarts . Heuristics make strong assumptions such as a particular parametric decay of path weights with length , often tailored to either homophilic or heterophilic graphs . Moreover , they can not be easily used on graphs with node or edge features . Instead of hard-coded structural features , link prediction algorithms based on GNNs like the VGAE ( Kipf & Welling , 2016b ) or SEAL Zhang & Chen ( 2018 ) use learned node-level representations that capture both the structure and the features , yielding unprecedented accuracy . Two strategies have proved successful for GNN-based link prediction . The first is to devise a score function which only depends on the two nodes that define a link . The second is to extract a subgraph around the focal link and solve link prediction by ( sub ) graph classification ( Zhang & Chen , 2017 ) . In both cases the node representations are either learned in an unsupervised way ( Kipf & Welling , 2016b ; Mavromatis & Karypis , 2020 ) or computed by a GNN trained joinly with the link classifier ( Zhang & Chen , 2018 ; Zhang et al. , 2020 ) . Algorithms based on two-body score functions achieved state-of-the-art performance when they appeared , but real networks brim with many-body correlations . Graphs in biochemistry , neurobiology , ecology , and engineering ( Milo et al. , 2004 ) exhibit motifs—distinct patterns occuring much more often than in random graphs . Motifs may correspond to basic computational elements and relate to function ( Shen-Orr et al. , 2002 ) . A multilayer GNN generates node representation by aggregating information from neighbors , capturing to some extent many-body and long-range correlations . This dependence is however indirect and complex topological properties such as frequency of motifs are smoothed out . The current state-of-the-art link prediction algorithm , SEAL ( Zhang & Chen , 2018 ) , is based on subgraph classification and thus may account for higher-order correlations . Unlike in vanilla graph classification where a priori all links are equaly important , in graph classification for link prediction the focal link plays a special role and the relative placement of other links with respect to it matters . SEAL takes this into account by labeling each node of the subgraph by its distance to the focal link . This endows the GNN with first-order topological information and helps distinguish relevant nodes in the pooled representation . Zhang and Chen show that node labeling is essential for SEAL ’ s exceptional accuracy . Nevertheless , important structural motifs are represented indirectly ( if at all ) by such labeling schemes , even less so after average pooling . We hypothesize that a major bottleneck in link prediction comes from suboptimal pooling which fails to account for topology . 1.2 Our contribution . Following Zhang & Chen ( 2018 ) , we approach link prediction from via subgraph classification . Instead of encoding relative topological information through node labels , we design a new pooling mechanism called WalkPool . WalkPool extracts higher-order structural information by encoding node representations and graph topology into random-walk transition probabilities on some effective latent graph , and then using those probabilities to compute features we call walk profiles . WalkPool generalizes loop-counting ideas used to build expressive graph models ( Pan et al. , 2016 ) . Computing expressive topological descriptors which are simultaneously sensitive to node features is the key difference between WalkPool and SEAL . Using normalized probabilities tuned by graph attention mitigates the well-known issue in graph learning that highly-connected nodes may bias predictions . Transition probabilities and the derived walk profiles have simple topological interpretations . WalkPool can be applied to node representations generated by any unsupervised graph representation models or combined with GNNs and trained in end-to-end . It achieves stateof-the-art results on all common link prediction benchmarks . Our code and data are available online at https : //anonymous.4open.science/r/WalkPooling_Ano . 2 Link prediction on graphs . We consider an observed graph with N nodes ( or vertices ) , Go = ( V , Eo ) , with V being the node set and Eo the set of observed links ( or edges ) . The observed link set Eo is a subset Eo ⊂ E∗ of the set of all true links E∗ . The target of link prediction is to infer missing links from a candidate set Ec , which contains both true ( in E∗ ) and false ( not in E∗ ) missing links . Problem 1 . ( Link prediction ) Design an algorithm LearnLP that takes an observed graph Go ⊂ G and produces a link predictor LearnLP ( Go ) = Π : V × V → { True , False } which accurately classifies links in Ec . Well-performing solutions to Problem 1 exploit structural and functional regularities in the observed graph to make inferences about unobserved links . Path counts and random walks For simplicity we identify the N vertices with integers 1 , . . . , N and represent G by its adjacency matrix , A = ( aij ) Ni , j=1 ∈ { 0 , 1 } N×N with aij = 1 if { i , j } ∈ Eo and aij = 0 otherwise . Nodes may have associated feature vectors ( xi ∈ RF , i ∈ { 1 , . . . , N } ) ; we collect all feature vectors in the feature matrix X = [ x1 , · · · , xN ] T ∈ RN×F . Integer powers of the adjacency matrix reveal structural information : [ Aτ ] ij is the number of paths of length τ connecting nodes i and j. WalkPool relies on random walk transition matrices . For an adjacency matrix A the transition matrix is defined as P = D−1A where D = diag ( d1 , . . . , dN ) and di = ∑ j aij = |N ( i ) | is the degree of the node i . Thus the probability pij = d −1 i aij that a random walker at node i transitions to node j is inversely proportional to the number of neighbors of i . An extension to graphs with non-negative edge weights W = ( wij ) is straightforward by replacing aij by wij . Powers of P are interpretable : [ Pτ ] ij is the probability that a random walker starting at node i will reach node j in τ hops . As we show in Section 3.2 , transition probabilities in WalkPool are determined as coefficients of an attention mechanism applied to learned node features . The node features are in turn extracted by a parameterized function fθ , which distills input attributes and to an extent graph topology into an embedding ( a vector ) for each node . The feature extractor fθ : { 0 , 1 } N×N ×RN×F → RN×F ′ takes as input the adjacency matrix ( which encodes the graph topology ) and the input feature matrix , and outputs a distilled node feature matrix . It should thus be equivariant to node ordering in A which is satisfied by GNNs . 3 WalkPool for link prediction by subgraph classification . We now describe WalkPool which directly leverages higher-order topological correlations without resorting to link labeling schemes and without making strong a priori assumptions . WalkPool first samples the k-hop subgraph Gk { i , j } ⊂ G o enclosing the target link ; and then computes random-walk profiles for sampled subgraphs with and without the target link . Random walk profiles are then fed into a link classifier . Computing walk profiles entails 1 . Feature extraction : Z = fθ ( A , X ) , where fθ is a GNN ; 2 . Transition matrix computation : P = AttentionCoefficientsθ ( Z ; G ) ; 3 . Walk profiles : Extract entries from Pτ for 2 ≤ τ ≤ τc related to the focal link . We emphasize that we do not use attention to compute per-node linear combinations of features like Veličković et al . ( 2017 ) ( which is analogous to Vaswani et al . ( 2017 ) ) , but rather interpret the attention coefficients as random walk transition probabilities . The features Z may be obtained either by an unsupervised GNN such as VGAE ( Kipf & Welling , 2016b ) , or they may be computed by a GNN which is trained jointly with WalkPool . 3.1 Sampling the enclosing subgraphs 2 34 1 1 2 1 1 4 1 1 3 1 1 3 2 sons beyond diminishing returns : a large k greatly increases memory and compute demands . The size of a 2-hop neighborhood is graph-dependent , but for the highly-connected E.coli and PB datasets we already need to enforce a maximum number of nodes per hop to fit the subgraphs in memory . Let d ( x , y ) be the shortest-path distance between nodes x and y . The k-hop enclosing subgraph Gk { i , j } for { i , j } is defined as the subgraph induced from G o by the set of nodes Vk { i , j } = { x ∈ V : d ( x , i ) ≤ k or d ( x , j ) ≤ k } . Then Gk { i , j } = ( V k { i , j } , E k { i , j } ) , where { x , y } ∈ E k i , j when x , y ∈ Vk { i , j } and { x , y } ∈ E o . We omit the dependence on k and write G { i , j } = ( V { i , j } , E { i , j } ) for simplicity . We fix an arbitrary order of the nodes in V { i , j } and denote the corresponding adjacency matrix by A { i , j } . Without loss of generality , we assume that under the chosen order of nodes , the nodes i and j are labeled as 1 and 2 so that the candidate link { i , j } is mapped to { 1 , 2 } in its enclosing subgraph . We denote the corresponding node feature matrix Z { i , j } , with values inherited from the full graph ( the rows of Z { i , j } are a subset of rows of Z ) . For the candidate set of links Ec , we construct a set of enclosing subgraphs Sc = { G { i , j } : { i , j } ∈ Ec } , thus transforming the link prediction problem into classifying these k-hop enclosing subgraphs . For training , we sample a set of known true and false edges Et and construct its corresponding enclosing subgraph set St = { Gi , j : ( i , j ) ∈ Et } .

Summary: The paper proposes a new link prediction method called WP, based on treating link prediction as a subgraph classification problem, where the link of interest is surrounded by the subgraph. Existing GNN architectures could be used for learning node embeddings. With node embeddings, a latent subgraph could be built using self-attention, where attention scores are used for computing the subgraph edge weights. With the latent subgraph, a set of features could be computed to be used as the features for the link of interest. The features will be used for subsequent link prediction. Essentially, the proposed WP method is a feature extractor for links (or its surrounding subgraph).