Overcoming barriers to the training of effective learned optimizers

1 INTRODUCTION . Much of the success of modern deep learning has been driven by a shift from hand-designed features carefully curated by human experts , to domain-agnostic methods that can learn features from large amounts of data . By leveraging large-scale datasets with flexible models , we are now able to rapidly learn powerful features for new problem settings that often generalize to novel tasks . While learned features outperform hand-designed features on numerous tasks ( Krizhevsky et al. , 2012 ; Berner et al. , 2019 ; Vinyals et al. , 2019 ; Piech et al. , 2015 ) , we continue to use hand-designed optimization algorithms ( such as gradient descent , momentum , and so on ) for training models . These hand-designed update rules benefit from decades of optimization research but still require extensive expert supervision in order to be used effectively in machine learning . For example , they fail to flexibly adapt to new problem settings and require careful tuning of learning rate schedules and momentum timescales for different model architectures and datasets ( Choi et al. , 2019 ) . In addition , most do not leverage alternative sources of information beyond the gradient , such as the validation loss . By separating the optimization target ( training loss ) from the broader goal ( generalization ) , classic methods require more careful tuning of regularization and/or data augmentation strategies by the practitioner . To address these drawbacks , recent work on learned optimizers aims to replace hand-designed optimizers with a parametric optimizer , trained on a set of tasks , that can then be applied more generally . Recent work in this area has focused on either augmenting existing optimizers to adapt their own hyperparameters ( Daniel et al. , 2016 ; Xu et al. , 2017 ; 2019 ) , or developing more expressive learned optimizers to replace existing optimizers entirely ( Andrychowicz et al. , 2016 ; Wichrowska et al. , 2017 ; Lv et al. , 2017 ; Metz et al. , 2018 ; 2019a ; b ; Gu et al. , 2019 ) . These latter models take in problem information ( such as the current gradient of the training loss ) and iteratively update parameters . However , to date , learned optimizers have proven to be brittle and ineffective at generalizing across diverse sets of problems . Our work identifies fundamental barriers that have limited progress in learned optimizer research and addresses several of these barriers to train effective optimizers : 1 . Computational scale : Training a learned optimizer is costly . When training the optimizer , a single training step requires applying the optimizer to a training task for some number of unrolled steps . This work utilizes massive parallel computing infrastructure to scale the number of unrolled steps an order of magnitude larger than in previous work . 2 . Training tasks : Deep learning requires large training datasets . For learned optimizers to be effective , we similarly need a large dataset of optimization tasks on which to train . We build off of the TaskSet dataset ( Metz et al. , 2020 ) and construct a dataset of more than a six thousand diverse optimization tasks commonly found in machine learning . We show how this large and diverse task distribution is critical for training optimizers that generalize . 3 . Inductive bias of optimizer architecture : The parameterization of the learned optimizer and the task information fed to it both strongly affect performance . We propose a new hierarchical learned optimizer architecture that incorporates additional task information ( such as validation loss ) , and show that it outperforms previous learned optimizer architectures . By addressing these barriers , we develop learned optimizers that exceed prior work in scale , robustness , and out of distribution generalization . As a final test , we show that the learned optimizer can be used to train new learned optimizers from scratch ( analogous to “ self-hosting ” compilers ( Hart and Levin , 1962 ) ) . 2 PRELIMINARIES . Training a learned optimizer is a bilevel optimization problem that contains two loops : an inner loop that applies the optimizer to solve a task , and an outer loop that iteratively updates the parameters of the learned optimizer ( Franceschi et al. , 2018 ) . We use the inner- and outer- prefixes throughout to be explicit about which optimization loop we are referring to . That is , the inner-loss refers to a target task ’ s loss function that we wish to optimize , and the outer-loss refers to a measure of the optimizer ’ s performance training the target task ( inner-task ) . Correspondingly , we refer to the optimizer parameters as outer-parameters , and the parameters that the optimizer is updating as inner-parameters . Outer-optimization refers to the act of finding outer-parameters that perform well under some outer-loss . For a given inner-task , we apply the learned optimizer for some number of steps ( unrolling the optimizer ) . Ideally , we would unroll each target task until some stopping criterion is reached , but this is computationally infeasible for even moderate scale machine learning tasks . Each outer-iteration requires unrolling the optimizer on a target task . Short ( truncated ) unrolls are more computationally efficient , but suffer from truncation bias Wu et al . ( 2016 ) ; Metz et al . ( 2019b ) in that the outer-loss surface computed using truncated unrolls is different ( and may have different minima ) than the fully unrolled outer-loss . 3 METHODS : ADDRESSING THE THREE BARRIERS TO LEARNED OPTIMIZERS . 3.1 OUTER-OPTIMIZATION . To train the optimizer , we minimize an outer-loss that quantifies the performance of the optimizer . This is defined as the mean of the inner-loss computed on the inner-validation set for some number of unrolled steps , averaged over inner-tasks in the outer-training taskset . Although this outer-loss is differentiable , it is costly to compute the outer-gradient ( which involves backpropagating through the unrolled optimization ) . In addition , the outer-loss surface is badly conditioned and extremely non-smooth Metz et al . ( 2019b ) , making it difficult to optimize . We deal with these issues by using derivative-free optimization–specifically , evolutionary strategies ( ES ) Rechenberg ( 1973 ) –to minimize the outer-loss , obviating the need to compute derivatives through the unrolled optimization process . Previous work has used unrolled derivatives ( Andrychowicz et al. , 2016 ; Wichrowska et al. , 2017 ; Metz et al. , 2019b ) , and was thus limited to short numbers of unrolled steps ( e.g . 20 in Andrychowicz et al . ( 2016 ) and starting at 50 in Metz et al . ( 2019b ) ) . Without ES , the gradient estimates we obtain are extreamly high variance to the point that no training occurs . Using ES , we are able to use considerably longer unrolls . Initial unroll lengths were chosen to balance communication cost between parallel workers ( updating optimizer parameters ) with the computational cost of unrolling on individual workers ( estimating the local gradient with ES ) . We start outer-training by sampling unroll steps uniformly from 240-360 . When performance saturates with these biased gradients , continue training with an unbiased , but more expensive , gradient estimator ( see Appendix B ) . ES and PES have an additional benefit , in that optimizing with ES smooths the underlying loss function . This smoothing helps stabilize outer-training ( Metz et al. , 2019b ) . We set the standard deviation of the Gaussian distribution used by the ES algorithm ( which also controls how much the outer-loss is smoothed ) to 0.01 . To deal with the high variance of the ES estimate of the gradient , we use antithetic sampling and train in in parallel using 1024 multi-core CPU workers . While using more workers increases training speed , we find 1024 to be the point where performance gains become sub-linear . For more details see Appendix B . 3.2 TASK DISTRIBUTIONS . To train the optimizer , we need to define a set of inner-tasks to use for training . The choice of training tasks is critically important for two reasons : it determines the ability of the optimizer to outer-generalize ( i.e . the learned optimizer ’ s performance on new tasks ) , and it determines the computational complexity of outer-training . For improved outer-generalization , we would like our inner-problems to be representative of tasks we care about . In this work , these are state-of-the-art machine learning models such as ResNets ( He et al. , 2016a ) or Transformers ( Vaswani et al. , 2017 ) . Unfortunately , directly utilizing these large scale models is computationally infeasible , therefore we outer-train on proxy tasks for speed ( Zoph et al. , 2018 ) . In order to outer-train a learned optimizer capable of generalizing to new optimization tasks , we utilize an outer-training task set consisting of around 6,000 tasks designed after Metz et al . ( 2020 ) . These tasks include RNNs ( Hochreiter and Schmidhuber , 1997 ; Chung et al. , 2014 ) , CNNs ( LeCun , 1998 ) , masked auto regressive flows ( Papamakarios et al. , 2017 ) , fully connected networks , language modeling , variational autoencoders ( Kingma and Welling , 2013 ) , simple 2D test functions , quadratic bowls , and more . For tasks that require them , we additionally sample a dataset , batch size , network architecture , and initialization scheme . To keep outer-training efficient , we ensure that all tasks take less than 100 milliseconds per-training step . For each task that makes use of a dataset , we create four splits of the data to prevent leakage : training data , which we compute gradients on and use to update the inner-parameters ; inner-validation data , which is used to compute validation losses used by the learned optimizer ; outer-validation data , which is used to update the weights of the learned optimizer ; and test data , which is used to test an already trained learned optimizer . Because loss values vary in magnitude , when outer-training we normalize these outer-loss values by the best loss achieved by a baseline optimizer and the initial loss value . Note this normalization is not used during inner-training . 3.3 OPTIMIZER ARCHITECTURE . Designing a learned optimizer architecture requires balancing computational efficiency and expressivity . Past work in learned optimizers has shown that incorporating inductive biases based on existing optimization techniques such as momentum or second moment accumulation leads to better performance ( Wichrowska et al. , 2017 ; Metz et al. , 2019b ) . The optimizer we use in this work consists of a hierarchical optimizer similar to ( Wichrowska et al. , 2017 ) ( Figure 1 ) . A per-tensor LSTM is run on features computed over each parameter tensor . This LSTM then forwards information to the other tensors ’ LSTMs as well as to a per-parameter feedforward neural network . The per-parameter feedforward network additionally takes in information about gradients and parameter value , and outputs parameter updates . Additional outputs are aggregated and fed back into the per-tensor network . This information routing allows for communication across all components . For per-parameter features we leverage effective inductive biases from hand-designed optimizers , and use a variety of features including the gradient , the parameter value , and momentum-like running averages of both . All features are normalized , in a fashion similar to that in RMSProp ( Tieleman and Hinton , 2012 ) , or relative to the norm across the full tensor . For per-tensor features we use a variety of features built from summary statistics computed from the current tensor , the tensor ’ s gradient , and running average features such as momentum and second moments . We also include information about the tensor ’ s rank and shape . We also feed global features into the per-tensor LSTM , such as training loss and validation loss , normalized so as to have a relatively consistent scale across tasks . To compute a weight update , the per-parameter MLP outputs two values , ( a , b ) , which are used to update inner-parameters : wt+1 = wt + exp ( a ) b . See Appendix C for many more details .

Summary: The authors propose to use a combination of Andrychowicz et al. LSTM based approach and Metz et.al feed forward network to learn an optimizer that is useful accross any task. The authors propose to use a few thousand tasks as developed in Metz et. al. for the purpose of training the optimizer. They focus on evolutionary strategies (ES) in lieu of unrolled optimization to keep the problem tractable.