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 introduces a searching framework of neural architectures search by modifying the objective function to optimize the aggregated performance over the neighborhood of an architecture. From the observation that flat-minima architecture $\alpha$ generalizes better than sharp-minima architecture (Zela et al. (ICLR2020), the author proposes the objective function considering the neighborhood to enforce the flat minima. The author supports their method by providing ablation studies and architecture performances from CIFAR-10/100, ImageNet, and NAS-BENCH-201.