[Paper Review] The large learning rate phase of deep learning: the catapult mechanism

Aitor Lewkowycz, Yasaman Bahri+3 more|arXiv (Cornell University)|Mar 4, 2020

Stochastic Gradient Optimization Techniques37 references60 citations

TL;DR

The paper introduces three learning-rate phases (lazy, catapult, divergent) in gradient descent, presents a solvable finite-width model showing catapult dynamics that lead to flatter minima, and validates predictions across practical deep networks with empirical evidence that optimal performance often lies in the large learning-rate catapult phase.

ABSTRACT

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.

Motivation & Objective

• Motivate and characterize how initial learning rate profoundly affects deep network training and generalization.
• Introduce a tractable finite-width model that predicts three distinct learning-rate regimes.
• Bridge theory and practice by validating phase predictions in realistic deep networks across architectures.
• Show that optimal performance often occurs in the large learning-rate (catapult) phase.
• Dissociate flatness and generalization dynamics from SGD noise, focusing on learning-rate-driven effects.

Proposed method

• Derive exact gradient descent updates for a one-hidden-layer linear network with large but finite width and MSE loss.
• Identify and analyze three learning-rate regimes using the top NTK eigenvalue as a curvature proxy.
• Extend the analysis to the full model with d-dimensional inputs and m training samples, deriving analogous update dynamics.
• Conduct empirical experiments across fully-connected, convolutional, and residual networks to test phase predictions.
• Estimate the practical maximum learning rate using an architecture-dependent constant c_act, with ReLU ~12 per the experiments.

Experimental results

Research questions

• RQ1 What are the dynamical phases of gradient descent under varying initial learning rates in wide but finite networks?
• RQ2 How does the learning rate influence kernel curvature, particularly the NTK top eigenvalue, during training?
• RQ3 Can large learning rates be stabilized to converge to flatter minima, and how does this affect generalization?
• RQ4 Do theoretical phase predictions hold in realistic architectures and SGD settings?
• RQ5 What is the empirical relationship between architecture, nonlinearity, and the maximum stable learning rate?

Key findings

• Three learning-rate phases exist: lazy (eta < 2/lambda_0), catapult (2/lambda_0 < eta < eta_max), and divergent (eta > eta_max).
• In the catapult phase, initial loss rises with a rapid curvature decrease, then converges to a flatter minimum than in the lazy phase.
• The maximum stable learning rate is approximately eta_max = c_act./lambda_0, with c_act depending on the nonlinearity (≈4 in theory, ~12 for ReLU in practice).
• Empirical results across CNNs, ResNets, and fully-connected nets align with phase boundaries and show peak performance in the catapult phase.
• Optimal performance often occurs in the large learning-rate catapult phase, consistent across architectures and training budgets.
• After the catapult, the model behavior resembles linear dynamics with a nearly constant kernel, suggesting restoration of linear-like dynamics.