Fast MNAS: Uncertainty-aware Neural Architecture Search with Lifelong Learning

Sampling-based neural architecture search ( NAS ) always guarantees better convergence yet suffers from huge computational resources compared with gradient-based approaches , due to the rollout bottleneck – exhaustive training for each sampled generation on proxy tasks . This work provides a general pipeline to accelerate the convergence of the rollout process as well as the RL learning process in samplingbased NAS . It is motivated by the interesting observation that both the architecture and the parameter knowledge can be transferred between different experiments and even different tasks . We first introduce an uncertainty-aware critic ( value function ) in PPO to utilize the architecture knowledge in previous experiments , which stabilizes the training process and reduces the searching time by 4 times . Further , a life-long knowledge pool together with a block similarity function is proposed to utilize lifelong parameter knowledge and reduces the searching time by 2 times . It is the first to introduce block-level weight sharing in RL-based NAS . The block similarity function guarantees a 100 % hitting ratio with strict fairness . Besides , we show a simply designed off-policy correction factor that enables “ replay buffer ” in RL optimization and further reduces half of the searching time . Experiments on the MNAS search space show the proposed FNAS accelerates standard RL-based NAS process by ∼10x ( e.g . ∼256 2x2 TPUv2 * days / 20,000 GPU * hour→ 2,000 GPU * hour for MNAS ) , and guarantees better performance on various vision tasks . 1 INTRODUCTION . Neural architecture search ( NAS ) has made great progress in different tasks such as image classification ( Tan & Le , 2019 ) and object detection ( Tan et al. , 2019b ) . And usually , there are four commonly used NAS algorithms : differentiable , one-shot , evolutional , and reinforcement learning ( RL ) based method . The RL-based method , due to its fair sampling and training processes , has often achieved a great performance among different tasks . However , one of the biggest challenges of it is the high demand for computing resources , which makes it hard to follow by the research community . RL-based NAS consumes a large number of computing powers on two aspects : a ) the need for sampling a large number of architectures to optimize the RL agent and b ) the tedious training and testing process of these samples on proxy tasks . For example , the originator of NAS ( Zoph & Le , 2016 ) requires 12,800 generations of architecture and current state-of-the-art MNAS ( Tan et al. , 2019a ) and MobileNet-V3 ( Howard et al. , 2019 ) require 8000 or more generations to find the optimal architectures . Besides , each generation is usually trained for 5 epochs . All in all , it costs nearly 64 TPUv2 devices for 96 hours or 20,000 GPU hours on V100 for just one single searching process . With such a severe drawback , researchers start looking for other options like differential ( Liu et al. , 2018b ; Chen et al. , 2019 ) , or one-shot based ( Bender , 2019 ; Guo et al. , 2019 ) method for NAS . The one-shot family has drawn lots of attention recently due to its efficiency . It applies a single super-network based search space with that all the architectures , also called sub-networks , share parameters with the super-network during the training process . In this way , the training process is condensed from training thousands of sub-networks into training a super-network . However , this share-weight strategy may bring problems for the performance estimation of sub-networks . For example , two sub-networks may propagate conflicting gradients to their shared components , and the shared components may converge to favor one of the sub-networks and repel the other randomly . This conflicting phenomenon may result in instability of the search process and inferior final architectures , compared with RL-based methods . In this work , we seek to combine the privilege of RL-based methods and one-shot methods , by leveraging the knowledge from previous NAS experiments . The proposed method is based on two key observations : First , the optimal architectures for different tasks have common architecture knowledge . Second , the parameter knowledge can also be transferred across experiments and even tasks . Based on the observations , for transferable architecture knowledge , we develop Uncertainty-Aware Critic ( UAC ) to learn the architecture-performance joint distribution from other experiments even other tasks in an unbiased manner , utilizing the transferability of the structural knowledge , which reduces the sample ’ s training time by 50 % and the result is shown in Figure 1 ( with UAC ) ; For the transferable parameter knowledge , we propose Lifelong Knowledge Pool ( LKP ) to restore the block-level parameters and fairly share them to new samples ’ initialization , which speeds up each samples ’ convergence for 2 times , as shown in Figure 1 ( with LKP ) ; Finally , we also developed an Architecture Experience Buffer ( AEB ) with a significant off-policy correctness factor to store the old models for reusing in RL optimization , with half of the time saved . And this is shown in Figure 1 ( with AEB ) . Under the strictly same environment as MNAS and MobileNet-v3 , FNAS speed up the searching process by 10× and the performances are even better . 2 REVISITING SAMPLING-BASED NAS . 2.1 NAS FAMILY . From the perspective of how to derive the performance estimation of an architecture , NAS methods can be split into two categories , sampling-based and share-weight based . Sampling-based methods usually sample many architectures from the search space and train them independently . Based on the performance of these well-trained architectures , several ways can be utilized to fetch the best one , such as Bayesian optimization ( Kandasamy et al. , 2018 ) , evolutionary algorithm ( Real et al. , 2019 ) , and training an RL agent ( Zoph & Le , 2016 ) . The main drawback of these methods is a huge time and resource consumption of training the sampled architectures . To alleviate this issue , a common practice is to shorten the training epochs and use proxy networks with fewer filters and cells ( Zoph et al. , 2018 ; Tan et al. , 2019a ) . Liu et al . ( Liu et al. , 2018a ) propose to train a network to predict the final performance . Different from these methods , we leverage the accumulated knowledge to accelerate the training process . Instead of training many architectures independently , the second kind of methods resorts to train a super-network and estimate the performance of architectures with weights shared from the supernetwork ( Bender , 2019 ; Wu et al. , 2019 ; Liu et al. , 2018b ; Chen et al. , 2019 ; Xu et al. , 2019 ; Cai et al. , 2019 ; Stamoulis et al. , 2019 ; Guo et al. , 2019 ) . With the easy access of performance estimation , DARTS ( Liu et al. , 2018b ) proposes a gradient-based method to search for the best architecture in an end-to-end manner . However , as pointed in ( Li & Talwalkar , 2019 ) , the performance estimation based on shared weights may be unreliable . Chen et al . ( Chen et al. , 2019 ) proposes to progressively shrink the search space so that the estimation can be more and more accurate . Cai et al . ( Cai et al. , 2019 ) introduces a shrinking based method to train the supernet so as to generate networks of different scales without retraining . We also share weights between architectures but in different ways . We construct a general weight pool with many trained architectures , and when we want to train a new architecture , we initialize it by the trained architectures in the pool . In this way , the number of training epochs can be reduced without harming reliability . 3 KNOWLEDGE BETWEEN NAS EXPERIMENTS IS TRANSFERABLE . RL-based NAS consumes a lot of computing resources . MNAS , as it ’ s said before , trains 8,000 models for the agent to converge , which costs 20,000 GPU hours on V100 . And all the samples trained for one experiment will not be used anymore . However , the active differentiable-based NAS demonstrates that with various weight-sharing techniques , the NAS algorithms can be accelerated a lot . In this section , we will show that the knowledge of previous searched experiments can be reused by the following two observations , which helps the follow-up experiment greatly . 3.1 ARCHITECTURE KNOWLEDGE CAN BE TRANSFERRED . Optimal architectures for different tasks have common architecture knowledge . One always holds the assumption that the performance of a model is consistent among different tasks . A common practice of this assumption is applying good ImageNet ( Deng et al. , 2009 ) models to COCO object detection ( Lin et al. , 2014 ) as the backbone . In NAS , however , this assumption needs to be carefully checked as the huge search space of it requires the hypothesis to be well generalized . Here , We statistically verify this assumption . In Figure 4 , we sample 100 optimal models on face recognition and ImageNet classification tasks respectively . For each model , we firstly expand each digit of its embedding ( i.e . a 35 dimension vector ) to one-hot : e.g . from `` 4 '' to `` 1000 '' ; `` 3 '' to `` 0100 '' . In this way , the original embedding is expanded to 112-dim from 35-dim . After that , we calculate the expectation for each digit of this expended vector among the top 100 optimal architectures on the face recognition and ImageNet classification respectively . Shown in Figure 4 , we get an observation that the operators can be divided into two spaces , i.e . exclusive space , where the probability difference is large and common space , where the probability difference is small . Many previous works ( Liu et al. , 2018a ; Kokiopoulou et al. , 2019 ; Luo et al. , 2018 ; Wen et al. , 2019 ; Luo et al. , 2020 ) use a predictor to predict a model ’ s performance to speed up the NAS process . However , as the predictor requires thousands of samples to train , they usually implement in a progressive ( Liu et al. , 2018a ) or semi-supervised manner ( Luo et al. , 2020 ) . Inspired by the interesting observation above , we implement it in a unified way where different tasks ’ samples are used together to train a unified value network to predict a model ’ s performance . When searching architecture on a new task , we just use directly the unified network trained by the old data and keep updating it in the new task during the search process , which speeds up the convergence of the value network . Showing in Figure 5a , when transferring a value network trained on ImageNet to face recognition task , the network converges much faster . 3.2 PARAMETER KNOWLEDGE CAN BE TRANSFERRED . Initializing the network by ImageNet pre-trained models and training the model on other tasks has nearly been a standard way as it can always speed up the convergence process . However , pretraining has been ignored in the NAS area as it may break the rank of different models . In our experiments , we observe that the trained checkpoint , we call it parameter knowledge , can help us to get the real rank faster than training from scratch . Besides , this feature holds regardless of the data distribution . We randomly sample 50 models and train them on ImageNet in two ways : from scratch or by initializing with parameter knowledge from face experiment . Then , we compare the rank correlation with real rank ( i.e . fully trained rank ) along the training process . Showing in Figure 5b and 5c , with parameter knowledge from face experiment , one gets more accurate rank in fewer epochs .

Summary: The paper propose a few improvements to the sampling-based NAS using RL: 1) an uncertainty-aware critic to decide whether the sample needs to be trained; 2) a life-long knowledge pool to initialize the sample that needs training; and 3) an architecture experience buffer to reuse old samples for RL training. The experiments are done on ImageNet, facial recognition and transferability on object detection. The proposed methods are compared with related works. Finally the paper finishes with ablation studies on both the effectiveness and transferability of the proposed modules.