AutoLoss: Learning Discrete Schedule for Alternate Optimization

1 INTRODUCTION . Many machine learning ( ML ) problems involve iterative alternate optimization of different objectives { ` m } Mm=1 w.r.t different sets of parameters { θn } Nn=1 until a global consensus is reached . For instances , in training generative adversarial networks ( GANs ) ( Goodfellow et al. , 2014 ) , parameters of the generator and the discriminator are alternately updated to an equilibrium ; in many multi-task learning problems ( Argyriou et al. , 2007 ) , one usually has to alternate the optimization of different taskspecific objectives on corresponded data , until the target task performance is maximized . In these processes , one needs to determine which objective ` m and which set of parameters θn to choose at each step , and subsequently , how many iteration steps to perform for the subproblem minθn ` m. We refer to this as determining an optimization schedule ( or update schedule ) . While extensive research has been focused on developing better optimization algorithms or update rules ( Kingma & Ba , 2014 ; Bello et al. , 2017 ; Duchi et al. , 2011 ; Sutskever et al. , 2013 ) , how to select optimization schedules has remained less studied . When the objective is complex ( e.g . nonconvex or combinatorial ) and the parameters to be optimized are high-dimensional , the optimization schedule can directly impact the quality of convergence . However , we hypothesize that the schedule is learnable in a data-driven way , with the following empirical evidence : ( 1 ) The optimization of many ML models is sensitive to the update schedule . For examples , the updates of the generator and the discriminator in GANs are carefully reconciled to avoid otherwise model collapse or gradient vanishing ( Goodfellow et al. , 2014 ; Radford et al. , 2015 ) ; In solving many multi-task learning or regularizer-augmented objectives , the optimization target L is a combination of multiple task-specific objectives . It is desirable to weight each objective differently as L = ∑M m=1 λm ` m , while different values of { λm } Mm=1 result in different ( local ) optima . This indicates that different loss terms shall not be treated equally , and achieving the best downstream task performance requires optimizing every ` m to different extents . ( 2 ) Previous research and practice have suggested that there do exist optimization schedules that are more probable to produce better convergence than random ones , e.g . Arjovsky et al . ( 2017 ) and Salimans et al . ( 2016 ) suggest that keeping the steps of updating the generator and discriminator of GANs at K : 1 ( K > 1 ) leads to faster and more stable training of GANs . Based on the hypothesis , in this paper , we develop AutoLoss , a generic meta-learning framework to automatically determine the optimization schedule in iterative and alternate optimization processes . AutoLoss introduces a parametric controller attached to an alternate optimization task . The controller is trained to capture the relations between the past history and the current state of the optimization process , and the next step of the decision on the update schedule . It takes as input a set of status features , and decides which objectives from { ` m } Mm=1 to optimize , and which set of parameters from { θn } Nn=1 to update . The controller is trained via policy gradient to maximize the eventual outcome of the optimization ( e.g . downstream task performance ) . Once trained , it can guide the optimization of task models to achieve higher quality of convergence faster , by predicting better schedules . To evaluate the effectiveness of AutoLoss , we instantiate it on four typical ML tasks : d-ary quadratic regression , classification using a multi-layer perceptron ( MLP ) , image generation using GANs , and neural machine translation based on multi-task learning . We propose an effective set of features and reward functions that are suitable for the controllers ’ learning and decisions . We show that , on all four tasks , the AutoLoss controller is able to capture the distribution of better optimization schedules that result in higher quality of convergence on the corresponding task than strong baselines . For examples , on quadratic regression with L1 regularization , it learns to detect the potential risk of overfitting , and incorporates L1 regularization when necessary , helps the task model converge to better results that can hardly be achieved by optimizing linear combinations of objective terms . On GANs , the AutoLoss controller learns to balance the training of generator and discriminator dynamically , and report both faster per-epoch convergence and better quality of generators after convergence , compared to fixed heuristic-driven schedules . On machine translation , it automatically learns to resemble human-tuned update schedules while being more flexible , and reports better perplexity results . In summary , we make the following contributions in this paper : ( 1 ) We present a unified formulation for iterative and alternate optimization processes , based on which , we develop AutoLoss , a generic framework to learn the discrete optimization schedule of such processes using reinforcement learning ( RL ) . To our knowledge , this is the first framework that tries to learn the optimization schedule in a data-driven way . ( 2 ) We instantiate AutoLoss on four ML tasks : d-ary regression , MLP classification , GANs , and NMT . We propose a novel set of features and reward functions to facilitate the training of AutoLoss controllers . ( 3 ) We empirically demonstrate AutoLoss ’ efficacy : it delivers higher quality of convergence for all four tasks on synthetic and real dataset than strong baselines . Training AutoLoss controller has acceptable overhead less than most hyperparameter searching methods ; the trained AutoLoss controller is generalizable – it can guide and improve the training of a new task model with different specifications , or on different dataset . 2 RELATED WORK . Alternate Optimization . Many ML models are trained using algorithms with iterative and alternate workflows , such as EM ( Moon , 1996 ) , stochastic gradient descent ( SGD ) ( Bottou , 2010 ) , coordinate descent ( Wright , 2015 ) , multi-task learning ( Zhang & Yang , 2017 ) , etc . AutoLoss can improve these processes by learning a controller in a data-driven way , and figuring out better update schedules using this controller , as long as the schedule does affect the optimization goal . In this paper , we focus mostly on optimization problems , but note AutoLoss is applicable to alternate processes that involve non-optimization subtasks , such as sampling methods ( Griffiths & Steyvers , 2004 ; Ma et al. , 2015 ) . Meta learning . Meta learning ( Andrychowicz et al. , 2016 ; Maclaurin et al. , 2015 ; Wang et al. , 2016 ; Finn et al. , 2017 ; Chen et al. , 2016 ) has drawn considerable interest and been recently applied to improve the optimization of ML models ( Ravi & Larochelle , 2016 ; Li & Malik , 2016 ; Bello et al. , 2017 ; Fan et al. , 2018 ) . Among these works , the closest to ours are Li & Malik ( 2016 ) , Bello et al . ( 2017 ) , Fan et al . ( 2018 ) . Li & Malik ( 2016 ) propose learning to optimize to directly predict the gradient values at each step of SGD . Since the gradients are continuous and usually high-dimensional , directly regressing their values might be difficult , and the learned gradient regressor is nontransferable to new models or tasks . Differently , Bello et al . ( 2017 ) propose to learn better gradient update rules based on a domain specific language . The learned rules outperform manually designed ones and are generalizable . AutoLoss differs from both – instead of learning to generate values of updates ( gradients ) , AutoLoss focuses on producing better scheduling of updates . Therefore AutoLoss can model other classes of problems such as scheduling the generator and discriminator training in GANs ( Larsen & Sønderby ) , or even go beyond optimization problems . In Fan et al . ( 2018 ) , a learning to teach framework is proposed that a teacher model , trained by optimization metadata , can guide the learning of student models . AutoLoss instantiates the framework in the sense that the teacher model ( controller ) produces better schedules for the task model ( student ) optimization . Also of note is another line of works that apply RL to enable automatic machine learning ( AutoML ) , such as adaptive step size controller ( Daniel et al. , 2016 ) , device placement optimization ( Mirhoseini et al. , 2017 ) , neural architecture search ( Baker et al. , 2016 ; Zoph & Le , 2016 ) , etc . While addressing different problems , AutoLoss ’ controller is trained in a similar way ( Peters & Schaal , 2008 ) for sequential and discrete predictions . 3 AUTOLOSS . Background . In most ML tasks , given observed data D , we aim to minimize an objective function L ( D ; Θ ) with respect to the parameters Θ of the model that we use to characterize the data . Solving this minimization problem involves finding the optima of Θ ( denoted as Θ∗ ) , which we usually resort to a variety of de facto optimization methods ( Boyd & Vandenberghe , 2004 ) if close-formed solutions are unavailable . In the rest of the paper , we will focus on two typical classes of optimization workflows which many modern ML model solvers would fall into : iterative and alternate optimization . Iterative optimization methods look for the optimal parameter Θ∗ in an iterative-convergent way , by repeatedly updating Θ until certain stopping criteria is reached . Specifically , at iteration t , the parameters Θ are updated from Θ ( t ) to Θ ( t+1 ) following the update equation Θ ( t+1 ) = Θ ( t ) + · ∆L ( D ( t ) ; Θ ( t ) ) , where we denote ∆L as the function that calculates update values of Θ depending on L , D ( t ) ⊆ D as a subset of D used at iteration t and a scaled factor . Many widely-adopted algorithms ( Bottou , 2010 ; Boyd et al. , 2003 ) fall into this family , e.g . in the case for SGD , ∆L reduces to deriving the gradient updates∇Θ ( we skip optional steps such as momentum or projection for clarity ) , D ( t ) is a stochastic batch , and is the learning rate . To describe alternate optimization , we notice the objective L is usually composed of multiple different optimization targets : L = { ` m } Mm=1 , and we want Θ∗ to minimize a certain combination of them . For example , when fitting a regression model with mean square error ( MSE ) , appending an L1 loss helps obtain sparsity ; in this case , L is written as a linear combination of MSE and L1 terms . Similarly , the parameters Θ in many cases are also composable , e.g . when the model has multiple components with independent sets of parameters . If we decompose Θ = { θn } Nn=1 , an alternate optimization ( in our definition ) contains multiple steps , where each step t involves choosing ` mt ∈ L , θnt ∈ Θ , which we will call as determining an optimization action ( notated as a ) , and update θnt w.r.t . ` mt . Further , we note that many ML optimization tasks in practice are both iterative and alternate , such as the training process of GANs , where the updates of generator and discriminator parameters are alternated , each with a few iterations of stochastic updates , until equilibrium . We therefore present iterative and alternate optimization with the following unified formulation : for t = 1→ T , choose ( ` mt , θnt ) = aqt ∈ A , update θ ( t+1 ) nt = θ ( t ) nt + ·∆ nt ` mt , ( 1 ) where A = { aq } Qq=1 denotes the task-specific action space that defines all legitimate pairs of loss and parameter to choose from ; ∆nt ` mt are update values of θnt w.r.t . ` mt . Eq . 1 reduces to the vanilla form of iterative optimization when A = { ( L , Θ ) } . AutoLoss . Given the formulation in Eq . 1 , our goal is to determine aqt , i.e . which losses to optimize and what parameters to update at each t , in order to maximize the downstream task performance . We introduce a meta model , which we call controller , to be distinguished from the task model used in the downstream task . The controller is expected to learn during its exploration of task model optimization processes , and is able to decide how to update once sufficient knowledge has been accumulated . Specifically , we let the controller make sequential decisions at each step t ; it scans through the past history and the current states of the process ( described as a feature vectorX ( t ) ∈ RK ) , and predicts a one-hot vector Y ( t ) ∈ { 0 , 1 } |A| , i.e . aq ∈ |A| will be selected if the qth entry of Y ( t ) is 1 . We model our controller as a conditional distribution p ( y|x ; φ ) parameterized by φ1 , where we denote y and x as the |A|-dim decision variable and K-dim feature variable , respectively . At each step t , we sample Y ( t ) ∼ p ( y|x = X ( t ) ; φ ) , and perform updates following Eq . 1 and Y ( t ) . Parameter learning . The parameters of the controller φ is trained to maximize the performance of the optimization task given sampled sequences of decisions within T steps , notated as Y = { Y ( t ) } Tt=1 . Accordingly , we introduce the training objective of the controller as J ( φ ) = EY∼p ( y|x ; φ ) [ R ( Y ) |L , Θ ] , where R ( · ) is the reward function that evaluates the final task 1The other alternative is to condition the decision at the t step on the decision made at the t− 1 step , though we choose a simpler one to highlight the generic idea behind AutoLoss . performance after applying the schedule Y for its optimization . Since the decision process involves non-differentiable sampling , we learn the parameters using REINFORCE ( Williams , 1992 ) , where the unbiased policy gradients at each updating step of the controller are estimated by sampling S sequences of decisions { Ys } Ss=1 ( for all experiments we set S = 1 ) and compute ∇φJ ( φ ) = 1 S S∑ s=1 [ ( R ( Ys ) −B ) · ∇φ T∑ t=1 log p ( Y ( t ) s |X ( t ) s ; φ ) |L , Θ ] , ( 2 ) where Y ( t ) s is the tth decision in Ys . To reduce the variance , we introduce a baseline term B in Eq . 2 to stabilize the training ( similar to Pham et al . ( 2018 ) ) , where B is defined as a moving average of received reward : B ( h+1 ) ← ηB ( h ) + ( 1 − η ) R ( h ) with η as a decay factor . Whenever applicable , R ( Ys ) −B is further clipped to a given range . See detailed training algorithms in Appendix A.1 .

Summary: This paper proposes a meta-learning solution for problems involving optimizing multiple loss values. They use a simple (small mlp), discrete, stochastic controller to control applications of updates among a finite number of different update procedures. This controller is a function of heuristic features derived from the optimization problem, and is optimized using policy gradient either exactly in toy settings or in a online / truncated manor on larger problems. They present results on 4 settings: quadratic regression, MLP classification, GAN, and multi-task MNT. They show promising performance on a number of tasks as well as show the controllers ability to generalize to novel tasks.