Tasks, stability, architecture, and compute: Training more effective learned optimizers, and using them to train themselves — part 4/5
Cleaned-text code points [47463, 63307). Bibliography omitted here; source.md includes it.
Then the the loss features for use with both training and validation loss. We use a lower decay constant on the validation as it is updated less frequently (once every 10 steps). Once again these could be outer-learned but in this work we leave them fixed. ⬇
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’. ⬇
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.
⬇
Next we look to the update performed. We split up into two components – first a validation state, and second a training state.
⬇
The validation update simply consists of updating the two components that make use of the loss value. ⬇
The the function applied to the training gradients is much more complicated. Much like the validation features it updates the various rolling statistics including the two loss data structures, the rolling momentum / rms terms, as well as the gradient clipping state. ⬇
Next, a sequence of features related to loss values are computed. ⬇
Next we call into the learned optimizer function in “theta_mod” which does the bulk of the computation. ⬇
We then populate the next state with all the outputs. ⬇
### C.4 “ThetaMod”: The learned optimizer outer-parameters
All of the meta-learned outer-parameters are stored in a class called ThetaMod. Its constructor defines various dimensions of the different Sonnet modules.
⬇
This constructor creates a shared sonnet module for the per layer LSTM. This is constructed here as it is needed in more than one method.
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_mlp_size more values which are fed back into the LSTM.
Third, we construct two linear projections that map from the RNN output to either the per parameter feed forward network, or back to the global features.
Finally, we initialize the two variables used by the rolling gradient clipping terms. We parameterize these variables on a log scale (possibly subtracting 1) so as to roughly match the scaling of the other neural network weights.
⬇
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 as well as various pieces of hidden information for the next iteration. ⬇
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. ⬇
Next we compute the per tensor features. This is a rank two tensor, consisting of the number of tensors, and a feature dimension. ⬇
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. ⬇
We then linearly project the LSTM outputs to be input to the per parameter feed foward network, and unstack them to form a a list of rank 1 features – one per tensor.
⬇
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 are initialized to accumulate results. ⬇
We pass these per parameter features through these features into a per parameter MLP which produces a single output.
We split this output into three pieces: A direction, a magnitude, and features which are be fed back to the per tensor LSTM after being reduced along the parameter dimension. ⬇
A step can then be computed by exponentiating the magnitude and multiplying by a 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.
We then subtract step from the previous parameter value. While it is possible just to predict a new parameter directly, learning deltas is significantly easier to learn and results in much more stable training.
Finally, we collect all the information needed to be returned from the per parameter LSTM and output. ⬇
With the for loop over, we then return the necessary info. Before this, however, we must aggregate information from the per tensor LSTM to be able to pass it to the next layer’s global features after applying a linear projection from the LSTM output and reducing along the tensor dimension.
⬇
Now we will discuss the different features methods. First, computing the global features. ⬇
This function simply returns the concatenation of the Two loss features, and the aggregated features from the per layer lstm.
Next we look to the per tensor features. We return a 2D tensor with the leading dimension a “batch size” being the number of tensors, and the second dimension representing the feature dimension.
⬇
We compute a large number of features. At this point have not extensively ablated each one and we 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.
⬇
We begin by iterating over each tensor.
⬇
First, we compute features based on the rolling momentum values. We compute features that represent the log magnitude, the sign, and the log variance. ⬇
We compute a similar set of features for the second moment features. ⬇
A similar set for the inner-parameters: ⬇
Next we include the magnitude of the weight norm.
As well as the gradient norm:
Next we add a feature if the tensor is a scalar or not. ⬇
The shape of the underlying tensor is also a useful feature. The length of this shape varies though depending on the rank. As such we first pad the tensor to be at least rank 4 and then pass the log shape shifted by -1. This shift is to keep this value roughly zero mean.
⬇
The features from the previous execution’s MLP are included. These are to enable communication back from the MLP to the per layer network.
This dictionary of features is then flattened (adding appropriate dimensions if needed) and appended to the list of inputs. ⬇
Once the stacked_inputs is filled, we zip and stack and concat the values forming a number of layers by number of features tensor. ⬇
The last set of features is the per parameter 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.
⬇
As with the per layer features, we add a series of features. First we add raw features for the gradients, parameter values, momentum, and second moment accumulators, and past weights. ⬇
We then add some derived features such as the log abs parameter value:
Finally we add terms based on the rsqrt of the second moment terms inspired by how Adam normalizes. ⬇
These values are concatenated and normalized by a second moment normalizer computed across the num tensor dimension. We additionally clip the output of this normalizer. This is crucial for stable training as some of these features will produce large magnitudes in some channels.
⬇
Next, we embed the current inner-training step with sinusoids of different frequencies.
⬇
We compute a features based on the number of tensors: ⬇
As well as features about the number of parameters in the current tensor: ⬇
Next, we compute some statistics about the log norm of the weight matrix: ⬇
Next, we tile the input from the LSTM so as to add a number of parameters dimension. ⬇
Finally we concat and return all these features. ⬇
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 outer-parameterized values shown bellow. ⬇
## Appendix D Adam8p
This matches the adam8p optimizer described in [ 16 ] . The eight hyper-parameters are: the learning rate, α \alpha , first and second moment momentum, β 1 \beta_{1} , β 2 \beta_{2} , the numerical stability term, ϵ \epsilon , ℓ 2 \ell_{2} and ℓ 1 \ell_{1} regularization strength, and learning rate schedule constants λ exp_decay \lambda_{\text{exp\_decay}} and λ linear_decay \lambda_{\text{linear\_decay}} .
ϕ ( 0 ) = \displaystyle\phi^{(0)}= problem specified random initialization (1) m ( 0 ) = \displaystyle m^{(0)}= 0 \displaystyle 0 (2) v ( 0 ) = \displaystyle v^{(0)}= 0 \displaystyle 0 (3) g ( t ) = \displaystyle g^{(t)}= d d ϕ ( t ) ( f ( x , ϕ ( t ) ) + ℓ 2 ‖ ϕ ( t ) ‖ 2 2 + ℓ 1 | | ϕ ( t ) | | 1 ) \displaystyle\frac{d}{d\phi^{(t)}}(f(x;\phi^{(t)})+\ell_{2}||\phi^{(t)}||^{2}_{2}+\ell_{1}||\phi^{(t)}||_{1}) (4) m ( t ) = \displaystyle m^{(t)}= β 1 m ( t − 1 ) + g ( t ) ( 1 − β 1 ) \displaystyle\beta_{1}m^{(t-1)}+g^{(t)}(1-\beta_{1}) (5) v ( t ) = \displaystyle v^{(t)}= β 2 v ( t − 1 ) + ( g ( t ) ) 2 ( 1 − β 2 ) \displaystyle\beta_{2}v^{(t-1)}+(g^{(t)})^{2}(1-\beta_{2}) (6) m ^ ( t ) = \displaystyle\hat{m}^{(t)}= m ( t ) 1 − β 1 t + 1 \displaystyle\dfrac{m^{(t)}}{1-\beta_{1}^{t+1}} (7) v ^ ( t ) = \displaystyle\hat{v}^{(t)}= v ( t ) 1 − β 2 t + 1 \displaystyle\frac{v^{(t)}}{1-\beta_{2}^{t+1}} (8) u ( t ) = \displaystyle u^{(t)}= m ^ ( t ) v ^ ( t ) + ϵ \displaystyle\dfrac{\hat{m}^{(t)}}{\sqrt{\hat{v}^{(t)}}+\epsilon} (9) s linear ( t ) = \displaystyle s_{\text{linear}}^{(t)}= max ( 1 − t λ linear_decay , 0 ) \displaystyle\text{max}(1-t\lambda_{\text{linear\_decay}},0) (10) s exp ( t ) = \displaystyle s_{\text{exp}}^{(t)}= exp ( − t λ exp_decay ) \displaystyle\text{exp}(-t\lambda_{\text{exp\_decay}}) (11) ϕ ( t + 1 ) = \displaystyle\phi^{(t+1)}= α s linear ( t ) s exp ( t ) u ( t ) \displaystyle\alpha s_{\text{linear}}^{(t)}s_{\text{exp}}^{(t)}u^{(t)} (12)
We sample learning rate logritmically between 1e-8 and 10, beta1 and beta2 we parametrize as 1 − x 1-x and sample logrithmically between 1e-4 and 1 and 1e-6 and 1 respectively. For learning rate schedules we sample linear decay between 1e-7, 1e-4 logrithmically and exponential decay logrithmically between 1e-3, 1e-6. We sample both ℓ 1 \ell_{1} and ℓ 2 \ell_{2} logrithmcally between 1e-8, 1e1.
## Appendix E Performance Table
We show numerical performance measurements of our learned optimizer and baseline optimizers. Each value is the mean or median over 100 different inner-training tasks.
train test Optimizer mean median mean median learned 0.077325 0.024167 0.113242 0.028928 global adam lr (1 Trial) 0.226807 0.105513 0.261613 0.092532 global adam8p, RS (1 Trial) 0.240924 0.109519 0.257368 0.100959 global nadamw, RS (1 Trial) 0.215176 0.098448 0.229359 0.086703 per task adam lr (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 (1k Trial) 0.039348 0.028254 0.043193 0.029897 per task adam8p (10 Trial) 0.278858 0.132965 0.264282 0.128594 per task adam8p (100 Trial) 0.095372 0.046926 0.107177 0.046672 per task adam8p (1k Trial) 0.044144 0.028692 0.060663 0.028813 per task adam8p + opt_list(2k Trial) 0.018403 0.023350 0.025419 0.025555
## Appendix F Experimental details for figures
### F.1 Outer training of different learned architectures
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 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.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 [ 14 ] we use RMS terms which improve performance over that reported in [ 14 ] by a small amount.
For the L2LBGDGD model we copy Andrychowicz et al. [9] using a two hidden layer, 20 unit GRU [ 26 ] with the same input gradient processing.
For the LSTM model, we make use of a model inspired by [ 9 ] . We make use of a 64 dimension LSTM hidden size to produce two outputs, a direction and a sign which are combined similar to [ 14 ] .
Performance on 100 tasks (each with five random seeds) is recorded over the course of training. We show an 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 we only train with truncation’s of size 240-360. We use a learning rate of 3 ∗ 10 − 3 3*10^{-3} for all models.
### F.3 Large scale transfer 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 3 with a stride of 2. This aggressive down sampling was 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 a (number of output channels, number of bottleneck channels, stride).
At the output we apply batch norm, reduce over the spatial dimensions, and project to the the appropriate number of classes.
Our goal with these experiments are not to seek peak performance. Instead show that our learned optimizer is capable of 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 11 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 10^{-5} and 3 ∗ 10 − 3 3*10^{-3} 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 a new optimizer there are no hyper parameters. We found we must also include the original gradient clipping used for the Adam models (clipping all values between -0.1, 0.1). Without this, our learned optimizer is not able to make any progress. Understanding this is an interesting to us as we 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, we use the training loss for both. This is computed as the mean over the normalized losses computed over a single batch.
## Appendix G Self Optimization Extended
## Appendix H Infrastructure Used
The types of learning systems we discuss here are more complex than the average deep learning training setup. We are training not one, but many different kinds of models on different datasets with wildly different properties. As a 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 a trivial task as it requires inner-training a large number of different models for a large number of inner-iterations.