PaperlayerView exact revision on arXiv

Source-linked AI summary

A Loss Curvature Perspective on Training Instability in Deep Learning

Justin Gilmer, Behrooz Ghorbani, Ankush Garg, Sneha Kudugunta, Behnam Neyshabur, David Cardoze, George Dahl, Zachary Nado, Orhan Firat

arXiv:2110.04369v1cs.LGcs.AI

TL;DR

The paper examines why neural-network training configurations can diverge and how loss curvature changes with learning rates and training choices. Across experiments, initialization, warmup, and normalization enable higher learning rates by reducing curvature, while the authors note limits to the analysis for adaptive optimizers and broader Hessian effects.

  • Problem

    Training can diverge after small configuration changes, and selecting settings that balance stability with rapid progress is difficult and often requires extensive hyperparameter searches.

    Source: §1 · p. 1

  • Method

    The authors run large-scale experiments tracking loss sharpness while varying learning rate, warmup, initialization, and architectural choices.

    Source: §1 · p. 2

  • Results

    Across experiments, initialization, learning-rate warmup, and normalization enabled higher learning rates without divergence by reducing λ1, which also adapted to the learning rate.

    Source: §8 · pp. 9–10

  • Takeaways & Limitations

    Well-optimized ResNet-50 and WideResNet 28-10 models showed little warmup benefit at tested rates, but small architecture modifications could produce high curvature and early instability.

    Source: §8 · p. 10

    The analysis focuses primarily on SGD with momentum; the authors leave the Adam analogue of λ1 ≤2/η and a deeper analysis of adaptive optimization for future work.

    Source: §7 · p. 9

Abstract

from arXiv · show

In this work, we study the evolution of the loss Hessian across many classification tasks in order to understand the effect the curvature of the loss has on the training dynamics. Whereas prior work has focused on how different learning rates affect the loss Hessian observed during training, we also analyze the effects of model initialization, architectural choices, and common training heuristics such as gradient clipping and learning rate warmup. Our results demonstrate that successful model and hyperparameter choices allow the early optimization trajectory to either avoid -- or navigate out of -- regions of high curvature and into flatter regions that tolerate a higher learning rate. Our results suggest a unifying perspective on how disparate mitigation strategies for training instability ultimately address the same underlying failure mode of neural network optimization, namely poor conditioning. Inspired by the conditioning perspective, we show that learning rate warmup can improve training stability just as much as batch normalization, layer normalization, MetaInit, GradInit, and Fixup initialization.

1 Introduction

The paper studies why training configurations become unstable and how curvature changes with learning rates, initialization, and architecture. It argues that several successful interventions help by keeping or moving optimization into lower-curvature regions that tolerate larger learning rates.

1 Introduction

  • Reliable training remains difficult because small configuration changes can cause divergence, while selecting stable yet rapid optimization settings is hard and often requires extensive tuning.The introduction notes that hyperparameter-tuning protocols can have an outsized influence on results [3] [28], and successes often rely on large searches [20].Source: §1 · p. 1
  • The experiments vary learning rate, warmup period, initialization, and architecture, and identify reducing λ1 as a primary benefit of these choices.The study also examines normalization location and warmup schedules.Source: §1 · p. 2
  • SGD+momentum training is stable only when its trajectory primarily remains where λ1 ≲2/η, corroborating theoretical predictions and prior empirical observations.Here η is the learning rate; the paper reports this as a large-scale empirical confirmation.Source: §1 · p. 2
  • For architectures without normalization, successful initialization strategies primarily reduce curvature early in training, enabling larger learning rates.The paper identifies MetaInit, GradInit, and Fixup among the initialization strategies considered.Source: §1 · p. 2
  • Warmup gradually reduces λ1 and offers benefits similar to better initialization; warmup alone can match key initialization advances, [37], [38].The paper connects warmup’s mechanism to the dynamical stability model of Wu et al..Source: §1 · p. 3
  • Large loss curvature can impair scaling at large batch sizes, while interventions that improve conditioning can substantially improve data-parallelism use.Source: §1 · p. 3

2 Related Work

Prior work links curvature to normalization, learning rates, and training stability, but leaves open how these relationships extend across initialization and warmup choices. The paper reports both familiar curvature dynamics and cases that depart from existing stability accounts.

2 Related Work

  • Batch Normalization has been linked to improved loss-surface smoothness, especially at higher learning rates, but the paper also notes instability and high curvature early in some BatchNorm runs.The latter cases are reported in Section 4.Source: §2 · p. 3
  • SGD models show progressive sharpening, but some models train despite starting above λ1=2.0/η, while others diverge despite starting at or below that threshold.The paper notes that its experimental setup violates assumptions of the cited catapult model, and that the observed divergence below the threshold is not predicted by it.Source: §2 · p. 3
  • Larger learning rates are associated with flatter observed curvature, and warmup and gradient clipping can reduce λ1 further, allowing some models to train at even larger rates.The paper connects this mechanism to preconditioning benefits discussed by Jastrzebski et al. [13].Source: §2 · p. 3

3 Experimental Setup

The experiments span image classification, language modeling, and machine translation, with architecture and normalization variants selected for each setting. Models are trained across learning rates, and curvature is estimated with Lanczos iterations.

3 Experimental Setup

  • Benchmarks include CIFAR-10 and ImageNet for image classification, LM1B for language modeling, and WMT for neural machine translation.CIFAR-10 experiments use WideResNet and DenseNet with and without Batch Normalization; NMT is tuned on WMT’16 EN-DE validation data and evaluated on WMT’14 EN-DE test BLEU.Source: §3 · p. 3
  • The study varies initialization and normalization choices, including scaled default initialization, MetaInit, and three Layer Norm settings for six-layer Transformer models.The Transformer settings are pre-Layer Norm, post-Layer Norm, and no Layer Norm.Source: §3 · p. 3
  • Models use cosine learning-rate decay unless stated otherwise; warmup is linear from zero to the target rate, and Hessian top eigenvalues are estimated with Lanczos.Lanczos iteration counts vary by architecture, with details deferred to the appendix.Source: §3 · p. 4

4 Early Training Instability and the Loss Hessian

Across models and datasets, successful training requires the optimization trajectory to enter a region where the largest Hessian eigenvalue is near or below 2/η. Initialization and architecture affect early curvature, but warmup can help high-curvature models reach flatter regions, and a stable start alone does not guarantee successful training.

4 Early Training Instability and the Loss Hessian

  • Across datasets and models, successful training occurs only after optimization enters the stable region λ1 ≤ 2/η; divergent models are outside it shortly before divergence.The measurements include curvature at initialization and during training, with divergent runs recorded just before loss becomes NaN.Source: §4 · p. 4, Figure 2 · p. 5
  • For CIFAR-10 WideResnet without batch normalization, η > .1 diverges by default, while warmup or MetaInit enables training at η = .1.Warmup progressively lowers curvature; MetaInit starts in a flatter region. The same model without warmup diverges within three steps.Source: §4 · p. 4, Figure 3 · p. 6
  • DenseNet batch-normalized variants start with higher curvature than non-BN variants, yet warmup makes them trainable at η > .1 when non-BN variants diverge.At initialization, these experiments do not show the commonly described BN smoothness benefit.Source: §4 · p. 4
  • On ImageNet, preactivation ResNet-50 variants start at higher curvature and diverge when η > .5, while the post-activation variant trains at learning rates up to 10.The comparison also considers reversing the ReLU and batch-normalization order in the preactivation block.Source: §4 · p. 4
  • Some models diverge despite starting where λ1 < 2/η, as curvature can rise rapidly in early training.Examples include pre- and post-layernorm Transformers and a MetaInit WideResnet; some diverge before step 100.Source: §4 · p. 4, Figure 3 · p. 6

5 The Interaction between Learning Rate Warmup, Initialization and Curvature

Learning-rate warmup and initialization can both guide optimization toward lower curvature, connecting them through the trajectory they produce rather than only their starting points. Warmup gradually raises η while curvature stays near the stability boundary, and mid-training conditioning is largely determined by the learning rate.

5 The Interaction between Learning Rate Warmup, Initialization and Curvature

  • Warmup gradually raises η while optimization moves toward lower curvature; during much of the ramp, λ1 remains near 2/η.This relationship is observed across different warmup lengths.Source: §5 · p. 6, Figure 4 · p. 7
  • On a non-BN WideResnet at η = .1, default initialization diverges within three steps without warmup, whereas warmup enables training and MetaInit enables training without warmup.MetaInit begins at low curvature; with warmup, curvature progressively decreases as the learning rate reaches its peak.Source: §5 · p. 7
  • GradInit and gradient clipping both reduce curvature, but GradInit’s flatter starting point lasts only a few SGD steps before curvature approaches the warmup-trained model’s.The paper reports that learning-rate warmup alone can match key advances from initialization methods in Tables 1 and 2.Source: §5 · p. 7
  • After 3000 SGD steps, MetaInit’s gradient quotient is almost completely determined by the learning rate rather than the initialization method.The gradient quotient is the conditioning measure directly optimized by MetaInit; its initialization advantage is short-lived.Source: §5 · p. 7

6 The Effects of Curvature on Batch Size Scaling

In WideResnet variants, higher curvature limits the training-speed gains available from larger batches, alongside limiting the usable learning rate. The low-curvature model scales nearly linearly with batch size, whereas the high-curvature model sees minimal speed improvement and remains constrained to a small learning rate.

6 The Effects of Curvature on Batch Size Scaling

  • The study sweeps learning rate and batch size for low-, moderate-, and high-curvature WideResnet variants, measuring steps to 85% validation accuracy and the optimal learning rate.Steps are normalized to the value measured at batch size 64.Source: §6 · p. 8
  • The low-curvature WideResnet shows almost linear training-speed gains as batch size grows, while the high-curvature variant gains little.The high-curvature model’s optimal η is almost always the largest non-divergent value, indicating that curvature restricts larger learning rates.Source: §6 · pp. 8–9
  • For the low-curvature model, η* increases linearly with batch size; for the high-curvature model, η* stays around 3 × 10^-3.The high-curvature model’s η* is almost always its largest non-divergent learning rate.Source: §6 · p. 9
  • Because large λ1 limits the maximum usable learning rate, it can constrain effective use of data parallelism; warmup or better initialization may improve batch-size scaling.An optimal learning rate near the largest non-divergent value is proposed as a diagnostic for curvature-limited training speed.Source: §6 · pp. 8–9

7 Limitations

The analysis focuses primarily on SGD with momentum, leaving deeper analysis of Adam for future work. Its measurements also focus on initialization’s influence on the top Hessian eigenvalue mid-training, potentially missing effects on the broader eigenspectrum.

7 Limitations

  • The analysis focuses primarily on SGD with momentum; the analogue of λ1 ≤ 2/η for Adam is unclear, and deeper analysis is left for future work.The authors report evidence that curvature adaptation and warmup’s push toward flatter regions also occur with Adam-trained Transformers.Source: §7 · p. 9
  • Measurements focus on initialization’s influence on λ1 mid-training, so possible lasting effects on the broader Hessian eigenspectrum may be missed.The authors cite a possible effect on λk/λ1 for smaller eigenvalues as an example.Source: §7 · p. 9

8 Conclusion

The authors conclude that initialization, warmup, and normalization can support higher learning rates by reducing λ1 during training. They also report that λ1 adapts to the learning rate, while high curvature can limit batch-size scaling and warmup does not benefit every well-optimized model.

8 Conclusion

  • Initialization, warmup, and normalization enable higher learning rates without divergence by reducing λ1 during training.Two investigated popular variants, ResNet-50 and WideResNet 28-10, did not benefit from warmup and had small λ1 values at the learning rates considered.Source: §8 · pp. 9–10
  • Low-curvature WideResNet scales nearly linearly with batch size, while the high-curvature model shows almost no speedup at larger batch sizes.The low-curvature model’s optimal learning rate increases almost linearly with batch size; the high-curvature model’s optimal rate remains at its largest non-diverging value.Source: Figure 6 · p. 10
  • Even seemingly trivial architecture modifications can produce large λ1 and early instability, so reducing the learning rate or abandoning the modification may be a naive response.The authors suggest investigating initialization, warmup, gradient clipping, or normalization-layer placement.Source: §8 · p. 10
  • λ1 adapts to the learning rate across cases both below and above 2.0/η, contrary to conventional advice to choose step size from measured λ1.The authors call for more work to understand this adaptation and its implications for optimizer design, step-size selection, and initialization.Source: §8 · p. 10

9 Societal Impact

The paper discusses potential computational and environmental consequences of its findings, including possible emissions from the experiments and potential inference power savings from normalization-free networks. A miscellaneous figure caption states that test-time behavior mirrors training-loss dynamics.

9 Societal Impact

  • Normalization-free networks can improve inference-time power consumption in large-scale deployments, while warmup is suggested as an alternative to computationally expensive initialization strategies.The authors note that computational savings can sometimes lead to increased overall usage.Source: §9 · pp. 10–11
  • The experiments required numerous runs that can sometimes result in significant carbon emissions, although the cloud resources used ran on carbon-free power that was not yet 24/7 carbon-free and still used offsets.The authors note that their experiments were numerous enough to reach useful conclusions.Source: §9 · p. 11
  • The authors state that the work contributes to understanding methods that may be used for benevolent or harmful purposes, but they do not believe harmful uses are most likely given the results’ fundamental nature.Source: §9 · p. 11

A Miscellaneous Figures

  • The figure caption says test-time behavior of the WideResNet models closely mirrors their training-loss dynamics.Source: Figure 7 · p. 14

B Performance of models in Figure 2

Across the reported models, high curvature limits the use of higher learning rates and final performance; some models that begin below the stability threshold nevertheless enter a higher-curvature region before diverging. The appendix describes the Lanczos estimates and experimental configurations used for the eigenspectrum analysis and related model studies.

B Performance of models in Figure 2

  • High curvature limits use of higher learning rates and final performance across the reported settings; some models enter λ1 > 2.0/η before diverging around step 90.Figure 13 follows a WideResNet 28-10 that began with λ1 < 2.0/η but quickly left the stable region.Source: B Performance of models in Figure 2 · p. 14, Figure 13 · p. 16

C Details on Computing the Hessian Eigenspectrum via Lanczos

  • Lanczos iterations estimate the Hessian’s top eigenvalue using Hessian-vector products computed via Pearlmutter’s trick.The experiments use 45 steps for WMT/LM1B and 40 for image models, except ResNet with ReLU→BN ordering, which uses 200 steps because of a small eigengap; vectors are stored in float64 and reorthogonalized at each step.Source: C Details on Computing the Hessian… · pp. 14–15

D.1 Neural Machine Translation

  • The NMT experiments use Transformer-Base and Transformer-Wide models, sweeping warmup, learning rate, clipping, and initialization scaling while optimizing validation loss and reporting test BLEU.Transformer-Base uses batch size 1024 and Transformer-Wide uses batch size 512; both are trained for 60 epochs.Source: §D.1 · p. 15

D.2 DenseNets

  • The appendix specifies the ResNet-50 (without BatchNorm) setup as SGD with Nesterov momentum, gradient clipping, and 1000-step linear warmup.For Table 1, DenseNet-100 uses the GradInit codebase with default initialization and gradient clipping of norm 6.Source: §D.2 · pp. 16–17

E Training Details for Figure 2

Training details vary by architecture: the experiments use WideResNet, DenseNet, ResNet-50, and Transformer models with architecture-specific batch sizes and schedules.

E Training Details for Figure 2

  • WideResNet-28-10 models used batch size 1024 for 300 epochs, MixUp, and 1000-step linear warmup followed by cosine decay.Source: E Training Details for Figure 2 · p. 17
  • DenseNet models used batch size 512 for 200 epochs, SGD with momentum, and 1000-step warmup where applicable followed by cosine decay.The setup also specifies weight decay, L2 regularization, growth rate 32, and reduction 0.5.Source: E Training Details for Figure 2 · p. 17
  • ResNet-50 models used batch size 2048 with SGD and Nesterov momentum, 1000-step linear warmup followed by cosine decay, and label smoothing.The experiments also used flip-and-crop augmentation.Source: E Training Details for Figure 2 · p. 17
  • Transformer models on LM1B used batch size 1024 with SGD and Nesterov momentum, six layers, eight heads, and the ResNet learning-rate schedule.The setup used embedding dimension 512, MLP hidden dimension 1024, and attention dropout 0.1.Source: E Training Details for Figure 2 · p. 17

F Curvature Adaptation with the Adam Optimizer

The appendix tests whether curvature-based stability insights extend to Adam using a preconditioned-Hessian eigenvalue approximation. Instabilities coincide with an empirical threshold crossing, while the authors leave broader exploration of adaptive optimizers to future work.

F Curvature Adaptation with the Adam Optimizer

  • The stability approximation does not account for adaptive preconditioning or momentum, and the authors leave further study of adaptive optimizers to future work.They report empirical evidence that the approximation helps explain optimization stability.Source: F Curvature Adaptation with the Adam… · pp. 17–18
  • Adam training instabilities coincide with the preconditioned Hessian’s maximum eigenvalue crossing the empirical 40/η threshold.Figure 14 plots the eigenvalue and training loss for models with different warm-up lengths.Source: F Curvature Adaptation with the Adam… · pp. 17–18, Figure 14 · p. 18
  • The preconditioned maximum eigenvalue adapts to the warm-up schedule, and the empirical 40/η threshold fits the data better than 2/η.Source: F Curvature Adaptation with the Adam… · p. 18
  • Longer warmup closely corresponds to a slower change in λ1 as the optimization trajectory moves toward regions of reduced curvature.Source: Figure 15 · p. 18
  • With batch size increased from 32 to 2048, mid-training curvature rises from well below toward the 2.0/η approximation.The WideResNet without batch normalization used the same learning schedule and 1000 warmup steps across batch sizes.Source: Figure 16 · p. 19

G Compute Resources Used

  • Experiments mostly used v2 cloud TPUs; ResNet-50 and DenseNet Figure 2 experiments used v3 TPUs, GradInit used a V100 GPU, and machine translation used v3 TPUs.The Figure 2 experiments ran in parallel using up to 50 v2 TPUs over a few days.Source: G Compute Resources Used · p. 20
Loading 2110.04369v1…