GeNAS: Neural Architecture Search with Better Generalization

1 Introduction

Recently, Neural Architecture Search (NAS) [Liu et al., 2018b; Liu et al., 2018a; Hong et al., 2020] has evolved to achieve remarkable accuracy along with the development of humandesigned networks [He et al., 2016; Dosovitskiy et al., 2020] on the image recognition task. Several NAS methods [Zoph et al., 2018; Chu et al., 2020; Zhang et al., 2021; Liu et al., 2018b; Xu et al., 2019; Hong et al., 2020] further have demonstrated generalization ability (generalizability) of automatically designed networks with test accuracy and transfer performance onto the other datasets. For the widespread leverage of architectures found by NAS on the various tasks such as object detection [Lin et al., 2014] and segmentation [Cordts et al., 2016] (task-generalizability), investigating generalizability of each architecture candidate is prerequisite and indispensable. Despite its importance, quantitative measurement of generalizability during the architecture search process is still under-explored. In this paper, we aim to find an optimal proxy measurement to discriminate generalizable architectures during the search process. 1 2

Previous NAS algorithms including the pioneering differentiable search method, DARTS [Liu et al., 2018b] and evolutionary search method, SPOS [Guo et al., 2020] use validation performance as a proxy measure for the generalizability as follows:

where a and A denote an architecture candidate and the entire search space, respectively, and S(·) represents a measurement function which is broadly defined by accuracy [Guo et al., 2020] or negative of loss value [Liu et al., 2018b] on a validation dataset. Although these performance-based search (PBS) methods find the optimal architecture for generalization on the validation set, they show poor generalizability on the test set and other tasks, caused by overfitting on validation set [Zela et al., 2019; Oymak et al., 2021]. In addition, PBS methods represent a large discrepancy between the validation accuracy and ground truth test accuracy provided by NAS benchmark [Dong and Yang, 2020] as shown in [Guo et al., 2020; Zhang et al., 2021].

1 Code is available at https://github.com/clovaai/GeNAS.

2 Extended paper (including the appendix) is available at https://arxiv.org/abs/2305.08611.

To search generalizable architectures, several literatures [Shu et al., 2019; Zhang et al., 2021] empirically observe that architectures with fast convergence during training have a tendency to show better generalizability on test set. Based on the empirical connection between convergence speed and generalization, RLNAS [Zhang et al., 2021] proposed an Angle-Based Search (ABS) method, which estimates angle between initial and final network parameters after convergence of the model (i.e. convergence speed) as a proxy performance measure during the search process. However, we argue that ABS still has a large headroom for better generalization in terms of flat (wide) local minima, which has been considered as one of the key signals for inspecting generalizability of a trained network [Keskar et al., 2016; Zhang et al., 2018; Pereyra et al., 2017; Cha et al., 2020; He et al., 2019]. Intuitively, since the architecture with flat loss minima has widely low loss values around the minimum, it can achieve a low generalization error even if the loss surface is shifted due to the distribution gap from the test dataset. Since ABS only concerns the angle between initial model wights and trained ones in terms of convergence speed, found architectures can not be guaranteed to have flat local minima, as shown in Figure 1. Specifically, architectures not chosen by ABS (i.e. small angle) might have better generalizability based on the flat property of loss minima. Table 1 corroborates that angle is indeed in short of correlation with flatness of local

minima.

To explicitly design a search proxy measure that has a high correlation with the generalizability of the found model, we propose a flatness-based search method, namely FBS, which excavates a well-generalizable architecture by measuring the flatness of loss surface. FBS can find out robust architecture with low generalization error on shifted data distribution (e.g. test data, out-of-distribution datasets, downstream tasks) by inspecting both the depth and flatness of loss curvature near local minima through injecting random noise. In addition, FBS can be either replaced or incorporated with other stateof-the-art search measures to enhance performance as well as generalizability.

Consequently, building upon our search method FBS, we propose a novel flatness-based NAS framework, namely GeNAS, to exactly discriminate generalizability of architectures during searching. We show the effectiveness of the proposed GeNAS for both cases when using FBS solely or integrated into the conventional architecture score measurements such as PBS and ABS. Specifically, our GeNAS achieves comparable or even better performances on several NAS benchmarks compared to PBS- and ABS-based search methods [Liu et al., 2018b; Zhang et al., 2021; Xu et al., 2019; Guo et al., 2020; Chu et al., 2020; Chen et al., 2019; Hong et al., 2020]. Furthermore, we also show that the proposed FBS can be combined with conventional search metrics (e.g. PBS, ABS), inducing significant performance gain. Finally, we also demonstrate that our FBS can well-generalize on various data distribution shifts, as well as on multiple downstream tasks such as object detection and semantic segmentation.

Our contributions can be summarized as follows:

• We firstly demonstrate that the flatness of local minima can be employed to quantify generalizability of architecture in NAS domain, which only had been a means of confirming the generalizability after training a neural network.
• We propose a new architecture search proxy measure, flatness of local minima, well-suited for finding architectures with better generalization, which can replace or even significantly enhance the search performance of the existing search proxy measures.
• The found architecture induced by our FBS demonstrates the state-of-the-art performance on various search spaces and datasets, even showing great robustness on data distribution shift and better generalization on various downstream tasks.

2 RelatedWork

2.1 Neural Architecture Search

Early NAS methods are based on the reinforcement learning (RL) [Baker et al., 2016; Zoph et al., 2018], which train the agent network to choose better architecture. The RLbased methods require the test accuracy of each candidate network for reward value, so training every candidate network from scratch is also required to measure that. For this reason, it is not feasible on a large-scale dataset such as ImageNet [Krizhevsky et al., 2012]. To solve this problem, the weight-sharing NAS methods are introduced [Liu et al., 2018b; Xu et al., 2019; Xie et al., 2018; Guo et al., 2020; Zhang et al., 2021]. The weight-sharing NAS generally uses the SuperNet, which contains all operations in objective search space, and chooses several operations from the SuperNet to decide the final architecture, which is called SubNet. Among these weightsharing NAS frameworks, [Liu et al., 2018b; Xu et al., 2019] introduced a gradient-based architecture search method, where they jointly train the architecture parameters with weight parameters using gradient descent. After training, the final architecture is decided according to the architecture parameters. Meanwhile, the one-shot NAS methods [Guo et al., 2020; Bender et al., 2018; Brock et al., 2017] pointed out the critical drawback of these gradient-based search methods as there exists a strongly coupled and biased connection between SuperNet weight parameters and its architecture parameters; only a small subset of SuperNet weight parameters with large architecture parameter value will be intensely optimized, leaving the others trained insufficiently. Therefore, [Guo et al., 2020; Bender et al., 2018; Brock et al., 2017] sequentially decoupled the optimization process for SuperNet and architecture parameters, showing superior search performance over the gradientbased search methods. Inspired by these breakthroughs and its flexibility of introducing various search proxy measures, we construct our GeNAS based on the one-shot NAS framework.

2.2 Architecture Search ProxyMeasure

During the search time, it is hard to check the actual test performance of each architecture candidate when it is trained from scratch, so the proxy measure has to be employed for the candidate evaluation. Several approaches proposed to predictively discriminate well-trained neural networks without any training by inspecting either the correlation between the linear maps of variously augmented image [Mellor et al., 2021] or spectrum of Neural Tangent Kernel (NTK) [Chen et al., 2021]. Although these training-free search proxy measures significantly reduced the search costs within even four GPU hours, actual test performance was inferior to that of traininginvolved search proxy measures such as validation accuracy and loss. Meanwhile, ABS methods [Zhang et al., 2021; Hu et al., 2020] introduced a new search proxy measure, angle, for indicating the generalizability of a neural network architecture, showing search accuracy improvement [Zhang et al., 2021] over conventional search proxy measures such as validation accuracy [Guo et al., 2020]. Since ABS method only investigates the convergence speed of an architecture, [Zhang et al., 2021] successfully searched a well-trainable architecture under ground truth label absent during SuperNet training. However, searching with randomly-distributed label still shows large performance gap (about 0.15 Kendall's Tau score gap on NAS-Bench-201) to that of searching with the ground-truth label. Therefore, in order to fulfill higher test generalization of a searched architecture, we train SuperNet and searched architecture under ground-truth label setting.

2.3 Flatness of Local Minima

The flatness of loss landscape near local minima has been considered as a key signal for representing generalizability. [Keskar et al., 2016; Jastrz˛ebski et al., 2017; Hoffer et al., 2017] empirically observed that appropriate training hyperparameters such as batch size, learning rate, and the number of training iterations can implicitly enable a model to have wide and flat minima, enhancing test generalization performance. [Chaudhari et al., 2019] further explicitly drives a neural network model to the flat minima through an entropy-regularized SGD. Several works also promoted the flat local minima in terms of regularization during training using Label Smoothing [Pereyra et al., 2017] and Knowledge Distillation [Zhang et al., 2018], enjoying test performance gain. Based on these empirical connections between test generalizability and flatness of local minima of a neural network, we investigate the role of flatness of minima on the architecture search process. [Zela et al., 2019] has in common with our work in that they also employed a flatness of local minima during the architecture search process, but in an indirect manner. They proposed an early-stopping search process to prevent overfitting on the validation set when the approximated sharpness of local minima exceeds the threshold. Similarly, [Chen and Hsieh, 2020; Wang et al., 2021] tackled to alleviate fluctuating loss surface and accuracy caused by the discretization of architecture parameters in DARTS [Liu et al., 2018b]. We point out these previous similar works lack general usage on various NAS frameworks since they heavily depend on the DARTS [Liu et al., 2018b] framework. Meanwhile, our method can be applied to any architecture search framework without dependence on the architecture parameters of DARTS, such as evolutionarybased search algorithm [Guo et al., 2020].

3 Method

3.1 GeNAS: Generalization-Aware NASWith Flatness of Local Minima

GeNAS is aimed to search for network architectures with better generalization performance. To this end, we introduce a procedure for quantitatively estimating the flatness of an architecture's converged minima as a search proxy measure Fval(·) as follows:

Namely, we select the maximal flat architecture a∗ by evaluating flatness of every SubNet extracted from the pre-trained SuperNet W ∗A. From the previous studies [Zhang et al., 2018; Cha et al., 2020] that empirically investigated the landscape of converged local minima, the neural networks having flat local minima where the changes of the validation loss around the local minima are relatively small show better generalization performance at the test phase. Based on these simple but effective empirical connections, we introduce a novel method that searches for the architecture with maximal loss flatness around converged minima which can be formulated as below, following [Zhang et al., 2018]:

where L(θ) denotes validation loss value given weight parameter θ abbreviating W ∗A(a), and N(σi) denotes random Gaussian noise with its mean and standard deviation being 0 and σi, respectively. Namely, we inspect the flatness of minima near converged weight parameter space by injecting random Gaussian noise. The hyper-parameters σ denote the range for inspection of flatness and t denotes the number of perturbations. To perturb the weight parameters, we use unidirectional random noise, much more cost-efficient than recent flatness measuring approaches using Hessian [Yao et al., 2019] and bidirectional random noise [He et al., 2019] which can induce a considerable amount of computational cost. We observe that our choice is sufficient to discriminate architectures with high generalization performance.

Eq (2) and (3) would find architecture a∗ having the flattest local minima in the entire search space, but a∗ might have sub-optimal local minima far from the global minimum. In Figure 2, bottom-K architectures with the lowest ground test accuracy given by NAS-Bench-201 show the flattest local minima with relatively large loss values compared to the middle-K and top-K architectures. Therefore, naive investigation of the flatness of an architecture comes to achieve such sub-optimal architecture in terms of loss value. Note that the top-K architectures have the lowest loss values compared to middle and bottom architectures, equipped with flatness near converged minima. Correspondingly, considering both the flatness of loss landscape and the depth of minima is essential for excavating a generalizable architecture. For this reason, we add an additional term on Eq (3) to search for architectures with deep minima, along with flatness as follows:

Here, σ1 denotes the smallest perturbation degree among σ, hence the second term inspects how low the loss value nearest converged minima is. The term α denotes the balancing coefficient term between flat and deep minima, which is set to 1 unless specified.

3.2 SearchingWith CombinedMetrics

Recent works [Hosseini et al., 2021; Mellor et al., 2021] adopted a combined search metric for enhancing the performance of the resultant architecture. [Hosseini et al., 2021] employed an integrated search metric where the conventional cross-entropy loss over a clean image is combined with approximately measured adversarial robustness lower bound to enhance test accuracy of both clean images and adversarially attacked images. Inspired by the weak correlation between the existing search metrics (e.g. angle) and flatness (Table 1), we target to explicitly fulfill the large headroom of conventional search metrics to find better generalizable architectures in terms of our proposed flatness-based search measure (Eq (4)). Formally, we combine the existing metrics with flatness as a search proxy measure as follows:

Figure 3: Test loss curvatures of architectures found by Angle, Angle+Accuracy, Angle+Flatness.

where S denotes conventional search metrics such as angle and validation accuracy, γ is a balancing parameter between the existing metric and flatness, and β is a normalization term, which is fixed as σ-1 1 , for matching scale of flatness termwith the existing search metric.

4 Experiments

We first evaluate our proposed GeNAS framework on widely used benchmark dataset, ImageNet with DARTS [Liu et al., 2018b] search space. Furthermore, we thoroughly conduct ablation studies with regard to the components of GeNAS on NAS-Bench-201 [Dong and Yang, 2020] benchmark. We refer the reader to the appendix for more experimental details. For better confirming robust generalization effect with regard to data distribution shift, we evaluate the found architectures on ImageNet variants (ImageNetV2 [Recht et al., 2019],-A,-O [Hendrycks et al., 2021]). Furthermore, we test the transferability of our excavated architectures onto other task domains, object detection, and semantic segmentation, with MS-COCO [Lin et al., 2014] and Cityscapes [Cordts et al., 2016] dataset.

4.1 ImageNet

Searching on CIFAR-10

We analyze the transferability of architectures found on small datasets such as CIFAR-10 and CIFAR-100 onto ImageNet. Specifically, we search architectures with 8 normal cells (i.e., stride = 1) and 2 reduction cells (i.e., stride = 2) on CIFAR10/100, and transfer these normal / reduction cell architectures onto ImageNet by training from scratch and evaluating top-1 accuracy on ImageNet validation set. We compare our proposed FBS with other search metrics on CIFAR-10 in the upper part of Table 2. As a stand-alone search metric, flatness measure shows the best search performance among the other metrics including accuracy and angle with comparable FLOPs (∼= 0.6G) and parameters, when transferring searched architecture from CIFAR-10 onto ImageNet. Furthermore, when the angle is combined with flatness, loss landscape of found architecture becomes to be flatter and deeper as shown in Figure 3. As a result, search performance is further improved by 0.36% top-1 accuracy without any increase of either FLOPs or parameters. Also, the accuracy-based proxy measure also achieves performance gain when flatness is combined. The results show that our proposed flatness search metric indeed serves as a powerful search proxy measure for finding welltransferable architectures and also enhances the other search metrics to have a stronger ability to find architectures with better test generalization performance.

Searching on CIFAR-100

In middle part of Table 2, we analyze transferability of architectures found on CIFAR-100 onto ImageNet. The results show that flatness consistently reports significantly superior search performance even with fewer flops and parameters compared to ABS or PBS metrics, about 1.05% and 0.68% better top-1 accuracy, respectively. Furthermore, when flatness is appended onto angle and accuracy as a search proxy measure, top-1 accuracy drastically increases by 0.72% and 0.48%, respectively, which was consistently shown in CIFAR-10.

Searching on ImageNet

In the bottom part of Table 2, we directly search architectures on ImageNet and evaluate validation accuracy on ImageNet to compare in-domain search performance. Similar to the trend of the transfer experiments, our flatness metric achieves the best search performance compared to the existing search metrics and improves generalizability of them.

Comparison With SOTA NAS Methods

In Table 3, our GeNAS clearly represents large headroom compared to the other state-of-the-art NAS methods. Especially in comparison with SDARTS [Chen and Hsieh, 2020] which is a similar approach to GeNAS by using an implicit regularization for smoothing accuracy landscape, our GeNAS outperforms with a comparable number of FLOPs. Table 2 and 3 results show that our proposed flatness search metric indeed serves as a powerful search proxy measure for finding well-transferable architectures and also enhances the other search metrics to have a stronger ability to find architectures with better test generalization performance.

4.2 Generalization Ability

For a more sophisticated investigation of generalization ability, we analyze GeNAS in terms of robustness towards data distribution shift and transferability onto various downstream tasks in Table 4.

Distribution Shift Robustness

To measure robustness towards data distribution shift, we evaluate our found architectures on ImageNet variants, ImageNetV2 matched frequency [Recht et al., 2019] and ImageNetA [Hendrycks et al., 2021], where the test-set is distinct from the original ImageNet validation set. The results demonstrate superior robustness compared to the other NAS methods. Our GeNAS widens the performance gap especially when the distribution shift is severe as in ImageNet-A, which has extremely confusing examples.

Task Generalization

Object Detection We evaluate the generalization capability of architectures found by GeNAS on the downstream task, specifically object detection. We firstly re-train architectures found on CIFAR-100 onto ImageNet, and finetune on MS-COCO [Lin et al., 2014] dataset. For training, we adopt the default training strategy of RetinaNet [Lin et al., 2017] from Detectron2 [Wu et al., 2019]. We only replace the backbone network of RetinaNet for analyzing the sole impact of architectures found by each NAS method. The result shows that our GeNAS framework guided by the flatness measure clearly achieves the best AP scores. In case of RLNAS (angle) combined with flatness as a search metric, AP is enhanced by about 0.61%, without an increase of FLOPs or number of parameters.

Semantic Segmentation We also test the generalization of our GeNAS on Semantic Segmentation task with Cityscapes [Cordts et al., 2016] dataset. Based on the DeepLab-v3 [Chen et al., 2017], we only replaced the backbone network and trained with MMSegmentation [Contributors, 2020] framework. The results demonstrate the effectiveness of our flatness-guided architectures with a large performance margin. Consistently, our flatness guidance ensures a large performance gain, about 1.21%, when added onto anglebased search.

4.3 Ablation Study

To better analyze our proposed FBS-based GeNAS framework, we conduct an ablation study of each component and hyperparameters consisting of GeNAS.

Flatness Range

We analyze the effect of range of inspecting flatness near converged local minima in Table 5. The results demonstrate that searching flat architectures within too small area near converged minima (1st row in Table 5) is not sufficient for discriminating generalizable architectures. When σ is increased to {2e - 3, 1e - 2, 2e - 2}, Kendall's Tau is considerably improved, while further widening the flatness inspection range (4th row in Table 5) only significantly degrades the search performance on various datasets.

Deep and Low Minima

We further investigate the effect of searching architectures equipped with not only flatness but also the depth of loss landscape near converged minima. Specifically, we adjust α in Eq (4), where α = 0 denotes searching with only flatness of local minima. Results on Table 6 demonstrate that as α value increases from zero to one, search performance is drastically enhanced, indicating the indispensability of searching with both flatness and depth of minima. Note that α = 0 case can search out a sub-optimal architecture that has highly flat loss curvature but its loss values near local minima are too high, as shown in Figure 2. When α is further increased to α > 1, Kendall's Tau rank correlation starts to decrease, denoting that searching with largely depending on the depth of converged minima is not optimal for discriminating better generalizable architectures.

Perturbation Methodology

To quantitatively measure flatness of loss landscape, all the weight parameters of a given network are perturbed with random direction following Gaussian distribution as in Eq (4). Here, we investigate the effect of perturbation positions and directions. In Table 7, perturbing only weight parameters of target search cells (i.e. excluding stem conv layer and final fully-connected layer) only harms Kendall's Tau. Moreover, with regard to the perturbation directions, strongly perturbing the given models' parameters across the hessian eigenvectors [Yao et al., 2019] suffers from a slight decrease of Kendall's Tau (Table 7) with large computational overhead induced by approximation of hessian.

Effect of Flatness on ABS

We analyze the effect of integrating flatness on ABS. Specifically, we adjust γ in Eq (5), which balances the coefficient concerning the ratio of flatness to angle term. In Table 8, integrating flatness with a small proportion to angle mildly improves top-1 accuracy. As γ increases, top-1 accuracy of searched architecture gradually increases to reach 0.72% improvement over γ = 0 (ABS) case.

4.4 Search Cost Analysis

In Figure 4, we compare the required architecture search time of GeNAS with the other SOTA NAS frameworks. We measured the execution time spent for the SuperNet training and the search process, using a single NVIDIA V100 GPU. Our required search time is competitive to the other NAS methods while exhibiting shortened time compared to the other flatness-based search method (i.e., SDARTS-ADV).

5 Conclusion

This paper demonstrates that the flatness of local minima can be directly employed as a proxy of discriminating and searching for generalizable architectures. Based on the quantitative benchmark experiments on various search spaces and datasets, we demonstrate the superior generalizability of our flatnessbased search over conventional search metrics, while showing comparable or even better search performance compared to recent state-of-the-art NAS frameworks. We further analyze the insufficient generalizability of conventional search metrics in terms of the flatness of local minima. Consequently, integrating conventional search metrics with our proposed flatness measure can further lead to significantly boosting search performance. We also demonstrate superior generalization capability of GeNAS on the downstream object detection and semantic segmentation tasks while showing great robustness with regard to the data distribution shift.