[Paper Review] A Generalizable Approach to Learning Optimizers

Diogo Almeida, Clemens Winter+2 more|arXiv (Cornell University)|Jun 2, 2021

Advanced Neural Network Applications53 references4 citations

TL;DR

This paper proposes Learned Hyperparameter Optimizers (LHOPTs), a generalizable approach that learns to update optimizer hyperparameters—such as learning rates—using an LSTM controller, unitless features, relative actions, and a novel reward function. The method achieves 2x speedups on ImageNet and 2.5x speedups on large-scale language modeling, generalizing across tasks with over 5 orders of magnitude difference in compute, outperforming tuned Adam and baseline optimizers.

ABSTRACT

A core issue with learning to optimize neural networks has been the lack of generalization to real world problems. To address this, we describe a system designed from a generalization-first perspective, learning to update optimizer hyperparameters instead of model parameters directly using novel features, actions, and a reward function. This system outperforms Adam at all neural network tasks including on modalities not seen during training. We achieve 2x speedups on ImageNet, and a 2.5x speedup on a language modeling task using over 5 orders of magnitude more compute than the training tasks.

Motivation & Objective

• To address the lack of generalization in existing learned optimizers, which fail on real-world tasks outside their training distribution.
• To develop a framework that prioritizes generalization over task-specific performance by learning hyperparameter updates instead of model parameter updates.
• To decouple the policy update frequency from the inner optimization step count, enabling broader applicability.
• To design features, actions, and reward functions that enhance robustness and transferability across diverse neural network architectures and modalities.
• To demonstrate that learned optimizers can generalize to real-world problems with orders-of-magnitude differences in compute and model scale.

Proposed method

• The system uses an LSTM controller to update optimizer hyperparameters at a user-defined frequency, independent of inner optimization steps.
• It employs unitless, normalized features derived from training statistics (e.g., gradient norms, loss trends) to improve robustness across tasks.
• Actions are defined as relative changes to hyperparameters (e.g., multiplicative scaling), ensuring invariance to absolute values.
• A novel reward function is designed to optimize for final model performance, not intermediate metrics, promoting generalization.
• The policy is trained via reinforcement learning on a diverse set of small-scale tasks, including MNIST, CIFAR-10, and small language modeling tasks.
• Hyperparameter schedules learned on small tasks are transferred to large-scale models without retraining, enabling zero-shot generalization.

Experimental results

Research questions

• RQ1 Can a learned optimizer generalize to tasks significantly larger and different from its training distribution?
• RQ2 Does learning hyperparameter updates instead of parameter updates improve generalization and computational efficiency?
• RQ3 How effective is the proposed reward function in promoting transferability across diverse neural network architectures and modalities?
• RQ4 To what extent do unitless features and relative actions enhance robustness to distributional shift?
• RQ5 Can a single hyperparameter schedule learned on a small task achieve substantial speedups on large-scale, well-tuned models?

Key findings

• LHOPTs achieve a 2x speedup on ImageNet training compared to tuned AdamW, with no hyperparameter tuning on the target task.
• On a large-scale language modeling task, LHOPTs deliver a 2.5x speedup despite being trained on tasks averaging only 170 GPU-seconds.
• The method generalizes across more than 5 orders of magnitude in compute, from small tasks to models requiring up to 300 GPU-days.
• LHOPTs outperform both MLPerf baselines and tuned Adam on speech recognition and neural collaborative filtering, tasks not seen during training.
• Computational overhead is minimal—under 20% on Transformers and less than 2% on ResNets—due to efficient feature computation and constant-size policy network.
• The approach generalizes even when trained on only 28x28 images and without exposure to attention mechanisms, suggesting task diversity is less critical than previously thought.