Revisiting Virtual Nodes in Graph Neural Networks for Link Prediction

1 INTRODUCTION . Link prediction is an important task to complete graphs that are missing edges in various domains : citation networks ( Kipf & Welling , 2016 ) , social networks ( Adamic & Adar , 2003 ) , medical drug interaction graphs ( Abbas et al. , 2021 ) , or knowledge graphs ( KGs ) ( Ji et al. , 2021 ) . Numerous kinds of models have been proposed to solve the link prediction problem , ranging from KG-specific predictors ( Ji et al. , 2021 ) to graph neural networks ( GNNs ) ( Kipf & Welling , 2016 ; Zhang & Chen , 2018 ) . Over dense biomedical networks , GNNs turned out to work especially well ( Hu et al. , 2020 ) . In this work , we focus on graph neural networks for link prediction . Many of the popular GNNs are based on the message-passing scheme , which computes node embeddings based on iteratively aggregating the features of ( usually direct/one-hop ) neighbor nodes along the graph edges ( Gilmer et al. , 2017 ) . Interestingly , best performance is usually obtained by only considering two to three hops of neighbors ( i.e. , 2-3 layers in the GNN ) . One main reason identified for this is over-smoothing , the problem that node representations become indistinguishable when the number of layers increases ( Li et al. , 2018 ) . The exponentially-growing amount of information has also been suggested as one issue connected to capturing long-range dependencies ( Alon & Yahav , 2021 ) . While it is likely that link prediction most often depends on the local node neighborhood , it is not beyond imagination that there are critical long-range dependencies ( e.g. , complex chains of drug-drug or drug-protein interactions ) . Hence , using a small number of layers to overcome the above problems results in under-reaching . There have been several recent proposals to overcome under-reaching . On the one hand , several works propose techniques that allow for larger numbers of GNN layers ( Xu et al. , 2018 ; Wu et al. , 2019 ; Liu et al. , 2020 ; Chen et al. , 2020 ; Sun et al. , 2021 ; Zhou et al. , 2020 ; Li et al. , 2020a ) . However , although ( Chen et al. , 2020 ) show that over-smoothing happens particularly in dense graphs , the link prediction experiments in these works consider citation or recommendation networks , but not the especially dense biomedical ones . And our experiments over the latter suggest that the reported results are not generalizable to the more challenging biomedical data . On the other hand , there are approaches that adapt the message-passing scheme to consider neighbors beyond the one-hop neighborhood : based on graph diffusion ( Atwood & Towsley , 2016 ; Klicpera et al. , 2019a ; Abu-El-Haija et al. , 2019 ; Xu et al. , 2019a ; Ma et al. , 2020 ; Klicpera et al. , 2019b ) and other theories ( Morris et al. , 2019 ; You et al. , 2019 ) . However , most of these models are relatively complex and , in fact , in our experiments over the challenging graphs from the Open Graph Benchmark ( OGB ) ( Hu et al. , 2020 ) , several ran out of memory . Moreover , the majority has not considered link prediction , while this problem was recently shown to be more difficult than node classification ( Zhang et al. , 2020 ) . In this paper , we propose a simple but elegant solution to under-reaching based on the concept of virtual nodes ( Gilmer et al. , 2017 ; Li et al. , 2017 ; Pham et al. , 2017 ; Ishiguro et al. , 2019 ) . Virtual nodes are well known to often improve the graph classification performance of graph neural networks , where an artificial virtual node is added to every graph and connected to all nodes in the graph . While the virtual nodes were originally thought as representations of the entire graph , they also provide shortcuts for message passing between nodes along the graph edges . Surprisingly , the impact of virtual nodes for the link prediction problem has not been investigated yet . The reason for this might be that the often very large and heterogeneous “ network ” graphs in link prediction are of very different nature and require novel/adapted solutions ( e.g. , a protein interaction network may easily contain millions of nodes , whereas a molecule to be classified contains usually less than fifty ) . We explore application and effects of virtual nodes in link prediction theoretically and empirically : • We propose to use multiple virtual nodes in the link prediction scenario and describe a graph-based technique to connect them to the graph nodes . Consider Figure 1 . In a nutshell , we use a graph clustering algorithm to determine groups of nodes in the graph that belong together and then connect these nodes to a common virtual node . In this way , under-reaching is decreased because clustered nodes can share information easily ; at the same time , the nodes are spared of unnecessary information from unrelated nodes ( i.e. , in contrast to the single virtual node model ) . • We also investigate alternative methods to determine the virtual node connections ( e.g. , randomization in clustering ) and compare to the original model with a single virtual node . • We theoretically investigate the benefit of using ( multiple ) virtual nodes in terms of two aspects : influence score and the expressiveness in learning a structural link representation . • We conducted extensive experiments over challenging datasets of different type , provide ablation studies that confirm the superiority of our proposed techniques , analyze the results in detail , and provide first guidelines about how to use virtual nodes with different types of data and GNNs . • Most importantly , we show that our virtual node extensions most often yield rather stable performance increases and allow standard GNNs to compete with complex state-of-the-art models that also try to improve message passing , as well as with the models leading the OGB leaderboards . 2 RELATED WORK . We give an overview on approaches that are similar from a technical perspective ; for a more detailed summary , see Appendix A . For a more general overview of the large and diverse field of link prediction , we refer to good summaries in recent works ( Martínez et al. , 2016 ; Zhang et al. , 2020 ) . Deeper GNNs . Several techniques address over-smoothing and hence allow for constructing deeper GNNs to solve under-reaching . These models range from the simple but efficient message propagation in SGC ( Wu et al. , 2019 ; Liu et al. , 2020 ) and APPNP ( Klicpera et al. , 2019a ) and connections in JKNet ( Xu et al. , 2018 ) , to more advanced proposals ( Chen et al. , 2020 ; Sun et al. , 2021 ; Zhou et al. , 2020 ; Li et al. , 2020a ) such as the differentiable aggregation functions in DeeperGCN ( Li et al. , 2020a ) . However , although ( Chen et al. , 2020 ) show that over-smoothing happens particularly in dense graphs , the experiments in most of these works consider citation or recommendation networks , but not the especially dense and important biomedical ones . And our experiments over the latter suggest that the reported results are not generalizable to the more challenging biomedical data . Beyond One-Hop Neighbors . Recently , graph diffusion methods are used in various ways to determine the message targets and thus extend standard message passing beyond the one-hop neighborhood . Atwood & Towsley ( 2016 ) use k-hop random walks to extend the node features . APPNP ( Klicpera et al. , 2019a ) applies personalized PageRank to propagate the node predictions generated by a neural network . Other models concatenate ( Abu-El-Haija et al. , 2019 ) or aggregate ( Xu et al. , 2019a ; Ma et al. , 2020 ) node embeddings in every layer using a diffusion-based transition matrix . The diffusion-based graph neural network ( GDC ) ( Klicpera et al. , 2019b ) aggregates information from multiple neighborhood hops at each layer by sparsifying a generalized form of graph diffusion . Subsequent works use diffusion methods on multiple scales ( Liao et al. , 2019 ; Luan et al. , 2019 ; Xhonneux et al. , 2020 ) and attention ( Wang et al. , 2020 ) . Morris et al . ( 2019 ) take higher-order graph structures at multiple scales into account during message passing based on the k-dimensional Weisfeiler and Leman graph algorithm . All the above approaches are relatively complex , many terminated with memory errors in our experiments , and few have been evaluated for link prediction . Virtual Nodes . To the best of our knowledge , virtual nodes have only been considered in the context of graph classification so far , where a single virtual node ( also called supernode ) is added to the graph to be classified and connected to all graph nodes ( Gilmer et al. , 2017 ; Li et al. , 2017 ; Pham et al. , 2017 ; Ishiguro et al. , 2019 ) . Note that the original idea was to compute a graph embedding in parallel with the node embeddings and even connected the virtual node only in one direction ( i.e , via edges from the graph nodes ) instead of bidirectionally ( Li et al. , 2017 ) . There are some GNNs which point out special nodes that we could consider as “ virtual ” . Fey et al . ( 2020 ) propose a GNN for molecule graph classification which clusters certain nodes within a molecule using a structure-based , molecule-specific algorithm and then applies message passing within and between these clusters . The graph-partition based message passing from Liao et al . ( 2018 ) also used clustering , but it just divides the original messages into inter- and intra-cluster . Our approach creates new “ paths ” in the graph and we theoretically demonstrate its expressiveness . P-GNN ( You et al. , 2019 ) assigns nodes to random clusters ( “ anchor-sets ” ) and then creates a message for each node for every anchor-set , while ignoring the message passing from original direct neighbors . Our virtual nodes represent an alternative means to aggregate messages from multiple graph nodes which are not necessarily direct neighbors . We also explore the idea of similar random assignments in our context , but show that more elaborate techniques generally work better . Most importantly , we do not propose a specific , new GNN but a new technique for augmenting existing graph neural networks . Although it is a well-known trick , the advantage of using virtual nodes has never been theoretically investigated nor fully understood . We focus on link prediction and considerably extend the virtual node technique . There are commonalities in the advantages of using virtual nodes for graph classification and link prediction , but their role in link prediction is to improve the representation of the link instead of the graph ( nodes ) . We analyze theoretically and empirically how they improve GNN performance . 3 PRELIMINARIES . Link Prediction . We consider an undirected graphG = ( V , E ) with nodes V and edgesE ⊆ V ×V . Note that this basic choice is only for ease of presentation . All our techniques work for directed graphs and , with simple adaptation , also for graphs with labelled edges . We assume V to be ordered and may refer to a node by its index in V . For a node v ∈ V , Nv denotes the set of its neighbors . Given two nodes , the link prediction task is to predict whether there is a link between them . Message-Passing Graph Neural Networks . In this paper , we usually use the term graph neural networks ( GNNs ) to denote GNNs that use message passing as described by Gilmer et al . ( 2017 ) . These networks compute for every v ∈ V a node representation h ` v at layer ` ∈ [ 1 , 2 , . . . , k ] , by aggregating its neighbor nodes based on a generic aggregation function and then combine the obtained vector with h ` −1v as below ; h 0 v are the initial node features . h ` v = COMBINE ` ( h ` −1v , AGGREGATE ` ( { h ` −1u | u ∈ Nv } ) ) ( 1 ) Link prediction with GNNs is usually done by combining ( e.g. , concatenating ) the final representations hLu , h L v , of the nodes u , v under consideration and passing them through several feed-forward layers with a final sigmoid function for scoring . We follow this approach . We further use [ 1 , n ] to denote an interval [ 1,2 , . . . , n ] .

Summary: This paper investigates using virtual nodes in graph neural networks for link prediction. Specifically, the authors use a graph clustering algorithm to determine groups of nodes in the graph and adopt multiple virtual nodes in the graph for the link prediction senario. They also theoretically investigate the effect of using virtual nodes for link prediction. Experiments conducted on six datasets  provide insights and guidelines about using virtual nodes for link prediction.