Generalizing Few-Shot NAS with Gradient Matching

1 INTRODUCTION . In recent years , there has been a surge of interest in Neural Architecture Search ( NAS ) ( Stanley & Miikkulainen , 2002 ; Zoph & Le , 2017 ; Pham et al. , 2018 ; Real et al. , 2019 ; Liu et al. , 2019 ) for its ability to identify high-performing architectures in a series of machine learning tasks . Pioneering works in this field require training and evaluating thousands of architectures from scratch , which consume huge amounts of computational resources ( Miikkulainen et al. , 2019 ; Zoph & Le , 2017 ; Zoph et al. , 2018 ) . To improve the search efficiency , One-Shot NAS ( Pham et al. , 2018 ; Liu et al. , 2019 ; Bender et al. , 2018 ) proposes to train a single weight-sharing supernet ( Shot ) that encodes every architecture in the search space as a sub-path , and subsequently uses this supernet to estimate the performance of the underlying architectures efficiently . The supernet is represented as a directed acyclic graph ( DAG ) , where each edge is associated with a set of operations . In the One-Shot supernet , child models share weight when their paths in the DAG overlap . This way , One-Shot methods manage to cut the search cost down to training a single supernet model while still achieving state-of-the-art performances . Despite the search efficiency , training the weight-sharing supernet also induces coupled optimization among child models . Consequently , the supernet suffers from degenerated search outcomes due to inaccurate performance estimation , especially on top architectures ( Bender et al. , 2018 ; Yu et al. , 2020b ; Pourchot et al. , 2020 ; Zhang et al. , 2020b ; Zhao et al. , 2021 ) . To reduce the level of weightsharing , Few-Shot NAS ( Zhao et al. , 2021 ) proposes to split the One-Shot supernet into multiple independent sub-supernets via edge-wise exhaustive partitioning . Concretely , it assigns every op- eration on the selected edges to a separated sub-supernet , with weight-sharing enabled only within each sub-supernet ( Figure 1 ( a ) ) . This way , architectures containing different operations on the split edge are divided into different sub-supernets , thereby disabling weight-sharing between them . Although the performance of Few-Shot NAS surpasses its One-Shot counterpart , its splitting schema - naively assigning each operation on an edge to a separated sub-supernet - is not ideal . Some operations might behave similarly to each other and thus can be grouped into the same sub-supernet with little harm . In such cases , dividing them up could be a waste of precious search budgets while bringing little benefit to the performance . On the other hand , the gain from dividing dissimilar operations considerably out-weights that of dividing similar ones . The above analysis necessitates the design of an effective splitting criterion to distinguish between these two cases . Consider two child models that contain the same operations on all but the to-be-split edge of the supernet . Intuitively , they should not use one copy of weight if their training dynamics at the shared modules are dissimilar . Concretely , when these two networks produce mismatched gradients for the shared weight , updating the shared module under them would lead to a zigzag SGD trajectory . As a result , the performance estimation of these networks might not reflect their true strength . This can be further motivated by viewing the supernet training as multi-criteria optimization ( Fliege & Svaiter , 2000 ; Brown & Smith , 2005 ) , where each criterion governs the optimization of one child model ; Only similar update directions from different objectives could reduce the loss of all child networks . Otherwise , the performance of some child networks ( objectives ) would deteriorate . Based on these inspirations , we propose to directly use the gradient matching score ( GM ) as the supernet splitting criterion . The splitting decision can be made via graph clustering over the sub-supernets , with graph links weighted by the GM scores ( Figure 1 ( b ) ) . Utilizing the proposed splitting schema , we generalize the supernet splitting of Few-Shot NAS to support arbitrary branching factors ( number of children at each node of the partition tree ) . With a much lower branching factor , we could afford to split three times more edges compared with Few-Shot NAS under the same search budget , and achieve superior search performance . We conduct extensive experiments on multiple search spaces , datasets , and base methods to demonstrate the effectiveness of the proposed method , codenamed GM-NAS . Despite its simplicity , GMNAS consistently outperforms its One-Shot and Few-Shot counterparts . On DARTS Space , we achieve a test error of 2.34 % , ranked top among SOTA methods . On the MobileNet Space , GMNAS reaches 19.7 % Top-1 test error on ImageNet , surpassing previous comparable methods . 2 RELATED WORK . One-Shot NAS with Weight-Sharing aims at addressing the high computational cost of early NAS algorithms ( Bender et al. , 2018 ; Liu et al. , 2019 ; Li & Talwalkar , 2020 ) . Concretely , One-Shot NAS builds one single supernet that includes all child models in the search space as sub-paths , and allows the child models to share weight when their paths overlap . Architecture search can be conducted by training the supernet once and use it as the performance estimator to derive the best child model from the supernet . As a result , One-Shot NAS reduces the search cost down to training one single model ( the supernet ) . Despite the search efficiency , Weight-Sharing technique adopted by One-Shot NAS also posts a wide range of inductive biases due to coupled optimization , such as operation co-adaptation ( Bender et al. , 2018 ; Li et al. , 2019a ) , poor generalization ability ( Zela et al. , 2020 ; Chen et al. , 2021 ) , and distorted ranking correlations especially among top architectures ( Zhang et al. , 2020b ; Bender et al. , 2018 ; Zhao et al. , 2021 ) . As a result , the performance estimation of child models from the supernet could be inaccurate , and thus degrades the search results . Several lines of methods have been proposed to address the degenerated search performance caused by weight-sharing . Search space pruning methods identify and progressively discard poorperforming regions of the search space , so that other models do not need to share weight with candidates from these regions ( Liu et al. , 2018 ; Li et al. , 2019a ; Chen et al. , 2021 ; Li et al. , 2020a ; Hu et al. , 2020b ) . Distribution learning methods aim at inferring a sampling distribution that biases towards top performers ( Xie et al. , 2019 ; Chen et al. , 2021 ; Hu et al. , 2020a ; Dong & Yang , 2019 ) . Recently , there emerges a new orthogonal line of work that directly reduces the level of weightsharing by partitioning the search space into multiple sub-regions , with weight-sharing enabled only among models inside each sub-regions ( Zhang et al. , 2020b ; Zhao et al. , 2021 ) . Zhang et al . ( 2020b ) shows on a reduced search space that this treatment improves the ranking correlation among child models . Zhao et al . ( 2021 ) further proposes Few-Shot NAS : an edge-wise exhaustive partitioning schema that splits the One-Shot supernet into multiple sub-supernets , and obtains significantly improved performance over One-Shot baselines . Our work generalizes Few-Shot NAS to arbitrary branching factors by utilizing gradient matching score as the splitting criteria and formulating the splitting as a graph clustering problem . 3 METHOD . 3.1 FROM ONE-SHOT NAS TO FEW-SHOT NAS . One-Shot NAS One-Shot NAS represents the search space as a directed acyclic graph ( DAG ) , where each node denotes a latent feature map and each edge ( e ) contains operations o from the set O ( e ) . This way , every child model in the search space can be represented as one path in the DAG ( a.k.a . supernet ) . The search process is conducted by first optimizing the supernet once and subsequently using it as the performance estimator to derive the best architectures . One-Shot NAS induces the maximum level of weight-sharing : a child model shares weights with every other model in the search space as long as they include the same operation ( s ) on some edge ( s ) . Although weightsharing supernet drastically reduces the search cost , it also leads to inaccurate performance estimation of child models , as pointed out by several previous works ( Zela et al. , 2020 ; Bender et al. , 2018 ; Zhang et al. , 2020b ; Zhao et al. , 2021 ; Wang et al. , 2021b ) . Few-Shot NAS To address the aforementioned issue in One-Shot NAS , Zhao et al . ( 2021 ) proposes Few-Shot NAS that leverages supernet splitting to reduce the level of weight-sharing in OneShot supernet . Few-Shot NAS divides the One-Shot supernet into multiple sub-supernets ( each corresponding to a “ shot ” ) , where weight-sharing occurs only among the child models that belong to the same sub-supernet . Concretely , it adopts an edge-wise splitting schema : It first randomly selects a target compound edge from the supernet and then assigns each operation on the target edge into a separated sub-supernet , while keeping the remaining edges unchanged . As a result , child models containing different operations on the target edge are assigned to different sub-supernets , and therefore do not share weight with each other . As shown in Figure 1 ( a ) , this supernet splitting schema forms a partition tree of supernets , with a branching factor B equals to the number of operations on the edge ( |O ( e ) | ) . 3.2 GENERALIZED SUPERNET SPLITTING WITH ARBITRARY BRANCHING FACTORS . Due to the exhaustive splitting schema , Few-Shot NAS suffers from a high branching factor per split . Suppose we perform splits on T edges , the number of leaf-node sub-supernets becomes |O ( e ) |T . On DARTS Space , where the supernet contains 7 operations per edge and 14 edges per cell , splitting merely two edges lead to 72 = 49 sub-supernets . Consequently , conducting architecture search over these many sub-supernets induces prohibitively large computational overhead . For this reason , Few-Shot NAS could only afford to split very few edges ( in fact , one single split in most of its experiments ) . This could be suboptimal for many popular NAS search spaces as the supernets usually contain multiple edges/layers ( 14 for DARTS Space and 22 for MobileNet Space ) , and the decision on a single edge might not contribute much to the performance . Given a fixed search budget , measured by the total number of leaf-node sub-supernets , there exists a trade-off between the branching factor and the total number of splits . Exhaustive splitting can be viewed as an extreme case with a maximum branching factor . However , naively separating all operations on an edge might be unnecessary : Since some operations might behave similarly , splitting them into separated sub-supernets wastes a lot of search budgets while enjoying little benefit . In such case , we could group these operations into the same sub-supernet with minor sacrifice . This reduces the branching factor , allowing us to split more edges for a predefined budget and improve the search performance on large NAS search spaces with many edges/layers . We term this splitting schema as Generalized Supernet Splitting with arbitrary branching factors . We provide a formal formulation for the generalized supernet splitting with arbitrary branching factors . At each split , we select an edge e and divide the operations on it into B disjoint partitions : O ( e ) = ⋃ be=1···B O ( e ) be . B is thus the branching factor . For Few-Shot NAS , B = |O ( e ) | , and O ( e ) be is simply a unit set with one operation from edge e. When B < |O ( e ) | , more than one operations on the target edge will be assigned to the same sub-supernet , i.e . |O ( e ) be | > 1 . Let E denote the set of all edges and Et be the set of partitioned edges after t splits , then any sub-supernet generated after the t-th split can be uniquely determined by the partitioning set Pt = { ( e , O ( e ) be ) |e ∈ Et } ; Pt contains tuples ( e , O ( e ) be ) that record which operations on edge e get assigned to this sub-supernet . Let APt denote a sub-supernet with partition set Pt , then APt = ( ⋃ ( e , O ( e ) be ) ∈Pt O ( e ) be ) ⋃ ( ⋃ e∈E\Et O ( e ) ) . ( 1 )

Summary: This paper aims at improving few-shot NAS by proposing a gradient-guided schema to partition the supernet into sub-supernets during the search. The proposed splitting criterion is based on the cosine similarity between the gradients from different child models. Namely, child models that have similar gradients are more likely to be grouped together without splitting, so as to save computational costs and to allow the splitting of more layers than the vanilla few-shot NAS, achieving better performance. Extensive experiments show its superiority over the few-shot NAS.