Source-linked AI summary
Learned Optimizers that Scale and Generalize
Olga Wichrowska, Niru Maheswaranathan, Matthew W. Hoffman, Sergio Gomez Colmenarejo, Misha Denil, Nando de Freitas, Jascha Sohl-Dickstein
TL;DR
Learned optimizers face limited generalization and high memory and computation costs, restricting their use on larger and new optimization problems. This paper introduces a hierarchical RNN optimizer trained on diverse small tasks with optimization-informed features, and reports generalization to large ImageNet models for early training. Its wall-clock performance remains worse than simpler optimizers for small minibatches, with the gap decreasing as batch size increases.
Problem
Learned optimizers have been limited by poor generalization to new problems and high memory and computation costs when scaling.
Method
The paper combines a hierarchical RNN matched to target parameters with optimization-informed features and an ensemble of small, diverse meta-training tasks.
Results
The optimizer generalizes from small meta-training problems to early training on large ResNet and Inception models on ImageNet.
Takeaways & Limitations
The results support scaling and generalization of RNN-based optimizers to new large problems, while wall-clock performance lags behind simpler optimizers.
Takeaways & Limitations
For small minibatches, the learned optimizer significantly underperforms ADAM and RMSProp in wall-clock time, although the overhead can decrease with larger minibatches.
Abstract
from arXiv · showhide
Learning to learn has emerged as an important direction for achieving artificial intelligence. Two of the primary barriers to its adoption are an inability to scale to larger problems and a limited ability to generalize to new tasks. We introduce a learned gradient descent optimizer that generalizes well to new tasks, and which has significantly reduced memory and computation overhead. We achieve this by introducing a novel hierarchical RNN architecture, with minimal per-parameter overhead, augmented with additional architectural features that mirror the known structure of optimization tasks. We also develop a meta-training ensemble of small, diverse optimization tasks capturing common properties of loss landscapes. The optimizer learns to outperform RMSProp/ADAM on problems in this corpus. More importantly, it performs comparably or better when applied to small convolutional neural networks, despite seeing no neural networks in its meta-training set. Finally, it generalizes to train Inception V3 and ResNet V2 architectures on the ImageNet dataset for thousands of steps, optimization problems that are of a vastly different scale than those it was trained on. We release an open source implementation of the meta-training algorithm.
1. Introduction
The paper targets learned optimizers that have struggled to generalize to new problems and scale efficiently. It combines diverse small-task meta-training with a hierarchical RNN and optimization-informed features to address these barriers.
- Learned RNN optimizers previously failed to generalize to new problems, sustain progress over long runs, and scale because of memory and computation costs.These limitations motivate the paper’s focus on both generalization and efficiency.
- The meta-training set is an ensemble of small tasks with diverse loss landscapes.The tasks are designed to support learning across varied optimization conditions.
- The hierarchical RNN reduces memory and compute overhead while capturing inter-parameter dependencies.Its architecture is designed to remain efficient while representing relationships among parameters.
- Optimization-informed features include adaptive input and output scaling, multiple momentum timescales, and a combination of Nesterov momentum with RNN attention.These features incorporate strategies motivated by successful hand-designed optimizers and recurrent networks.
- The meta-optimization pipeline uses an objective encouraging exact convergence and samples training horizons from a heavy-tailed distribution.These changes target the optimizer’s behavior across optimization durations.
2. Related work
Related work develops meta-learning methods for learning update rules, optimizers, architectures, and reinforcement-learning procedures. Earlier learned optimizers established transfer across some problem dimensions but remained focused on related task classes.
- Meta-learning research includes systems that modify their own weights and end-to-end differentiable approaches for learning update strategies.These lines of work frame learning itself as an object of optimization.
- Bengio et al. learned parameterized local neural-network update rules, producing an algorithm in the form of a local update rule.Runarsson and Jonsson extended this approach to more complex update models.
- Andrychowicz et al. learned optimizers that integrate gradient observations over time and can apply component-wise across problems of different dimensionality.Their transfer experiments focused on problems of the same class, despite considering different datasets and model structures.
- Li and Malik learned optimizer meta-parameters with policy search, framing learning to optimize by reinforcement learning.
- Other work applies meta-learning to neural-architecture generation, few-shot learning, and reinforcement-learning tasks.These approaches span controller RNNs for architecture descriptions, short-horizon few-shot tasks, and bandit or cognitive-science examples.
3. Architecture
The paper uses a hierarchical RNN learned optimizer whose shared architecture adapts to target problems while reducing per-parameter overhead. Optimization-inspired features provide gradient-history information and dynamic update control intended to improve generalization.
- Hierarchical architecture: A hierarchical RNN optimizer shares meta-parameters across target problems while adapting its structure to each problem’s parameters.It receives gradients and gradient-derived quantities, then outputs parameter updates.
- Hierarchical architecture: Small Parameter RNNs handle individual parameters, Tensor RNNs capture inter-parameter dependencies, and a Global RNN captures inter-tensor dependencies.Tensor and Global RNN states are communicated through averaged hidden states and bias inputs.
- Hierarchical architecture: The hierarchy enables low per-parameter computational cost by keeping Parameter RNNs small while larger Tensor and Global RNNs retain problem-level information.The Tensor and Global RNNs provide coordination across parameters and tensors.
- Optimization-inspired features: The optimizer incorporates dynamic scaling, multi-timescale momentum, and an attention mechanism combining learned updates with Nesterov-style extrapolation.These features expose gradient magnitude, change, noise, and curvature-related information to the learned optimizer.
- Optimization-inspired features: Update outputs are decomposed into directions and step lengths, supporting parameter-scale invariance and dynamically adjusted learning rates.The optimizer controls step length multiplicatively rather than outputting it directly.
- Scaling considerations: The learned optimizer’s memory cost is O(NP + NP KP + NT KT + KG), typically reducing to O(NP + NP KP).Keeping the Parameter RNN latent size KP small keeps memory overhead small.
4. Meta-training
The optimizer is meta-trained on a diverse ensemble of small tasks designed to capture common loss-landscape and gradient properties. Its meta-objective and training procedure are designed to encourage precise convergence and generalization to long optimization runs.
- Meta-training setup: The meta-training process trains an RNN optimizer with shared meta-parameters across an ensemble of target optimization tasks.Meta-training uses a standard optimizer to update the learned optimizer’s meta-parameters.
- Meta-training task corpus: The ensemble combines toy landscape pathologies, convex losses, logistic regression, minibatch and noisy objectives, sparse gradients, parameter-scale changes, and varying steepness profiles.The task corpus includes two-dimensional benchmark functions, quadratic bowls, logistic regression, minibatch variants, noise, slow optimization, sparse gradients, rescaling, and monotonic loss transformations.
- Meta-training task corpus: The meta-training corpus uses small problems to capture commonly encountered optimization properties while avoiding the memory demands of training directly on large real-world problems.The task ensemble was chosen to address prior failures to generalize beyond the meta-training problem.
- Meta-objective: The meta-objective averages log loss across training problems, encouraging exact convergence and dynamic learning-rate adjustment from gradient history.Small loss values receive greater weight under the logarithm while still providing a signal at every optimization step.
- Meta-optimization: Heavy-tailed sampling of partial unrollings and optimization steps is used to encourage generalization to long training runs.Meta-learning gradients are computed by backpropagation through partial optimization unrolls, including second derivatives.
- Meta-optimization: The optimizers were meta-trained for at least 40M meta-iterations using asynchronous RMSProp across 1000 workers with a learning rate of 10^-6.Each meta-iteration samples a task, runs the learned optimizer, computes a meta-gradient, and updates the meta-parameters.
5. Experiments
Experiments show that the learned optimizer matches or outperforms standard optimizers on training-set tasks and generalizes to unseen network types and larger models. It is robust to learning-rate initialization, but has wall-clock and long-run limitations.
- Generalization to new problem types: The optimizer remains comparable to ADAM and RMSProp on MNIST networks with ReLU activations and minibatches of size 64, including iteration counts not seen during meta-training.The test problems include a two-layer fully connected network and a two-layer convolutional network.
- Performance on training set problems: The learned optimizer matches or outperforms ADAM and RMSProp on problem types from the meta-training set.Figure 3 reports this result on sample problems in the training corpus.
- Generalization to new problem types: It performs comparably to ADAM, RMSProp, and SGD with momentum on small convolutional and fully connected networks despite seeing neither layer type during meta-training.The comparison covers both final loss and iterations to convergence.
- Generalization to large networks: The learned optimizer stably trains Inception V3 and ResNet V2 for the first 10K to 20K steps with performance similar to traditional optimizers tuned for each problem.These experiments use networks substantially larger and different from the meta-training tasks.
- Generalization to large networks: Later in training, the learned optimizer stops making effective progress and the Inception V3 loss approaches approximately 6.5.The paper identifies addressing this long-run behavior as future work.
- Learning-rate robustness: The optimizer is more robust to initial learning rate than other optimizers on a randomly generated quadratic problem.Figure 5 compares training curves across different learning-rate initializations.
- Wall-clock comparison: For small minibatches, the learned optimizer significantly underperforms ADAM and RMSProp in wall-clock time.The paper reports that constant minibatch overhead can become smaller as minibatch size increases.
6. Conclusion
RNN-based optimizers meta-trained on small problems can scale to early training on large ResNet and Inception models on ImageNet. Their wall-clock disadvantage versus simpler optimizers decreases as batch size increases.
- RNN-based optimizers meta-trained on small problems scale to early training on large ResNet and Inception models on ImageNet.
- As batch size increases, the Learned optimizer’s total computation time approaches ADAM.
- The hierarchical architecture reduces memory overhead and enables communication across parameters.
- The meta-training ensemble uses small optimization problems capturing common and diverse properties of loss landscapes.
B. Additional details of RNN architecture
The hierarchical RNN uses a target-problem-matched architecture while sharing meta-parameters across problems, allowing learned updates with low per-parameter overhead.
- A meta-trainable linear projection from average rescaled gradients reduces information stored in the Parameter RNN hidden state.
C.1. Heavy-tailed distribution over training steps
Meta-training varies the lengths of optimization unrolls and training runs, while the Inception V3 experiment used a newer optimizer codebase.
- The total number of unrolls is drawn from an exponential distribution with scale 50 plus a constant offset of 1.
- The number of training iterations within each unroll is drawn from an exponential distribution with scale 200 and a constant offset of 50.
- The Inception V3 experiment used a slightly newer version of the learned optimizer codebase.
D.1. Parameter noise during training
Parameter noise is added during training to keep meta-loss informative over long unrolls and make the learned optimizer robust to noise. The update design also changes gradient normalization and removes attention-direction normalization.
- D.1. Parameter noise during training: Small Gaussian noise is added to parameters during each optimization step because small meta-training problems can become nearly optimized early.
- D.1. Parameter noise during training: The noise effectively moves the loss landscape and provides a more informative learning signal after many unrolls.
- D.1. Parameter noise during training: The noise scale α is drawn from a log uniform distribution between 10^-10 and 10^-2 for each problem.
- D.1. Parameter noise during training: The average gradient is normalized using a running estimate of the root-mean-square magnitude of the average gradient.
- D.1. Parameter noise during training: The fastest-timescale average gradient is defined as the raw gradient, and the average gradient is normalized from the immediately faster timescale.
- D.1. Parameter noise during training: The parameter and attention update directions are no longer forced to be normalized, while updates remain decomposed into a learning rate and a step.
D.4. More stable meta-training hyper-parameters
Meta-loss gradients are asymmetrical and heavy-tailed, which can bias RMSProp and Adam updates. The meta-optimizer RMSProp therefore uses a higher mean-square-gradient momentum term, γ = 0.999 instead of 0.9.
- Asymmetrical, heavy-tailed meta-loss gradients can bias RMSProp and Adam updates.Both optimizers underweight extremely rare, extremely large gradients.
- RMSProp and Adam underweight contributions from extremely rare, extremely large gradients.
- γ = 0.999 replaces 0.9 for the meta-optimizer RMSProp's mean-square-gradient momentum term.This change is intended to reduce the bias caused by the gradient distribution.