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: This paper presents a method for architecture search in deep neural networks in order to identify scaled-down networks that can operate on resource limited hardware. The approach taken in this paper is different from other approaches, which train a single big model then fine tune smaller models for specific hardware, or use distillation to train progressively smaller models. Here, a single model is trained in a manner that allows subsequent slicing to smaller models without additional training required. The authors employ a variety of strategies to get this to work well, including specific initialization techniques, regularization methods, learning schedules, and a coarse-to-fine optimization method to obtain the smaller models. The authors demonstrate SotA performance on ImageNet relative to other techniques.