Understanding and correcting pathologies in the training of learned optimizers

By Luke Metz, Niru Maheswaranathan, Jeremy Nixon, C. Freeman, Jascha Narain Sohl-Dickstein. 2018

168 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Optimization & Training

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

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).”
When MAML Can Adapt Fast and How to Assist When It Cannot (2019, background): “[31] Luke Metz, Niru Maheswaranathan, Jeremy Nixon, Daniel Freeman, and Jascha SohlDickstein.”
A Generalizable Approach to Learning Optimizers (2021, methodology): “Most existing methods attempt to learn a general update rule from scratch: learning parameter updates [3, 46, 33, 34].”
Gradients are Not All You Need (2021, methodology): “In the context of continuous control this has been explored in depth as a solution to high variance gradients arising from chaos in Parmas et al. [2018], and in the context of learned optimizers in Metz et al. [2019].”
Learning to learn by gradient descent by gradient descent (2019, methodology): “Therefore this paper [2] is focused on ﬁnding the optimal unrolled steps to perform truncated backpropagation.”
Introducing Symmetries to Black Box Meta Reinforcement Learning (2021, background): “Furthermore, it was shown that meta-loss landscapes are difficult to navigate and the search distribution helps in smoothing those (Metz et al. 2019).”
Practical tradeoffs between memory, compute, and performance in learned optimizers (2022, methodology): “…Wichrowska et al., 2017), reinforcement learning (Li and Malik, 2017a;b), and more advanced training procedures (Lv et al., 2017; Maheswaranathan et al., 2019; Metz et al., 2019a; Chen et al., 2020; Vicol et al., 2021; Metz et al., 2021) leveraging both Evolution Strategies (ES), and gradients.”
Bootstrapped Meta-Learning (2021, background): “As a result most works optimise performance after K applications of the update rule and assume that this yields improved performance for the remainder of the learner’s lifetime [10, 34, 35].”
Brax - A Differentiable Physics Engine for Large Scale Rigid Body Simulation (2021, background): “Differentiating through long trajectories is an active area of research[38, 21, 18] and is known to be difficult to optimize[39, 40], thus we defer more advanced differentiable algorithms to future releases.”
MetaInit: Initializing learning by learning to initialize (2019, methodology): “While this approach can be used to tune the parameter norms, it has a signiﬁcant computational cost, can lead to overﬁtting, and gradients taken through optimization can be very poorly conditioned [40].”
Symbolic Discovery of Optimization Algorithms (2023, methodology): “Our work lies in the area of AutoML and meta-learning that includes learning to learn [1, 5, 63, 64, 78, 96, 97], neural architecture search [14, 15, 57, 71, 79, 87, 93, 94, 98, 106] and hyperparameter optimization [25, 43, 44, 55], etc.”