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

Aitor Lewkowycz; Yasaman Bahri; Ethan Dyer; Jascha Sohl-Dickstein; Guy Gur-Ari

Source-linked AI summary

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

Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, Guy Gur-Ari

arXiv:2003.02218v1Mar 4, 2020stat.MLcs.LG

TL;DR

The paper asks how learning rates shape deep-network dynamics, learned representations, and generalization beyond what existing infinite-width theory explains. It analyzes solvable network dynamics and tests their predictions in practical settings, finding distinct learning-rate phases, flatter minima at large rates, and often optimal performance in the catapult phase.

Problem
The interplay between training dynamics, learned-network properties, and generalization remains largely open, while infinite-width theory is insufficient for practical finite-width networks trained with large learning rates.
Method
The paper develops solvable training dynamics for a neural-network model and compares its predictions with experiments across practical deep-learning settings.
Results
The model’s predictions agree well with experiments, identifying distinct learning-rate phases and showing that large-learning-rate training can converge to flatter minima with optimal performance often occurring in that regime.
Takeaways & Limitations
Large initial learning rates can provide a training mechanism for accessing flatter minima and improved performance without relying on SGD noise.
Takeaways & Limitations
The results are established primarily for MSE loss and sufficiently wide networks, while extension to softmax classification and optimizers such as momentum remains unresolved.

Abstract

from arXiv · show

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.

1. Introduction

The paper addresses gaps in explaining how learning rates shape deep-network training and generalization. It develops a dynamical account of distinct learning-rate phases, including a large-learning-rate regime that reaches flatter minima and often achieves optimal performance.

The interplay between training dynamics, learned-network properties, and generalization remains largely unresolved despite deep learning’s empirical success.
The paper identifies a mechanism in which large initial learning rates, rather than SGD noise, drive convergence toward flatter minima, including under full-batch gradient descent.
Existing infinite-width and linearized theories do not fully capture practical finite-width networks trained with large learning rates or explain their performance gaps.
Training is characterized by three phases: lazy, catapult, and divergent, analytically derived for wide finite-width linear networks and tested in nonlinear architectures.
In the catapult phase, loss initially grows while curvature decreases, after which training converges to a flatter minimum; optimal performance is often achieved in this range.
The stable large-learning-rate range is bounded by ηmax = cact./λ0; empirically cact. ≈12 for ReLU networks, versus cact. = 4 in the theoretical model.

2. Theoretical results

The theoretical model separates training into learning-rate regimes with distinct dynamics. In the catapult phase, finite-width effects reduce curvature after an initial loss increase, enabling convergence where infinite-width linearized dynamics diverge.

Model and setup: The model analyzes gradient descent for a large but finite-width, one-hidden-layer linear network and extends conclusions to more general settings.The setup includes d-dimensional inputs, m training samples, and general labels; the full analysis is provided in supplementary material.
Lazy phase: For η < ηcrit = 2/λ0, the kernel curvature remains constant and training behaves as a linearized NTK model.The function and loss converge to zero in O(1) steps, with only O(n^-1) kernel corrections at finite width.
Catapult phase: For ηcrit < η < ηmax, the loss initially grows, then curvature decreases until ηλt ≲ 2, after which the loss converges to zero.This finite-width catapult effect occurs after approximately log(n) steps, when the function magnitude reaches order n^1/2 and kernel updates become non-negligible.
Catapult phase: In the catapult phase, the final kernel is smaller than its initialization value by an order n^0 amount, indicating flatter minima than in the lazy phase.The large learning rate causes non-negligible kernel updates and breaks the linearized approximation even at large width.
Empirical validation: Empirical dynamics reproduce the predicted phases: curvature is constant in the lazy phase, decreases in the catapult phase, and diverges above ηmax.Figure 2 reports these behaviors for the warmup model and marks learning rates where training diverges empirically.
Divergent phase: For η > ηmax = 4/λ0, positive kernel updates accelerate function growth and the loss diverges.The model identifies ηmax as its maximum stable learning rate; the bound includes a neglected term whose effect is noted as non-exact.

3. Experimental results

Experiments support the model’s predicted learning-rate phases across diverse deep-learning settings. Optimal performance generally occurs in the catapult phase, above ηcrit and near ηmax, although fixed-step evaluation favors large learning rates.

Experimental validation: The experiments test predicted phase transitions, including ηcrit and ηmax, across practical deep-learning settings.The theoretical framework is not guaranteed beyond its analyzed model, so these values are treated as empirical predictions.
Experimental validation: Different learning-rate phases appear across architectures, nonlinearities, parameterizations, and regularization choices.The experiments use MSE loss, sufficiently wide networks, and SGD.
Early-time dynamics: For ηcrit < η < 4/λ0, curvature decreases during training to below 2/η, matching the predicted catapult behavior.For η < ηcrit, curvature remains learning-rate independent and approximately constant; in the Wide ResNet case it initially rises before stabilizing.
Early-time dynamics: ReLU networks exhibit ηmax ≈12/λ0, exceeding the theoretical 4/λ0 maximum and depending only weakly on architectural details.The empirical replacement is ηmax = cact./λ0, with cact. ≈12 for ReLU networks.
Early-time dynamics: The loss initially increases in the catapult phase before converging, with this transient lasting fewer than 10 steps.This behavior confirms another model prediction.
Generalization performance: Across tested settings, catapult-phase performance matches or exceeds lazy-phase performance, and fixed-compute optima occur above ηcrit near ηmax.This pattern holds under both fixed compute and fixed physical-time comparisons, though fixed-step evaluation biases results toward large learning rates.

4. Additional properties of the model

After the catapult excursion, the model predicts a return to approximately linear dynamics with a constant kernel. Experiments show that post-transition linearized models can match nonlinear performance over a range of learning rates.

Restoration of linear dynamics: After the catapult excursion, logits return to O(1), the NTK stops changing, and dynamics become approximately linear with a constant kernel.This is a central additional prediction of the model at large width.
Restoration of linear dynamics: Experiments evolve the model until the catapult effect ends, then compare its dynamics with a model linearized at that point.The test targets the conjectured asymptotic linearization behavior.
Restoration of linear dynamics: At fixed width, linear and nonlinear networks have matching accuracy above the transition up to 4/λ0.Additional evidence for this asymptotic linearization behavior is reported in the Supplement.
Non-perturbative phase transition: In the small-learning-rate phase, infinite-width dynamics follow a linear random-features model, while finite-width corrections can be treated perturbatively.These corrections are organized as an expansion around infinite width.
Non-perturbative phase transition: When η > ηcrit, the leading linear dynamics diverge and logits grow to O(n^1/2), making all expansion terms comparable.Therefore, truncating the perturbative series cannot describe catapult-phase growth.

5. Discussion

The paper explains how large learning rates let deep networks reduce curvature, access flatter minima, and outperform linearized models through a distinct catapult phase. Its theoretical transitions agree quantitatively with realistic experiments, while several extensions remain open.

Learning-rate mechanism: For MSE loss, linear model training diverges above ηcrit = 2/λ0, but deep networks can stabilize larger rates by reducing curvature.The catapult effect begins with increasing loss before convergence to a small final value.
Learning-rate phases: The theoretical model predicts sharp, non-perturbative transitions from a lazy phase to a catapult phase where only the full model converges.Finite-width corrections to the linearized model cannot describe behavior beyond the critical learning rate.
Empirical agreement: The model derives transition learning rates from initialization curvature and agrees quantitatively with empirical dynamics across realistic settings.For ReLU networks, the conjectured maximum learning rate is approximately 12/λ0 and is confirmed in many cases.
Mechanistic interpretation: Catapult dynamics reduce misalignment between feedforward activations and backpropagated gradients by shrinking their orthogonal components.The paper hypothesizes that this may benefit generalization by reducing sensitivity to orthogonal out-of-distribution activation perturbations.
Performance implications: Large-learning-rate regularization is independent of stochastic-gradient noise, and optimal performance is often found above ηcrit where linearized models are unstable.This helps explain why performance gaps between nonlinear networks and their wide linear counterparts can persist at arbitrarily large widths.
Open questions: The model does not fully explain ReLU maximum rates, softmax classification, or optimizer-dependent behavior such as momentum effects.These cases are identified as open questions for future work.

A. Experimental details

The experiments use JAX-based implementations, standardized loss and kernel normalizations, fixed initialization seeds for learning-rate comparisons, and architecture-specific settings across MNIST and CIFAR tasks. Supplementary experiments measure curvature and examine phase behavior, accuracy, and physical-time comparisons.

Implementation: Experiments use JAX and the Neural Tangents Library for implementation.
Normalization: The normalized MSE loss averages squared prediction errors over classes, batch examples, and one-target labels.
Normalization: The NTK is normalized so its eigenvalues match the non-zero Fisher-information eigenvalues.
Curvature measurement: Top NTK eigenvalues are estimated with Lanczos’ algorithm on small data batches and averaged across batches.The experiments do not focus on fluctuations in the top eigenvalue between batches.
Experimental controls: Learning-rate comparisons use the same weight initialization seed and usually one initialization, while σw and σb denote width-independent initialization coefficients.
Supplementary analyses: Supplementary figures visualize three-phase dynamics, test accuracy, training accuracy, and learning-rate behavior under modified experimental settings.The figures include longer training runs, different regularization settings, and physical-time comparisons.
Architectures and datasets: The study evaluates fully connected, convolutional, and Wide ResNet architectures with SGD, including MNIST, CIFAR-10, and CIFAR-100 settings.Captions specify variations in depth, nonlinearity, normalization, batch size, regularization, and augmentation-related settings.

B.3. Comparison of learning rates for different L2 regularization for WRN28-10 on CIFAR10

CIFAR-10 experiments test whether the learning-rate advantage persists across L2 regularization choices. Larger learning rates remain favorable, and performance peaks in the catapult phase, though the gap narrows with larger regularization.

No L2 regularization: Without L2 regularization, larger learning rates perform better even without learning-rate decay, although convergence takes much longer.
Comparison setup: The CIFAR-100 comparison uses WRN28-10 with ηcrit ≈ 0.4 and evaluates test accuracy across learning rates and physical-time schedules.
No L2 regularization: Removing L2 regularization lengthens convergence but has little effect on performance in the reported experience.
Larger L2 regularization: With L2 = 0.001, optimal physical-time performance still peaks in the catapult phase, although the difference is smaller.

B.4. Training accuracy plots

Training-accuracy plots compare learning rates under fixed steps and fixed physical time. Smaller rates can have higher training accuracy at equal physical time yet still perform worse when the comparison fixes the number of steps.

Training-accuracy comparisons: Training-accuracy plots cover WRN CIFAR-10 and CIFAR-100 experiments under both fixed-step and fixed-physical-time comparisons.
Training-accuracy comparisons: Smaller learning rates have higher training accuracy at the same physical time but still perform worse when compared after a fixed number of steps.

C.1. ReLU activations for the simple model

ReLU nonlinearities broaden the stable learning-rate regime but introduce distinct failure modes and optimizer-dependent kernel dynamics. Across normalizations and activations, the critical behavior remains tied to learning-rate and curvature scales.

ReLU trainability: ReLU networks remain trainable up to η = 12/λ0, beyond which single-sample neurons die while deeper networks can diverge.The simple model avoids loss divergence through neuron death, whereas deeper multi-sample networks exhibit loss divergence after approximately this threshold.
Momentum: Momentum makes the linear model stable for η < 2/[λ0(1 + γ)], but ReLU networks instead show non-constant curvature dynamics below their effective critical rate.For ReLU networks, strong kernel time dependence appears once η exceeds ηeff,crit = ηcrit/(1−γ).
Regularization: L2 regularization leaves the early rearrangement dynamics roughly unchanged, even while the maximum Hessian eigenvalue decreases slowly.The authors do not expect regularization to affect early-time dynamics, but note that it may influence later behavior after strong rearrangement.
Activation dependence: For tanh activation, ηmax is closer to the simple-model prediction 4/λ0.This contrasts with the broader trainability regime observed for deep ReLU networks.
Normalization: NTK normalization preserves phase behavior because the relevant quantity is ηλt, independent of normalization.The same setup is evaluated for Wide ResNet networks with NTK normalization.

D.1. Full model analysis

The finite-width analysis identifies lazy and catapult phases separated by ηλ0 = 2, with curvature dynamics explaining both convergence and instability. In the catapult phase, curvature decreases until training reaches a flatter minimum.

Lazy phase: When ηλ0 < 2, error components shrink, kernel updates remain order n^-1, and the network effectively evolves as a linear model.This is the lazy phase associated with fixed-kernel infinite-width behavior.
Catapult phase: When 2 < ηλ0 < 4, the error grows exponentially and aligns with the top kernel eigenvector, initiating the catapult dynamics.The top-eigenvector component dominates after order n^0 steps.
Curvature reduction: Curvature decreases once error growth makes kernel updates order n^0, then convergence begins after λt falls below 2/η.Gradient descent ultimately reaches a global minimum with lower curvature than at initialization.
Phase transition: At ηcrit, final curvature is constant below the transition but non-constant above it in the modified large-width limit, making the transition non-analytic.Consequently, finite-order Taylor expansion from the small-learning-rate phase cannot describe the large-learning-rate phase.
Finite-width effects: The usual infinite-width approximation fails above ηcrit because the dynamics require large finite-width effects rather than perturbative corrections.The authors argue that perturbative finite-width expansions do not capture this regime.

E.2. Critical exponents

Convergence slows near the critical learning rate according to a power law with exponent -1 on both sides of the transition. Empirical results support this predicted critical behavior and examine the later return of linear dynamics.

Below the transition: Near ηcrit from below, convergence time scales as t* ∼ (ηcrit − η)^−1, giving critical exponent -1.This follows from the slowly shrinking error when ηλ0 = 2 − ϵ.
Above the transition: Near ηcrit from above, early exponential error growth dominates convergence time and again yields exponent -1.The early phase lasts t̂ ∼ ϵ^-1 log √n before curvature updates become substantial.
Empirical verification: Figure S12 empirically verifies that convergence time diverges near ηcrit with measured exponents close to -1.The experiment uses the warmup model at width 16,000.
Return of linear dynamics: After the initial catapult dynamics, a model linearized after 10 steps tracks nonlinear performance, while NTK changes decrease as width increases.Figure S13 studies two-class MNIST with 100 samples per class.