Practical tradeoffs between memory, compute, and performance in learned optimizers

By Luke Metz, C. Freeman, James Harrison, Niru Maheswaranathan, Jascha Narain Sohl-Dickstein. 2022

42 citations recorded by the source corpus.

Optimization plays a costly and crucial role in developing machine learning systems. In learned optimizers, the few hyperparameters of commonly used hand-designed optimizers, e.g. Adam or SGD, are replaced with flexible parametric functions. The parameters of these functions are then optimized so that the resulting learned optimizer minimizes a target loss on a chosen class of models. Learned optimizers can both reduce the number of required training steps and improve the final test loss. However, they can be expensive to train, and once trained can be expensive to use due to computational and memory overhead for the optimizer itself. In this work, we identify and quantify the design features governing the memory, compute, and performance trade-offs for many learned and hand-designed optimizers. We further leverage our analysis to construct a learned optimizer that is both faster and more memory efficient than previous work. Our model and training code are open source.

Subfields: Artificial Intelligence, Optimization & Training, ML Systems & Hardware

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Optimization & Training

In ML Systems & Hardware

Citations in the corpus

Learning to Learn with Generative Models of Neural Network Checkpoints (2022, background): “They are typically trained with unrolled optimization (Hochreiter et al., 2001; Younger et al., 2001; Gregor & LeCun, 2010; Andrychowicz et al., 2016; Ravi & Larochelle, 2017; Wichrowska et al., 2017; Lv et al., 2017; Metz et al., 2019; 2022) or reinforcement learning (Li & Malik, 2016; 2017).”
Learning to Optimize: A Primer and A Benchmark (2021, background): “The work [152] identifies and quantifies the key designing challenges with respect to the memory, computing, and performance tradeoffs for learned and classic optimizers, and further leverages their analysis to construct a learned optimizer that is both faster and more memory efficient.”