Tasks, stability, architecture, and compute: Training more effective learned optimizers, and using them to train themselves — part 1/5

Cleaned-text code points [0, 15917). Bibliography omitted here; source.md includes it.

# Tasks, stability, architecture, and compute: Training more effective learned optimizers, and using them to train themselves

###### Abstract

Much as replacing hand-designed features with learned functions has revolutionized how we solve perceptual tasks, we believe learned algorithms will transform how we train models. 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.

## 1 Introduction

Much of the success of modern deep learning has been driven by a shift from hand-designed features carefully curated by human experts, to domain-agnostic methods that can learn features from large amounts of data. By leveraging large-scale datasets with flexible models, we are now able to rapidly learn powerful features for new problem settings that often generalize to novel tasks. While learned features outperform hand-designed features on numerous tasks [ 1 , 2 , 3 , 4 ] , we continue to use hand-designed optimization algorithms (such as gradient descent, momentum, and so on) for training models.

These hand-designed update rules benefit from decades of optimization research but still require extensive expert supervision in order to be used effectively in machine learning. For example, they fail to flexibly adapt to new problem settings and require careful tuning of learning rate schedules and momentum timescales for different model architectures and datasets [ 5 ] . In addition, most do not leverage alternative sources of information beyond the gradient, such as the validation loss. By separating the optimization target (training loss) from the broader goal (generalization), classic methods require more careful tuning of regularization and/or data augmentation strategies by the practitioner.

To address these drawbacks, recent work on learned optimizers aims to replace hand-designed optimizers with a parametric optimizer, trained on a set of tasks, that can then be applied more generally. Recent work in this area has focused on either augmenting existing optimizers to adapt their own hyperparameters [ 6 , 7 , 8 ] , or developing more expressive learned optimizers to replace existing optimizers entirely [ 9 , 10 , 11 , 12 , 13 , 14 , 15 ] . These latter models take in problem information (such as the current gradient of the training loss) and iteratively update parameters. However, to date, learned optimizers have proven to be brittle and ineffective at generalizing across diverse sets of problems.

Our work identifies fundamental barriers that have limited progress in learned optimizer research and addresses several of these barriers to train effective optimizers: 1. Computational scale : Training a learned optimizer is costly. When training the optimizer, a single training step requires applying the optimizer to a training task for some number of unrolled steps. This work utilizes massive parallel computing infrastructure to scale the number of unrolled steps an order of magnitude larger than in previous work.

2. Training tasks : Deep learning requires large training datasets. For learned optimizers to be effective, we similarly need a large dataset of optimization tasks on which to train the optimizer. We build off of the TaskSet dataset [ 16 ] and construct a dataset of more than a thousand diverse optimization tasks commonly found in machine learning. We show how this large and diverse task distribution is critical for training optimizers that generalize.

3. Inductive bias of optimizer architecture : The parameterization of the learned optimizer and the task information fed to it both strongly affect performance. In this work, we propose a new hierarchical learned optimizer architecture that incorporates additional task information (such as validation loss), and show that it outperforms previous learned optimizer architectures.

By addressing these barriers, we develop learned optimizers that exceed prior work in scale, robustness, and out of distribution generalization. As a final test, we show that the learned optimizer can be used to train new learned optimizers from scratch (analogous to “self-hosting” compilers [ 17 ] ). 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.

## 2 Preliminaries

Training a learned optimizer is a bilevel optimization problem that contains two loops: an inner loop that applies the optimizer to solve a task, and an outer loop that iteratively updates the parameters of the learned optimizer [ 18 ] . We use the inner- and outer- prefixes throughout to be explicit about which optimization loop we are referring to. That is, the inner-loss refers to a target task’s loss function that we wish to optimize, and the outer-loss refers to a measure of the optimizer’s performance training the target task (inner-task). Correspondingly, we refer to the optimizer parameters as outer-parameters , and the parameters that the optimizer is updating as inner-parameters . Outer-optimization refers to the act of finding outer-parameters that perform well under some outer-loss .

For a given inner-task, we apply the learned optimizer for some number of steps ( unrolling the optimizer). Ideally, we would unroll each target task until some stopping criterion is reached, but this is computationally infeasible for even moderate scale machine learning tasks. Each outer-iteration requires unrolling the optimizer on a target task. Short (truncated) unrolls are more computationally efficient, but suffer from truncation bias [ 19 , 14 ] in that the outer-loss surface computed using truncated unrolls is different (and may have different minima) than the fully unrolled outer-loss.

## 3 Methods: Addressing the three barriers to learned optimizers

### 3.1 Outer-Optimization

To train the optimizer, we minimize an outer-loss that quantifies the performance of the optimizer. This is defined as the mean of the inner-loss computed on the inner- validation set for some number of unrolled steps, averaged over inner-tasks in the outer-training taskset. Although this outer-loss is differentiable, it is costly to compute the outer-gradient (which involves backpropagating through the unrolled optimization). In addition, the outer-loss surface is badly conditioned and extremely non-smooth [ 14 ] , making it difficult to optimize.

We deal with these issues by using derivative-free optimization–specifically, evolutionary strategies (ES) [ 20 ] –to minimize the outer-loss, obviating the need to compute derivatives through the unrolled optimization process. Previous work has used unrolled derivatives [ 9 , 10 , 14 ] , and was thus limited to short numbers of unrolled steps (e.g. 20 in Andrychowicz et al. [9] and starting at 50 in Metz et al. [14] ). Using ES, we are able to use considerably longer unrolls. Initial unroll lengths were chosen to balance communication cost between parallel workers (when updating optimizer parameters) with the computational cost of unrolling on individual workers (when estimating the local gradient with ES). We start outer-training by sampling unroll steps uniformly from 240-360 steps. When performance saturates, we continue training with Persistent Evolotionary Strategies (PES) [ 21 ] . PES provides an unbiased estimate of the outer-gradient over the entire inner task, but at the cost of higher variance gradients.

ES and PES have an additional benefit, in that optimizing with ES smooths the underlying loss function. This smoothing helps stabilize outer-training [ 14 ] . We set the standard deviation of the Gaussian distribution used by the ES algorithm (which also controls how much the outer-loss is smoothed) to 0.01. To deal with the high variance of the ES estimate of the gradient, we use antithetic sampling and train in in parallel using 1024 multi-core CPU workers. While using more workers increases training speed, we find 1024 to be the point where performance gains become sub-linear. For more details see Appendix B .

### 3.2 Task distributions

To train the optimizer, we need to define a set of inner-tasks to use for training. The choice of training tasks is critically important for two reasons: it determines the ability of the optimizer to outer-generalize (i.e. the learned optimizer’s performance on new tasks), and it determines the computational complexity of outer-training. For improved outer-generalization, we would like our inner-problems to be representative of tasks we care about. In this work, these are state-of-the-art machine learning models such as ResNets [ 22 ] or Transformers [ 23 ] . Unfortunately, directly utilizing these large scale models is computationally infeasible, therefore we outer-train on proxy tasks for speed [ 24 ] .

In order to outer-train a learned optimizer capable of generalizing to new optimization tasks, we utilize an outer-training task set consisting of around 6,000 tasks designed after Metz et al. [16] . These tasks include RNNs [ 25 , 26 ] , CNNs [ 27 ] , masked auto regressive flows [ 28 ] , fully connected networks, language modeling, variational autoencoders [ 29 ] , simple 2D test functions, quadratic bowls, and more. For tasks that require them, we additionally sample a dataset, batch size, network architecture, and initialization scheme. To keep outer-training efficient, we ensure that all tasks take less than 100 milliseconds per-training step. For each task that makes use of a dataset, we create four splits of the data to prevent leakage: training data, which we compute gradients on and use to update the inner-parameters; inner-validation data, which is used to compute validation losses used by the learned optimizer; outer-validation data, which is used to update the weights of the learned optimizer; and test data, which is used to test an already trained learned optimizer. Because loss values vary in magnitude, when outer-training we normalize these outer-loss values by the best loss achieved by a baseline optimizer and the initial loss value. Note this normalization is not used during inner-training.

### 3.3 Optimizer architecture

Designing a learned optimizer architecture requires balancing computational efficiency and expressivity. Past work in learned optimizers has shown that incorporating inductive biases based on existing optimization techniques such as momentum or second moment accumulation leads to better performance [ 10 , 14 ] . The optimizer we use in this work consists of a hierarchical optimizer similar to [ 10 ] (Figure 1 ). A per-tensor LSTM is run on features computed over each parameter tensor. This LSTM then forwards information to the other tensors’ LSTMs as well as to a per-parameter feedforward neural network. The per-parameter feedforward network additionally takes in information about gradients and parameter value, and outputs parameter updates. Additional outputs are aggregated and fed back into the per-tensor network. This information routing allows for communication across all components.

For per-parameter features we leverage effective inductive biases from hand-designed optimizers, and use a variety of features including the gradient, the parameter value, and momentum-like running averages of both. All features are normalized, in a fashion similar to that in RMSProp [ 30 ] , or relative to the norm across the full tensor. For per-tensor features we use a variety of features built from summary statistics computed from the current tensor, the tensor’s gradient, and running average features such as momentum and second moments. We also include information about the tensor’s rank and shape. We also feed global features into the per-tensor LSTM, such as training loss and validation loss, normalized so as to have a relatively consistent scale across tasks. To compute a weight update, the per-parameter MLP outputs two values, ( a , b ) (a,b) , which are used to update inner-parameters: w t + 1 = w t + exp ⁡ ( a ) ​ b w^{t+1}=w^{t}+\exp(a)b . See Appendix C for many more details.

## 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: a component-wise LSTM optimizer from Andrychowicz et al. [9] (L2LBGDGD), a modification of this LSTM with the decomposed direction and magnitude output from Metz et al. [14] (LSTM), the fully connected optimizer from Metz et al. [14] (FF), as well as the proposed learned optimizer in this work (§ 3.3 ) (LSTM_FF). As shown in Figure 2 (a), the proposed architecture achieves the lowest outer-training loss and achieves this in 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 of tasks. Previous work only compared a proposed learned optimizer against hand-designed (baseline) optimizers.

Next, we explored how increasing the number of inner-tasks used when training an optimizer affects final performance. To do this, we randomly sampled subsets of tasks from the full task set, while evaluating performance on a common held-out set of tasks. Figure 2 (b) shows that increasing the number of tasks leads to large improvements.

### 4.2 Comparisons with hand-designed optimizers

We compare against three baseline optimizers: AdamLR , which is the Adam optimizer [ 31 ] with a tuned learning rate. Adam8p , which is a version of the Adam optimizer with eight tunable hyperparameters: learning rate, β 1 \beta_{1} , β 2 \beta_{2} , and ϵ \epsilon , plus additional ℓ 1 \ell_{1} and ℓ 2 \ell_{2} regularization penalties, and a learning rate schedule parameterized with a linear decay and exponential decay. See Appendix D for more details. Our final baseline optimizer, called opt_list , consists of the NAdam optimizer [ 32 ] with “AdamW” style weight decay [ 33 ] , cosine learning rate schedules [ 34 ] and learning rate warm up (See Metz et al. [16] for more info). Instead of tuning these with some search procedure, however, we draw them from a sorted list of hyperparameters provided by [ 16 ] for increased sample efficiency.

Evaluation of optimizers, let alone learned optimizers, is difficult due to different tasks of interest, hyperparameter search strategies, and compute budgets [ 5 , 35 ] . 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 an “off-the-shelf” 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 a new machine learning model or dataset. Second, we consider a “finely tuned” limit, where a practitioner has a large compute budget devoted to tuning hyperparameters of a traditional optimizer for a particular problem of interest.