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 goal of a learned optimizer is to replace a human-designed optimizer with a parametric optimizer. However, prior learned optimizers were ineffective at generalizing to a diverse set of tasks. This paper investigates how to learn a useful optimizer by increasing computational scale, building a large, diverse training dataset, and designing the learned optimizer's architecture.