Speedy Performance Estimation for Neural Architecture Search

1 Introduction . Reliably estimating the generalisation performance of a proposed architecture is crucial to the success of Neural Architecture Search ( NAS ) but has always been a major bottleneck in NAS algorithms [ 13 ] . The traditional approach of training each architecture for a large number of epochs and evaluating it on validation data ( full training ) provides a reliable performance measure , but requires prohibitively large computational resources on the order of thousands of GPU days [ 55 , 39 , 56 , 38 , 13 ] . This motivates the development of methods for speeding up performance estimation to make NAS practical for limited computing budgets . A popular simple approach is early-stopping , which offers a low-fidelity approximation of generalisation performance by training for fewer epochs [ 27 , 14 , 25 ] . However , if we stop training early and evaluate the set of models on validation data , their relative performance ranking may not correlate well with the performance ranking of the fully-trained models [ 53 ] , i.e . their relative performance on the test set after the entire training budget has been used . Another line of work focuses on learning curve extrapolation [ 10 , 23 , 3 ] , which trains a surrogate model to predict the final generalisation performance based on the initial learning curve or meta-features of the architecture . However , the training of the surrogate often requires hundreds of fully evaluated architectures to achieve satisfactory extrapolation performance and the hyper-parameters of the surrogate also need to be optimised . Alternatively , the idea of weight sharing is adopted in one-shot NAS methods to speed up evaluation ∗Equal contribution . Correspondence to robin @ robots.ox.ac.uk . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . [ 36 , 28 , 47 ] . Despite leading to significant cost-saving , weight sharing heavily underestimates the true performance of good architectures and is unreliable in predicting the relative ranking among architectures [ 49 , 52 ] . A more recent group of estimators claim to be zero-cost [ 32 , 1 ] . Yet , their performance is often not competitive with the state of the art , inconsistent across tasks and can not be further improved with greater training budgets . In view of the above limitations , we propose a simple model-free method , Training Speed Estimation ( TSE ) , which provides a reliable yet computationally cheap estimate of the generalisation performance ranking of a set of architectures . Our method is inspired by recent empirical and theoretical results linking training speed and generalisation [ 17 , 30 ] and measures the training speed of an architecture by summing the training losses of the commonly-used SGD optimiser during training . We empirically show that our estimator can outperform strong existing approaches to predict the relative performance ranking among architectures , and can remain effective for a variety of search spaces and datasets . Moreover , we verify its usefulness under different NAS settings and find it can speed up query-based NAS approaches significantly as well improve the performance of one-shot and differentiable NAS . Our code is available at https : //github.com/rubinxin/TSE . 2 Method . Motivation The theoretical relationship between training speed and generalisation is described in a number of existing works . Stability-based generalisation bounds for SGD [ 17 , 29 ] bound the generalisation gap of a model based on the number of optimisation steps used to train it . These bounds predict that models which train faster obtain a lower worst-case generalisation error . In networks of sufficient width , a neural tangent kernel-inspired complexity measure can bound both the worst-case generalisation gap and the rate of convergence of ( full-batch ) gradient descent [ 2 , 6 ] . However , these bounds can not distinguish between models that are trained for the same number of steps but attain near-zero loss at different rates as they ignore the training trajectory . We instead draw inspiration from another approach which incorporates properties of the trajectory taken during SGD , seen in the information-theoretic generalisation bounds of [ 33 ] and [ 34 ] . These bounds depend on the variance of gradients computed during training , a quantity which provides a first-order approximation of training speed by quantifying how well gradient updates computed for one mini-batch generalize to other points in the training set . In view of this link , the empirical and theoretical findings in recent work [ 15 , 43 ] , which show that a notion of the variance of gradients over the training data is correlated with generalisation , become another piece of supporting evidence for training speed as a measure of generalisation . Finally , [ 30 ] prove that , in the setting of linear models and infinitely wide deep models performing Bayesian updates , the marginal likelihood , which is a theoretically-justified tool for model selection in Bayesian learning , can be bounded by a notion of training speed . This notion of training speed is defined as a sum of negative log predictive likelihoods , terms that resemble the losses on new data points seen by the model during an online learning procedure . Maximising the marginal likelihood is also equivalent to minimising a PAC-Bayesian bound on the generalisation error of a model , as shown by [ 16 ] . In particular , this suggests that using the TSE approach for model selection in Bayesian models is equivalent to minimizing an estimate of a PAC-Bayes bound . Because the NAS settings we consider in our experiments use neural networks rather than Bayesian models and due to space constraints , we defer the statement and proof of this result to Appendix A . Training Speed Estimation The results described above suggest that leveraging a notion of training speed may benefit model selection procedures in NAS . Many such notions exist in the generalisation literature : [ 18 ] count the number of optimisation steps needed to attain a loss below a specified threshold , while [ 17 ] consider the total number of optimisation steps taken . Both measures are strong predictors of generalisation after training , yet neither is suitable for NAS , where we seek to stop training as early as possible if the model is not promising . We draw inspiration from the Bayesian perspective and the PAC-Bayesian bound discussed above and present an alternate estimator of training speed that amounts to the area under the model ’ s training curve . Models that train quickly attain a low loss after few training steps , and so will have a lower area-under-curve than those which train slowly . This addresses the shortcomings of the previous two methods as it is able to distinguish models both early and late in training . Definition 1 ( Training Speed Estimator ) . Let ` denote a loss function , fθ ( x ) the output of a neural network f with input x and parameters θ , and let θt , i denote the parameters of the network after t epochs and i mini-batches of SGD . After training the network for T epochs2 , we sum the training losses collected so far to get the following Training Speed Estimate ( TSE ) : TSE = T∑ t=1 [ 1 B B∑ i=1 ` ( fθt , i ( Xi ) , yi ) ] ( 1 ) where l is the training loss of a mini-batch ( Xi , yi ) at epoch t and B is the number of training steps within an epoch . This estimator weights the losses accumulated during every epoch equally . However , recent work suggests that training dynamics of neural networks in the very early epochs are often unstable and not always informative of properties of the converged networks [ 24 ] . Therefore , we hypothesise that an estimator of network training speed that assigns higher weights to later epochs may exhibit a better correlation with the true generalisation performance of the final trained network . On the other hand , it is common for neural networks to overfit on their training data and reach near-zero loss after sufficient optimisation steps , so attempting to measure training speed solely based on the epochs near the end of training will be difficult and likely suffer degraded performance on model selection . To verify whether it is beneficial to ignore or downplay the information from early epochs of training , we propose two variants of our estimator . The first , TSE-E , treats the first few epochs as a burn-in phase for θt , i to converge to a stable distribution P ( θ ) and starts the sum from epoch t = T − E + 1 instead of t = 1 . In the case where E = 1 , we start the sum at t = T and our estimator corresponds to the sum over training losses within the most recent epoch t = T . TSE-E = T∑ t=T−E+1 [ 1 B B∑ i=1 ` ( fθt , i ( Xi ) , yi ) ] , TSE-EMA = T∑ t=1 γT−t [ 1 B B∑ i=1 ` ( fθt , i ( Xi ) , yi ) ] The second , TSE-EMA , does not completely discard the information from the early training trajectory but takes an exponential moving average of the sum of training losses with γ = 0.9 , thus assigning higher weight to the sum of losses obtained in later training epochs . We empirically show in Section 4.2 that our proposed TSE and its variants ( TSE-E and TSE-EMA ) , despite their simple form , can reliably estimate the generalisation performance of neural architectures with a very small training budget , can remain effective for a large range of training epochs , and are robust to the choice of hyperparameters such as the summation window E and the decay rate γ . However , our estimator is not meant to replace the validation accuracy at the end of training or when the user can afford large training budget . In those settings , validation accuracy remains as the gold standard for evaluating the true test performance of architectures . Ours is just a speedy performance estimator for NAS , aimed at giving an indication early in training about an architecture ’ s generalisation potential under a fixed training set-up . Our choice of using the training loss , instead of the validation loss , to measure training speed is an important component of the proposed method . While it is possible to formulate an alternative estimator , which sums the validation losses of a model early in training , this estimator would no longer be measuring training speed . In particular , such an estimator would not capture the generalisation of gradient updates from one mini-batch to later mini-batches in the data to the same extent as TSE does . Indeed , we hypothesise that once the optimisation process has reached a local minimum , the sum over validation losses more closely resembles a variance-reduction technique that estimates the expected loss over parameters sampled via noisy SGD steps around this minimum . We show in Figure 1 and Appendix C that our proposed sum over training losses ( TSE ) outperforms the sum over validation losses ( SoVL ) in ranking models in agreement with their true test performance . 3 Related Work . Various approaches have been developed to speed up architecture performance estimation , thus improving the efficiency of NAS . Low-fidelity estimation methods accelerate NAS by using the 2T can be far from the total training epochs Tend used in complete training validation accuracy obtained after training architectures for fewer epochs ( namely early-stopping ) [ 27 , 14 , 56 , 53 ] , training a down-scaled model with fewer cells during the search phase [ 56 , 38 ] , or training on a subset of the data [ 22 ] . However , low-fidelity estimates underestimate the true performance of the architecture and can change the relative ranking among architectures [ 13 ] . This undesirable effect on relative ranking is more prominent when the cheap approximation set-up is dissimilar to the full training procedure [ 53 ] . As shown in Fig . 1 below , the validation accuracy at early epochs of training suffers low rank correlation with the final test performance . Another class of performance estimation methods trains a regression model to extrapolate the learning curve from what is observed in the initial phase of training or predict the final test accuracy purely based on architecture structure . Regression model choices that have been explored include Gaussian processes with a tailored kernel function [ 41 , 19 ] , an ensemble of parametric functions [ 10 ] , tree-based models [ 7 , 20 , 50 ] , a Bayesian neural network [ 23 ] and a ν-support vector machine regressor ( ν-SVR ) [ 3 ] which achieves state-of-the-art prediction performance [ 45 ] . Although these model-based methods can often predict the performance ranking better than their model-free early-stopping counterparts , they require a relatively large amount of fully evaluated architecture data ( e.g . 100 fully evaluated architectures in [ 3 ] ) to train the regression surrogate properly and optimise the model hyperparameters in order to achieve a good prediction performance . The high computational cost of collecting the training set makes such model-based methods less favourable for NAS unless the practitioner has already evaluated hundreds of architectures on the target task . Moreover , both low-fidelity estimates and learning curve extrapolation estimators are empirically developed and lack theoretical motivation . Weight sharing is employed in one-shot or gradient-based NAS methods to reduce computational costs [ 36 , 28 , 47 ] . Under the weight-sharing setting , all architectures are considered as subnetworks of a supernetwork . Only the weights of the supernetwork are trained while the architectures ( subnetworks ) inherit the corresponding weights from the supernetwork . This removes the need for retraining each architecture during the search and thus achieves a significant speed-up . However , the weight sharing ranking among architectures often correlates poorly with the true performance ranking [ 49 , 52 , 54 ] , meaning architectures chosen by one-shot NAS are likely to be sub-optimal when evaluated independently [ 54 ] . In Section 4.4 , we demonstrate that we improve the performance of weight sharing in correctly ranking architectures by combining our estimator with it . Recently , several works propose to estimate network performance without training by using methods from the pruning literature [ 1 ] or examining the covariance of input gradients across different input images [ 32 ] . Such methods incur near-zero computational costs but their performances are often not competitive and do not generalise well to larger search spaces , as shown in Section 4.2 below . Moreover , these methods can not be improved with additional training budget . Another line of work studies the effect of architecture encoding on speeding up the search [ 48 , 44 , 40 ] . However , learning better latent representation for architectures is orthogonal to the performance estimation . Apart from the above mentioned performance estimators used in NAS , many complexity measures have been proposed to analyse the generalisation performance of deep neural networks . [ 18 ] provides a rigorous empirical analysis of over 40 such measures . This investigation finds that sharpness-based measures [ 31 , 21 , 35 , 12 ] ( including PAC-Bayesian bounds ) obtain a good correlation with test set performance , but their estimation require adding randomly generated perturbations to the network parameters and the magnitude of the perturbations needs to be carefully optimised with additional training , making them unsuitable performance estimators for NAS . Optimisation-based complexity measures , which counts the number of steps required to reach a certain loss value , also perform well in predicting generalisation . However , as discussed in Section 2 , it is closely related to our approach but not as easy to deploy as our estimators under the NAS setting .

Summary: This paper proposes a simple yet reliable estimator for fully trained network performance inspired by the PAC-Bayes bounds. Based on the observation in previous work that networks that train faster have better generalization performance, they propose that small areas under the training loss curve early in training correlate with better final performance. They propose several variations of this estimator with varying methods of weighting more recent epochs more than early epochs. They demonstrate the effectiveness of their estimators to provide good ranking correlation performance with the final accuracy with limited training on several search spaces/datasets such as NB-201-cifar10/cifar100/imagenet16 and Resnet/Resnext-cifar10. They also demonstrate that their estimator provides better rankings and find better performing networks compared to using the validation accuracy when applied to the problem of selecting a final architecture from a trained super-network for one-shot NAS. They also demonstrate that it can more efficiently find high performing networks compared to validation accuracy when applied to Regularized Evolution NAS and Bayesian Optimization NAS.