DiffAutoML: Differentiable Joint Optimization for Efficient End-to-End Automated Machine Learning

1 INTRODUCTION . While deep learning has achieved remarkable progress in various tasks such as computer vision and natural language processing , it usually requires tremendous human involvement to design and train a satisfactory deep model for one task ( He et al. , 2016 ; Sandler et al. , 2018 ) . To alleviate such burden on human users , a dozen of AutoML algorithms are proposed in recent years to enable training a model from data automatically without human experiences , including automated data augmentation ( DA ) , neural architecture search ( NAS ) , and hyper-parameter optimization ( HPO ) ( e.g. , Chen et al. , 2019 ; Cubuk et al. , 2018 ; Mittal et al. , 2020 ) . These AutoML components are usually developed independently . However , implementing these AutoML components for a specific task in separate stages not only suffers from low efficiency but also leads to sub-optimal results ( Dai et al. , 2020 ; Dong et al. , 2020 ) . How to achieve full-pipeline “ from data to model ” automation efficiently and effectively is still a challenging problem . One main difficulty for achieving automated “ from data to model ” is how to combine different AutoML components ( e.g. , NAS and HPO ) appropriately for a specific task . Optimizing these components in a joint manner is an intuitive solution but usually suffers from the enormous and impractical search space . Dai et al . ( 2020 ) and Wang et al . ( 2020 ) introduced pre-trained predictors to achieve the joint optimization of NAS and HPO , and the joint optimization of NAS and automated model compression , respectively . For a new coming task , however , it is usually burdensome to pre-train such a predictor . On the other hand , Dong et al . ( 2020 ) investigated the joint optimization between NAS and HPO via differentiable architecture and hyper-parameter search spaces . Automated DA is seldom considered in the joint optimization of AutoML components . Nevertheless , our experimental results showed that different data augmentation protocols may result in different optimal architectures ( see Sec . 4.3 for more details ) . Based on these considerations , it is worthy to investigating the joint optimization among automated DA , NAS and HPO . Another main challenge to achieve automated “ from data to model ” is the end-to-end searching and training of models without parameter retraining . Even considering only one AutoML component , many NAS algorithms require two stages including searching and retraining ( e.g. , Liu et al. , 2018 ; Xie et al. , 2019 ) . Automated DA methods such as Lim et al . ( 2019 ) also needed to retrain the model parameters when the DA policies were searched . In these cases , whether the searched architectures or DA policies would perform well after retraining is questionable , due to the inevitable difference of training setup between the searching and retraining stages . Recently , Hu et al . ( 2020 ) developed a differentiable NAS methods to provide direct NAS without parameter retraining . Other AutoML components , including HPO and automated DA , are seldom considered in the task-specific end-toend AutoML algorithms . Considering the above challenges , we propose DiffAutoML , a differentiable joint optimization solution for efficient end-to-end AutoML . In DiffAutoML , end-to-end NAS optimization is realized in a differentiable one-level manner with the help of stochastic architecture search and the approximation for the gradient of architecture parameters . Meanwhile , the DA and HPO are regarded as dynamic schedulers , which adapt themselves to the update of network parameters and network architecture . Specifically , Differentiable relaxation is used in DA optimization in an one-level way while the hyper-gradient is used to in HPO in a two-level way . With this differentiable method , DiffAutoML can effectively deal with the huge search space and the low optimization efficiency caused by this joint optimization problem . To summarize , our main contributions are as follows : • We propose a well-defined AutoML problem , i.e. , task-specific end-to-end AutoML framework , which aims to achieve automated “ from data to model ” without human involvement . • We first jointly optimize three different AutoML components , including automated DA , NAS and HPO , in a differentiable search space with high efficiency . • Experiment results show that , compared with optimizing modules in sequence , one-stage DiffAutoML can effectively realize the co-optimization of different modules . • Extensive experiments also reveal mutual influence among different AutoML components , i.e. , the change of settings for one module may greatly influence the optimal results of another module , justifying the necessity of end-to-end joint optimization . 2 RELATED WORKS . In this section , we briefly introduce some related AutoML algorithms , including automated DA , NAS and HPO . A more detailed introduction is provided in Appendix D. Automated Data Augmentation Data augmentation is commonly used to increase the data diversity and thus to improve the generalization performance of deep learning models ( DeVries & Taylor , 2017 ; Zhang et al. , 2017 ; Yun et al. , 2019 ) . Several prior works have been proposed to automate the search of data augmentation policies , using methods ranging from reinforcement learning ( Cubuk et al. , 2018 ; Zhang et al. , 2020 ) , evolutionary algorithm ( Ho et al. , 2019 ) , Bayesian optimization ( Lim et al. , 2019 ) , gradient-based approaches ( Lin et al. , 2019 ) to simple grid search ( Cubuk et al. , 2020 ) . Cubuk et al . ( 2020 ) point out that the optimal data augmentation policy depends on the model size . This indicates that fixing an augmentation policy when searching for neural architectures may lead to sub-optimal solutions , thus motivating joint optimization . Neural Architecture Search Recent advances in NAS have demonstrated state-of-the-art performance outperforming human experts ’ design on a variety of tasks ( Zoph & Le , 2017 ; Cai et al. , 2018 ; Liu et al. , 2018 ; Kandasamy et al. , 2018 ; Pham et al. , 2018 ; Real et al. , 2018 ; Zoph et al. , 2018 ; Ru et al. , 2020b ) . Especially the development of gradient-based one-shot methods ( Liu et al. , 2018 ; Chen et al. , 2019 ; Guo et al. , 2019 ; Xie et al. , 2019 ; Hu et al. , 2020 ) have significantly reduced the computational costs of NAS . To further improve the search efficiency , most NAS methods search architectures in a low-fidelity set-up ( e.g. , fewer training epochs , smaller architectures ) and retrain the optimal architecture using the full set-up before deployment . This separation of search and evaluation , however , usually suffer from sub-optimal results ( Hu et al. , 2020 ) . End-to-end NAS strategies ( Xie et al. , 2019 ; Hu et al. , 2020 ) are thereby developed to return read-to-deploy networks at the end of the search . Our work also proposes an end-to-end solution . Hyper-parameter Optimization Tuning hyper-parameters for deep learning models has also attracted lots of research attention . Current HPO methods address various set-ups involving multifidelity evaluations , parallel computation , transfer-learning ( Mittal et al. , 2020 ) or a mixed variable types ( Ru et al. , 2020a ) . Although many HPO strategies have been adapted to form NAS strategies , methods that jointly optimize both architectures and hyper-parameters are rarely seen except the ones discussed below . Joint Optimization of AutoML Components . Conventional neural architecture search methods perform search under a fixed set of training hyper-parameters and apply or search for a separate set of hyper-parameters when retraining the best architecture found . Such search protocol may lead to sub-optimal results ( Zela et al. , 2018 ; Dong et al. , 2020 ) as it neglects the influence of training hyper-parameters on architecture performance and ignores superior architectures under alternative hyper-parameter values ( Dai et al. , 2020 ) . In view of this , several works have been proposed to jointly optimise architecture structure and training hyper-parameters ( Dai et al. , 2020 ; Wang et al. , 2020 ; Dong et al. , 2020 ) . Zela et al . ( 2018 ) introduces the use of multi-fidelity Bayesian optimization to search over both the architecture structure and training hyper-parameters . Dai et al . ( 2020 ) trains an accuracy predictor to estimate the network performance based on both the architecture and training hyper-parameters , and then uses evolutionary algorithm to perform the search . Both these methods are sample-based and requires a relatively large number of architecture and hyperparameter evaluations to finetune predictor or obtain good recommendations . A more related work is AutoHAS ( Dong et al. , 2020 ) , which introduces a differentiable approach in conjunction with weight sharing for the co-optimization task . Dong et al . ( 2020 ) demonstrates empirically that such differentiable one-shot approach achieve superior efficiency over sampled-based methods . Our proposed method differs from AutoHAS in mainly two aspects : first , AutoHAS updates the entire supernetwork at each optimization step while our method trains only a subnetwork . Thus , our method is much more memory efficient than AutoHAS . Second , we further extend the co-optimization scope from NAS and training hyper-parameters to also include data augmentation hyper-parameters . To the best of our knowledge , we are the first to jointly optimize all three aspects and we demonstrate the advantage of such practice in Sec . 4 . 3 METHODOLOGY . Consider a dataset D = { ( xi , yi ) } Ni=1 , where N is the size of this dataset , and yi is the label of the input sample xi . We aim to train a neural network f ( · ) , which can achieve the best accuracy on the test dataset Dtest . Multiple AutoML components are considered , including automated DA , NAS , and HPO . Let τ , α , η , and θ represent the data augmentation parameters , the architecture parameters , the hyper-parameters , and the objective neural network parameters , respectively . This problem can be formulated as argmin τ , α , η , θ L ( τ , α , η , θ ; D ) s.t . ci ( α ) ≤ Ci , i = 1 , ... , γ , ( 1 ) whereL ( · ) represents the loss function ; D denotes the input data ; ci ( · ) and γ refer to the formula and count of resource constraints Ci , such as storage cost and computational cost . We adopt different loss functions to optimize the corresponding parameters , which will be introduced in the following subsections . Considering the huge search space , it is challenging to achieve the co-optimization of τ , α , η , and θ within one-stage without parameter retraining . In this work , we propose to use the differentiable method to provide a computationally efficient solution . See Fig . 1 as an illustration . 3.1 OPTIMIZATION OF DATA AUGMENTATION PARAMETERS . For every mini-batch of training data Btr = { ( xk , yk ) } n tr k=1 with batch size n tr , we conduct data augmentation to increase the diversity of the training data . We consider K data augmentation operations , and each training sample is augmented by a transformation consisting of two successive operations ( Cubuk et al. , 2018 ; Lim et al. , 2019 ) . Each operation is associated with a magnitude that is uniformly sampled from [ 0 , 10 ] . The data augmentation parameter τ represents a probability distribution over the augmentation transformations . For t-th iteration , we sample ntr transformations according to τ t with Gumbel-Softmax reparameterization ( Maddison et al. , 2016 ) and to generate the corresponding augmented samples in the batch . Given a sampled architecture ( see details in Sec . 3.2 ) , the loss function for each augmented sample is denoted by Ltr ( f ( αt , θt ; Tk ( xk ) ) ) , where Tk represents the selected transformation . In order to relax τ to be differentiable , we regard pk ( τ t ) , the probability of sampling the transformation Tk , as an importance weight for the loss function Ltr ( f ( αt , θt ; Tk ( xk ) ) ) . The objective of data augmentation is to minimize the following loss function : LDA ( τ t ) = − ntr∑ k=1 pk ( τ t ) Ltr ( f ( αt , θt ; Tk ( xk ) ) ) , ( 2 ) With this loss function , DiffAutoML intends to increase the sampling probability of those transformations that can generate samples with high training losses . By sampling such transformations , DiffAutoML can pay more attention to relatively “ hard ” samples and increase model robustness against difficult samples . Instead of training a controller to generate adversarial augmentation policies via reinforcement learning ( Zhang et al. , 2020 ) , we search for the probability distribution of augmentation transformations directly via gradient-based optimization . In this way , the optimization of data augmentation is very efficient and hardly increases the computing cost .

Summary: This paper shows how the complex autoML pipeline for neural networks can be trained in an end-to-end manner by combining existing methods. By using backpropagatable discrete sampling methods (Gumbel softmax), input transformed by data augmentation is seamlessly embedded in full backpropagation flow. And a differentiable architecture search algorithm is used, which also incorporates architecture search in full backpropagation flow. On top of this differentiable procedure, an alternating optimization is introduced to train network parameters and hyperparameters.