SUMNAS: Supernet with Unbiased Meta-Features for Neural Architecture Search

1 INTRODUCTION . Recent Neural Architecture Search ( NAS ) algorithms often train an over-parameterized network called a supernet to obtain supreme sub-models by sharing parameters rapidly . In this process , the shared parameters are generally trained to reach a state where better architecture discovery is possible through comparing the sub-models . Supernet training is usually conducted by sampling one or more sub-models and updating them with their gradients . Since parameters are shared , learning not to be biased towards a particular sub-model is crucial in training a supernet . However , a parameter update due to sampling a specific sub-model often causes a bias for the sampled sub-model ; the supernet forgets the knowledge learned from the previous sampling because of the biased update , and other sub-models share the same parameters can be degraded and underrated . Benyahia et al . ( 2019 ) referred to this phenomenon as multi-model forgetting , and many current architecture search strategies have been designed without considering this phenomenon . To mitigate this problem , it is important to make the supernet parameters learn unbiased features that are globally suitable for the sub-models . Therefore , we propose a meta-learning-based approach that enables the supernet to learn unbiased meta-features . We adopt model-agnostic meta-learning ( MAML ) ( Finn et al. , 2017 ; Nichol et al. , 2018 ) during supernet training . MAML is designed to learn the meta-features suitable for multiple tasks , and the trained meta-parameters are then quickly adapted to an unseen task through few-shot learning while reusing the learned meta-features . We apply this concept of MAML to supernet training by assuming the multiple tasks in MAML as learning for multiple sub-models in a supernet . We call the proposed supernet training algorithm Supernet with Unbiased Meta-Features for Neural Architecture Search ( SUMNAS ) , given that it introduces the meta-learning principle to the multimodel forgetting problem . SUMNAS consists of two stages : supernet training and heuristic search with sub-model evaluation . In the supernet training stage , the parameters learn the meta-features for multiple sub-models , which can be considered as MAML ’ s training meta-features for multiple tasks . The meta-parameters obtained from the supernet training can then be directly used without additional training for comparison between sub-models during the evaluation phase . To the best of our knowledge , this work is the first to utilize the meta-learning capability—learning unbiased meta-features—for accurately ranking sub-models in a supernet . There have been approaches ( Shaw et al. , 2019 ; Lian et al. , 2020 ; Wang et al. , 2020 ; Elsken et al. , 2020 ) that adopt meta-learning principles to train a supernet using various tasks and search for the best architecture for an unseen task with a few data instances , which is the few-shot learning variant of NAS . They train the model parameters such that the models are robust to unseen datasets . On the other hand , we take a fundamentally different view of a task . A task usually refers to a dataset , but we show in Section 3 that each sub-model within a search space can be regarded as a separate “ task ” that the supernet has to adapt to . We therefore apply meta-learning principles to make the model parameters robust to several different sub-models . We evaluate SUMNAS with qualitative and quantitative experiments on the CIFAR10 and ImageNet datasets . We show that the architecture rankings SUMNAS predicts have a stronger correlation with the accurate rankings as compared to prior NAS algorithms that use a supernet . Besides , we observe better architecture search performance when an existing search methodology is applied to SUMNAS and show that SUMNAS parameters properly learn meta-features by investigating the performances of the sub-models . 2 PRELIMINARIES . One-shot NAS : In the context of a supernet that shares parameters with its sub-models , the majority of NAS approaches use the supernet as a proxy to indirectly predict the performance of sub-models . With the performance oracle , differentiable techniques ( Liu et al. , 2018 ; Cai et al. , 2019 ) , reinforcement learning ( Pham et al. , 2018 ) , and heuristic methods ( Guo et al. , 2019 ; Chu et al. , 2019 ) are used to determine the best architecture . Although the weight sharing has significantly improved search speed , it is not easy to provide accurate performance indicators when using a supernet as a proxy for sub-model performance . One of the most difficult challenges to obtaining an accurate performance oracle is overriding knowledge learned by previously sampled sub-models . Researchers have recently introduced the practice of repeatedly updating parameters through one or more sampled sub-models ( Guo et al. , 2019 ; Chu et al. , 2019 ; Li et al. , 2020 ) . However , Benyahia et al . ( 2019 ) have shown that catastrophic forgetting occurs in the sequential learning of sampled sub-models . Catastrophic forgetting is a phenomenon where a neural network forgets previously learned knowledge when learning new information . In sampling-based supernet training , repetitive sub-model sampling and training keep introducing new information into shared parameters . The knowledge learned with the previously sampled sub-models is forgotten due to training subsequently sampled sub-models without considering former training . Knowledge overriding causes the predicted performance of the architectures to fluctuate depending on the frequency or the sequence of the sampled architectures , which leads to inaccurately predicted rankings . The forgetting problem in the stochastic training of shared supernet parameters is termed multi-model forgetting . To alleviate this problem , Benyahia et al . ( 2019 ) and Zhang et al . ( 2020 ) add a regularizer to prevent weights from deviating too far from the posterior distribution learned from previously trained sub-models . Model-agnostic meta-learning : Model-agnostic meta-learning ( MAML ) is a process of learning initialization for few-shot learning of an unseen task . Let the ith task be Ti and θ′i be suitable to model parameters for Ti . MAML provides an objective function to obtain optimal meta-parameters by minimizing the loss function for each task Ti : θmeta = argmin θ ETi∼p ( T ) [ LTi ( fθ∗i ) ] ( 1 ) where f is the network function consisting of θ′ , and L is the loss function for each task Ti . The L term in equation 1 allows the θ∗i calculated by few-shot learning with initialization θ to learn a suitable feature for each task Ti , and θmeta minimizes the expected loss corresponding to the tasks . As a consequence , the learned meta feature is generalized across tasks , allowing for fast learning of new tasks . Finn et al . ( 2017 ) solve the objective function by repeatedly sampling tasks and updating parameters with gradient descents . They also suggest a first-order approximation to avoid the computation of second-order derivatives . Nichol et al . ( 2018 ) propose a variation of gradient descent for metalearning , which is called Reptile . The first-order MAML carries out task optimization using only the last gradient calculated by the inner loop step , ignoring the second-order derivative . However , Reptile uses the average of the gradients calculated over multiple inner loop steps to update the parameters . These averaged values reflect the generalization of the gradients for the data . 3 SUPERNET WITH UNBIASED META-FEATURES FOR NEURAL ARCHITECTURE SEARCH . To tackle the multi-model forgetting problem , we suggest a new supernet training strategy based on meta-learning , Supernet with Unbiased Meta-Features for Neural Architecture Search ( SUMNAS ) . As mentioned , learning unbiased features that are suitable for sub-models in the supernet is essential to alleviate the knowledge overriding . We take the idea of MAML , which learns meta-features suitable for multiple tasks , and apply it to supernet training so that SUMNAS learns such unbiased meta-features of sub-models . In this section , we describe our approach and explain how it can mitigate the multi-model forgetting problem . We then describe our supernet training algorithm . 3.1 META-FEATURE LEARNING . The approach we introduce to improve robustness against the multi-model forgetting problem is learning unbiased meta-features . SUMNAS has two stages : supernet training and heuristic search with sub-model evaluation , which are similar to existing sampling-based NAS algorithms ( Chu et al. , 2019 ; Guo et al. , 2019 ; Li et al. , 2020 ) , as presented in Figure 1 . First , SUMNAS trains parameters to learn meta-features ( Figure 1 [ a ] ) , which are not biased to a specific sub-model and are suitable for sub-models , so training them does not overwrite previously trained knowledge . Afterward , the parameters are used to evaluate a specific sub-model during the search ( Figure 1 [ b ] ) . Vanilla sampling-based supernet training optimizes strictly shared supernet parameters on sampled sub-models , and it can be expressed as follows : θs = argmin θ Ef∼p ( f ) [ LT ( fθ ) ] , ( 2 ) where f is a sub-model sampled from a distribution of sub-models p ( f ) , and T is a given dataset . The parameter θ is strictly shared among sub-models , and it has to learn model-specific features of all of the sub-models . The training for one or more sampled sub-models does not consider other ones that have been previously sampled , and simply overrides knowledge acquired from previous training to optimize currently sampled sub-models . In our approach , similar to previous works , parameters are shared during supernet training . However , the parameters are trained to learn unbiased meta-features . Here , the unbiased meta-features make fair comparison possible by providing each sub-model with an appropriate level of optimized parameters that are not overfitted for a specific sub-model . Our meta-learning approach has been formulated as follows : θs = argmin θ Ef∼p ( f ) [ LT ( fθ∗ ) ] s.t . θ∗ = Af ( L , T , θ ) , ( 3 ) whereAf is an adaptation function optimizing the parameters for the sub-model f with a few inputs . With this optimization , parameters naturally learn meta-features . This idea is based on MAML , which targets few-shot learning of an unseen task using adaptable parameters . MAML trains parameters to learn meta-features suitable for multiple tasks and uses the meta-parameters to initialize few-shot learning . Training meta-parameters for a weight-shared supernet is similar to training with MAML . We map a task used in meta-learning for few-shot learning to a sub-model in NAS using a supernet , as seen by comparing Equation 1 with Equation 3 . In other words , the difference is that MAML minimizes the expectation of the loss across tasks , but SUMNAS minimizes the expectation of the loss across sub-models . Interestingly , we notice that learning the meta-features in a supernet can be viewed as MAML of each operator in the supernet . We can convert Equation 3 into the following equation : θmeta = argmin θ Eo∼p ( o ) [ Ef∼p ( f |fi=o ) [ LI ( f i+1 : θ∗ ◦ o ) ] ] s.t . I = p ( f0 : i−1θ∗ ( x ) , y| { x , y } ∼ T ) θ∗ = Af ( L , T , θ ) ( 4 ) where f i : j is a composition of sampled operators from i-th layer to j-th layer . A sampled sub-model f0 : i−1 feeds an intermediate representation f0 : i−1 ( x ) to the operator o . There is a distribution of the intermediate representation fed to the operator , and the distribution is determined by a submodel sampled from p ( f ) . The operator also has a target distribution of outputs to learn determined by the sampled sub-model . In the MAML , examples fed to a meta-learning model also follow a distribution , and this distribution is a task Ti sampled from p ( T ) . Thus , inputs and targets of both a multi-task meta-learning model and an operator in a supernet follow sampled distributions . We also aim to train an operator that works well on a set of intermediate representations provided by a sampled sub-model , just as MAML learns to work well on a set of examples given by a sampled task . That is , a model and a task in MAML can be mapped to an operator and a distribution of intermediate representations fed to the operator in supernet training . In other words , we can regard the distribution of intermediate representations as to the task for the operator , although the task for the operator is parameterized and optimizable differently from tasks in MAML . This implies that we can use any MAML algorithms without alteration to solve Equation 4 . Moreover , it supports the assumption that the success of the meta-learning algorithms for few-shot learning will be transferred to NAS . However , supernet training and MAML have a fundamental difference . MAML targets to train parameters for an unseen task , but supernet training aims to obtain parameters for pre-defined submodels , which are known at train time , and participate in the training process . Therefore , after our supernet training , the supernet has knowledge of the sub-models , so additional training , such as fine-tuning for each sub-model is not mandatory . In our empirical results , the additional fitting during evaluation may improve the performance of each sub-model , but the ranking performance stays more or less the same .

Summary: The paper proposed an improved training strategy for oneshot NAS supernetworks. The key idea is to view the training of each subnetwork as a "task", and then to apply an MAML/Reptile-style meta-learning scheme to ensure efficient cross-task adaptivity. Experiments on NAS-Bench-201 and ImageNet show improved calibration between the supernet’s predictions and the architectures' true rankings.