GRAPH NEURAL NETWORKS FOR LINK PREDICTION WITH SUBGRAPH SKETCHING

Abstract

Many Graph Neural Networks (GNNs) perform poorly compared to simple heuristics on Link Prediction (LP) tasks. This is due to limitations in expressive power such as the inability to count triangles (the backbone of most LP heuristics) and because they can not distinguish automorphic nodes (those having identical structural roles). Both expressiveness issues can be alleviated by learning link (rather than node) representations and incorporating structural features such as triangle counts. Since explicit link representations are often prohibitively expensive, recent works resorted to subgraph-based methods, which have achieved state-of-the-art performance for LP, but suffer from poor efficiency due to high levels of redundancy between subgraphs. We analyze the components of subgraph GNN (SGNN) methods for link prediction. Based on our analysis, we propose a novel full-graph GNN called ELPH (Efficient Link Prediction with Hashing) that passes subgraph sketches as messages to approximate the key components of SGNNs without explicit subgraph construction. ELPH is provably more expressive than Message Passing GNNs (MPNNs). It outperforms existing SGNN models on many standard LP benchmarks while being orders of magnitude faster. However, it shares the common GNN limitation that it is only efficient when the dataset fits in GPU memory. Accordingly, we develop a highly scalable model, called BUDDY, which uses feature precomputation to circumvent this limitation without sacrificing predictive performance. Our experiments show that BUDDY also outperforms SGNNs on standard LP benchmarks while being highly scalable and faster than ELPH.

1. INTRODUCTION

Link Prediction (LP) is an important problem in graph ML with many industrial applications. For example, recommender systems can be formulated as LP; link prediction is also a key process in drug discovery and knowledge graph construction. There are three main classes of LP methods: (i) heuristics (See Appendix C.1) that estimate the distance between two nodes (e.g. personalized page rank (PPR) (Page et al., 1999) or graph distance (Zhou et al., 2009) ) or the similarity of their neighborhoods (e.g Common Neighbors (CN), Adamic-Adar (AA) (Adamic & Adar, 2003) , or Resource Allocation (RA) (Zhou et al., 2009) ); (ii) unsupervised node embeddings or factorization methods, which encompass the majority of production recommendation systems (Koren et al., 2009; Chamberlain et al., 2020) ; and, recently, (iii) Graph Neural Networks, in particular of the Message-Passing type (MPNNs) (Gilmer et al., 2017; Kipf & Welling, 2017; Hamilton et al., 2017) .foot_0 GNNs excel in graph-and node-level tasks, but often fail to outperform node embeddings or heuristics on common LP benchmarks such as the Open Graph Benchmark (OGB) (Hu et al., 2020) . There are two related reasons why MPNNs tend to be poor link predictors. Firstly, due to the equivalence of message passing to the Weisfeiler-Leman (WL) graph isomorphism test (Xu et al., 2019; Morris et al., 2019) , standard MPNNs are provably incapable of counting trian-



GNNs are a broader category than MPNNs. Since the majority of GNNs used in practice are of the message passing type, we will use the terms synonymously.

