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 work is a further exploration in the direction of Few-Shot NAS. It proposes a method of partioning supernet based on gradient-matching score. Compared with the method based on exhaustive-spltting, the method proposed in this paper can achieve better results. Its contributions are as follows:  1. Point out the problem based on exhaustive-splitting method; 2. A partioning method based on gradient-matching score is proposed, and this method can achieve good results on multiple search spaces and datasets.