Scaling Up Neural Architecture Search with Big Single-Stage Models

Neural architecture search ( NAS ) methods have shown promising results discovering models that are both accurate and fast . For NAS , training a one-shot model has became a popular strategy to approximate the quality of multiple architectures ( child models ) using a single set of shared weights . To avoid performance degradation due to parameter sharing , most existing methods have a two-stage workflow where the best child model induced from the one-shot model has to be retrained or finetuned . In this work , we propose BigNAS , an approach that simplifies this workflow and scales up neural architecture search to target a wide range of model sizes simultaneously . We propose several techniques to bridge the gap between the distinct initialization and learning dynamics across small and big models with shared parameters , which enable us to train a single-stage model : a single model from which we can directly slice high-quality child models without retraining or finetuning . With BigNAS we are able to train a single set of shared weights on ImageNet and use these weights to obtain child models whose sizes range from 200 to 1000 MFLOPs . Our discovered model family , BigNASModels , achieve top-1 accuracies ranging from 76.5 % to 80.9 % , surpassing all state-of-the-art models in this range including EfficientNets . 1 INTRODUCTION . Designing network architectures that are both accurate and efficient is crucial for deep learning on edge devices . It is well known that a single neural network architecture can require more than an order of magnitude more inference time if it is deployed on a slower device ( Yu et al. , 2018 ) . This makes it appealing to not only search for architectures that are optimized for specific devices , but also to ensure that a range of models can be deployed effectively . In the past , Neural Architecture Search ( NAS ) methods ( Zoph & Le , 2016 ; Zoph et al. , 2018 ; Real et al. , 2018 ) have shown to be excellent at optimizing for a single device and latency target ( Tan et al. , 2019 ) . However , if we wish to target a large array of devices , it becomes prohibitively expensive and time-consuming to run a separate search for each one . A possible solution to this is to use basic scaling heuristics such as the EfficientNet family ( Tan & Le , 2019 ) . However , this loses out on opportunities to optimize models specialized for the diverse performance characteristics of individual devices . Another option would be to use efficient architecture search methods , e.g. , Pham et al . ( 2018 ) ; Bender et al . ( 2018 ) ; Liu et al . ( 2018b ) . However , to target multiple devices , we must run many searches and retrain all of the searched models from scratch . In this work , we search over a big single-stage model that contains both small child models ( ∼200 MFLOPs , comparable to MobileNetV3 ) and big child models ( ∼1 GFLOPs , comparable to EfficientNet B2 ) . Different from existing one-shot methods ( Bender et al. , 2018 ; Liu et al. , 2018b ; Brock et al. , 2018 ; Pham et al. , 2018 ) , our trained single-stage model offers a much wider coverage of model capacities , and more importantly , all child models are trained in a way such that they simultaneously reach excellent performance at the end of the search phase , without requiring a separate retraining step . Architecture selection can be then carried out via a simple coarse-to-fine selection strategy . Once an architecture is selected , we can obtain a child model by slicing the single-stage model for instant deployment w.r.t . the given constraints such as memory footprint and/or runtime latency . The workflow is illustrated in Figure 1 . The success of our method heavily relies on training a high-quality single-stage model , which is challenging on its own . For example , we find the training loss explodes if the single-stage model is not properly initialized , and bigger child models start to overfit before smaller ones plateau . It is particularly nontrivial to simultaneously retain good performance on every individual child model due to aggressive parameter sharing during architecture search . We address these challenges through a combination of techniques , including an improved sampling strategy and efficient inplace distillation , and substantially stabilize the single-stage model training through better initialization , learning rate schedule and regularization . Effectiveness of the proposed solutions is backed up by ablation studies . With the proposed techniques , we are able train a single-stage model on ImageNet and obtain a family of child models that simultaneously surpass all the state-of-the-art models in the range of 200 to 1000 MFLOPs , including EfficientNets B0-B2 ( 1.6 % more accurate under 400 MFLOPs ) , without retraining or finetuning the child models upon the completion of search . One of our child models achieves 80.9 % top-1 accuracy at 1G FLOPs ( four times less computation than a ResNet-50 ) . 2 RELATED WORK . Earlier NAS methods ( Zoph & Le , 2016 ; Zoph et al. , 2018 ; Liu et al. , 2017 ; 2018a ; Real et al. , 2018 ) train thousands of candidate architectures from scratch ( on a smaller proxy task ) and use their validation performance as the feedback to an algorithm that learns to focus on the most promising regions in the search space . More recent works have sought to amortize the cost by training a single over-parameterized one-shot model . Each architecture in the search space uses only a subset of the operations in the one-shot model ; these child models can be efficiently ranked by using the shared weights to estimate their relative accuracies ( Brock et al. , 2018 ; Pham et al. , 2018 ; Bender et al. , 2018 ; Liu et al. , 2018b ; Cai et al. , 2018 ; Wu et al. , 2019 ) . As a complementary direction , resource-aware NAS methods are proposed to simultaneously maximize prediction accuracy and minimize resource requirements such as latency , FLOPs , or memory footprints ( Tan et al. , 2019 ; Cai et al. , 2019 ; Wu et al. , 2019 ; Stamoulis et al. , 2019 ; Guo et al. , 2019 ; Yu & Huang , 2019a ) . All the aforementioned approaches require two-stage training : Once the best architectures have been identified ( either through the proxy tasks or using a one-shot model ) , they have to be retrained from scratch to obtain a final model with higher accuracy . In most of these existing works , a single search experiment only targets a single resource budget or a narrow range of resource budgets at a time . To alleviate these issues , Cai et al . ( 2019 ) proposed a progressive training approach ( OFA ) concurrently with our work . The idea is to pre-train a single full network and then progressively distill it to obtain the smaller networks . Moreover , a channel sorting procedure is required to progressively construct the smaller networks . In our proposed BigNAS , however , all the child models in the single-stage model are trained simultaneously , allowing the learning of small and big networks to mutually benefit each other . During the training , we always keep lower-index channels in each layer and lower-index layers in each stage for our child models , eliminating the sorting procedure . While OFA focuses on a limited range of model sizes ( for example , models around 300 MFLOPs ) , our BigNAS is able to handle a wider set of models ( from 200 MFLOPs to 1 GFLOPs ) and offers a better coverage over diverse deployment scenarios and varied resource budgets . Our work shares high-level similarities with slimmable networks ( Yu et al. , 2018 ; Yu & Huang , 2019b ; a ) in terms of training a single shared set of weights which can be used for many child models . However , while slimmable networks are specialized to vary the number of channels only , we are able to handle a much larger space where many architectural dimensions ( kernel and channel sizes , network depths , input resolutions ) are searched simultaneously , subsuming and outperforming the scaling heuristics in EfficientNets ( Tan & Le , 2019 ) . 3 ARCHITECTURE SEARCH WITH SINGLE-STAGE MODELS . Our proposed method consists of two steps : 1 . We train a big single-stage model from which we can directly sample or slice different architectures as child models for instant inference and deployment . In contrast to previous works ( Brock et al. , 2018 ; Pham et al. , 2018 ; Bender et al. , 2018 ; Liu et al. , 2018b ; Stamoulis et al. , 2019 ; Guo et al. , 2019 ) , our training is single-stage as it does not require finetuning the sampled architectures or retraining them from scratch at the end of search . 2 . Architecture selection using a simple coarse-to-fine selection method to find the most accurate model under the given resource constraints ( for example , FLOPs , memory footprint and/or runtime latency budgets on different devices ) . In the following , we will first systematically study how to train a high-quality single-stage model from five aspects : network sampling during training , inplace distillation , network initialization , convergence behavior and regularization . Then we will present a coarse-to-fine approach for efficient resource-aware architecture selection . 3.1 TRAINING A HIGH-QUALITY SINGLE-STAGE MODEL . Training a high-quality single-stage model is important and highly non-trivial due to the distinct initialization and learning dynamics of small and big child models . In this section , we first generalize two techniques originally introduced by Yu & Huang ( 2019b ) to simultaneously train a set of highquality networks with different channel numbers , and show that both can be extended to handle a much larger space where the architectural dimensions , including kernel sizes , channel numbers , input resolutions , network depths are jointly searched . We then present three additional techniques to address the distinct initialization and learning dynamics of small and big child models . Sandwich Rule In each training step , given a mini-batch of data , the sandwich rule ( Yu & Huang , 2019b ) samples the smallest child model , the biggest ( full ) child model and N randomly sampled child models ( N = 2 in our experiments ) . It then aggregates the gradients from all sampled child models before updating the weights of the single-stage model . As multiple architectural dimensions are included in our search space , the “ smallest ” child model is the one with lowest input resolution , thinnest width , shallowest depth , and smallest kernel size ( the kernel of the depthwise convolutions in each inverted residual block ( Sandler et al. , 2018 ) ) . The motivation is to improve all child models in our search space simultaneously , by pushing up both the performance lower bound ( the smallest child model ) and the performance upper bound ( the biggest child model ) across all child models . Inplace Distillation During the training of a single-stage model , inplace distillation ( Yu & Huang , 2019b ) takes the soft labels predicted by the biggest possible child model ( full model ) to supervise all other child models . The benefit of inplace distillation comes for free in our training setting , as we always have access to the predictions of the largest child model in each gradient update step thanks to the sandwich rule . We note that all child models are only trained with the inplace distillation loss , starting from the first training step to the end of the training . During training , input images are randomly cropped as a preliminary data augmentation step . When distilling a high-resolution teacher model into a low-resolution student model , we find that it is helpful to feed the same image patches into both the teacher and the student . In our data preparation , we first randomly crop an image with a fixed resolution ( on ImageNet we use 224 ) , and then apply bicubic interpolation to the same patch to transform it into all target resolutions ( e.g. , 192 , 288 , 320 ) . In this case , soft labels predicted by the biggest child model ( the teacher ) are more compatible with the inputs seen by other child models ( the students ) . Therefore this can serve as a more accurate distillation signal . Our preliminary results show that sampling different patches even in a same image leads to ∼ 0.3 % drop on top-1 accuracy for child models . Initialization Previous weight initialization methods , such as “ He Initialization ” ( He et al. , 2015 ) , are deduced from fixed neural networks where the number of input units n ( the fan-in ) is constant . The principal motivation of these initialization methods is to keep the variance of the responses in each layer unchanged , so that the forward information signals and the backward gradients will not be reduced or magnified exponentially as the network goes deeper . For example , He et al . ( 2015 ) suggested to initialize the variance of the weights as 2n for convolutions with ReLU activations . However , the above is ill-fitted for initializing a single-stage model , where n is no longer a constant across different child models with varied kernel sizes and input channels . This issue is exaggerated when we train bigger and deeper single-stage models . In practice , we find the training loss of a single-stage model explodes when we use the optimized learning rates for training a normal network . The training starts to work when we reduce the learning rate to 30 % , but it leads to much worse results ( ∼ 1.0 % top-1 accuracy on ImageNet ) . Identifying this issue is critical while the solution is quite simple . As all child models in our search space are residual networks , we initialize the output of each residual block ( before skip connection ) to an all-zeros tensor by setting the learnable scaling coefficient γ = 0 in the last Batch Normalization ( Ioffe & Szegedy , 2015 ) layer of each residual block , ensuring identical variance before and after each residual block regardless of the fan-in . This initialization is originally mentioned in ( Goyal et al. , 2017 ) which improves accuracy by ∼ 0.2 % in their setting , yet is more critical in our setting ( improving by ∼ 1.0 % ) due to the above analyzed initialization issue . We also additionally add a skip connection between each stage transitions when either resolutions or channels differ ( using 2 × 2 average pooling and/or 1 × 1 convolution if necessary ) to explicitly construct an identity mapping ( He et al. , 2016b ) . Convergence Behavior In practice , we find that big child models converge faster while small child models converge slower . Figure 2a shows the typical learning curves during the training of a single-stage model , where we plot the validation accuracies of a small and a big child model over time . This reveals a dilemma : at training step t when the performance of big child models peaks , the small child models are not fully-trained ; and at training step t′ when the small child models have better performance , the big child models already overfitted . To address this issue , we put our focus on the learning rate schedule . We first plot the optimized and widely used exponentially decaying learning rate schedule for MobileNet-series ( Howard et al. , 2017 ; Sandler et al. , 2018 ; Howard et al. , 2019 ) , MNasNets ( Tan et al. , 2019 ) and EfficientNets ( Tan & Le , 2019 ) in Figure 2b . We introduce a simple modification to this learning rate schedule , named exponentially decaying with constant ending , which has a constant learning rate at the end of training when it reaches 5 % of the initial learning rate ( Figure 2b ) . It brings two benefits . First , with a slightly larger learning rate at the end , the small child models learn faster . Second , the constant learning rate at the end alleviates the overfitting of big child models as the weights oscillate . Regularization Big child models tend to overfit the training data whereas small child models tend to underfit . In previous work , Bender et al . ( 2018 ) apply the same weight decay to all child models regardless whether they are small or big . In EfficientNet , Tan & Le ( 2019 ) linearly increase dropout ( Srivastava et al. , 2014 ) ratio as it moves from the smallest EfficientNet-B0 to the biggest EfficientNet-B7 . This becomes even more complicated in the context of training big single-stage models , due to the interplay among the small child models and big child models with shared parameters . Nevertheless , we introduce a simple rule that is surprisingly effective for this problem : regularize only the biggest ( full ) child model ( i.e. , the only model that has direct access to the ground truth training labels because of inplace distillation ) . We apply this rule to both weight decay and dropout , and empirically demonstrate its effectiveness in our experiments . Batch Norm Calibration Batch norm statistics are not accumulated when training the singlestage model as they are ill-defined with varying architectures . After the training is completed , we re-calibrate the batch norm statistics ( Yu & Huang , 2019b ) for each sampled child model for deployment without retraining or finetuning any network parameters .

Summary: The authors propose a search for neural architectures with different resource requirements by training a single model only. Furthermore, models found at the end of the search require no additional post-processing and are ready for deployment. A weight-sharing technique is applied to make this happen. The authors argue that there are multiple adaptions required to make it work. This includes the child model sampling strategy, use of model distillation, weight initialization, learning rate schedule, regularization and batch norm calibration.