NASI: Label- and Data-agnostic Neural Architecture Search at Initialization

1 INTRODUCTION . The past decade has witnessed the wide success of deep neural networks ( DNNs ) in computer vision and natural language processing . These DNNs , e.g. , VGG ( Simonyan & Zisserman , 2015 ) , ResNet ( He et al. , 2016 ) , and MobileNet ( Howard et al. , 2017 ) , are typically handcrafted by human experts with considerable trials and errors . The human efforts devoting to the design of these DNNs are , however , not affordable nor scalable due to an increasing demand of customizing DNNs for different tasks . To reduce such human efforts , Neural Architecture Search ( NAS ) ( Zoph & Le , 2017 ) has recently been introduced to automate the design of DNNs . As summarized in ( Elsken et al. , 2019 ) , NAS conventionally consists of a search space , a search algorithm , and a performance evaluation . Specifically , the search algorithm aims to select the best-performing neural architecture from the search space based on its evaluated performance via performance evaluation . In the literature , various search algorithms ( Luo et al. , 2018 ; Zoph et al. , 2018 ; Real et al. , 2019 ) have been proposed to search for architectures with comparable or even better performance than the handcrafted ones . However , these NAS algorithms are inefficient due to the requirement of model training for numerous candidate architectures during the search process . To improve the search inefficiency , one-shot NAS algorithms ( Dong & Yang , 2019 ; Pham et al. , 2018 ; Liu et al. , 2019 ; Xie et al. , 2019 ) have trained a single one-shot architecture and then evaluated the performance of candidate architectures with model parameters inherited from this fine-tuned one-shot architecture . So , these algorithms can considerably reduce the cost of model training , but still require the training of the one-shot architecture . This naturally leads to the question whether NAS is realizable at initialization such that model training can be completely avoided during the search process ? To the best of our knowledge , only a few efforts to date have been devoted to developing NAS algorithms without model training empirically ( Mellor et al. , 2020 ; Park et al. , 2020 ; Abdelfattah et al. , 2021 ; Chen et al. , 2021 ) . This paper presents a novel NAS algorithm called NAS at Initialization ( NASI ) that can completely avoid model training to boost search efficiency . To achieve this , NASI exploits the capability of a Neural Tangent Kernel ( NTK ) ( Jacot et al. , 2018 ; Lee et al. , 2019a ) in being able to formally characterize the performance of infinite-wide DNNs at initialization , hence allowing the performance of candidate architectures to be estimated and realizing NAS at initialization . Specifically , given the estimated performance of candidate architectures by NTK , NAS can be reformulated into an optimization problem without model training ( Sec . 3.1 ) . However , NTK is prohibitively costly to evaluate . Fortunately , we can approximate it1 with a similar form to gradient flow ( Wang et al. , 2020 ) ( Sec . 3.2 ) . This results in a reformulated NAS problem that can be solved efficiently by a gradient-based algorithm via additional relaxation with Gumbel-Softmax ( Jang et al. , 2017 ; Maddison et al. , 2017 ) ( Sec . 3.3 ) . Interestingly , NASI is shown to be label- and data-agnostic under mild conditions , which thus implies the transferability of architectures selected by NASI over different datasets ( Sec . 4 ) . We will firstly empirically demonstrate the improved search efficiency and the competitive search effectiveness achieved by NASI in NAS-Bench-1Shot1 ( Zela et al. , 2020b ) ( Sec . 5.1 ) . Compared with other NAS algorithms , NASI incurs the smallest search cost while preserving the competitive performance of its selected architectures . Meanwhile , the architectures selected by NASI from the DARTS ( Liu et al. , 2019 ) search space over CIFAR-10 consistently enjoy the competitive or even outperformed performance when evaluated on different benchmark datasets , e.g. , CIFAR-10/100 and ImageNet ( Sec . 5.2 ) , indicating the guaranteed transferability of architectures selected by our NASI . In Sec . 5.3 , NASI is further demonstrated to be able to select well-performing architectures on CIFAR-10 even with randomly generated labels or data , which strongly supports the label- and data-agnostic search and also the guaranteed transferability achieved by our NASI . 2 RELATED WORKS AND BACKGROUND . 2.1 NEURAL ARCHITECTURE SEARCH . A growing body of NAS algorithms have been proposed in the literature ( Zoph & Le , 2017 ; Liu et al. , 2018 ; Luo et al. , 2018 ; Zoph et al. , 2018 ; Real et al. , 2019 ) to automate the design of neural architectures . However , scaling existing NAS algorithms to large datasets is notoriously hard . Recently , attention has thus been shifted to improving the search efficiency of NAS without sacrificing the generalization performance of its selected architectures . In particular , a one-shot architecture is introduced by Pham et al . ( 2018 ) to share model parameters among candidate architectures , thereby reducing the cost of model training substantially . Recent works ( Chen et al. , 2019 ; Dong & Yang , 2019 ; Liu et al. , 2019 ; Xie et al. , 2019 ; Chen & Hsieh , 2020 ; Chu et al. , 2020 ) along this line have further formulated NAS as a continuous and differentiable optimization problem to yield efficient gradient-based solutions . These one-shot NAS algorithms have achieved considerable improvement in search efficiency . However , the model training of the one-shot architecture is still needed . More recently , a number of algorithms have been proposed to estimate the performance of candidate architectures without model training . For example , Mellor et al . ( 2020 ) have explored the correlation between the divergence of linear maps induced by data points at initialization and the performance of candidate architectures heuristically . Meanwhile , Park et al . ( 2020 ) have approximated the performance of candidate architectures by the performance of their corresponding Neural Network Gaussian Process ( NNGP ) with only initialized model parameters , which is yet computationally costly . Abdelfattah et al . ( 2021 ) have investigated several training-free proxies to rank candidate architectures in the search space , while Chen et al . ( 2021 ) intuitively adopt theoretical aspects in deep networks ( e.g. , NTK ( Jacot et al. , 2018 ) and linear regions of deep networks ( Raghu et al. , 2017 ) ) to select architectures with a good trade-off between its trainability and expressivity . Our NASI significantly advances this line of work in ( a ) providing theoretically grounded performance estimation by NTK ( compared with ( Mellor et al. , 2020 ; Abdelfattah et al. , 2021 ; Chen et al. , 2021 ) ) , ( b ) guaranteeing the transferability of its selected architectures with its provable label- and data-agnostic search under mild conditions ( compared with ( Mellor et al. , 2020 ; Park et al. , 2020 ; Abdelfattah et al. , 2021 ; Chen et al. , 2021 ) ) ) and ( c ) achieving SOTA performance in a large search space over various benchmark datasets ( compared with ( Mellor et al. , 2020 ; Park et al. , 2020 ; Abdelfattah et al. , 2021 ) ) . 2.2 NEURAL TANGENT KERNEL ( NTK ) . Let a dataset ( X , Y ) denote a pair comprising a set X of m n0-dimensional vectors of input features and a vector Y ∈ Rmn×1 concatenating the m n-dimensional vectors of corresponding output values , respectively . Let a DNN be parameterized by θt ∈ Rp at time t and output a vector 1More precisely , we approximate the trace norm of NTK . f ( X ; θt ) ∈ Rmn×1 ( abbreviated to ft ) of the predicted values of Y . Jacot et al . ( 2018 ) have revealed that the training dynamics of DNNs with gradient descent can be characterized by an NTK . Formally , define the NTK Θt ( X , X ) ∈ Rmn×mn ( abbreviated to Θt ) as Θt ( X , X ) , ∇θtf ( X ; θt ) ∇θtf ( X ; θt ) > . ( 1 ) Given a loss function Lt at time t and a learning rate η , the training dynamics of the DNN can then be characterized as ∇tft = −η Θt ( X , X ) ∇ftLt , ∇tLt = −η ∇ftL > t Θt ( X , X ) ∇ftLt . ( 2 ) Interestingly , as proven in ( Jacot et al. , 2018 ) , the NTK stays asymptotically constant during the course of training as the width of DNNs goes to infinity . NTK at initialization ( i.e. , Θ0 ) can thus characterize the training dynamics and also the performance of infinite-width DNNs . Lee et al . ( 2019a ) have further revealed that , for DNNs with over-parameterization , the aforementioned training dynamics can be governed by their first-order Taylor expansion ( or linearization ) at initialization . In particular , define f lin ( x ; θt ) , f ( x ; θ0 ) +∇θ0f ( x ; θ0 ) > ( θt − θ0 ) ( 3 ) for all x ∈ X . Then , f ( x ; θt ) and f lin ( x ; θt ) share similar training dynamics over time , as described formally in Appendix A.2 . Besides , following the definition of NTK in ( 1 ) , this linearization f lin achieves a constant NTK over time . Given the mean squared error ( MSE ) loss defined as Lt , m−1 ‖Y − f ( X ; θt ) ‖22 and the constant NTK Θt = Θ0 , the loss dynamics in ( 2 ) above can be analyzed in a closed form while applying gradient descent with learning rate η ( Arora et al. , 2019 ) : Lt = m−1 ∑mn i=1 ( 1− ηλi ) 2t ( u > i Y ) 2 , ( 4 ) where Θ0 = ∑mn i=1 λi ( Θ0 ) uiu > i , and λi ( Θ0 ) and ui denote the i-th largest eigenvalue and the corresponding eigenvector of Θ0 , respectively . 3 NEURAL ARCHITECTURE SEARCH AT INITIALIZATION . 3.1 REFORMULATING NAS VIA NTK . Given a loss function L and model parameters θ ( A ) of architecture A , we denote the training and validation loss as Ltrain and Lval , respectively . NAS is conventionally formulated as a bi-level optimization problem ( Liu et al. , 2019 ) : minA Lval ( θ∗ ( A ) ; A ) s.t . θ∗ ( A ) , arg minθ ( A ) Ltrain ( θ ( A ) ; A ) . ( 5 ) Notably , model training is required to evaluate the validation performance of each candidate architecture in ( 5 ) . The search efficiency of NAS algorithms ( Real et al. , 2019 ; Zoph et al. , 2018 ) based on ( 5 ) is thus severely limited by the cost of model training for each candidate architecture . Though recent works ( Pham et al. , 2018 ) have considerably reduced this training cost by introducing a one-shot architecture for model parameter sharing , such a one-shot architecture requires training and hence incurs the training cost . To completely avoid this training cost , we exploit the capability of NTK for characterizing the performance of DNNs at initialization . Specifically , Sec . 2.2 has revealed that the training dynamics of an over-parameterized DNN can be governed by its linearization at initialization . With the MSE loss , the training dynamics of such linearization are further determined by its constant NTK . Therefore , the training dynamics and hence the performance of a DNN can be characterized by the constant NTK of its linearization . However , this constant NTK is computationally costly to evaluate . To this end , we instead characterize the training dynamics ( i.e. , MSE ) of DNNs in Proposition 1 using the trace norm of NTK at initialization , which can be efficiently approximated . For simplicity , we use this MSE loss in our analysis . Other widely adopted loss functions ( e.g. , cross entropy with softmax ) can also be applied , as supported in our experiments . Note that throughout this paper , the parameterization and initialization of DNNs follow that of Jacot et al . ( 2018 ) . For a L-layer DNN , we denote the output dimension of its hidden layers and the last layer as n1 = · · · = nL−1 = k and nL = n , respectively . Proposition 1 . Suppose that ‖x‖2 ≤ 1 for all x ∈ X and Y ∈ [ 0 , 1 ] mn for a given dataset ( X , Y ) of size |X | = m , a given L-layer neural architecture A outputs ft ∈ [ 0 , 1 ] mn as predicted labels of Y with the corresponding MSE loss Lt , λmin ( Θ0 ) > 0 for the given NTK Θ0 w.r.t . ft at initialization , and gradient descent ( or gradient flow ) is applied with learning rate η < λ−1max ( Θ0 ) . Then , for any t ≥ 0 , there exists a constant c0 > 0 such that as k →∞ , Lt ≤ mn2 ( 1− ηλ ( Θ0 ) ) q + ( 6 ) with probability arbitrarily close to 1 where q is set to 2t if t < 0.5 , and 1 otherwise , λ ( Θ0 ) , ( mn ) −1 ∑mn i=1 λi ( Θ0 ) , and , 2c0 √ n/ ( mk ) ( 1 + c0 √ 1/k ) . Its proof is in Appendix A.3 . Proposition 1 implies that NAS can be realizable at initialization . Specifically , given a fixed and sufficiently large training budget t , in order to select the best-performing architecture , we can simply minimize the upper bound of Lt in ( 6 ) over all the candidate architectures in the search space . Here , Lt can be applied to approximated Lval since both strong theoretical ( Mohri et al. , 2018 ) and empirical ( Hardt et al. , 2016 ) justifications in the literature have shown that training and validation loss are generally highly related . Hence , ( 5 ) can be reformulated as minA mn 2 ( 1− ηλ ( Θ0 ( A ) ) ) + s.t . λ ( Θ0 ( A ) ) < η−1 . ( 7 ) Note that the constraint in ( 7 ) is derived from the condition η < λ−1max ( Θ0 ( A ) ) in Proposition 1 , and η and are typically constants2 during the search process . Following the definition of trace norm , ( 7 ) can be further reduced into maxA ‖Θ0 ( A ) ‖tr s.t . ‖Θ0 ( A ) ‖tr < mnη−1 . ( 8 ) Notably , Θ0 ( A ) only relies on the initialization of A . So , no model training is required in optimizing ( 8 ) , which achieves our objective of realizing NAS at initialization . Furthermore , ( 8 ) suggests an interesting interpretation of NAS : NAS intends to select architectures with a good trade-off between their model complexity and the optimization behavior in their model training . Particularly , architectures containing more model parameters will usually achieve a larger ‖Θ0 ( A ) ‖tr according to the definition in ( 1 ) , which hence provides an alternative to measuring the complexity of architectures . So , maximizing ‖Θ0 ( A ) ‖tr leads to architectures with large complexity and therefore strong representation power . On the other hand , the complexity of the selected architectures is limited by the constraint in ( 8 ) to ensure a well-behaved optimization with a large learning rate η in their model training . By combining these two effects , the optimization of ( 8 ) naturally trades off between the complexity of the selected architectures and the optimization behavior in their model training for the best performance . Appendix C.1 will validate such trade-off . Interestingly , Chen et al . ( 2021 ) have revealed a similar insight of NAS to us .

Summary: This paper proposes a training-free NAS method called NASI, which exploits the Neural Tangent Kernel (NTK) to characterize the performance of the candidate architectures at initialization. To alleviate the costly evaluation for NTK, the authors apply a similar form to gradient flow to approximate NKT. Moreover, they combined their NTK trick with gradient-based NAS algorithm via Gumbel-Softmax to solve NAS problem efficiently. The experiment results on various benchmarks illustrate the effect of NASI.