Differentiable Graph Optimization for Neural Architecture Search

1 INTRODUCTION . Neural Architecture Search ( NAS ) methods achieve great success and outperform hand-crafted models in many deep learning applications , such as image recognition , object detection and natural language processing ( Zoph et al. , 2017 ; Liu et al. , 2019 ; Ghiasi et al. , 2019 ; Chen et al. , 2020 ) . Due to the expensive cost of training-evaluating a neural architecture , the key challenge of NAS is to explore possible good candidates effectively . To cope with this challenge , various methods have been proposed , such as reinforcement learning ( RL ) , evolution algorithm ( EA ) , Bayesian optimization ( BO ) and weight-sharing strategy ( WS ) , to perform efficient search ( Zoph & Le , 2016 ; Real et al. , 2019 ; Hutter et al. , 2011 ; Liu et al. , 2019 ; Guo et al. , 2019 ) . While the weight-sharing strategy improves overall efficiency by reusing trained weights to reduce the total training cost , zeroth-order algorithms like RL , EA and BO employ black-box optimization , with the goal of finding optimal solutions with fewer samples . However , the search space of NAS is exponentially growing with the increasing number of choices . As a result , such huge combinatorial search spaces lead to insufficient exploitation of black-box learning framework ( Luo et al. , 2018 ) . Another line of research has been focused on formulating the NAS search space as graph structures , typically directed acyclic graphs ( DAGs ) , and then the search target is cast as choosing an optimal combination of the nodes and edges in the graph structure ( Pham et al. , 2018 ; Liu et al. , 2019 ; Xie et al. , 2019 ) . However , existing methods tend to perform the optimization in the indirect manner using black-box optimization . In contrast , we aim to explicitly model the search space as graphs and optimize graph structures directly . We thus propose Graph Optimized Neural Architecture Learning ( GOAL ) , a novel NAS approach combined with graph learning for efficient exploitation , as briefly shown in Fig.1 . Unlike other black-box approaches , we use a differentiable surrogate model to directly optimize the graph structures . The surrogate model takes a graph structure corresponds to a neural architecture as input , and predicts a relative ranking score as the searching signal . We then apply gradient descent on the input graph structure to optimize the corresponding architecture , which attempts to obtain a better predicted ranking score . As we optimize the surrogate model and the architectures iteratively , the optimal architectures could be typically obtained after a few iterations . In particular , to cope with the difficulty of using gradient-based optimization on the discrete graph structure , we adapt the proximal algorithm for allowing us to optimize discrete variables in a differentiable manner ( Parikh et al. , 2014 ; Bai et al. , 2019 ; Yao et al. , 2019 ) . We build the surrogate model with Edge Conditional Convolution ( ECC ) ( Simonovsky & Komodakis , 2017 ) , a variant of Graph Convolutional Network ( GCN ) ( Kipf & Welling , 2017 ) , to handle the graph representation of various search spaces . Our empirical results on existing NAS benchmarks with different search spaces demonstrate that , GOAL outperforms exist state-of-the-art black-box optimization baselines and neural surrogate model based methods by a large margin in terms of the search efficiency . The main contributions of this paper are summarized as follows : • We propose a differentiable surrogate model for ranking neural architectures based on GNN , which takes advantage of the graph structure of neural architectures and guides the architecture search efficiently . • We present GOAL , a novel gradient-based NAS sample-efficient approach with the assistance of the proposed surrogate model . Comparing to exist algorithms with GNN surrogates , GOAL makes full use of the learned representation by jointly optimizing the GNN model and candidate architectures and performing efficient exploitation within graphstructured search spaces . • Our empirical results demonstrate that the GOAL significantly outperforms existing stateof-the-art methods in various search spaces settings . 2 RELATED WORKS . 2.1 NEURAL ARCHITECTURE SEARCH . From the earliest explorations on automatic neural network designing to recent NAS trends , the NAS problem developed from hyper-parameter optimization , becoming a more challenging task due to the inherent complexity of its search space ( Bergstra et al. , 2013 ; Elsken et al. , 2018 ) . Existing popular approaches include various kinds of algorithms : reinforcement learning ( Zoph & Le , 2016 ; Zoph et al. , 2017 ; Pham et al. , 2018 ) , evolution algorithm ( Real et al. , 2019 ) , Bayesian optimization ( Falkner et al. , 2018 ; White et al. , 2019 ) , monte carlo tree search ( Wang et al. , 2019b ; a ) , gradient based methods ( Liu et al. , 2019 ; Luo et al. , 2018 ) , etc . There are also some works employ surrogate models to predict the performance of architectures before training to reduce the cost of architecture evaluation ( Liu et al. , 2018 ; Wen et al. , 2019 ; Wang et al. , 2019b ) . Some most recently parallel works even tries to improve black-box optimization methods like Bayesian optimization by efficient surrogate predictors ( White et al. , 2019 ; Shi et al. , 2019 ) . Existing gradient-based methods usually employ weight-sharing based relaxation or use encoder-decoder to optimize a continuous hidden space ( Liu et al. , 2019 ; Luo et al. , 2018 ) . These approximations cause biased model criterion and generation , which can harm the final performance ( Yu et al. , 2019 ; Yang et al. , 2020 ) . In contrast , our method directly optimizes the discrete architectures , avoids the biased model criterion . Due to the complex search settings and expensive evaluation cost , NAS works are hard to evaluate and reproduce fairly ( Li & Talwalkar , 2019 ; Yang et al. , 2020 ) . NAS-Bench-101 and NAS-Bench-201 are proposed to provide fair and easy benchmarks for NAS algorithms ( Dong & Yang , 2020 ; Ying et al. , 2019 ) . Taking advantage of the development of graph neural networks ( GNNs ) , several recent works show that the power of GNNs can also benefit the NAS task ( Shi et al. , 2019 ; Wen et al. , 2019 ) . However , these works only employ GNNs as a powerful black-box predictor . In contrast , we make full use of the learned representation of GNN by back-propagating through the GNN model and performing gradient-based optimization on the input graphs . 2.2 GRAPH STRUCTURE LEARNING . Graph neural network ( GNN ) is a kind of neural model for extracting features from graph-structured data ( Zhou et al. , 2018 ) . The most popular family of GNNs is the Graph Convolutional Network ( GCN ) , which propagates the features of nodes in the spectral or spatial domain ( Kipf & Welling , 2017 ; Morris et al. , 2019 ) . Since the vanilla GCN can only handle monomorphic graphs with single type of edge , many variants have been proposed to handle more complex graphs ( Schlichtkrull et al. , 2018 ; Simonovsky & Komodakis , 2017 ) . As the normal GNN pipeline requires an available fine-grained graph structure , recent works present approaches to optimize the graph structure and GNN model jointly ( Chen et al. , 2019 ; Franceschi et al. , 2019 ) . The graph structures could be either constructed from scratch or fine-tuned from a sub-optimal graph with the GNN model . We follow this manner , iteratively optimize the graph of neural architectures , guided by the GNN model . 3 PRELIMINARY : NEURAL ARCHITECTURE SEARCH SPACES . While our method can be potentially applied more generally , in this paper we focus on the most commonly used cell-based search spaces of convolutional neural networks ( CNNs ) ( Zoph et al. , 2017 ; Pham et al. , 2018 ; Liu et al. , 2019 ; Ying et al. , 2019 ; Dong & Yang , 2020 ) . A cell is a basic building block of architectures represented as a DAG , which consists of several basic neural operations and their connectivity relationship as shown in Fig.2 . Popular basic operations include different kinds of convolutions and poolings , such as 1 × 1 convolution , 3 × 3 convolution , 3 × 3 max pooling , etc . A complete CNN architecture is thus formed by stacking several repeats of the same cell . There are different kinds of representations to form an architecture by a DAG . Fig.2 shows two typical representations . It is non-trivial to convert from one to each other , since the corresponding architectures formed by different representations do not coincide . The cell-based search spaces are usually heuristically constrained by limiting the max number of nodes and edges , or the degree of each node , which corresponds to the number of operations and connections in a cell ( Liu et al. , 2019 ; Ying et al. , 2019 ) . This brings different feasible sets in our optimization in Sec.5.1 . 4 LEARNING TO RANK ARCHITECTURES BY GRAPH . We first design the GNN model which guides our optimization on neural architectures . Denoting a neural architecture setting as α , we train a differentiable model f ( α ) which predicts a score ỹ to indicate the real performance y , e.g . accuracy on the image classification task , of α . We build the model as GNN by Edge Conditional Convolution ( ECC ) layers ( Simonovsky & Komodakis , 2017 ) . While the GNN model itself is fully differentiable , it is possible to obtain gradient w.r.tα from f ( α ) to optimize the architecture . 4.1 GRAPH REPRESENTATION OF NEURAL ARCHITECTURES . As described in Sec.3 , mainstream NAS works search architectures in search spaces which could be represented as directed acyclic graphs ( DAGs ) . We form these search spaces as a unified perspective . Denoting G = ( V , E , T , A ) as the graph corresponding to an architecture α , where V is the set of nodes , E is the set of edges , T and A denote the types of nodes and edges respectively . For a predefined search space , we can form each architecture in the space by fixing V and E as a fullyconnected DAG , where each node i is connected to each node j for any i < j ≤ |V | , and assign proper T and A . As the connectivity between nodes can also be denoted as an edge type ( i.e . denoting as a disconnected type , as shown in the left of Fig.2 ) , we can determine both operation types and connectivity of the graph by T and A . The neural architecture setting α is therefore formed by α = ( T , A ) , consisting of the types of nodes and types of edges for graph structure . In particular , we form each component of T and A as a one-hot vector , which encodes a categorical selection . To build the surrogate model dealing with various nodes and edges types in the graph representation of neural architectures , we employ the ECC , a variant of GCN , to extract the features hi of each node i in the graph ( Simonovsky & Komodakis , 2017 ) . For each node i in the directed graph G , the l-th layer of a L-layers ECC propagates features h ` −1j of the previous layer from node j ∈ N ( i ) to node i , where N ( i ) denotes the predecessors set of node i . The ECC learns a kernel generator ke : A 7→ Rd ` ×d ` −1 , which maps the edge attribute space A to a real-valued d ` × d ` −1 filter for the graph convolution , where d ` indicates the number of channels in the ` -th convolution layer . Thus the graph convolution is formed by : h ` i = Θ ` h ` −1i + ∑ j∈N ( i ) k ` e ( aji ) h ` −1 j ( 1 ) where aji is the attribute , i.e . categorical selection as one-hot vector , of the edge from node j to node i , Θ is a learnable parameter . For high-level prediction upon the whole graph , we compute the feature vector hG of the whole graph G by taking the weighted sum of each node features of the last layer ( Li et al. , 2016 ) : hG = ∑ i exp ( whi ) ∑ j exp ( whj ) hi ( 2 ) where w is a learnable parameter to compute the weight for features of each node . Followed by fully-connected layers that take hG as input , we can get the final prediction ỹ as an indicator score of the input architecture .

Summary: The authors address the Neural Architecture Search problem. At the core of their contribution is an architectural improvement; performance prediction of a considered architecture is much better when using a particular graph neural network on (softened) architecture topology. The rest of the NAS pipeline is naturally built around this observation and the final performance looks quite strong.