PredNAS: A Universal and Sample Efficient Neural Architecture Search Framework

To the best of our knowledge, [51] is the pioneering work of NAS. In [51], a recurrent network, also named controller, is used to propose new network architectures for evaluation. The controller is trained by reinforcement learning to maximize the performance of sampled networks on specific task. Though [51] showed the proposed method could design state-of-the-art models, the search cost was prohibitive. Several following works [52, 35, 27, 40] tried to mitigate this problem by reducing the search space. Different from searching the whole network on target dataset, [52] proposed to search stackable cells on proxy tasks. PNAS [35] adopted a progressive search approach to reduce the search cost. [40] focused on the design of search space and they discovered several principles to progressively simply the search space. Similarly, ABS [27] presented an angle-based metric to drop candidates during search.

To alleviate the high cost of evaluation of each sampled architecture in NAS, some other works [1, 39, 36, 22, 3, 2, 47] utilized the weight-sharing mechanism. The key idea is to allow the samples to share weights to reduce the training cost of each architecture. The first work in this spirit is ENAS [39]. Then one-shot methods [1, 22, 36] proposed to train a SuperNet capable of enumerating the child models in the search space for fast model evaluation. SPOS [22] adopted an uniform sampling strategy for the training of SuperNet and used evolutionary algorithm for architecture search. DARTS [36] and DSO-NAS [49] formulated NAS as an optimization problem and optimized the architecture parameters by gradient descent. However, SuperNet based methods may suffer from the problem of unreliable ranking [9] and performance collapse [48, 4, 5]. FairNAS [9] proposed a strict fairness sampling and training strategy to improve the ranking correlation of SPOS. [48] proposed to early stop the search process to handle the instability problem of DARTS. SmoothDARTS [4] found the rounding step of deriving the discrete architecture from continuous optimization in DARTS could introduce large performance drop. They proposed a perturbation-based regularization to smooth the loss landscape of DARTS.

In the context of hyperparameter optimization, early work adopted probabilistic models [15] or Bayesian neural networks [29] to estimate the performance of neural network. The subsequent work tried to predict the accuracy of neural networks by neural networks. Peephole [13] used an LSTM [25] network to integrate the information of different layers following the network topology and adopted a Multiple Layer Perceptron (MLP) to predict the accuracy of the input network at a specific training epoch. [45] proposed to use Graph Convolutional Networks (GCNs) [28] to represent the connections of network architecture. The network operations are represented as one-hot codes and the topology of the neural network is formulated by a adjacency matrix. FBNetV3 [11] adopted a simple MLP to search both architectures and training recipes jointly. Additionally, they proposed a pre-training approach to improve the sample efficiency of predictor. Recently, instead of improving the regression accuracy of predictor, ReNAS [46] and AceNAS [50] adopted rank based loss for reliable predictions. The closest work to our PredNAS is NAO[38], which adopted gradient to do architecture optimization. However, NAO utilized encoder-decoder framework to do transformation between architecture and network embedding. Hundreds of training samples and an extra structure reconstruction loss are needed for the training. In our framework, the design of network encoding and projection function allows us to do architecture embedding transformation without encoder and decoder. The training samples used in our work (30) in much less than NAO (600), and we show our PredNAS can directly search on large-scale tasks such as ImageNet without transferring the searched architecture from CIFAR to ImageNet.

Designing the structure of backbone network on image classification task and then transfer the ImageNet pretrained model to downstream tasks is a de facto paradigm in computation vision community. However, recent works show that this paradigm may be sub-optimal due to the gap between classification task and downstream applications. DetNAS [7] adopted the technique of one-shot SuperNet for object detection backbone search. Auto-DeepLab [34] constructed a hierarchical architecture search space for semantic segmentation task and adopted gradient-based method for search. SpineNet [18] proposed a scale-permuted model and showed the learned backbone could achieve better performance than regular scale-decreased models on both object detection and image classification tasks. Other works adopted NAS for component search, such as the architecture of feature pyramid [21] and prediction head [44] in object detection. In this paper, we adopted our PredNAS on a very efficient object detection framework, YOLOX [20], and show we can explore better architectures based on a large search space proposed by us.

Given the architecture search space $\mathbf{\Omega}$ and target resource budget $c$, NAS can be formulated as the following optimization problem:

where ${\mathcal{F}}({\bm{a}})$ is the performance indicator of architecture ${\bm{a}}$ on a specific dataset and ${\mathcal{C}}({\bm{a}})$ is the corresponding resource cost such as FLOPs, latency or parameters. For simplicity, we denote ${\bm{a}}$ as the encoding of its corresponding architecture. Following [43], we are more interested in finding multiple Pareto-optimal [12] solutions instead of a single architecture with highest accuracy. We construct a formulation of weighted sum of objectives to approximate Pareto optimal solutions:

where $\alpha$ is a tunable weight. Then the key problem becomes: 1) How could we approximate ${\mathcal{F}}(\cdot)$ and ${\mathcal{C}}(\cdot)$ with limited training data? 2) How could we optimize the above problem efficiently? Our answer is surprisingly simple: neural network with gradient descent¹¹1Though the formulation of ${\mathcal{C}}(\cdot)$ with FLOPs (or parameters) guided constraints can be derived directly, it is difficult to formulate latency or power guided computational complexity in a parametric form. For a general formulation, we adopt a predictor to approximate ${\mathcal{C}}(\cdot)$.. In this paper, we adopt two neural networks, the main predictor and the auxiliary predictor to approximate ${\mathcal{F}}(\cdot)$ and ${\mathcal{C}}(\cdot)$ respectively. Then the gradient of ${\bm{a}}$ respect to ${\mathcal{F}}({\bm{a}})$ and ${\mathcal{C}}({\bm{a}})$ can be computed by back-propagation easily for the optimization of ${\bm{a}}$.

Defining search space is the first step in NAS. Generally, the search space in NAS can be divided into two categories[16] as shown in Fig. 1(left): (1) Topology Search Space (TSS) adopted in DARTS[36] and NAS-Bench-201[17], concerning the connection topology and the associated operations on the connections; (2) Size Search Space (SSS) focuses on the width, depth or other parameters with the same topology [22, 40, 2]. Following [11] and [45], we adopt different network encoding methods to represent the architectures sampled from these two search space. For TSS, an adjacency matrix is used to represent the topology of architecture, and the operations on each connections are encoded as an operation matrix which consists of one-hot vectors. As for SSS, we simply concatenate the values of network depths, channel widths and other architecture parameters into a vector.

We train two predictors to approximate ${\mathcal{F}}(\cdot)$ and ${\mathcal{C}}(\cdot)$: (1) main predictor ${\mathcal{P}}_{m}$, estimating the performance of given architectures, e.g. accuracy or mean average precision (mAP); (2) auxiliary predictor ${\mathcal{P}}_{aux}$, predicting the resource cost of given architectures. These two predictors share the same network structure. For encoding of network structure, it highly depends on the search space of each dataset, so we leave it in the experiment section. For TSS, we adopt Graph Neural Network (GCN) [28] to handle graph-structured training samples following [45]. As for SSS, we use a simple MLP.

The top of Fig. 2 shows the training process. For simplicity, we consider FLOPs constraints and size search space for illustration in the following paper. We first randomly sample $N$ ($N<=50$) architectures from search space, and then collect the training samples as a dataset of $({\bm{a}},acc.,\text{FLOPs})$ triplets. The performance of architecture can be obtained via training the architecture on the target task or a proxy task to reduce the overhead. Following [11], we adopt the Huber loss to train both the main and auxiliary predictors.

In the search stage, most of the predictor based methods [45, 11, 50] adopt predictors only for architecture evaluation. They usually randomly sample a large number (about $10^{6}$) of architectures, and then take the networks with top-K performances to retrain. In this paper, we argue that these methods essentially utilize naive random shooting method for optimization, and ignore the valuable differentiable property of the predictor. Actually, we can update the sampled architecture according to the gradient of predictors. The bottom of Fig. 2 shows the search procedure of our PredNAS. Given an initial architecture ${\bm{a}}$, the main predictor ${\mathcal{P}}_{m}$ and the auxiliary predictor ${\mathcal{P}}_{aux}$, we update ${\bm{a}}$ iteratively by projected gradient ascent:

where $\eta$ is the learning rate, $t$ is the number of iterations, $\alpha$ is a trade-off between performance and resource constraints and $\mathbf{P}_{\mathbf{\Omega}}$ is a projection function which projects the updated architecture embedding back onto search space $\mathbf{\Omega}$. We formulate $\mathbf{P}_{\mathbf{\Omega}}$ as the following optimization problem:

In size search space, the output of this projection function can be obtained by rounding and clipping. As for topology search space, we select the operation with maximum probability as the chosen operation. Updating ${\bm{a}}^{(t)}$ with gradient $\frac{\partial{{\mathcal{P}}_{m}({\bm{a}}^{(t)})}}{\partial{{\bm{a}}^{(t)}}}$ increases the performance of ${\bm{a}}^{(t)}$ while $\frac{\partial{{\mathcal{P}}_{aux}({\bm{a}}^{(t)})}}{\partial{{\bm{a}}^{(t)}}}$ reduces the computational complexity. With a suitable $\alpha$, the gradient-based search strategy will optimize the architecture ${\bm{a}}$ towards to higher performance while reducing the resource cost. Given a number of initial architectures, we can explore new networks with higher performances based on the proposed gradient based search method efficiently. To obtain architectures under a hard resource constraints, we adopt a grid search strategy to tune the value of $\alpha$ and all the models under the target constraints during update will be added into the model pool.

Alg. 1 describes the whole process of search. After training predictors, we randomly sample $T_{max}$ architectures from search space. Then we update each architecture following Eqn.(3) until reaching max iteration $t_{max}$. We collect all models within constraints to a model pool and sort them by the predictor scores. Finally, we will retrain top-K of them on a specific task and select the best model as the final result.

In size search space, we adopt a 3-layer MLP followed by one fully-connected layer as our predictor. As for topology search space, we use a 3-layer GCNs followed by two fully-connected layers. We utilize GeLU [24] as activation function and set the dropout as 0.05 in the last fully-connected layers and GCNs layers. Table I summarizes the structures of predictors. Following [45], each layer in GCNs can be formulated as:

where $l$ indicates the index of GCNs layers. ${\mathbf{J}}\in\mathbb{R}^{I\times I}$ is the adjacency matrix, $I$ means the number of nodes. The operation matrix is ${\mathbf{V}}_{0}\in\mathbb{R}^{I\times D}$, $D$ is the number of operation candidates. ${\mathbf{W}}_{l}^{1}$ and ${\mathbf{W}}_{l}^{2}$ are trainable weights in GCNs.

NAS-Bench-201 [17] consists of 15625 neural cell candidates with different topologies on CIFAR-10, CIFAR-100 [30] and ImageNet-16-120 [8]. We introduce how we encode the architectures firstly. Following [45], we formulate the structure of cell as a densely connected Directed Acyclic Graph (DAG) and adopt an adjacency matrix and operation matrix to represent the structure. The DAG contains an input node, an output node and up to 6 interior nodes. Each interior node can be one of the following operations: (1) zeroize, (2) skip connection, (3) 1-by-1 convolution, (4) 3-by-3 convolution and (5) 3-by-3 average pooling layer. So the number of nodes in the DAG is $6+2=8$ and the number of candidate options per node is $5+2=7$²²2Following [45], the input and output are also considered as candidate options.. Then the topology of cell can be represented by a $8\times 8$ adjacency matrix. The operation matrix is a $8\times 7$ matrix which consists of 8 one-hot vectors. Since there have zeroize and skip connection operations in the candidate operation options, the adjacency matrix can be fixed to a graph with maximum connections in the search space, and only the operation matrix is updated.

We adopt the same structure of GCN proposed in [45] as our predictors. Since this benchmark usually take budget (the number of evaluated models) into consideration, we do not use the resource predictor. As for search, we set the number of iterations to 200 and decay the learning rate by a factor of 2. A softmax operation is applied on each row of the operation matrix to ensure that the sum of each row is 1. After gradient ascent, we select the operator with maximum probability in each row as the chosen operator.

We train main predictors with 30 samples and query performances of the top-40 networks in the model pool, which indicates that PredNAS trains $30+40=70$ models to get the final performance. Tabel II shows the results of our method. With the same number of query samples, we achieve better performance than BOHB[19], REA[41], RL[51] and a very strong baseline AceNAS[50] on all the three datasets. On ImageNet-16-120 dataset, our PredNAS achieve about $0.6\%-1.5\%$ higher accuracy than other methods. For fair comparison, we also report the results with 110 query samples following AceNAS. The results are comparable. These results validate the sampling efficiency of our method under low query budget. As shown in the following experiments, this advantage becomes significant when enlarging the search space.

We validate our method on ImageNet dataset and compare the searched networks with RegNetX models. RegNetX models are a series networks with different FLOPs obtained from the RegNet search space [40]. RegNet is obtained by progressively shrinking an unconstrained search space AnyNet by human heuristics. The search space of AnyNet contains 16 variables including the number of blocks, block width, bottleneck ratio and the number of group in 4 stages. Table III summarizes the search space of AnyNet, which includes about $(16\cdot 128\cdot 3\cdot 6)^{4}\approx 10^{18}$ architectures. The authors of [40] proposes several prior knowledge to progressively refine the AnyNet search space and sample 500 models every step to validate whether the refinement is worthy. They finally explore a RegNet search space which contains about $10^{8}$ models after 5 times refinement. Then each of the RegNetX models is obtained by picking the best model from 25 random architectures sampled from the RegNet search space with a specific FLOPs constraints. In this paper, we argue that this strategy still needs intensive human knowledge and trial and error, which is contradict to the spirit of AutoML in fact. Consequently, we adopt our PredNAS directly on the unconstrained AnyNet search space, and compare the searched architectures to RegNetX models.

We random sample 30 models from the AnyNet search space and train each model for 10 epochs on the ImageNet dataset [14] following the proxy task setting in [40]. As [40] provides RegNetX models with a variety of FLOPs, we adjust the weight of auxiliary predictor $\alpha$ to get the models with different target FLOPs during search. For example, the $\alpha$ is set to 1 for 200MF FLOPs regime while 0.2 for 3.2GF FLOPs.

For each FLOPs regime, we choose top-30 models from the model pool and pick the best one according to its performance on the proxy task. Then we train the best models with 100 epochs following [40]. Finally, we sampled $30+30\times 10=330$ models to get 10 architectures with different FLOPs, which is much less than 2750, the number of models sampled in RegNet [40].

As shown in Table IV, we achieve comparable performances to RegNetX from 200MF to 12GF by directly adopting our method on the AnyNet search space. We found some structures of models we explore are not consistent with the design principles introduced in RegNet. For example, PredNAS-600MF is deeper than RegNetX-600MF (see Table III), and the bottleneck ratio is not strict to 1. This reveal the fact that human hand-crafted design may still introduce inductive biases which should be avoided. Other architectures of searched model can be found in Appendix.

Another interesting finding is that we find there still exists a non-negligible gap between proxy and full train setting. PredNAS usually finds significantly better model than RegNetX in the proxy setting, however in full training, the conclusion contradicts. To further validate the gap between proxy and full train setting, we show another series of models PredNAS^† in Table IV. For each FLOPs regime, We select top-3 models according to their performances on the proxy task, and retrain all of them with 100 epochs. Then the models with best performances are selected as PredNAS^†. The architectures with best performances on full train setting may yield worse results on the proxy task. Consequently, we believe designing a better proxy setting that is more consistent with full training results is a crucial task for efficient NAS, however it is beyond the scope of this work and we leave it for future work.

We also adopt our method on the MSCOCO [33] dataset. The baseline is YOLOX [20], a highly efficient method proposed very recently. We propose a new search space based on YOLOX. The network architecture in YOLOX consists of three modules: backbone, neck and decoupled detection heads. The backbone and neck are composed of the same type of basic blocks. The searchable dimensions are depth (i.e. the number of blocks), widths and expand ratio of blocks. There are three decoupled detection heads in YOLOX, and each of them is responsible for a level of FPN [32] feature. The decoupled detection head contains convolution layers (shared convs) followed by two parallel branches for classification (cls head) and regression (reg head), respectively. We search for the depths and widths of convolution layers in each head. The summary of our proposed YOLOX search space is shown in Table V. This search space includes about $10^{100}$ possible architectures.

For such huge search space, we merely randomly sample 40 models to build training data. We train each sampled model following [20], but using 6% train data to alleviate the training cost.

For each target FLOPs, we select the best architecture from top-30 models in the model pool.

Table VI shows the results of our searched models. With FLOPs constraints, our models (PredNAS-S/M/L) consistently achieve better performances than YOLOX-S/M/L on the COCO val2017 set. Especially, comparing with YOLOX-S, our PredNAS-S achieves 42.5% AP, improving YOLOX-S by 2.0%. This result shows the effectiveness of our method on small scale models in object detection task. We further analyze the discrepancy of architectures between YOLOX-S and PredNAS-S (see Table V). PredNAS-S has a deeper backbone and computation resources are allocated more on the shared convolutions in head while less on the classification and regression heads.

We conduct experiments on NAS-Bench-201 with different N and K. N is the number of training samples for predictor and K is the number of top models selected for evaluation. The top of Table VIII shows the results with different number of training samples. Surprisingly, we found a small number of N is enough for training a predictor which can distinguish the performance of architectures coarsely. With K = 40, training predictor with only 10 samples can achieve better results than several baseline methods shown in Table II. The bottom of Table VIII shows the results with different K. We found K is more important in our experiments. With the increase of K, the performances are improving stably. This ablation shows that we can train a predictor with few samples and adopt it to do coarse model selection. Then the evaluation of top-K models can be seen as a refinement procedure.

RegNet. Table IX shows the searched models of PredNAS and PredNAS^† in RegNet search space. Following to RegNet, we increase the upper bound of search space to get large scale models during search. Therefore, some dimensions in the searched models may outrange the AnyNet search space described in the main paper. For example, the width of PredNAS-12GF in the last stage is 1368, larger than 1024. Interestingly, most of the searched models have increasing widths as the RegNet suggested.