Understanding and correcting pathologies in the training of learned optimizers

Luke Metz; Niru Maheswaranathan; Jeremy Nixon; C. Daniel Freeman; Jascha Sohl-Dickstein

Source-linked AI summary

Understanding and correcting pathologies in the training of learned optimizers

Luke Metz, Niru Maheswaranathan, Jeremy Nixon, C. Daniel Freeman, Jascha Sohl-Dickstein

arXiv:1810.10180v5Oct 24, 2018cs.NEstat.ML

TL;DR

Learned optimizers are difficult to train because truncated unrolling creates biased gradients while long unrolls can produce exploding gradients. The paper combines two unbiased gradient estimators for a variational objective and dynamically reweights them, enabling a learned optimizer to outperform hand-designed optimizers in wall-clock time and test loss on targeted convolutional-network tasks.

Problem
Learned optimizers face biased gradients under short truncations and exploding gradients under long unrolls, limiting their practical training.
Method
The paper uses a variational objective with reparameterization and evolutionary-strategy gradient estimators whose contributions are dynamically reweighted.
Results
The learned optimizer achieves better test loss and faster wall-clock time than SGD+Momentum, RMSProp, and ADAM on the targeted task distribution.
Takeaways & Limitations
Stable outer-training enables learned optimizers to be trained on longer unrolls and to compete with hand-designed optimizers in wall-clock performance.
Takeaways & Limitations
The study focuses on a restricted family of tasks, leaving the limits of learned-optimizer generalization across tasks for future work.

Abstract

from arXiv · show

Deep learning has shown that learned functions can dramatically outperform hand-designed functions on perceptual tasks. Analogously, this suggests that learned optimizers may similarly outperform current hand-designed optimizers, especially for specific problems. However, learned optimizers are notoriously difficult to train and have yet to demonstrate wall-clock speedups over hand-designed optimizers, and thus are rarely used in practice. Typically, learned optimizers are trained by truncated backpropagation through an unrolled optimization process resulting in gradients that are either strongly biased (for short truncations) or have exploding norm (for long truncations). In this work we propose a training scheme which overcomes both of these difficulties, by dynamically weighting two unbiased gradient estimators for a variational loss on optimizer performance, allowing us to train neural networks to perform optimization of a specific task faster than tuned first-order methods. We demonstrate these results on problems where our learned optimizer trains convolutional networks faster in wall-clock time compared to tuned first-order methods and with an improvement in test loss.

1. Introduction

Learned optimizers promise task-specific improvements but are difficult to train because existing approaches scale poorly or suffer from biased and exploding gradients. The paper addresses these pathologies with dynamically combined unbiased estimators and demonstrates faster wall-clock optimization with better test loss on targeted convolutional-network tasks.

Motivation: Learned optimizers aim to exploit task-specific structure for performance improvements over general optimizers.They complement hand-designed methods that are tuned to specific problems through hyperparameter search.
Training challenges: Black-box training methods scale poorly with the number of optimizer parameters.The paper identifies evolutionary algorithms, random search, reinforcement learning, and Bayesian optimization as examples.
Training challenges: First-order training requires costly backpropagation through iteratively applied learned updates, motivating truncated backpropagation.Truncation permits multiple outer-gradient computations over shorter segments but introduces the central training difficulty studied here.
Proposed approach: The proposed method dynamically combines reparameterization and evolutionary-strategy gradient estimators to avoid exploding gradients and support longer unrolls.Both estimators are unbiased for the variational objective.
Results: 10k inner-parameter updates can be outer-trained with a simpler small MLP optimizer on more complex inner problems than prior work.The method avoids extensive normalization and annealing tricks previously used for stability.
Results: The learned optimizer achieves better test loss and faster wall-clock time than SGD+Momentum, RMSProp, and ADAM on the targeted task distribution.The comparison is reported for small convolutional networks on image classification.

2. Unrolled optimization for learning optimizers

Learned optimizer training is framed as a bilevel problem requiring gradients through an unrolled inner optimization. Long unrolls can produce exploding, highly irregular outer landscapes, while short truncations can yield biased updates.

Problem framework: The inner problem repeatedly updates target weights, while the outer problem adjusts learned-optimizer parameters to minimize an average training or validation loss.The learned update rule maps weights and optimization information at iteration t to iteration t+1.
Unrolled optimization: Unrolled optimization computes outer gradients by backpropagating through repeated optimizer applications, whereas truncation computes cheaper gradients over shorter windows.The truncation length trades computational cost against gradient quality.
Exponential explosion of gradients with increased sequence length: Products of Hessian-dependent factors can make outer-gradient norms grow exponentially with the number of gradient-descent steps.For quadratic losses, the outer-gradient is a degree-T matrix polynomial and can grow exponentially when the relevant eigenvalue exceeds 1.
Exponential explosion of gradients with increased sequence length: Increasing unrolling steps can make the outer loss nearly discontinuous around momentum values that switch trajectories between local minima.The toy system moves between minima, oscillates at larger momentum, and develops extremely large gradient norms near transition points.
Exponential explosion of gradients with increased sequence length: Small learning-rate changes in a two-layer MLP produce increasingly divergent parameter trajectories as training proceeds.The experiment varies Adam learning rates from 0.1469 to 0.1484 while holding randomness fixed.
Truncation bias: Severely truncated backpropagation can produce biased outer updates: in one Adam learning-rate experiment, truncated gradients were anti-correlated with the true outer gradient.The resulting learning rate decreased despite initialization near the optimum, increasing outer loss.

3. Towards stable training of learned optimizers

The method smooths the outer objective and combines two unbiased gradient estimators with inverse-variance weighting. This addresses exploding gradients and changing estimator variance while retaining lower-bias estimates from longer truncations.

The method optimizes a smoothed outer-loss rather than the original high-frequency outer objective.The smoothing variance is fixed at 0.01 in all experiments.
Two unbiased estimators are constructed for the smoothed loss using reparameterization and evolutionary-strategy gradients.The estimators are denoted g_rp and g_es.
Inverse-variance weighting merges the estimators, addressing variance differences that can reach 20 orders of magnitude.The merged estimator has variance at most the lower of the two component variances.
The reparameterization estimator’s variance grows during outer-training, whereas the ES estimator’s variance remains constant.The reparameterization estimator changes from approximately five orders of magnitude better to approximately twenty orders of magnitude worse than ES.
Antithetic sampling and data parallelism are used to reduce estimator variance and compute multiple estimator samples.Each normal draw is used with both positive and negative signs; one sample costs two forward and two backward unrolled passes.
Longer truncations provide lower-bias gradient estimates and avoid exploding gradients, while shorter truncations can suffice early in outer-training.Longer truncations are computationally expensive, motivating the full outer-training procedure’s adaptive use of truncation lengths.

4. Experiments

The experiments train a small learned optimizer on convolutional image-classification tasks and compare it with tuned first-order baselines. It achieves faster or lower-loss optimization on held-out tasks, shows some out-of-distribution generalization, and depends on the combined estimator, truncation curriculum, and momentum features.

Experimental setup: The optimizer is a one-hidden-layer MLP with approximately 1k parameters, applied independently to each target parameter.Its outputs define an unnormalized update direction and per-parameter log learning rate.
Experimental setup: Training uses three-layer convolutional networks on 32x32x3 image-classification tasks, with 10,000 inner iterations and ImageNet classes split into training and test sets.The optimizer is outer-trained on sampled class subsets and evaluated on held-out classes.
Baselines: The learned optimizer is compared with learning-rate-tuned Adam, RMSProp, and Momentum, plus an eight-parameter Adam baseline optimized by random search.The additional baseline tunes Adam parameters, decay coefficients, and regularization.
Held-out performance: When trained against validation loss, the learned optimizer reaches a lower test loss within the measured interval, whereas training-loss optimization converges faster on training loss but performs poorly on test loss.The training-objective behavior is attributed to never observing validation loss during outer-training.
Generalization: Across 100 outer-test tasks, the learned optimizer outperforms the tuned Adam baseline on the majority of tasks and generalizes to several dissimilar architectures and datasets.The out-of-distribution example uses a six-layer convolutional network trained on MNIST.
Ablations: Ablations show that the combined gradient estimator, increasing unroll schedule, and momentum inputs are important for performance.Short fixed unrolls converge poorly, long fixed unrolls have high variance, and removing momentum prevents convergence.

5. Discussion

The discussion frames exploding gradients and truncation bias as central obstacles to learned-optimizer training, addressed by a combined estimator and truncation curriculum. The demonstrated gains apply to a restricted task family, leaving cross-task generalization as an open question.

Training pathologies: Truncated optimization creates biased gradients, while long unrolls can produce exploding gradients during learned-optimizer training.These are identified as the two principal training pathologies.
Proposed remedy: The method combines reparameterization and evolutionary-strategy estimators for a variational outer-objective and uses a truncation-step curriculum.The combined estimator is intended to address both bias and exploding-gradient difficulties.
Empirical consequence: The learned optimizer is reported to be faster in wall-clock time than existing optimizers on the demonstrated tasks.The discussion presents this as the principal empirical consequence of the training approach.
Scope: The experiments focus on a restricted family of tasks, and future work will examine when learned optimizers generalize across tasks.The stated scope boundary concerns the task distribution rather than the optimizer architecture alone.
Future directions: The authors also suggest that the outer-gradient estimator may apply to other long-time-dependence recurrent problems beyond meta-learning.They name neural turning machines and neural GPUs as examples.

A. Derivation of the unrolled gradient

The unrolled-gradient derivation expresses how an outer loss after T inner updates changes with a learning-rate parameter. It reveals that repeated Hessian products create the exploding-gradient pathology and gives a one-step interpretation based on adjacent-gradient correlation.

Setup: For gradient descent with learning rate θ, the derivation seeks the gradient of the final inner loss after T steps with respect to θ.The inner parameters are w, the loss is ℓ, and g(t) and H(t) denote the gradient and Hessian.
Unrolled gradient: The recursive sensitivity of the final parameters to θ expands into products involving the Hessian at every inner iteration.This recursion is the source of the long-horizon dependency in the unrolled gradient.
Pathology: These repeated Hessian products explain where exploding outer-gradients arise in unrolled optimization.The derivation explicitly connects the recursive parameter sensitivity to the pathology.
One-step intuition: With one unrolled step, the gradient is the negative inner product between the current and previous gradients.Positive correlation implies increasing the learning rate under the corresponding update direction.
Illustration: Figure 9 illustrates how varying learning rate and momentum can make outer-optimization landscapes increasingly pathological as inner step count grows.The figure uses a toy one-dimensional inner loss and plots the final parameter as a function of training steps.

C. Outer-Training Algorithm

The outer-training algorithm samples tasks and truncation lengths, estimates gradients using antithetic evolutionary-strategy perturbations alongside the other estimator, and updates the optimizer parameters with SGD. It combines estimator samples using empirical variance and mean information.

Outer-training loop: Outer-training initializes optimizer parameters, samples a task and inner-loop initialization, then repeatedly processes truncations of the inner optimization.Each worker performs these operations in parallel during outer-training.
Truncation sampling: For each truncation, the algorithm samples an evolutionary-strategy perturbation and chooses the truncation length according to the current outer-training iteration.It then generates positive and negative inner trajectories under the perturbed optimizer parameters.
Gradient estimation: The algorithm evaluates paired outer objectives on the two trajectories and forms an evolutionary-strategy gradient sample from their difference.The objectives use either training or validation data from the inner problem.
Gradient estimation: Gradient samples from the reparameterization and evolutionary-strategy estimators are stored until a batch is available.The current inner parameters are advanced using the endpoint of one sampled trajectory.
Parameter update: For each estimator, the algorithm computes empirical means and variances, uses them to form a combined gradient, and updates θ with SGD.The update is θ ← θ − αg_combined.

D.1. Architecture

The optimizer operates independently on each parameter using a compact MLP that combines momentum, gradient, weight, and time features. Its training uses normalized features, exponential output scaling, and progressively longer unrolls.

D.1. Architecture: A single-hidden-layer MLP with 32 units updates each parameter independently from momentum, gradient, weight, and transformed iteration features.Momentum uses decay values [0.5, 0.9, 0.99, 0.999, 0.9999], while time uses nine logarithmically spaced timescales from 3 to 300k.
D.1. Architecture: The optimizer uses second-moment normalization for non-time features to preserve directionality, then produces exponentially combined outputs from two learned scalings.The output form is exp(λexpo1)λlino2.
D.1. Architecture: The inner target is a three-layer convolutional network with 32 units per layer, strides [2,2,1], ReLU activations, and ten output classes.The final convolution is average-pooled and followed by a linear projection.
D.1. Architecture: Training increases unrolling from 50 to 10k inner steps over 5k outer iterations, with random length variation of ±20%.Outer parameters use Adam with batch size 128 and learning rates 0.003 for training and 0.0003 for validation objectives.

E. Additional inner loop problem learning curves

Additional learning curves evaluate the learned optimizer on outer-training and outer-validation distributions, with tuned regularized baselines and diverse inner problems. The experiments also examine transfer across architectures, depths, and parameter counts.

E. Additional inner loop problem learning curves: The additional curves compare learned optimizers against Momentum, Adam, RMSProp, and tuned regularized variants over 0–10,000 inner steps.The regularized baselines are hyperparameter-searched over learning rates, schedules, epsilon, and L1/L2 regularization.
E. Additional inner loop problem learning curves: Figures 10 and 11 provide additional outer-validation and outer-training problems for the learning-curve evaluation.These figures extend the reported comparisons beyond the primary curves.
E. Additional inner loop problem learning curves: The study focuses on three-layer convolutional networks trained on ten-class subsets of 32x32 ImageNet, then tests transfer to varied architectures, depths, and parameter counts.The transferred models include roughly 16× more parameters, and the authors describe generalization as promising but imperfect.
E. Additional inner loop problem learning curves: The evaluated inner problems include fully connected MNIST networks, convolutional MNIST networks, and a ten-way 32x32 ImageNet classification task.Architectures range from two to six hidden or convolutional layers and use 32 or 128 units per layer.

G. Inner-loop training speed

The learned optimizer has higher per-batch computation cost than Adam and SGD, especially on GPU, but the paper reports that its performance gains still produce faster wall-clock optimization on the tested models.

G. Inner-loop training speed: The learned optimizer’s computation is linear in the number of model parameters and smaller than the cost of computing gradients.Its bulk computation consists of two batched matrix multiplications of sizes features×32 and 32×2.
G. Inner-loop training speed: 177 batches per second for the learned optimizer versus 278 for Adam and 358 for SGD on an Nvidia Titan X GPU.These measurements use TensorFlow implementations without fused operations and random input data.
G. Inner-loop training speed: Despite slower computation, the learned optimizer remains considerably faster in wall-clock time because its performance gains exceed the execution slowdown.For the paper’s wall-clock figures, step-versus-performance curves are rescaled by steps per second rather than directly timed.

H. Ablation learning curves

The ablation learning curves summarize performance variation across runs for the configurations described in Section 4.5.

H. Ablation learning curves: The plots show median performance with the 25th–75th percentile shaded and individual runs as thinner solid lines.The median is drawn as a thick line bordered in black.

I. Additional Truncation Bias Experiments

Additional experiments confirm that truncation bias depends on the outer-optimizer configuration and persists across settings. The learned optimizer also shows stable performance beyond its outer-training horizon, while remaining below the strongest tuned Adam baseline.

I. Additional Truncation Bias Experiments: Truncated gradients produce significant truncation bias across multiple Adam and SGD configurations and outer-optimizers.The experiments vary learning rates, beta1 usage, and outer-optimizer choice.
I. Additional Truncation Bias Experiments: The batch-normalized task family uses the same meta-training procedure and learned optimizer architecture while targeting validation loss.
I. Additional Truncation Bias Experiments: The learned optimizer outperforms learning-rate-tuned Adam but does not outperform the 8-parameter tuned Adam baseline.The authors are unsure of the gap's source and suspect hyperparameter tuning could improve the result.
I. Additional Truncation Bias Experiments: After outer-training on 10k inner iterations, the learned optimizer maintains consistent test and train loss through 100k inner iterations.The optimizer does not completely minimize training loss, and test performance therefore remains consistent far beyond the outer-training regime.