Tasks, stability, architecture, and compute: Training more effective learned optimizers, and using them to train themselves — part 2/5
Cleaned-text code points [15917, 31601). Bibliography omitted here; source.md includes it.
For the “off-the-shelf” limit, we consider the scenario where a practitioner has access to a limited number of optimization runs (trials) ( ≤ 10 \leq 10 ) for a particular problem. Thus, we select a first set of baseline hyperparameters using a single, default value (denoted global in Fig 3 ) across all of the tasks. We 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 a 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, we consider task-specific hyperparameter tuning, where the hyperparameters for each baseline optimizer are selected individually for each task (denoted per-task in Fig 3 ).
We plot a histogram over tasks showing the difference in performance between the each baseline optimizer and the learned optimizer in each row of 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 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 on 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 ℓ 2 \ell_{2} 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 a simple 2D loss surface, with a continuum of solutions along the diagonal: f ( x , y ) = 1 2 ( x − y ) 2 f(x,y)=\frac{1}{2}\left(x-y\right)^{2} . Although any point along the x = y x=y diagonal is a global minimum, we wanted to see if the learned optimizer would prefer any particular solution within that set.
Figure 4 shows the resulting training trajectories along the 2D loss surface from many starting points. For a 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 it pushes the parameters towards a solution with small norm: ( 0 , 0 ) (0,0) .
### 4.4 Generalization along different task axes
Next, we 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-10 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 learned optimizer to a baseline optimizer, Adam, over a grid of eight different learning rates logarithmically spaced from 10 − 5.5 10^{-5.5} to 10 − 2 10^{-2} .
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 of the best learning rate tuned Adam optimizer. On CIFAR-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.
### 4.5 Generalization to large-scale problems
We test the learned optimizer on large-scale machine learning tasks. We use two ResNet V2 [ 36 ] architectures: a 14 layer residual network trained on CIFAR-10 [ 37 ] ; and a 35 layer residual network trained on 64x64 resized ImageNet [ 38 ] . We train both networks with the learned optimizer and compare the performance to learning rate tuned Adam and Momentum (Fig. 6 ). Details in Appendix D . For CIFAR-10, we find that the learned optimizer achieves similar performance as the baselines but does not overfit later in inner-training. For ImageNet, we find that the learned optimizer performs slightly worse.
Note that our baselines only include learning rate tuning. More specialized hyperparameter configurations, designed specifically for these tasks, such as learning rate schedules and data augmentation strategies will perform better. An extensive study of learned optimizer performance on a wide range of state-of-the-art models is a subject for future work.
### 4.6 Learned optimizers training themselves
Finally, we performed an experiment to test if a learned optimizer can be use to train new learned optimizers. Figure 8 shows that this “self-optimized” training curve is 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 of 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 is continued beyond 10k iterations, learned optimizer performance worsens (see Appendix G ), suggesting that more inner-iterations are needed when outer-training.
## 5 Discussion
In this work, we train a learned optimizer using a larger and more diverse set of training tasks, better optimization techniques, an improved optimizer architecture, and more compute than previous work. The resulting optimizer outperforms hand-designed optimizers constrained to a single set of hyperparameters, and performs comparably to hand designed optimizers after a modest hyperparameter search. Further, it demonstrates some ability to generalize to optimization tasks unlike those it was trained on. Most dramatically, it demonstrates an ability to train itself.
However, the learned optimizer we develop is not without limitations. Below we summarize areas for future work.
Generalization: While it performs well for tasks like those in TaskSet, we do not yet fully understand its outer-generalization capabilities. It 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 5 .
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 as is the case with BERT [ 39 ] .
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 efficient, is memory inefficient requiring > 5 × >\!5\times more storage per-parameter than Adam. We do not believe this poses a fundamental barrier, but modifications similar to those in Shazeer and Stern [40] will be needed to train larger models.
Outer-training task distribution: We have shown that training on larger tasksets leads to better generalization. We have not studied which tasks should be included so as to aid outer-generalization.
## Acknowledgements
We would like to thank Alex Alemi, Eric Jang, Diogo Moitinho de Almeida, Timothy Nguyen, Alec Radford, Ruoxi Sun, Paul Vicol and Wojciech Zaremba for discussion related to this work as well as the Brain Team for providing a supportive research environment. We would also like to thank the authors of Numpy [ 41 , 42 , 43 ] , Seaborn [ 44 ] , and Matplotlib [ 45 ] .
## Broader Impact
Machine learning training and inference is a major energy consumer, and training models likely dominates most individual machine learning researcher’s carbon emissions [ 46 ] . By meta-learning optimizers, we hope to amortize the cost of training ML models, and thus reduce the environmental impact of training a single model. We hope to achieve this both by reducing the need for extensive hyperparameter tuning as models are developed, and making training more efficient for the best hyperparameters.
With this work, and more generally with meta-learning algorithms, we hope to provide researchers access to more performant easier to use optimizers for their problems. Improving technology to do machine learning will accelerate its impact, for better or worse. We believe machine learning technologies will be beneficial to humanity on the whole, and thus by improving the ability to optimize models we are moving towards this goal.
Our optimizer is trained on a mixture of different tasks and datasets. It is our goal to construct an outer-training distribution that learns relevant inductive biases for machine learning tasks of interest. However, we currently do not have good methods to decipher exactly what types of inductive biases a learned optimizer might learn. Thus, future work is needed to explore what kinds of information is absorbed by learned optimizers. In the meantime, care should be taken when employing learned optimizers as the inductive biases of a particular learned optimizer may not be appropriate for an end user’s goals.
In fact, it has been suggested that learned optimizers could be a source of alignment drift in the development of Artificial General Intelligence (AGI) – leading to AGI that does not perform as its creators hoped [ 47 ] . While we do not feel that is an immediate danger from the present work, it is a consideration that should be kept in mind as learned optimizers become increasingly powerful.
## Appendix A Extended Related Work
We categorize progress in learned learned optimizers into three categories: parameterizations, data distributions, and outer-training methodology. We include a review of these previous efforts herein.
### A.1 Parameterizations
Broadly, authors attack the problem of learned optimization at a high level using different parameterizations. We review these different choices below.
#### 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 method. These methods impose a strong inductive bias. All inner-learning must be done via manipulation of existing methods. By enforcing a constrained structure, these method 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-generalization. These methods often additionally make use off 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 et al. [6] explores using a 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. [7] learns a small LSTM network with inputs of current training loss, and predicts a learning rate for SGD. Xu et al. [8] uses an LSTM or an MLP with features computed during training and outputs a scalar which is multiplied by the previous learning rate.
#### A.1.2 Symbolic parameterizations
Existing learning rules are often symbolic in nature. All first order optimizers used today are symbolic and expressible by a small collection of mathematical operations. Bello et al. [48] take inspiration from this and also parameterize optimizers as collections of mathematical primitives. This parameterization can lead to interpretable optimizers. For example, Bello et al. [48] 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 to resemble program synthesis which is notoriously difficult.
#### A.1.3 Continuous parameterizations
The last family of parameterization is based on continuously parameterized function approximations. Older work often makes use of simple / linear combinations of features while more recent work makes use of neural networks. There are a few classes of update rule parameterizations.
Bengio et al. [49] proposes a 7 parameter, and 16 parameter update rule parameterized by mixing various biologically inspired learning signals. Runarsson and Jonsson [50] 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. [9] makes use of an LSTM [ 25 ] possibly with global average pooling to enable concepts like ℓ 2 \ell_{2} gradient clipping to be implemented [ 51 ] . They also explore Neural Turing Machine [ 52 ] like parameterizations specifically designed for low rank memory updates so that it can learn algorithms like LBFGS [ 53 , 54 ] . In contrast, Metz et al. [14] makes use of a per-parameter MLP instead of the LSTM.
One critical design choice when designing neural network parameterized learned optimizers is input features. When training neural network models it is critical that the inputs be similar scales to aide in optimization. Andrychowicz et al. [9] trains on raw gradients which are scaled by either decomposing each scalar into a sign and a magnitude before feeding into an LSTM. Lv et al. [11] uses the gradient and momentum normalized by the rolling average of gradients squared (similar to the rescaling done by Adam). Wichrowska et al. [10] also makes use of momentum except using multiple timescales normalized in a similar way. In addition to these momentum terms, Metz et al. [14] other features such as weight value and additionally applies a normalization based on the second moment of features.
In [ 9 , 11 , 12 ] all methods employ a per-parameter update. The updates are independent of the number of parameters 1 1 1 [ 14 ] uses normalization across parameters so this is not strictly true . While expressive, this is expensive as no computation can be shared. Wichrowska et al. [10] improve upon this by additionally having per-layer, and a global LSTM.