SINGLE PATH ONE-SHOT NEURAL ARCHITECTURE SEARCH WITH UNIFORM SAMPLING

1 INTRODUCTION . Deep learning automates feature engineering and solves the weight optimization problem . Neural Architecture Search ( NAS ) aims to automate architecture engineering by solving one more problem , architecture design . Early NAS approaches ( Zoph et al. , 2018 ; Zhong et al. , 2018a ; b ; Liu et al. , 2018a ; Real et al. , 2018 ; Tan et al. , 2018 ) solves the two problems in a nested manner . A large number of architectures are sampled and trained from scratch . The computation cost is unaffordable on large datasets . Recent approaches ( Wu et al. , 2018a ; Cai et al. , 2018 ; Liu et al. , 2018b ; Xie et al. , 2018 ; Pham et al. , 2018 ; Zhang et al. , 2018c ; Brock et al. , 2017 ; Bender et al. , 2018 ) adopt a weight sharing strategy to reduce the computation . A supernet subsuming all architectures is trained only once . Each architecture inherits its weights from the supernet . Only fine-tuning is performed . The computation cost is greatly reduced . Most weight sharing approaches use a continuous relaxation to parameterize the search space ( Wu et al. , 2018a ; Cai et al. , 2018 ; Liu et al. , 2018b ; Xie et al. , 2018 ; Zhang et al. , 2018c ) . The architecture distribution parameters are jointly optimized during the supernet training via gradient based methods . The best architecture is sampled from the distribution after optimization . There are two issues in this formulation . First , the weights in the supernet are deeply coupled . It is unclear why inherited weights for a specific architecture are still effective . Second , joint optimization introduces further coupling between the architecture parameters and supernet weights . The greedy nature of the gradient based methods inevitably introduces bias during optimization and could easily mislead the architecture search . They adopted complex optimization techniques to alleviate the problem . The one-shot paradigm ( Brock et al. , 2017 ; Bender et al. , 2018 ) alleviates the second issue . It defines the supernet and performs weight inheritance in a similar way . However , there is no architecture relaxation . The architecture search problem is decoupled from the supernet training and addressed in a separate step . Thus , it is sequential . It combines the merits of both nested and joint optimization approaches above . The architecture search is both efficient and flexible . The first issue is still problematic . Existing one-shot approaches ( Brock et al. , 2017 ; Bender et al. , 2018 ) still have coupled weights in the supernet . Their optimization is complicated and involves sensitive hyper parameters . They have not shown competitive results on large datasets . This work revisits the one-shot paradigm and presents a new approach that further eases the training and enhances architecture search . Based on the observation that the accuracy of an architecture using inherited weights should be predictive for the accuracy using optimized weights , we propose that the supernet training should be stochastic . All architectures have their weights optimized simultaneously . This gives rise to a uniform sampling strategy . To reduce the weight coupling in the supernet , a simple search space that consists of single path architectures is proposed . The training is hyperparameter-free and easy to converge . This work makes the following contributions . 1 . We present a principled analysis and point out drawbacks in existing NAS approaches that use nested and joint optimization . Consequently , we hope this work will renew interest in the one-shot paradigm , which combines the merits of both via sequential optimization . 2 . We present a single path one-shot approach with uniform sampling . It overcomes the drawbacks of existing one-shot approaches . Its simplicity enables a rich search space , including novel designs for channel size and bit width , all addressed in a unified manner . Architecture search is efficient and flexible . Evolutionary algorithm is used to support real world constraints easily , such as low latency . Comprehensive ablation experiments and comparison to previous works on a large dataset ( ImageNet ) verify that the proposed approach is state-of-the-art in terms of accuracy , memory consumption , training time , architecture search efficiency and flexibility . 2 REVIEW OF NAS APPROACHES . Without loss of generality , the architecture search spaceA is represented by a directed acyclic graph ( DAG ) . A network architecture is a subgraph a ∈ A , denoted as N ( a , w ) with weights w. Neural architecture search aims to solve two related problems . The first is weight optimization , wa = arg min w Ltrain ( N ( a , w ) ) , ( 1 ) where Ltrain ( · ) is the loss function on the training set . The second is architecture optimization . It finds the architecture that is trained on the training set and has the best accuracy on the validation set , as a∗ = arg max a∈A ACCval ( N ( a , wa ) ) , ( 2 ) where ACCval ( · ) is the accuracy on the validation set . Early NAS approaches perform the two optimization problems in a nested manner ( Zoph & Le , 2016 ; Zoph et al. , 2018 ; Zhong et al. , 2018a ; b ; Baker et al. , 2016 ) . Numerous architectures are sampled from A and trained from scratch as in Eq . ( 1 ) . Each training is expensive . Only small dataset ( e.g. , CIFAR 10 ) and small search space ( e.g , a single block ) are affordable . Recent NAS approaches adopt a weight sharing strategy ( Cai et al. , 2018 ; Liu et al. , 2018b ; Wu et al. , 2018a ; Xie et al. , 2018 ; Bender et al. , 2018 ; Brock et al. , 2017 ; Zhang et al. , 2018c ; Pham et al. , 2018 ) . The architecture search space A is encoded in a supernet1 , denoted as N ( A , W ) , where W is the weights in the supernet . The supernet is trained once . All architectures inherit their weights directly from W . Thus , they share the weights in their common graph nodes . Fine tuning of an architecture is performed in need , but no training from scratch is incurred . Therefore , architecture search is fast and suitable for large datasets like ImageNet . Most weight sharing approaches convert the discrete architecture search space into a continuous one ( Wu et al. , 2018a ; Cai et al. , 2018 ; Liu et al. , 2018b ; Xie et al. , 2018 ; Zhang et al. , 2018c ) . Formally , space A is relaxed to A ( θ ) , where θ denotes the continuous parameters that represent the distribution of the architectures in the space . Note that the new space subsumes the original one , A ⊆ A ( θ ) . An architecture sampled from A ( θ ) could be invalid in A . 1 “ Supernet ” is used as a general concept in this work . It has different names and implementation in previous approaches . An advantage of the continuous search space is that gradient based methods ( Liu et al. , 2018b ; Cai et al. , 2018 ; Wu et al. , 2018a ; Véniat & Denoyer , 2018 ; Xie et al. , 2018 ; Zhang et al. , 2018c ) is feasible . Both weights and architecture distribution parameters are jointly optimized , as ( θ∗ , Wθ∗ ) = arg min θ , W Ltrain ( N ( A ( θ ) , W ) ) . ( 3 ) After optimization , the best architecture a∗ is sampled from A ( θ∗ ) . Note that it could be invalid in A . If so , it is validated ( e.g. , by binarization of θ ( Liu et al. , 2018b ) ) . It then inherits the weights from Wθ∗ and is fine-tuned . Optimization of Eq . ( 3 ) is challenging . First , the weights of the graph nodes in the supernet depend on each other and become deeply coupled during optimization . For a specific architecture , it inherits certain node weights from W . While these weights are decoupled from the others , it is unclear why they are still effective . Second , joint optimization of architecture parameter θ and weights W introduces further complexity . Solving Eq . ( 3 ) inevitably introduces bias to certain areas in θ and certain nodes in W during the progress of optimization . The bias would leave some nodes in the graph well trained and others poorly trained . With different level of maturity in the weights , different architectures are actually non-comparable . However , their prediction accuracy is used as guidance for sampling inA ( θ ) ( e.g. , used as reward in policy gradient ( Cai et al. , 2018 ) ) . This would further mislead the architecture sampling . This problem is in analogy to the “ dilemma of exploitation and exploration ” problem in reinforcement learning . To alleviate such problems , existing approaches adopt complicated optimization techniques ( see Table 7 for a summary ) . Nevertheless , there lacks a comprehensive evaluation of their effectiveness ( Li & Talwalkar , 2019 ) . Task constraints Real world tasks usually have additional requirements on a network ’ s memory consumption , FLOPs , latency , energy consumption , etc . These requirements only depends on the architecture a , not on the weights wa . Thus , they are called architecture constraints in this work . A typical constraint is that the network ’ s latency is no more than a preset budget , as Latency ( a∗ ) ≤ Latmax . ( 4 ) Note that it is challenging to satisfy Eq . ( 2 ) and Eq . ( 4 ) simultaneously for most previous approaches . Some works augment the loss function Ltrain in Eq . ( 3 ) with soft loss terms that consider the architecture latency ( Cai et al. , 2018 ; Wu et al. , 2018a ; Xie et al. , 2018 ; Véniat & Denoyer , 2018 ) . However , it is hard , if not impossible , to guarantee a hard constraint like Eq . ( 4 ) . 3 OUR SINGLE PATH ONE-SHOT APPROACH . As analyzed above , the coupling between architecture parameters and weights is problematic . This is caused by joint optimization of both . To alleviate the problem , a natural solution is to decouple the super net training and architecture search in two sequential steps . This leads to the so called one-shot approaches ( Brock et al. , 2017 ; Bender et al. , 2018 ) . In general , the two steps are formulated as follows . Firstly , the supernet weight is optimized as WA = arg min W Ltrain ( N ( A , W ) ) . ( 5 ) Compared to Eq . ( 3 ) , the continuous parameterization of search space is absent . Only weights are optimized . Secondly , architecture searched is performed as a∗ = arg max a∈A ACCval ( N ( a , WA ( a ) ) ) . ( 6 ) During search , each sampled architecture a inherits its weights from WA as WA ( a ) . The key difference of Eq . ( 6 ) from Eq . ( 1 ) and ( 2 ) is that architecture weights are ready for use . Evaluation of ACCval ( · ) only requires inference . Thus , the search is very efficient . The search is also flexible . Any adequate search algorithm is feasible . The architecture constraint like Eq . ( 4 ) can be exactly satisfied . Search can be repeated many times on the same supernet once trained , using different constraints ( e.g. , 100ms latency and 200ms latency ) . These properties are absent in previous approaches . These make the one-shot paradigm attractive for real world tasks . One problem in Sec . 2 still remains . The graph nodes ’ weights in the supernet training in Eq . ( 5 ) are coupled . It is unclear why the inherited weights WA ( a ) are still good for an arbitrary architecture a . The recent one-shot approach ( Bender et al. , 2018 ) attempts to decouple the weights using a “ path dropout ” strategy . During an SGD step in Eq . ( 5 ) , each edge in the supernet graph is randomly dropped . The random chance is controlled via a dropout rate parameter . In this way , the coadaptation of the node weights is reduced during training . Experiments in ( Bender et al. , 2018 ) indicate that the training is very sensitive to the dropout rate parameter . This makes the supernet training hard . A carefully tuned heat-up strategy is used . In our implementation of this work , we also found that the validation accuracy is very sensitive to the dropout rate parameter . Single Path Supernet and Uniform Sampling . Let us restart to think about the fundamental principle behind the idea of weight sharing . The key to the success of architecture search in Eq . ( 6 ) is that , the accuracy of any architecture a on a validation set using inherited weightWA ( a ) ( without extra fine tuning ) is highly predictive for the accuracy of a that is fully trained . Ideally , this requires that the weight WA ( a ) to approximate the optimal weight wa as in Eq . ( 1 ) . The quality of the approximation depends on how well the training loss Ltrain ( N ( a , WA ( a ) ) ) is minimized . This gives rise to the principle that the supernet weightsWA should be optimized in a way that all architectures in the search space are optimized simultaneously . This is expressed as WA = arg min W Ea∼Γ ( A ) [ Ltrain ( N ( a , W ( a ) ) ) ] , ( 7 ) where Γ ( A ) is a prior distribution of a ∈ A . Note that Eq . ( 7 ) is an implementation of Eq . ( 5 ) . In each step of optimization , an architecture a is randomly sampled . Only weights W ( a ) are activated and updated . So the memory usage is efficient . In this sense , the supernet is no longer a valid network . It behaves as a stochastic supernet ( Véniat & Denoyer , 2018 ) . This is different from ( Bender et al. , 2018 ) . To reduce the co-adaptation between node weights , we propose a supernet structure that each architecture is a single path , as shown in Fig.3 . Compared to the path dropout strategy in ( Bender et al. , 2018 ) , the single path strategy is hyperparameter-free . We compared the two strategies within the same search space ( as in this work ) . Note that the original drop path in ( Bender et al. , 2018 ) may drop all operations in a block , resulting in a short cut of identity connection . In our implementation , it is forced that one random path is kept in this case since our choice block does not have an identity branch . We randomly select sub network and evaluate its validation accuracy during the training stage . Results in Fig.1 show that drop rate parameters matters a lot . Our single path strategy corresponds to using drop rate 1 . It works the best , which also verifies the benefits of weight decoupling by our single path strategy . The prior distribution Γ ( A ) is important . In this work , we empirically find that uniform sampling is good . This is not much of a surprise . A recent work also finds that purely random search is competitive to several SOTA NAS approaches ( Li & Talwalkar ( 2019 ) ) . We also experimented with a variant that samples the architectures uniformly according to their constraints , named uniform constraint sampling . Specifically , we randomly choose a range , and then sample the architecture repeatedly until the FLOPs of sampled architecture falls in the range . This is because a real task usually expects to find multiple architectures satisfying different constraints . In this work , we find the uniform constraint sampling method is slightly better . So we use it by default in this paper . We note that sampling a path according to architecture distribution during optimization is already used in previous weight sharing approaches ( Pham et al. , 2018 ; Véniat & Denoyer , 2018 ; Wu et al. , 2018a ; Cai et al. , 2018 ; Xie et al. , 2018 ; Zhang et al. , 2018c ; Yao et al. , 2019 ; Dong & Yang , 2019 ; Stamoulis et al. , 2019 ) . The difference is that , the distribution Γ ( A ) is a fixed priori during our training ( Eq . ( 7 ) ) , while it is learnable and updated ( Eq . ( 3 ) ) in previous approaches ( e.g . RL ( Pham et al. , 2018 ) , policy gradient ( Véniat & Denoyer , 2018 ; Cai et al. , 2018 ) , Gumbel Softmax ( Wu et al. , 2018a ; Xie et al. , 2018 ) , APG ( Zhang et al. , 2018c ) ) . As analyzed in Sec . 2 , the latter makes the supernet weights and architecture parameters highly correlated and optimization difficult . Comprehensive experiments in Sec . 4 show that our approach achieves better results than the SOTA methods . Note that there is no such theoretical guarantee that using a fixed prior distribution is inherently better than optimizing the distribution during training . Our better result likely indicates that the joint optimization in Eq . ( 3 ) is too difficult for the existing optimization techniques . Supernet Architecture and Novel Choice Block Design . Choice blocks are used to build a stochastic architecture . Fig . 3 illustrates an example case . A choice block consists of multiple architecture choices . For our single path supernet , each choice block only has one choice invoked at the same time . A path is obtained by sampling all the choice blocks . The simplicity of our approach enables us to define different types of choice blocks to search various architecture variables . Specifically , we propose two novel choice blocks to support complex search spaces . Channel Number Search . We propose a new choice block based on weight sharing , as shown in Fig . 4 . The main idea is to preallocate a weight tensor with maximum number of channels , and the system randomly selects the channel number and slices out the corresponding subtensor for convolution . With the weight sharing strategy , we found that the supernet can converge quickly . In detail , assume the dimensions of preallocated weights are ( max c out , max c in , ksize ) . For each batch in supernet training , the number of current output channels c out is randomly sampled . Then , we slice out the weights for current batch with the form Weights [ : c out , : c in , : ] , which is used to produce the output . The optimal number of channels is determined in the search step . Mixed-Precision Quantization Search . In this work , We design a novel choice block to search the bit widths of the weights and feature maps , as shown in Fig . 5 . We also combine the channel search space discussed earlier to our mixed-precision quantization search space . During supernet training , for each choice block feature bit width and weight bit width are randomly sampled . They are determined in the evolutionary step . See Sec . 4 and Fig . 5 for details . Evolutionary Architecture Search . For architecture search in Eq . ( 6 ) , previous one-shot works ( Brock et al. , 2017 ; Bender et al. , 2018 ) use random search . This is not effective for a large search space . This work uses an evolutionary algorithm . Note that evolutionary search in NAS is used in ( Real et al. , 2018 ) , but it is costly as each architecture is trained from scratch . In our search , each architecture only performs inference . This is very efficient . The algorithm is elaborated in Algorithm 1 . For all experiments , population size P = 50 , max iterations T = 20 and k = 10 . For crossover , two randomly selected candidates are crossed to produce a new one . For mutation , a randomly selected candidate mutates its every choice block with probability 0.1 to produce a new candidate . Crossover and mutation are repeated to generate enough new candidates that meet the given architecture constraints . Before the inference of an architecture , the statistics of all the Batch Normalization ( BN ) ( Ioffe & Szegedy , 2015 ) operations are recalculated on a random subset of training data ( 20000 images on ImageNet ) . It takes a few seconds . This is because the BN statistics from the supernet are usually not applicable to the candidate nets . This is also referred in ( Bender et al. , 2018 ) . Fig . 2 plots the validation accuracy over generations , using both evolutionary and random search methods . It is clear that evolutionary search is more effective . Experiment details are in Sec . 4 . The evolutionary algorithm is flexible in dealing with different constraints in Eq . ( 4 ) , because the mutation and crossover processes can be directly controlled to generate proper candidates to satisfy the constraints . Previous RL-based ( Tan et al. , 2018 ) and gradient-based ( Cai et al. , 2018 ; Wu et al. , 2018a ; Véniat & Denoyer , 2018 ) methods design tricky rewards or loss functions to deal with such constraints . For example , ( Wu et al. , 2018a ) uses a loss function CE ( a , wa ) · α log ( LAT ( a ) ) β to balance the accuracy and the latency . It is hard to tune the hyper parameter β to satisfy a hard constraint like Eq . ( 4 ) . Summary . The combination of single path supernet , uniform sampling training strategy , evolutionary architecture search , and rich search space design makes our approach simple , efficient and flexible . Table 7 in Appendix performs a comprehensive comparison of our approach against previous weight sharing approaches on various aspects . Ours is the easiest to train , occupies the smallest memory , best satisfies the architecture ( latency ) constraint , and easily supports large datasets . Extensive results in Sec . 4 verify that our approach is the state-of-the-art .

Summary: This paper presents a new one-shot NAS approach. Parameter updating and structure updating are optimized separately. In the process of parameter optimization, different from the previous methods, the network samples the structure according to the uniform probability distribution. This sampling method can avoid the coupling caused by optimizing the structural parameters at the same time. After training the supernet, the network uses the genetic algorithm to search the structure.