Overcoming barriers to the training of effective learned optimizers
Authors: Luke Metz
Niru Maheswaranathan
C. Daniel Freeman
Ben Poole
Jascha Sohl-Dickstein
Venue:  (2021)
Keywords: learned optimizers
meta-learning

Abstract: In this work we focus on general-purpose learned optimizers capable of training a wide variety of problems with no user-specified hyperparameters. We introduce a new, neural network parameterized, hierarchical optimizer with access to additional features such as validation loss to enable automatic regularization. Most learned optimizers have been trained on only a single task, or a small number of tasks. We train our optimizers on thousands of tasks, making use of orders of magnitude more compute, resulting in optimizers that generalize better to unseen tasks. The learned optimizers not only perform well, but learn behaviors that are distinct from existing first order optimizers. For instance, they generate update steps that have implicit regularization and adapt as the problem hyperparameters (e.g. batch size) or architecture (e.g. neural network width) change. Finally, these learned optimizers show evidence of being useful for out of distribution tasks such as training themselves from scratch.

Under review as a conference paper at ICLR 2021
OVERCOMING BARRIERS TO THE TRAINING OF EFFEC- TIVE LEARNED OPTIMIZERS
Anonymous authors Paper under double-blind review
ABSTRACT
In this work we focus on general-purposs learned optimizers capable of training a wide variety of problems with no user-specified hyperparameters. We introduce a new, neural network parameterized, hierarchical optimizer with access to additional features such as validation loss to enable automatic regularization. Most learned optimizers have been trained on only a single task, or a small number of tasks. We train optimizers thousands of tasks, making use of orders of magnitude more compute, resulting in optimizers that generalize better to unseen tasks. The learned optimizers not only perform well, but learn behaviors that are distinct from existing first order optimizers. For instance, they generate update steps that have implicit regularization and adapt the problem hyperparameters (e.g. batch size) or architecture (e.g. neural network width) change. Finally, these learned optimizers show evidence of being useful for out of distribution tasks such as training themselves from scratch.
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 comain-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 ctt al. 2019 Vinyals et al.) 2019 Piech et al., (015),, 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 flcxibly 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.) (019).. 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 hand-designed opti- mizers 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., (017]] 2019), or developing more expressive learned optimizers to replace existing optimizers entirely (Andrychowicz.c 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 prob- lem 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 costly. When training the optimizer, single training step requires applying the optimizer to training task for some number

---

Under review as a conference paper at ICLR 2021
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 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 hier-
and the task information fed to it both strongly affect performance. We propose archical 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, robust- ness, and out of distribution generalization. As final test, we show that the learned optimizer can be used to train new learned optimizers from scratch (analogous to "self selff-hottngg compilers (Hart and Levin 1962))
PRELIMINARIES
Training a learned optimizer a bilevel optimization problem that contains two loops: an inner loop that applies the optimizer to solve task, and an outer loop that iteratively updates the parameters of the learned optimizer (Franceschi et al. 2018). We use the imner- 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-parameten and the parameters that the optimizer is updating as inner-parumelers. 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 computationally infeasible for even moderate scale machine learning tasks. Each outer- ittertiteration requires unrolling the optimizer on target task. Short (truncated) unrolls are more computationally efficient, but suffer from truncation bias Wu al. (2016); Metz al. (2019)) 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-OPTIMIZAN
To train the optimizer, minimize 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 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 free optimizationnsspeffccllllll specifically, evolutionary strategies Rechhenbbg (1973) minimize the outer-loss, obviating the need to compute derivatives through the unrolled optimization process. Previous work has used unrolled derivatives (Andrychowiez et al. Wichrowsk et al. 2017 Metz et al. (0019)), 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
2

---

Under review as a conference paper at ICLR 2021
(a)
(b)
train valid loss tensor shape gradient norm
LSTM
gradient
gradient
LSTM
AW
momentum
FF
AW
gradient momentum second moment
FF
AW
Figure 1: (a) The learned optimizer architecture proposed AnAryycchoiizz et al. (2016) consisting of a per-parameter LSTM. (b) The learned optimizer architecture proposed in Metz et al. (2019b) consisting of per-parameter fully connected feeedforwwrrd FF) neural network with additional input features. (c) The learned optimizer architecture proposed in this work consisting of a per-tensor LSTM which exchanges information with a per-parameter feed forward neural network (FF). The LSTMs associated with each tensor additionally share information with each other (shown in gray).
these biased gradients, continue training with an unbiased, but more expensive, gradient estimator (see Appendix
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 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 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-generalizzaion,, inner-problems to be representative of tasks we care about. In this work, these are state-of-the-ar machine learning models such 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 2018).
would like
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 al. (2020). These tasks include RNNs (Hochreiter and Schmidhuber 1997 | 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). 3), 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- outer-training efficient, less than 100 milliseconds per-trame ver-training step. For each task that makes use of dataset, we create four splits of the data prevent leakage: training data, which the inner- parameters; inner-validation data, which is used to compute validation losses used by the learned optimizer; ouler-validation data, which 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 baseline optimizer and the initial loss value. Note this normalization is not used during inner-training.
ensure that all tasks take
compute gradients and use to update
3.3 OPTIMIZER ARCHITECTURE
Designing a learned optimizer architecture requires balancing computational efficiency and expressiv- ity. 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 perfor- mance (Wichrowska al. 2017] Metz et al. 2019b). The optimizer of a hierarchical optimizer similar to (Wichrowska et al,, 2017) (Figure [T. A per-tensor LSTM is on features computed over each parameter tensor. This LSTM then forwards information the
in this work consists
3

---

Under review as a conference paper at ICLR 2021
(a)
L2LBGDCD
LSTM FF
LFTMLFF
0.5
0.40
outer-loss on test
0.35
training tasks
availation loss
0.30
0.25
0.2
0.20
0.1
0.15
5000 10000
15000
109
101
outer-training Reertiooe
Figure 2: (a) Our proposed learned optimizer has greater sample efficiency than existing methods. On the x-axis we show outer-updates performed to outer-train each learned optimizer. On the y-axis we show outer-loss. Each point consists of using single outer parameter value train 100 models for 10k inner-iterations, each with five random initializations. We show the average validation performance post-normalization averaged over all tasks, seeds, and inner- aner-training steps. (b) Outer- generalization performance increases with an increasing number of outer-traine snter-training tasks. We show the mean performance achieved after 10000 outer-training iterations (between 10500 and 9500) averaged over four random seeds while varying outer training set size. We show each seed as a star, and standard deviation as the shaded region. We find poor guteraeenerllization for smaller numbers of tasks. As increase the number of outer-training tasks performance becomes more consistent, and lower loss values are achieved.
other tensors' LSTMs as well as to a per- per-parameter feedforward neural network. The per 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 desigddd optimizers, and use variety of features including the gradient, the parameter value, and momentum-like like running averages of both. All features are normalized, in a fashion similar to that RMSProp (Tieleman and variety
Hinton 2012), relative to the 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 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 weight update, the per- per-parameter MLP outputs two values, (a, b). which are used to update inner- paremeeaamtters: w w' + exp(a)b. See Appendix C for many more details.
across the full tensor. For per-tensor features
4 RESULTS
4.1 COMPARING LEARNED OPTIMIZER ARCHITECTURES AND TRAINING TASK SET SIZES
First, we show experiments comparing the performance of different learned optimizer architectures from the literature. We trained: component-wise LSTM optimizer from Andrychowicz al. (2016) (L2LBGDGD), a modification of this LSTM with the decomposed direction and magnitude output from Metz et al. (2019b (LSTM), the fully connected optimizer from Metz et al. (2019)) (FF), as well as the proposed learned optimizer in this work ($3.3 (LSTM_ (F). As shown in Figure a). the proposed architecture achieves the lowest outer-training loss and achieves this the fewest outer-training steps. To the best of our knowledge, this is the first published comparison across different learned optimizer architectures, on the same suite tasks. Previous work only compared a proposed learned optimizer against hand-designed (baseline) optimizers.
Next, we explored how increasing the number of inner-tase used when training an optimizer affects final performance. To do this. randomly sampled subsets of tasks from the full task set, while evaluating performance on a common held-out set of tasks. Figure (b)) shows that increasing the number of tasks leads large improvements.
4

---

Under review as a conference paper at ICLR 2021
Global Adam LR Trial)
Global Adam8p, Trial)
Global NAdamW, RS Trial)
Por Task Adam LR (14 Trial)
OptLiet (10 Trial)
Per Task Opt. List (1000 Trial)
Per Task AdamBp (10 Trial)
Trial)
Por Task Adwm@pp@pp.......... List (2000 Trial)
-0.6 -44
-0.2
0.0
0.2
6.4
Leamned Optimizer Improve ment Over Baseline (higher better)
Figure 3: Learned optimizer performance as compared to battery of different haselines. For each baseline, we show box plot over 100 different tasks. We show both training (solid color) which the optimizer has seen, and test set (hatched light color) which is sampled from the same distribution but never used for outer-training. Values the right of the dashed line indicate the learned optimizer outperforming the corresponding baseline. The learned optimizer outperforms any single baseline optimizer with fixed hyperparameter configuration (marked global, with a single Trial). We also outperform per-task tuning of baselines, when tuning is done only over modest number of hyperparameter configurations (for instance, learning rate tuned Adam which uses 14 trials). Note that the learned optimizer undergoes no per-task tuning and only makes use of one trial.
4.2 COMPARISONS WITH HAND-DESIGNO OPTIMIZERS
We compare against three baseline optimizers: AdamLR, which is the Adam optimizer Kingma and Ba (0114) with tuned learning rate. Adam8p. which version of the Adam optimizer with eight tunable hyperparameters: learning rate, 31, 82, and 6, plus additional & and & regularization penalties, and learning rate schedule parameterized with a linear decay and exponential decay. See Appendix D for more details. Our final baseline optimizer, called opr_oisttt list, consists of the NAdam optimizer (Dozat) 2016) with "AdamW" style weight decay (Loshchilov and Hutter 2017). cosine learning rate schedules (Loshchilov and Hutter 2016) and learning rate warm up (See Metz et al. (2020) for more info). Instead of tuning these with some search procedure, however, we draw them from sorted list of hyperparameters provided by (Metz et al. 2020) for increased sample efficiency.
Evaluation of optimizers, let alone learned optimizers, difficult due to different tasks of interest, hyperparameter search strategies, and compute budgets (Choi al. 2019 Sivaprasad et al. 2019). We structure our comparison by exploring two scenarios for how a machine learning practitioner might go about tuning the hyperparameters of an optimizer. First, we consider "off-the-shell" limit, where a practitioner performs a small number of optimizer evaluations using off-the-shelf methods (for instance, tuning learning rate only). This is typically done during exploration of new model dataset. Second, we consider a "finely tuned" limit. where practitioner has large compute budget devoted tuning hyperparameters of a traditional optimizer for a particular problem of interest.
For the "off-the-shell" limit, we consider the scenario where practitioner has access to limited number of optimization runs (trials) (S 10) for a particular problem. Thus, we select first set of baseline hyperparameters using single, default value (denoted global in Fig 3 across all of the tasks. use random search (RS) using 1000 different hyperparameter values to find the global value that performs best on average for all tasks. Practitioners often tune models with small number of evaluations. As such, we include comparisons to per-task tuned learning rate tuned Adam, the first 10 entries of opt list, and 10 hyperparameter evaluations obtained from random search using the adam8p hyperparmeterization..
For the "finely tuned" limit, consider task-specific hyperparameter tuning, where the hyperparam- eters for each baseline optimizer are selected individually for each task (per-task in Fig
We plot histogram over tasks showing the difference in performance between the each baseline optimizer and the learned optimizer in each row Fig 3 First, we note that the distribution is broad, indicating that for some tasks the learned optimizer is much better, whereas for others, the baseline
5

---

Figure 4: Implicit regularization in the learned optimizer. Panels show optimization trajectories over 2D loss surface, fzz y) (x y)² that has continuum of solutions along the diagonal. Left: Trajectorics the Adam optimizer go to the nearest point on the diagonal. Right: Trajectories of the Under review as a conference paper at ICLR 2021
Adam Optimizer
Learned Optimizer
1.0
1.0
Funetion (loss) value
0.5
0.5
y-coordinato
y-coordinate
0.0
0.0
-0.6
-0.5
-1.0
-1.0
x-coordinate
x-coordinate
learned optimizer go towards the diagonal, but also decay towards single solution the origin, if there were an added regularization penalty on the magnitude each parameter.
optimizer(s) are better. On average, we see small but consistent performance improvements over baseline optimizers, especially in the "off-the-shelf" scenario. We attribute this to the diverse set of tasks used for training the learned optimizer.
4.3 UNDERSTANDING OPTIMIZER BEHAVIOR
To better understand the behavior of the learned optimizer, we performed probe experiments where we compared trajectories of the learned optimizer simple loss surfaces against baseline optimizers. The goal of these experiments was to generate insight into what the learned optimizer has learned.
In machine learning, many problems benefit from including some kind of regularization penalty, such as an & penalty on the weights or parameters of a model. We explored whether the learned optimizer (which was trained to minimize validation loss) had any implicit regularization, beyond what was specified as part of the loss. To test this, we ran optimizers on simple 2D loss surface, with a continuum of solutions along the diagonal: ((x,y) = 2 (x - y)² Although any point along the y diagonal is global minimum, we wanted to if the learned optimizer would prefer any particular solution within that set. Ssoursswswwws the resulting training trajectories along the 2D loss surface from many starting points. For baseline optimizer (left), the trajectories find the nearest point on the diagonal. However, we find that the learned optimizer has learned some implicit regularization, in that pushes the parameters towards solution with small norm: (0, 0).
4.4 GENERALIZATION ALONG DIFFERENT TASK AXES
Next, wondered whether the learned optimizer was capable of training machine learning models which differed across different architectural and training hyperparameters. To test this type of generalization, we trained fully connected neural networks on CIFAR-IO and MNIST, and swept three model hyperparameters: the number of hidden units per layer (network width), the batch size, and the number of training examples (dataset size, formed by subsampling the full dataset). For each sweep, we compare the leamed optimizer to baseline optimizer, Adam, over grid of eight different learning rates logarithmically spaced from 10" 10-5.5 10
The results of these experiments are in Fig. 5] As we vary the number of hidden units (left column) or batch size (middle column), the learned optimizer generalizes outside of the range of hidden units used during training (indicated by the shaded regions). In addition, the learned optimizer matches the performance the best learning rate tuned Adam optimizer. On CIFAR-10 10, as we move further away from the outer-training task distribution, the learned optimizer diverges. For dataset size (right column), we find that the learned optimizer is more sensitive to the amount of data present (performance drops off more quickly as the dataset size decreases). These experiments demonstrate the learned optimizer is capable of adapting to some aspects of target tasks which differ from its outer-training distribution, without additional tuning.

---

Under review as a conference paper at ICLR 2021
MNIST sverage test loes
2.5
Num hidden units
2.5
Betch size
2.5
Dalasel
2.0
2.0
2.0
1.5
1.5
1.5
10-
0.5
0.5
0.5
0.0
0.0
0.0
10
10
10P
10P
10P
10
10a
10P
10
10P
10
10P
10P
10
10
CIFAR-10 average test loss
2.4
2.4
2.4
2.2
2.2
22
2.0
2.0
20
1.8-
1.8
18
1.5
1.5
109
102
100
102
10P
104 10
leamed
outer tain mo
104
10+
10
10
Adam learning
Figure 5: We show outer-generalization of the learned optimizer in a controlled setting by varying hyperparameters of the task being trained. We vary parameters around two types of models. Top row: A two hidden layer, 32 unit fully connected network trained on MNIST with batch size 128. Bottom row: A two layer hidden layer, 64 unit fully connected network trained on CIFAR-10 with batch size 128. We vary the width of the underlying network, batch size, and dataset size in each column. Each point on each curve shows the mean test loss averaged over 10k inner steps. We show median performance over five random task initializations. Error bars denote one standard deviation away from the median. Color denotes different learning rates for Adam log spaced every half order of magnitude for Adam with purple representing 10-6 and yellow at 10-2 We find the learned optimizer is able to generalize outside of the outer eeter-training distribution (indicated by the shaded patch) in some cases.
6.0
7.0
7.5
1.5
Lawemed
6.5
Learredd
7:0
learreed
6.5
6.0
0.5
5.5
0.0
5.0
5.0
0
2000
4000
8000
8000
10000
2000
4000
6000
8000
10000
10
104 10
10" 10* 10a
10 10 100
inmer-trainig step
inner-training step
(c)
learning
Figure 6: Learned optimizers are able to optimize ResNet models. (a) Inner-learning curves using the learned optimizer to train a small ResNet on CIFAR-10. We compare to learning rate tuned Adam and Momentum. Solid lines denote test performance, dashed lines train. (b) Inner-learning curves using the learned optimizer to train a small ResNet on 64x64 resized ImageNet. We compare to learning rate tuned Adam and Momentum. (c) Performance averaged between 9.5k and 10k inner-iterations for Adam and Momentum as function of learning rate, on the same task as in (b). Despite requiring no tuning, the learned optimizer performs similarly to these baselines after tuning.
4.5 GENERALIZATION TO LARGE-SCALL PROBLEMS
We test the learned optimizer on large-scale machine learning tasks. We use two ResNet V2 (He al. 2016b) architectures: a 14 layer residual network trained on CIFAR- 10 (Krizhessky al. 2009); and a 35 layer residual network trained on 64x64 resized ImageNet (Russakovsky et al. 2015). We train both networks with the learned optimizer and compare the performance to learning rate tuned Adam and Momentum (Fig. 6). Details in Appendix DD For CIFAR- optimizer achieves similar performance as the baselines but does not overfit later in inner- training. For ImageNet, find that the learned optimizer performs slightly worse. We expect modifying the underlying outer-training task distribution to include more problems similar to this evaluation problem will improve performance.
find that the learned
Note that our baselines only include learning rate tuning. More specialized hyperparameter configu- rations, designed specifically for these tasks, such as learning rate schedules and data augmentation
7

---

Under review as a conference paper at ICLR 2021
0.40
Matris 5) Adumth-33) Logimod
0.35
traning
0.30
0.25
0.20
0.15
0.10
2000
4000
6000
5000
10000
Outer Irsining itenation
Figure 7: The learned optimizers can be used to train themselves about as efficiently as hand designed methods. On the x-axis show number of weight updates done to the learned optimizer. On the y-axis we show outer-loss. Each point consists of inner-training 100 models, each with five random initializations trained for 10k inner-iterations. We show the average validation performance post-normalization averaged over all tasks, seeds, and inner-training steps. Each line represents a different randomly initialized learned optimizer. In orange we show Adam with learning rate of 3 10 the same value used to train the optimizers in this work. In green, we show Adam with learning rate of 10-3, the learning rate that performed best in this 10k outer-iteration regime. In blue we show our learned optimizer.
strategies will perform better. An extensive study of learned optimizer performance wide range of state-of-the of-the--rt models is subject for future work.
4.6 LEARNED OPTIMIZERS TRAINING THEMSELVES
Finally, we performed an experiment to test if learned optimizer can be use to train new learned optimizers. F shows that this "self-optimized" training curve similar to the training curve using our hand-tuned training setup (using the Adam optimizer). We interpret this as evidence of unexpectedly effective generalization, as the training a learned optimizer is unlike anything in the set of training tasks used to train the optimizer. We show outer-training for 10k outer-training iterations, matching the number of inner-iterations used when outer- training the learned optimizer. If outer-training continued beyond 10k iterations, learned optimizer performance worsens (see Appendix G). suggesting that more inner-iterations are needed when outer-training
shows
DISCUSSION
In this work, train learned optimizer using larger and more diverse of training tasks, better optimization techniques, an improved optimizer architecture, and more compute than previous work. single of hyper-
The resulting optimizer outperforms hand-designed optimizers constrained parameters, and performs comparably to hand designed optimizers after a modest hyperparameter search. Further, demonstrates some ability to generalize to optimization tasks unlike those it was trained on. Most dramatically, it demonstrates an ability to train itself.
The learned optimizer we develop is not without limitations.
Generalization: While it performs well for tasks like those in TaskSet, we do not yet fully understand its outer-gert senerreeellliitiinn capabilities. shows promising ability to generalize to out of distribution problems such as training itself, and ResNet models, but it does not outperform competing algorithms in all settings on simple "out of distribution" optimization tasks, like those seen in Figure S
Optimization / Compute: Currently it takes significant compute expenditure to train a learned optimizer at this scale, resulting in a nontrivial carbon footprint. Ideally training can be run once and the resulting learned optimizer can be used broadly is the case with BERT (Devlin et al. 2018).
Learned optimizer architectures: We have shown there is considerable improvement to be had by incorporating better inductive biases both in outer-optimization speed and capacity. Improving these architectures, and leveraging additional features will hopefully lead to more performant learned optimizers. Additionally, the current optimizer, while compute efficiint, is memory inefficient requiring > 5x more storage per-par parameter than Adam. Modifications similar to those in Shazeer and Stern 2018) will be needed to train larger models.
8

---

Under review as a conference paper at ICLR 2021
Outer-training task distribution: We have shown that training larger tasksets leads to better generalization. have not studied which tasks should be included so as to aid outer suter-generalization.

---

Under review as a conference paper at ICLR 2021
REFERENCES
Alex Krizhevsky, Ilya Sutskever, and Geofrey E. Hinton. Imagenet classification with deep convolu- tional neural networks. In Advances in neural information processing systems, pages 1097-1105, 2012.
Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemystaw Dębiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Chris Hesse, et al. Dota 2 with large scale deep reinforcement learning. arXiv preprint arXiv:1912.06680 2019.
Oriol Vinyals, Igor Babuschkin, Junyoung Chung, Michael Mathieu, Max Jaderberg, Wojciech M Czarnecki. Andrew Dudzik, Aja Huang, Petko Georgiev, Richard Powell, et al. Alphastar: Mastering the real-time strategy game starcraft ii. DeepMind blog, page 2, 2019.
Chris Piech, Jonathan Bassen, Jonathan Huang, Surya Ganguli, Mehran Sahami, Leonidas Guibas, and Jaseha Sohl-Dickstein Deep knowledge tracing. In Advances in neural information processing systems, pages 005--13,222 2015.
Dami Choi, Christopher J Shallue, Zachary Nado, Jaehoon Lee, Chris J Maddison, and George E Dahl. On empirical comparisons of optimizers for deep learning. arXiv preprint arXiv: 10.00446, 2019.
Christian Daniel, Jonathan Taylor, and Sebastian Nowozin. Learning step size controllers for robust neural network training. In Thirtieth AAAI Conference on Artificial Intelligence, 2016.
Chang Xu, Tao Qin, Gang Wang, and Tie-Yan Liu. Reinforcement learning for learning rate control. arXiv preprint arXiv: 1705. 159, 2017.
Zhen Xu, Andrew M Dai, Jonas Kemp, and Luke Metz. Learning an adaptive learning rate schedule. 2019.
arXiv preprint arXiv:
Marein Andrychowiez, Misha Denil, Sergio Gomez, Matthew Hoffman, David Pfau, Tom Schaul, and Nando de Freitas. Learning to learn by gradient descent by gradient descent. In Advances in Neural Information Processing Systems, pages 3981-3989, 2016.
Olga Wichrowska, Niru Maheswaranathan, Matthew W Hoffman, Sergio Gomez Colmenarejo, Misha Denil, Nando de Freitas, and Jascha Sohl-Dickstein. Learned optimizers that scale and generalize. International Conference on Machine Learning, 2017.
Kaifeng Lv, Shunhua Jiang, and Jian Li. Learning gradient descent: Better generalization and longer horizons. arXiv preprint arXiv:1703.088
2017.
Luke Metz, Niru Maheswaranathan, Brian Cheung, and Jascha Sohl-Dickstein. Learning unsupervised learning rules. arXiv preprint arXiv: 00222, 2018
Luke Metz, Niru Maheswaranathan, Jonathon Shlens, Jascha Sohl-Dickstein, and Ekin Cubuk. Using learned optimizers to make models robust input noise. arXiv preprint arXiv:1906.000 2019a.
Luke Metz, Niru Maheswaranathan, Jeremy Nixon, Daniel Freeman, and Jascha Sohl-Dickstein. Understanding and correcting pathologies in the training of learned optimizers. In International Conference on Machine Learning, pages 4556-4565, 2019b.
Keren Gu, Sam Greydanus, Luke Metz, Niru Maheswaranathan, and Jascha Sohl-Dickstein. Meta- learning biologically plausible supervepe update rules. bioRxiv, 2019.
Luke Metz, Niru Maheswaranathan, Ruoxi Sun, C Daniel Freeman, Ben Poole, and Jascha Sohl- Dickstein. Using a thousand optimization tasks to learn hyperparameter search strategies. arXiv preprint arXiv:2002.11887, 2020.
M Levin T Hart and Mike Levin. Ai memo 39-the new compiler. Technical report, Technical report, MIT, 1962.
Luca Franceschi, Paolo Frasconi, Saverio Salzo, and Massimilano Pontil. Bilevel programming for hyperparameter optimization and meta-learning. arXiv preprint arXiv:/8000 049/0, 2018.
10

---

Under review as a conference paper at ICLR 2021
Yuhuai Wu, Mengye Ren, Renjie Liao, and Roger B Grosse. Understanding short-horizon bias in stochastic meta-optimization. pages 478 778-487,2 2016.
Ingo Rechenberg. Exoluuronsstaatgii-oppiiiiiinene technisher systeme nach prinzipien der biologis- chen evolution. 1973.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 201 6a.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pages 5998-6008, 2017.
Barret Zoph, Vijay Vasudevan, Jonathon Shlens, and Quoe V Le. Learning transferable architectures for scalable image recognition. Proceedings of the IEEE conference on computer vision and pattern recognition, 2018.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural compulation, 9(8): 1735-1780, 1997.
Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv: 2014.
Yann LeCun. The mnist database of handwritten digits. http://yann. lecun. com/exdb/nmist// 1998.
George Papamakarios, Theo Pavlakou, and lain Murray. Masked autoregressive flow for density estimation. In Advances in Neural Information Processing Systems, pages 2338-2347, 2017.
Diederik Kingma and Max Welling. arXiv:1312.6114 2013.
Auto-encoding variational bayes. arXiv preprint
Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5-rmsprop: Divide the gradient by running average of its recent magnitude. COURSERA: Neural networks for machine learning, 4(2):26 31, 2012.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6998 6980, 2014.
Timothy Dozat. Incorporating nesterov momentum into adam. 2016.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05111, 2017.
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv: 00983, 2016.
Prabhu Teja Sivaprasad, Florian Mai, Thijs Vogels, Martin Jaggi, and François Fleuret. On the tunability of optimizers in deep learning. arXiv preprint arXiv: 19/0. 1758, 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In European conference on computer vision, pages 630-645. Springer, 2016b.
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar-10 and cifar- 100 datasets. URl: http:://ww.. cs. toronto. edulkrizJcifar html, 6, 2009.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV), 115
13):211-252, 2015. doi: 10.1007/s| 1233-015-08661-
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: /0/0. 04.054 2018.
11

---

Under review as a conference paper at ICLR 2021
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. arXiv preprint
2018.
Irwan Bello, Barret Zoph, Vijay Vasudevan, and Quoe Le. Neural optimizer search with reinforcement ppf
learning. 2017. URL
Samy Bengio, Yoshua Bengio, Jocelyn Cloutier, and Jan Geesei. On the optimization of a synaptic learning rule. In Preprints Conf. Optimality in Artificial and Biological Neural Networks, pages 6-8. Univ. of Texas, 1992.
Thomas Philip Runarsson and Magnus Thor Jonsson. Evolution and design of distributed learning rules. In Combinations of Evolutionary Computation and Neural Networks, 2000 IEEE Symposium on, pages 59-63, IEEE, 2000.
Yoshua Bengio. Nicolas Boulanger-Lewanoosske and Razvan Pascanu. Advances in optimizing recurrent networks. In 2013 IEEE International Conference on Acoustics, Speech and Signal Processing, pages 8624-8628. IEFE, 2013.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv preprint arXiv:1410.5401, 2014.
Jorge Nocedal. Updating quasi-newton matrices with limited storage. Mathematics of computation. 35(151):773-78,, 1980.
Dong C Liu and Jorge Nocedal. the limited memory bfgs method for large scale optimization. Mathematical programmingg 45(1-3):503-528 1989.
Nicolas Heess, Srinivasan Sriram, Jay Lemmon, Josh Merel, Greg Wayne, Yuval Tassa, Tom Erez, Ziyu Wang, SM Eslami, Martin Riedmiller, et al. Emergence of locomotion behaviours in rich environments. arXiv preprint arXXv:1707.002288 2017.
Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, and John Schulman. Quantifying generalization in reinforcement learning. arXiv preprint
2018.
Bobak Shahriari, Kevin Swersky, Ziyu Wang, Ryan Adams, and Nando De Freitas. Taking the human out of the loop: A review of bayesian optimization. Proceedings of the IEEE, 104(1):
148-175, 2015.
Daniel Golovin, Benjamin Solnik, Subhodeep Moitra, Greg Kochanski, John Karro, and D Sculley. Google vizier: A service for black-box optimization. In International Conference on Knowledge Discovery and Data Mining, 2017.
Ke Li and Jitendra Malik. Learning to optimize. International Conference on Learning Representa- tions, 2017a
Ke Li and Jitendra Malik. Learning to optimize neural nets. arXiv preprint 1rXiv:1703.0044), 201 7b.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.06377 2017.
Jan Peters, Katharina Mulling, and Yasemin Altun. Relative entropy policy search. In Twenty-Fourth AAAI Conference on Artificial Intelligence, 2010.
Yoshua Bengio. Gradient-based optimization of hyperparameters. Neural computation, 12(8): 1889-1900, 2000.
Atilim Gunes Baydin and Barak A Pearlmutter. Automatic differentiation of algorithms for machine learning. arXiv preprint 1rXiv:1404.74556 2014.
Dougal Maclaurin, David Duvenaud, and Ryan Adams. Gradient-based hyperparameter optimization through reversible learning. In International Conference on Machine Learning, pages 2113-2122 2015.
Paul Werbos. Backpropagation through time: what it does and how to do it. Proceedings of the IEFF, 18010):150-1156) 1990.
12

---

Under review as a conference paper at ICLR 2021
Corentin Tallec and Yann Ollivier. Unbiasing truncated backpropagation through time. arXiv preprint arXiv:170S.08209, 2017.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. arXiv preprint arXiv:1703.03400, 2017.
Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-learning algorithms. arXiv preprint (rXiv:1803.0299, 2018.
Barak Pearlmutter. An investigation of the gradient descent process in neural networks. PhD thesis, Carnegie Mellon University Pittsburgh, PA, 1996.
Paavo Parmas, Carl Edward Rasmussen, Jan Peters, and Kenji Doya. Pipps: Flexible model-based policy search robust the curse of chaos. In International Conference on Machine Learning, pages 4062-4071, 2018.
Joe Staines and David Barber. Variational optimization. arXiv preprint arXiv: 1rXiv:122.44007,000 2012.
LA Rastrigin. About convergence of random search method in extremal control of multi-parameter systems, Avtomat, Telemekh, 24(11):1467-1473, 1963.
Yurii Nesterov and Vladimir Spokoiny. Random gradient-free minimization convex functions. Technical report, Université catholique de Louvain, Center for Operations Research and Economet- rics (CORE), 2011.
Krzysztof Choromanski. Mark Rowland, Vikas Sindhwani, Richard E Turner, and Adrian Weller. Structured evolution with compact architectures for scalable policy optimization. arXiv preprint arXiv:1804.02339, 2018.
Tim Salimans, Jonathan Ho, Xi Chen, Szymon Sidor, and Ilya Sutskever, Evolution strategies scalable alternative to reinforcement learning. arXiv preprint arXiv: 2017.
Niru Maheswaranathan, Luke Metz, George Tucker, Dami Choi, and Jascha Sohl-Dickstein. Guided evolutionary strategies: Augmenting random search with surrogate gradients. In International Conference on Machine Learming, pages 4264-4273 PMLR, 2019.
Oriol Vinyals, Charles Blundell, Timothy Lillicrap, Daan Wierstra, et al. Matching networks for one shot learning. Advances in neural information processing systems, pages 3630-3638, 2016.
Rein Houthooft, Richard Chen, Phillip Isola, Bradly C Stadie, Filip Wolski, Jonathan Ho, and Pieter Abbeel. Evolved policy gradients. arXiv preprint 1802. 04821, 2018.
Louis Kirsch, Sjoerd van Steenkiste, and Jürgen Schmidhuber. Improving generalization in meta reinforcement learning using learned objectives. In International Conference on Learning Repre- sentations, 2018.
Junhyuk Oh, Matteo Hessel, Wojciech Czarnecki, Zhongwen Xu, Hado van Hasselt, Satinder Singh, and David Silver. Discovering reinforcement learning algorithms. arXiv preprint arXiv:2007.08799 2020.
Paul Vicol and et. al. Persisted evolutionary strategies. In Preparation, 2020.
Sanjay Ghemawat, Howard Gobioff, and Shun-Tak Leung. The google file system. In Proceedings of the nineteenth ACM symposium on Operating systems principles, pages 29-43, 2003.
Martín Abadi, Paul Barham, Jiammin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. Tensorflow: A system for large-scale machine learning. In OSDI, volume 16, pages 265-283, 2016.
13

---

Under review as a conference paper at ICLR 2021
A EXTENDED RELATED WORK
We categorize progress in learned optimizers into three categories: parameterizations, data distribu- tions, and outer-training methodology. We include a review of these previous efforts herein. Finally we include discussion on other, non-optimizer, learned learning algorithms.
A.1 PARAMETERIZATIONS
Broadly, authors attack the problem of learned optimization at high level using different parameteri- zations.
A.1.1 CONTROLLER BASED PARAMETERIZATION
One class of learned optimizer parameterizes some function, usually as a neural network, that returns hyperparameters for an existing hand designed bethod.awwk / not precise These parameterizations impose strong inductive bias towards the hand designed optimizers they parameterize. All inner- learning must be done via manipulation existing methods. By enforcing a constrained structure, these methods limit the types of learning rules expressible. This limitation does inject strong priors into the model making both outer optimization easier, as well as produce better outer- seeerreeeraaizttion These methods often additionally make of derived features from the currently training model. These features include things like loss values, variance of function evaluations, and variance of weight matrices.
There are a number of existing works that explore different types of features and architectures. Daniel al. (016) explores using simple, linear policy which maps from hand designed features to log learning rate which is then used with SGD, RMSProp or momentum. Xu et al. (2017) learns a small LSTM network with inputs of current training loss, and predicts a learning rate for SGD. Xu al.)
(2019) uses an LSTM an MLP with features computed during training and outputs a scalar which multiplied by the previous learning rate.
A.1.2 SYMBOLIC PARAMETERIZATTINS
Existing learning rules are often expressed symbolically by a small number of mathematical ex- pressions operations. All first order optimizers used today symbolic and expressible by small collection of mathematical operations. Bello et al. (017)) take inspiration from this and also pa- rameterize optimizers as collections of mathematical primitives. benifit of this learned optimizer parameterization is that they are more interpretable. For example, [Bello et al. (2017) show common subpatterns discovered based on multiplication of the sign of the gradient and sign of the momentum value. Not all computations lend themselves to a symbolic regression. Finding symbolic formula to solve increasingly more difficult tasks starts quickly turns into a combinatorical search and thus difficult to perform (analogous to the difficulties encountered in program synthesis).
A.1.3 MORE GENERAL PARAMETERIZATION
The last family of parameterization is based directly predicting the update step via more general function approximators. Older work often makes use of simple linear combinations of features while more recent work makes use of neural networks. There are few classes of update rule parameterizations.
Bengio al.) (1992) proposes a parameter, and 16 parameter update rule parameterized by mixing various biologically inspired learning signals. Runarsson and Jonsson (2000) use a simple continuous parameterization that operates on error signals and produces changes in inner-parameters.
Recently, there has been renewed interest in learned optimizers, in particular using neural network paramerizations of the learning rules. [Andrychowicz et al. (2016) makes use an LSTM(Hochreiter and Schmidhuber) 1997) possibly with global average pooling enable concepts like lg gradient clipping to be implemented (Bengio et al. 2013). They also explore Neural Turing Machine (Graves) et al.) 2014) like parameterizations specifically designed for low rank memory updates that can learn algorithms like LBFGS (Nocedal 1980; Liu and Nocedal 1989). In contrast, Metz et al. (2019h) makes of per-parameter ML.P instead of the LSTM.
14

---

Under review as a conference paper at ICLR 2021
One critical design choice when designing neural network parameterized learned optimizers is to determine what imput features should be used. awk When training neural network models it is critical that the inputs be similar scales for optimization be successful, Andrychowicz et al. (2016) trains on raw gradients which are scaled by either decompasing each scalar into a sign and a magnitude before feeding into LSTM. [Lv al. (2017) uses the gradient and momentum normalized by the rolling average of gradients squared (similar to the rescaling done by Adam), Wiichrosska et al. (2017) also makes use of momentum except using multiple timescales normalized in a similar way. addition to these momentum terms, Metz et al. (2019b) other features such as weight value and additionally applies a normalization based on the second moment of features.
In Andrychowicz et al. (2016); Ly et al. (2017); Metz et al. (2018) all methods perform updates that are computed independently for each parameters While expressive, this expensive as no computation can be shared. Wichrowska et al. (2017) improve upon this by additionally having per-layer, and global LSTM.
When designing learned optimizer architectures, there are number of design decisions to keep in mind. One must balance compute cost the learned optimizer with expressibility. Often this shifts models to be considerably smaller (e.g. <64 hidden units) than most applications of deep learning. Wichrowska et al. (2017). for example makes use an 8-hidden unit LSTM per perammetee.r. Selection of features to feed into the learned optimizer is also critical. While traditionally deep learning involves learning every feature, this learning comes at the cost of increased compute which often not feasible. This paragraph needs some work. First sentence, and "selection of features feed into the learned..."
A.2 DATA DISTRIBUTION OF TASKS
There agreed upon standards when defining datasets for training learned optimizers. Instead the community trains on what ever datasets are available or what ever best approximates a desired goal (training particular model, creating more general optimizers). Constructing large distributions of tasks is labor intensive and thus not often done.
Metz et al. 2019b) draws inspiration from the few shot learning literature and constructs 10 way classification problems sampling from different classes on imagenet.
Wichrowska et al. (017) leverages large distribution of synthetic tasks. These tasks are designed to represent different types of loss surfaces that might be found in loss surfaces.
TaskSet, (Metz al. 2020 dataset of tasks specifically designed for learned optimizer research. We use this dataset throughout our work.
There is balance between performance of the tasks, and ability outer-train. Selecting the types of problems we want to train on is often not enough. Additionally, outer-training on the closest task possible will not produce the best optimizer nor even converge. Training on distributions with increased variation smooths the outer-loss surface and makes exploration simpler. This mirrors phenomena found RL (Heess al. 2017 Cobbe al. 2018
A.3 OUTER- OOPTIMIZATIO METHODS
The outer-optimization problem consists of finding a particular set of weights, or configuration of learned optimizer. A number of different strategies have been proposed.
A.3.1 HYPER PARAMETER OPTIMIZATION
One of the most common family of outer-learning methods used taken from the family of hyper- parameter optimization methods. In the context of learning optimizers, this can be seen as finding optimizer hyper parameters, e.g. learning rate, over particular task instance. Numerous methods exist do this ranging from Bayesian hyper parameter optimization (Shahriari et al. (015). to grid search, to genetic algorithms. Golovin et al. (2017) for more complete description.
Metz. al. 2019b) uses normalization across parameters this not strictly
15

---

Under review as a conference paper at ICLR 2021
These methods do not work in all settings. The types of outer-learning problems encountered in learned optimizers is different in that the measure of performance (outer-objective) is often an expectation over some distribution of tasks and thus not easily to evaluate. Additionally, the amount of outer-parameten often larger than simply finding a few hyperparameters. Despite this, we include this here to show the similarity learned optimizer research.
A.3.2 REINFORCEMENT LEARNING
Learned optimizers can naturally be cast into a sequential decision process. The state the system consists of the inner-parameter values, action space is the steps taken, and the reward is achieving low loss in the future. A number of works have thus taken this viewpoint. Li and Malik (2017a makes of the guided policy search algorithm. Xu et al. (2019) uses PPO PPoSuchulman al. (2017). Daniel et al. (2016) make use of Relative Entropy Policy Search (Peters et al. 2010). The exact algorithm used function of the underlying parameterization.
A.3.3 NEURAL ARCHITECTURE SEARCH
Instead of learning the policy directly, Bello et al. (017) makes use reinforcement learning (PPO (Schulman et al. 2017) to learn controller which produces the symbolic learned optimizer. This is distinct from 3.2222 as does not leverage the sequential nature of the inner problem. Instead, treats the environment a bandit problem.
A.3.4 BACKPROPOGATION / GRADIENT BASED
Computing derivatives through leaning procedures has been explored in the context of hyperparameter optimization in (Bengio) 2000 Baydin and Pearlmutter] 2014 Maclaurin et al. (015)) Andrychowicz al. (016) was the first make use of gradient based learning for this application.
To train a learned optimizer, ideally, one would compute the derivative of final performance with respect to the optimizer parameters. This is often referred to as unrolling the entire training procedure into one large graph then running reverse mode automatic differentiation. Not only is this often too expensive to do in practice, the resulting outer-loss surface can be poorly conditioned (Metz et al. 2019b). In practice, approximations often made.
One common family of approximations truncated backpropogation through time. The core idea is to break apart the long, sequential unrolled computations into shorter sequences and thus propagating error back through the entire sequence [Werbos] 1990 Tallec and Ollivier) (2017). This approximation used widely in neural network parameterized learned optimizers (Andrychowiez et al.) 2016 Wichrowska et al. 2017 Lv et al.) 2017]] Metz et al. 2019b). Unlike in language modeling, truncated backprop has been shown to lead dramatically worse solutions for meta-learning applications (Wu
et al. 2016 Metz et al., 2019b).
second family of approximations involve not computing any second order or higher terms. This is often refereed to as affirst order" method(Finn et al. Nichol Nichol et al., (018). Andrychowicz (2016) uses only these first order gradients where as subsequent work, (Wichrowska al. 2017 computes the full gradient. The trade offs between these two optimization methods has been discussed in the few shot learning literature in context of MAML Reptile (Finn et al. 2017
Nichol et al. 2018).
Computing gradients through iterative, non-linear, dynamics has been shown to cause chaotic dynamics. Pearlmutter (1996); Parmas et al. (2018); Maclaurin et al. (2015) showed high sensitivity to learning rate with respect to performance after multiple steps of unrolled optimization. [Mez et al. (2019b) shows this issue for neural network parameterized learned optimizers and proposes an alternative objective based on variational optimization Staines and Barber 2012) and multiple gradient estimators similar to Parmas et al. (2018).
Despite improvements, gradient based training learned optimizers is still difficult, and requires great
care.
16

---

Under review as a conference paper at ICLR 2021
A.3.5
EVOLUTIONARY STRATEGIES
alternative way estimate gradients is with black box method such as Evolutionary Strategies(Rass Rechenberg, 1973 Nesterov and Spokoiny [001T; Choromanski et al. [018 Salimans ct al. 2017). These methods are memory efficient, requiring no storage of intermediate states, but can produce high variance estimates. In the case of learned optimizer optimization, however, for certain problems these methods can result in lower variance gradient estimates (Metz et al. 20022000bb.. Hybrid approaches that leverage both gradients and ES have been such as Guided ES (Maheswaranathan et al,, 2019) have also been proposed for meta-optimization. This work leverages one of the simplest forms of evolutionary strategies described in (Salimans et al. 2017) which uses fixed standard deviation normal distribution for the population.
irigin 1963
A.4 OTHER LEARNED LEARNING ALGORIHHMS
There exists a wide variety of other meta-learning / learned learning algorithms. One domain that has received a lot of attention is few shot learning (Vinyals et al. (0066). The methods that perform well usually tend to learn more transferable features, than learning algorithms.
second domain where learned learning algorithms has been explored extensively reinforcement learning. We will not survey all methods, but choose three that most related in the types of challenges they address. Evolved Policy Gradients(Houthoof et al. 2018), explores learning a reinforcement learning algorithm by first learning loss function then optimizing policy with respect to this loss. With this parameterization they are able to show generalization to solve novel, and out of distribution tasks. MetaGenRL (Kirsch et al. 2018) is another method focuses generalizing even further out of distribution enabling transfer across environments with different action and observation spaces. Finally Learned Policy Gradient (Oh et al. 2020) adds learned components to actor critic methods to learn both how the value function is updated, and the policy is updated. They also present transfer results, outer-training on toy environments and transferring Atari.
B OUTER OPTIMIZATION DETAILS
In this work, as with Metz et al. 2019b), we use asynchronous, batched training. Each task has different complexity, thus will produce outer-gradient estimates a different rate. We use asyncronous minibatched training as synchronous training with these heterogenious workloads would be too slow and wasteful. We tie the outer batch size to the number of workers. To prevent stale gradients, we additionally throw away all outer-gradicn estimates that are from more than five outer-iterations away from the current weights.
We optimize all models with Adam and sweep learning rates between 3e-5 and 3e-3 for all experiments. The optimal learning rate changes depending on how long outer- srater-trainng occurs. We have done a preliminary exploration into learning rate schedules but have not yet been able to improve this constant schedule. For all outer-training experiments, we always run more than one random optimizer initialization. Due the relatively small number of units, and biased gradient estimators, performance is dependant on random seed. For all experiments we use gradient clipping of 0. applied to each weight independently. Without this clipping no training occurs. This surprises us as our gradient estimator is evolutionary strategies which will not typically have exploding gradients. Upon further investigation, however, the outer-gradien variance is much larger without this clipping.
When computing outer- graddegsa,,sts,, we follow Metz et al. 2019bb and compute am outer-loss based on multiple mini batches of data. our work we use five. Inner-training always uses a single minibatch of inner-training data as well as single batch of inner-validation data when used.
We first train with 240-360 length unrolls over a max of 10k inner-steps. While training we compute and record traer-ggaining 10k steps for 100 tasks sampled from the outer-training distribution and save outer-parameters every hour. While training we monitor performance across all outer-learning rates and all seeds the outer-training distribution. When training plateaus, we manually look through these evaluations and select a candidate set of optimizers to further train with an increasing truncation schedule. Not all optimizers fine tune in the same way despite having the same performance on the outer-training data so selecting more than one is critical. At this point we are unsure where this phenomenon comes from.
17

---

Under review as a conference paper at ICLR 2021
Next fine tune theses models with gradient estimator that does not have short-horizon bias. We explored two methods. The first method is based on an increasing truncation schedule. We tested linearly increasing truncation length from 300- 10k inner-steps outer-steps. We find the faster increase, 10k steps, performs best. Second, we explored fine tuning with Persisted Evolutionary Strategies unbiased gradient estimator (Vicol and al. 2020). We found this achieved similar final performance but achieved it in half the time. When fine tuning we also make use of different learning rates. We find higher learning rates make progress faster, but can be unstable in that the performance varies as function of outer-training step. Additionally, the learning rate chosen depends on the learning rate used previously in the first training phase. Before finetuning, we "warm up' the Adam internal rolling statistics. While this might not be strictly required, ensures that there is no decrease performance early outer-training. This can done by simply setting the outer-leamiig rate to zero for the first 300 outer-iterations.
the course of 30k 10k
C LEARNED OPTIMIZER ARCHITECTURE DETAILS
In this section we describe the detailed learned optimizer architecture. For ease of understanding we opt to show mix of pseudo-code based on python and textual descriptions as opposed to mathematical expressions. Finally we chose to describe optimizer series of stateless and pure functions for clarity.
We used this architecture for all of our experiments except of Fig 2b which used an older version of architecture with additional features which where dropped from the final version.
C.I HIGH LEVEL STRUCTURE OF THE OPTIMIZER
The learned optimizer has two main components: function that maps from some set of inputs, state, and parameters to some new state and new parameters.
class Opt iimmmiii mizer: def next_state sinputs: Inputs, state: State,
parameters: Parans
-> (State, Parans);
And function produce an initial state from the given amner-arrametes.
def initia aiitial_ssatee state (self, f, params:
Params)
State;
Parameters stored dictionary of different tensors keyed by name.
Params Dict [Text,
Tensor]
For convenience, we also define gradients to be the same type as Params:
Grads Params Dict [Text, Tensor]
The state consists of multiple values that we will discuss in detail further, For now, however, we list the full state with high level comments.
State namedtuple( "State",
current inner training iteration saraining_step" : Int
# the statistics for the rolling averages.
"rolling_features" RellingFettureState, Hidden state of the 1stm bestm__idddn_sstteee # Activations passed from the MLP to the LSTM.
LSTMiiddenSttte,
"from_mlp": list [Tensor [shape (from_mip_siiee ]], Activations from 1.TM passed to 1.MM
18

---

Under review as a conference paper at ICLR 2021
"from 1stm": Tensori shape= (from, 1stm size,) Rolling statistics of the train loss value "train__oss__ccom LossAccunState,
# Roiling statistics of the valid loss value "valid_loss_acnn"
LossAccunState,
State to manage inner-gradinet clipping. "dynamic_clip": Rollingclipstate, # Parameter values from the nearest 100 steps in the past.
1)
Both from_ larh_sstm_sss_ size, hro_mlp_siie hyper parameters set as part of the learned optimizer control how much information is sent from the MLP or from the LSTM.
The input the next_state function consists of inner-gradients, computed inner-training batch of data, as well as optionally validation data which is passed in every 10 iterations. We choose to not pass in validation data every steps for computational efficiency.
Inputs namedtuple "Inputs",
gradient from task. This is same type as the parameters. "inner grads": =inner_grads"; Grads, training loss computed from mini batch. "train_loss"; float, validation loss from a mind batch of validation data. "valid_loss"; Optional [float],
1)
Each task specifies function that samples parameter initialization, as well as to produce outputs.
class Task:
def initial_params (self) Params
def inputs_for_paraaee (self, params)
Grads, inner-training gradients float, inner-training loss Optional [float]) # optional inner-training validation
For outer-training each task also includes second loss function which computes the task's loss on the outer-validation split of data. Note this uses different validation set of data than the previous function.
def val valid_outer_loss (self.
params:
Params)
float
Inner training / application of the learned optimizer looks like:
task SampleTask() ()
optimizer Optimizer() ()
params task initialParams ()
state wtwiimiier.iitial_ssse.. (state)
for i in range (N) Length of inner training
saliidddta=ld inputs GetInputFromTaste (params, rith_valid data= (i 10 == 0))
params, state optimizer next__sttteeeetee (inputs, params, state)
Computing the outer objective and outer-gradients from inner-initialization looks like the following:
task SampleTask()
optimizer Optimizer() ()
When outer-training with ES, perturb optimizer weights.
params task www.iiiaaaaaaams.........
state optimizer. initial _stae (state)
19

---

Under review as a conference paper at ICLR 2021
outer_valid_lesses
for in range (N) Length of a trunction
inputs Get MetInttFromTass (params, with_valid_data $ 10 -- 0))
params, state opt (ixt_ssa params, state) for in zange outer_valide wtw.ooossssseeeeddddo. (task. wwwii loss (params))
outer_ mean (outerr wet_r__aaid____sseese_e
When outer-training with ES, must perturb the optimizer's parameters before computing outer_ioss.
C.2 UTILITIES/ / COMPONENTS
First we will describe the individual components and utilities used, then we will go on to describe the full update function.
C.2.1 ROLLING FEATURES
These are moving average of gradients and second moments computed similarly to Adam RMSProp. The rolling state consists of tensors containing momentum values (ms) and second moment values (rms):
RoLLingFetturesState collections.
colLections..aneetuullee. www.dtuneeeeeeeeeeem....e ("Ro11 Lngguuuttttttetettatttt
("ms": Dict [Text, Tensor], "rms": Dict [Text, Tensor]])
The Text represent names that map to the tensor of the corresponding shape from the parameters. The values are the same shape the corresponding inner- parameterre with an additional axis appended to keep track of multiple different decay values. While its possible to outer-learn these values, we fix them at 0.5, 0.9, 0.99, 0.999, 0.9999.
To update these we construct helper class.
Class RollingState:
def iiiiit (self, self welf.dee se.Ffddecyyvluueeee decay_values
We define an initial value which is simply simply zeros:
def initial ialtittllttteeeeeeeee (self, params: Params) RellIngFeatureState d_dimm 1en (decay_values)
ms (k: tf. zeros shape + [p_dims]]
rms - [k: tf. zeros ([v.shape + [n_dims]]
return Roll inallinggaateeeettaae (ms-ms, rms-rms)
To update these values we follow procedure similar to RMSProp update equations:
def update state (self, state: RoLLingPeatureStatee, grads: Grads) >>
RolLingfeatureState:
ret_ms (}
ret_rms (eys for k state
g grads [k]
s state
ret ret_ms [k] np. np. ti. [k])
ret_rms [k] np. 1p.
[k]
for di, decay in enueerate saraaayvvauess)
ret_mms [k] [..., di] - ret [k] ...., dil decay (1 decay
) grad
20

---

Under review as a conference paper at ICLR 2021
di ret [k] ..., di] decay
decay) grad+*2
return RollingFeaturestate (ms-ret__ss rms-ret_rnss)
C.2.2 LOSSACCUM
This represents how get loss information into our learned optimizer. Loss values have pre- determined scale and span many orders of magnitude. As such, we must somehow standardize them that the inputs bounded and able to be easily used by neural networks. We get around this by keeping track of the normalized mean and variance of the mini-batch losses.
lossAcccutttte
col illcttinns upw.nnnnnam ("oosssptttlllssssccussssttttt loossAcuustatee"
("mean": float, # rollinn ing mean loss value
"var": float, rolling second moment loss value "updates": int # Number updates performed so far 1)
The class that manages these states is parameterized by the decay of the rolling window.
class LossAccum:
initt (self, decay): self.decay decay
The initial state is simply zeros.
def Enitial_state (se11) LossAccunState:
return LossAccumStatee
mean-tf. constant (0., Stype-tf.floatie
car-tf.conttan (0., (yre-tt..foott2).
updates-tff.onnstam (0, dtype-tff.nt6.l)
To compute updates we do rolling mean and variance computations.
def next_state (self, state: RellingAccumstate, loss: float): ->
LossAccumState
new_mean self.decay state.men mean + (1.0 self.decay)
new_var self.decay state.var var
1.0 delay) tquare (new_mean - loss)
new_updates = state. state. updates +
return licssAccumState meean-new_ment
nar-nex_var, updates-nee__pdatee
Two additional functions used to normalize loss values for use in neural networks. First, we have a "corrected" mean (similar to what is done by Adam) for a given AccumState.
def corrected mean (self, state: AccumState) -> float: = 1. (1 self. Sal...dcay=*tt.ttt float (state updates) 1e-8) return state, mean
Second, we have function that weights one loss by different AccumState. This is eventually used weight the validation loss accum against the training loss state allowing the learned optimizer to detect overfitting.
def we weight llosttlosss___lloool (self, state: AccumState, loss: float) -> float: 1. (1 self. cor_mean state.mean
(state. updates) 1e-8)
cor_var state.var +
= = (less cor cor_mean) tf.rsqrt (cor_ (corsqt 1e-8)
return Ef.ccip_by_vvauee (1, -5, 5)
21

---

Under review as a conference paper at ICLR 2021
C.2.3 ROLLINGCLIPSTATE
Gradient clipping a often used technique in deep learning. We found large benefits by applying some form of learned clipping (clipping inner-gradient values) in our learned optimizer. We cannot simply select a default value because gradient norms vary across problem. As such we meta-learn pieces of dynamic gradient clipping algorithm. This is our first iteration of this concept and expect large gains can be obtained with better scheme.
This algorithm is stateful thus also needs state container.
RollingCl ipstate
(int, the number times this has been updated, float rolling average of mean of squared gradient values.
This class is parameterized by the decay constant of the rolling average and the multiplier to determine when clipping should start. Both of these values outer-learned with the rest of the learned optimizer.
class Roll RollingGradClip.
def init (self, alpha=0.99, clip mult=10) self self.alpha-alpha. self.cllp.mult-ciip.mmuu
The initial states are initialized to
initial_state
initial statstttteee.eee (self) -> RollingClipstate:
return (tf. constant (1, dtype=tf.float32),
float32),
constantt (1.0,
(1-self..
We provide a normalization function that both updates the RollingClipState and provides clipped gradients.
def nex__stacc_anLLoomaallle (sclf,
state: RollingClipstateo
grads: Grads) ->
RollingclipStaate, Grads:
def _ (state: Bollingclipstatte grads: Grads) summary_ops = []
t, snd state
clip_amount amount = (snd / (1-self.alphav (ewltt)s)s............... cllpp mult
clipgs ltf.elip by value (g, -clip amount, clip_
amount) for g in grads, values
return clipgs
end state
clipped_graaee self. self._normalize (state, grads)
mean_square list
(g)) for in
seipped_gradse
newennd_moment tg. (1c-8 tf reduce mean (mean_squarr__iist) next_snd and self. + sew_snn__mment (1. self alpha))pha) return (t+1, next_snd), clipped_grads
C.3 LEARNED OPTIMIZER SPECIFIC PUTTING IT ALL TOGETHER
The learned optimizer has two main components. The first is the Optimizer class that manages everything surrounding inner-learning. This function has no learnable outer- parameters. The second is what we call "theta mod" which contains all of the outer-variables and functions that learning.
First construct the optimizer class with the corresponding theta_moo (see bellow).
22

---

Under review as a conference paper at ICLR 2021
class Optimizer:
def init (self, theta_hoaal self.theta_med theta_mod
Next the rolling features at fixed, hard coded intervals. These could be outer-learned but in this work they are fixed.
self decays=10.5, 0.9, 0.99, 0.999, 0. 9999]]
RollingFeatures
Then the the loss features for use with both training and validation loss. We use a lower decay constant the validation it is updated less frequently (once every 10 steps). Once again these could be outer-learned but in this work we leave them fixed.
self. self.train loss self.valie loss_accun accum LossAccur (0.
LossAccur (0.95 95)
Finally we construct the gradient clipping utility. Unlike the previous parameters, we do outer-learn these and pass in two variables off of the theta_mod!
self. relffrolling.hle clip
Roll RollingGradClip(
self. theta mod. rolling clip alpha, self. rolling awtp__lppplplleeeeeeeeee mult)
Next, we need an initial state for the optimizer defined in C.1 This mostly consists of obtaining initial states from the various components of the learned optimizer.
initialState (self, params : Parans) State:
shapes Iv. shape. ass.phaeaa_lstttttesees_eessesse list () for v in params.valuee
return State(
training_step-0,
(shapes),
lstm tetpppa..e..cccmaaaa....a..aaaaa
(len (shapes)
Next we look the update performed. We split up into two components first validation state, and
second a training state.
def noxt_paraam__nd__stttellss.
params: Params, current_state: State, train_less: float, train_grads: Grads, valid_less: Optional (float]) -> (Params,
State)
if valid_loss:
current_ state
vossd_loss)
params, next_state self.noxt_state_treinn (parans, current_state train_loss, train_grads)
return params, next_state
The validation update simply consists of updating the two components that make use of the loss value.
def next_state_ val iddtiin (self, current_state: State,
valid_loas, float) State:
23

---

Under review as a conference paper at ICLR 2021
next valivaliid... loss self.valid loss accum. state purrent_state.raliclooos.oo valid_loss)
return current_sstaee__replaaaa valid valeeloesaccuurroo.eaallooo....o
The the function applied to the training gradients is much complicated. Much like the validation features it updates the various rolling statistics including the two loss data structures, the rolling momentum rms terms, well as the gradient clipping state.
def sntr_state_traiinninn (seif,
params: Params, current_state:
current state: State,
loss: float,
grads: Grads) -> (Params, State)
grads 1ist (grads. next_train_los__ecnn
current perrent_staae.traiine
wttttttrannnsssssooe pess_accuu, loss)
next_dynamic_cllpp grads =
(cuzzen_staaee.dyyaiicec. grads)
rex__rellin__fetuuee rot..rooling_fcatuuee mtttussssnctttttttee
grads)
Next, sequence of features related to loss values are computed.
train_loss feat = self.train lcas wcccm..eighttttte..... _oss/
next_train__oss_acom loss)
valic_less self http:p/oe.....ooo....ooooooonnnr
valic_loss_accum)
valid_loss_feat
next_ next_train_loss_acom
valid_loss)
Next we call into the learned optimizer function "theta which does the bulk of the computa- tion.
next next_params, params, from_lstm, from_mlp, lstm_hidden_state
self tel.fmthetammdd.cmom
train_loss_fent feat, cellidlloss_fed current_state.from
current_state..rrmm lstm, roxt_roolingg_faturees grads, params,
current_stae..triinnnn
We then populate the next state with all the outputs.
next_state = Stato
Banling__atursssssstooooonnnnnn cralniigg.tppparrre.n....eeaaagg from
1,
lsnnhhideen_.tteeesssm..iddeeee from__ssttffrom_stta
from_mlp=ffrom_mpp
Lrain train loss accum,
next partms neat_state
24

---

Under review as a conference paper at ICLR 2021
C.4 "THETAMOD": THE LEARNED OPTIMIZER PAREE-PRRAMETERS
All of the meta-leamed outer-parametent are stored in a class called ThetaMod. Its constructor defines various dimensions of the different Sonnet modules.
class ThetaMod:
init (self,
$ Size of the features from the per paraneter mip. from from_mip sap_mmmpppsa___saa size=16,
# size of the features from the per tensor LSTM. croo__sst_siiee1e # Size of features from the 1stm being passed to the per param MLP.
1stm 1stm_to_ff-1 to_ff-17,
LSTM hidden size.
lstm_hidden_size-
# Multiplier on output to improve conditioning Seep_multpplierr...o
# Multiplier octput to improve condition tindiigonng
magnitude_rate-0.001
**kwargs):
self.step multiplier = step multiplier
self .ww.nitude rate = magnitude rate
This constructor creates shared sonnet module for the per layer LSTM. This is constructed here as it is needed in more than one method.
self,
rnn rnn Ini.iaaletatteL.nttssss....bbddd
Second, we construct the per parameter feed forward network. This is a tiny network that operates on each parameter. It outputs two values for producing a step, then from_mij from_mlp_stt mlp_size more values which are fed back into the LSTM.
sel f. ffmod - snt net...tt.............. MMiP ([32, 32]
[2 sel f. = ="PerParamMil
name
Third, we construct linear projections that map from the RNN output to either the per parameter feed forward network, or back to the global features.
self. ron_to_ff snt. Linear (1stm_to_ff, name= name="ran_ee "atn__nnn_tt"""" relf.ron_to_in = mat.Lincar (fron_lstm_sizo, name="rnn_to_ibe
Finally, initialize the two variables used by the rolling gradient clipping terms. We parameterize these variables on log scale (possibly subtracting 1) so as to roughly match the scaling of the other neural network weights.
def get pet_log_param_variabe (name,
mitiaalvalue
rinus_one=False)
if minus_ one: initial_value - - initial_value
init_log np. log (initial_value)
tff.getvvriiablet name, initiallzer=tf..ons (shape=[]) tf.constant (
init_log, stype=tf.float333
trainable-True)
ret tf. if minus_one:
ret = 1 ret
return ret
25

---

Under review as a conference paper at ICLR 2021
self.rolling_cia alpha get log param variable nane=>roliig__lip__lllp initial_value-0.99 minus_one-Truc)
self, self.rooling_clii._iit ge__lo__aaran_aariiee
nange""rolling__lip.m.m.....
initial_value=10.10 minus_one-=alle
The heart of this function takes the features passed in from the optimizer, applies some neural network based learning algorithm, then produces the next weight value well various pieces of hidden information for the next iteration.
def compute_update
[self,
Training loss feature train_loss_fes feat:: Tensor [shape- () Validation loss feature
valid_loss loss : Tensor s I,
Information passed to each layer of the LSTM from the paran MLP.
from_mlp: List [Tensor [shape- strom_ml_siiee Information aggregated across LSTM layers to be passed to all LSTM layers.
#rem_lstm Tensor [shape- (frem latrror_sstm_iieee The current LSTM state.
1stm_state: 1sTMHiedenState,
Features from the momentum / rolling second moments. Dolling_features BollingFeatureSee
Gradients from current training mini-batch grads: Grads,
Current parameters
params: Params,
current inner-iteration traininggster int,
the output parameter value. (Params,
# next from_lstm
Tensor (shape-
next from_mlp
List Tensor shape= (fron next 1stm hidden state. LSTMaidtienState))
We first compute the global feature vector. These consist of features related to the loss curves, as well as the aggregated data from the previous LSTM execution. These features form a rank 1 tensor.
@lobal_feature sel f. compue.mmeeeggleaee 1_eatureeseesseeeee (train_loss_feat
Molid_loss__fi from_1stm)
Next we compute the per tensor features. This rank two tensor, consisting of the number of tensors, and feature dimension.
per wwweensorfeattrreee
self self seaff..ommpre_etnss__eeeerrr (grads,
params, from mlp,
features)
We tile the global features to be num tensors by num global features then concat these to the per tensor features and pass them into the per layer LSTM.
26

---

Under review as a conference paper at ICLR 2021
rnn inputs tf. concatenate (1 [epr_tensoo mansor__eatures,
tž.tile (global teatures, Iper_ tensor_ feaerrrettttteeeeeee shape 0 =)) 1stm_out, heat__sstm_hidden__e
self. rnn (rnn_inputs, 1stm_state)
We then linearly project the LSTM outputs to be input to the per parameter feed foward network, and unstack them form list of rank features per tensor.
If inputs self, rnn to Ef (1stm_outs EE inputs = t..unstcck (EE inputs)
Next we iterate over each parameter, and first compute per parameter features. The per parameter features consist of gradients, parameter values, parameter values 100 steps in the past, momentum values, second moment values, activations from the per tensor LSTM. Additionally there are two global features included: the training step and the number of tensors total. Two empty lists initialized to accumulate results.
next__o _st_sssssssss__sssssss___ssssss_ next params
for wi in range (len (params)
feats compute_per_pa laattrffeuutsssee (grads [wil,
param [aram[wi],
holling_featuree
[nlllingfeeatueer rms [wi],
ff_inputs [wi]
training_step,
len (params))
We pass these parameter features through these features into per parameter MIP which produces a single output.
output self f. ffmod (feats)
We split this output into three pieces: direction, a magnitude, and features which are be fed back to the per tensor LSTT after being reduced along the parameter dimension.
direction output 1;, 0:11
magnitude output I;, 1:2]
te_1stm - tf. reduce_mean (cutput [:, 2:1, axis-0)
step can then be computed by exponentiating the magnitude and multiplying by direction. Rescaling these quantities is critical. Without this rescaling, the magnitude of the updates is extremely large, resulting in chaotic training dynamics which makes outer-learning difficult.
step rection
tf. ex. (magni tude self self.masee magni tude_ rate) seli -step_multiplier
We then subtract step from the previous parameter value. While it is possible just to predict parameter directly, learning deltas significantly casier to learn and results much more stable training.
next_ paran
param
wii
step
Finally, collect all the information needed be returned from the per parameter LSTT and output.
next next_paraass append (next_param) next_to_ append (next_to_1stm)
With the for loop over, then return the necessary info. Before this, however, must aggregate information from the per tensor LSTM to be able to pass it to the next layer's global features after applying linear projection from the LST output and reducing along the tensor dimension.
27

---

Under review as a conference paper at ICLR 2021
frem_lstm refreeuce,e.oo meee (self ton (1stm_out), axis=0)
return next _ext__params_ stom_istm,, next to_ _stms,, next histt_hidden_stallee
Now will discuss the different features methods. First, computing the global features.
def compute_globe _elll_eeettrsss (self, train lecph_sss____e valid_loss_feadt Tensor (shape== ()
feat: feat:
Tensor censorsshapes () ]
from_lstm: Tensor [shape= (n_ )
): -> Tensor [shape= (t_gllobalfeettures
This function simply returns the concatenation of the Two loss features, and the aggregated features from the per layer Istm.
return tf. concatenate ( [|train Straiin_oss________ foss_feet,, valid loss feat ) from 1stm] )
Next look to the per tensor features. We return 2D tensor with the leading dimension "batch size" being the number tensors, and the second dimension representing the feature dimension.
def
(self,
grads; Grads,
params: Params,
from_mls List [Tenser (shape- (fromm (frammpptfrronmm__ssieee ]], velling_featured RollingfeatureState
): -> Tensor shape= (a_tensors, n_features)]
We compute large number of features. At this point have not extensively ablated each one and expect many are duplicate or unneeded. Before we begin listing features, however, we make use of a helper function that computes the clipped log abs of tensors.
def clip log mag tf. log (1e-8 a.s (v) return tf. el ip p_bpy__lllllll (mag, -5, 5)
We begin by iterating over each tensor.
stackcd_inputs [] for i in range (len (params))
inputs (}
First, we compute features based on the rolling momentum values. We compute features that represent the log magnitude, the sign, and the log variance.
tf. reduce
(relling features.m....... i1) clip _og_abs (mean
inputs "mean_ ms _as_mmmmm_s_ ____mag" inputs | "mean_nee ms sign"| | = sign (mean ms)
var_ms = tf. www.rreddeaaeeeee meanceee_eaaa_e (t squareessueeeeaeeoeeeeeeereeeee. (
mean_ms)
inputs ["var_ms") clip_log_abs (var_ms)
We compute a similar set of features for the second moment features.
mean_rma rms tf. reduce_deeeaee meceeeeeeeaeee wtw.lliiengeauuuesssoo rms (1])
inputsl aneean_ms" = elip_log_ate (mean_rms)
inputa ["mean_signt
tf tf.sigm (mean_rms)
var_rms tf. www.reddcc_eee maannce__eeaa_eeee (t. ftaaantt.uuuuauaaae
square ( ris
mean_rms))
inputs | "var rms"
clip abs
similar set for the inner-parameters:
28

---

Under review as a conference paper at ICLR 2021
parans @eatuucceeeaaaeeeee
mean_y = tf. reduce
mean (v)
inputs ["mean___aa"]] clip_log_abs (mean_v)
inputs [ - tf. sign (mean_v)
var_v re. (tf (tf.equare (v me)n_v))
inputs "var_v"] slip_log_abs (var_v)
Next we include the magnitude of the weight norm.
inputs snorn_weight?2
clip_log_abs (tf..orm (pormppaama
As well as the gradient norm:
inputs "g_norm"] clip_log_abs nor.
norm (grads (i]))
Next we add feature if the tensor is scalar not.
len (v, www.shaappsssseeee.c ()) 0:
is_scalar tf.constant (1.0) else:
is sc_acr co.contt (-1.0)
inputs "is_ "rss_sscaarr _s_scalar
The shape of the underlying tensor is also a useful feature. The length of this shape varies though depending the rank. As such first pad the tensor to be least rank 4 and then pass the log shape shifted by -1. This shift is to keep this value roughly zero mean.
extra_dims [1.] (4 len (v. shape.
())
shape_stack tf concat [[tf.te_ff claceffottttooooooeoooooooooooeo ft... (v)),
tt.stack axis=0)
for j in range ( : shift so 10 is neg one to be better scaled.
inputs (oouts("shan ["shape_*d" àd" jl tf.log (stape_stack) [3]
1.0
The features from the previous execution's 's MLP are included. These are enable communication back from the MLP to the per layer network.
input "from_mip""
from (tom_mlpli)
This dictionary of features is then flattened (adding appropriate dimensions if needed) and appended the list of inputs.
values sorted values values (inputs)
reshaped tf cf.cxpand_dime (v, 0) if len (v. www.shhc..ss_lsttt.o.o.o.o values
else v
for v in vnluts.values ()
stacked_inpute append (reshaped)
Once the stacked inputs is filled, we zip and stack and concat the values forming a number of layers by number of features tensor.
(v) for v in zip ( return co.concat (stacked_features) axis-1)
stacked_featuree -
(=stacked_inputs))
The last of features is the per parameter features.
def compute_per_nt aaaatfueuuresssss (se1f. grad: Tensor,
29

---

Under review as a conference paper at ICLR 2021
param: Tensor, ms: Tensor,
rms: Tensor,
ff_inputs: Tenser (shape- (fun_ff ) training_step: int, num num_params: int
Tensor [shape- Lnnum_parass,.
num_ sum_features)
Many of these inputs are of the shape of the underlying tensor. Because many of these operations are per parameter we construct a flat version of all of these.
flat_g tf. reshape (grad, [-1, 11)
flat_y tf. reshape (param, [-1, 11)
flat ms reshape (ms, [-1, ms. shape (-11)
flat__mm rms = tf. reshape (rms, [-1, rms. www.e (-11)
As with the per layer features, add series of features. First we add raw features for the gradients, parameter values, momentum, and second moment accumulators, and past weights.
inps = ()
inpa ["flat_g"] flat_g
inps ["flat_v"] flat_v
inpa [mm"] m
inps "rms"]
We then add some derived features such as the log abs parameter value:
inps [" ["log_abs_v"] tf. (flat_v) 1e-8)
Finally we add terms based on the rsqrt of the second moment terms inspired by how Adam normalizes.
rsqrt = rsqrt (rms 1e-6)
rms rms_scaled_g
rsqrt inps (i"msm_sscall som__scaled_g" rms_scaled_g
inps ["rsqrt"] rsqrt
These values are concatenated and normalized by second moment normalizer computed across the num tensor dimension. We additionally clip the output this normalizer. This crucial for stable training some of these features will produce large magnitudes in some channels.
inp t.. concat (sorted_ peattedd_alleeee velues (inps), 1)
inp = inp tf.rsqrt (1e-8 ttf.reduc__meet wewluuceeeeaaeeeeeee (tf. square (inp), axis=0,
keep_dins" dins= True)
inp - cf.clippby_yaluee (inp * 0.5, -1, 1)
Next, embed the current inner-training step with sinusoids of different frequencies.
def sin Lin_enbedding (x)
pio_projj in 11, 3, 10, 30, 100, 300, 1000, 3000, 10000, 30000, 100000]:
s - tof.to_float mix_proj- append (tf.sin tf.to_ float )x)) return st.stack (mix (mix_proo)
(i) / float (np.pi))
step = utils.sin utiils..ii.emmeeddnngn (training_ste)) stack_step tf.tile(
tf. reshas (step, [1, -1)), tf.stack (tf, shape (flat_g) (0), 11))
30

---

Under review as a conference paper at ICLR 2021
We compute features based the number of tensors:
tensooss tf. (float (Len (ann
Stack_nun_tensors tf.tile(
tf. reshape pog_nnum_tenscreee [1, 11),
tf. tt.sk Itf. shape (flat_g) (0), 111)
As well features about the number of parameters in the current tensor:
loegn_weight tf. 1og (t float (tf. shape (flat_v) [0]))
stack _aack_log_nwweii tf. reshape niohh__eight, [1, 1]),
tf.tile(
tf.stack ([tf. (flat_g) [0], 111)
Next, we compute some statistics about the log norm of the weight matrix:
Log_nore tf. log (tf. (flat_v) stack_ fack_log_norm tf.tile(
1e-8)
reshape (log_norm, [1, 11),
tf.stack ([tf. shape (flat_g) [0], 111)
Next, we tile the input from the LSTM so as to add a number of parameters dimension.
Ff inp tfltile(
tf. reshape (ff_inputs [1, -11), tf.stack [tt. shape (flat_g) [0), 11))
Finally concat and return all these features.
return tf. concat (I
inp, stack_stee) stack_ ptt__uu___e_ea_ee com_tensor, stack loo___log_oomm, stack_log_n_neghht ff_inp
1, axis-1)
We have discussed how new values of LSTM hidden state, activations from the LSTMs, and activations from the MLP are produced, but we have not yet shown the initial values. In all cases they are initialized from puter-parameteriize values shown bellow.
8ant reuse__vve_e def = self. initial state (n, trainable=True)
n)
return
8ent. rosse_evs__arr def inittal__fron__stmmccllff return tf tf.get_variablet
name"Innttallffrooleett shape- dtype-tf.float.32, trainable-True)
cent.reusee_arrablle
initial from lalaammmmpppsse (mp n)
s = tf.get_variablee
name="inittal__fronnnlpp
ctype=tf.float322, trainable-True)
return [3 for in range (n))
31

---

Under review as a conference paper at ICLR 2021
D ADAM8P
This matches the adam8p optimizer described in (Metz et al. 2020). The eight the learning rate, a, first and second moment momentum, 31. 32, the numerical stability term, 6, 8 and E regularization strength, and learning rate schedule constants dexp,decay and Alinear decay-
are:
(6(0) problem specified random initialization
(1)
(0) =0
(0) -0
(3)
(() d
dølc) (f (10) )+ + 11 (11
(4)
m (1) = 3mm (1-1) + g0 (1- -B1)
(5)
(1) (1 (g (0) B2)
(1 (l))
B1+1
(1)
((0)
(8)
B++ m (1)
(1)
(O()) +e
8 linear max (1
(4)
DAtineur decay: 0)
(10)
s
expp tteppppue decay)
(11)
(t+ 1) -assiiner (1) ue
(12)
We sample learning rate logritmically between 1c-8 and 10, betal and beta2 we parametrize as 1 I and sample logrithmically between lc-4 and and Ie-6 and I respectively. For learning rate schedules we sample linear decay between Ic-7, le-4 logrithmically and exponential decay logrithmically between 1e-3, le-6. We sample both & and & logrithmcally between 1e-8, lel.
E PERFORMANCE TABLE
We show numerical performance measurements of our learned optimizer and baseline optimizers. Each value the mean or median over 100 different inner-training tasks.
train mean
test
Optimizer
median
mean
median
learned
0.077325
0.024167
0.113242
0.028928
global adam Ir (1 Trial)
0.226807
0.105513
0.261613
0.092532
global adam&p, RS (1 Trial)
0.240924
0. 10199519
0.257368
0.100959
global nadamw, RS (1 Trial)
0.215176
0.098448
0.229359
0.086703
per task adam Ir (14 Trial)
0.150443
0.067373
0.152028
0.063194
per task opt_ list (10 Trial)
0.082756
0.039647
0.086195
0.038541
per task opt list ( (100 Trial)
0.041560
0.028851
0.047224
0.031042
per task opt list (Ik Trial)
0.039348
0.028254
0.043193
0.029897
task adamöp (10 Trial)
0.278858
0. 03132965
0.264282
0.
per task adam8p (100 Trial)
0.095372
0.046926
0.107177
0.046672
per task adam8p (Ik Trial)
0.044144
0.028692
0.060663
0.028813
per task adam8p + apt_list(2k Trial)
0.018403
0.023350
0.025419
0.025555
32

---

Under review as a conference paper at ICLR 2021
F EXPERIMENTAL DETAILS FOR FIGURES
F.1 OUTER TRAINING OF DIFFERENT LEARNED ARCIITECTURES
All models are trained with truncated evolutionary strategies using the parameters described in the main text. Instead of using the two stage schedule we only train with truncation's 's of size 240-360. We use the same, fixed learning rate for the FF and LSTM_FF model of 0.000500. For LSTM we found a lower learning rate of 0. 0.000100 to perform better. Ideally we would include an extensive hyperparameter comparison, but the computational cost is prohibitive.
For the FF model we make use of a two hidden layer 32 unit MLP. In addition to the features described in (Metz al. 2019b> we use RMS terms which improve performance over that reported in (Metz ct al. 2019b by a small amount.
For the L2LBGDGD model we copy Andrychowicz et al. (2016) using two hidden layer, 20 unit GRU (Chung et al. 2014) with the same input gradient processing.
For the LSTM model, make of model inspired by (Andrychowicz et al. 2016). We make use of a 64 dimension LSTM hidden size to produce two outputs, direction and a sign which are combined similar to (Metz al. 2019b).
Performance on 100 tasks (each with five random seeds) is recorded over the course of training. We show exponential moving average of this data in the figure.
F.2 OUTER TRAINING WITH DIFFERENT SIZED MODELS
All models are trained with truncated evolutionary strategies using parameters described in the main text. Instead of using the two stage schedule only train with truncation's of size 240-360. We use a learning rate of 3 10-3 for all models.
F.3 LARGE SCALE TRANSPER WITH RESNET MODELS
We make use of two ResNetV2 models. Before applying the residual blocks, we pass the inputs through a 64 unit, 7x7 convolutional kernel with stride 2 followed by max pooling of size with a stride of 2. This aggressive down sampling originally designed for larger image models but was also left in despite our small sized models.
We then apply some number of residual blocks parameterized by (number of output channels, number of bottleneck channels, stride).
At the output apply batch norm, reduce the spatial dimensions, and project to the the appropriate number of classes.
Our goal with these experiments are not seek peak performance. Instead show that learned optimizer is capable generalizing to vastly different distributions of models.
CIFAR-10 Resnet This model makes use of four residual blocks: [(128, 32, 1). (128, 32, 2), (128, 64, 1), (128, 64, 2)).
Imagenet Resnet This model makes use of residual blocks: [(128, 32, 1), (128, 32, 2), (256, 64, 1)*2, (256, 64, 2), (512, 128, 1)*3, (512, 128, 2), (1024, 256, 1)*2].
F.4 SELF OPTIMIZATION DETAILS
For our baseline learned optimizer trained with Adam, we search over two learning rates, 10-5 and 3 10 and select the best one (1e-3) for the given 10k outer-training steps as well as selected the optimizer used to train the best optimizer (3e-5).
When training an optimizer with also include the original gradient clipping used for the Adam models (clipping all values between 0.1, 0) Without this, our learned optimizer is not able to make any progress. Understanding this is
optimizer there no hyper parameters. We found must
33

---

Under review as a conference paper at ICLR 2021
an interesting to would have suspected the dynamic inner-clipping should have addressed this.
For technical reasons, the learned optimizer does not have easy access validation losses. Instead, use the training loss for both. This computed as the mean over the normalized losses computed over single batch.
G SELF OPTIMIZATION EXTENDED
0.40
Adam(3e-5) Learned ad
Outer Outertrainin loss
0.35
0.30
0.15
0.10
5000
10000
15000 20000 25000 30000
Outer training iteration
Figure 8: When using learned optimizer to train itself, find good performance early in training, but diverge after 10k outer-iterations. On the x-axis we show number of weight updates done to the learned optimizer. On the y-axis show outer-loss. Each point consists of inner-training 100 models, each with five random initializations trained for 10k inner-iterations. We show the average validation performance post-normane normalzaaliinnonn averaged over all tasks, seeds, and inner-training steps. Each line represents a different randomly initialized learned optimizer. In orange we show Adam with learning rate of 3e-5 the same value used train the optimizers in this work. In green, we show Adam with learning rate of 1e-3, the learning rate that performed best in this 10k outer-iteration regime. In blue we show our learned optimizer. Light colors represent instantaneous outer-losses while solid denotes smoothing done via an exponential moving average.
34

---

Under review as a conference paper at ICLR 2021
H INFRASTRUCTURE USED
The types of learning systems we discuss here are more complex than average deep learning training setup. We are training not one, but many different kinds of models on different datasets with wildly different properties. As result, the static graph paradigm of traditional deep learning models distributed one graph and using MPI / allreduce style primitives is insufficient. Finally, evaluating the performance of an optimizer is not trivial task as it requires inner-training a large number of different models for large number of inner-iterations.
Developing custom training and evaluation infrastructure has been critical to the success of this family of work. We describe the training infrastructure employed, discuss techniques use for evaluation and monitoring. At this time, we are unable to release a fully running open source implementation due to the use of internal tools.
H.1 OUTER-TRAINING CLUSTER
The cluster consists of six kinds of jobs workers, learner, gradient storage, summary aggregators, evaluation chief, evaluation worker. Additionally make use of distributed file system (Ghemawat
et al. 2003).
Learner (1x)
Grab batches of
Gradient storage (1x) Stores ES naise and losses
3x (eval meta-train, meta-test, + other out of distribution)
outer-gradients
Update purar-parameters
Saves checkpoints
Eval Chief
Coordinate
Worker (512x-4096x)
inner-training on various task distributions
Distributed file system
Perform unrolls on sampled problems. Send outer-gradient info.
Stores checkpoints of optimizers
Stores results from evaluations
Update weights from Meta-Trainer
Stores summaries
Eval workers (300x)
Compute summaries models
Run inner-training with optimizer on task
Distributed GRPC based summary aggregator (16x)
Figure 9: High level / system diagram for infrastructure used. See Appendix H for more information.
Workers Workers are jobs that sample outer-task, build a tensorflow (Abadi et al. (016) graph and provides some way to estimate outer-gradients. In this work, we leverage antithetic ES sampling with shared randomness where possible. In practice, this means start two inner-training procedures starting from the same initializations, using the same minibatches of data, but leveraging two slightly differnt outer-parameter values. Before each outer- sittietr estimate new outer- pilter-prameters. well as the current outer-iteration are obtained by RPC call to the learner. This gradient estimate, the current outer-iteration (e.g. global step) along with the type of problem selected sent to the key value storage worker keyed with some unique identifier.
Gradient Storage The gradient storage job is a distributed dictionary which stores information computed from the worker. At this point, given the low frequency of updates only single machine is required. This a separate machine from the learner to control the number queries required sent to the learner.
35

---

Under review as a conference paper at ICLR 2021
Learner The learner polls the gradient storage for batches of data. If the outer-training iteration is within some fixed amount of time to the current outer iteration the data is deleted. In this work we use 5. This prevents stale gradients from being used to update the outer parameters. This has the negative effect of potentially ignoring the slowest tasks. At this point, however, we find that this rejection threshold is quite small (<1%) and thus safely ignored. When batch of outer-information is obtained, the learner aggregates outer-gradients, and performs an update of outer-optimization, in our case Adam. Current loss values for each task family are also recorded in dictionary, and aggregated across all tasks seen thus far. We find this averaged signal, as well as simply just logging out the loss from the current batch, to be useful for evaluating model performance as it is lower variance and is stationary in time. This job is also responsible for writing the meta--arrseeters to disk every 10 min for use later and for evaluation.
Summary aggregators When training a standard model, users monitor values - e.g. activation values, losses, or accuracies. In this work we are training thousands of different neural network tasks at time and would ideally like be able to obtain more insight into what is going on. Existing tooling such as TensorFlow's event files are inadequate.
We work around this by introducing a distributed service meant just for summary aggregation. This service exposes a RPC api that takes batches of summaries and logs them to disk in a custom binary format for fast visualization later. We shard request based on the name of the summary. We additionally log out subsampled data for case of viewing after training.
In this work, we process roughly 10k summaries a second split across 24 machines. We found being able to do in-depth post-hoc analysis of model training to be critical to diagnosing issues and uncovering bugs. For example, looking at gradient variance for each task in the outer-training set let us diagmose which tasks are diverging and why. To ensure that summary computation does not cause increased computational overhead we make use of both stochastic logging as well as both batching of RPC calls summary aggregators, and sending these RPC in parallel with computation.
Evaluations: Evaluating a learned optimizer is not nearly simple say performing forward passes on some test set of minibatches. Evaluation consists of an expectation over training many different tasks. To obtain a low variance estimate we employ a large amount of compute to evaluate the checkpoints saved out by the learner. To lower variance, we pick a fixed set of tasks to evaluate. In addition the meta-training task distribution (which we pick subset of 100 to evaluate), we also evaluate on a test set of tasks which is iid to the meta-train but never seen when outer-training as well as an extremely out of distribution set of data not close to tasks in the outer-training distribution. By monitoring a variety of different problems, with different similarities to the meta-training set we obtain a better sense of outer- sener-eeneralization. These evaluations take time, around an hour of compute per task, so it must parallelized.
Originally, we tried to avoid this complexity and obtain some measure of outer-loss directly from the training cluster. While this works to some extent, it is noisy, and doesn't capture what we actually care about - performance of some fixed weights of a learned optimize when applied to a target problem for extended period of time. Averaging over the meta-training cluster also doesn't 't capture variance from outer-weight to weigh.
The evaluation of one task also requires orchestration. We employ evaluation chief, and evaluation workers for each task set. While its possible to do this type of evaluation offline, we find running online yields much better workflow and thus speeds up research.
Evaluation Chief: The evaluation chief is responsible for watching directories distributed filesystem and to enqueue evaluation tasks. Workers use RPC to request, and to send back results (e.g. learning curves). When set of tasks is complete for a given outer-parameter checkpoint the results are written out to disk for later analysis.
Evaluation workers: Evaluation workers request task configurations from the chief over RPC. This configuration is parsed and converted to a TensorFlow graph which is then run. Over the course of running, various signals including training loss, validation loss, and test loss, are recorded. Upon completion these results are sent back over RPC the chief to be aggregated and written to disk.
36

---

Under review as a conference paper at ICLR 2021
H.2 COMPUTATIONAL EXPENSE
At this point, training of learned optimizer is quite expensive. It roughly entails 60K CPU cores for around month, around 5k CPU years. Much as neural architecture search has been dramatically optimized and improved, we hope learned optimizers will obtain a similar outer-training speedups. The energy and environmental impact of this work also worth noting. These models take on the order of 200 megawatt hours of power to outer-train. We believe that these methods can be used to speed up development of many existing models once outer-trained. By meta-training a hyper parameter free model, one has to do no hyper parameter tuning of the optimizer for example. In general, one should weight the costs outer-training against the potential savings achieved. In the future we expect these models will be trained once and applied everywhere.
Finally, we would like to highlight that of this work was performed on CPU. tests on GPU (without optimizing) are ill-suited for the small workloads we perform here (serially training lots of small networks). We believe it is possible to vectorize the inner-training of neural networks (e.g. training N networks instead of just 1 per worker) but have not explored this route yet.
found performance be similar. Modern accelerator hardware, TPU, GPU,
37


Reviews:
[{"review_id": "peTyZOxb7gq", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary\n\nThis paper attempts to address the fundamental barriers of learned optimization. The authors identify three barriers: computational requirements, number of training tasks and lack of inductive bias. A “large-scale” evaluation and comparison of learned optimizers is then carried out using many (1024) multi-core CPUs. A simple modification of an existing learned optimizer is also proposed that involves adding more input features. Unfortunately the results don’t seem to reveal any new insights.  \n\nStrengths\n- The paper proposes a new, simple hierarchical learned optimizer that outperforms existing learned optimizers. The proposed model is very simple in theory but the implementation seems to still require quite a bit of “hand-engineering” in terms of selecting features etc.\n\n\n- The experimental investigation reveals some interesting (albeit unsurprising) insights of large scale training of learned optimizers. These include things like training with more tasks improves performance, that learned optimization performs well in the hyper-parameter regime in which it was trained, that it learns some form of regularization and that it outperforms Adam when a non-optimal learning rate is used.\n\n\nConcerns:\n- My main concern with the paper is that some claims are over-blown. Although it is not clear at all that the current generation of learned optimizers can outperform hand-crafted optimizers, the paper makes misleading claims that can easily be taken out of context. Statements like, “We see this final accomplishment as being analogous to the first time a compiler is complete enough that it can be used to compile itself.” and “we believe learned algorithms will transform how we train models\" are too strong given the current evidence of the performance of the learned optimizers. I would suggest the authors tone down these claims.\n\n\n- The proposed hierarchical learned optimizer (as well as existing ones) seem to be more fragile than hand-crafted approaches such as Adam. For example, on CIFAR-10 in Figure 5 the learned optimizer fails even for batch sizes in the training regime. Is there any reason why this might be the case, especially considering that it has access to all the same information as Adam and Adam’s “hand-crafted” operations are quite simple? \n\n\n- The disadvantages of the proposed learned optimizer still seems to outweigh the benefits. For example, the “careful tuning of learning rate schedules and momentum timescales” is traded instead for the selection of design of a sufficient range of tasks on which to train the optimizer. This seems to be a far more difficult task than just tuning a few hyperparameters. In addition, although hand-crafted optimizers “do not leverage alternative sources of information beyond the gradient”, the learned optimizers do not do much better and just seem to learn very simple regularization strategies. Currently the discussion on the advantages and disadvantages is completely separate. I think these need to be contrasted and compared on the grounds of what properties a user would prefer in an optimizer.\n\n\n- The contribution of the paper in terms of new insight or knowledge is not clear. The “large-scale” training on a wide range of tasks and many unrolled steps is interesting but I’m not sure what new insights can be inferred from this? Furthermore, the hierarchical optimizer seems to be a small improvement of the mode proposed in (Wichrowska et al., 2017) with some additional input information (like validation loss).\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "maBKcvqbl0g", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nStrengths: \n+ This paper thoroughly examines the challenges of how to train a learned optimizer.\n\nWeaknesses:\n+ Training the learned optimizer is fairly complex and computationally expensive, which will prevent broader adoption and makes the paper difficult to reproduce.\n\nQuestions:\n1) What conditions are necessary for a learned optimizer to replace a standard optimizer like Adam?\n2) Could a learned optimizer work with noisy, quantized gradients? e.g. deploying a learned optimizer in a federated learning environment\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "XbxgzNQAycE", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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. \n\nThe task is a challenging one and the authors propose a hierarchical optimizer similar to Wichrowska et. al. I wish the authors invested more thought into the architecture of the optimizer beyond assimilating features of previous work. \n\nMy main concern with the paper is lack of generalization capability of the proposed approach. The authors could have done a more thorough job evaluating the architecture by performing ablation studies, i.e., different permutation of inputs to the LSTM and feed-forward modules and understanding its effect. Fo e.g., how useful is the second moment or parameter values to FF; tensor shape/losses to the LSTM. I am also interested in understanding the trade-off involved in using ES over unrolled optimization. There should be some experiment discussing this trade-off as it seems to be a critical improvement over prior approaches. The sub optimal performance on imagenet is concerning.The authors should also provide guidance on task selection for optimizer training. This guidance could be coupled to the sub-par performance on imagenet. \n\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper aims to address several challenges in learning neural network-based optimization algorithms by increasing the #unrolled steps, increasing the #training tasks, and exploring new parameterizations for the learning optimizer. The authors demonstrated the effectiveness of applying persisted Evolution Stratergies and backdrop through over 10,000 inner-loop steps can improve the performance of the learned optimizer. Empirical experiments showcased incorporating LSTM to the previous state-of-the-art improve their training performance. 

There are a lot of interesting ideas in the paper. However, packaging them together and only glance over each idea briefly unfortunately dilutes the contribution and the novelty of the work. There are still some major concerns echoed among the reviewer:

1) The proposed hierarchical optimizer seems interesting. It is one of the major contributions of the paper. But, its architecture was only briefly mentioned in Sec 3.3. Its motivation, implementation and the corresponding engineering choices remain unclear by just reading the main text. Some of the details were discussed in the appendix but it would be of great interest if authors could give some intuition on which subset of the tasks the proposed architecture gives the most improvement / failure among the 6000 tasks.

2) Training the optimizer on a diverse set of tasks is crucial for the learned optimizer to generalize. One of the paper's contributions is to further expand the task dataset from the prior work Metz et al., (2020). The authors have conducted very thorough experiments on this new dataset, which is amazing. I would argue there are even enough results for another standalone paper. However, there is surprisingly little detail on how the newly proposed dataset differs from the prior TaskSet dataset. What are the new optimization problems? How are they different from the family of tasks in TaskSet? A TSNE plot of the tasks similar to Figure 1 from Metz et al. (2020) could provide more intuition for the reader and highlight the contribution. 

Overall, if the authors could provide more insight into their experiments and the proposed methods, it would help the readers greatly to see the novelty and the contribution of the paper. The current version of the paper will need additional development and non-trivial modifications to be broadly appreciated by the community. 


Decision: Reject