AutoGEL: An Automated Graph Neural Network with Explicit Link Information

1 Introduction . As one of ubiquitous data structures , graph G ( E , V ) contains the node-set V = { v1 , · · · , vn } and edge-set E = { e ( vi , vj ) : vi , vj 2 V } , which can represent a lot of real-world data sets , such as social networks [ Bu et al. , 2018 ] , physical systems [ Sanchez-Gonzalez et al. , 2018 ] , protein-protein interaction networks [ Yue et al. , 2020 ] . In recent years , Graph Neural Networks ( GNNs ) have been introduced for various graph tasks and achieve unprecedented success , such as node classification [ Kipf and Welling , 2016a , Hamilton et al. , 2017 ] , link prediction [ Vashishth et al. , 2019 , Li et al. , 2020 ] , and graph classification [ Niepert et al. , 2016 , Zhang et al. , 2018 ] . Generally , GNNs encode G ( V , E ) into the d-dimensional vector space ( e.g. , V 2 R|V |⇥d ) that preserves similarity in the original graph . Despite the great success , those GNNs are restricted to specific instances within GNN design space [ You et al. , 2020 ] . Different graph tasks usually require different GNN architectures [ Gilmer et al. , 2017 ] . For example , compared with the node classification task , GNNs for graph classification introduces an extra readout phase to obtain graph embeddings . However , architecture design for these GNNs remains a challenging problem due to the diversity of graph data sets . Given a graph task , a GNN architecture performs well on one data set does not necessarily imply that it is also suitable for other data sets [ You et al. , 2020 ] . Some pioneer works have been proposed to alleviate the above issue in GNN models . They introduce Neural Architecture Search ( NAS ) [ Elsken et al. , 2019 ] approaches to automatically design suitable ⇤Corresponding author 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . GNN architecture for the given data set ( i.e. , AutoGNN ) [ Zhou et al. , 2019 , Gao et al. , 2020 , Jiang and Balaprakash , 2020 , Zhao et al. , 2021 ] . Architectures identified by these AutoGNNs rival or surpass the best handcrafted GNN models , demonstrating the potential of AutoGNN towards better GNN architecture design . Unfortunately , the existing AutoGNNs are mainly designed for the node classification and graph classification task . Their designs do not include edge embeddings , i.e. , modeling and organizing link information in an implicit way . First , it is difficult for existing AutoGNNs to handle another important graph task on the edge-level , link prediction ( LP ) task . Second , lack of edge embedding makes them inexpressive to leverage the complex link information , such as direction information of edges and different edge types in multi-relational graphs . Especially , various edge types could impose different influence for encoding nodes into embeddings , which can further improve the model performance on the node-level and graph-level tasks . Therefore , a new AutoGNN is desired to model link information explicitly on various data sets . To bridge the aforementioned research gap , we propose AutoGEL , a novel AutoGNN framework with Explicit Link information , which can handle the LP task and improve performance of AutoGNNs on other graph tasks . Specifically , AutoGEL explicitly learns the edge embedding in the message passing framework to model the complex link information , and introduces the several novel design dimensions into the GNN search space , enabling a more powerful GNN to be searched for any given graph data set . Moreover , AutoGEL adopts a robust differentiable search algorithm to guarantee the effectiveness of searched architectures and control the computational footprint . We summarize the contributions of this work as follows : • The design of existing AutoGNNs follows an implicit way to leverage and organize the link information , which can not handle the LP task and limits the performance of AutoGNNs on other graph tasks . In this paper , we present a novel method called AutoGEL to solve these issues through explicitly modeling the link information in the graphs . • AutoGEL introduces several novel design dimensions into the GNN search space at both the intra-layer and inter-layer designs , so as to improve the task performance . Moreover , motivated by one robust NAS algorithm SNAS , AutoGEL upgrades the search algorithm adopted in existing AutoGNNs to further guarantee the effectiveness of final derived GNN . • The experimental results demonstrate that AutoGEL can achieve better performance than manually designed models in the LP task . Furthermore , AutoGEL shows excellent competitiveness with other AutoGNN works on the node and graph classification tasks . 2 Related Work . 2.1 General Message Passing Framework . The majority of GNNs follow the neighborhood aggregation schema [ Gilmer et al. , 2017 ] , i.e. , the Message Passing Neural Network ( MPNN ) , which is formulated as : m k+1 v = AGGk ( { Mk ( h k v , h k u , e k vu ) : u 2 N ( v ) } ) , hk+1 v = ACTk ( COMk ( { h k v , m k+1 v } ) ) , ( 1 ) ŷ = R ( { hL v |v 2 G } ) , ( 2 ) where k denotes k-th layer , N ( v ) denotes a set of neighboring nodes of v , hk v , hk u denotes hidden embeddings for v and u respectively , ek vu denotes features for edge e ( v , u ) ( optional ) , mk+1 v denotes the intermediate embeddings gathered from neighborhood N ( v ) , Mk denotes the message function , AGGk denotes the neighborhood aggregation function , COMk denotes the combination function between intermediate embeddings and embeddings of node v itself from the last layer , ACTk denotes activation function . Such message passing phase in ( 1 ) repeats for L times ( i.e. , k 2 { 1 , · · · , L } ) . For graph-level tasks , it further follows the readout phase in ( 2 ) where information from the entire graph G is aggregated through readout function R ( · ) . 2.2 Automated Graph Neural Networks ( AutoGNN ) . In recent years , AutoGNN has emerged as a promising direction towards better graph neural architecture design [ Zhou et al. , 2019 , Gao et al. , 2020 , Jiang and Balaprakash , 2020 , You et al. , 2020 , Zhao et al. , 2021 , Ding et al. , 2021 ] . To enable a powerful GNN architecture to be searched , AutoGNNs first propose the GNN search space in the intra-layer level , i.e. , providing common choices for several important operators in one MPNN layer ( see ( 1 ) and ( 2 ) ) . Here we summarize candidate choices for those operators : • Message Function Mk : Existing AutoGNNs mainly focus on the node-level and graph-level task , thus edge embeddings are often not available . Mk ( hv , hu , evu ) is reduced to Mk ( hv , hu ) . Typically , M ( hv , hu ) = avuWhu , where avu denotes the attention scores , and W 2 Rd⇥d denotes the linear transformation matrix . Note that NAS-GCN [ Jiang and Balaprakash , 2020 ] only takes the edge feature evu as input without learning edge embeddings . We next denote the edge embedding to he for distinguishing . • Aggregation AGGk : It controls the way to aggregate message from nodes ’ neighborhood . It can be any differentiable and permutation invariant functions , usually AGGk 2 [ sum , mean , max ] . And sum ( · ) = P u2N ( v ) Mk ( hv , hu ) , mean ( · ) = P u2N ( v ) Mk ( hv , hu ) /|N ( v ) | , and max ( · ) denotes channel-wise maximum across the node dimension . • Combination COMk : It determines the way to merge messages between neighborhood and node itself . In literature , COMk is selected from [ concat , add , mlp ] , where concat ( · ) = [ hkv , mk+1v ] , add ( · ) = hk v +mk+1 v , and mlp ( · ) = MLP ( hk v +mk+1 v ) ( MLP is Multi-layer Perceptron ) . • Activation ACTk : [ identity , sigmoid , tanh , relu , elu ] are some of the most commonly used activation functions [ Gao et al. , 2020 ] . • Graph Pooling : The pooling operator has been introduced for the graph classification task , such as [ global pool , global attention pool , f latten ] [ Jiang and Balaprakash , 2020 , Wei et al. , 2021 ] . In addition to the above intra-layer operators , several works [ You et al. , 2020 , Zhao et al. , 2021 , Jiang and Balaprakash , 2020 ] propose the idea of searching layer connectivity to combine hidden representations of different layers in a better way , i.e. , inter-layer design . More details will be discussed in Sec . 3.1.2 . After the search space design , AutoGNNs adopt various search algorithms to search the optimal architecture from the search space . AGNN [ Zhou et al. , 2019 ] and GraphNAS [ Gao et al. , 2020 ] follow the reinforcement learning ( RL ) [ Williams , 1992 ] way to search architectures . They utilize a recurrent neural network controller for architecture sampling , and update the controller to maximize the expected performance of sampled architectures . NAS-GCN [ Jiang and Balaprakash , 2020 ] adopt evolutionary algorithm ( EA ) , where new architectures are generated by performing mutation from parent architectures and the population , i.e. , the best performing architectures , are iteratively updated . However , both RL and EA algorithms require a large number of architectures to be sampled , which is inherently computational expensive . To improve the search efficiency , SANE [ Zhao et al. , 2021 ] adopts a deterministic differentiable search algorithm DARTS [ Liu et al. , 2018 ] , where a supernet containing all candidate operators is constructed and architecture parameters are jointly optimized with network parameters through gradient descent . Unfortunately , it has been discussed in SNAS [ Xie et al. , 2018 ] that DARTS suffer from the unbounded bias issue towards its objective , which limits the performance of the final derived model . In Tab . 1 , we summarize existing AutoGNNs from several perspectives : graph , MPNN space , search algorithm , and task scenario . 2.3 GNNs for Link Prediction Task . Even with the great effort invested into the construction and maintenance of networks , many graph data sets still remain incomplete [ Schlichtkrull et al. , 2018 ] . Therefore , the link prediction ( LP ) task is one of the most crucial problems in the graph-structured data , which aims to recover those missing links in a graph [ Zhang et al. , 2019 , 2020a ] , i.e. , predicting the missing node in e ( vi , ? ) where ? denotes the target node that has the potential link with vi . Recently , GNN models have been introduced to handle the LP task ( abbreviated to GLP models ) , which can be roughly categorized based on its application scenarios : GLP on homogeneous graphs ( i.e. , only one type of nodes and edges [ Yang et al. , 2016 ] ) , and multi-relational graphs ( i.e. , graphs with multiple edge types [ Toutanova and Chen , 2015 ] ) .

Summary: This paper proposes a neural architecture search method for using GNNs for link prediction. The proposed method explicitly models edge information in the search process, and proposes several novel design and search space choices for AutoGNN. Through a DARTS-like differentiable search method, the proposed method achieves state-of-the-art performance on various link prediction and node/graph classification tasks.