Efficient Graph Neural Architecture Search

Recently , graph neural networks ( GNN ) have been demonstrated effective in various graph-based tasks . To obtain state-of-the-art ( SOTA ) data-specific GNN architectures , researchers turn to the neural architecture search ( NAS ) methods . However , it remains to be a challenging problem to conduct efficient architecture search for GNN . In this work , we present a novel framework for Efficient GrAph Neural architecture search ( EGAN ) . By designing a novel and expressive search space , an efficient one-shot NAS method based on stochastic relaxation and natural gradient is proposed . Further , to enable architecture search in large graphs , a transfer learning paradigm is designed . Extensive experiments , including node-level and graph-level tasks , are conducted . The results show that the proposed EGAN can obtain SOTA data-specific architectures , and reduce the search cost by two orders of magnitude compared to existing NAS baselines . 1 INTRODUCTION . Recent years have witnessed the success of graph neural networks ( GNN ) ( Gori et al. , 2005 ; Battaglia et al. , 2018 ) in various graph-based tasks , e.g. , recommendation ( Ying et al. , 2018a ) , chemistry ( Gilmer et al. , 2017 ) , circuit design ( Zhang et al. , 2019 ) , subgraph counting ( Liu et al. , 2020 ) , and SAT generation ( You et al. , 2019 ) . To adapt to different graph-based tasks , various GNN models , e.g. , GCN ( Kipf & Welling , 2016 ) , GAT ( Veličković et al. , 2018 ) , or GIN ( Xu et al. , 2019 ) , have been designed in the past five years . Most existing GNN models follow a neighborhood aggregation ( or message passing ) schema ( Gilmer et al. , 2017 ) , as shown in the left part of Figure 1 , which is that the representation of a node in a graph is learned by iteratively aggregating the features of its neighbors . Despite the broad applications of GNN models , researchers have to take efforts to design proper GNN architectures given different tasks by imposing different relational inductive biases ( Battaglia et al. , 2018 ) . As pointed out by Battaglia et al . ( 2018 ) , the GNN architectures can support one form of combinatorial generalization given different tasks , i.e. , graphs . Then a natural and interesting question can be asked : Can we automatically design state-of-the-art ( SOTA ) GNN architectures for graph-based tasks ? A straightforward solution is to adopt the neural architecture search ( NAS ) approaches , which have shown promising results in automatically designing architectures for convolutional neural networks ( CNN ) ( Zoph & Le , 2017 ; Pham et al. , 2018 ; Liu et al. , 2019a ; Tan & Le , 2019 ; You et al. , 2020a ) . However , it is nontrivial to adopt NAS to GNN . The first challenge is to define the search space . One can design a dummy search space to include as many as possible the related parameters , e.g. , aggregation functions , number of layers , activation functions , etc. , on top of the message passing framework ( Eq . ( 1 ) ) , however , it leads to quite a large discrete space , for example , 315,000 possible GNN architectures are generated by including just 12 types of model parameters in You et al . ( 2020b ) ) , which is challenging for any search algorithm . The second challenge is to design an effective and efficient search algorithm . In the literature , reinforcement learning ( RL ) based and evolutionary based algorithms have been explored for GNN architecture search ( Gao et al. , 2020 ; Zhou et al. , 2019 ; Lai et al. , 2020 ; Nunes & Pappa , 2020 ) . However , they are inherently computationally expensive due to the stand-alone training manner . In the NAS literature , by adopting the weight sharing strategy , one-shot NAS methods are orders of magnitude more efficient than RL based ones ( Pham et al. , 2018 ; Liu et al. , 2019a ; Xie et al. , 2019 ; Guo et al. , 2019 ) . However , the one-shot methods can not be directly applied to the aforementioned dummy search space , since it remains unknown how to search for some model parameters like number of layers and activation functions by the weight sharing strategy . Therefore , it is a challenging problem to conduct effective and efficient architecture search for GNN . In this work , we propose a novel framework , called EGAN ( Efficient GrAph Neural architecture search ) , to automatically design SOTA GNN architectures . Motivated by two well-established works ( Xu et al. , 2019 ; Garg et al. , 2020 ) that the expressive capabilities of GNN models highly rely on the properties of the aggregation functions , a novel search space consisting of node and layer aggregators is designed , which can emulate many popular GNN models . Then by representing the search space as a directed acyclic graph ( DAG ) ( Figure 1 ( c ) ) , we design a one-shot framework by using the stochastic relaxation and natural gradient method , which can optimize the architecture selection and model parameters in a differentiable manner . To enable architecture search in large graphs , we further design a transfer learning paradigm , which firstly constructs a proxy graph out of the large graph by keeping the properties , and then searches for GNN architectures in the proxy graph , finally transfers the searched architecture to the large graph . To demonstrate the effectiveness and efficiency of the proposed framework , we apply EGAN to various tasks , from node-level to graph-level ones . The experimental results on ten different datasets show that EGAN can obtain SOTA data-specific architectures for different tasks , and at the same time , reduce the search cost by two orders of magnitude . Moreover , the transfer learning paradigm , to the best of our knowledge , is the first framework to enable architecture search in large graphs . Notations . Let G = ( V , E ) be a simple graph with node features X ∈ RN×d , where V and E represent the node and edge sets , respectively . N represents the number of nodes and d is the dimension of node features . We use N ( v ) to represent the first-order neighbors of a node v in G , i.e. , N ( v ) = { u ∈ V| ( v , u ) ∈ E } . In the literature , we also create a new set Ñ ( v ) is the neighbor set including itself , i.e. , Ñ ( v ) = { v } ∪ { u ∈ V| ( v , u ) ∈ E } . 2 RELATED WORKS . GNN was first proposed in ( Gori et al. , 2005 ) and in the past five years different GNN models ( Kipf & Welling , 2016 ; Hamilton et al. , 2017 ; Veličković et al. , 2018 ; Gao et al. , 2018 ; Battaglia et al. , 2018 ; Xu et al. , 2019 ; 2018 ; Liu et al. , 2019b ; Abu-El-Haija et al. , 2019 ; Wu et al. , 2019 ; Zeng et al. , 2020 ; Zhao & Akoglu , 2020 ; Rong et al. , 2020 ) have been designed , and they rely on a neighborhood aggregation ( or message passing ) schema ( Gilmer et al. , 2017 ) , which learns the representation of a given node by iteratively aggregating the hidden features ( “ message ” ) of its neighbors . Besides , in Xu et al . ( 2018 ) ; Chen et al . ( 2020 ) , the design of residual networks ( He et al. , 2016a ; b ) are incorporated into existing message passing GNN models . Battaglia et al . ( 2018 ) pointed out that the GNN architectures can provide one form of combinatorial generalization for graph-based tasks , and Xu et al . ( 2019 ) ; Garg et al . ( 2020 ) further show that the expressive capability of existing GNN models is upper bounded by the well-known Weisfeiler-Lehman ( WL ) test . It will be an interesting question to explore GNN architectures with better combinatorial generalization , thus neural architecture search ( NAS ) can be a worthwhile approach for this consideration . NAS ( Baker et al. , 2017 ; Zoph & Le , 2017 ; Elsken et al. , 2018 ) aims to automatically find SOTA architectures beyond human-designed ones , which have shown promising results in architecture design for CNN and Recurrent Neural Network ( RNN ) ( Liu et al. , 2019a ; Zoph et al. , 2018 ; Tan & Le , 2019 ) . Existing NAS approaches can be roughly categorized into two groups according to search methods ( Bender et al. , 2018 ) : i.e . the stand-alone and one-shot ones . The former ones tend to obtain the SOTA architecture from training thousands of architectures from scratch , including reinforcement learning ( RL ) based ( Baker et al. , 2017 ; Zoph & Le , 2017 ) and evolutionary-based ones ( Real et al. , 2019 ) , while the latter ones tend to train a supernet containing thousands of architectures based on the weight sharing strategy , and then extract a submodel as the SOTA architecture at the end of the search phase ( Pham et al. , 2018 ; Bender et al. , 2018 ; Liu et al. , 2019a ; Xie et al. , 2019 ) . The difference of training paradigms leads to that one-shot NAS methods tend to be orders of magnitude more efficient than the RL based ones . Recently , there are several works on architecture search for GNN , e.g. , RL based ones ( Gao et al. , 2020 ; Zhou et al. , 2019 ; Lai et al. , 2020 ) , and evolutionary-based ones ( Nunes & Pappa , 2020 ; Jiang & Balaprakash , 2020 ) , thus all existing works are computationally expensive . Franceschi et al . ( 2019 ) proposes to jointly learn the edge probabilities and the parameters of GCN given a graph , thus orthogonal to our work . Besides , Peng et al . ( 2020 ) proposes a one-shot NAS method for GCN architectures in the human action recognition task . In Li et al . ( 2020 ) , a sequential greedy search method based on DARTS ( Liu et al. , 2019a ) is proposed to search for GCN architectures , however , the tasks are more focused on vision related tasks , with only one dataset in conventional node classification task . In this work , to the best of our knowledge , for conventional node-level and graph-level classification tasks , we are the first to design a one-shot NAS method for GNN architecture search , which is thus , in nature , more efficient than existing NAS methods for GNN . In Appendix A.3 , we give more detailed discussions about the comparisons between EGAN and more recent GNN models . 3 THE PROPOSED FRAMEWORK . 3.1 THE DESIGN OF SEARCH SPACE . As introduced in Section 2 , most existing GNN architectures rely on a message passing framework ( Gilmer et al. , 2017 ) , which constitutes the backbone of the designed search space in this work . To be specific , a K-layer GNN can be written as follows : the l-th layer ( l = 1 , · · · , K ) updates hv for each node v by aggregating its neighborhood as h ( l ) v = σ ( W ( l ) · Φn ( { h ( l−1 ) u , ∀u ∈ Ñ ( v ) } ) ) , ( 1 ) where h ( l ) v ∈ Rdl represents the hidden features of a node v learned by the l-th layer , and dl is the corresponding dimension . W ( l ) is a trainable weight matrix shared by all nodes in the graph , and σ is a non-linear activation function , e.g. , a sigmoid or ReLU . Φn is the key component , i.e. , a pre-defined aggregation function , which varies across different GNN models . Thus a dummy search space is to include as many as possible related parameters in Eq . ( 1 ) . However , it leads to a very large search space , making the search process very expensive . In this work , motivated by two well-established works ( Xu et al. , 2019 ; Garg et al. , 2020 ) , which show that the expressive capabilities of GNN models highly rely on the properties of aggregation functions , we propose to search for different aggregation functions by simplifying the dummy search space . For other parameters , we do simple tuning in the re-training stage , which is also a standard practice in existing NAS methods ( Liu et al. , 2019a ; Xie et al. , 2019 ) . Then the first component of the proposed search space is the node aggregators , which consists of existing GNN models . To improve the expressive capability , we add the other component , layer aggregators , to combine the outputs of node aggregator in all layers , which have been demonstrated effective in JK-Network ( Xu et al. , 2018 ) . Then we introduce the proposed search space , as shown in Figure 1 ( c ) , in the following : • Node aggregators : We choose 12 node aggregators based on popular GNN models , and they are presented in Table 7 in Appendix A.1 . The node aggregator set is denoted by On . • Layer aggregators : We choose 3 layer aggregators as shown in Table 7 in Appendix A.1 . Besides , we have two more operations , IDENTITY and ZERO , related to skip-connections . Instead of requiring skip-connections between all intermediate layers and the final layer in JK-Network , in this work , we generalize this option by proposing to search for the existence of skip-connections between each intermediate layer and the last layer . To connect , we choose IDENTITY , and ZERO otherwise . The layer aggregator set is denoted by Ol and the skip operation set by Os . To further guarantee that K-hop neighborhood can always be accessed , we add one more constraint that the output of the node aggregator in the last layer should always be used as the input of the layer aggregator , thus for a K-layer GNN architecture , we need to search K− 1 IDENTITY or ZERO for the skip-connection options .

Summary: The paper proposes a framework for efficient architecture search for graphs. This is done by combining a differentiable DARTS-like architecture encoding with a transfer learning method, that searches on smaller graphs with similar properties, and then transfers to the target graphs. The experiments show that EGAN matches or exceeds both hand-designed and NAS-designed GNNs. Moreover, the method is very fast to run.