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: Authors revise the one-shot NAS algorithm in this work. One-shot NAS that employs a supernet to share the weights between subnets is an efficient NAS algorithm. Authors develop a new training paradigm to train the supernet sufficiently. Specifically, they uniformly sample a single path from supernet at each iteration to make the training effective and stable.