Neural Architecture Search by Learning Action Space for Monte Carlo Tree Search

1 INTRODUCTION During the past two years , there has been a growing interest in Neural Architecture Search ( NAS ) that aims to automate the laborious process of designing neural networks . Architectures found by NAS have achieved remarkable results in image classification ( Zoph and Le ( 2016 ) ; Real et al . ( 2018 ) ) , object detection and segmentation ( Ghiasi et al . ( 2019 ) ; Chen et al . ( 2019 ) ; Liu et al . ( 2019 ) ) , as well as other domains such as language tasks ( Luong et al . ( 2018 ) ; So et al . ( 2019 ) ) . Starting from hand-designed discrete model space and action space , NAS utilizes search techniques to explore the search space and find the best performing architectures with respect to a single or multiple objectives ( e.g. , accuracy , latency , or memory ) , and preferably with minimal search cost . However , one common issue faced by the previous works on NAS is that the action space needs to be manually designed . The action space proposed by Zoph et al . ( 2018 ) involves sequential actions to construct a network , such as selecting two nodes , and choosing their operations . Other prior works including reinforcement learning based ( Zoph and Le ( 2016 ) ; Baker et al . ( 2016 ) ) , evolution-based ( Real et al . ( 2018 ; 2017 ) ) , and MCTSbased ( Wang et al . ( 2019 ) ; Negrinho and Gordon ( 2017 ) ) approaches , all use manually designed action spaces . As suggested in Sciuto et al . ( 2019 ) ; Li and Talwalkar ( 2019 ) ; Xie et al . ( 2019 ) , action space design alone can be critical to network performance . Furthermore , it is often the case that manually designed action space is not related to the performance that needs to be optimized . In Sec 2 , we demonstrate an example where subtly different action space can lead to significantly different search efficiency . Finally , unlike games that generally have a predefined action space ( e.g. , Atari , Chess and Go ) , in NAS , it is the final network that matters rather than the specific path of the search , which gives a large playground for action space learning . Based on the above observations , we propose Latent Action Neural Architecture Search ( LaNAS ) that learns the action space to maximize search efficiency for performance metrics . Previous methods typically construct an architecture from an empty network by sequentially applying predefined actions , e.g . adding layers , setting the kernel or depth , resulting in a state space that provides undiscriminating rewards to inform the search ( Fig . 2a sequential ) . LaNAS takes a dual approach and treats each action as a linear constraint that intersects with the search space Ω to bifurcate it into a good and bad region that contains networks with similar performance metric ( Fig . 1 ) . Once multiple actions are recursively applied to the entire search space , the search space is clearly separated ( e.g . Fig . 2a global , Fig . 4 ) into regions that provide more informative rewards , thereby greatly improving the search efficiency . To achieve this goal , LaNAS iterates between learning and searching stage . In the learning stage , each action in LaNAS is learned to partition the model space into high-performance and low-performance regions , to achieve accurate performance prediction . In the searching stage , LaNAS applies MCTS on the learned action space to sample more model architectures . The learned actions provide an informed guide for MCTS to sample from performance-promising regions , while the adaptive exploration in MCTS collects more data to progressively refine the learned actions to zoom into more promising regions for further improving the sample efficiency . The iterative process is jump-started by first collecting a few random samples . We show that LaNAS yields a tremendous acceleration on a diverse set of benchmark tasks , including publicly available NASBench-101 ( 420,000 NASNet models trained on CIFAR-10 ) ( Ying et al . ( 2019 ) ) , our self-curated ConvNet-60K ( 60,000 plain VGG-style ConvNets trained on CIFAR-10 ) , and LSTM-10K ( 10,000 LSTM cells trained on PTB ) . Our algorithm consistently finds the best performing architecture on all three tasks with at least an order of fewer samples ( on average ) than Random Search , Regularized Evolution , MCTS , Neural Architecture Optimization and Bayesian Optimization . In the open domain search scenario , our algorithm finds a network that achieves 98.0 % accuracy on CIFAR-10 and 75.0 % top1 accuracy ( mobile setting ) on ImageNet in only 803 samples , using 33× fewer samples and achieving higher accuracy than AmoebaNet ( Real et al . ( 2018 ) ) . Moreover , we empirically demonstrate that the learned latent actions can transfer to a new search task to further boost efficiency . Finally , we provide empirical observations to illustrate the search dynamics and analyze the behavior of our approach . We also conduct various ablation studies in together with a partition analysis to provide guidance in determining search hyper-parameters and deploying LaNAS in practice . 2 A MOTIVATING EXAMPLE . To demonstrate the importance of action space in NAS , we start with a motivating example . Consider a simple scenario of designing a plain Convolutional Neural Network ( CNN ) for CIFAR-10 image classification . The primitive operation is a Conv-ReLU layer . Free structural parameters that can vary include network depth L = { 1 , 2 , 3 , 4 , 5 } , number of filter channels C = { 32 , 64 } and kernel size K = { 3 × 3 , 5 × 5 } . This configuration results in a search space of 1,364 networks . To perform the search , there are two natural choices of the action space : sequential and global . sequential comprises actions in the following order : adding a layer l , setting kernel size Kl , setting filter channel Cl . The actions are repeated L times . On the other hand , global uses the following actions instead : { Setting network depth L , setting kernel size K1 , ... , L , setting filter channel C1 , ... , L } . For these two action spaces , MCTS is employed to perform the search . Note that both action spaces can cover the entire search space but have very different search trajectories . Fig . 2 ( a ) visualizes the search for these two action spaces . Actions in global clearly separates desired and undesired network clusters , while actions in sequential lead to network clusters with a mixture of good or bad networks in terms of performance . As a result , the overall distribution of accuracy along the search path ( Fig . 2 ( b ) ) shows concentration behavior for global , which is not the case for sequential . We also demonstrate the overall search performance in Fig.2 ( c ) . As shown in the figure , global finds desired networks much faster than sequential . This observation suggests that changing the action space can lead to very different search behavior and thus potentially better sample efficiency . In other words , an early exploration on the network depth is critical . Increasing the depth is an optimization direction that can potentially lead to better model accuracy . One might come across a natural question from this motivating example . Is it possible to find a principle way to distinguish a good action space from a bad action space in NAS ? Is it possible to learn an action space such that it can best fit the performance metric to be optimized ? 3 LEARNING LATENT ACTIONS . In this section , we describe LaNAS , which comprises two phases : ( 1 ) learning phase , and ( 2 ) search phase . Fig . 3 presents a high level description of LaNAS , of which the corresponding algorithms are further described in Alg.1 . 3.1 LEARNING PHASE . In the learning phase at iteration t , we have a dataset Dt = { ( ai , vi ) } obtained from previous explorations . Each data point ( ai , vi ) in Dt has two components : ai represents network attributes ( e.g. , depth , number of filters , kernel size , connectivity , etc ) and vi represents the performance metric estimated from training ( or from pre-trained dataset such as NASBench-101 ) . Our goal is to learn a good action space from Dt to guide future exploration as well as to find the model with the desired performance metric efficiently . Starting from the entire search space Ω , the learnt actions recursively ( and greedily ) split it into smaller regions such that the estimation of performance metric becomes more accurate . This helps us prune away poor regions as soon as possible and increase the sample efficiency of NAS . In particular , we model the recursive splitting process as a tree . The root node corresponds to the entire model space Ω , while each tree node j corresponds to a region Ωj ( Fig . 1 ) . At each tree node j , we partition Ωj into disjoint regions Ωj = ∪k∈ ( j ) Ωk , such that on each child region Ωk , the estimation of performance metric V ( Ωk ) is the most accurate ( or equivalently , has lowest variance ) . At each node j , we learn a classifier that embodies an latent action to split the model space Ωj . The linear classifier takes the portion of the dataset that falls into its own region Dt ∩ Ωj , The performance of a region Ωj , V ( Ωj ) , is estimated by 1/N ∑ i∈Dt∩Ωj vi . To minimize the variance of V ( Ωk ) for all child nodes , we learn a linear regressor fj that minimizes ∑ i∈Dt∩Ωj ( fj ( ai ) − vi ) 2 . Once learned , the parameters of fj and V̂ ( Ωj ) form a linear constraint that bifurcates Ωj into a good region ( > V̂ ( Ωj ) ) and a bad region ( ≤ V̂ ( Ωj ) ) for sampling . For convenience , the left child always represents the good region . The partition threshold V̂ ( Ωj ) , combined with fj , forms two latent actions at node j , going left ( fj ( ai ) > V̂ ( Ωj ) ) and going right ( fj ( ai ) ≤ V̂ ( Ωj ) ) . Note that we need to initialize each node classifier properly with a few random samples to establish initial boundary in the search space . An ablation study on the number of samples for initialization is provided in Fig . 7d . 3.2 SEARCH PHASE . Once actions are learned , the search phase follows . The search uses the learned actions to sample more architectures ai as well as their performance vi , and store ( ai , vi ) in dataset Dt to refine the action space in the next learning phase . Note that in the search phase , the tree structure and the parameters of those classifiers are fixed and static . Instead , the search phase learns to decide which region Ωj on tree leaves to sample ai , with a proper balance between the exploration of less known Ωj and exploitation of promising Ωj . Following the construction of the classifier at each node , a trivial go-left greedy based search strategy can be used to exclusively exploit the most promising Ωk defined by the current action space . However , this is not a good strategy since it only exploits the current action space , which is learned from the current samples and may not be optimal . There can be good model regions that are hidden in the right ( or bad ) leaves that need to be explored . In order to overcome this issue , we use Monte Carlo Tree Search ( MCTS ) as the search method , which has the characteristics of adaptive exploration , and has shown superior efficiency in various tasks . MCTS keeps track of visiting statistics at each node to achieve the balance between exploitation of existing good region and exploration of new region . In lieu of MCTS , our search phase also has select , sampling and backpropagate stages . LaNAS skips the expansion stage in regular MCTS since the connectivity of our search tree is static . When the action space is updated , previous sampled networks and their performance metrics in Dt are reused and redirected to ( maybe different ) nodes in initializing visitation counts n ( s ) and node values v ( s ) . When there is no learned action space , we random sample the model space to get jump started . Algorithm 1 : LaNAS search procedure . Data : specifications of the search domain 1 Function get_ucb ( X̄next , nnext , ncurt ) 2 c = 0.5 3 return X̄next nnext + 2c √ 2log ( ncurt ) nnext if nnext 6= 0 else +∞ 4 begin 5 while acc < target do 6 for n ∈ Tree.N do 7 n.g.train ( ) 8 for i = 1→ # selects do 9 leaf , path = ucb_select ( root ) 10 constraints = get_constraints ( path ) 11 network = sampling ( constraints ) 12 acc = network.train ( ) 13 back − propagate ( network , acc ) Algorithm 2 : Subroutines in Alg . 1 . 1 Function get_constraints ( s_path ) 2 constraints = [ ] 3 for node ∈ s_path do 4 W , b = node.g.params ( ) , X̄ = node.X̄ 5 if node on left then constraints.add ( Wa + b ≥ X̄ ) else constraints.add ( Wa + b < X̄ ) 6 return constraints 7 Function ucb_select ( c = root ) 8 path = [ ] 9 while c not leaf do 10 path.add ( c ) 11 lucb = get_ucb ( c.left.X̄ , c.left.n , c.n ) , rucb = get_ucb ( c.right.X̄ , c.right.n , c.n ) 12 if lucb > rucb then c = c.left else c = c.right 13 return path , c

Summary: This paper introduces a Neural Architecture Search algorithm that attempts to solve the problems of the existing NAS only utilizing manually designed action space (not related to the performance). The paper proposes LaNAS which is based on an MCTS algorithm to partition the search space into tree nodes by the performance in the tree structure. The performance of the method is shown in the NASBench-101 dataset and Cifar-10 open domain search.