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

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

When designing learned optimizer architectures, there are a number of design decisions to keep in mind. One must balance compute cost of the learned optimizer with expressibility. Often this shifts models to be considerably smaller than those used in supervised learning. Wichrowska et al. [10] , for example makes use of a 8-hidden unit LSTM per-parameter. 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 is often not feasible.

### A.2 Data distribution of tasks

There is no agreed upon standards when defining datasets for training learned optimizers. The community is adhoc, training on what ever dataset is available or what ever best suites the goals (training a particular model, creating more general optimizers). Constructing large distributions of tasks is labor intensive and thus not often done.

Metz et al. [14] draws inspiration from the few shot learning literature and constructs 10 way classification problems sampling from different classes on imagenet.

Wichrowska et al. [10] leverages a large distribution of synthetic tasks. These tasks are designed to represent different types of loss surfaces that might be found in loss surfaces.

TaskSet, [ 16 ] is a dataset of tasks specifically designed for learned optimizer research. We use this dataset throughout our work.

There is a balance between performance of the tasks, and ability to 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 in RL [ 55 , 56 ] .

### A.3 Outer-Optimization methods

The outer-optimization problem consists of finding a particular set of weights, or configuration of a learned optimizer. A number of different strategies have been proposed.

#### A.3.1 Hyper parameter optimization

One of the most common outer-learning methods used is hyper-parameter optimization. In the context of learning optimizers, this can be seen as finding optimizer hyper parameters, e.g. learning rate, over a particular task instance. Numerous methods exist to do this ranging from Bayesian hyper parameter optimization [ 57 ] , to grid search, to genetic algorithms. See Golovin et al. [58] for a more complete description.

The types of outer-learning problems encountered in learned optimizers is different in that the evaluation function is often an expectation over some distribution. Additionally, the amount of outer-parameters is often larger than simply finding a few hyperparameters. Never the less, we include this here to show the similarity to learned optimizer research.

#### A.3.2 Reinforcement learning

Learned optimizers can naturally be cast into a sequential decision process. The state of the system consists of the inner-parameter values, action space is the steps taken, and the reward is achieving a low loss in the future. A number of works have thus taken this viewpoint. Li and Malik [59] , Li and Malik [60] makes use of the guided policy search algorithm. Xu et al. [8] uses PPO Schulman et al. [61] , Daniel et al. [6] make use of Relative Entropy Policy Search [ 62 ] . The exact algorithm used is a function of the underlying parameterization.

#### A.3.3 Neural Architecture Search Style

Instead of learning the policy directly, Bello et al. [48] makes use of reinforcement learning (PPO [ 61 ] to learn a controller which produces the symbolic learned optimizer. This is distinct from § A.3.2 as it does not leverage the sequential nature of the inner problem. Instead, it treats the environment as a bandit problem.

#### A.3.4 Backpropogation / Gradient based

Gradient based methods leverage local perturbations in parameter space. Computing derivatives through leaning procedures has been explored in the context of hyperparameter optimization in [ 63 , 64 , 65 ] . Andrychowicz et al. [9] was the first to make use of gradient based learning for this application.

To train a learned optimizer, ideally, one would compute the derivative of the entire training run with respect to optimizer parameters. This is often referred to as unrolling the entire training procedure into one large graph then running reverse mode automatic differentiation on this. Not only is this often too expensive to do in practice, the resulting loss surface can be poorly conditioned [ 14 ] . As such approximations are often made.

One common family of approximation is based on truncated backpropogation through time. The core idea is to break apart long unrolled computations into shorter sequences and thus not propagating any error back through the entire unrolled computation graph pieces [ 66 , 67 ] . This approximation is used widely in neural network parameterized learned optimizers [ 9 , 10 , 11 , 14 ] . Unlike in language modeling, truncated backprop has been shown to lead to dramatically worse solutions for meta-learning applications [ 19 , 14 ] .

A second family of approximations involve first order gradient calculations. Andrychowicz et al. [9] does use the first order gradient calculation where as subsequent work, [ 10 ] does and computes the full gradient. The trade offs between these two gradient estimators has been discussed in the few shot learning literature in the context of MAML / Reptile [ 68 , 69 ] .

Computing gradients through iterative, non-linear, dynamics has been shown to cause chaotic dynamics. Pearlmutter [70] , Maclaurin et al. [65] showed high sensitivity to learning rate with respect to performance after multiple steps of unrolled optimization. Metz et al. [14] shows this issue for neural network parameterized learned optimizers and proposes a solution based on variational optimization and multiple gradient estimators.

Despite improvements, there are a lot of considerations that must be taken into account for doing gradient based training.

#### A.3.5 Evolutionary Strategies

A alternative way to estimate gradients is to use black box method such as Evolutionary Strategies [ 71 , 20 , 72 , 73 , 74 ] . These methods are memory efficient, requiring no storage of intermediate states, but can suffer from high variance. In the case of learned optimizer optimization, however, these methods can result in lower variance gradient estimators [ 14 ] . Hybrid approaches that leverage both gradients and ES have been such as Guided ES [ 75 ] have also been proposed for meta-optimization. This work leverages one of the simplest forms of evolutionary strategies as described in [ 74 ] which uses a fixed standard deviation.

## Appendix B Outer Optimization Details

In this work, as with Metz et al. [12] , we use asynchronous, batched training. Each task has a different complexity, thus will produce outer-gradient estimates at 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-gradient estimates that are from more than five outer-iterations away from the current weights.

We optimize all models with Adam. We sweep learning rates between 3e-5 and 3e-3 for all experiments. We find the optimal learning rate is very sensitive and changes depending on how long outer-training occurs. We have preliminary explorations into learning rate schedules but have not yet been able to improve on this constant schedule. For all outer-training experiments, we always run more than one random seed. Due to 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.1 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-gradient variance is much larger without this clipping.

When computing outer-gradients, we follow [ 14 ] and compute a outer-loss based on multiple mini batches of data. In our work we use 5. Note inner-training always uses a single minibatch of inner-training data as well as a 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 logged out 10k length unrolls from 100 tasks sampled from the outer-training distribution and saved outer-parameters every hour. While training we monitor performance across all outer-learning rates and all seeds on 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.

Next we fine tune theses models in an unbiased fashion. We explored two methods. First, based on an increasing truncation schedule. We tested linearly increasing truncation length from 300-10k steps over the course of 30k or 10k steps. We find the faster increase, 10k steps, performs best. Second, we explored fine tuning with Persisted Evolutionary Strategies – an unbiased gradient estimator [ 21 ] . 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 a 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, it ensures that there is no decrease in performance early in outer-training. This can be done by simply setting the outer-learning rate to zero for the first 300 outer-iterations.

## Appendix C Learned Optimizer Architecture Details

In this section we describe the detailed learned optimizer architecture. For ease of understanding we opt to show a mix of pseudo-code based on python and textual descriptions as opposed to mathematical expressions. Finally we chose to describe our optimizer as a series of stateless and pure functions for clarity.

We used this architecture for all of our experiments except of Fig 2 b which used an older version of the architecture with additional features which where dropped from the final version.

### C.1 High level structure of the optimizer

The learned optimizer has two main components: a function that maps from some set of inputs, a state, and parameters to some new state and new parameters. ⬇

And a function to produce an initial state from the given inner-parameters.

Parameters are stored as a dictionary of different tensors keyed by name.

For convenience, we also define gradients to be the same type as Params:

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. ⬇

Both from_lstm_size, from_mlp_size are hyper parameters set as part of the learned optimizer to control how much information is sent from the MLP or from the LSTM.

The input to the next_state function consists of inner-gradients, computed on a 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. ⬇

Each task specifies a function that samples parameter initialization, as well as to produce outputs. ⬇

For outer-training each task also includes a second loss function which computes the task’s loss on the outer-validation split of data. Note this uses a different validation set of data than the previous function.

Inner training / application of the learned optimizer looks like: ⬇

Computing the outer objective and outer-gradients from inner-initialization looks like the following: ⬇

### 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 a 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):

⬇

The Text represent names that map to the tensor of the corresponding shape from the parameters. The values are the same shape as the corresponding inner-parameter 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 a helper class. ⬇

We define an initial value which is simply simply zeros: ⬇

To update these values we follow a procedure similar to RMSProp update equations: ⬇

#### C.2.2 LossAccum

This represents how we get loss information into our learned optimizer. Loss values have no pre-determined scale and span many orders of magnitude. As such, we must somehow standardize them so that the inputs are 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.

⬇

The class that manages these states is parameterized by the decay of the rolling window.

⬇

The initial state is simply zeros.

⬇

To compute updates we do rolling mean and variance computations. ⬇

Two additional functions are 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. ⬇

Second, we have a function that weights one loss by a different AccumState. This is eventually used to weight the validation loss accum against the training loss state allowing the learned optimizer to detect overfitting.

⬇

#### C.2.3 RollingClipState

Gradient clipping is 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 a dynamic gradient clipping algorithm. This is our first iteration of this concept and we expect large gains can be obtained with a better scheme.

This algorithm is stateful thus also needs some state container.

⬇

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 are outer-learned with the rest of the learned optimizer.

⬇

The initial states are initialized to 1. ⬇

We provide a normalization function that both updates the RollingClipState and provides clipped gradients. ⬇

### 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 we are learning.

First we construct the optimizer class with the corresponding theta_mod (see bellow). ⬇

Next the rolling features at fixed, hard coded intervals. These could be outer-learned but in this work they are fixed. ⬇