Z I C O : Z ERO - SHOT NAS VIA INVERSE C OEFFICIENT OF V ARIATION ON G RADIENTS
Anonymous authors
Paper under double-blind review
A BSTRACT
Neural Architecture Search (NAS) is widely used to automatically design the neural network with the best performance among a large number of candidate architectures.
To reduce the search time, zero-shot NAS aims at designing training-free proxies that can predict the test performance of a given architecture.
However, as shown recently, none of the zero-shot proxies proposed to date can actually work consistently better than a naive proxy, namely, the number of network parameters (#Params).
To improve this state of affairs, as the main theoretical contribution, we first reveal how some specific gradient properties across different samples impact the convergence rate and generalization capacity of neural networks.
Based on this theoretical analysis, we propose a new zero-shot proxy, ZiCo , the first proxy that works consistently better than #Params.
We demonstrate that ZiCo works better than State-Of-The-Art (SOTA) proxies on several popular NAS-Benchmarks (NASBench101, NATSBench-SSS/TSS, TransNASBench-101) for multiple applications (e.g., image classification/reconstruction and pixel-level prediction).
Finally, we demonstrate that the optimal architectures found via ZiCo are as competitive as the ones found by one-shot and multi-shot NAS methods, but with much less search time.
For example, ZiCo-based NAS can find optimal architectures with 78.1%, 79.4%, and 80.4% test accuracy under inference budgets of 450M, 600M, and 1000M FLOPs on ImageNet within 0.4 GPU days.
I NTRODUCTION
During the last decade, deep learning has achieved great success in many areas, such as computer vision and natural language modeling Krizhevsky et al.
; Liu & Deng (2015); Huang et al.
; He et al.
; Dosovitskiy et al.
; Brown et al.
Vaswani et al.
In recent years, neural architecture search (NAS) has been proposed to search for optimal architectures, while reducing the trial-and-error (manual) network design efforts Baker et al.
; Zoph & Le (2017); Elsken et al.
Moreover, the neural architectures found via NAS show better performance than the manually-designed networks in many mainstream applications Real et al.
Gong et al.
; Xie et al.
; Wu et al.
; Wan et al.
(2020); Li & Talwalkar (2020);
Kandasamy et al.
; Yu et al.
(2020b)
; Liu et al.
(2018b)
; Cai et al.
; Zhang et al.
(2019a);
Zhou et al.
; Howard et al.
Despite these advantages, most of the existing NAS approaches involve a time-consuming and resource-intensive search process.
For example, multi-shot NAS uses a controller or an accuracy predictor to conduct the search process and it requires training of multiple networks; thus, multi-shot
NAS is extremely time-consuming (typically thousands of GPU hours)
Real et al.
; Chiang et al.
Alternatively, one-shot NAS merges all possible networks from the search space into a supernet and thus only needs to train the supernet once Dong & Yang (2019); Zela et al.
Chen et al.
; Cai et al.
; Stamoulis et al.
; Chu et al.
; Guo et al.
; this enables one-shot NAS to find a good architecture with much less search time.
Though the oneshot NAS has significantly improved the time efficiency of NAS, training is still required during the search process.
In the last few years, the zero-shot approaches have been proposed to liberate NAS from training entirely Wu et al.
; Zhou et al.
(2022; 2020); Ingolfsson et al.
; Tran & Bae (2021); Do &
Luong (2021); Tran et al.
Shu et al. (2022b).
Essentially, zero-shot NAS utilizes some proxy that can predict the test performance of a given network without training .
Moreover, the design of the proxy in zero-shot NAS is usually based on some theoretical analysis of deep networks.
Hence, zero-shot approaches can not only significantly improve the time efficiency of NAS, but also deepen the theoretical understanding on why certain networks work well.
Nonetheless, as revealed in Ning et al.
; White et al.
(2022), the zero-shot proxies proposed to date cannot work consistently better than a naive proxy, namely, the number of parameters (#Params); in fact, #Params often achieves the best performance on most of popular NAS benchmarks.
These results may undermine the effectiveness of zero-shot NAS approaches.
To address the limitations of existing zero-shot proxies, we target the following key questions :
How do some specific gradient properties,
i
e., mean value and standard deviation across different samples, impact the training convergence of neural networks?
Can we use these two gradient properties to design a new theoretically-grounded proxy that works better than #Params consistently?
To this end, we first theoretically analyze how the mean value and standard deviation of gradients across different training batches impact the training convergence of neural networks.
Based on our analysis, we propose ZiCo , a new proxy for zero-shot NAS.
We demonstrate that, compared to all existing proxies (including #Params), ZiCo has either a higher or at least on-par correlation with the test accuracy on popular NAS-Benchmarks (NASBench101, NATS-Bench-SSS/TSS) for multiple datasets (CIFAR10/100, ImageNet16-120).
Finally, we demonstrate that ZiCo enables a zero-shot NAS framework that can efficiently find the network architectures with highest test accuracy compared to other zero-shot baselines.
In fact, our ZiCo-based zero-shot NAS framework achieves competitive FLOPs-accuracy tradeoffs compared to multiple one-shot and multi-shot NAS, but with much lower time costs.
To summarize, we make the following major contributions :
• We theoretically reveal how the mean value and variance of gradients across multiple samples impact the training convergence and generalization capacity of neural networks.
• We propose a new zero-shot proxy, ZiCo, that works better than existing proxies on popular NAS-Benchmarks (NASBench101, NATS-Bench-SSS/TSS, TransNASBench-101) for multiple applications (e.g. image classification/reconstruction and pixel-level prediction).
• We demonstrate that our proposed zero-shot NAS achieves competitive test accuracy with representative one-shot and multi-shot NAS with much less search time.
The rest of the paper is organized as follows.
We discuss related work in Section 2.
In Section 3, we introduce our theoretical analysis.
We introduce our proposed zero-shot proxy (ZiCo) and the NAS framework in Section 4.
Section 5 validates our analysis and presents our results with the proposed zero-shot NAS.
We conclude the paper in Section 6 with remarks on our main contribution.
R ELATED WORK 2. Z ERO - SHOT NAS
The goal of zero-shot NAS is to rank the accuracy of various candidate network architectures without training , such that we can replace the expensive training process in NAS with some computationefficient proxies Xiang et al. (2021a)
; Javaheripi et al.
; Bhardwaj et al.
; Li et al.
Hence, the quality of the proxy determines the effectiveness of zero-shot NAS.
Several works use the number of linear regions to approximately measure the expressivity of a deep neural network Mellor et al.
; Chen et al.
(2021b)
; Bhardwaj et al.
Alternatively, most of the existing proxies are derived from the gradient of deep networks.
For example, Synflow, SNIP, and GraSP rely on the gradient w.r.t the parameters of neural networks; they are proved to be the different approximations of Taylor expansion of deep neural networks Abdelfattah et al.
; Lee et al.
(2019b);
Tanaka et al.
; Wang et al.
Moreover, the Zen-score approximates the gradient w.r.t featuremaps and measures the complexity of neural networks Lin et al.
; Sun et al.
Furthermore, Jacob cov leverages the Jacobian matrix between the loss and multiple input samples to quantify the capacity of modeling the complex functions Lopes et al.
Though zero-shot
NAS can significantly accelerate the NAS process, it has been revealed that the naive proxy #Params generally works better than all the proxies proposed to date Ning et al.
; White et al.
These limitations of existing proxies motivate us to look for a new proxy that can consistently work better than #Params and address the limitations of existing zero-shot NAS approaches.
2. K ERNEL M ETHODS AND C ONVERGENCE A NALYSIS
Kernel methods are widely explored to analyze the convergence property of networks trained with gradient descent Neal (1996);
Williams (1996); Du et al. (2019a); Lu et al.
Allen-Zhu et al.
; Hanin & Nica (2020); Golikov et al.
For example, the training of wide neural networks is proved to be equivalent to the optimization of a specific kernel function Arora et al.
(2019a)
; Lee et al. (2019a); Chizat et al.
Arora et al. (2019b); Cho & Saul (2009).
Moreover, given the networks with specific width constraints, researchers proved that the training convergence of networks can be described by some corresponding kernels and the convergence rates of training are highly coupled with the eigenvalues of the kernel-based covariance matrix Mei et al.
Zhang et al. (2019b); Garriga-Alonso et al.
Du et al. (2019b).
In our work, we extend such kernel-based analysis to reveal the relationships between the gradient properties and the training convergence for neural networks.
C ONVERGENCE AND G ENERALIZATION VIA G RADIENT A NALYSIS
We consider the mean value and standard deviation of gradients across different samples and first explore how these two metrics impact the training convergence of linear regression tasks.
3. L INEAR R EGRESSION
Inspired by Du et al.
(2019b), we use the training set S with M samples as follows:
[Equation]
where R is a positive constant and || · || denotes the L 2 -norm of a given vector; x i ∈ R
d is the i th input sample and normalized by its L 2 -norm (i.e., || x i ||
= 1 )
, and y i is the corresponding label.
We define the following linear model f = a T x optimized with an MSE-based loss function L :
[Equation]
where a ∈ R d is the initial weight vector of f .
We denote the gradient of L w.r.t to a as g ( x i ) when taking ( x i , y i ) as the training sample:
[Equation]
We denote the j th element of g ( x i ) as g j ( x i ) .
We compute the mean value ( µ j ) and standard deviation ( σ j ) of g j ( x i ) across all training samples as follows:
[Equation]
Theorem 3.1.
We denote the updated weight vector as ˆ a and denote (cid:80) ij [ g j ( x i )] 2 = G .
Assume we use the accumulated gradient of all training samples and learning rate η to update the initial weight vector a , i.e., ˆ a = a − η (cid:80) i g ( x i ) .
If the learning rate 0 < η < 2 , then the total training loss is bounded as follows:
[Equation]
In particular, if the learning rate η = 1 M , then L (ˆ a ) is bounded by:
[Equation]
We provide the proof in Appendix A and the experimental results to validate this theorem in Sec 5.2.
Remark 3.1 Intuitively, Theorem.
3.1 tells us that the higher the gradient absolute mean across different training samples, the lower the training loss the model converges to; i.e., the network converges at a faster rate.
Similarly, if ηM < 1 , the smaller the gradient standard deviation across different training samples/batches, the lower the training loss the model can achieve.
3. MLP S WITH R E LU
In this section, we generalize the linear model to a network with ReLU activation functions.
We primarily consider the standard deviation of gradients in the Gaussian kernel space.
We still focus on the regression task on the training set S defined in Eq.
We consider a neural network in the same form as Du et al.
(2019b):
[Equation]
where m is the number of output neurons of the first layer; s r is the r th element in the output weight vector s ; W ∈ R m × d is the input weight matrix, and w r ∈ R d is the r th row weight vector in W .
For training on the dataset S with M samples defined in Eq.
1, we minimize the following loss function:
[Equation]
Following the common practice Du et al. (2019b), we fix the second layer ( s ) and use gradient descent to optimize the first layer ( W ) with a learning rate η :
[Equation]
where W ( t − 1) denote the input weight matrix after t − 1 training steps; w r ( t ) denote the r th row weight vector after t training steps.
Definition 1.
(Gram Matrix)
A Gram Matrix H ( t ) ∈
R M × M on the training set { ( x i , y i ) , i = 1 , ..., M } after t training steps is defined as follows:
[Equation]
where I is the indicator function and I {A} =
if and only if event A happens.
We denote the λ min ( H ) as the minimal eigenvalue of a given matrix H .
We denote the λ 0 = λ min ( H ( ∞ )) .
Theorem 3.2.
Given a neural network with ReLU activation function optimized by minimizing Eq.
8, we assume that each initial weight vector { w r (0) , r = 1 , ..., n } is i.i.d.
generated from N (0 , I ) and the gradient for each weight follows i.i.d.
N (0 , σ ) , where the σ is measured across different training steps.
For some positive constants δ and ϵ , if the learning rate η satisfies η <
[Equation]
then with with probability at least (1 − δ )(1 − ϵ ) , the following holds true: for any r ∈ [ m ] , ||
w r (0) − w r ( t ) || ≤ C = ηtσ (cid:112) Φ(1 − ϵ ) , and at training step t the Gram matrix H ( t ) satisfies:
[Equation]
Φ( · ) is the inverse cumulative distribution function for a d -degree chi-squared distribution χ 2 ( d ) .
We provide the proof in Appendix B. We now introduce the following conclusion from Du et al.
(2019b) to further help our analysis.
Lemma 1.
Du et al.
(2019b)
Assume we set the number of output neurons of the first layer m =
Ω( M 6 λ 40 δ 3 ) and we i.i.d.
initialize w r ∼ N (0 , I ) and s r ∼ uniform [ {− 1 , 1 } ] , for r ∈ [ m ] .
When minimizing the loss function in Eq.
8 on the training set S in Eq.
1, with probability at least 1 − δ over the initialization, the training loss after t training steps is bounded by:
[Equation] [Equation]
Theorem 3.3.
Under the assumptions of Theorem 3.2 and Lemma 1, with probability at least (1 − δ )(1 − ϵ ) , the following inequality holds true:
[Equation]
The proof consists of replacing λ min ( H ( t )) in Eq.
12 with its lower bound given by Theorem 3.2.
Remark 3.2 Theorem.
3.3 shows that after some training steps t , the network with a smaller standard deviation ( σ ) of gradients will have a smaller training loss; i.e., the network has a faster convergence rate at each training step.
We further validate this theorem in Sec. 5.2.
We provide the generalization capacity analysis and experimental results in Appendix C.
3. S UMMARY OF OUR THEORETICAL ANALYSIS
Theorem 3.1 and Theorem 3.3 tell us that the network with high training convergence speed should have high absolute mean values and low standard deviation values for the gradient, w.r.t the parameters across different training samples/batches.
Inspired by these theoretical insights, we next propose a proxy that jointly considers both absolute mean and standard deviation values.
N EW ZERO - SHOT PROXY AND NAS FRAMEWORK
In this section, we first define our proxy (ZiCo) and then introduce our zero-shot NAS framework.
Following the standard practice, we consider convolutional neural networks (CNNs) as candidate networks.
4. P ROPOSED ZERO - SHOT PROXY : Z I C O
Definition 2.
Given a neural network with D layers and loss function L , the Zero-shot inverse
Coefficient of Variation (ZiCo) is defined as follows:
[Equation]
where Θ denote the initial parameters of the given network; θ l denote the parameters of the l th layer of the network, and ω represents each element in θ l ; X i and y i are the i th input batch and corresponding labels from the training set; N is number of training batches used to compute ZiCo.
We incorporate log to stabilize the computation by regularizing the extremely large or small values.
Of note, our metric is applicable to general CNNs; i.e., there’s no restriction w.r.t.
the neural architecture when calculating ZiCo.
As discussed in Section 3.3, the networks with higher ZiCo tend to have better convergence rates.
Hence, the architectures with higher ZiCo are better architectures.
We remark that the loss values in Eq.
14 are all computed with the initial parameters Θ ; that is, we never update the value of the parameters when computing ZiCo for a given network (hence it follows the basic principle of zero-shot NAS, i.e., never train, and only use the initial parameters).
In practice, two batches are enough to make ZiCo achieve the SOTA performance among all previously proposed accuracy proxies (see Sec. 5.5).
Hence, we use only two input batches ( N = 2 ) to compute
ZiCo; this makes ZiCo highly time efficient for a given network.
4. Z I C O - BASED ZERO - SHOT NAS
INPUT: Number of search steps T
Inference budget B , Search space S
Set of input batch Z = { ( X i , y i ) , i = 1 , 2 }
Population size E , Initial network F 0 ∈ S OUTPUT:
Optimal network F P
SEARCH:
Initialize F = { F 0 } for i = 1 to T do
Randomly sample network F t from F F i = randomly mutated architecture based on F t from S if F i meets the inference budget B then
Compute ZiCo for F i on Z by Eq.
Add F i to F if | F | > E then
Remove network with the smallest ZiCo from F end if end if end for F P = the network of the highest ZiCo in F .
We use an Evolutionary Algorithm (EA) to conduct the zero-shot NAS because it is concise and easy to implement 1 As shown in Algorithm 1, we search for the neural architectures with the highest ZiCo within the search space, given a specific budget B (e.g., FLOPs).
We repeat the search T times; at each search step, we randomly select a structure from the candidate set F and mutate its architectures (e.g., kernel size, block type, number of blocks, and layer width) to generate a new network F i ∈ S .
If the generated network F i meets the inference budget B , we calculate its ZiCo on Z and add F i to the candidate set F .
We remove the network with the smallest ZiCo from F , if the number of architectures in F exceeds the threshold
E .
After T steps, we select the network with the largest ZiCo as the final (
[Table]
[Figure]
E XPERIMENTAL R ESULTS 5. E XPERIMENTAL S ETUP
We conduct the following types of experiments: ( i ) Empirical validation of Theorem 3.1 and Theorem 3.3, ( ii ) Evaluation of the proposed ZiCo on multiple NAS benchmarks, ( iii ) Illustration of
ZiCo-based zero-shot NAS on ImageNet.
For the experiments ( i ) to validate Theorem 3.
we optimize a linear model as in Eq.
2 on the
MNIST dataset, the mean gradient values and the standard deviation vs. the total training loss.
Moreover, we also optimize the model defined by Eq.
7 on MNIST and report the training loss vs. the standard deviation in order to validate Theorem 3.2.
For experiments ( ii ), we compare our proposed ZiCo against existing proxies on three mainstream
NAS benchmarks: NATSBench is a popular cell-based search space with two different search spaces: (1) NATSBench-TSS consisting of 15625 total architectures with different cell structures trained on CIFAR10, CIFAR100, and ImageNet16-120 (Img16-120) datasets, which is just renamed from NASBench-201 Dong & Yang (2020); (2) NATSBench-SSS contains includes 32768 architectures (which differ only in the width values of each layer) and is also trained on the same three above datasets Dong et al.
NASBench101 provides users with 423k neural architectures with their test accuracy on CIFAR10 dataset; the architectures are built by stacking the same cell multiple times Ying et al. TransNASBench- 101-Mirco contains 4096 networks with different cell structures on various downstream applications (see Appendix E.2) Duan et al.
For experiments ( iii ), we use ZiCo to conduct the zero-shot NAS (see Algorithm 1) on ImageNet.
We first use Algorithm 1 to find the networks with the highest ZiCo under various FLOPs budgets.
We conduct the search for 100k steps; this takes 10 hours on a single NVIDIA 3090 GPU (i.e., 0.
GPU days).
Then, we train the obtained network with the exact same training setup as Lin et al.
Specifically, we train the neural network for 480 epochs with the batch size 512 and input resolution 224.
We also use the distillation-based training loss functions by taking Efficient-B3 as the teacher.
Finally, we set the initial learning rate as 0.1 with a cosine annealing scheduling scheme.
5. V ALIDATION OF T HEOREM 3.1 AND T HEOREM 3.
To empirically validate Theorem 3.1, we first create the training set S by normalizing randomly sampled 1000 training samples in MNIST and normalizing them with their L 2 -norm.
We compute the gradient w.r.t.
the network parameters for each individual training sample.
Next, as discussed in
Theorem 3.1, we use the accumulated gradient over these samples to update the network parameters with learning rate η = 1 .
Then, we calculate the square sum of mean gradients and the total training loss.
We repeat the above process 1000 times on the same S .
As shown in Fig. 1(a), we plot the total training loss vs. square sum of mean gradients as defined in Eq.
Clearly, the networks with the higher square sum of mean gradients values tend to have lower training loss.
In comparison, Fig. 1(b) shows that networks with a lower square sum of variance value tend to have lower training loss values, which coincides with the conclusion drawn from Eq.
These results empirically validate our Theorem 3.1.
[Table]
[Figure]
Moreover, to optimize a two-layer MLP with ReLU activation functions as defined in Eq.
7, we use the entire training set of MNIST and apply the gradient descent (Eq. 9) to update the weights.
We set the batch size as 256 and measure the standard deviation of gradients ( σ ) w.r.t.
parameters across different training batches.
We set a very small learning rate η = 10 − 8 to satisfy the assumption in
Theorem 3.2.
We plot the training loss after one training epoch vs. standard deviation of gradients ( σ ) in Fig. 1(c).
Clearly, the results show that if a network has a lower gradient standard deviation, then it tends to have lower training loss values, and thus, a faster convergence rate.
These results empirically prove our claims in Theorem 3.3.
5. Z I C O VS . OTHER PROXIES ON NAS BENCHMARKS
We first calculate the correlation coefficients between various proxies and the test accuracy on
CIFAR10, CIFAR100, and ImageNet16-120 datasets for NATSBench.
As shown in Table.
ZiCo achieves the highest correlation with the real rest accuracy, except for CIFAR10/100 on
NATSBench-SSS.
Moreover, we observe that most of the previously proposed proxies work well for some specific scenarios, but do not generalize well to other scenarios.
For example, though Synflow is slightly better than ZiCo for CIFAR10/100 on NATSBench-SSS, it has poor correlation scores in Img16-120.
Similarly, Zen-score performs well for Img16-120 on NATSBench-SSS, but it doesn’t work well on NATSBench-TSS.
In contrast, ZiCo has either the highest or second highest correlation coefficients under all these scenarios in Table 1.
We provide more results in Appendix E.
For NASBench101, as shown in Fig. 2, ZiCo has a significantly higher correlation score with the real test accuracy than all the other proxies, except Zen-score.
For example, ZiCo has a 0.46 Kendall’s τ [Table] [Table] score, while #Params is only 0.31.
In general, ZiCo has the highest correlation coefficients among all existing proxies for various search spaces and datasets of NATSBench and NASBench101.
Beside the correlation coefficients, we also report the optimal architectures found with various proxies.
As shown in Table 2, the architectures found via ZiCo have the highest test accuracy on all these three datasets.
To our best knowledge, ZiCo is the first proxy that shows a consistently higher correlation coefficient compared to #Params.
The above results validate the effectiveness of our proposed ZiCo; thus, ZiCo can be directly used to search for optimal networks for various budgets.
Next, we describe the search results in detail.
5. Z I C O ON I MAGE N ET
Search Space We use the commonly used MobileNetv2-based search space where the candidate networks are built by stacking multiple Inverted Bottleneck Blocks (IBNs) with SE modules Sandler et al.
; Pham et al.
Lin et al.
As for each IBN, the kernel size of the depthwise convolutional layer is sampled from { 3,5,7 } and the expansion ratio is randomly selected from { 1,2,4,6 } .
We primarily consider ReLU as the activation function.
We use standard Kaiming Init to initialize all linear and convolution layers for every candidate networks He et al.
More details of the search space are given in Appendix D.
We use Algorithm 1 to search networks under various FLOPs budgets (450M, 600M, and 1000M) within the above search space.
As shown in Table 3, ZiCo outperforms most previous NAS approaches by a large margin.
For example, when the FLOPs budget is around 450M, ZiCo achieves 78.1% Top-1 accuracy, which is competitive with one of the SOTA NAS methods (DONNA), but
[Figure]
with fewer FLOPs and 648 × faster search speed Moons et al.
Moreover, if the FLOPs is 600M, ZiCo achieves 2.6% higher Top-1 Accuracy than the latest one-shot NAS method (MAGICAT) with a 3 × reduction in terms of search time Xu et al.
To make further comparison with #Params, we also use #Params as the proxy and Algorithmto conduct the search under a 450M FLOPs budget.
As shown in Table 3, the obtained network by #Params has a 14.6% lower accuracy than ours (63.5% vs. 78.1%).
Hence, even though the correlations for ZiCo and #Params in Table 1 and the optimal networks in Table 2 are similar for small-scale datasets, ZiCo significantly outperforms naive baselines like #Params for large datasets like ImageMet.
To conclude, ZiCo achieves SOTA results for Zero-Shot NAS and outperforms naive methods, existing zero-shot proxies, as well as several one-shot and multi-shot methods.
We remark that these results demonstrate two benefits of our proposed ZiCo: ( i ) Lightweight computation costs.
As discussed in Sec 3, during the search process, to evaluate a given architecture, we only need to conduct the backward propagation twice (only takes 0.3s on an NVIDIA 3090 GPU).
The computation efficiency and exemption of training enable ZiCo to significantly reduce the search time of NAS.
( ii ) High correlation with the real test accuracy.
As demonstrated in Sec 5.3, ZiCo has a very high correlation score with real accuracy for architectures from various search spaces and datasets.
Hence, ZiCo can accurately predict the test accuracy of diverse neural architectures, thus helping find the optimal architectures with the best test performance.
5. A BLATION STUDY
Number of batches We randomly select 2000 networks from NATSBench-TSS on CIFAR dataset and compute ZiCo under varying number of training batches ( N in Eq. 14) from { 2,...,10 } .
We then calculate the correlation score between ZiCo computed under different N values and the real test accuracy.
Fig. 3(a) shows that using two batches to compute ZiCo generates the highest score.
Hence, in our work, we always use two batches ( N = 2 ) to compute ZiCo since it is both accurate and time-efficient.
Batch size We pick 2000 networks from NATSBench-TSS on CIFAR100 and compute ZiCo with two batches under varying batch size { 1,2,4,8,16,32,64,128 } .
We then calculate the correlation score between ZiCo computed under various batch sizes and the real test accuracy.
As shown in Fig. 3(b), batch size 64 is enough to stabilize the coefficient.
Therefore, we set the batch size as 128 and use two batches to compute ZiCo. We provide more ablation studies in Appendix F.
C ONCLUSION
In this work, we have proposed ZiCo, a new SOTA proxy for zero-shot NAS.
As the main theoretical contribution, we first reveal how the mean value and standard deviation of gradients impact the training convergence of a given architecture.
Based on this theoretical analysis, we have shown that ZiCo works better than all zero-shot NAS proxies proposed so far on multiple popular NAS-Benchmarks (NASBench101, NATSBench-SSS/TSS) for multiple datasets (CIFAR10/100, ImageNet16-120).
In particular, we have demonstrated that ZiCo is consistently better than (#Params) and existing zeroshot proxies.
Moreover, ZiCo enables us to find architectures with competitive test performance to representative one-shot and multi-shot NAS methods, but with much lower search costs.
For example, ZiCo-based NAS can find the architectures with 78.1%, 79.4%, and 80.4% test accuracies under 450M, 600M, and 1000M FLOPs budgets on ImageNet within 0.4 GPU days.
A P ROOF OF T HEOREM 3.
Theorem 3.1 We denote the updated weight vector as ˆ a and (cid:80) ij [ g j ( x i )] 2 = G .
Assume we use the accumulated gradient of all training samples and learning rate η to update the initial weight vector a , i.e., ˆ a = a − η (cid:80) i g ( x i ) .
If the learning rate 0 < η < 2 , then the total training loss is bounded as follows:
[Equation]
In particular, if the learning rate η = 1 M , then L (ˆ a ) is bounded by:
[Equation]
Proof.
Given each training sample ( x i , y i ) the gradient of L w.r.t to a when taking ( x i , y i ) as the input is as follows:
[Equation]
We note that:
[Equation]
Then the total training loss among all training samples is given by:
[Equation] [Equation]
By using Eq.
18, we can rewrite Eq.
19 as follows:
[Equation]
Recall the assumption that ˆ a = a − η (cid:80)
i g ( x i ) ; we rewrite Eq.
20 as follows:
[Equation]
According to the Cauchy–Schwarz inequality and || x i ||
= 1 , the total training loss is bounded by:
[Equation]
Since (cid:80) Mi =1 12 (ˆ a T x i − y i ) 2 is always non-negative
[Equation]
the above upper bound of training loss satisfies:
[Equation]
Note that, if 0 < η < 2 , then η (2 − η ) > 0 .
Therefore, the larger (cid:80) j µ 2 j term would make the upper bound of training loss in Eq.
22 closer to 0.
In other words, the higher the gradient absolute mean values across different training samples/batches, the lower the training loss values the model converges to; i.e., the network converges at a faster rate.
In particular, if η = 1 M , the Eq.
22 can be rewritten as:
[Equation]
This completes our proof.
B P ROOF OF T HEOREM 3.
Theorem 3.2 Given a neural network with ReLU activation function optimized by minimizing Eq.
8, we assume that each initial weight vector { w r (0) , r = 1 , ..., n } is i.i.d.
generated from N (0 , I ) and the gradient for each weight follows an i.i.d.
N (0 , σ ) .
For some positive constants δ and ϵ , if the learning rate η satisfies η <
[Equation]
2Φ [Equation] 1 − ϵ ) tσ , then with with probability at least (1 − δ )(1 − ϵ ) , the following holds true: for any r ∈ [ m ] , ||
w r (0) − w r ( t ) || ≤ C = ηtσ (cid:112) Φ(1 − ϵ ) , and at training step t the Gram matrix H ( t ) satisfies:
[Equation]
Φ( · ) is the inverse cumulative distribution function for a d -degree chi-squared distribution χ 2 ( d ) .
Proof.
We first compute the probability of || w r (0) − w r ( t ) || ≤ C .
Based on the assumption w i (0) , i = 1 , ..., n } follows i.i.d.
N (0 , I ) and the gradient for each weight follows i.i.d.
N (0 , σ ) , considering the weight updating rule defined in Eq.
9, each element in w r (0) − w r ( t ) follows a i.i.d.
N (0 , ηtσ ) .
Therefore, || w r (0) − w r ||
η 2 t 2 σfollows the chi-squared distribution with d degrees of freedom χ 2 ( d ) .
[Equation]
Given an input sample x
i and a weight vector w r ( t ) from W ( t ) , we define the following event:
[Equation]
If || w r ( t ) − w r (0) ||
≤ C holds true, x
[Equation] [Equation]
Eq.
28 tells us that if || w r (0) || is larger than || w r ( t ) − w r (0)
|| , then x T i
w r (0) determines the sign value of x T i w r ( t ) ; in other words, x T i w r ( t ) always has the same sign values with x T i w r (0) ; i.e., I { x T i w r (0) ≥ 0 } =
I { x T i w r ( t ) ≥ 0 } .
That is, if || w r ( t ) − w r (0) || ≤ C
and I { x T i w r (0) ≥ 0 } ̸ =
I { x T i w r ( t ) ≥ 0 } hold true, then ||
w r (0) || ≤ C .
Therefore, the probability of event
A ir :
[Equation]
By anti-concentration inequality of Gaussian distribution Du et al.
(2019b), we have:
[Equation]
Therefore, if any weight vector w 1 , ..., w m satisfies || w r (0) − w r ( t ) || ≤ C , we can bound the entry-wise deviation on the Gram matrix H ( t ) at training step t : for any ( i, j ) ∈ [ n ] × [ n ] :
E
[Equation]
where the expectation is summing over the initial weight w (0) .
Hence, considering all the elements in H
[Equation]
we have:
[Equation]
Therefore, by Markov’s inequality, given the probability 1 − δ , we get:
[Equation]
In Du et al.
(2019b)
[Equation] the authors prove that, given a small perturbation K :
[Equation]
, then λ
[Equation]
In our case, K in Eq.
34 is given by 2 M 2 √ 2 C √ πδ .
Therefore,
[Equation]
We replace the term η in Eq.35 with η ’s upper bound given in the assumption of Theorem 3.2, i [Equation] e., η <
[Equation]
, we can get that λ min ( H ( t )) is always larger than 0; that is:
[Equation]
This completes our proof.
C G ENERALIZATION C APACITY VIA G RADIENT A NALYSIS C. G ENERALIZATION VS . E IGENVALUE OF THE H ESSIAN MATRIX
Several prior works show that the generalization capacity of a neural network is highly correlated with its sharpness of the loss function Keskar et al. (2017b;a); Li et al.
; Liang et al.
Usually, a flatter loss landscape leads to a better generalization capacity.
Moreover, it has also been shown that the largest eigenvalue of the Hessian matrix of loss can be used to describe the sharpness of the loss landscape Sagun et al.
(2018); more precisely:
{+Proposition C.
+}
The lower the largest eigenvalue of the Hessian matrix, the higher the generalization capacity of the network.
[Lewkowycz et al.
; Sagun et al.
Next, we analyze how the gradient of a neural network impacts the largest eigenvalues of the Hessian matrix and the generalization capacity of the given network.
C. G RADIENT VS . E IGENVALUE OF H ESSIAN M ATRIX
In the following analysis, we use the same setup of Sec 3.2 and focus on the analysis of a two-layer
MLP with ReLU activation functions.
Specifically, for the model defined in Sec 3.2 (see Eq. 7), authors in Du et al.
(2019b) show that the value of each element in the Hessian matrix is the same as the one in its Gram Matrix (Eq. 10):
[Equation]
where < v 1 , v 2 > represents the inner product of two vectors v 1 and v 2 .
Next, we show how the gradient statistics impact the largest values of the Hessian matrix.
Theorem C.1.
Given a neural network with ReLU activation function optimized by minimizing
Eq. 8, we assume that each initial weight vector { w r (0) , r = 1 , ..., n } is i.i.d.
generated from
N (0 , I ) and the gradient for each weight follows an i.i.d.
distribution N (0 , σ ) .
For some positive constants δ and ϵ , if the learning rate η satisfies η <
[Equation]
with with probability at least (1 − δ )(1 − ϵ ) , the following holds true: for any r ∈ [ m ] , ||
w r (0) − w r ( t ) || ≤ C = ηtσ (cid:112) Φ(1 − ϵ ) , and at training step t , the Hessian matrix H ( t ) satisfies:
[Equation]
Φ( · ) is the inverse cumulative distribution function for a d -degree chi-squared distribution χ 2 ( d ) .
The proof is similar to the proof of Theorem 3.2 (see Appendix B).
We provide the entire proof below.
Proof.
We first compute the probability of || w r (0) − w r ( t ) || ≤ C .
Based on the assumption that { w i (0) , i = 1 , ..., n } follow i.i.d.
N (0 , I ) and the gradient of each weight follows i.i.d.
N (0 , σ ) , considering the weight updating rule defined in Eq.
9 with learning rate η , each element in w r (0) − w r ( t ) follows an i.i.d.
N (0 , ηtσ ) .
{+Therefore, || w r (0) − w r || 
+}
η 2 t 2 σfollows a chi-distribution with d degrees of freedom χ 2 ( d ) :
[Equation]
Given an input sample x
i and a weight vector w r ( t ) from W ( t ) , we define the following event:
[Equation]
If || w r ( t ) − w r (0) ||
≤ C holds true, then:
[Equation]
Eq.
41 implies that if || w r (0) || is larger than || w r ( t ) − w r (0)
|| , then x T i
w r (0) determines the sign value of x T i w r ( t ) .
In other words, x T i w r ( t ) always has the same sign values as x T i w r (0) ; that is, I { x T i w r (0) ≥ 0 }
=
I { x T i w r ( t ) ≥ 0 } .
Hence, if || w r ( t ) − w r (0) || ≤ C
and I { x T i w r (0) ≥ 0 } ̸ =
I { x T i w r ( t ) ≥ 0 } hold true, then ||
w r (0) || ≤ C .
Therefore, the probability of event
A ir :
[Equation]
By the anti-concentration inequality of a Gaussian distribution Du et al.
(2019b), we have:
[Equation]
[Equation]
Therefore, if any weight vector w 1 , ..., w m satisfies || w r (0) − w r ( t ) || ≤ C , we can bound the entry-
[Equation]
deviation on the Gram matrix H ( t ) at the training
[Equation]
t :
[Equation]
( i
[Equation]
[Equation]
[Equation]
We note that all the samples in the training set S (Eq. 1) are normalized with their L2-norm.
Hence, we have both || x i
||
= 1 and ||
x j ||
= 1 .
Therefore, using the Cauchy–Schwarz inequality, the above equation is bounded as follows:
[Equation]
where the expectation is over the initial weight w r (0) , r = { 1 , ..., m } .
Hence, considering all the elements in H , we have:
[Equation]
Therefore, by the Markov’s inequality, given the probability 1 − δ , we get:
[Equation]
Based on the matrix perturbation theory Bauer & Fike (1960); Eisenstat & Ipsen (1998), given a small perturbation K : if
[Equation]
In our case, K in Eq.
48 is given by 2 M 2 √ 2 C √ πδ ; that is:
[Equation]
This completes our proof.
Remark C.2 Theorem.
C.1 shows that after some training steps t , the network with a smaller standard deviation ( σ ) of gradients will have a lower largest eigenvalues of the Hessian matrix;
i
[Figure]
e., the network has a flatter loss landscape rate at each training step.
Therefore, based on Proposition C.1, the model will generalize better.
We further validate this theorem in the following section.
[Figure]
C. E XPERIMENTAL RESULTS : V ALIDATION OF T HEOREM C.
To empirically validate Theorem C.
[Figure]
we first create the training set S by normalizing the training samples in MNIST with their L2-norm.
Next, we optimize a two-layer MLP with ReLU activation functions as defined in Eq.
We use the entire training set of MNIST and apply the gradient descent (Eq. 9) to update the weights.
We vary the batch size as { 64,128,256 } and measure the standard deviation of gradients ( σ ) w.r
[Figure]
t. parameters across different training batches.
A very small learning rate of η = 10 − 8 is set to satisfy the assumption in Theorem C.1.
Fig.
4 demonstrates the training loss after one epoch vs. standard deviation of gradients ( σ ).
Clearly, the results show that if a network has a lower gradient standard deviation, then it tends to have lower test loss values, and thus, a better generalization capacity.
These results empirically prove our claims in Theorem C.1.
Search space We use the commonly used MobileNetv2-based search space where the candidate networks are built by stacking multiple Inverted Bottleneck Blocks (IBNs) with SE modules Sandler et al.
; Pham et al.
Lin et al.
(2021); all the SE modules share the same se ratio as 0.25.
For each IBN, we vary the kernel size of the depth-wise convolutional layer from { 3,5,7 } and sample the expansion ratio from { 1,2,4,6 } .
We primarily consider ReLU as the activation function.
For each point-wise convolutional layer, the range of the number of channels is from 8 to 1024 with a step size of 8.
We use standard Kaiming Init to initialize all linear and convolution layers for every candidate networks He et al.
Search details : We use Algorithm 1 (main paper) to conduct the zero-shot search.
We repeat the search 10 5 times with the population size E = 512 .
For each of the candidate architectures, we compute ZiCo with two batches randomly sampled from the training set of ImageNet with batch size 128.
In total, it takes 10 hours on a single NVIDIA 3090 GPU for 10 5 search steps.
Training details : We use the same data augmentations configurations as in Pham et al.
(2018): mix-up, label-smoothing, random erasing, random crop/resize/flip/lighting, and AutoAugment.
We use the SGD optimizer with momentum 0.9 and weight decay 4e-5.
We take EfficientNet-B3 as a teacher network and use the knowledge distillation method to train the network.
We set the initial learning rate as 0.1 and used the cosine annealing scheme to adjust the learning rate during training.
We train the obtained network 480 epochs, which takes 83 hours on a 40-core Intel Xeon CPU and 8 NVIDIA 3090 GPU-powered server. Clearly, our proposed ZiCo is the only proxy that works consistently better than #Params and is generally the best among all these proxies. ‡ Both TE-NAS (Chen et al. (2021b))
and NASI (Chen et al.
(2021b)) use NTK (Jacot et al.
as the accuracy proxy to build their own search algorithms. [Table] [Table] E S UPPLEMENTARY RESULTS ON NAS B ENCHMARKS E. C OMPARISON WITH MORE PROXIES
In this section
[Table]
we further provide the comparison between our proposed ZiCo and more proxies proposed recently: KNAS (Xu et al. , NASWOT (Lopes et al. (2021)), GradSign ( Zhang & Jia (2022)) [Table] and NTK (TE-NAS Chen et al.
(2021b), NASI Shu et al.
(2022a)).
To compute the correlations, we use the official code released by the authors of the above papers to obtain the values of these proxies 2 .
As shown in Table 4, our proposed ZiCo performs better than all these proxies.
For example, NASWOT and GradSign achieve a similar correlation score as ZiCo on NATSBench-TSS; however, ZiCo has a significantly higher correlation score than these two proxies on NATSBenchSSS.
[Table]
E. C OMPARISON ON T RANS NAS-B ENCH -101-M ICRO
In this section
[Table]
we compare our proposed ZiCo against existing proxies on more diverse tasks.
We compare our proposed ZiCo against existing proxies on one mainstream NAS benchmark TransNASBench-101 Duan et al.
We pick the largest search space TransNAS-Bench-101-Micro which contains 4096 total architectures with different cell structures.
We compare ZiCo with various proxies under the following four tasks:
• Scene Classification.
Scene classification is a 47-class classification task that predicts the room type in the image.
• Jigsaw.
In the Jigsaw task, the input image is divided into nine patches and shuffled based on one of 1,000 predefined permutations.
The target here is to classify which permutation is used.
• Autoencoding.
Autoencoding is a pixel-level prediction task that encodes an input image into a low-dimension embedding vector and then reconstructs the raw image from the vector.
• Surface Normal.
Similar to autoencoding, surface normal is a pixel-level prediction task that predicts surface normal statistics.
As shown in Table 5, ZiCo consistently works well on Scene Classification, Jigsaw, and Surface
Normal; ZiCo has only 0.01 or 0.02 lower correlation scores than the highest scores.
Though
Fisher works better than ZiCo on Autoencoding, ZiCo has significantly higher correlation scores than Fisher on the remaining three tasks.
One possibility why Fisher works best on Autoencoding is that Autoencoding is an image-to-image task; Fisher is the only proxy that is built on the gradient w.r.t.
feature maps and thus can better extract the information between the input and output images.
Although Fisher works better than ZiCo on Autoencoding (we are still second best), ZiCo has a significantly higher score on the remaining tasks.
As shown in the main paper, we again note that existing proxies do not achieve a high correlation on all tasks consistently.
Table 6 demonstrates the test accuracy of the best architectures found using various proxies on each of the above tasks in TransNAS-Bench-101-Micro.
Once again, we see that ZiCo significantly outperforms existing proxies on all tasks except Autoencoding, where we trail Fisher by only 0.
SSIM.
Nonetheless, ZiCo is second best on the Autoencoding task.
Note that, similar to the correlation results in Table 5, other proxies do not consistently achieve high accuracy. For instance, while methods like Synflow or Zenscore achieve results close to ours on Scene Classification and Surface Normal, they produce poor results on other tasks like Jigsaw. Therefore, ZiCo consistently performs well on highly different tasks.
[Table] [Table]
E. I LLUSTRATION OF V ARIOUS P ROXIES VS . R EAL T EST A CCURACY
We provide some illustration figures of real test accuracy vs. various proxies on NATSBench-SSS search space for CIFAR10 (Fig. 5) and ImageNet16-120 datasets(Fig. 6).
We also show the same illustrative results (real test accuracy vs. various proxies) on NASBench101 search space in Fig.
F A BLATION STUDY F. I MPACT OF M EAN AND STD
We randomly select 2000 networks from NATSBench-SSS on CIFAR10, CIFAR100, and Img16120 datasets and compute the following proxies: (i) Mean value of gradients only; (ii) Standard deviation (STD) value of gradients only; (iii) Combination of mean and std value, i.e.,
our proposed
ZiCo .
We then calculate the correlation coefficients between these proxies and the real test accuracy.
As shown in Table.
7, our proposed ZiCo performs better on these three datasets than either using mean only or STD only.
Therefore, our proposed ZiCo is a better-designed proxy than using mean or STD individually.
F. S EARCH ALGORITHMS : ZERO - COST PT
In this section, we demonstrate that our proposed ZiCo can be combined with other search algorithms.
We take the Zero-Cost-PT (Zero-PT) as an example Xiang et al.
(2021b) because it is specifically designed for zero-shot proxies and is very time-efficient.
Essentially, Zero-PT first integrates all candidate networks into a supernet and assigns learnable weights to each candidate operation (same as one-shot NAS).
Then Zero-PT uses the zero-cost proxy instead of the training accuracy to update the weights for each candidate operation.
The final architecture is generated by selecting the operations with the highest weight values.
We combine different accuracy proxies with Zero-PT under the NASBench-201 and report the optimal architectures found with various proxies 3 .
As shown in Table 8, the architectures found via
ZiCo have the highest test accuracy except for Img16-120 datasets (ZiCo is the second best on
Img16-120)).
[Table] F. T RAINING RECIPE : WITHOUT DISTILLATION In this section [Table]
we train the obtained network under various FLOPs budgets with the exact same training setup as Xu et al.
; Cai et al.
Specifically, we train the neural network for 150 epochs with batch size 512 and input resolution 224.
We train the network without knowledge distillation and do not use advanced data augmentation methods (e
[Table]
g., mixup, RandAugment, etc).
Finally, we set the initial learning rate as 0.4 with a cosine annealing scheduling scheme.
Moreover, we train EfficientNets and the previous SOTA zero-shot NAS approach (Zen-score) under the same setup.
As shown in Table
[Table]
ZiCo outperforms all of the previous zero-shot NAS approaches.
For example, when the FLOPs budget is around 600M, ZiCo achieves 77.1% Top-1 accuracy, which is 1.0% and 1.6% higher than previous SOTA zero-shot NAS methods, i.e., Zen-score, and TE-NAS, respectively.
Moreover, ZiCo finds a model with similar accuracy as EfficientNet-B1, but with 100M fewer FLOPs and much less search cost.
Overall, compared to the regular one-shot or multi-shot
NAS methods, ZiCo achieves comparable or higher test accuracy with 5-9500 × less search time.
F. S EARCH SPACE : DARTS
In this section, we use ZiCo to conduct the zero-shot NAS on the DARTS search space.
We first use
Algorithm 1 to find the networks with the highest ZiCo without FLOPs budgets on the CIFAR dataset.
We conduct the search for 100k steps; this takes 0.7 hours on a single NVIDIA 3090 GPU (i.e., 0.03 GPU days).
Then, we train the obtained network with the exact same training setup as the original DARTS paper Liu et al.
(2019) 4 ; specifically, we train the neural network for 600 epochs with a batch size of 128.
We only use the standard data augmentation (normalization, cropping, and random flipping) together with the cutout tricks.
We don’t use knowledge distillation or any other
[Table]
advanced data augmentation tricks.
Finally, we set the initial learning rate as 0.025 with a cosine annealing scheduling scheme.
We repeat the same experiments for Zen-score.
As shown in Table 10, ZiCo outperforms previous zero-shot NAS approaches, e.g, Zen-score and
TE-NAS.
Moreover, compared to the regular one-shot or multi-shot NAS methods, ZiCo achieves comparable or higher test accuracy with at least 10 × less search time.
[Figure] [Figure] [Figure]