Year: 2021

Abstract:
Efficient performance estimation of architectures drawn from large search spaces is essential to Neural Architecture Search. One-Shot methods tackle this challenge by training one supernet to approximate the performance of every architecture in the search space via weight-sharing, thereby drastically reducing the search cost. However, due to coupled optimization between child architectures caused by weight-sharing, One-Shot supernet’s performance estimation could be inaccurate, leading to degraded search outcomes. To address this issue, Few-Shot NAS reduces the level of weight-sharing by splitting the One-Shot supernet into multiple separated sub-supernets via edge-wise (layer-wise) exhaustive partitioning. Since each partition of the supernet is not equally important, it necessitates the design of a more effective splitting criterion. In this work, we propose a gradient matching score (GM) that leverages gradient information at the shared weight for making informed splitting decisions. Intuitively, gradients from different child models can be used to identify whether they agree on how to update the shared modules, and subsequently to decide if they should share weight. Compared with exhaustive partitioning, the proposed criterion significantly reduces the branching factor per edge. This allows us to split more edges (layers) for a given budget, resulting in substantially improved performance as NAS search spaces usually include dozens of edges (layers). Extensive empirical evaluations of the proposed method on a wide range of search spaces (NASBench-201, DARTS, MobileNet Space), datasets (cifar10, cifar100, ImageNet), and search algorithms (DARTS, SNAS, RSPS, ProxylessNAS, OFA) demonstrate that it significantly outperforms its Few-Shot and One-Shot counterparts while surpassing previous comparable methods in terms of the accuracy of derived architectures.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 METHOD

heading: 3.1 FROM ONE-SHOT NAS TO FEW-SHOT NAS
text: 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)|).

heading: 3.2 GENERALIZED SUPERNET SPLITTING WITH ARBITRARY BRANCHING FACTORS
text: 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)

heading: 3.3 SUPERNET SPLITTING VIA GRAPH MIN-CUT WITH GRADIENT MATCHING SCORE
text: Generalized Supernet Splitting necessitates the design of a splitting criterion for grouping operations, which leads to a question as to how to decide which networks should or should not share weight. A naive strategy is to perform the random partition, with the underlying assumption that the weight-sharing between all child models are equally harmful. However, empirically we find that this treatment merely matches exhaustive splitting in terms of their performance, and in some cases even worsens the results. To demonstrate this, we compare random partition with the Few-Shot NAS (exhaustive) on NASBench201 (Dong & Yang, 2020) with four operations (skip, conv 1x1, conv 3x3, avgpool 3x3)1. For random partitioning, we split
two edges on its supernet (T = 2), and divide the operations on each edge into two groups randomly (B = 2). This leads to four sub-supernets, same as Few-Shot NAS with one single split. As shown in Table 1, random split degrades the search performance for both continuous and sampling-based One-Shot NAS. The comparison result reveals that the weight-sharing between some child models are much more harmful than others, and therefore need to be carefully distinguished.
If not all weight-sharing are equally harmful, how should we decide whether two models should be split into separated sub-supernets? Intuitively, networks should not share weight if their training dynamics mismatch at the shared modules. In other words, they disagree on how to update the shared weight. In this case, the angles between gradient directions produced by two networks at the shared modules might be large, leading to a zigzag pattern of SGD trajectory. As a result, the
1Using 4 operations allow us to match the number of supernets between random and exhaustive split.
performance estimation of these models in the supernet could be inaccurate, thereby degrading the performance of downstream architecture search task.
To show this, we evaluate the performance of a single model A when updated together with another network Asim that produces similar gradients at shared weight, and compare the performance with the same model updated together with a network Adissim that produces dissimilar gradients. We generate (A,Asim, Adissim) by sampling from the NASBench-201 search space, in a way that they have the same operation (and thus share weights) on all but one edges (More on this in Appendix B.2). We then proceed to train (A,Asim) together and similarly (A,Adissim) together via Random Sampling with Parameter Sharing (RSPS) (Li & Talwalkar, 2020), and record A’s performance under these two cases. As shown in Table 2, A achieves a much lower loss when updated together with Asim, where gradients from two child models are similar at the shared module. The result indicates that weight-sharing could be more harmful between networks with dissimilar training dynamics.
Inspired by the above analysis, we propose to measure the harmfulness of weight-sharing between child models (hence sub-supernets) directly via gradient matching (GM). Concretely, consider two operations from an edge; if these operations, when enabled separately, lead to drastically different gradients on the shared edges, we give them higher priority for being split into different sub-supernets during the supernet partitioning. The entire splitting schema
can be formulated as a graph clustering problem: Given a (sub-)supernet APt−1 , we evaluate the gradient of the supernet when each operation o is enabled on edge et separately, and then compute the cosine similarity between every pair of these gradients:
GM(APt−1 |et=o, APt−1 |et=o′ ) = COS [ ∇wsL(APt−1 |et=o;ws),∇wsL(APt−1 |et=o′ ;ws) ] (2)
where APt−1 |et=o means to enable operation o on edge et of supernet APt−1 , and ws is the shared weight. This leads to a densely connected graph (Figure 1 (b)) where the vertices are operations on et and links between them are weighted by GM score computed in Eqn. (2). Therefore, supernet partitioning can be conducted by performing graph clustering on this graph, with the number of clusters equal to the desired branching factor. There exist many applicable algorithms for solving this problem, but |O(e)| is usually small, we perform graph min-cut via brute-force search to divide the operations (supernets) into B balanced groups. For B = 2, it can be written as:
U = argmin U⊆O(e) ∑ o∈U,o′∈O(e)\U GM(APt−1 |et=o, APt−1 |et=o′ ), (3)
s.t. b|O(e)|/2c ≤ |U| ≤ d|O(e)|/2e.
where {U ,O(e) \ U} are the obtained partitions. The proposed splitting schema substantially improves the search performance over exhaustive split and random split, as evidenced by Table 1.

heading: 3.4 THE COMPLETE ALGORITHM
text: Edge selection using gradient matching score Apart from smartly selecting which operations should be grouped within an edge, we can also use the gradient matching score to select which edge to split next. More specifically, the graph min-cut algorithm produces a cut cost for each edge - the sum of gradient matching scores of the cut links on the gradient similarity graph. This can serve as an edge importance measure for determining which edge to split on next, as a lower cut cost indicates that splitting on this edge first might relieve the adverse effect of weight-sharing to a larger extend. Empirically, we find that this edge scoring measure reduces the variance over random selection used in Few-Shot NAS, and also improves the performance (Section 5.1).
Supernet splitting with restart During the supernet splitting phase, we warmup the supernet’s weight (w) for a few epochs before each split, in order to collect more accurate gradient information along the optimization trajectory. Since the effect of weight-sharing has already kicked in during this phase, we re-initialize all the leaf-node sub-supernets after the final split is completed, and conduct
architecture search over them from scratch. The complete algorithm is summarized in Algorithm 1 in the Appendix.

heading: 4 EXPERIMENTS
text: In this section, we conduct extensive empirical evaluations of the proposed method across various base methods (DARTS, RSPS, SNAS, ProxylessNAS, OFA) and search spaces (NASBench-201, DARTS, and MobileNet Space). Experimental results demonstrate that the proposed GM-NAS consistently outperforms its Few-Shot and One-Shot counterparts.

heading: 4.1 NASBENCH-201
text: We benchmark the proposed method on the full NASBench-201 Space (Dong & Yang, 2020) with five operations (none, skip, conv 1x1, conv 3x3, avgpool 3x3). We run the search phase of each method under four random seeds (0-3) and reports the mean accuracy of their derived architectures, as well as the standard deviation to capture the variance of the search algorithms. For our method, we split the operations on each edge into two groups (one group with three operations, another group with two operations), and cut two edges in total, amounting to four sub-supernets. Note that this is one supernet less than Few-Shot NAS, which uses five sub-supernets due to its exhaustive splitting schema.
Still, the proposed method improves over Few-Shot NAS by a large margin. As shown in Table 3, the architectures derived from GM DARTS achieve an average accuracy of 93.72% on CIFAR-10, leading to an improvement of 5.17% over Few-Shot DARTS and 39.24% over DARTS. To further test the generality of the proposed method over various One-Shot NAS algorithms, we also compare GM-NAS with Few-Shot NAS on sampling-based methods such as SNAS and RSPS. As before, our method consistently out-performs Few-Shot NAS and the One-Shot base methods by a substantial margin. Notably, when combined with the proposed method, SNAS matches the previous SOTA result (DrNAS) on CIFAR-10 and CIFAR-100.

heading: 4.2 DARTS SPACE
text: We further investigate the performance of the proposed method on the DARTS search space. To encourage fair comparisons with prior arts, we follow the same search and retrain settings as the original DARTS (Liu et al., 2019). Similar to the experiments on NASBench-201, we also run the search phase of each method under four random seeds (0-3) and report the best and average accuracy of all derived architectures, as well as the error bar to capture the variance of the search algorithms 2. For GM-NAS, we select three edges (T = 3) in total based on the edge importance
2Note that previous methods usually pick the best architecture from up to 10 search runs and report the standard deviation of only the evaluation phase. Such reporting schema does not capture the variance of the search algorithm, and thus biases toward highly unstable methods.
measure introduced in Section 3.4 and split the operations on each selected edge into two groups (B = 2). This leads to eight sub-supernets, comparable to the seven supernets used in the FewShot NAS baseline. We also restrict our total search cost to match that of Few-Shot NAS for fair comparisons. We refer the readers to Appendix C for more details about the settings.
As shown in Table 4, GM-NAS consistently outperforms Few-Shot NAS on both variants of DARTS and also SNAS. For instance, GM DARTS (1st) achieves a 2.35% test error rate, 0.13% lower than Few-Shot DARTS (1st). In addition, GM-NAS also achieves significantly better average test accuracy (Avg column) than Few-Shot and One-Shot NAS, which shows the robustness of our search algorithm under different random seeds. Notably, the best test error we obtain across different base methods is 2.34% (GM-SNAS), ranking top among prior arts.

heading: 4.3 MOBILENET SPACE
text: In addition to the cell-based search spaces, we also evaluate the effectiveness of GM-NAS on the MobileNet Space. Following Few-Shot NAS, we apply GM-NAS to two sampling-based methods - ProxylessNAS (Cai et al., 2019) and OFA (Cai et al., 2020). To match the total number of subsupernets of GM-NAS with Few-Shot NAS, we select two layers (T = 2) to perform the supernet partitioning, and divide the operations into two groups (B = 2) for the first edge, and three groups (B = 3) on the second edge.
The results are summarized in Table 5. When applied to ProxylessNAS, GM-NAS achieves a Top-1 test error rate of 23.4%, out-performing both Few-Shot and One-Shot versions by 0.7% and 1.5%, respectively. On OFA, we obtain a 19.7% Top-1 test error, surpassing all comparable methods within 600-FLOPs latency. The strong empirical results demonstrate GM-NAS’ ability to effectively reduce the harm of weight-sharing in One-Shot Algorithms on large-scale search spaces and datasets, compared with Few-Shot NAS.

heading: 5 ABLATION STUDY
text: In this section, we conduct extensive ablation studies on GM-NAS. Similar to Section 3, we use CIFAR-10 and NASBench-201 space with four operations as it allows us to align the number of supernets with Few-Shot NAS and establish proper comparisons.

heading: 5.1 THE EFFECT OF DIFFERENT COMPONENTS IN THE PROPOSED METHOD
text: Gradient similarity measures We examine the effect of different similarity measures for computing the gradient matching score on the proposed method. We compare cosine similarity with l2 distance, as well as per-filter cosine similarity (Zhao & Bilen, 2021) that computes cosine similarity for each convolution filters separately. As shown in Table 6, per-filter cosine similarity performs similarly to cosine
similarity while l2 distance is not as effective. We therefore adopt cosine similarity for its simplicity.
Edge selection using gradient matching score Instead of randomly selecting edges to partition the supernet on, GM-NAS leverages the cut cost from graph min-cut to decide which edge to partition on. Empirically, we find that this technique reduces the variance and also improves the search performance compared with random selection. GM-NAS with random edge selection obtains 93.58% mean accuracy with a standard derivation of 0.39 on NASBench-201. In contrast, GM-NAS with the proposed edge selection achieves 93.95% accuracy with only a faction of the variance (0.08).
Warmup epochs during the Splitting phase To obtain accurate gradient information, we warmup the supernet for a few epochs during the supernet partitioning phase. As shown in Figure 2, our proposed method is robust to a wide range of warmup epochs. Note that the variance increases without warmup (warmup epoch = 0) due to noisy gradients at initialization, indicating that proper supernet warmups are necessary for GM-NAS.
Restart After the final split is completed, we re-initialize the weights of sub-supernets before conducting architecture search on them. Intuitively, restarting the sub-supernets eliminate the negative effect of weight-sharing estab-
lished during the splitting phase. To verify this, we test GM-NAS without restart while keeping the number of supernet training epochs identical to the GM-NAS baseline. We obtain a mean accuracy of 91.04% on CIFAR-10, which is 2.91% lower than the original GM-NAS, which shows the necessity of supernet restart.

heading: 5.2 RANKING CORRELATION
text: We also evaluate the ranking correlation (measured by Spearman correlation) among top architectures of the proposed method. This is a particularly important measure for effective NAS algorithms as they have to navigate the region of top models to identify the best candidates (Abdelfattah et al., 2021). For Few-Shot NAS, we split two edges, amounting to 42 = 16 sub-supernets. For GM-NAS, due to a smaller branching factor (B = 2), we could afford to split four edges while keeping the total number of supernet the same as Few-Shot NAS. We train each sub-supernet using Random Sampling with Parameter Sharing (RSPS) (Li & Talwalkar, 2020), same as the original Few-Shot NAS paper. Table 7 summarizes the results. The proposed method obtains 0.532 Spearman correlation among top 1% architectures, much higher than Few-Shot NAS (0.117). The improved ranking correlation also justifies the superior performance of GM-RSPS over Few-Shot RSPS reported in Section 4.1.

heading: 6 CONCLUSION
text: In this paper, we demonstrate that gradient similarity can effectively measure the harm of weightsharing among child models. We propose a novel Gradient Matching NAS (GM-NAS) - a generalized supernet splitting schema that utilizes gradient matching score as the splitting criterion and formulates supernet partitioning as a graph clustering problem. Extensive empirical results across multiple prevailing search spaces, datasets, and base methods show that GM-NAS consistently achieves stronger performance than its One-Shot and Few-Shot counterparts, revealing its potential to play an important role in weight-sharing Neural Architecture Search methods.
ETHICS STATEMENT
We do not aware of any potential ethical concerns regarding our work.
REPRODUCIBILITY STATEMENT
We provide a copy of our code in the supplementary material, including both search and retrain phase for our method and the reproduced baseline, to ensure reproducibility on all search spaces. Our experimental setting is stated in Section 4, and hyperparameters are described in the Appendix D. Furthermore, we also include discussions on the reproducibility of relevant baselines on DARTS and MobileNet Space in Appendix D.

heading: A PSEUDOCODE FOR GM-NAS
text: Algorithm 1 summarizes the main pipeline of the proposed Generalized Supernet Splitting with Gradient Matching (GM-NAS). The pseudocode for our partitioning algorithm is provided separately in Algorithm 2 for ease of reference. Note that the main difference between GM-NAS and Few-Shot NAS lies in the supernet splitting phase: Compared with Few-Shot NAS’ exhaustive partitioning, GM-NAS leverages the well-motivated gradient matching score and the graph min-cut algorithm for making much more informed splitting decisions.
Algorithm 1: Main Pipeline Input: The set of supernet A = {A}, warmup epochs warm epo for supernet splitting, number of splits T , Branching factor B per split. // supernet splitting phase
1 A = {A0}; 2 forall t = 1, · · · , T do 3 A′ = {}; 4 forall A ∈ A do 5 Train A for warm epo epochs;
6 A′ insert←−−−− GM SPLIT (A, branch factor = B); 7 end 8 A = A′; 9 end // search phase
10 reinitialize w in A; 11 perform architecture search on A using corresponding base methods;
Algorithm 2: Supernet Splitting with Gradient Matching, GM SPLIT (A,B) Input: Supernet A, Branching factor B
1 forall unsplitted edge e on A do 2 forall operation o on e do 3 temporarily enable o while disabling other operations on e; 4 evaluate A’s gradient on other shared edges, averaged over M mini-batches; 5 end 6 calculate the gradient matching scores based on Eqn.(2); 7 compute e’s importance score and operation partition based on Eqn.(3); 8 end 9 Select the edge (e∗) with the best importance score to split A on;
10 From e∗, partition A into B sub-supernets {Ab}b=1···B ; 11 return {Ab}b=1···B

heading: B COMPLEMENTARY ANALYSIS
text: B.1 WHAT GRADIENT TELLS US ABOUT WHICH OPERATIONS SHOULD (NOT) SHARE WEIGHTS?
On NASBench-201, we find that gradient matching score always decides to assign conv 1x1 and conv 3x3 into one partition (sub-supernet), and Skip and AvgPool 3x3 into another. This behavior is reasonable: High-performing architectures on NASBench-201 usually contains more parametric operations; And weight-sharing between these high performers and the rest of poor architectures could be quite harmful to the former. GM-NAS breaks the weight-sharing between high performers and the rest of the architectures. This has been shown to improve the performance estimation of top models, leading to better search performance (Zhang et al., 2020b).
This is no longer the case on the more complex DARTS Space, where the number of parametric operations is less correlated with the search performance. We observe that three of the four convolution operations are usually assigned to one group, while the other group consists of three non-parametric operations and one convolution. And the specific convolution operation assigned to the second group varies.
For MobileNet space consisting of structured convolution blocks, we did not observe any particular patterns in terms of how GM-NAS decides to partition the supernet. Still, the superior performance of GM-NAS over Few-Shot NAS and One-Shot NAS in Table 5 demonstrates that GM-NAS can make effective partitioning decisions on this search space.
B.2 MORE DETAILS ON THE PRELIMINARY EXPERIMENT IN TABLE 2
We provide further details for the experiment in Table 2. Given a target architecture A, we construct Asim (and Adissim) by changing the operation on a randomly selected edge of A to another operation so that the cosine similarity between gradients computed from A and from Asim (Adissim) at shared weight are large (small). Concretely, cos(∇wsL(A;ws),∇wsL(Asim;ws)) > 0.7 for Asim and cos(∇wsL(A;ws),∇wsL(Adissim;ws)) < 0.3 for Adissim, where ws is their shared weight. Note that the gradient similarities are computed by averaging over 100 mini-batches after updating A together with Asim (or Adissim) for 2 epochs, in order to obtain a more accurate estimation. We then record and compare the training losses of A when 1) it is updated together with Asim and 2) it is updated together with Adissim for 20 epochs. The above process is repeated for 50 randomly sampled A, and the mean training losses and gradient similarities are reported in Table 2. As expected, the training loss of A is much lower in case (1) when A shares weights with Asim, since the training dynamics of A and Asim are similar at the shared weight.

heading: C SUPERNET SELECTION
text: In this section, we address the problem of deriving the final architecture from the set of partitioned sub-supernets, which we term Supernet Selection. Recall that Few-Shot NAS (and our method) splits the supernets into N sub-supernets, and then performs architecture search over these sub-supernets by training them independently from scratch before deriving a single final architecture from them.
Few-Shot NAS argues that the architecture selected from the sub-supernet with the lowest validation loss is typically the best among architectures derived from all sub-supernets. However, empirically we find that this is often not the case, especially on irregular search spaces like the cell-based DARTS Space. As shown in Table 8, for Few-Shot NAS, the architecture selected from the sub-supernet with lowest validation loss (or accuracy) does not match the top architecture from all sub-supernets (”max” entry in the Table). This result aligns with previous findings that the supernet’s performance is unrelated to the final subnetwork accuracy (Li et al., 2020a). There could be multiple potential reasons for this behavior. For example, the sub-supernet that hosts the top child architecture might also contain many mediocre architectures, so the performance of this sub-supernet as a whole might not necessarily top other sub-supernets. Moreover, after the supernet splitting phase, architecture search is performed on these sub-supernets independently, so their performance are also subject to randomness in the search phase.
One could always evaluate each architecture derived fromN sub-supernets and pick the best one out of them, which leads to extra overhead asN increases. To solve this, we propose to adopt Successive Halving (Karnin et al., 2013; Jamieson & Talwalkar, 2016; Li et al., 2019b) to reduce the supernet selection cost. Successive Halving progressively discards half of the poor architectures following a predefined schedule, and stops until only one candidate is left. For DARTS Space, we set the schedule to (30, 100, 600), so that the overhead of retraining N = 8 architectures is only 2.5× that of retraining a single architecture (30 ∗ 8 + 70 ∗ 4 + 500 ∗ 2 = 1520 epochs). Empirically, we find that varying this schedule has a negligible effect on the recall of successive halving schedules, which aligns with previous discoveries on other NAS search spaces (Wang et al., 2021a). As shown in Table 8, successive halving produces essentially the same results for all methods. We summarize the detailed successive halving procedure in Algorithm. 3.
Note that we only perform successive halving for the DARTS Space; and we apply it to both GMNAS and Few-Shot NAS for fair comparisons. For MobileNet Space, we follow Few-Shot NAS and
use the validation loss for supernet selection, as it produces good enough results and speedup the experiments.
Algorithm 3: Supernet Selection via Successive Halving Input: A candidate pool |P| with N child networks derived from N sub-supernets, checkpoint schedules ckpts = {epo1, epo2, ..., epoT } Result: a single network from P
1 epoch epo = 1; 2 while |P| > 1 do 3 train each network in P for one epoch 4 epo += 1; 5 if epo ∈ ckpts then 6 calculate the validation accuracy for each network in P ; 7 discard the bottom half of networks from P based on their validation accuracy; 8 end 9 return the architecture left in P;
D IMPLEMENTATION DETAILS
D.1 NASBENCH-201
For experiments in Section 4.1, we set the warmup epoch to 15 for DARTS and SNAS, and 20 for RSPS as its single path nature requires longer training (Zhang et al., 2020a). We then perform architecture search on each sub-supernet for 30 epochs following the same protocol of the corresponding base algorithms. The search is conducted on three datasets separately for four random seeds. For all methods, we select the best architecture from the sub-supernets as the final architecture.
D.2 DARTS SPACE
Supernet Partition During the supernet (with 8 cells) splitting phase, we set the warmup epoch to 2, the number of splits to 3, and the branching factor to 2. After the supernet splitting phase is finished, we conduct architecture search on the generated sub-supernets for 15 epochs. This way, the total number of epochs is 2 ∗ 1 + 2 ∗ 2 + 2 ∗ 4 + 15 ∗ 8 = 134, similar to Few-Shot NAS (train each of the 7 sub-supernets for 20-25 epochs according to the author).
Supernet Selection After searching architectures on each sub-supernet, we apply the successive halving (Appendix C) to select the top-performed child network derived from eight sub-supernets.
Retraining settings (architecture evaluation) To establish a fair comparison with prior arts, we strictly follow the retrain settings of DARTS (Liu et al., 2019) to evaluate the searched architecture. Concretely, we stack 20 cells to compose the final derived architecture and set the initial channel number as 36. The derived architecture is trained from scratch with a batch size 96 for 600 epochs. We use SGD with an initial learning rate of 0.0025, a momentum of 0.9, and a weight decay of 3× 10−4, and a cosine learning rate scheduler. In addition, we also deploy the cutout regularization with length 16, drop-path with probability 0.3, and an auxiliary tower of weight 0.4.
Reproducing Few-Shot NAS Baseline Since Few-Shot NAS does not release its search code for the DARTS Space, we follow the search settings in the original Few-Shot NAS paper to reproduce its result: We randomly select one edge and split the supernet into 7 sub-supernets, and conduct architecture search on each sub-supernet for 20 epochs. For a fair comparison with GM-NAS, we also use the successive halving (Appendix C) to select the best architecture from these sub-supernets. The reproduced results we obtain for Few-Shot DARTS is 2.48% (the best column in Table 4). This is worse than the one reported in Few-Shot NAS (2.31%) because Few-Shot NAS uses its own retraining protocol to train and evaluate their search architecture, rather than the DARTS’ protocol widely adopted by previous methods (confirmed with the author). Evaluating their released best architecture under DARTS protocol results in 2.44% test error, similar to our reproduced result (2.48%). Nonetheless, GM DARTS achieves 2.35% error rate on DARTS, substantially lower than both of these numbers.
D.3 MOBILENET
GM ProxylessNAS / OFA For both ProxylessNAS and OFA, we set the warmup epoch of GMNAS to 40. We perform split the supernet twice in total, and set the branching factor to 2 for the first split and 3 for the second split, so that the resulting number of sub-supernets match that of Few-shot NAS. After the supernet partitioning phase, for GM ProxylessNAS, we use the same search settings as ProxylessNAS, except for setting the number of search epochs to 40; For GM OFA, we also follow the same search settings as the original paper.
Supernet Selection As mentioned in Appendix C, for Mobilenet Space, we select the sub-supernet with the lowest validation loss and derive the best architecture from this supernet, same as Few-Shot NAS.
Retraining settings (architecture evaluation) For ProxylessNAS, we follow the settings of FewShot ProxylessNAS (Zhao et al., 2021; Cai et al., 2019) to train our discovered architecture. Since the architecture evaluation codes for OFA (finetune) and Few-Shot OFA (retrain) are not released, we also use the setting of Few-Shot ProxylessNAS to train and evaluate the derived architecture of OFA (reproduced), Few-Shot OFA (reproduced), and GM-OFA.
Reproducing OFA and Few-Shot OFA baseline OFA (Cai et al., 2020) does not release the predictor training and child network finetune code. So we follow the settings of Few-Shot ProxylessNAS to reproduce the OFA result. Our reproduced accuracy is 79.7%, comparable to the one (80.0%) reported in the OFA paper.
Few-Shot NAS (Zhao et al., 2021) also does not release their code and the searched architectures on OFA. We try our best to reproduce their result by communicating with the author, but are still
not able to reproduce the reported 80.5% accuracy on ImageNet. One potential reason could be that Few-Shot NAS uses an unreleased powerful teacher network to train the derived architectures. Since we could not obtain this teacher network, we follow the Few-Shot ProxylessNAS setting described above to train and evaluate searched architectures for OFA, Few-Shot OFA, and our methods. The reproduced numbers are reported in Table 5.
E SEARCHED ARCHITECTURES

heading: F ABLATION STUDY ON THE NUMBER OF SPLITS
text: We conducted extra ablation study on how the number of splits affect the search outcome of the proposed method using NASBench-201 and CIFAR-10. As showing in Figure 17, GM+DARTS reaches a near oracle search performance with only two splits. Splitting one more time pushes the performance further, but may not worth the extra search cost.

Review summaries:
This paper improves the sub-supernet splitting strategy in few-shot NAS with a gradient matching (GM) score. During supernet optimization, the GM explicitly computes the "agreements" among different operations on an edge in the supernet, which is then treated as the criterion to split those most mismatched operations into different sub-supernets. Although the basic intuition makes some good points and is easy to grasp, I have several concerns on the methodology and the experimental designs of this paper. 

This paper proposed a method for few-shot NAS tasks. It found some operations may behave similarly and can be grouped into one and others may not. Those patterns should be considered to save the computation as well as improve accuracy. It proposed a graph construction and spitting pipeline to achieve the goal. Specifically, they used gradient matching and cosine distance to construct the graph among layers. Then, splitting the graph by solving a graph clustering problem.   It conducted experiments on three benchmarks (cifar10, cifar100, and ImageNet) over three search spaces (NASBench-201, DARTS, MobileNet Space) and compared with various baselines. The results show the strength of the proposed method. 

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.

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.