Neural Architecture Search on ImageNet in Four GPU Hours: A Theoretically Inspired Perspective

1 INTRODUCTION . The recent development of deep networks significantly contributes to the success of computer vision . Thanks to many efforts by human designers , the performance of deep networks have been significantly boosted ( Krizhevsky et al. , 2012 ; Simonyan & Zisserman , 2014 ; Szegedy et al. , 2015 ; He et al. , 2016 ; Xie et al. , 2017 ) . However , the manual creation of new network architectures not only costs enormous time and resources due to trial-and-error , but also depends on the design experience that does not always scale up . To reduce the human efforts and costs , neural architecture search ( NAS ) has recently amassed explosive interests , leading to principled and automated discovery for good architectures in a given search space of candidates ( Zoph & Le , 2016 ; Brock et al. , 2017 ; Pham et al. , 2018 ; Liu et al. , 2018a ; Chen et al. , 2018 ; Bender et al. , 2018 ; Gong et al. , 2019 ; Chen et al. , 2020a ; Fu et al. , 2020 ) . As an optimization problem , NAS faces two core questions : 1 ) “ how to evaluate ” , i.e . the objective function that defines what are good architectures we want ; 2 ) “ how to optimize ” , i.e . by what means we could effectively optimize the objective function . These two questions are entangled and highly non-trivial , since the search spaces are of extremely high dimension , and the generalization ability of architectures can not be easily inferred ( Dong & Yang , 2020 ; Dong et al. , 2020 ) . Existing NAS methods mainly leverage the validation set and conduct accuracy-driven architecture optimization . They either formulate the search space as a super-network ( “ supernet ” ) and make the training loss differentiable through the architecture parameters ( Liu et al. , 2018b ) , or treat the architecture selection as a sequential decision making process ( Zoph & Le , 2016 ) or evolution of genetics ( Real et al. , 2019 ) . However , these NAS algorithms suffer from heavy consumption of both time and GPU resources . Training a supernet till convergence is extremely slow , even with many effective heuristics for sampling or channel approximations ( Dong & Yang , 2019 ; Xu et al. , 2019 ) . Approximated proxy inference such as truncated training/early stopping can accelerate the search , but is well known to introduce search bias to the inaccurate results obtained ( Pham et al. , 2018 ; Liang et al. , 2019 ; Tan et al. , 2020 ) . The heavy search cost not only slows down the discovery of novel architectures , but also blocks us from more meaningfully understanding the NAS behaviors . On the other hand , the analysis of neural network ’ s trainability ( how effective a network can be optimized via gradient descent ) and expressivity ( how complex the function a network can represent ) has witnessed exciting development recently in the deep learning theory fields . By formulating neural networks as a Gaussian Process ( no training involved ) , the gradient descent training dynamics can be characterized by the Neural Tangent Kernel ( NTK ) of infinite ( Lee et al. , 2019 ) or finite ( Yang , 2019 ) width networks , from which several useful measures can be derived to depict the network trainability at the initialization . Hanin & Rolnick ( 2019a ; b ) ; Xiong et al . ( 2020 ) describe another measure of network expressivity , also without any training , by counting the number of unique linear regions that a neural network can divide in its input space . We are therefore inspired to ask : • How to optimize NAS at network ’ s initialization without involving any training , thus significantly eliminating a heavy portion of the search cost ? • Can we define how to evaluate in NAS by analyzing the trainability and expressivity of architectures , and further benefit our understanding of the search process ? Our answers are yes to both questions . In this work , we propose TE-NAS , a framework for trainingfree neural architecture search . We leverage two indicators , the condition number of NTK and the number of linear regions , that can decouple and effectively characterize the trainability and expressivity of architectures respectively in complex NAS search spaces . Most importantly , these two indicators can be measured in a training-free and label-free manner , thus largely accelerates the NAS search process and benefits the understanding of discovered architectures . To our best knowledge , TE-NAS makes the first attempt to bridge the theoretical findings of deep neural networks and real-world NAS applications . While we intend not to claim that the two indicators we use are the only nor the best options , we hope our work opens a door to theoretically-inspired NAS and inspires the discovery of more deep network indicators . Our contributions are summarized as below : • We identify and investigate two training-free and label-free indicators to rank the quality of deep architectures : the spectrum of their NTKs , and the number of linear regions in their input space . Our study finds that they reliably indicate the trainability and expressivity of a deep network respectively , and are strongly correlated with the network ’ s test accuracy . • We leverage the above two theoretically-inspired indicators to establish a training-free NAS framework , TE-NAS , therefore eliminating a drastic portion of the search cost . We further introduce a pruning-based mechanism , to boost search efficiency and to more flexibly trade-off between trainability and expressivity . • In NAS-Bench-201/DARTS search spaces , TE-NAS discovers architectures with a strong performance at remarkably lower search costs , compared to previous efforts . With just one 1080Ti , it only costs 0.5 GPU hours to search on CIFAR10 , and 4 GPU hours on ImageNet , respectively , setting the new record for ultra-efficient yet high-quality NAS . 2 RELATED WORKS . Neural architecture search ( NAS ) is recently proposed to accelerate the principled and automated discovery of high-performance networks . However , most works suffer from heavy search cost , for both weight-sharing based methods ( Liu et al. , 2018b ; Dong & Yang , 2019 ; Liu et al. , 2019 ; Yu et al. , 2020a ; Li et al. , 2020a ; Yang et al. , 2020a ) and single-path sampling-based methods ( Pham et al. , 2018 ; Guo et al. , 2019 ; Real et al. , 2019 ; Tan et al. , 2020 ; Li et al. , 2020c ; Yang et al. , 2020b ) . A one-shot super network can share its parameters to sampled sub-networks and accelerate the architecture evaluations , but it is very heavy and hard to optimize and suffers from a poor correlation between its accuracy and those of the sub-networks ( Yu et al. , 2020c ) . Sampling-based methods achieve more accurate architecture evaluations , but their truncated training still imposes bias to the performance ranking since this is based on the results of early training stages . Instead of estimating architecture performance by direct training , people also try to predict network ’ s accuracy ( or ranking ) , called predictor based NAS methods ( Liu et al. , 2018a ; Luo et al. , 2018 ; Dai et al. , 2019 ; Luo et al. , 2020 ) . Graph neural network ( GNN ) is a popular choice as the predictor model ( Wen et al. , 2019 ; Chen et al. , 2020b ) . Siems et al . ( 2020 ) even propose the first large-scale surrogate benchmark , where most of the architectures ’ accuracies are predicted by a pretrained GNN predictor . The learned predictor can achieve highly accurate performance evaluation . However , the data collection step - sampling representative architectures and train them till converge - still requires extremely high cost . People have to sample and train 2,000 to 50,000 architectures to serve as the training data for the predictor . Moreover , none of these works can demonstrate the cross-space transferability of their predictors . This means one has to repeat the data collection and predictor training whenever facing an unseen search space , which is highly nonscalable . The heavy cost of architecture evaluation hinders the understanding of the NAS search process . Recent pioneer works like Shu et al . ( 2019 ) observed that DARTS and ENAS tend to favor architectures with wide and shallow cell structures due to their smooth loss landscape . Siems et al . ( 2020 ) studied the distribution of test error for different cell depths and numbers of parameter-free operators . Chen & Hsieh ( 2020 ) for the first time regularizes the Hessian norm of the validation loss and visualizes the smoother loss landscape of the supernet . Li et al . ( 2020b ) proposed to approximate the validation loss landscape by learning a mapping from neural architectures to their corresponding validate losses . Still , these analyses can not be directly leveraged to guide the design of network architectures . Mellor et al . ( 2020 ) recently proposed a NAS framework that does not involve training , which shares the same motivation with us towards training-free architecture search at initialization . They empirically find that the correlation between sample-wise input-output Jacobian can indicate the architecture ’ s test performance . However , why does the Jacobian work is not well explained and demonstrated . Their search performance on NAS-Bench-201 is still left behind by the state-of-the-art NAS works , and they did not extend to DARTs space . Meanwhile , we see the evolving development of deep learning theory on neural networks . NTK ( neural tangent kernel ) is proposed to characterize the gradient descent training dynamics of infinite wide ( Jacot et al. , 2018 ) or finite wide deep networks ( Hanin & Nica , 2019 ) . Wide networks are also proved to evolve as linear models under gradient descent ( Lee et al. , 2019 ) . This is further leveraged to decouple the trainability and generalization of networks ( Xiao et al. , 2019 ) . Besides , a natural measure of ReLU network ’ s expressivity is the number of linear regions it can separate in its input space ( Raghu et al. , 2017 ; Montúfar , 2017 ; Serra et al. , 2018 ; Hanin & Rolnick , 2019a ; b ; Xiong et al. , 2020 ) . In our work , we for the first time discover two important indicators that can effectively rank architectures , thus bridging the theoretic findings and real-world NAS applications . Instead of claiming the two indicators we discover are the best , we believe there are more meaningful properties of deep networks that can benefit the architecture search process . We leave them as open questions and encourage the community to study . 3 METHODS . The core motivation of our TE-NAS framework is to achieve architecture evaluation without involving any training , to significantly accelerate the NAS search process and reduce the search cost . In section 3.1 we present our study on two important indicators that reflect the trainability and expressivity of a neural network , and in section 3.2 we design a novel pruning-based method that can achieve a superior trade-off between the two indicators . 3.1 ANALYZING TRAINABILITY AND EXPRESSIVITY OF DEEP NETWORKS . Trainability and expressivity are distinct notions regarding a neural network ( Xiao et al. , 2019 ) . A network can achieve high performance only if the function it can represent is complex enough and at the same time , it can be effectively trained by gradient descent .

Summary: Training-free NAS is a promising direction. This work demonstrates that two theoretically inspired indicators: the spectrum of NTK and number of linear regions, strongly correlate with the network’s performance, and can be leveraged to reduce the search cost and decouple the analysis of the network’s trainability and expressivity. The authors thus proposed TE-NAS to rank architectures by analyzing the two indicators. Without involving any training, TE-NAS can achieve competitive NAS performance within minimum search time.