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

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

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 we 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, a learner, gradient storage, summary aggregators, evaluation chief, evaluation worker. Additionally we make use of a distributed file system [ 76 ] .

##### Workers

Workers are jobs that sample a outer-task, build a tensorflow [ 77 ] 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 we 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-gradient estimate new outer-parameters, as 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 is 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 a single machine is required. This is a separate machine from the learner to control the number of queries required sent to the learner.

##### 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 a 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 a 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-parameters 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 a time and would ideally like to 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 ease 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 diagnose 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 to summary aggregators, and sending these RPC in parallel with computation.

##### Evaluations:

Evaluating a learned optimizer is not nearly as simple as 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 to the meta-training task distribution (which we pick a 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-generalization. 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 an extended period of time. Averaging over the meta-training cluster also doesn’t capture variance from outer-weight to outer-weight.

The evaluation of one task also requires orchestration. We employ a evaluation chief, and evaluation workers for each task set. While its possible to do this type of evaluation offline, we find running online yields a much better workflow and thus speeds up research.

##### Evaluation Chief:

The evaluation chief is responsible for watching directories on a distributed filesystem and to enqueue evaluation tasks. Workers use RPC to request, and to send back results (e.g. learning curves). When a 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 to the chief to be aggregated and written to disk.

### H.2 Computational expense

At this point, training of a learned optimizer is quite expensive. It roughly entails 60K CPU cores for around a month, or 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 is 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 of 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 all of this work was performed on CPU. In tests on GPU (without optimizing) we found performance to be similar. Modern accelerator hardware, TPU, GPU, 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.

## Instructions for reporting errors

We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the methods listed below:

Click the "Report Issue" ( ) button, located in the page header.

Tip: You can select the relevant text first, to include it in your report.

Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability should not be a barrier to accessing research. Thank you for your continued support in championing open access for all.

Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .