Geometry-Aware Gradient Algorithms for Neural Architecture Search

1 INTRODUCTION . Neural architecture search has become an important tool for automating machine learning ( ML ) but can require hundreds of thousands of GPU-hours to train . Recently , weight-sharing approaches have achieved state-of-the-art performance while drastically reducing the computational cost of NAS to just that of training a single shared-weights network ( Pham et al. , 2018 ; Liu et al. , 2019 ) . Methods such as DARTS ( Liu et al. , 2019 ) , GDAS ( Dong & Yang , 2019 ) , and many others ( Pham et al. , 2018 ; Zheng et al. , 2019 ; Yang et al. , 2020 ; Xie et al. , 2019 ; Liu et al. , 2018 ; Laube & Zell , 2019 ; Cai et al. , 2019 ; Akimoto et al. , 2019 ; Xu et al. , 2020 ) combine weight-sharing with a continuous relaxation of the discrete search space to allow cheap gradient updates , enabling the use of popular optimizers . However , despite some empirical success , weight-sharing remains poorly understood and has received criticism due to ( 1 ) rank-disorder ( Yu et al. , 2020 ; Zela et al. , 2020b ; Zhang et al. , 2020 ; Pourchot et al. , 2020 ) , where the shared-weights performance is a poor surrogate of standalone performance , and ( 2 ) poor results on recent benchmarks ( Dong & Yang , 2020 ; Zela et al. , 2020a ) . Motivated by the challenge of developing simple and efficient methods that achieve state-of-the-art performance , we study how to best handle the goals and optimization objectives of NAS . We start by observing that weight-sharing subsumes architecture hyperparameters as another set of learned parameters of the shared-weights network , in effect extending the class of functions being learned . This suggests that a reasonable approach towards obtaining high-quality NAS solutions is to study how to regularize and optimize the empirical risk over this extended class . While many regularization approaches have been implicitly proposed in recent NAS efforts , we focus instead on the question of optimizing architecture parameters , which may not be amenable to standard procedures such as SGD that work well for standard neural network weights . In particular , to better-satisfy desirable properties such as generalization and sparsity of architectural decisions , we propose to constrain architecture parameters to the simplex and update them using exponentiated gradient , which has favorable convergence properties due to the underlying problem structure . Theoretically , we draw upon the mirror descent meta-algorithm ( Nemirovski & Yudin , 1983 ; Beck & Teboulle , 2003 ) to give convergence guarantees when using any of a broad class of such geometry-aware gradient methods to optimize the weight-sharing objective ; empirically , we show that our solution leads to strong improvements on several NAS benchmarks . We summarize these contributions below : 1 . We argue for studying NAS with weight-sharing as a single-level objective over a structured function class in which architectural decisions are treated as learned parameters rather than hyperparameters . Our setup clarifies recent concerns about rank disorder and makes clear that proper regularization and optimization of this objective is critical to obtaining high-quality solutions . 2 . Focusing on optimization , we propose to improve existing NAS algorithms by re-parameterizing architecture parameters over the simplex and updating them using exponentiated gradient , a variant of mirror descent that converges quickly over this domain and enjoys favorable sparsity properties . This simple modification—which we call the Geometry-Aware Exponentiated Algorithm ( GAEA ) —is easily applicable to numerous methods , including first-order DARTS Liu et al . ( 2019 ) , GDAS Dong & Yang ( 2019 ) , and PC-DARTS ( Xu et al. , 2020 ) . 3 . To show correctness and efficiency of our scheme , we prove polynomial-time stationary-point convergence of block-stochastic mirror descent—a family of geometry-aware gradient algorithms that includes GAEA—over a continuous relaxation of the single-level NAS objective . To the best of our knowledge these are the first finite-time convergence guarantees for gradient-based NAS . 4 . We demonstrate that GAEA improves upon state-of-the-art methods on three of the latest NAS benchmarks for computer vision . Specifically , we beat the current best results on NAS-Bench-201 ( Dong & Yang , 2020 ) by 0.18 % on CIFAR-10 , 1.59 % on CIFAR-100 , and 0.82 % on ImageNet16-120 ; we also outperform the state-of-the-art on the DARTS search space Liu et al . ( 2019 ) , for both CIFAR-10 and ImageNet , and match it on NAS-Bench-1Shot1 ( Zela et al. , 2020a ) .1 Related Work . Most optimization analyses of NAS show monotonic improvement ( Akimoto et al. , 2019 ) , asymptotic guarantees ( Yao et al. , 2020 ) , or bounds on auxiliary quantities disconnected from any objective ( Noy et al. , 2019 ; Nayman et al. , 2019 ; Carlucci et al. , 2019 ) . In contrast , we prove polynomial-time stationary-point convergence on a single-level objective for weight-sharing NAS , so far only studied empirically ( Xie et al. , 2019 ; Li et al. , 2019 ) . Our results draw upon the mirror descent meta-algorithm ( Nemirovski & Yudin , 1983 ; Beck & Teboulle , 2003 ) and extend recent nonconvex convergence results Zhang & He ( 2018 ) to handle alternating descent . While there exist related results ( Dang & Lan , 2015 ) the associated guarantees do not hold for the algorithms we propose . Finally , we note that a variant of GAEA that modifies first-order DARTS is related to XNAS ( Nayman et al. , 2019 ) , whose update also involves exponentiated gradient ; however , GAEA is simpler and easier to implement.2 Furthermore , the regret guarantees for XNAS do not relate to any meaningful performance measure for NAS such as speed or accuracy , whereas we guarantee convergence on the ERM objective . 2 THE WEIGHT-SHARING OPTIMIZATION PROBLEM . In supervised ML we have a dataset T of labeled pairs ( x , y ) drawn from a distribution D over input/output spaces X and Y . The goal is to use T to search a function class H for hw : X 7→ Y parameterized by w ∈ Rd that has low expected test loss ` ( hw ( x ) , y ) when using x to predict the associated y on unseen samples drawn from D , as measured by some loss ` : Y × Y 7→ [ 0 , ∞ ) . A common way to do so is by approximate ( regularized ) empirical risk minimization ( ERM ) , i.e . finding w ∈ Rd with the smallest average loss over T , via some iterative method Alg , e.g . SGD . 2.1 THE BENEFITS AND CRITICISMS OF WEIGHT-SHARING FOR NAS . NAS is often viewed as hyperparameter optimization on top of Alg , with each architecture a ∈ A corresponding to a function class Ha = { hw , a : X 7→ Y , w ∈ Rd } to be selected by using validation data V ⊂ X × Y to evaluate the predictor obtained by fixing a and doing approximate ERM over T : min a∈A ∑ ( x , y ) ∈V ` ( hwa , a ( x ) , y ) s.t . wa = Alg ( T , a ) ( 1 ) Since training individual sets of weights for any sizeable number of architectures is prohibitive , weight-sharing methods instead use a single set of shared weights to obtain validation signal about many architectures at once . In its most simple form , RS-WS ( Li & Talwalkar , 2019 ) , these weights 1Code to obtain these results has been made available in the supplementary material . 2XNAS code does not implement search and , as with previous efforts ( Li et al. , 2019 , OpenReview ) , we can not reproduce results after correspondence with the authors . XNAS ’ s best architecture achieves an average test error of 2.70 % under the DARTS evaluation , while GAEA achieves 2.50 % . For details see Appendix C.4 . are trained to minimize a non-adaptive objective , minw∈Rd Ea ∑ ( x , y ) ∈T ` ( hwa , a ( x ) , y ) , where the expectation is over a fixed distribution over architectures A . The final architecture a is then chosen to maximize the outer ( validation ) objective in ( 1 ) subject to wa = w. More frequently used is a bilevel objective over some continuous relaxation Θ of the architecture space A , after which a valid architecture is obtained via a discretization step Map : Θ 7→ A ( Pham et al. , 2018 ; Liu et al. , 2019 ) : min θ∈Θ ∑ ( x , y ) ∈V ` ( hw , θ ( x ) , y ) s.t . w ∈ arg min u∈Rd ∑ ( x , y ) ∈T ` ( hu , θ ( x ) , y ) ( 2 ) This objective is not significantly different from ( 2 ) , since Alg ( T , a ) approximately minimizes the empirical risk w.r.t . T ; the difference is replacing discrete architectures with relaxed architecture parameters θ ∈ Θ , w.r.t . which we can take derivatives of the outer objective . This allows ( 2 ) to be approximated via alternating gradient updates w.r.t . w and θ. Relaxations can be stochastic , so that Map ( θ ) is a sample from a θ-parameterized distribution ( Pham et al. , 2018 ; Dong & Yang , 2019 ) , or a mixture , in which case Map ( θ ) selects architectural decisions with the highest weight in a convex combination given by θ ( Liu et al. , 2019 ) . We overview this in more detail in Appendix A . While weight-sharing significantly shortens search ( Pham et al. , 2018 ) , it draws two main criticisms : • Rank disorder : this describes when the rank of an architecture a according to the validation risk evaluated with fixed shared weights w is poorly correlated with the one using “ standalone ” weights wa = Alg ( T , a ) . This causes suboptimal architectures to be selected after shared weights search ( Yu et al. , 2020 ; Zela et al. , 2020b ; Zhang et al. , 2020 ; Pourchot et al. , 2020 ) . • Poor performance : weight-sharing can converge to degenerate architectures ( Zela et al. , 2020a ) and is outperformed by regular hyperparameter tuning on NAS-Bench-201 ( Dong & Yang , 2020 ) . 2.2 SINGLE-LEVEL NAS AS A BASELINE OBJECT OF STUDY . Why are we able to apply weight-sharing to NAS ? The key is that , unlike regular hyperparameters such as step-size , architectural hyperparameters directly affect the loss function without requiring a dependent change in the model weights w. Thus we can distinguish architectures without retraining simply by changing architectural decisions . Besides enabling weight-sharing , this point reveals that the goal of NAS is perhaps better viewed as a regular learning problem over an extended class HA = ⋃ a∈AHa = { hw , a : X 7→ Y , w ∈ Rd , a ∈ A } that subsumes the architectural decisions as parameters of a larger model class , an unrelaxed “ supernet. ” The natural approach to solving this is by approximate empirical risk minimization , e.g . by approximating continuous objective below on the right using a gradient algorithm and passing the output θ through Map to obtain a valid architecture : min w∈Rd , a∈A ∑ ( x , y ) ∈T ` ( hw , a ( x ) , y ) ︸ ︷︷ ︸ discrete ( unrelaxed ) supernet ( NAS ERM ) min w∈Rd , θ∈Θ ∑ ( x , y ) ∈T ` ( hw , θ ( x ) , y ) ︸ ︷︷ ︸ continuous relaxation ( supernet ERM ) ( 3 ) Several works have optimized this single-level objective as an alternative to bilevel ( 2 ) ( Xie et al. , 2019 ; Li et al. , 2019 ) . We argue for its use as the baseline object of study in NAS for three reasons : 1 . As discussed above , it is the natural first approach to solving the statistical objective of NAS : finding a good predictor hw , a ∈ HA in the extended function class over architectures and weights . 2 . The common alternating gradient approach to the bilevel problem ( 2 ) is in practice very similar to alternating block approaches to ERM ( 3 ) ; as we will see , there are established ways of analyzing such methods for the latter objective , while for the former convergence is known only under very strong assumptions such as uniqueness of the inner minimum ( Franceschi et al. , 2018 ) . 3 . While less frequently used in practice than bilevel , single-level optimization can be very effective : we use it to achieve new state-of-the-art results on NAS-Bench-201 ( Dong & Yang , 2020 ) . Understanding NAS as single-level optimization—the usual deep learning setting—makes weightsharing a natural , not surprising , approach . Furthermore , for methods—both single-level and bilevel— that adapt architecture parameters during search , it suggests that we need not worry about rank disorder as long as we can use optimization to find a single feasible point that generalizes well ; we explicitly do not need a ranking . Non-adaptive methods such as RS-WS still do require rank correlation to select good architectures after search , but they are explicitly not changing θ and so have no variant solving ( 3 ) . The single-level formulation thus reduces search method design to well-studied questions of how to best regularize and optimize ERM . While there are many techniques for regularizing weight-sharing—including partial channels ( Xu et al. , 2020 ) and validation Hessian penalization ( Zela et al. , 2020a ) —we focus on the second question of optimization .

Summary: This paper introduces the geometry-aware framework that can be adapted to any existing weight-sharing NAS methods optimized over gradient descent. The authors focus on the aspect of optimizing the architecture parameters to overcome the criticism of weight-sharing methods. The author's method relies on the mirror descent supporting their methods with a theoretical guarantee for the fast convergence. The author also supports their methods on various datasets such as CIFAR-10, ImageNet, NAS-Bench-201 (Dong & Yang), and NAS-Bench-1Shot (Zela et al.).