S INGLE S HOT N EURAL A RCHITECTURE S EARCH VIA D IRECT S PARSE O PTIMIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Recently Neural Architecture Search (NAS) has aroused great interest in both academia and industry, however it remains challenging because of its huge and non-continuous search space.
Instead of applying evolutionary algorithm or reinforcement learning as previous works, this paper proposes a Direct Sparse Optimization NAS (DSO-NAS) method.
In DSO-NAS, we provide a novel model pruning view to NAS problem.
Speciﬁcally, we start from a completely connected block, and then introduce scaling factors to scale the information ﬂow between operations.
Next, we impose sparse regularizations to prune useless connections in the architecture.
Lastly, we derive an efﬁcient and theoretically sound optimization method to solve it.
Our method enjoys both advantages of differentiability and efﬁciency, therefore can be directly applied to large datasets like ImageNet.
Particularly, on CIFAR-10 dataset, DSO-NAS achieves an average test error 2.84%, while on the ImageNet dataset DSO-NAS achieves 25.4% test error under 600M
FLOPs with 8 GPUs in 18 hours.
I NTRODUCTION
With no doubt, Deep Neural Networks (DNN) have been the engines for the AI renaissance in recent years.
Dating back to 2012, DNN based methods have refreshed the records for many AI applications, such as image classiﬁcation (Krizhevsky et al.
; Szegedy et al.
; He et al.
speech recognition (Hinton et al.
; Graves et al.
and Go Game
(Silver et al.
Considering its amazing representation power, DNNs have shifted the paradigm of these applications from manually designing the features and stagewise pipelines to end-to-end learning.
Although DNNs have liberated researchers from such feature engineering, another tedious work has emerged – “network engineering”.
In most cases, the neural networks need to be designed based on the speciﬁc tasks, which again leads to endless hyperparameters tuning and trails.
Therefore, designing a suitable neural network architecture still requires considerable amounts of expertise and experience.
To democratize the techniques, Neural Architecture Search (NAS) or more broadly, AutoML has been proposed.
There are mainly two streams for NAS:
The ﬁrst one is to follow the pioneering work Zoph & Le (2017), which proposed a reinforcement learning algorithm to train a Recurrent
Neural Network (RNN) controller that generates coded architectures (Zoph et al.
; Pham et al.
The second one is the evolutionary algorithm, which iteratively evaluates and proposes new models for evaluation (Real et al.
; Stanley & Miikkulainen (2002)).
Despite their impressive performance, the search processes are incredibly resource-hungry and unpractical for large datasets like ImageNet, though some acceleration methods have been proposed (Zhong et al.
; Pham et al.
Very recently, DARTS (Liu et al.
(2018b)) proposed a gradient-based method in which the connections are selected by a softmax classiﬁer.
Although DARTS achieves impressive performance with great acceleration, its search space is still limited to ﬁx-length coding and blocksharing search as in previous works.
In this work, we take another view to tackle these problems.
We reformulate NAS as pruning the useless connections from a large network which contains the complete network architecture hypothesis space.
Thus only one single model is trained and evaluated.
Since the network structure is directly optimized during training, we call our method Direct Sparse Optimization NAS (DSONAS).
We further demonstrate that this sparse regularized problem can be efﬁciently optimized by a modiﬁed accelerated proximal gradient method opposed to the inefﬁcient reinforcement learning or revolutionary search.
Notably, DSO-NAS is much simpler than the existing search methods as it uniﬁes the neural network weight learning and architecture search into one single optimization problem.
DSO-NAS does not need any controller (Zoph & Le (2017); Zoph et al.
; Pham et al.
or performance predictor (Liu et al.
(2018a)) or relaxation of the search space (Zoph & Le (2017);
Zoph et al.
; Pham et al.
; Liu et al.
(2018b)).
As a result of the efﬁciency and simplicity, DSO-NAS ﬁrst demonstrate that NAS can be directly applied to large datasets like
ImageNet with no block structure sharing.
Our experiments show that DSO-NAS can achieve 2.84% average test error on CIFAR-10, as well as top-1 error 25.4% on ImageNet with FLOPs (the number of multiply-adds) under 600M.
In summary, our contributions can be summarized as follows:
• We propose a novel model pruning formulation for neural architecture search based on sparse optimization.
Only one model needs to be trained during the search.
• We propose a theoretically sound optimization method to solve this challenging optimization problem both effectively and efﬁciently.
• We demonstrate the results of our proposed method are competitive or better than other NAS methods, while signiﬁcantly simplifying and accelerating the search process.
R ELATED W ORKS
In this section, we brieﬂy review two research ﬁelds that may be related to our work.
2. N ETWORK P RUNING
Network pruning is a widely used technique for model acceleration and compression.
The early works of pruning focus on removing unimportant connections (LeCun et al.
(1990); Hassibi & Stork (1993)
; Han et al.
; Guo et al.
Though connection level pruning can yield effective compression, it is hard to harvest actual computational savings because modern GPU cannot utilize the irregular weights well.
To tackle this issue, a signiﬁcant amount of works on structure pruning have been proposed.
For neuron level pruning, several works prune the neurons directly by evaluating the importance of neurons based on speciﬁc criteria (Hu et al.
; Li et al.
(2017); Mariet & Sra (2016) ; Liu et al.
More generally, Wen et al.
(2016) proposed the sparse structure learning.
They adopted group sparsity on multiple structures of networks, including ﬁlter shapes, channels and layers.
Recently, Huang & Wang (2018) proposed a simpler way for structure pruning.
They introduced scaling factors to the outputs of speciﬁc structures (neural, groups or block) and apply sparse regularizations on them.
After training, structures with zero scaling factors can be safely removed.
Compared with (Wen et al.
, the proposed method is more effective and stable.
In this work, we extend (Huang & Wang (2018)) into a more general and harder case, neural architecture search.
2. N EURAL A RCHITECTURE S EARCH
Recently, there has been growing interest in developing methods to generate neural network architecture automatically.
One heavily investigated direction is evolutionary algorithm (Meyer-Lee et al.
; Miller et al.
Real et al.
; Stanley & Miikkulainen (2002)).
They designed the modiﬁcations like inserting layers, changing ﬁlter sizes or adding identity mapping as the mutations in evolution.
Not surprisingly, their methods are usually computationally intensive and less practical in large scale.
Another popular direction is to utilize reinforcement learning with an RNN agent to design the network architecture.
The pioneering work (Zoph & Le (2017)) applies an RNN network as the controller to sequentially decide the type and parameters of layers.
Then the controller is trained by RL with the reward designed as the accuracy of the proposed model.
Although it achieves remarkable results, it needs 800 GPUs to get such results, which is not affordable for broad applications.
Based on this work, several methods have been proposed to accelerate the search process by limiting the search space (Zoph et al.
(2018)), early stopping with performance prediction (Zhong et al.
(2018)), progressive search (Liu et al.
(2018a))
or weight sharing (Pham et al.
Despite their success, the aforementioned methods treat the search of network architecture as a black-box optimization problem. Moreover, the search spaces of them are limited due to the ﬁxed-length coding of architecture.
More recently, Ashok et al. ; He et al. (2018) integrated reinforcement learning into model compression. For efﬁcient search, the action space is constrained as pruning ratio in He et al.
Though they only applied this method on channel selection, it also can be seen as an architecture search approach.
Similarly, V´eniat & Denoyer (2018) also constrained the search space by predeﬁning it as a super-network, which is composed of a set of layers connected together.
In order to learn budget-aware network structure, they proposed a budget aware objective function and solved it by policy gradient inspired algorithm.
Though the building of the search space is similar to our method, the network structure is limited to ResNet fabric and convolutional neural fabrics.
Our most related work is a gradient based method DARTS (Liu et al.
(2018b)).
In DARTS, a special parameter a is applied on every connection and updated during training process.
A Softmax classiﬁer is then applied to select the connection to be used for nodes.
However, the search space of DARTS is also limited: every operation can only have exact two inputs; the number of nodes are ﬁxed within a block.
P ROPOSED M ETHOD
In this section, we will elaborate the details of our proposed method.
We will start with the intuition and motivations, then followed by the design of search space and the formulation of our method.
Lastly, we will describe the optimization and training details.
3. M OTIVATIONS
The idea of DSO-NAS follows the observation that the architecture space of neural network (or a micro structure in it) can be represented by a completely connected Directed Acyclic Graph (DAG).
Any other architecture in this space can be represented by a sub-graph of it.
In other words, a speciﬁc architecture can be obtained by selecting a subset of edges and nodes in the full graph.
Prior works (Zoph & Le (2017), Liu et al.
(2018a), Liu et al.
(2018b)) focus on searching the architecture of two types of blocks, convolution block and reduction block.
Following the idea of micro structure searching, we adopt the complete graph to represent the search space of an individual block.
Then the ﬁnal network architecture can be represented by a stacking of blocks with residual connections.
Fig.
1 illustrates an example DAG of a speciﬁc block, whose nodes and edges represent local computation O and information ﬂow, respectively.
[-O (5) ( (cid:80) 
-]
For a DAG with T nodes, the output of i th node h
( i ) can be calculated by transforming the sum of all the outputs of the predecessors, h ( j ) , j < i , by the local operation
O ( i ) , namely:
[Equation]
Then the structure search problem can be reformulated as an edge pruning problem.
In the search procedure, we remove useless edges and nodes in the full DAG, leaving the most important structures.
To achieve this goal, we apply scaling factors on every edge to scale the output of each node.
Then Eqn.
1 can be modiﬁed to:
[Equation]
[Figure]
where λ ( i ) ( j ) is the scaling factor applied on the information ﬂow from node j to i .
Then we apply sparse regularizations on the scaling parameters to force some of them to be zero in the search.
Intuitively, if λ ( i ) ( j ) is zero, the corresponding edge can be removed safely and isolated nodes can also be pruned as no contribution is made.
3. S EARCH S PACE
DSO-NAS can search the structure of each building block in DNN, and then share it for all the blocks in the DNN, just as all previous works did.
It can also directly search the whole network structure without block structure sharing, while still keeping a competitive searching time.
In the following, we will discuss the search space of each individual block ﬁrst, and then specify the search space of macro-structure.
A block consists of M sequential levels which are composed of N different kinds of operations.
In each block, every operation has connections with all the operations in the former levels and the input of the block.
Also, the output of the block is connected with all operations in the block.
Then for each connection, we scale its output by a multiplier λ , and impose a sparse regularization on it.
After optimization, the ﬁnal architecture is generated by pruning the connections whose corresponding λ are zero and all isolated operations.
The procedure of block search is illustrated in Fig.
Formally, the output of the j -th operation in the i -th layer of the b -th block h ( b,i,j ) is computed as:
[Equation]
where O ( b,i,j ) is the transformation of the j -th operation in the i -th layer of the b -th block, λ
( i,j ) ( b,m,n ) is the scaling factor from node h ( b,m,n ) to h ( b,i,j ) , and O ( b − 1) is the output of the ( b − 1) -th block.
Here we denote h ( b, 0 , 0) =
O ( b − 1) as the input node and h ( b,M +1 , 0) =
O ( b ) as the output node of the b -th block, respectively.
The operation in the m -th layer may have ( m − 1) N + 1 inputs.
Note that the connections between operations and the output of block are also learnable.
The output of the b -th block O ( b ) is obtained by applying a reduction operation (concatenation followed by a convolution with kernel size 1 × 1) R to all the nodes that have contribution to the output:
[Equation]
where identity mapping is applied in case all operations are pruned.
The structure of whole network is shown in Fig.
3: a network consists of S stages with B convolution blocks in every stage.
Reduction block is located at the end of each stage except the last stage.
We try two search spaces: (1) the share search space where λ is shared among blocks.
(2) the full search space where λ in different blocks are updated independently.
and
• Separable convolution with kernel 3 ×• Separable convolution with kernel 5 ×• Average pooling with kernel 3 ×• Max pooling with kernel 3 ×
As for reduction block, we simply use convolutions with kernel size 1 × 1 and 3 × 3, and apply them with a stride of 2 to reduce the size of feature map and double the number of ﬁlters.
The output of reduction block is the sum of these two convolutions.
The task of searching blocks therefore reduces to learning λ on every edge, which can be formulated as:
[Equation]
where x i and y i are input data and label respectively, K denotes the number of training samples,
W represents the weights of network.
δ and γ represent the weights of regularizations, respectively.
3. O PTIMIZATION AND T RAINING
The sparse regularization of λ induces great difﬁculties in optimization, especially in the stochastic setting in DNN.
Though heuristic thresholding could work, the optimization is unstable and without theoretical analysis.
Fortunately, a recently proposed method Sparse Structure Selection (SSS) (Huang & Wang (2018)) solved this challenging problem by modifying a theoretically sound optimization method Accelerated Proximal Gradient (APG) method, it reformulates the original APG to avoid redundant forward and backward pass in calculating the gradients:
[Equation]
where t is the number of iterations, S η ( t ) γ represents the soft-threshold operator as S α ( z )
i = sign ( z i )( | z i | − α ) + , η ( t ) represents gradient step size and µ is the momentum.
In (Huang & Wang (2018)), the authors named it as APG-NAG.
The weights W and λ are updated using NAG and
APG-NAG jointly on the same training set.
However, APG-NAG cannot be directly applied in our algorithm since DNN usually overﬁts the training data in some degree.
Different from pruning, whose the search space is usually quite limited, the search space in NAS is much more diverse and huge.
If the structure is learned on such overﬁtting model, it will generalize badly on the test set.
To avoid this problem, we divide training data into two parts and update W and λ in two separate sets.
This conﬁguration guarantees that the λ (i.e. network structure) is learned on a different subset of training data which is not seen during the learning of W .
Therefore, the sample distribution in the structure learning is more similar to that in testing, which may lead to better performance.
3. I NCORPORATING D IFFERENT B UDGETS
Hand-crafted network usually incorporates many domain knowledge.
For example, as highlighted in (Ma et al.
(2018)), memory access may be the bottleneck for lightweight network on GPU because the use of separable convolution.
Our method can easily consider these priors in the search by adaptively adjusting γ for each connection.
The ﬁrst example is to balance the FLOPs for each block.
As indicated in (Jastrzebski et al.
(2018)), most intense changes of the main branch ﬂow of ResNet are concentrated after reduction block.
In our experiments, we empirically ﬁnd that the complexity of the block after each reduction block is much higher than the others’ if all γ are ﬁxed.
Consequently, to balance the FLOPs among different blocks, we adjust the regularization weight for λ , namely γ t at iteration t adaptively according to the FLOPs of the block:
[Equation]
where FLOPs block represents the FLOPs of the completely connected block and FLOPs t , which can be calculated based on λ , represents the FLOPs of the kept operations at iteration t for a block.
Using this simple strategy, we can smooth the distribution of FLOPs by penalizing λ according to the FLOPs of the block.
We call this method Adaptive FLOPs in the following.
The second example is to incorporate speciﬁc computation budget such as Memory Access Cost (MAC).
Similarly, the γ applied on the n -th operation in m -th level at iteration t is calculated by:
[Equation]
where MAC t ( m,n ) represents the MAC of the n -th operation in the m level, and MAC max represents the maximum MAC in the network.
Using this strategy, DSO-NAS can generate architectures with better performance under the same budget of MAC.
We call this method Adaptive MAC in the following.
E XPERIMENTS
In this section, we ﬁrst introduce the implementation details of our method, then followed by the results of classiﬁcation task on CIFAR-10 and ImageNet datasets.
At last, we analyze each design component of our method in detail.
4. I MPLEMENTATION D ETAILS
The pipeline of our method consists of three stages:
Training the completely connected network for several epochs to get a good weights initialization.
Searching network architecture from the pretrained model.
Re-training the ﬁnal architecture from scratch and evaluating on test dataset.
In the ﬁrst two stages, the scaling parameters in batch normalization layers are ﬁxed to one to prevent affecting the learning of λ .
After step two, we adjust the number of ﬁlters in each operation by a global width multiplier to satisfy the computation budget, and then train the network from scratch as done in (Pham et al.
For benchmark, we test our algorithm on two standard datasets, CIFAR-10 (Krizhevsky & Hinton (2009)) and ImageNet LSVRC 2012 (Russakovsky et al.
We denote the model searched with and without block structure sharing as DSO-NAS-share and DSO-NAS-full , respectively.
In each block, we set the number of levels M = 4 , the number of operations N = 4 as four kinds of operations are applied for both CIFAR and ImageNet experiments indicated in Sec. 3.2.
For the hyper-parameters of optimization algorithm and weight initialization, we follow the setting of (Huang & Wang (2018)).
All the experiments are conducted in MXNet (Chen et al. with
NVIDIA GTX 1080Ti GPUs.
We will release our codes if the paper is accepted.
4. CIFAR
The CIFAR-10 dataset consists of 50000 training images and 10000 testing images.
As described in Sec.3.3, we divide the training data into two parts: 25000 for training of weights and restfor structure.
During training, standard data pre-processing and augmentation techniques (Xie et al.
(2017)) are adopted.
The mini-batch size is 128 on 2 GPUs and weight decay is set to 3e-4.
Firstly, [Figure] we pre-train the full network for 120 epochs, and then search network architecture until convergence (about 120 epochs). Constant learning rate 0.1 is applied in both pre-training and search stages.
The network adopted in CIFAR-10 experiments consists of three stages, and each stage has eight convolution blocks and one reduction block.
Adaptive FLOPs (see section 3.3) is applied during the search.
After the search, we train the ﬁnal model from scratch with the same setting of (Pham et al.
The searched models are trained for 630 epochs with NAG, where we change the learning rate following the cosine scheduler with l max = 0 .
05 , l min = 0 .
{+001 , T 0 = 10 , T mul = 
+}
(Loshchilov &
Hutter (2017)).
Additional improvements including dropout (Srivastava et al.
with probability 0.6, cutout (DeVries & Taylor (2017)) with size 16, drop path (Larsson et al.
with probability 0.5, auxiliary towers located at the end of second stage (Lee et al.
with weight 0.4.
Table 1 shows the performance of our searched models, including DSO-NAS-full and DSO-NASshare, where ”c/o” represents evaluate model with cutout technique.
We report the mean and standard deviation of ﬁve independent runs.
Due to limited space, we only show the block structure of DSO-NAS-share in Fig.4(a).
We also compare the simplest yet still effective baseline – random structure, both our DSO-NAS-share and DSO-NAS-full yield much better performance with less parameters.
Comparing with other state-of-the-art methods, our method demonstrates competitive results with similar or less parameters while costing only one GPU day. [Table]
4. ILSVRC
In the ILSVRC 2012 experiments, we conduct data augmentation based on the publicly available implementation of ’fb.resnet’ 1 .
Since this dataset is much larger than CIFAR-10, the training dataset is divided into two parts: 4 / 5 for training weights and 1 / 5 for training structure.
In the pre-training stage, we train the whole network for 30 epochs with constant learning rate 0.1, weight decay 4 × 10 − 5 .
The mini-batch size is 256 on 8 GPUs.
The same setting is adopted in the search stage, which costs about 0.75 days with 8 GPUs.
After the search, we train the ﬁnal model from scratch using NAG for 240 epochs, with batch size 1024 on 8 GPUs.
We set weight decay to 4 × 10 − 5 and adopt linear-decay learning rate schedule (linearly decreased from 0.5 to 0).
Label smoothing (Szegedy et al.
and auxiliary loss (Lee et al.
(2015)) are used during training.
There are four stages in the ImageNet network, and the number of convolution blocks in these four stages is 2, 2, 13, 6, respectively.
We ﬁrst transfer the block structure searched on CIFAR-10.
We also directly search the network architecture on ImageNet.
The ﬁnal structure generated by DSO-NAS-share is shown in Fig.4(b).
The quantitative results for ImageNet are shown in Table 2, where results with * are obtained by transferring the generated CIFAR-10 blocks to ImageNet.
[Table]
MobileNet (Howard et al
[Table] [Table]
ShufﬂeNet-v1 2x (Zhang et al.
[Table] [Table]
It is notable that given similar computation budget
[Table]
DSO-NAS achieves competitive or better performance than other state-of-the-art methods with less search cost [Table] except for MnasNet whose search space is carefully designed and different from other NAS methods.
The block structure transferred from CIFAR-10 dataset also achieves impressive performance, proving the generalization capability of the searched architecture.
Moreover, directly searching on target dataset (ImageNet) brings additional improvements.
4. A BLATION S TUDY
In this section [Table] we present some ablation analyses on our method to illustrate the effectiveness and necessity of each component.
4.4. T HE E FFECTIVENESS OF B UDGET A WARE S EARCH
With adaptive FLOPs technique, the weight of sparse regularization for each block will be changed adaptively according to Eqn. 9.
We ﬁrst show the distribution of FLOPs among different blocks in Fig. 5(a).
This strategy can prevent some blocks from being pruned entirely as expected.
We also show the error rates of different settings in Fig. 5(b) and Fig. 5(c). [Figure] [Figure]
It is clear that the networks searched with adaptive FLOPs technique are consistently better than the ones without under the same total FLOPs or parameters.
DSO-NAS can also search for architecture based on certain computational target, such as MAC discussed in Sec.
The results are shown in Fig.
It is obvious to see that DSONAS can generate architecture with higher accuracy under certain MAC budget, proving the effectiveness of adaptive MAC technique.
The method can similarly be applied to optimize many other computation budgets of interest, which we leave for further study.
4.4. O THER F ACTORS F OR S EARCHING A RCHITECTURE
We conduct experiments on different settings of our proposed architecture search method to justify the need of each component we designed.
The results are shown in Table
[Figure] [Table] ”Pre-train model” means whether we conduct step one in Sec. 4, while ”Split training” means whether to split the whole training set into two sets for weight and structure learning separately.
The
Ratio of W&S means the ratio of training sample for weight learning and structure learning.
As for the ratio of x : y , we update weight for x times and update λ for y times for every x + y iterations.
Note that we only pre-train the model on the weight learning set.
It is notable that the use of a separate set for structure learning plays an important role to prevent overﬁtting training data, and improve the performance by 0.2%.
The ratio of these two sets has minor inﬂuence.
Besides, a good initialization of weight is also crucial as random initialization of weights may lead to another 0.2% drop on accuracy under that same parameter budgets.
C ONCLUSIONS AND F UTURE W ORK
Neural Architecture Search has been the core technology for realizing AutoML.
In this paper, we have proposed a Direct Sparse Optimization method for NAS.
Our method is appealing to both academic research and industrial practice in two aspects: First, our uniﬁed weight and structure learning method is fully differentiable in contrast to most previous works.
It provides a novel model pruning view to the NAS problem.
Second, the induced optimization method is both efﬁcient and effective.
We have demonstrated state-of-the-art performance on both CIFAR and ILSVRC image classiﬁcation datasets, with affordable cost (single machine in one day).
In the future we would like to incorporate hardware features for network co-design, since the actual running speed of the same network may highly vary across different hardware because of cache size, memory bandwidth, etc.
We believe our proposed DSO-NAS opens a new direction to pursue such objective.
It could push a further step to AutoML for everyone’s use