The large learning rate phase of deep learning: the catapult mechanism

By Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Narain Sohl-Dickstein, Guy Gur-Ari. 2020

292 citations recorded by the source corpus.

The choice of initial learning rate can have a profound effect on the performance of deep networks. We present a class of neural networks with solvable training dynamics, and confirm their predictions empirically in practical deep learning settings. The networks exhibit sharply distinct behaviors at small and large learning rates. The two regimes are separated by a phase transition. In the small learning rate phase, training can be understood using the existing theory of infinitely wide neural networks. At large learning rates the model captures qualitatively distinct phenomena, including the convergence of gradient descent dynamics to flatter minima. One key prediction of our model is a narrow range of large, stable learning rates. We find good agreement between our model's predictions and training dynamics in realistic deep learning settings. Furthermore, we find that the optimal performance in such settings is often found in the large learning rate phase. We believe our results shed light on characteristics of models trained at different learning rates. In particular, they fill a gap between existing wide neural network theory, and the nonlinear, large learning rate, training dynamics relevant to practice.

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

Feature Learning in Infinite-Width Neural Networks (2020, background): “Other Related Works [24] proposed a toy model to study how large learning rate can induce a neural network to move out of the kernel regime in Ω(log(width)) time.”
Understanding the Role of Training Regimes in Continual Learning (2020, background): “The learning rate and batch size influence both the endpoint curvature and the whole trajectory [77, 17, 45].”
Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability (2021, background): “Lewkowycz et al. (2020) found that in neural network training, if the sharpness at initialization is larger than 2/η, then after becoming initially destabilized, gradient descent does not always diverge entirely (as the quadratic Taylor approximation would suggest), but rather sometimes “catapults”…”
Dataset Distillation with Infinitely Wide Convolutional Networks (2021, methodology): “…, and have been utilized in understanding trainability [Xiao et al., 2020], generalizations [Adlam and Pennington, 2020], training dynamics [Lewkowycz et al., 2020, Lewkowycz and Gur-Ari, 2020], uncertainty [Adlam et al., 2021], and demonstrated their effectiveness for smaller datasets…”
Do Wide and Deep Networks Learn the Same Things? Uncovering How Neural Network Representations Vary with Width and Depth (2020, background): “…to their corresponding kernels (Matthews et al., 2018; Lee et al., 2018; Jacot et al., 2018), but substantial differences exist between behavior in this inﬁnite width limit and the behavior of ﬁnite-width networks (Novak et al., 2018; Wei et al., 2019; Chizat et al., 2019; Lewkowycz et al., 2020).”
Learning where to learn: Gradient sparsity in meta and continual learning (2021, background): “Interestingly, similar effects have been reported in classic neural network training where both freezing layers throughout training (Raghu et al., 2017; Brock et al., 2017) and the use of large learning rates (Lewkowycz et al., 2020) seem to improve generalization.”
TRAK: Attributing Model Behavior at Scale (2023, unknown): “Prior works (Leclerc & Madry, 2020; Lewkowycz et al., 2020) suggest that neural network training can exhibit two stages or regimes: in the first stage, the features learned by the network evolve rapidly; in the second stage, the features remain approximately invariant and the overall optimization…”
Gradient Flow in Sparse Neural Networks and How Lottery Tickets Win (2020, background): “Given the importance of the early phase of training (Frankle, Schwab, and Morcos 2020; Lewkowycz et al. 2020), it is natural to ask about the difference between LTs and the solution they are derived from (i.e. pruning solutions).”
Deep Learning on a Data Diet: Finding Important Examples Early in Training (2021, background): “However, ﬁnite neural networks outperform their inﬁnite-width limits [20] and have different dynamics early in training [21].”
Explaining neural scaling laws (2021, background): “Furthermore, in the large width limit, randomly initialized neural networks become Gaussian Processes [9–11, 46–48], and in the low-learning rate regime [13, 49, 50] neural networks train as linear models at infinite width [12, 13, 51].”
Learning to Optimize: A Primer and A Benchmark (2021, background): “For example, it is well known that training over-parameterized networks will exhibit sharply distinct behaviors at small and large learning rates, and the two regimes usually co-exist in practical training and are separated by a phase transition [114; 115; 230].”
Adaptive Gradient Methods at the Edge of Stability (2022, background): “This behavior is similar to the “catapult” effect studied in [25] for non-adaptive optimizers.”