Neighborhood-Aware Neural Architecture Search

1 INTRODUCTION . The process of automatic neural architecture design , also called neural architecture search ( NAS ) , is a promising technology to improve performance and efficiency for deep learning applications ( Zoph & Le , 2017 ; Zoph et al. , 2018 ; Liu et al. , 2019 ) . NAS methods typically minimize the validation loss to find the optimal architecture . However , directly optimizing such an objective may cause the search algorithm to overfit to the search setting , i.e. , finding a solution architecture with good search performance but generalizes poorly to the test setting . This type of overfitting is a result of the differences between the search and test settings , such as the length of training schedules ( Zoph & Le , 2017 ; Zoph et al. , 2018 ) , cross-architecture weight sharing ( Liu et al. , 2019 ; Pham et al. , 2018 ) , and the usage of proxy datasets during search ( Zoph & Le , 2017 ; Zoph et al. , 2018 ; Liu et al. , 2019 ) . To achieve better generalization , we propose a novel NAS formulation that searches for “ flat-minima architectures ” , which we define as architectures that perform well under small perturbations of the architecture ( Figure 1 ) . One example of architectural perturbations is to replace a convolutional operator with a skip connection ( identity mapping ) . Our work takes inspiration from prior work on neural network training , which shows that flat minima of the loss function correspond to network weights with better generalization than sharp ones ( Hochreiter & Schmidhuber , 1997 ) . We show that flat minima in the architecture space also generalize better to a new data distribution than sharp minima ( Sec . 3.3 ) . Unlike the standard NAS formulation that directly optimizes single architecture performance , i.e. , α∗ = arg minα∈A f ( α ) , we optimize the aggregated performance over the neighborhood of an architecture : α∗ = arg min α∈A g ( f ( N ( α ) ) ) , ( 1 ) where f ( · ) is a task-specific error metric , α denotes an architecture in the search space A , N ( α ) denotes the neighborhood of architecture α , and g ( · ) is an aggregation function ( e.g. , the mean function ) . Note that we overload the notation of the error metric f ( · ) and define f ( · ) to return a set of errors when the input is a set of architectures in the neighborhood : f ( N ( α ) ) = { f ( α′ ) | α′ ∈ N ( α ) } . Common choices for f ( · ) are validation loss and negative validation accuracy . We will discuss more details of neighborhood N ( α ) and aggregation function g ( · ) in the following text . To implement our formulation , one must define the neighborhood N ( α ) and specify an aggregation function g ( · ) . How to define the neighborhood of an architecture is an open question . One possible method to obtain neighboring architectures is to perturb one or more operations in the architecture and the degree of perturbation defines the scope of the neighborhood . This method can be applied to sampling-based search algorithms , e.g. , random search and reinforcement learning . However , it can not be directly used to generate neighboring architectures for gradient-based search algorithms ( a.k.a , differentiable NAS ) , where the neighboring architectures themselves also need to be differentiable with respect to the architecture being learned . To address this issue , we propose a differentiable representation for the neighborhood of architectures , which makes the objective function differentiable and allows us to apply our formulation to gradient-based algorithms , e.g. , DARTS ( Liu et al. , 2019 ) . Properly choosing the aggregation function g ( · ) can help the search algorithm identify flat minima in the search space . Our choice of g ( · ) ( e.g. , mean ) is inspired by the definition of the flatness/sharpness of local minima in previous work ( Chaudhari et al. , 2017 ; Keskar et al. , 2017 ; Dinh et al. , 2017 ) . We summarize our contributions as follows : 1 . We propose a neighborhood-aware NAS formulation based on the flat minima assumption , and demonstrate a principled way to apply our formulation to existing search algorithms , including sampling-based algorithms and gradient-based algorithms . We empirically validate our assumption and show that flat-minima architectures generalize better than sharp ones . 2 . We propose a neighborhood-aware random search ( NA-RS ) algorithm and demonstrate its superiority over the standard random search . On NAS-Bench-201 ( Dong & Yang , 2020 ) , NA-RS outperforms the standard random search by 1.48 % on CIFAR-100 and 1.58 % on ImageNet-16-120 . 3 . We propose a differentiable neighborhood representation so that we can apply our formulation to gradient-based NAS methods . We augment DARTS ( Liu et al. , 2019 ) with our formulation and name the proposed method NA-DARTS . Our NA-DARTS outperforms DARTS by 1.18 % on CIFAR-100 and 1.2 % on ImageNet , and also performs better than or on par with state-of-the-art NAS methods . 2 RELATED WORK . Flat Minima . Hochreiter & Schmidhuber ( 1997 ) shows that flat minima of the loss function of neural networks generalize better than sharp minima . Flat minima are used to explain the poor generalization of large-batch methods ( Keskar et al. , 2017 ; Yao et al. , 2018 ) , where large-batch methods are shown to be more likely to converge to sharp minima . Chaudhari et al . ( 2017 ) propose an ob- jective function for training neural networks so that flat minima are preferred during optimization . Their objective can be interpreted as a weighted average of the ( transformed ) function values of data points around the local minima , which inspires us to consider mean as one of the aggregation functions . Previous work mentioned above focus on flat minima in the network weight space . However , we study flat minima in the architecture space , which is discrete and fundamentally different from the continuous weights studied in previous work . This makes it non-trivial to apply the flat minima idea to NAS . Zela et al . ( 2020 ) observes a strong correlation between the generalization error of the architecture found by DARTS ( Liu et al. , 2019 ) and the flatness of the loss function at the found architecture . They propose several regularization strategies to improve DARTS , such as early stopping before the loss curvature becomes too high . Our flat minima assumption is motivated by their observation and our method can be combined with their regularization strategies . NAS - Search Algorithm . Various search algorithms have been applied to solve NAS , including sampling-based and gradient-based algorithms . Representative sampling-based algorithms include random search ( Li & Talwalkar , 2019 ) , reinforcement learning ( Baker et al. , 2017 ; Zoph & Le , 2017 ; Zoph et al. , 2018 ; Zhong et al. , 2018 ) , Bayesian optimization ( Kandasamy et al. , 2018 ; Cao et al. , 2019 ) , evolutionary algorithms ( Xie & Yuille , 2017 ; Real et al. , 2017 ; 2019 ) , and sequential modelbased optimization ( Liu et al. , 2018 ) . To make NAS more computationally efficient , weight sharing across architectures is proposed to amortize the training cost of candidate architectures ( Pham et al. , 2018 ; Bender et al. , 2018 ) . Based on weight sharing , gradient-based algorithms are proposed to directly learn the architecture with gradient descent ( Liu et al. , 2019 ; Xie et al. , 2019 ) . Our focus is not proposing novel search algorithms but revisiting the standard NAS formulation . Our proposed formulation can be applied to both sampling-based algorithms and gradient-based algorithms . NAS - Search Space . Search space is crucial for the performance of NAS . One of the most widely used search spaces is the cell search space ( Zoph et al. , 2018 ) , which searches for a cell that can be stacked multiple times to form the entire network . Our proposed neighborhood-aware formulation is agnostic to the search space , and we specifically showcase our formulation on the cell search space . 3 NEIGHBORHOOD-AWARE FORMULATION . We propose a neighborhood-aware NAS formulation ( Eq . 1 ) to identify flat minima in the search space . Our formulation builds upon the assumption that flat-minima architectures usually generalize better than sharp ones . In this formulation , the optimal architecture is selected according to the aggregated performance g ( f ( N ( α ) ) ) of neighbors of an architecture , instead of the standard criterion , i.e. , single architecture performance f ( α ) only . We now introduce the neighborhood definition of an architecture N ( α ) and the aggregation function g ( · ) . 3.1 NEIGHBORHOOD DEFINITION AND CELL SEARCH SPACE . Formally defining the neighborhood requires a distance metric between architectures , which largely depends on how an architecture is represented and how the search space is constructed . We adopt the cell search space ( Zoph et al. , 2018 ) as it has been widely used in recent NAS methods ( Liu et al. , 2018 ; 2019 ) . Instead of the entire architecture , we search for a cell that can be stacked multiple times to form the entire architecture . The number of times the cell is stacked and the output layer are manually defined prior to the search . A cell is defined as a directed acyclic graph ( DAG ) consisting of n nodes . Each node represents a feature map . Each directed edge ( i , j ) ( 1 ≤ i < j ≤ n ) is associated with an operation used to transform the feature map at node i , and passes the transformed feature map to node j . The feature map at one node is the sum of all the feature maps on the incoming edges to this node : x ( j ) = ∑ ( i , j ) ∈E ∑m k=1 α ( i , j ) k ok ( x ( i ) ) , where E denotes the set of edges in the cell , x ( i ) is the feature map at node i , and ok is the kth operation among the m available operations . α ( i , j ) is a m-dim one-hot vector , indicating the operation choice for edge ( i , j ) . A cell is then represented by a set of variables α = { α ( i , j ) } . Note that α ( i , j ) being a one-hot vector means that only one operation is chosen for edge ( i , j ) . On a side note , the one-hot constraint on α ( i , j ) can be relaxed in differentiable NAS methods ( Liu et al. , 2019 ; Xie et al. , 2019 ) . We define the distance between two cells α and α′ as : dist ( α , α′ ) = ∑ ( i , j ) ∈E δ ( α ( i , j ) , α′ ( i , j ) ) , ( 2 ) where δ ( · , · ) is the total variation distance between two probability distributions : δ ( p , q ) = 1 2 ||p − q||1 = 1 2 ∑m k=1 |pk − qk| . Here p and q are both m-dim probability distributions . The total variation distance is symmetric and bounded between 0 and 1 . It also offers the following property : δ ( α ( i , j ) , α′ ( i , j ) ) = 0 implies that the two cells have the same operation at edge ( i , j ) and δ ( α ( i , j ) , α′ ( i , j ) ) = 1 implies that they have different operations at edge ( i , j ) . Note that instead of directly counting the edge differences , we adopt total variation distance to accommodate relaxed α that is later used in differentiable NAS methods ( Liu et al. , 2019 ; Xie et al. , 2019 ) . The neighborhood of a cell α is defined as : N ( α ) = { α′ | dist ( α , α′ ) ≤ d } , ( 3 ) where d is a distance threshold . Due to the property of the total variation distance , when d is an integer , the neighborhood contains all the cells that have at most d edges associated with different operations from α . For clarification , our definition of neighborhood includes the reference architecture α itself .

Summary: This paper proposes a neighbor-aware method in the neural architecture search (NAS). The paper states that by optimizing a neighbor of the neural network, it can search the result in a flat-minima, which is more stable than the sharp minima. The experiment results in further support that the proposed NA-RS and NA-DARTS outperform the current SOTA in various tasks.