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 provides a method for link prediction based on random walk based pooling method (WalkPool). It aggregates the probabilities of the adjacent structure of a node or link to design the node embeddings. Finally, it shows superiority of the method against baselines via AUC and AP.