Cream of the Crop: Distilling Prioritized Paths For One-Shot Neural Architecture Search

Houwen Peng, Hao Du, Hongyuan Yu, Qi Li, Jing Liao, Jianlong Fu

Introduction

Neural Architecture Search (NAS) is an exciting field which facilitates the automatic design of deep networks. It has achieved state-of-the-art performance on a variety of tasks, surpassing manually designed counterparts [e.g., 1, 2, 3]. Recently, one-shot NAS methods became popular due to low computation overhead and competitive performance. Rather than training thousands of separate models from scratch, one-shot methods only train a single large hypernetwork capable of emulating any architecture in the search space. The weights are shared across architecture candidates, i.e., subnetworks. Such strategy is able to reduce the search cost from thousands of GPU days to a few.

However, all architectures sharing a single set of weights cannot guarantee each individual subnetwork obtains sufficient training. Although one-shot models are typically only used to sort architectures in the search space, the capacity of weight sharing is still limited. As revealed by recent works , weight sharing degrades the ranking of architectures to the point of not reflecting their true performance, thus reducing the effectiveness of the search process. There are a few recent works addressing this issue from the perspective of knowledge distillation . They commonly introduce a high-performing teacher network to boost the training of subnetworks. Nevertheless, these methods require the teacher model to be trained beforehand, such as a large pretrained model or a third-party model . This limits the flexibility of search algorithms, especially when the search tasks or data are entirely new and there may be no available teacher models.

In this paper, we present prioritized paths to enable the knowledge transfer between architectures, without requiring an external teacher model. The core idea is that subnetworks can learn collaboratively and teach each other throughout the training process, and thus boosting the convergence of individual architectures. More specifically, we create a prioritized path board which recruits the subnetworks with superior performance as the internal teachers to facilitate the training of other models. The recruitment follows the selective competition principle, i.e., selecting the superior and eliminating the inferior. Besides competition, there also exists collaboration. To enable the information transfer between architectures, we distill the knowledge from prioritized paths to subnetworks. Instead of learning from a fixed model, our method allows each subnetwork to select its best-matching prioritized path as the teacher based on the representation complementary. In particular, a meta network is introduced to mimic this path selection procedure. Throughout the course of subnetwork training, the meta network observes the subnetwork’s performance on a held-out validation set, and learns to choose a prioritized path from the board so that if the subnetwork benefits from the prioritized path, the subnetwork will achieve better validation performance.

Such prioritized path distillation mechanism has three advantages. First, it does not require introducing third-party models, such as human-designed architectures, to serve as the teacher models, thus it is more flexible. Second, the matching between prioritized paths and subnetworks are meta-learned, which allows a subnetwork to select various prioritized paths to facilitates its learning. Last but not the least, after hypernetwork training, we can directly pick up the best performing architecture from the prioritized paths, instead of using either reinforcement learning or evolutional algorithms to further search a final architecture from the large-scale hypernetwork.

The experiments demonstrate that our method achieves clear improvements over the strong baseline and establishes state-of-the-art performance on ImageNet. For instance, with the proposed prioritized path distillation, our search algorithm finds a 481M Flops model that achieving 79.2% top-1 accuracy on ImageNet. This model improves the SPOS baseline by 4.5% while surpassing the EfficientNet-B0 by 2.9%. Under the efficient computing settings, i.e., Flops $\leq 100$ M, our models consistently outperform the MobileNetV3 , sometimes by nontrivial margins, e.g., 3.0% under 43M Flops. The architecture discovered by our approach transfers well to downstream object detection task, getting an AP of 33.2 on COCO validation set, which is superior to the state-of-the-art MobileNetV3. In addition, distilling prioritized paths allows one-shot models to search architectures over more challenging search space, such as the combinations of MBConv , residual block and normal 2D Conv, thus easing the restriction of designing a carefully constrained space.

Preliminary: One-Shot NAS

One-shot NAS approaches commonly adopt a weight sharing strategy to eschew training each subnetwork from scratch [13, 14, 15, 8, 16, among many others]. The architecture search space $\mathcal{A}$ is encoded in a hypernetwork, denoted as $\mathcal{N}(\mathcal{A},W)$ , where $W$ is the weight of the hypernetwork. The weight $W$ is shared across all the architecture candidates, i.e., subnetworks ${\alpha\in\mathcal{A}}$ in $\mathcal{N}$ . The search of the optimal architecture $\alpha^{*}$ in one-shot methods is formulated as a two-stage optimization problem. The first-stage is to optimize the weight $W$ by

where $\mathcal{L}_{\emph{train}}$ represents the loss function on training dataset. To reduce memory usage, one-shot methods usually sample subnetworks from $\mathcal{N}$ for optimization. We adopt the single-path uniform sampling strategy as the baseline, i.e., each batch only sampling one random path from the hypernetwork for training . The second-stage is to search architectures via ranking the performance of subnetworks ${\alpha\in\mathcal{A}}$ based on the learned weights $W_{\mathcal{A}}$ , which is formulated as

where the sampled subnetwork $\alpha$ inherits the weight from $W_{\mathcal{A}}$ as $w_{\alpha}$ , and $\emph{Acc}^{\alpha}_{\emph{val}}$ indicates the top-1 accuracy of the architecture $\alpha$ on validation dataset. Since that it is impossible to enumerate all the architectures ${\alpha\in\mathcal{A}}$ for evaluation, prior works resort to random search , evolution algorithms or reinforcement learning to find the most promising one.

Distilling Priority Paths for One-Shot NAS

The weight sharing strategy reduces the search cost by orders of magnitude. However, it brings a potential issue, i.e., the insufficient training of subnetworks within the hypernetwork. Such issue results in the performance of architectures ranked by the one-shot weight is weakly correlated with the true performance. Thus, the search based on the weight $W$ may not find a promising architecture. To boost the training of subnetworks, we present prioritized path distillation. The intuitive idea is to leverage the well-performing subnetwork to teach the under-trained ones, such that all architectures converge to better solutions. In the following, we first present the mechanism of prioritized path board, which plays a fundamental role in our approach. Then, we describe the search algorithm using the prioritized paths and knowledge transfer between architectures. The overall framework is visualized in Fig. 1.

2 Architecture Search with Prioritized Paths

Our solution to the insufficient training of subnetworks is to distill the knowledge from prioritized paths to the weight sharing subnetworks. Due to the large scale of the search space, the structure of subnetworks are extremely diverse. Some subnetworks may be beneficial to other peer architectures, while others may not or even harmful. Hence, we allow each subnetwork to find its best matching collaborator from the prioritized path board, such that the matched path can make up its deficiency. We propose to learn the matching between prioritized paths and subnetworks by a meta network $\mathcal{M}$ . Since there is no available groundtruth to measure the matching degree of two architectures, we use the learning state (i.e., validation loss) of subnetworks as the signal to supervise the learning of the meta network. The underlying reason is that if the gradient updates of the meta network encourage the subnetworks to learn from the selected prioritized path and achieve a small validate loss, then this matching is profitable.

The hypernetwork training with prioritized path distillation includes three iterative phases.

where $\rho$ is the output of the meta network $\mathcal{M}$ and represents the matching degree (the higher the better) between the prioritized path $\hat{\alpha}_{k}$ and the subnetwork $\alpha$ , $x$ indicates the training data, and $\theta$ denotes the weight of $\mathcal{M}$ . The input to the meta network $\mathcal{M}$ is the difference of the feature logits between the subnetworks $\mathcal{N}(\hat{\alpha}_{k},w_{\hat{\alpha}_{k}})$ and $\mathcal{N}(\alpha,w_{\alpha})$ . Such difference reflects the complementarity of the two paths. The meta network learns to select the prioritize path $\hat{\alpha}^{*}_{k}$ that is complementary to the current subnetwork $\alpha$ .

Phase 2: distilling knowledge from the prioritized path. With the picked prioritized path $\hat{\alpha}^{*}_{k}$ , we perform knowledge distillation to boost the training of the subnetwork $\alpha$ . The distillation is supervised by a weighted average of two different objective functions. The first objective function $\mathcal{L}_{\emph{CE}}$ is the cross entropy with the correct labels $y$ . This is computed using the logits in softmax of the subnetwork, i.e., $p(x,w_{\alpha})$ = $\emph{softmax}(\mathcal{N}(x,\alpha,w_{\alpha}))$ . The second objective function $\mathcal{L}_{\emph{KD}}$ is the cross entropy with the soft target labels and this cross entropy distills knowledge from the prioritized path $\hat{\alpha}^{*}_{k}$ to the subnetwork $\alpha$ . The soft targets $q(x,w_{\hat{\alpha}^{*}_{k}})$ are generated by a softmax function that converts feature logits to a probability distribution. We use SGD with a learning rate $\eta$ to optimize the objective functions and update the subnetwork weight $w_{\alpha}$ as

Phase 3: updating the meta network. Since there is no available groundtruth label measuring the matching degree and complementarity of two architectures, we resort to the loss of the subnetwork to guide the training of the matching network $\mathcal{M}$ . The underlying reason is that if one prioritized path $\hat{\alpha}^{*}_{k}$ is complementary to the current subnetwork $\alpha$ , then the updated subnetwork with the weight $w^{(t+1)}_{\alpha}$ can achieve a lower loss on the validation data. We evaluate the new weight $w^{(t+1)}_{\alpha}$ on the validatation data ( $x$ , $y$ ) using the cross entropy loss $\mathcal{L}_{\emph{CE}}(y,p(x,$ ${w^{(t+1)}_{\alpha}}))$ . Since ${w^{(t+1)}_{\alpha}}$ depends on $\rho$ via Eq. (5) while $\rho$ depends on $\theta$ via Eq. (4), this validation cross entropy loss is a function of $\theta$ . Specifically, dropping ( $x$ , $y$ ) from the equations for readability, we can write:

This dependency allows us to compute $\nabla_{\theta}\mathcal{R}$ to update ${\theta}$ and minimize $\mathcal{R}({\small{w_{\alpha}^{(t+1)}}})$ . The differentiation $\nabla_{\theta}\mathcal{R}$ requires computing the gradient of gradient, which is time-consuming, we thereby updates ${\theta}$ every $\tau$ iterations. In essence, the meta network observing the subnetwork’s validation loss to improve itself is similar to an agent in reinforcement learning performing on-policy sampling and learning from its own rewards . In implementation, we adopts one fully-connected layer with 1,000 hidden nodes as the architecture of meta network, which is simple and efficient.

The above three phases are performed iteratively to train the hypernetwork. The iterative procedure is outlined in Alg. 1. Thanks to the prioritized path distillation mechanism, after hypernetwork training, we can directly select the best performing subnetwork from the prioritized path board as the final architecture, instead of further performing search on the hypernetwork.

Experiments

In this section, we first present ablation studies dissecting our method on image classification task, and then compare our method with state-of-the-art NAS algorithms. The experiments on object detection and more challenging search space are performed to evaluate the generality and robustness.

2 Ablation Study

We dissect our method and evaluate the effects of each components. Our baseline is the single-path one-shot method, which trains the hypernetwork with uniform sampling and searches architectures by an evolution algorithm . We re-implement this algorithm in our codebase, and it achieves 76.3% top-1 accuracy on ImageNet, being superior to the original 74.7% reported in due to different search spaces (ShuffleUnits v.s. MBConv ). If we replace the evolution search with the proposed prioritized path mechanism, the performance is still comparable to the baseline, as presented in Tab. 1(#1 v.s. #2). This suggests the effectiveness of the prioritized paths. By comparing #2 with #4/#5, we observe that the knowledge distillation between prioritized paths and subnetworks is indeed helpful for both hypernetwork training and the final performance, even when the matching between prioritized paths and subnetworks is random, i.e. #4. The meta-learned matching function is superior to random matching by 1.3% in terms of top-1 accuracy on ImageNet. The ablation between #5 and #6 shows that the evolution search over the hypernetwork performs comparably to the prioritized path distillation, suggesting that the final paths left in the prioritized path board is the "cream of the crop".

We further perform a correlation analysis to evaluate whether the enhanced training of the hypernetwork can improve the ranking of subnetworks. To this end, we randomly sample 30 subnetworks and calculate the rank correlation between the weight sharing performance and the true performance of training from scratch. Unfortunately, training such many subnetworks on ImageNet is very computational expensive, we thus construct a subImageNet dataset, which only consists of 100 classes randomly sampled from ImageNet. Each class has 250 training images and 50 validation images (Image lists are released with the code). Its size is about 50 $\times$ smaller than the original ImageNet. The Kendall rank correlation coefficient on subImageNet is reported in Tab. 1. It is clear to observe that after performing prioritized path distillation, the ranking correlation is improved significantly, e.g., from the baseline 0.19 to 0.37 (#1 v.s. #5 in Tab. 1).

There are two hyperparameters in our method: one is the size of the prioritized path board and the other is the number of validation images for prioritized path selection in Eq. (3). The impact of these two hyperparameters are reported in Tab. 4.2 and 4.2 respectively. We observe that when the number of prioritized paths is increased, the performance of hypernetwork becomes better, yet bringing more search overhead. Considering the tradeoff, we empirically set the number of prioritized paths to 10 in the experiments. A similar situation is occurred on the number of val images. We randomly sample 2,048 images from the validation set (50k images in total) for prioritized path selection because it allows fast evaluations and keep a relatively high Kendall rank.

3 Comparison with State-of-the-Art NAS Methods

Fig. 2 presents the comparison of our method with state-of-the-arts under mobile settings on ImageNet. It shows that when the model Flops are smaller than 600M, our method consistently outperforms the recent MobileNetV3 and EfficientNet-B0/B1 . In particular, our method achieves 77.6% top-1 accuracy on ImageNet with 287M Flops, which is 1.0% higher than MobileNetV31.25x while using 1.2 $\times$ fewer Flops. Moreover, our method is flexible to search low complexity models, only requiring the users input a desired minimum and maximum Flops constraint. From Fig. 2(right), we can see that when the Flops are smaller than 100M, our models establishes new state-of-the-arts. For example, when using 43M Flops, MobileNetV3 is inferior to our model by 3.0%. Besides model complexity, we are also interested in inference latency.

As shown in Tab. 4, where we report the average latency of 1,000 runs, our method runs 1.1 $\times$ faster than EfficientNet-B0 and 1.7 $\times$ faster than MobileNetV3 on a single core of Intel Xeon CPU E5-2690. Also, the performance of our method is 1.3% superior to EfficientNett-B0 and 2.4% superior to MobileNetV3. This suggests our models are competitive when deployed on real hardwares.

Tab. 5 presents more comparisons. It is worth noting that there are few recent works leveraging knowledge distillation techniques to boost training . Compared to these methods, our prioritized path distillation is also superior. Specifically, DNA recruits EfficientNet-B7 , a very high-performance third-party model, as the teacher and achieves 78.4% top-1 accuracy (without using AutoAugment), while our method (Cream-L) gets a superior accuracy of 80.0% without using any other pretrained models. Our method performs comparably to the recent OFA yet taking much less time on hypernetwork training, i.e., 12 v.s. 53 GPU days. Thanks to the prioritized path mechanism, our method only need to evaluate $K$ =10 prioritized paths on the validation set and then select the best performing one. This procedure only takes 0.02 GPU days, which is 30 $\times$ faster than other approaches of using evolutional search algorithm, such as SPOS and OFA . The learned architectures are plotted in Appendix B.

4 Generality and Robustness

As shown in Tab. 7, our method surpasses MobileNetV2 by 4.9% while using fewer Flops. Compared to MnasNet , our method utilizes 19% fewer Flops while achieving 2.7% higher performance, suggesting the architecture has good generalization capacity when transferred to other vision tasks. If we further increase the model complexity, our method can achieve an AP of 36.8%, which is comparable to the recent Hit-Detector ( $36.9$ AP with $1839$ M Flops) but uses much less Flops.

A robust search algorithm should be capable of searching architectures over diverse search spaces. To evaluate this, we evaluate our method on more challenging space, i.e., the combinations of operators from different designed space, including MBConv , Residual Block and normal 2D convolutions. Due to limited space, we present the detailed settings of the new search spaces in Appendix C. As the results reported in Tab. 6, we observe that when the search space becomes more challenging, the performance of the baseline SPOS algorithm is degraded. In contrast, our method shows relatively stable performance, demonstrating it has potentials to search for architectures over more flexible spaces. The main reason is attributed to the prioritised path distillation, which improves the ranking correlation of architectures.

Related Work

Neural Architecture Search. Early NAS approaches search a network using either reinforcement learning or evolution algorithms . These approaches require training thousands of architecture candidates from scratch, leading to unaffordable computation overhead. Most recent works resort to the one-shot weight sharing strategy to amortize the searching cost . The key idea is to train a single over-parameterized hypernetwork model, and then share the weights across subnetworks. The training of hypernetwork commonly samples subnetwork paths for optimization. There are several path sampling methods, such as drop path , single path and multiple paths . Among them, single-path one-shot model is simple and representative. In each iteration, it only samples one random path and train the path using one batch data. Once the training process is finished, the subnetworks can be ranked by the shared weights. On the other hand, instead of searching over a discrete set of architecture candidates, differentiable methods relax the search space to be continuous, such that the search can be optimized by the efficient gradient descent. Recent surveys on architecture search can be found in .

Distillation between Architectures. Knowledge distillation is a widely used technique for information transfer. It compresses the "dark knowledge" of a well trained larger model to a smaller one. Recently, in one-shot NAS, there are few works leveraging this technique to boost the training of hypernetwork [e.g., 42], and they commonly introduce additional large models as teachers. More specifically, OFA pretrains the largest model in the search space and use it to guide the training of other subnetworks, while DNA employs the third-party EfficientNet-B7 as the teacher model. These search algorithms will become infeasible if there is no available pretrained model, especially when the search task and data are entirely new. The most recent work, i.e. BigNAS , proposes inplace distillation with a sandwich rule to supervise the training of subnetworks by the largest child model. Although this method does not reply on other pretrained models, it cannot guarantee the fixed largest model is the best teacher for all other subnetworks. Sometimes the largest model may be a noise in the search space. In contrast, our method dynamically recruits prioritized paths from the search space as the teachers, and it allows subnetworks to select their best matching prioritized models for knowledge distillation. Moreover, after training, the prioritized paths in our method can serve as the final architectures directly, without requiring further search on the hypernetwork.

Conclusions

In this work, motivated by the insufficient training of subnetworks in the weight sharing methods, we propose prioritized path distillation to enable knowledge transfer between architectures. Extensive experiments demonstrate the proposed search algorithm can improve the training of the weight sharing hypernetwork and find promising architectures. In future work, we will consider adding more constraints on prioritized path selection, such as both model size and latency, thus improving the flexibility and user-friendliness of the search method. The theoretical analysis of the prioritized path distillation for weight sharing training is another potential research direction.

Broader Impact

Similar to previous NAS works, this work does not have immediate societal impact, since the algorithm is only designed for image classification, but it can indirectly impact society. As an example, our work may inspire the creation of new algorithms and applications with direct societal implications. Moreover, compared with other NAS methods that require additional teacher model to guide the training process, our method does not need any external teacher models. So our method can be used in a closed data system, ensuring the privacy of user data.

Acknowledgements

We acknowledge the anonymous reviewers for their insightful suggestions. In particular, we would like to thank Microsoft OpenPAI v-team for providing AI computing platform and large-scale jobs scheduling support, and Microsoft NNI v-team for AutoML toolkit support as well as helpful discussions and collaborations. Jing Liao and Hao Du were supported in part by the Hong Kong Research Grants Council (RGC) Early Career Scheme under Grant 9048148 (CityU 21209119), and in part by the CityU of Hong Kong under APRC Grant 9610488. This work was led by Houwen Peng, who is the Primary Contact (✉ houwen.peng@microsoft.com).