Title: Why Warmup the Learning Rate? Underlying Mechanisms and Improvements
Venue: NeurIPS 2024
Accepted: yes
Keywords: Learning rate Warmup

Optimization Dynamics
Abstract: In modern deep learning, it is common to warm up the learning rate $\eta$, often by a linear schedule between $\eta_{\text{init}} = 0$ and a predetermined target $\eta_{\text{trgt}}$. In this paper, we show through systematic experiments with SGD and Adam that the overwhelming benefit of warmup arises from allowing the network to tolerate larger $\eta_{\text{trgt}}$ by forcing the network to more well-conditioned areas of the loss landscape. The ability to handle larger target learning rates in turn makes hyperparameter tuning more robust while improving the final performance of the network. We uncover different regimes of operation during the warmup period, depending on whether the network training starts off in a progressive sharpening or sharpness reduction phase, which in turn depends on the initialization and parameterization. Using these insights, we show how $\eta_{\text{init}}$ can be properly chosen by utilizing the loss catapult mechanism, which saves on the number of warmup steps, in some cases completely eliminating the need for warmup. We also suggest an initialization for the variance in Adam, which provides benefits similar to warmup.

Why Warmup the Learning Rate? Underlying Mechanisms and Improvements

Abstract

It is common in deep learning to warm up the learning rate η , often by a linear schedule between η init = 0 and a predetermined target η trgt. In this paper, we show through systematic experiments using SGD and Adam that the overwhelming benefit of warmup arises from allowing the network to tolerate larger η trgt by forcing the network to more well-conditioned areas of the loss landscape. The ability to handle larger η trgt makes hyperparameter tuning more robust while improving the final performance. We uncover different regimes of operation during the warmup period, depending on whether training starts off in a progressive sharpening or sharpness reduction phase, which in turn depends on the initialization and parameterization. Using these insights, we show how η init can be properly chosen by utilizing the loss catapult mechanism, which saves on the number of warmup steps, in some cases completely eliminating the need for warmup. We also suggest an initialization for the variance in Adam which provides benefits similar to warmup.

1 Introduction

One of the most important choices to make in gradient-based optimization is the learning rate (step size) η . If η is too small, then learning may take place too slowly or the model might get stuck in unfavorable regions of the loss landscape. If η is too large, training will typically diverge. In practice, it is common to pick a dynamical learning rate schedule η t [2, 4, 40, 26]. Modern learning rate schedules for deep learning typically consist of a warmup period where η t is increased linearly from zero to a target value η trgt over a warmup time T wrm [13, 34]. After the warmup period, it is common to eventually decay the learning rate, for example via a cosine decay schedule [34, 26, 40].

Given that warmup is standard in the practitioner's toolkit, it is important to understand it deeply and identify improvements. In modern settings, perhaps the earliest work to use warmup was [14], which used a small constant learning rate for the first few epochs of training and then switched to a larger learning rate. A linear warmup schedule was later introduced in [13]. The intuition given was that to scale the minibatch size in SGD by a factor of k , it is natural to also scale the learning rate by a factor of k , provided the model is not changing too rapidly and successive gradients are roughly aligned. However at the beginning of training, the model is changing rapidly, so it is natural to start with a lower learning rate and gradually increase it to the target value after the network has stabilized.

Other explanations suggest that since the network is initialized randomly, the gradient steps at the beginning of training are not meaningful, and thus it would be harmful to take large steps in such directions [40], so it makes sense to take smaller steps early in training. The analysis by [12] suggests that warmup primarily limits the magnitude of weight updates in the deeper layers, preventing large instabilities. It has also been suggested that the key benefit of warmup arises for adaptive optimizers, such as Adam: [23] argues that the variance of the adaptive learning rate is large during early training because the network has seen too few training samples; it is asserted that this large variance is harmful,

and that warmup acts as a variance reduction method by allowing the network to collect accurate statistics of the gradient moments before using larger learning rates. Alternatively, it is also sometimes stated that the initialization may start the model off at places in parameter space that are unstable, difficult to optimize, and easily lead to divergence, and that warmup can help alleviate this [40].

The above explanations are varied and do not clearly demonstrate why and to what extent warmup is necessary. A loss landscape perspective was given in [10] (and summarized in [26] Ch. 8), which argued that an important effect of warmup is to gradually reduce the sharpness (the top eigenvalue of the Hessian of the loss), thus causing the model to leave poorly conditioned areas of the loss landscape and move towards flatter regions which can tolerate larger learning rates. They argue that the mechanism for this is similar to the dynamical stability (catapult) mechanisms studied in [35, 22].

Our contributions. Here we perform extensive studies on the effect of learning rate warmup across architectures (FCNs, ResNets, and Transformers), initializations and parameterizations, datasets (CIFAR-10, CIFAR-100, TinyImageNet, WikiText), and for both SGD and Adam.

We demonstrate through systematic experiments that by far the primary benefit of learning rate warmup is to allow the network to tolerate larger learning rates than it otherwise would have. This builds on the observations of [10] by showing that any other benefits are marginal, disentangling the effect of warmup duration and target learning rate, and by extending the empirical evidence to include adaptive optimizers and Transformers.

For SGD, the maximal allowable learning rate is determined by the sharpness (the top eigenvalue of the Hessian of the loss). As we discuss in Section 4, we find that there are several qualitatively distinct regimes and mechanisms at play. These depend on whether the network starts off in a sharpness reduction or progressive sharpening phase [18, 19, 5], which in turn depends on the initialization and parameterization. We further find that the performance of the network is largely determined by the target learning rate. For a fixed target learning rate, increasing the warmup time provides only marginal benefit, which arises by keeping the network further away from the divergence (failure) boundary. The ability of the network to withstand a larger target learning rate in turn makes hyperparameter tuning of the target learning rate more robust, since the network responds well to a larger window of target learning rates, possibly explaining the popularity of warmup.

We then investigate Adam in detail, and show that the underlying mechanisms of warmup are similar to the SGD case, but with sharpness replaced by a preconditioned sharpness (the top eigenvalue of the pre-conditioned Hessian, defined below) . Our results disagree somewhat with prior results [23] on the underlying reason for warmup's benefits: We find that the key issue is not observing too few training samples, but rather that the pre-conditioned sharpness typically starts off at high values (even in the large batch case), causing considerable instabilities at high learning rates. Such instabilities, which may be retained in Adam's memory, can result in performance degradation and even training failures. Warmup mitigates such instabilities by gradually pushing down the preconditioned sharpness, enhancing performance, and preventing training failures. We propose a simple alternative initialization for Adam, which we refer to as GI-Adam, which provides benefits similar to warmup and consistently improves over standard Adam by inducing lower preconditioned sharpness at initialization, thus pushing the training failure boundary to higher target learning rates. This also demonstrates a different way to remove the bias correction of RMSProp with momentum.

Our analysis shows how much of the time spent during the warmup period is wasted. We show that this wasted time can be saved by making use of the catapult mechanism [22] to effectively estimate the initial sharpness scale by line search, providing a more principled choice of η init. Our experiments show that, depending on the target learning rate and initial sharpness, one can dramatically reduce the warmup time, and in some cases remove it altogether.

2 Notations and Preliminaries

SGD(-M): Given gradients g t := ∇ θ L ( θ t ) at step t , Stochastic Gradient Descent with momentum updates the parameters θ t using learning rate η t and momentum m t with coefficient β . The update equations are: m t +1 = g t + β m t and θ t +1 = θ t -η t m t +1 . β = 0 corresponds to SGD.

Adam: Adam [20] updates the parameters θ t according to the equations: m t +1 = β 1 m t +(1 -β 1 ) g t , v t +1 = β 2 v t +(1 -β 2 ) g 2 t , and θ t +1 = θ t -η t ˆ m t +1 √ ˆ v t +1 + ϵ = θ t -η t P t +1 m t +1 , where ˆ m t = m t 1 -β t 1

(a)

10 0

10 1

10 2

10 3

step

0

.

00

0

.

01

0

.

02

0

.

03

0

.

04

0

.

05

Training loss

T

wrm

1

64

256

1024

(b)

10 0

10 1

10 2

10 3

step

10

-

1

10 0

λ H

t

T

wrm

1

64

256

1024

(c)

10 0

10 1

10 2

10 3

step

10

-

2

10 0

10 2

Training loss

T

wrm

1

64

256

1024

(d)

10 0

10 1

10 2

10 3

step

10 1

10 2

λ H

t

T

wrm

1

64

256

1024

and ˆ v t = v t 1 -β t 2 are the bias-corrected moments, ϵ is a small scalar used for numerical stability and P t = (1 -β t 1 ) [ diag ( √ ˆ v t ) + ϵ I ] is the preconditioner.

Linear Warmup: This is defined by the schedule η t = η init +( η trgt -η init ) t / T wrm . The warmup rate is α := ( η trgt -η init ) / T wrm . T wrm = 1 corresponds to constant learning rate. Unless otherwise specified, we set η init = 0 when referring to linear warmup. We propose strategies for selecting η init in Section 6.

Sharpness: The sharpness is defined as the maximum eigenvalue of the Hessian of the loss λ H t := λ max ( ∇ 2 θ L ) , with subscript t indexing the training step. Adaptive optimizers, such as Adam, effectively perform gradient descent in a transformed space determined by ϕ := P 1 / 2 θ (for details, see Appendix B.1). Hence, their stability is determined by the largest eigenvalue of the pre-conditioned Hessian, denoted by λ P -1 H := λ max ( P -1 ∇ 2 θ L ) , rather than the sharpness itself.

Parameterizations in Neural Networks: The mechanism of warmup and its effectiveness is heavily influenced by the network parameterization (see Sections 4 and 5). Standard Parameterization (SP) [33] is a staple in common libraries [28, 3]. Another notable parameterization is the Neural Tangent Parameterization (NTP) [17], which along with SP resides in the kernel learning class at infinite width. Ref. [37] proposed Maximal Update Parameterization ( µ P) which exhibits feature learning at infinite width. Neural network parameterizations significantly impact training dynamics [19].

3 Overview of Training Instabilities and the Self-Stabilization Mechanism

One important underlying mechanism of warmup is intimately tied to training instabilities. These training instabilities, often referred to as 'catapults' [22, 5], arise when the learning rate η exceeds a critical threshold η c , where both η and η c generally change with time. The critical learning rate η c is influenced by a variety of factors, including the choice of optimizer [5, 6], mini-batch size [35, 6], and model properties such as depth, width, parameterization, and initialization [18, 19]. For a detailed overview of instability thresholds, see Appendix B.

When the instability threshold is exceeded ( η > η c ), two cases arise: (i) if the learning rate is higher than the instability threshold but smaller than a maximum stable learning rate (which varies with time), i.e., η c < η < η max, training stabilizes through a self-stabilization process and training continues, (ii) if the learning rate exceeds this maximum stable learning rate η > η max, training experiences severe instabilities. For SGD, these can result in training divergence, characterized by the loss increasing to infinity. For Adam, training may cease, resulting in a training failure, where the loss fails to improve significantly over its initial value, as we demonstrate in Section 5.

The self-stabilization mechanism of GD can be understood through both empirical observations (Figure 1) and a theoretical model. We first describe a model derived by Ref. [7] that effectively captures this phenomenon. The model assumes that the top eigenvector u changes slowly through

training and can be treated as constant and considers a cubic approximation of the GD dynamics around a reference point θ ∗ . The dynamics along the projection x t := u T ( θ t -θ ∗ ) is given by two coupled non-linear equations:

where α := -∇ λ H ·∇ L quantifies the instantaneous change in sharpness and β := ∥∇ λ H ∥ 2 controls the non-linear change in sharpness. In this model, an instability arises when η t > η c = 2 / λ H t . Ref. [7] considered a constant learning rate η and assumed progressive sharpening ( α > 0 ). In contrast, we consider a time-dependent learning rate and allow α to attain both positive and negative values in order to analyze different warmup mechanisms.

The self-stabilization mechanism manifests as a four-step process [22, 7]. Below, we describe the four steps of the self-stabilization mechanism using the above model and the T wrm = 64 trajectories illustrated in Figure 1(c, d):

(1) Approaching instability: Due to increasing learning rate and/or progressive sharpening, training approaches the instability threshold η t = η c = 2 / λ H t . In Figure 1(d), this occurs within the first 10 steps due to increasing learning rate.

(2) Blow up: On exceeding the instability threshold ( η > η c ), Equation (1) predicts exponential growth in x t , empirically manifesting as a sharp increase in loss, as observed in Figure 1(c).

(3) Sharpness reduction: For small enough learning rates, | x t | (and the loss) continues to increase until the higher-order term in the sharpness update equation causes a decrease in sharpness ( x t > α / β ). This is observed as an abrupt decrease in sharpness in Figure 1(d). If the sharpness fails to decrease over extended steps, it may result in training divergence (e.g., see T wrm = 1 case in the same figure).

(4) Return to stability: Once the sharpness has decreased appreciably so that η t λ H t < 2 , stability is restored and the loss eventually decreases.

While the self-stabilization process for more complex optimizers remains poorly understood, a qualitatively similar mechanism is observed in practice, as we will see in the later sections.

4 Warmup Mechanisms of Gradient and Adaptive Methods

This section analyzes the underlying mechanism of warmup through the lens of sharpness dynamics. A key finding is that warmup decreases sharpness in two ways: by allowing a natural sharpness reduction effect at early times and/or by forcing sharpness reduction through training instability at later times.

4.1 Stochastic Gradient Descent

Learning rate warmup is intrinsically tied to sharpness dynamics, as sharpness determines the instability threshold η c . As the learning rate is increased during warmup, training instabilities can be triggered. Assuming the warmup rate is not too high, these instabilities induce a temporary increase in the loss and a decrease in the sharpness to restore stability through the self-stabilization mechanism. Ultimately this allows the model to adapt to the increased learning rate. In other words, a primary goal of warmup is to gradually reduce sharpness, guiding training towards flatter regions that can accommodate training at higher learning rates [10].

However, digging deeper, we find that training has a 'natural' preference for sharpness evolution throughout the training course [19]. Before exceeding the instability threshold ( η < η c ) , training naturally experiences either a progressive increase or decrease in sharpness, as observed in Figure 1, which is unrelated to warmup. For instance, consider the sharpness trajectories with T wrm = 1024 in the above figure. In Figure 1(b), sharpness has a natural preference for increasing, whereas in Figure 1(d), it tends to decrease on its own. This natural sharpness evolution can defined as the sharpness evolution under gradient flow, corresponding to α in Equation (1). The interplay between this natural sharpness evolution and the deliberate intervention of warmup to reduce sharpness can result in completely distinct dynamics. Below, we use the model described by Equation (1) and the experiments in Figure 1 to describe these distinct dynamics.

(C1) Natural Progressive Sharpening ( α > 0 ; top row of Figure 1): The combined effect of naturally increasing sharpness while the learning rate is also being increased results in a 'head-on

collision" at which the instability threshold is exceeded ( η t > η c ). This causes the loss to increase, leading to a decrease in sharpness. Once the sharpness has decreased appreciably, the stability is restored ( η t < η c ). As training proceeds, both sharpness and learning rate continue to increase, again surpassing the instability threshold. This results in a persistent catapult cycle , characterized by η t ≈ 2 / λ H t ≈ η c , for the remainder of the warmup period, as seen in Figure 1(b).

(C2) Natural Sharpness Reduction ( α < 0 ; bottom row of Figure 1): The network is naturally already reducing its sharpness during early training. However, if the learning rate is increased sufficiently quickly, eventually the instability threshold will be reached (akin to a 'rear-end collision"), causing the loss to increase. For small enough learning rates, the increased loss induces a dramatically more pronounced decrease in sharpness than would naturally occur, ultimately restoring stability ( η t < η c ). To exceed the instability threshold again, the learning rate must significantly increase to account for the decreased sharpness, potentially requiring considerable training steps. Consequently, training experiences one or more separated catapults during the warmup phase, as seen in Figure 1(c, d). This contrasts with the progressive sharpening case, where training enters a continuous catapult cycle after reaching the instability threshold for the first time. Notably, training may eventually reach a very flat region of the landscape during warmup, with gradients pointing towards increasing sharpness (e.g., T wrm = 64 in Figure 1(d)). Upon reaching such a region, the dynamics aligns with the natural progressive sharpening scenario.

When natural sharpness reduction is significant (large negative α ), warmup may not need to actively reduce sharpness. Instead, it may 'piggy-back' on the inherent sharpness decrease, resulting in a completely different warmup mechanism, which does not rely on instabilities to facilitate training at higher learning rates.

The above two scenarios can be interpreted as cooperative or competitive dynamics between warmup and the natural evolution of sharpness. When training inherently undergoes sharpness reduction, it cooperates with warmup in decreasing sharpness. Conversely, if the natural trajectory of training is towards increasing sharpness, it opposes the warmup's effort, leading to a persistent cycle of catapults.

The Effect of Warmup Duration: Given a fixed target learning rate η trgt, increasing the warmup duration T wrm delays the point at which training exceeds the instability threshold η c , allowing the sharpness to evolve freely before reaching this point. In the sharpness reduction case, sharpness can significantly decrease by the time this threshold is reached, lowering the need for warmup to decrease sharpness actively. Consequently, increasing T wrm results in catapults that are both delayed and smaller in magnitude, as seen in Figure 1(d). As the catapults become less intense on increasing the warmup duration, the model can train at higher target learning rates without diverging. For extended warmup durations, warmup may not actively reduce sharpness in these sharpness reduction cases and instead it leverages the inherent sharpness decrease.

In the progressive sharpening case, increasing T wrm allows the sharpness to naturally increase. As a result, training exceeds the instability threshold for the first time at a relatively lower learning rate compared to the constant learning rate case. Although warmup has to now undertake more work in decreasing sharpness, it does so in a more gradual manner since increasing the warmup duration amounts to a lower warmup rate η trgt / T wrm . As a result, the fluctuations observed on exceeding the instability threshold are much smaller in magnitude, as seen in Figure 1(a, b).

Small vs. Large Initializations: So far, we have outlined different warmup mechanisms without describing specific conditions that typically exhibit them. Small initializations, such as those using maximal update parameterization ( µ P) [37] in the large width limit or appropriately using normalizing layers (e.g. standard Transformer architectures, see Figure 17 in Appendix E.5), are characterized by a small initial network output. Such initializations start in flat regions where gradients point toward increasing sharpness [19], placing them in the progressive sharpening category (C1). As we will see in Section 5, such initializations may not significantly benefit from warmup as they already start in a flat region. In contrast, large initializations, such as FCNS, CNNs, ResNets with Standard Parameterization (SP) initialized at criticality [29, 31] or Transformers with the last layer-norm removed, undergo an early sharpness reduction, categorizing them into sharpness reduction category (C2). As the primary effect of warmup is to reduce sharpness, we expect such large initializations to considerably benefit from warmup. Notably, large initializations can eventually undergo progressive sharpening at later training stages [18, 19] and adhere to the second mechanism, especially for prolonged warmups.

(a)

10 0

10 1

10 2

10 3

step

0

.

02

0

.

03

0

.

04

Training loss

T

wrm

1

64

256

1024

(b)

10 0

10 1

10 2

10 3

step

10 4

10 5

10 6

10 7

λ P

-

1

H

t

T

wrm

1

64

256

1024

(c)

10 0

10 1

10 2

10 3

step

0

.

1

0

.

2

0

.

3

λ H

t

(d)

10 0

10 1

10 2

10 3

step

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Training loss

T

wrm

1

64

256

1024

(e)

10 0

10 1

10 2

10 3

step

10 4

10 5

10 6

10 7

λ P

-

1

H

t

T

wrm

1

64

256

1024

(f)

10 0

10 1

10 2

10 3

step

10

20

30

40

λ H

t

Natural sharpness change provides an intuitive way of determining whether an initialization is 'small' or 'large': if training from a given initialization exhibits sharpness reduction, it suggests the existence of naturally flatter initializations that could be chosen instead. This observation is particularly helpful for analyzing Adam in the later sections and motivates modifications to improve it.

4.2 Stochastic Gradient Descent with Momentum (SGD-M)

The warmup mechanism of SGD-M, while at its core is similar to that of vanilla SGD, has a few subtleties. Here we summarize the major differences, leaving details to Appendix E.2.

During early training, the loss may decrease non-monotonically on incorporating momentum, even at small learning rates. Such oscillations are also observed when quadratic loss functions are optimized using GD with momentum [11]. These oscillations make it challenging to differentiate between warmup-induced catapults and fluctuations in loss due to the intrinsic effects of momentum. Nevertheless, we can still observe loss spikes correlated with an abrupt decrease in sharpness at large learning rates, as detailed in Appendix E.2.

Additionally, the instability threshold η c itself evolves differently during training. It changes from 2 / λ H 0 at initialization to (2+2 β ) / λ H t later in training. Moreover, the late-time instability threshold is significantly influenced by the batch size, exhibiting a much smaller value than SGD for the same batch size. These properties make it more challenging to analyze the training dynamics of SGD with momentum. Nonetheless, the fundamental warmup mechanisms closely mirror the vanilla SGD case. We leave a more detailed analysis of the early training dynamics of SGD-M for future studies.

4.3 Adaptive Gradient Methods (Adam)

Adaptive optimizers effectively perform gradient descent in a transformed space given by ϕ = P 1 / 2 θ , as we show in Appendix B.1. This analysis suggests that their local stability should be determined by the largest eigenvalue of ∇ 2 ϕ L = P -1 ∇ 2 θ L , which we refer to as the pre-conditioned Hessian. Indeed, this is what we observe in Figure 2, which shows the training loss, pre-conditioned sharpness λ P -1 H , and sharpness trajectories for full batch Adam. In these figures, sharpness is significantly smaller than its instability threshold (2+2 β 1 ) / η t ≈ 4000 , indicating that sharpness does not determine stability. Instead, loss catapults are associated with λ P -1 H exceeding its instability threshold.

The pre-conditioned sharpness starts high for both progressive sharpening (simpleµ P) and sharpness reduction (SP) scenarios considered in the previous section. For simplicity, we considered a simpler version of µ P, detailed in Appendix D.2.1. In particular, for µ P models, λ P -1 H 0 ∼ 10 5 despite being initialized in a flat region as measured by sharpness, while for SP models, λ P -1 H 0 ∼ 10 6 . These large initial values arise because P 1 = (1 -β 1 ) [ diag( g 2 0 ) + ϵ I ] and g 0 has components that are near zero.

(a)

8.6e-01

1.7e+00

3.4e+00

6.9e+00

1.4e+01

2.7e+01

5.5e+01

1.1e+02

2.2e+02

4.4e+02

8.8e+02

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

89.63

91.10

91.99

92.59

92.71

92.63

89.36

91.05

92.13

92.74

93.24

92.59

89.42

91.40

92.20

92.54

92.92

89.57

90.97

91.75

92.44

93.04

93.14

89.54

91.20

92.25

92.62

92.88

93.07

89.61

90.90

91.92

92.63

92.90

92.89

93.52

89.50

91.19

91.95

92.42

93.04

93.28

93.58

92.88

89.61

91.16

92.04

92.45

93.13

93.13

93.40

93.54

93.45

89.59

90.61

92.01

92.32

92.91

93.17

93.57

93.47

93.46

89.45

90.99

92.16

92.60

93.10

93.14

93.40

93.49

93.58

93.33

89.62

90.86

91.88

92.66

92.96

93.07

93.32

93.65

93.38

93.27

89.50

91.00

91.94

92.52

93.24

93.11

93.32

93.45

93.46

93.47

93.03

89.39

90.88

92.16

92.43

93.10

93.18

93.44

93.51

93.56

µ

P MSE SGD CIFAR-10

50

55

60

65

70

75

80

85

90

(b)

3.9e-01

7.8e-01

1.6e+00

3.1e+00

6.2e+00

1.2e+01

2.5e+01

5.0e+01

9.9e+01

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

92.55

93.11

93.23

93.03

92.38

91.81

10.37

92.39

93.24

93.16

93.01

92.57

92.01

90.79

10.37

92.88

92.94

93.40

92.99

92.59

92.40

13.04

92.61

92.76

93.22

93.21

92.77

92.26

10.37

92.72

93.34

93.28

93.42

92.47

91.75

87.94

10.37

92.61

92.93

92.99

92.91

92.43

85.32

10.37

92.74

92.92

93.34

93.02

92.58

10.24

92.61

92.86

93.29

93.00

92.97

92.41

10.37

92.75

93.11

92.87

92.86

92.24

91.71

90.33

92.34

93.09

93.06

93.02

93.03

10.24

91.62

10.37

92.30

93.30

93.29

92.99

92.81

90.95

91.95

92.45

92.68

93.29

93.14

92.99

92.89

92.29

91.46

19.21

92.31

92.97

92.98

93.06

93.27

93.04

92.50

µ

P XENT SGD CIFAR-10

50

55

60

65

70

75

80

85

90

(c)

4.0e-03

7.9e-03

1.6e-02

3.2e-02

6.4e-02

1.3e-01

2.5e-01

5.1e-01

1.0e+00

2.0e+00

4.1e+00

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

45.11

50.08

56.79

63.22

45.10

49.89

57.08

63.73

45.08

50.02

56.89

63.50

45.01

50.12

56.78

63.20

68.12

45.11

49.84

56.67

64.03

70.82

45.09

50.01

56.88

63.15

70.73

77.23

45.15

49.72

56.45

64.35

70.99

78.16

83.26

45.21

49.79

56.41

64.32

71.25

78.44

82.75

87.57

45.12

49.93

56.77

64.29

71.85

78.09

82.59

45.32

49.97

56.96

64.34

71.51

78.78

83.67

87.21

45.11

49.93

56.76

64.74

71.40

78.58

83.53

87.52

89.56

45.06

49.77

56.86

64.56

71.33

78.50

83.73

87.31

90.12

91.11

45.10

49.68

56.65

64.22

71.37

78.86

84.00

87.38

89.98

91.44

92.09

WRN-16-4 SP MSE SGD CIFAR-10

50

55

60

65

70

75

80

85

90

(d)

1.8e-03

3.7e-03

7.3e-03

1.5e-02

2.9e-02

5.9e-02

1.2e-01

2.3e-01

4.7e-01

9.4e-01

1.9e+00

3.7e+00

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

76.83

81.73

85.52

88.53

90.63

91.52

92.19

10.24

76.57

81.88

85.68

88.25

89.77

91.10

92.15

92.12

10.24

76.82

81.90

85.62

88.52

90.28

91.31

91.95

92.20

11.51

76.39

81.88

85.61

88.13

89.75

91.41

91.63

92.12

76.54

81.75

85.87

88.09

90.04

91.09

91.88

92.23

76.42

81.89

85.78

88.41

90.20

91.25

91.87

92.56

76.33

81.82

85.35

88.34

90.10

91.12

92.36

92.48

93.28

76.48

81.65

85.69

88.50

90.34

91.37

91.96

93.08

92.88

76.73

81.76

85.73

88.20

89.89

91.52

91.95

92.50

92.67

93.01

76.63

81.47

85.33

88.40

90.21

91.52

92.31

92.71

92.50

92.72

76.28

81.98

85.17

88.31

90.08

91.26

92.06

92.52

92.88

92.80

10.30

76.09

81.29

85.78

88.22

90.31

91.46

92.29

92.45

92.57

93.22

92.94

75.96

81.39

85.65

88.44

90.04

91.32

92.09

92.96

92.91

92.86

93.03

92.50

WRN-16-4 SP XENT SGD CIFAR-10

50

55

60

65

70

75

80

85

90

The large λ P -1 H 0 can lead to training failures if the learning rate does not start sufficiently small. We put forward modifications to improve Adam in Section 6; here we continue characterizing the warmup mechanisms of Adam.

Given that the pre-conditioned sharpness consistently starts high and decreases during early training, this behavior can be viewed as an extreme example of the natural sharpness reduction scenario (C2) described in the previous section. Training Adam at high initial learning rates without warmup can cause large catapults, as seen in Figure 2(d), potentially leading to training failures. Increasing the warmup duration allows the pre-conditioned sharpness to naturally decrease. This prevents the loss from spiking during early training and avoids training failures. In the later stages of training, the pre-conditioned sharpness may continue reducing or exhibit progressive sharpening. From here on, the dynamics follows the warmup mechanisms discussed in the previous sections, with sharpness replaced with pre-conditioned sharpness. Similar to the momentum case, Adam's stability threshold at late training times significantly decreases for smaller batch sizes [6], also shown in Appendix E.4.

5 Impact of Warmup on Training and Generalization

Here we investigate the impact of warmup on training efficacy and generalization by disentangling the role of η trgt and T wrm. Our key findings are that generalization capability is primarily determined by η trgt and that Adam is particularly sensitive to large catapults. The role of increasing T wrm is to (i) allow the network to tolerate larger η trgt, and (ii) move training further away from the divergence (failure) boundary, leading to a marginal improvement in generalization.

Experimental Setup: We consider WideResNets (WRNs) and Transformers (LM) parameterized in either SP or µ P. WRNs are trained on CIFAR-10, CIFAR-100, and Tiny-ImageNet, employing data augmentation. Transformers are trained on the next token prediction task using the WikiText dataset. These models are trained with MSE or cross-entropy (xent) loss functions using SGD or Adam optimizers for a fixed training budget of T = 10 5 steps unless otherwise specified. Training begins with a linear warmup phase from η init = 0 to η trgt over T wrm steps. After warmup, training continues at η trgt for the remaining training budget. In some cases, following the warmup period, we decrease the learning rate using cosine decay [24]. Further details are provided in Appendix D.

(a)

1.0e-04

3.0e-04

6.0e-04

8.0e-04

1.0e-03

3.0e-03

6.0e-03

1.0e-02

η

trgt

1

4

16

64

128

256

512

1024

2048

T

wrm

4.38

4.00

3.90

3.91

4.11

5.98

6.06

6.13

4.38

4.00

3.90

3.89

3.90

5.56

6.06

6.20

4.38

4.01

3.89

3.88

3.87

5.32

6.07

6.07

4.38

4.00

3.89

3.87

3.86

5.30

5.43

6.08

4.38

4.00

3.89

3.87

3.86

5.21

5.44

5.56

4.38

4.00

3.89

3.87

3.86

5.27

5.44

5.48

4.39

4.00

3.90

3.86

3.86

5.27

5.44

5.49

4.39

4.01

3.89

3.86

3.86

3.87

5.40

5.55

4.39

4.01

3.90

3.87

3.86

3.82

5.33

5.44

LM-4-768, wikitext-103, Crossentropy, Adamw

3

.

9

4

.

0

4

.

1

4

.

2

4

.

3

4

.

4

4

.

5

(b)

1.0e-04

3.0e-04

6.0e-04

8.0e-04

1.0e-03

3.0e-03

6.0e-03

1.0e-02

η

trgt

1

4

16

64

128

256

512

1024

2048

T

wrm

4.42

4.02

3.90

3.88

3.86

3.82

5.42

5.97

4.42

4.02

3.90

3.88

3.86

3.84

5.49

6.03

4.42

4.02

3.90

3.87

3.86

3.81

5.44

5.56

4.42

4.02

3.90

3.87

3.86

3.82

3.81

5.90

4.42

4.02

3.90

3.87

3.85

3.81

3.82

3.85

4.42

4.02

3.90

3.88

3.86

3.81

3.81

5.50

4.42

4.02

3.90

3.88

3.86

3.81

3.81

5.44

4.42

4.02

3.90

3.88

3.86

5.09

3.80

3.82

4.41

4.03

3.91

3.88

3.86

3.82

5.33

3.80

LM-4-768, wikitext-103, Crossentropy, GI-Adamw

3

.

9

4

.

0

4

.

1

4

.

2

4

.

3

4

.

4

4

.

5

5.1 Stochastic Gradient Descent (SGD)

Figure 3 presents heatmaps that show the best test accuracy achieved during training, plotted in the η trgt -T wrm plane for different parameterizations and loss functions. These phase diagrams of warmup also show the convergence-divergence boundary, with empty cells indicating training divergences, illustrating the interplay between warmup duration and the maximum trainable η trgt . Below, we discuss the crucial insights these results provide into warmup's role in training dynamics.

Longer Warmup Facilitates Training at Higher Learning Rates: These phase diagrams reveal that an extended warmup duration facilitates training at higher target learning rates. This benefit is particularly noticeable for large initializations (like SP) and MSE loss. In contrast, the advantage is less pronounced when using cross-entropy loss and smaller initializations (like µ P). The diminished benefit for µ P is likely due to its initialization in a relatively flat region of the loss landscape, which can already facilitate training at higher learning rates at initialization. This consistent increase in maximum η trgt with warmup durations can be understood through the lens of warmup mechanisms described in the previous section. As observed in Figure 1, when the warmup duration is increased, loss catapults occurring on surpassing the instability thresholds become milder. This effectively pushes the divergent boundary to higher learning rates.

Final Performance Primarily Depends on the Target Learning Rate: A closer look into these phase diagrams reveals that, slightly away from the divergent boundary, the test accuracy primarily depends on the target learning rate and nominally on the warmup duration. Based on the model performance, we can categorize these phase diagrams into two distinct cases: (i) models that fail to achieve optimal performance when trained with a constant learning rate (e.g., Figure 3(c)), and (ii) models that attain optimal performance without warmup (e.g., Figure 3(b)). The first scenario corresponds to models with large initializations. Increasing the warmup duration improves performance by facilitating training at higher learning rates. Yet, similar performance is observed for different warmup durations, suggesting that the primary gain comes from the target learning rate, rather than the duration itself. The second case arises for flat initializations, which can already train at large learning rates, and resultantly the optimal performance is already achieved without warmup. While increasing warmup duration facilitates training at even higher learning rates, it does not enhance performance. Nevertheless, it does broaden the range of optimal learning rates, reducing the need for precise tuning of the target learning rate, and making training more practical and robust. We conclude that warmup can serve two key purposes: (i) it can significantly improve model performance in large initialization cases, and (ii) extend the range of optimal target learning rates for small initializations, making it easier to tune the target learning rate. In Appendix F.2, we demonstrate that these results hold on incorporating momentum and employing cosine learning rate decay.

5.2 Adaptive Gradient Methods (Adam)

Figure 4(a) shows the warmup phase diagram of Adam. Increasing the warmup duration enables training at higher learning rates by allowing the pre-conditioned sharpness to decrease naturally, thereby reducing the severity of catapults. These large catapults, which may persist in Adam's memory, can lead to performance degradation and training failures. Thus, in addition to facilitating

(a)

1.4e-03

2.8e-03

5.6e-03

1.1e-02

2.2e-02

4.5e-02

9.0e-02

1.8e-01

3.6e-01

7.2e-01

1.4e+00

2.9e+00

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

2

3

3

3

1

1

1

1

3

6

7

7

1

1

1

1

6

11

14

15

1

1

1

1

12

22

27

30

31

1

1

1

1

24

44

54

60

62

63

1

1

1

1

48

88

108

119

123

126

1

1

1

1

96

176

216

237

246

251

1

1

1

1

192

352

432

473

492

502

507

1

1

1

1

384

704

864

945

984

1004

1014

1

1

1

1

768

1408

1728

1889

1968

2008

1

1

1

1

1536

2816

3456

3777

3936

4016

4056

4076

T

reach

1000

2000

3000

4000

(b)

1.4e-03

2.8e-03

5.6e-03

1.1e-02

2.2e-02

4.5e-02

9.0e-02

1.8e-01

3.6e-01

7.2e-01

1.4e+00

2.9e+00

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

0.0

0.0

0.0

0.0

0.0

0.0

0.0

0.0

-0.5

-0.5

-0.5

0.0

0.5

1.0

1.0

1.0

1.0

1.0

1.0

1.5

1.0

0.5

1.0

1.0

4.5

4.5

4.5

5.0

3.5

1.0

0.5

1.0

12.0

12.0

12.0

12.5

8.0

3.5

1.0

0.5

27.5

27.5

27.5

28.0

17.5

8.0

3.5

1.0

0.5

59.0

59.0

59.0

59.5

37.0

17.5

8.0

2.5

1.0

0.5

122.5

122.5

122.5

123.0

76.5

37.0

17.5

7.0

3.5

1.0

250.0

250.0

250.0

250.5

156.0

76.5

37.0

16.5

8.0

3.5

505.5

505.5

505.5

506.0

315.5

156.0

76.5

36.0

17.5

8.0

3.5

1017.0

1017.0

1017.0

1017.5

635.0

315.5

156.0

75.5

37.0

17.5

8.0

2040.5

2040.5

2040.5

2041.0

1274.5

635.0

315.5

155.0

76.5

37.0

4088.0

4088.0

4088.0

4088.5

2554.0

1274.5

635.0

314.5

156.0

76.5

37.0

17.5

T

save

0

1000

2000

3000

4000

training at higher rates similar to SGD, warmup further improves Adam's performance by addressing its vulnerability to large catapults, justifying its widespread use with Adam. Below, we discuss the distinct properties of Adam phase diagrams in detail.

Training Failures of Adam: Remarkably, we find that models trained with Adam always exhibit training failures rather than divergences where the loss grows without bound, as further demonstrated in Appendix G. In cases of training failure, we often observed that certain layers or residual blocks output zero, leading to vanishing gradients. This implies that the model gets stuck at a critical point and is unable to train further. Understanding this unexpected phenomenon requires further study.

Performance Degradation prior to Failure Boundary: Test accuracy in these phase diagrams declines well before the failure boundary, in contrast to SGD where optimal learning rates are observed near the divergence boundary. This discrepancy stems from Adam's property of retaining a memory of gradient magnitudes. At large learning rates, along with the loss, the gradients spike during early training, as seen in Figure 28 in Appendix G. While the gradients decrease after a few steps, the second moment v remains large for an extended period, leading to a small effective learning rate ηP -1 . As a result, training struggles to escape high-loss regions. Therefore, a longer warmup is more beneficial for Adam compared to SGD, as it is crucial to stay away from the failure boundary.

6 Improved Hyperparameter Initialization Schemes for Optimizers

Initial Learning Rate Selection for Warmup: Setting the initial learning rate to η init = 0 is common practice in warmup [27, 8]. Our analysis reveals that the primary effect of warmup is to facilitate training at higher learning rates by annealing sharpness (or pre-conditioned sharpness for Adam). From this perspective, starting with η init = 0 appears suboptimal, as it can significantly delay the learning rate from exceeding the instability threshold, thus delaying the primary effect of warmup.

An effective strategy involves setting η init = η c to induce loss increase and thereby sharpness decrease right from initialization. We introduce a straightforward search method that only uses forward passes to estimate the initial critical learning rate η c . The method consists of two stages: (i) an exponential search, starting from an initial guess η 0 , iteratively multiplies η 0 by a factor k > 1 until the loss increases. This identifies an interval [ η lwr , η uppr ] containing η c , (ii) a binary search further narrows down [ η lwr , η uppr ] by evaluating the loss at the midpoint η mid = ( η lwr + η uppr ) / 2 . If the loss increases, η uppr is updated to η mid; otherwise, η lwr is set to η mid. This process is repeated until the loss in the next step L ( θ 1 ) satisfies the condition L ( θ 1 ) < L ( θ 0 )(1 + δ ) , for some hyperparameter δ > 0 . For details, see Appendix B.3.

By setting η init = η c , training can achieve the target learning rate earlier. Consider the modified warmup schedule: η t = η init + η trgt ( t / T wrm ) , which attains η trgt in T reach = T wrm (1 -η c / η trgt ) steps, saving T wrm ( η c / η trgt ) steps. Incorporating the computational cost of additional forward passes T fp required for estimating η c ( ∼ 10 in number), and noting that one training step approximately equates to two forward passes, the net computational savings is T save = T wrm ( η c / η trgt ) -T fp / 2 . Figure 5 demonstrates how T reach and T save vary with the T wrm and η trgt. For η trgt < η c , the target learning rate is reached in a single step, nearly saving the entire duration of the warmup, whereas for ( η trgt > η c ) , starting η init ≳ η c can save up to half of the allocated warmup duration, although this saving diminishes on approaching the divergent/failure boundary.

(a)

10 0

10 1

10 2

10 3

step

0

.

1

0

.

2

0

.

3

Training loss

T

wrm

1

64

256

1024

(b)

10 0

10 1

10 2

10 3

step

10 4

10 5

10 6

10 7

λ P

-

1

H

t

T

wrm

1

64

256

1024

(c)

10 0

10 1

10 2

10 3

step

10

15

20

25

30

λ H

t

It is worth noting that there can be instances where there is no loss catapult at initialization. In our experiments, this only occurs for Transformers trained using SGD. In such scenarios, the prescribed approach is not applicable and one can resort to heuristics, such as setting the initial learning rate to a fraction of the maximum stable learning rate, such as η init = η trgt / 10 .

GI-Adam: Improving Adam by Initializing The Second Moment using Gradients

In Section 4.3, we observed that the pre-conditioned sharpness for Adam starts at a high value, even for low sharpness initializations like µ P, and can lead to training failures at large learning rates. We propose Gradient Initialized Adam (GI-Adam), which initializes the second moment using the gradient squared, v 0 = g 2 0 . In Appendix I.2, we show that a bias correction is not required when the second moment is initialized using the gradients. As a result, GI-Adam can be viewed as standard Adam with an automated warmup given by η t = η trgt √ 1 -β t 2 .

This simple trick reduces the initial pre-conditioned sharpness by around two orders of magnitude (more precisely by a factor of √ 1 -β 2 ) at initialization, preventing large catapults, as illustrated in Figure 6 (c.f. Figure 2(d-f)). Moreover, it consistently shows improvement over standard Adam across datasets and prevents training failures by pushing the training failure boundary to higher η trgt , as shown in Figure 4(b). We provide additional results for different datasets in Appendix F.3. In Appendix F.4, we show that GI-Adam consistently performs on par or better than RAdam [23] while offering a simple modification to Adam.

To further assess that the primary cause of instability during early training is the large pre-conditioned sharpness, we randomly initialize v 0 but with the same norm as the gradients at initialization. Like GI-Adam, this also results in improved performance as shown in Appendix I.3.

We further reduce the pre-conditioner size by removing bias correction for the first moment, referred to as Flat-Adam. As demonstrated in Appendix I.4, this modification eliminates the initial decrease in pre-conditioned sharpness. We leave a comprehensive evaluation of Flat-Adam for future work.

7 Discussion

Our analysis provides new insights into the role of warmup across optimizers and parameterizations. We found compelling evidence that the primary effect of warmup is to facilitate training at higher learning rates and stabilizing the training dynamics by keeping it away from the failure (divergence) boundary. Looking under the hood, we found a variety of underlying mechanisms, which also suggested several improvements for hyperparameter initialization. In Appendix A, we provide practical guidance for practitioners on choosing the warmup duration.

Our analysis also motivates a potential parameter-free warmup strategy, which we refer to as persistent catapult warmup . The central idea behind this strategy is to repeatedly induce catapults aimed to progressively reduce sharpness, thereby facilitating training at higher learning rates. We present encouraging preliminary results in Appendix C and defer further development to future work.

The maximum learning rate can be written as η max = c max / λ P -1 H . Here we showed how warmup effectively decreases λ P -1 H , which is a local measure of sharpness. There is also another possible effect of warmup, that it can cause an increase in c max, which can be viewed as a more non-local measure of sharpness. Further analysis is required to understand how warmup helps increase c max.

Limitations: Our experiments were conducted on relatively small-scale datasets and models, and further investigations are needed to understand the generalizability of our findings to larger-scale settings. For Adam, we did not explore the dependence on hyperparameters β 1 , β 2 , ϵ .

Acknowledgments and Disclosure of Funding

We thank Tianyu He, Darshil Doshi, Andrey Gromov, Dan Roberts, Jeremy Cohen, Jonas Geiping, Yuqing Wang and Alex Damian for discussions and comments on the draft. The authors acknowledge the University of Maryland supercomputing resources ( http://hpcc.umd.edu ) made available for conducting the research reported in this paper. This work is supported in part by NSF DMR-2345644 (MB).

References

Tiny imagenet challenge. https://cs231n.stanford.edu/reports/2017/pdfs/930. pdf .

Stephen P Boyd and Lieven Vandenberghe. Convex optimization . Cambridge university press, 2004.

James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.

Sébastien Bubeck et al. Convex optimization: Algorithms and complexity. Foundations and Trends® in Machine Learning , 8(3-4):231-357, 2015.

Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations , 2021.

Jeremy M. Cohen, Behrooz Ghorbani, Shankar Krishnan, Naman Agarwal, Sourabh Medapati, Michal Badura, Daniel Suo, David Cardoze, Zachary Nado, George E. Dahl, and Justin Gilmer. Adaptive gradient methods at the edge of stability, 2022.

Alex Damian, Eshaan Nichani, and Jason D. Lee. Self-stabilization: The implicit bias of gradient descent at the edge of stability. In The Eleventh International Conference on Learning Representations , 2023.

DeepMind, Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Claudio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, Iurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Laurent Sartran, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Miloš Stanojevi´ c, Wojciech Stokowiec, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020.

Emily Dinan, Sho Yaida, and Susan Zhang. Effective theory of transformers at initialization, 2023.

Justin Gilmer, Behrooz Ghorbani, Ankush Garg, Sneha Kudugunta, Behnam Neyshabur, David Cardoze, George Edward Dahl, Zachary Nado, and Orhan Firat. A loss curvature perspective on training instabilities of deep learning models. In International Conference on Learning Representations , 2022.

Gabriel Goh. Why momentum really works. Distill , 2017.

Akhilesh Gotmare, Nitish Shirish Keskar, Caiming Xiong, and Richard Socher. A closer look at deep learning heuristics: Learning rate restarts, warmup and distillation. In International Conference on Learning Representations , 2019.

Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in 1 hour. arXiv preprint arXiv:1706.02677 , 2017.

Kaiming He, X. Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , pages 770-778, 2016.

Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. Flax: A neural network library and ecosystem for JAX, 2020.

Dan Hendrycks and Kevin Gimpel. Bridging nonlinearities and stochastic regularizers with gaussian error linear units. CoRR , abs/1606.08415, 2016.

Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 31. Curran Associates, Inc., 2018.

Dayal Singh Kalra and Maissam Barkeshli. Phase diagram of early training dynamics in deep neural networks: effect of the learning rate, depth, and width. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

Dayal Singh Kalra, Tianyu He, and Maissam Barkeshli. Universal sharpness dynamics in neural network training: Fixed point analysis, edge of stability, and route to chaos. arXiv preprint arXiv:2311.02076 , 2023.

Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR) , San Diego, CA, USA, 2015.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.

Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.02218 , 2020.

Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. In International Conference on Learning Representations , 2020.

Ilya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts. In International Conference on Learning Representations , 2017.

Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In International Conference on Learning Representations , 2017.

Kevin P Murphy. Probabilistic machine learning: an introduction . MIT press, 2022.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, highperformance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'AlchéBuc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems , 32, 2019.

Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. In NIPS , 2016.

Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.

Daniel A. Roberts, Sho Yaida, and Boris Hanin. The Principles of Deep Learning Theory . Cambridge University Press, 2022. https://deeplearningtheory.com .

Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. In Katrin Erk and Noah A. Smith, editors, Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 1715-1725, Berlin, Germany, August 2016. Association for Computational Linguistics.

Jascha Narain Sohl-Dickstein, Roman Novak, Samuel S. Schoenholz, and Jaehoon Lee. On the infinite width limit of neural networks with a standard parameterization. ArXiv , abs/2001.07301, 2020.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Lei Wu, Chao Ma, et al. How sgd selects the global minima in over-parameterized learning: A dynamical stability perspective. Advances in Neural Information Processing Systems , 31, 2018.

Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 10524-10533. PMLR, 13-18 Jul 2020.

Greg Yang and Edward J. Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 11727-11737. PMLR, 18-24 Jul 2021.

Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hyperparameter transfer. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , 2021.

Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. 2017.

Aston Zhang, Zachary C Lipton, Mu Li, and Alexander J Smola. Dive into deep learning . Cambridge University Press, 2023.

Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. In International Conference on Learning Representations , 2018.

A Practical Guidance for Practitioners

How to Select the Warmup Duration? Given a target learning rate η trgt, if the training loss during the warmup period exhibits large instabilities (loss spikes), the warmup duration T wrm should be increased until such instabilities are sufficiently small. This effectively moves training away from the divergent / failure boundary, as illustrated in Figure 3. This is particularly crucial for Adam, as large instabilities can be detrimental and lead to considerable performance degradation without divergence, as discussed in Section 5.2.

How to Select the Target Learning Rate? As the primary effect of warmup is to anneal sharpness by increasing the learning rate beyond the instability threshold, it suggests that the target learning rate should be at least greater than the instability threshold at initialization.

When to Decay the Learning Rate? Figure 21 suggests that employing learning rate decay at small learning rates can result in performance degradation for a fixed training budget. Therefore, the learning rate should be decayed at large target learning rates only. The underlying intuition is that we use large target learning rates to train in a flat region of the landscape. However, these large learning rates restrict training to go into sharper regions of the basin and learning rate decay helps.

Leveraging µ P for Effecient Training: Our analysis suggests that the primary role of warmup facilitates training at higher learning rates by gradually reducing sharpness. Given this perspective, beginning training with flat initializations, such as µ P, is advantageous. These initializations might allow for achieving optimal performance without the need for warmup, as observed in Figure 3.

B Instability Thresholds

B.1 Instability Thresholds for Adaptive Optimizers

For adaptive optimizers, such as Adam, a pre-conditioner P t +1 multiplies the gradients in the update equation

Next, we define ϕ t := P 1 / 2 t +1 θ t , then the gradient and Hessian are ∇ ϕ L = P -1 / 2 ∇ θ L and ∇ 2 ϕ L := P -1 ∇ 2 θ L . As a result, the update equations for ϕ is given by

We observe that the update equation of an adaptive optimizer mirrors that of SGD under a change of variables. Using the stability arguments of SGD, we conclude that λ max ( P -1 t +1 ∇ 2 θ L ( θ t )) determines the stability.

B.2 Overview of Instability Thresholds

Lewkowycz et al. [22] showed that for wide networks in NTP/SP trained with MSE loss and SGD, this critical learning rate is 2 / λ H 0 early in training. Further investigation by Kalra and Barkeshli [18] demonstrated that sharpness reduction during early training causes η c to increase with depth and 1 / width . In such scenarios, η c can be as large as 40 / λ H 0 . Cohen et al. [5] demonstrated that sharpness at late training times for GD with momentum coefficient β oscillates above (2+2 β ) / η , suggesting η c ≳ (2+2 β ) / λ H t at late training times. Expanding on this, Cohen et al. [6] analyzed adaptive optimizers and found that for Adam, the pre-conditioned sharpness λ P -1 H oscillates around (2+2 β 1 ) / η (1 -β 1 ) at late training times. The instability threshold also depends on the mini-batch size [35] and is often observed to be smaller than their full batch counterparts [5, 6].

B.3 Estimating the Instability Threshold

This section describes the method for estimating the instability threshold η c at initialization (or generically, any point θ t ) using only forward passes. The method consists of two stages:

Exponential Search: An exponential search, starting from an initial guess η 0 , iteratively multiplies η 0 by a factor k = 2 until the loss increases. This identifies an interval [ η lwr , η uppr ] containing η c . The detailed algorithm is described in Algorithm 1. Unless specified, we use η 0 = 10 -4 as our initial guess. If the loss already increases at the initial guess η 0 , we set η init = η 0 .

Algorithm 1 Exponential Search

Binary search: A binary search further narrows down [ η lwr , η uppr ] by evaluating the loss at the midpoint η mid = ( η lwr + η uppr ) / 2 . If the loss increases, η uppr is updated to η mid; otherwise, η lwr is set to η mid. This process is repeated until the loss in the next step L ( θ 1 ) satisfies the condition L ( θ 1 ) < L ( θ 0 )(1 + δ ) , for some δ > 0 . The algorithm is detailed in Algorithm 2. In all our experiments, we set δ = 0 . 1 .

Algorithm 2 Binary Search

While both η 0 and δ are additional hyperparameters, the method does not heavily depend on these choices. A poor initial guess of η 0 would only take a few more iterations to find an interval [ η lwr , η uppr ] .

10 0

10 1

10 2

10 3

step

10

-

1

10 0

Training loss

T

wrm

64

256

1024

(a)

10 0

10 1

10 2

10 3

step

10 0

10 1

10 2

λ H

t

T

wrm

64

256

1024

(b)

Meanwhile, any small value of δ ∈ (0 , 1] is effective in finding η c , as small initial loss spikes have minimal impact on the overall dynamics. Note that for Adam, a small δ ( ∼ 0 . 01 ) has to be selected to ensure that we do not observe large catapults.

C Persistent Catapult Warmup

Our analysis also motivates a potential parameter-free warmup strategy, which we refer to as persistent catapult warmup . The central idea behind this strategy is to repeatedly induce catapults aimed to progressively reduce sharpness (or pre-conditioned sharpness), thereby facilitating training at higher learning rates. Given a target learning rate η trgt, the strategy consists of the following steps:

Start with a 'stable' reference point θ ∗ , defined as a point where the loss decreases in the next step and estimate the interval [ η lwr , η uppr ] containing η c , as described in Appendix B.3.

Induce a catapult by increasing the learning rate to η = η uppr .

Continue training and wait until the loss falls below the reference point, i.e., L ( θ t ) < L ( θ ∗ ) . This new point now becomes the stable reference point.

Repeat the above steps until the target learning is achieved, i.e., η = η trgt .

Here, the initial stable reference point is the model's initialization. The detailed algorithm is described in Algorithm 3 in Appendix C.

Figure 7 compares persistent catapult warmup (shown in black) with linear warmup. The persistent catapult warmup facilitates training at higher learning rates without the need to specify warmup duration. Since η c serves as an indicator of sharpness, persistent catapult warmup utilizes the local sharpness information to automatically determine the warmup rate, resulting in an adaptive non-linear warmup. This adaptive approach eliminates the need for manual tuning of the warmup duration, allowing for a more efficient and effective warmup.

Although persistent catapult warmup is a promising approach to warmup, it requires specifying how large a catapult should be induced, which introduces another hyperparameter. Nevertheless, persistent catapult warmup motivates the development of parameter-free warmup strategies that could simplify the training process. We leave further development of parameter-free warmup to future work.

D Experimental Details

This section provides additional experimental details. All models were implemented using the JAX [3], and Flax libraries [15]. The key results can be reproduced using the GitHub repo: https: //github.com/dayal-kalra/why-warmup .

Algorithm 3 Persistent Catapult Warmup

D.1 Datasets Details

D.1.1 Image Classification Tasks

We consider standard image classification datasets such as CIFAR-10, CIFAR-100 [21], and TinyImageNet [1]. The images are normalized to have zero mean and unit variance. For MSE loss, we use one-hot encoding for the labels.

Data augmentation: For various image classification tasks, we employ data augmentation techniques, applied in the following order: random horizontal flips, random cropping, and mixup [41].

D.1.2 Language Modeling Tasks

We consider the next token prediction task on the Wikitext-2 and Wikitext-103 datasets [25]. The Wikitext-2 dataset consists of ∼ 2 Mtokens, whereas the Wikitext-103 dataset has ∼ 0 . 1 B tokens. We use Byte Pair Encoding (BPE) tokenizer [32] with a Whitespace pre-tokenizer. For Wikitext-103 experiments, we considered a standard vocabulary size of 50 , 257 , whereas our initial Wikitext-2 experiments employed a smaller vocabulary size of 4096 .

D.2 Model Details

This section describes the models considered, including their parameterization and initialization details. We adopt parameterizations outlined in Table 9 of Ref. [38]. Unless otherwise specified, we employ ReLU non-linearities and initialize the weights with a truncated normal distribution 1 , with a variance σ 2 w = 2 . 0 in appropriate parameterizations (details below), except for the last layer, which has a weight variance of σ 2 w = 1 . 0 . All biases are initialized to zeros.

D.2.1 Parameterizations

Standard Parameterization (SP): For SP, the weights are initialized with truncated Gaussian distribution N (0 , σ 2 w / fanin ) and the biases are initialized to zero.

Maximal Update Parameterization ( µ P): For µ P, different schemes are employed for the intermediate and last layers. The intermediate layers are initialized using N (0 , σ 2 w / fanout ) and the layer outputs are scaled by the factor √ fanout / fanin . In comparison, the layer weights are initialized with N (0 , σ 2 w / fanin ) , and the final output is rescaled by the factor √ 1 / fanin . Conveniently, for SGD, the learning rate does not scale with width in the above µ P formulation. In comparison, for Adam, the learning rate corresponding to input, intermediate, and output layers are rescaled by the factors 1 / √ fanout , 1 / √ fanin and 1 / fanin . Since we are utilizing µ P only to obtain flat initializations, we omit the additional scaling of the learning rate for Adam in some experiments (e.g., Figure 2). As a result, the instability threshold is only dependent on the target learning rate η trgt during late training, rather than on the largest learning rate across layers. We refer to this parameterization as 'simpleµ P' for Adam.

D.2.2 Architectures

Fully Connected Networks (FCNs): We consider fully connected networks with a constant width of n and a depth of d layers. These networks are denoted by FCNd -n . Unless specified, we considered d = 4 layer FCNs with width n = 512 .

WideResNets (WRNs): We consider WideResNets [39] with d layers, S stages, and a widening factor of k , denoted by WRNd -k . The number of channels in each stage s ∈ [0 , S ) is given by 2 s × 16 × k , with the input layer having 16 channels. For example, WRN16 -4 consists of S = 3 stages, each with [2 , 2 , 2] layers, and the corresponding number of channels in each stage is [64 , 128 , 256] . In all our experiments, we use LayerNorm instead of BatchNorm.

Transformers: We consider Transformers with GPT-2 style architecture [30]. These models use sinusoidal positional embeddings [34] and are implemented in the Standard Parameterization (SP) with GELU activation [16]. We initialize all layers using the σ 2 w / fanin scheme, except for the embedding layers, as they do not involve matrix multiplication [9]. We consider both Pre-LN [36] and Post-LN [34] Transformer variants. We denote a Transformer with d blocks and an embedding dimension of n as LMd -n . Unless specified, the model has d = 4 blocks, embedding dimension n = 128 , context length T cntxt = 64 and are trained for 10 4 steps.

D.3 Optimization Details

D.3.1 Optimizers

SGD(-M): Given gradients g t at step t , Stochastic Gradient Descent with momentum (SGD-M) updates the parameters θ t using learning rate η t and momentum m t with coefficient β . The update equations are:

Here, β = 0 corresponds to SGD. In all experiments incorporating momentum, the default value of the coefficient is set to β = 0 . 9 .

Adam: Given gradients g t at step t , Adam [20] updates the parameters θ t using learning rate η t and the first two moments of the gradient m t and v t with their coefficients β 1 and β 2 , respectively. The equations governing the updates are:

where ˆ m t = m t 1 -β t 1 and ˆ v t = v t 1 -β t 2 are the bias-corrected moments, and ϵ is a small scalar used for numerical stability. The pre-conditioner for Adam is given by:

In all experiments, the default values are set to β 1 = 0 . 9 , β 2 = 0 . 999 , and ϵ = 10 -8 , unless otherwise specified.

GI-Adam: Given gradients g t at step t , GI-Adam updates the parameters like Adam but with the second moment initialized using the gradients at initialization, i.e., v 0 = g 2 0 . The equations governing the updates are:

D.3.2 Linear Warmup

Warmup linearly increases the learning rate from an initial value η init to a target value η trgt over T wrm training steps. The learning rate η t at step t is given by:

Here, α := ( η trgt -η init ) T wrm is referred to as the rate of warmup. Under the above definition, constant learning rate training corresponds to T wrm = 1 . T wrm = 1 corresponds to constant learning rate. Unless otherwise specified, we set η init = 0 when referring to linear warmup.

D.3.3 Learning Rate Decay

In several experiments, we employ learning rate decay following the warmup phase. Specifically, we use cosine learning rate decay, which is detailed below.

Cosine Decay: Towards the end of training, it is typical to reduce the learning rate to a small value. Cosine decay is a commonly used method for decaying the learning rate from an initial value of η trgt down to a value η min over T cos steps, according to the rule:

where ρ governs the rate of decay, with ρ = 1 being the standard. Note that with ρ = 0 , the learning rate is not decayed and instead maintained at η trgt. In the above expression, t counts the steps from the initiation of cosine decay and not the current training step. As per standard practice, we consider ρ = 1 and decay the learning rate to η min = η trgt / 10 .

where ˆ v t = v t 1 -β t 2 is the bias-correction second moment and ϵ is a small scalar used for numerical stability.

Flat-Adam: Given gradients g t at step t , Flat-Adam updates the parameters θ t using learning rate η t and the first two moments of the gradients m t and v t with their coefficients β 1 and β 2 , respectively. In contrast to Adam, we do not apply bias correction for the first moment and initialize the second moment at initialization using the gradients v 0 = g 2 0 . The equations governing the updates are:

where ˆ v t = v t 1 -β t 2 is the bias-correction second moment and ϵ is a small scalar used for numerical stability. The pre-conditioner for Flat-Adam is given by:

D.3.4 Target Learning Rate Sampling for Phase Diagrams

For SGD, target learning rates η trgt are exponentially sampled using the initial sharpness λ H 0 . Starting with η trgt = 1 / λ H 0 , subsequent rates are sampled until divergence as 2 x / λ H 0 for values of x increased in integer steps starting from zero. For WRNs trained with Adam, we sample target learning rates exponentially as η trgt = 2 x × 10 -5 , where x is incremented in integer steps starting from zero until training failure. For Transformers, we sample the learning rate in a similar fashion but starting from 10 -4 and increment x in steps of 0 . 5 .

D.4 Sharpness and Pre-conditioned Sharpness Measurement

We measured sharpness / pre-conditioned sharpness using the JAX implementation of the LOBPCG sparse eigenvalue solver with the tolerance set to 10 -9 and maximum number of iterations to n iter = 1000 . In most cases, the solver converges within 40 iterations. We performed these computations in float64, as the solver would not converge with float32 in some cases.

In certain instances, the pre-conditioned sharpness computation did not converge within 1000 solver iterations. Moreover, we observed that the solver converges on restarting it with a new initial guess of the eigenvector within 40 iterations. To address these edge cases, we employed the following method: if the solver did not converge within 100 iterations, we restarted it with a new initial guess for the eigenvector. We allowed for at most 10 restarts with the maximum number of iterations set to n iter = 1000 in the last attempt. In all reported cases, the solver converges using this method.

D.5 Additional Figure Details

Figure 1: Training trajectories of 4-layer FCNs with width n = 512 , trained on a 5k subset of CIFAR-10 using MSE loss and GD in (top) µ P with η trgt = 1 / λ H 0 , where λ H 0 ≈ 0 . 05 , and (bottom) SP with η trgt = 32 / λ H 0 , where λ H 0 ≈ 50 .

Figure 2: Training loss and sharpness trajectories of 4 layer FCNs with width n = 512 , in (top) µ P with learning rate η trgt = 0 . 003 and (bottom) SP with η trgt = 0 . 001 trained the CIFAR-10 dataset with MSE loss using full batch Adam with β 1 = 0 . 9 , β 2 = 0 . 999 and ϵ = 10 -8 . In these experiments, we use data augmentation as described in Appendix D.1.1.

Figure 3: Test accuracy heatmaps of WRN16 -4 trained on CIFAR-10 using different parameterizations and loss functions using SGD with a batch size B = 128 : (a) SP and MSE loss, (b) µ P and cross-entropy loss (c) SP and cross-entropy loss. All models are trained for 10 5 steps. In these experiments, we use data augmentation as described in Appendix D.1.1.

Figure 4: Test loss heatmaps of Pre-LN Transformers in SP trained on WikiText-103 with crossentropy loss using (a) Adam, and (b) Flat-Adam (introduced in Section 6) over Adam. The Transformer models have d = 4 blocks, embedding dimension n = 768 , a context length of T cnxt = 64 , and a batch size B = 64 . These experiments also employ cosine decay and weight decay with λ = 10 -4 .

Figure 5: Heatmaps showing (a) T reach, number of steps to reach η trgt, and (b) T save, the effective number of steps saved on setting η init = η c for WRN16 -4 in SP trained on CIFAR-10 with crossentropy loss using SGD with B = 128 for 10 4 steps. For a fair comparison with linear warmup, we choose η 0 = η trgt / T wrm as our initial guess.

Figure 6: Training loss and sharpness trajectories of FCNs in SP. The experimental setup is identical to Figure 2 but with GI-Adam instead of standard Adam.

D.6 Estimation of Computational Resources

The phase diagram experiments typically required about an hour on per run on an A100 GPU. Consequently, each phase diagram consumed approximately 100 A100 hours of computational time. With a total of 16 phase diagrams, this equates to 1600 A100 hours dedicated solely to phase diagram computations. Additionally, the warmup mechanism experiments, which were conducted over 2000

10 0

10 1

10 2

10 3

step

0

.

02

0

.

03

0

.

04

0

.

05

Training loss

T

wrm

1

64

256

1024

(a)

10 0

10 1

10 2

10 3

step

0

.

0

0

.

1

0

.

2

0

.

3

0

.

4

λ H

t

T

wrm

1

64

256

1024

(b)

10 0

10 1

10 2

10 3

step

10

-

1

10 0

10 1

10 2

Training loss

T

wrm

1

64

256

1024

(c)

10 0

10 1

10 2

10 3

step

10 0

10 1

10 2

10 3

λ H

t

T

wrm

1

64

256

1024

(d)

steps, required sharpness estimation. The FCN experiments required approximately 1200 A100 hours, while the WRN mechanism experiments consumed 1600 A100 hours. The experiments concerning the initial learning rate took about 20 A100 hours. This brings the total computational time amounted to approximately 4500 A100 hours. Preliminary experiments took about 1000 A100 hours. Hence, we estimate the total computational cost to be around 5500 A100 hours.

E Additional Results for Mechanisms of Warmup

This section presents additional trajectories for warmup mechanisms discussed in Section 4 covering various architectures, loss functions, and optimizers.

E.1 Stochastic Gradient Descent

Figure 8 shows that the warmup mechanisms for full batch GD are also observed in the SGD with a batch size B = 512 . The results for other optimizers in the mini-batch setting are discussed in their respective sections.

E.2 Stochastic Gradient Descent with Momentum

While the warmup mechanisms of SGD with momentum are fundamentally similar to those of vanilla SGD, three key differences arise, as discussed below.

First, the training loss can decrease loss in an oscillatory fashion during training [11]. To illustrate this, consider the full-batch GD with momentum. The middle row of Figure 9 demonstrates the sharpness reduction case with the learning rates well below the stability threshold ( η << η c ). Despite being far below these thresholds, the loss does not decrease monotonically but converges in an oscillatory fashion. This makes it challenging to differentiate between warmup-induced catapults and fluctuations in loss due to the intrinsic effects of momentum. Nevertheless, we can still observe loss spikes correlated with an abrupt decrease in sharpness at large learning rates, as seen in the bottom row of the same figure. Similar to the SGD case, we observe these catapults are delayed and become smaller in magnitude on increasing the warmup duration.

Next, the stability threshold η c for SGD with momentum evolves during training. For simplicity of explanations, we again consider the full batch GD case. The stability threshold η c for SGD with

10 0

10 1

10 2

10 3

step

10

-

5

10

-

3

10

-

1

Training loss

T

wrm

1

64

256

1024

(a)

10 0

10 1

10 2

10 3

step

10

-

1

10 0

λ H

t

T

wrm

1

64

256

1024

(b)

10 0

10 1

10 2

10 3

step

0

.

0

0

.

2

0

.

4

Training loss

T

wrm

1

64

256

1024

(c)

10 0

10 1

10 2

10 3

step

0

20

40

60

λ H

t

T

wrm

1

64

256

1024

(d)

10 0

10 1

10 2

10 3

step

10

-

5

10

-

3

10

-

1

Training loss

T

wrm

64

256

1024

(e)

10 0

10 1

10 2

10 3

step

10 0

10 1

λ H

t

T

wrm

64

256

1024

(f)

momentum changes from 2 / λ H 0 at initialization to (2+2 β ) / λ t late in training. At initialization, the momentum vector is set to zero m 0 = 0 , and the stability is given by vanilla GD threshold 2 / λ H 0 . As training progresses, the momentum m t increases in magnitude, and the instability threshold at late training time becomes (2+2 β ) / λ t . The bottom row of Figure 9 show the sharpness trajectories with both 2 / η t and (2+2 β ) / η t curves. For T wrm = 64 , the learning rate curve 2 / η t causes an abrupt decrease in sharpness, which is coupled with a loss spike. For longer warmup durations, the sharpness decreases before training exceeds the 2 / λ H t .

Finally, the instability threshold η c for SGD with momentum significantly decreases for smaller batch sizes. Figure 10 shows the training trajectories under the same setup as in Figure 8, but with momentum coefficient β = 0 . 9 . The late-time sharpness trajectories oscillate well below the (2+2 β ) / η t (and even below 2 / λ H t ), whereas the vanilla SGD counterpart oscillates on the 2 / η t curve. This indicates a strong dependence of the instability threshold on batch size.

Besides these three differences, we note that the warmup mechanisms of SGD with momentum are similar to the vanilla SGD case. We leave a thorough analysis of the early sharpness dynamics of SGD with momentum for future works.

E.3 Stochastic Gradient Descent and Cross-entropy Loss

The warmup mechanisms for models trained with cross-entropy loss exhibit trends similar to those observed with MSE loss with one crucial difference. Near convergence, sharpness first increases and then abruptly decreases. The decrease in sharpness towards the end of training is observed in previous studies analyzing SGD with fixed learning rate [5]. Additionally, we observe higher

10 0

10 1

10 2

10 3

step

0

.

01

0

.

02

0

.

03

0

.

04

0

.

05

Training loss

T

wrm

1

64

256

1024

(a)

10 0

10 1

10 2

10 3

step

0

.

0

0

.

1

0

.

2

0

.

3

0

.

4

0

.

5

λ H

t

T

wrm

1

64

256

1024

(b)

10 0

10 1

10 2

10 3

step

10

-

2

10

-

1

10 0

Training loss

T

wrm

64

256

1024

(c)

10 0

10 1

10 2

10 3

step

10 0

10 1

λ H

t

T

wrm

64

256

1024

(d)

10 0

10 1

10 2

10 3

step

0

.

5

1

.

0

1

.

5

2

.

0

2

.

5

Training loss

T

wrm

1

64

256

1024

(a)

10 0

10 1

10 2

10 3

step

10

-

1

10 0

10 1

λ H

t

T

wrm

1

64

256

1024

(b)

10 0

10 1

10 2

10 3

step

10

-

2

10

-

1

10 0

10 1

10 2

Training loss

T

wrm

1

64

256

1024

(c)

10 0

10 1

10 2

10 3

step

10 1

10 2

10 3

10 4

λ H

t

T

wrm

1

64

256

1024

(d)

fluctuations compared to the MSE loss case. Figure 11 shows trajectories of FCNs under different parameterizations trained on CIFAR-10 with cross-entropy loss using vanilla SGD. Meanwhile, Figure 12 shows the loss and sharpness trajectories of FCNs in SP trained on CIFAR-10 with cross-entropy loss using full batch GD with and without momentum.

10 0

10 1

10 2

10 3

step

10

-

2

10 0

10 2

Training loss

T

wrm

1

64

256

1024

(a)

10 0

10 1

10 2

10 3

step

10 0

10 1

10 2

10 3

λ H

t

T

wrm

1

64

256

1024

(b)

10 0

10 1

10 2

10 3

step

10

-

4

10

-

2

10 0

10 2

Training loss

T

wrm

1

64

256

1024

(c)

10 0

10 1

10 2

10 3

step

10 0

10 2

10 4

λ H

t

T

wrm

1

64

256

1024

(d)

10 0

10 1

10 2

10 3

step

0

.

01

0

.

02

0

.

03

0

.

04

0

.

05

Training loss

T

wrm

1

64

256

1024

(a)

10 0

10 1

10 2

10 3

step

10 4

10 5

10 6

10 7

λ P

-

1

H

t

T

wrm

1

64

256

1024

(b)

10 0

10 1

10 2

10 3

step

0

.

1

0

.

2

0

.

3

0

.

4

λ H

t

T

wrm

1

64

256

1024

(c)

10 0

10 1

10 2

10 3

step

0

.

0

0

.

5

1

.

0

1

.

5

Training loss

T

wrm

1

64

256

1024

(d)

10 0

10 1

10 2

10 3

step

10 4

10 5

10 6

10 7

λ P

-

1

H

t

T

wrm

1

64

256

1024

(e)

10 0

10 1

10 2

10 3

step

20

40

60

λ H

t

T

wrm

1

64

256

1024

(f)

E.4 Warmup Mechanisms of Adam

As discussed in Section 4.3, the instability threshold for Adam is determined by the pre-conditioned sharpness λ P -1 H and not by the sharpness itself. Moreover, training dynamics falls under the sharpness reduction case as the pre-conditioned sharpness starts off large and reduces considerably during the first few training.

Figure 13 shows the training trajectories of FCNs trained with Adam in the same setting as in Figure 2 but with a batch size of B = 512 . Similar to the SGD with momentum case, the late

10 0

10 1

10 2

10 3

step

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

Training loss

T

wrm

1

64

256

1024

(a)

10 0

10 1

10 2

10 3

step

10 3

10 4

10 5

10 6

10 7

λ P

-

1

H

t

T

wrm

1

64

256

1024

(b)

10 0

10 1

10 2

10 3

step

0

10

20

30

λ H

t

T

wrm

1

64

256

1024

(c)

10 0

10 1

10 2

10 3

step

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

2

.

5

Training loss

T

wrm

1

64

256

1024

(d)

10 0

10 1

10 2

10 3

step

10 3

10 4

10 5

10 6

10 7

λ P

-

1

H

t

T

wrm

1

64

256

1024

(e)

10 0

10 1

10 2

10 3

step

0

50

100

150

λ H

t

T

wrm

1

64

256

1024

(f)

time sharpness oscillates far below the instability threshold ( (2+2 β 1 ) / η t (1 -β 1 ) ), suggesting that the instability threshold heavily decreases with a smaller batch size. We note similar findings by Ref. [6].

Next, Figure 14 show the warmup mechanism of FCNs trained with cross-entropy loss using Adam under the full-batch setting. Similar to the SGD case, the pre-conditioned sharpness decreases towards the end of training.

E.5 Different Architectures and Datasets

10 0

10 1

10 2

10 3

step

0

.

03

0

.

04

0

.

05

Training loss

T

wrm

1

64

256

1024

(a)

10 0

10 1

10 2

10 3

step

10 0

10 1

10 2

10 3

λ H

t

T

wrm

1

64

256

1024

(b)

10 0

10 1

10 2

10 3

step

10

-

1

10 0

10 1

10 2

10 3

Training loss

T

wrm

1

64

256

1024

(c)

10 0

10 1

10 2

10 3

step

10 1

10 2

10 3

10 4

10 5

λ H

t

T

wrm

1

64

256

1024

(d)

10 0

10 1

10 2

10 3

step

1

.

9

2

.

0

2

.

1

2

.

2

2

.

3

Training loss

T

wrm

1

64

256

1024

(a)

10 0

10 1

10 2

10 3

step

10 0

10 1

10 2

10 3

λ H

t

T

wrm

1

64

256

1024

(b)

10 0

10 1

10 2

10 3

step

2

3

4

5

Training loss

T

wrm

1

64

256

1024

(c)

10 0

10 1

10 2

10 3

step

10 2

10 4

10 6

λ H

t

T

wrm

1

64

256

1024

(d)

In the previous sections, we confined our analysis to FCNs to thoroughly explore the effects of different optimizers and loss functions. This section expands on those results by demonstrating that the observed warmup mechanisms apply to ResNets and Transformers as well. The Resnet experiments also employ data augmentation as detailed in Appendix D.1.

Figures 15 and 16 show the training trajectories of WideResNets (WRNs) trained on CIFAR-10 with MSE and cross-entropy loss using SGD. These trajectories generally reflect the warmup mechanisms discussed in Section 4. However, certain additional features obscure the clarity of these mechanisms. Notably, we observed a significant sharpness spike on the first training step when using longer warmup durations, which automatically resolves in the subsequent step. The magnitude of this spike increases with longer warmup periods. Further analysis revealed that this phenomenon is associated with an initial increase in the first LayerNorm parameters, which also resolves automatically by the second step. Beyond this observation, the training trajectories align with the warmup mechanisms described in the main text.

Figure 17 illustrates the warmup mechanisms of Pre-LN Transformers trained on the WikiText-2 with SGD. The Pre-LN Transformer (top row) starts in a flat landscape region ( λ H 0 ∼ 5 ) and experiences progressive sharpening right from initialization. In contrast, when the last LayerNorm (just before the final linear layer) is removed (bottom row), the model starts training in a significantly sharper region, with the initial sharpness 100 times larger than the standard Pre-LN Transformer. This modified Pre-LN Transformer experiences a reduction in sharpness during the early stages of training.

Figure 18 presents the warmup mechanisms of Pre-LN Transformers trained on WikiText-2 using the Adam optimizer. Consistent with the results in the main text, the pre-conditioned sharpness exhibits a reduction early in training, despite the model initializing in a very flat region.

These experiments demonstrate that Transformers trained on language modeling tasks exhibit warmup mechanisms consistent with those discussed in the main text.

F Additional Phase Diagrams

This section presents further results related to the phase diagrams of warmup shown in Section 5.

10 0

10 1

10 2

10 3

step

5

10

15

20

25

30

Training loss

T

wrm

1

64

256

1024

(a)

10 0

10 1

10 2

10 3

step

10 0

10 1

10 2

10 3

λ H

t

T

wrm

1

64

256

1024

(b)

10 0

10 1

10 2

10 3

step

6

8

10

12

14

16

Training loss

T

wrm

1

64

256

1024

(c)

10 0

10 1

10 2

10 3

step

10 2

10 3

λ H

t

T

wrm

1

64

256

1024

(d)

10 0

10 1

10 2

10 3

step

4

5

6

7

8

9

Training loss

T

wrm

64

256

1024

(a)

10 0

10 1

10 2

10 3

step

10 4

10 5

10 6

λ P

-

1

H

t

T

wrm

64

256

1024

(b)

10 0

10 1

10 2

10 3

step

0

5

10

15

λ H

t

T

wrm

64

256

1024

(c)

F.1 Phase Diagrams for different Models and Datasets

Figure 19 shows the test accuracy heatmaps of WRN16 -4 trained on CIFAR-100 and Tiny-ImageNet. These models are trained using cross-entropy loss using SGD with a batch size of B = 128 . Additional phase diagrams for Adam are presented in Appendix F.3.

Figure 20(a) shows the test loss heatmaps of Pre-LN Transformer trained on the WikiText-2 dataset using SGD with a batch size B = 64 . Figure 20(b) shows the Pre-LN Transformer under the same setup except for the last layer LayerNorm removed. The standard Pre-LN Transformer starts off with a small sharpness, while the version without the last LN starts off with 100 times higher curvature and requires warmup to achieve good performance.

F.2 The Effect of Momentum and Learning Rate Decay

Figure 21 shows that incorporating momentum and cosine decay (for details, see Appendix D.3.3) minimally affects the warmup phase diagrams. While the conclusions regarding warmup presented in the main text remain unaffected, we note a few interesting observations.

3.6e-03

7.1e-03

1.4e-02

2.9e-02

5.7e-02

1.1e-01

2.3e-01

4.6e-01

9.1e-01

1.8e+00

3.7e+00

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

49.71

55.61

61.47

64.86

66.55

67.85

68.30

67.41

49.84

55.69

61.39

65.34

66.63

67.75

68.10

68.21

49.92

55.43

61.28

65.55

66.71

67.71

67.82

68.33

50.02

55.65

61.27

64.99

67.23

68.53

68.22

67.78

65.42

50.00

55.41

61.92

65.58

67.45

67.94

67.98

69.15

67.70

1.13

50.17

55.63

61.54

66.12

66.82

68.25

68.15

68.36

69.33

1.13

49.79

55.95

61.11

65.13

67.41

67.76

67.96

68.42

68.30

50.05

55.75

60.77

66.15

67.85

68.32

69.10

68.19

68.40

49.95

55.34

61.75

65.92

68.27

68.29

67.81

68.17

68.51

49.73

55.37

61.34

65.03

67.77

68.03

68.33

68.14

68.23

67.95

1.13

49.66

55.12

61.41

65.43

67.35

68.38

68.75

68.76

68.67

68.10

49.35

54.89

61.19

65.04

67.24

68.13

68.26

68.86

68.38

67.93

49.50

55.00

60.91

65.10

66.97

68.18

68.47

68.53

69.25

68.22

67.35

WRN-16-4 SP XENT SGD CIFAR-100

50

.

0

52

.

5

55

.

0

57

.

5

60

.

0

62

.

5

65

.

0

67

.

5

(a)

5.2e-03

1.1e-02

2.1e-02

4.2e-02

8.4e-02

1.7e-01

3.3e-01

6.7e-01

1.3e+00

2.7e+00

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

30.81

36.34

40.98

44.72

47.96

49.48

50.32

0.56

30.61

36.53

40.75

44.90

47.08

50.04

50.95

49.37

30.74

36.55

41.24

44.74

47.79

48.96

51.09

30.10

30.65

36.48

40.94

45.07

47.54

50.11

49.33

50.08

0.56

30.73

36.52

41.55

45.30

48.19

49.17

50.05

51.70

0.56

30.72

36.89

41.63

44.93

47.70

49.65

50.24

50.95

51.34

0.56

30.80

36.52

41.45

44.80

47.35

49.87

49.84

49.90

50.50

30.79

36.50

41.25

44.82

47.98

49.85

50.59

50.28

50.90

30.34

36.78

41.40

45.64

47.90

50.19

50.73

49.96

49.98

0.56

30.57

36.34

41.55

45.19

48.15

49.88

50.29

50.79

49.83

43.48

30.60

36.30

41.44

45.74

47.51

50.11

50.43

50.17

50.60

30.46

36.21

41.57

45.41

47.92

49.77

50.69

50.45

50.20

30.28

35.95

41.34

45.42

48.46

50.06

50.87

50.51

50.73

49.34

WRN-16-4 SP XENT SGD TINY-IMAGENET

25

30

35

40

45

50

(b)

2.4e-01

3.4e-01

4.9e-01

6.9e-01

9.7e-01

1.4e+00

1.9e+00

2.7e+00

η

trgt

1

4

16

64

128

256

512

1024

2048

4096

T

wrm

4.58

4.45

4.35

4.27

4.58

4.45

4.35

4.31

4.96

4.58

4.45

4.35

4.27

4.20

4.14

4.58

4.45

4.35

4.26

4.19

4.13

4.09

4.59

4.45

4.35

4.27

4.19

4.14

4.10

4.59

4.45

4.35

4.26

4.20

4.13

4.10

4.08

4.60

4.46

4.35

4.26

4.20

4.15

4.10

4.08

4.61

4.47

4.36

4.27

4.20

4.14

4.11

4.07

4.64

4.49

4.39

4.29

4.21

4.16

4.11

4.08

4.72

4.55

4.43

4.33

4.26

4.19

LM-4-128 XENT SGD WIKITEXT-2

4

.

25

4

.

50

4

.

75

5

.

00

5

.

25

5

.

50

5

.

75

6

.

00

(a)

3.8e-03

5.4e-03

7.6e-03

1.1e-02

1.5e-02

2.2e-02

3.0e-02

4.3e-02

6.1e-02

8.6e-02

1.2e-01

1.7e-01

2.4e-01

η

trgt

1

4

16

64

128

256

512

1024

2048

4096

T

wrm

5.93

5.83

5.74

5.67

5.59

5.51

5.39

5.23

5.93

5.83

5.74

5.67

5.59

5.51

5.39

5.23

5.93

5.83

5.74

5.67

5.59

5.51

5.39

5.23

5.00

5.93

5.83

5.75

5.67

5.59

5.51

5.40

5.24

5.01

4.80

5.94

5.83

5.75

5.67

5.59

5.51

5.40

5.25

5.03

4.82

4.66

5.94

5.84

5.75

5.67

5.59

5.51

5.40

5.25

5.05

4.83

4.67

4.56

5.94

5.84

5.75

5.67

5.60

5.52

5.41

5.26

5.06

4.84

4.68

4.57

5.95

5.85

5.76

5.68

5.60

5.52

5.42

5.28

5.08

4.86

4.70

4.58

5.97

5.86

5.77

5.69

5.62

5.54

5.44

5.31

5.13

4.90

4.73

4.60

4.51

6.01

5.90

5.81

5.72

5.65

5.57

5.48

5.37

5.21

5.00

4.79

4.65

4.55

LM-4-128 XENT SGD WIKITEXT-2

4

.

6

4

.

8

5

.

0

5

.

2

5

.

4

5

.

6

5

.

8

6

.

0

(b)

First, the divergent boundary shifts leftward on incorporating momentum, indicating that momentum permits smaller target learning rates without warmup, and warmup helps SGD-M more. Meanwhile, cosine decay has a minimal effect on the divergent boundary.

Additionally, we observe a performance enhancement by incorporating momentum, especially at small learning rates. In contrast, a decaying learning rate beyond warmup degrades performance at small learning rates while improving at higher ones. Finally, incorporating both momentum and cosine decay leads to further enhancement, indicating a synergistic interaction between the two.

F.3 Phase Diagrams of Adam and GI-Adam

Figures 23 to 25 compare the warmup phase diagrams of Adam and GI-Adam of WRNs trained on CIFAR-100, Tiny-ImageNet and of Transformers trained on WikiText-2 and Wikitext-103 dataset. Similar to the results shown in the main text, GI-Adam enhances performance over standard Adam by pushing the failure boundary.

F.4 Comparison with RAdam

Figures 26 and 27 and table 1 compare the performance of GI-Adam with RAdam. We observe that GI-Adam performs on par or better than RAdam while offering a simpler modification to Adam.

1.8e-03

3.7e-03

7.3e-03

1.5e-02

2.9e-02

5.9e-02

1.2e-01

2.3e-01

4.7e-01

9.4e-01

1.9e+00

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

76.69

81.63

85.28

88.66

89.75

91.21

91.67

91.95

76.67

81.95

85.88

88.58

90.21

90.97

91.86

91.93

76.70

81.78

85.61

88.24

89.72

91.10

91.60

92.15

10.37

76.46

81.76

86.01

88.45

90.02

91.45

91.76

76.43

82.01

85.87

87.95

90.20

90.98

92.00

92.41

91.84

76.32

81.80

85.86

88.17

89.91

91.47

91.76

92.21

83.34

76.32

81.92

86.10

88.25

90.27

91.00

91.83

92.58

93.10

92.47

76.43

82.05

85.95

88.59

90.05

91.05

92.24

92.48

92.42

20.60

76.59

82.02

85.45

88.30

90.45

91.29

92.05

92.34

92.68

10.30

76.18

82.27

85.42

88.13

90.08

91.45

91.87

92.30

93.17

76.01

82.11

85.68

88.55

90.21

91.34

92.17

92.60

92.64

93.13

76.23

82.29

85.74

88.40

90.33

91.22

92.41

92.59

92.76

93.16

93.08

75.65

82.15

85.76

88.34

89.90

91.25

92.25

92.58

92.96

92.89

WRN-16-4 SP XENT SGD CIFAR-10

50

55

60

65

70

75

80

85

90

(a)

1.8e-03

3.7e-03

7.3e-03

1.5e-02

2.9e-02

5.9e-02

1.2e-01

2.3e-01

4.7e-01

9.4e-01

1.9e+00

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

89.52

91.12

92.03

92.47

92.79

92.51

10.37

89.65

91.18

92.08

92.64

92.73

92.39

10.37

89.52

91.04

92.02

93.01

92.97

92.64

10.37

89.73

90.98

91.77

92.62

92.70

92.89

89.29

90.02

91.14

92.24

92.77

93.07

92.85

89.62

90.92

92.13

92.91

93.47

93.13

10.37

89.62

91.03

92.07

92.79

93.43

93.93

93.49

10.30

89.73

90.88

92.15

92.70

93.25

93.61

93.60

93.70

10.37

90.07

91.20

92.12

92.55

93.20

93.42

93.40

93.41

26.94

10.37

89.76

91.01

92.08

92.89

93.02

93.66

93.80

93.48

92.98

21.15

10.30

89.80

91.18

92.02

92.66

93.37

93.44

93.60

93.64

93.29

90.40

10.30

89.52

91.16

92.43

92.89

93.18

93.57

93.76

93.67

93.52

91.59

89.71

91.27

92.40

92.79

93.35

93.51

93.69

93.58

93.36

92.37

89.72

WRN-16-4 SP XENT SGD-M CIFAR-10

50

55

60

65

70

75

80

85

90

(b)

1.8e-03

3.7e-03

7.3e-03

1.5e-02

2.9e-02

5.9e-02

1.2e-01

2.3e-01

4.7e-01

9.4e-01

1.9e+00

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

71.06

77.31

82.40

86.60

88.93

90.44

91.19

91.52

71.03

77.07

82.80

86.74

89.22

90.38

91.37

91.50

70.97

77.23

82.51

86.66

88.80

89.99

91.21

91.81

91.98

71.14

77.20

82.47

86.38

88.91

90.60

91.31

70.99

77.08

82.59

86.42

88.75

90.59

91.49

92.47

91.79

71.02

76.98

82.66

86.52

89.17

90.57

91.90

92.35

91.11

70.91

76.93

82.32

86.67

88.96

90.42

91.41

92.15

92.89

10.30

70.96

76.90

82.50

86.64

88.87

90.60

91.54

92.44

92.50

92.95

10.30

70.86

76.98

82.46

86.44

89.07

90.31

91.64

92.40

93.04

70.84

77.08

82.88

86.32

88.89

90.76

91.81

92.20

92.85

71.10

77.36

82.63

86.52

89.00

90.57

91.69

92.43

92.98

93.37

71.04

77.63

83.00

86.46

89.11

90.40

91.91

92.43

93.07

92.97

93.29

70.89

77.19

82.87

86.68

88.85

90.53

92.01

92.55

93.05

93.11

WRN-16-4 SP XENT SGD CIFAR-10

50

55

60

65

70

75

80

85

90

(c)

1.8e-03

3.7e-03

7.3e-03

1.5e-02

2.9e-02

5.9e-02

1.2e-01

2.3e-01

4.7e-01

9.4e-01

1.9e+00

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

88.38

90.51

91.48

92.41

92.56

92.76

10.37

87.94

90.06

91.66

92.52

92.55

92.74

10.37

88.23

89.95

91.28

92.20

92.91

93.06

10.37

88.18

89.90

91.64

92.63

92.63

92.57

88.44

10.30

88.58

90.69

91.64

92.54

92.93

92.66

88.29

90.13

91.57

92.44

93.18

93.25

10.37

88.26

90.20

91.58

92.82

92.97

93.84

93.57

10.30

88.48

90.08

91.28

92.57

93.29

93.89

93.94

93.73

10.37

88.33

90.40

91.77

92.68

93.14

93.65

93.99

93.83

88.21

90.07

91.45

92.70

93.34

93.66

93.81

93.87

93.43

21.15

10.30

88.09

90.35

91.71

92.42

93.72

93.79

93.95

93.90

93.46

91.55

10.30

88.04

90.56

91.90

92.82

93.00

93.90

93.78

94.03

93.61

92.53

87.84

90.15

92.02

92.90

93.44

93.64

94.11

93.84

94.15

93.27

90.87

WRN-16-4 SP XENT SGD-M CIFAR-10

50

55

60

65

70

75

80

85

90

(d)

8.0e-05

1.6e-04

3.2e-04

6.4e-04

1.3e-03

2.6e-03

5.1e-03

1.0e-02

2.0e-02

4.1e-02

8.2e-02

1.6e-01

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

90.03

92.41

93.37

93.71

94.07

93.93

93.38

91.76

10.37

89.85

92.32

93.34

93.87

94.05

94.17

93.19

91.78

10.37

90.06

92.37

93.59

93.82

94.18

94.05

93.49

91.83

10.37

89.92

92.24

93.52

93.95

93.88

94.16

93.94

91.18

87.73

10.37

89.77

92.22

93.62

94.06

94.10

94.37

93.50

92.59

89.36

10.37

89.79

92.29

93.35

93.99

94.21

94.51

93.72

92.43

90.85

10.37

90.12

92.06

93.66

93.88

94.19

94.41

94.46

93.18

90.79

88.59

10.24

89.82

92.31

93.72

93.86

94.40

94.48

94.23

93.42

91.64

87.81

37.11

89.84

92.37

93.59

93.93

94.26

94.47

94.20

93.53

91.88

89.13

78.53

10.30

89.95

92.39

93.72

94.07

94.62

94.48

94.29

93.67

92.52

89.82

83.35

26.59

89.95

92.53

93.51

94.10

94.32

94.73

94.22

93.57

92.64

90.74

84.43

55.67

89.79

92.27

93.66

94.22

94.43

94.50

94.90

94.09

92.85

90.91

87.66

40.88

90.01

92.24

93.74

93.94

94.41

94.70

94.37

94.20

93.30

91.92

88.92

50.30

WRN-16-4 SP XENT Adam CIFAR-10

50

55

60

65

70

75

80

85

90

(a)

8.0e-05

1.6e-04

3.2e-04

6.4e-04

1.3e-03

2.6e-03

5.1e-03

1.0e-02

2.0e-02

4.1e-02

8.2e-02

1.6e-01

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

89.99

92.50

93.57

94.37

94.36

94.49

94.08

93.28

91.83

89.82

10.24

89.96

92.46

93.67

94.07

94.22

94.53

94.30

93.64

92.30

90.38

10.24

89.90

92.41

93.64

93.96

94.26

94.56

94.57

93.46

91.96

90.24

10.24

90.24

92.60

93.56

94.11

94.38

94.39

94.00

93.36

92.10

89.96

10.24

90.03

92.35

93.41

93.95

94.57

94.55

94.13

93.73

92.24

89.99

86.99

10.30

89.89

92.53

93.63

94.08

94.26

94.33

94.20

94.09

92.48

90.11

87.21

31.29

89.75

92.25

93.49

94.01

94.33

94.43

94.16

93.81

92.29

90.64

87.39

51.49

90.01

92.58

93.60

94.29

94.36

94.36

94.17

93.64

92.56

90.24

87.73

36.99

89.92

92.46

93.62

93.97

94.26

94.64

94.24

93.62

92.79

90.75

87.35

56.95

89.87

92.39

93.35

94.31

94.53

94.46

94.55

94.01

92.77

90.95

88.09

42.54

90.16

92.43

93.43

94.32

94.32

94.56

94.37

93.90

93.03

91.46

88.39

47.30

90.18

92.49

93.66

94.09

94.95

94.49

94.47

94.01

93.14

91.36

88.73

51.38

89.91

92.57

93.57

94.41

94.57

94.63

94.39

94.26

93.32

91.90

89.10

70.77

WRN-16-4 SP XENT GI-Adam CIFAR-10

50

55

60

65

70

75

80

85

90

(b)

G Non-divergence of Adam

Figure 28 shows that, despite experiencing catastrophic instabilities during early training, Adam does not diverge well beyond the training failure boundary. While Adam can recover from these

8.0e-05

1.6e-04

3.2e-04

6.4e-04

1.3e-03

2.6e-03

5.1e-03

1.0e-02

2.0e-02

4.1e-02

8.2e-02

1.6e-01

3.3e-01

6.6e-01

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

62.60 68.45 71.47 70.91 71.92 70.03

1.13

62.91 68.51 70.78 71.13 71.49 70.66

1.13

62.73 68.70 70.67 71.12 71.21 70.91 67.03

1.13

62.90 68.96 70.99 71.74 71.30 70.87 67.33

1.13

62.30 68.63 70.36 71.29 71.65 71.00 68.84 65.76

1.13

62.88 68.77 71.05 70.95 71.44 71.23 68.82 60.81 59.40

1.13

62.74 68.67 70.86 71.06 71.97 70.51 69.81 64.92 45.54

1.13

62.36 68.17 71.50 71.30 72.53 71.13 69.84 64.08 52.06

1.13

62.06 68.98 71.18 71.79 71.74 71.90 70.00 68.07 49.37 40.52

1.13

62.38 68.53 71.28 71.73 71.78 71.37 70.82 67.68 63.84 37.85

2.68

62.43 68.66 70.95 72.23 72.02 72.37 71.13 68.01 64.99 56.37 23.57

1.13

1.13

1.13

62.42 68.43 71.22 72.19 72.32 72.40 71.85 69.51 67.26 63.21 39.87

6.90

2.62

1.13

62.29 68.89 71.43 72.33 73.29 72.71 72.98 70.93 68.55 65.15 58.98 11.69

5.02

1.13

WRN-16-4 SP XENT Adam CIFAR-100

50

55

60

65

70

(a)

8.0e-05

1.6e-04

3.2e-04

6.4e-04

1.3e-03

2.6e-03

5.1e-03

1.0e-02

2.0e-02

4.1e-02

8.2e-02

1.6e-01

3.3e-01

6.6e-01

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

61.79 68.74 70.62 72.26 72.03 71.97 71.33 67.20

1.13

61.81 68.40 71.49 71.83 72.60 72.01 69.85 67.06

1.13

61.71 68.75 70.90 71.88 72.54 72.30 70.48 68.52 59.90

1.13

62.10 68.58 71.19 71.98 72.17 71.81 70.01 68.30 64.72

1.13

62.33 68.65 71.14 71.93 72.46 72.70 70.67 67.78 62.60 30.77

1.13

62.16 68.76 71.08 71.92 72.27 72.51 71.27 67.88 63.22 44.18

1.13

62.27 68.69 70.97 71.84 71.87 71.88 71.39 69.52 64.53 41.13

1.95

62.04 68.80 71.17 72.07 72.06 71.83 71.38 69.02 66.72 59.14

3.90

62.30 68.34 70.95 71.97 71.99 72.43 71.31 69.39 67.00 61.58 27.38

1.13

62.51 68.50 71.05 71.61 72.67 72.22 71.60 69.08 67.04 62.20 45.88

3.31

1.13

62.36 68.63 70.89 72.04 73.05 72.17 71.68 70.41 67.02 63.72 54.91

4.49

2.28

1.13

62.48 68.87 71.69 72.10 71.92 73.12 71.89 70.50 67.66 64.52 57.90

8.38

5.48

3.23

62.34 68.68 71.88 72.85 72.72 72.77 72.42 70.89 69.06 66.06 60.41 14.67

9.52

4.22

WRN-16-4 SP XENT GI-Adam CIFAR-100

50

55

60

65

70

(b)

8.0e-05

1.6e-04

3.2e-04

6.4e-04

1.3e-03

2.6e-03

5.1e-03

1.0e-02

2.0e-02

4.1e-02

8.2e-02

1.6e-01

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

40.10

46.79

51.35

52.98

53.74

0.63

40.19

46.81

50.99

52.91

54.08

52.38

0.56

40.11

46.59

51.07

53.08

54.30

52.89

0.56

40.38

46.95

50.51

53.13

53.02

52.86

0.56

40.15

46.91

51.38

52.90

54.45

53.66

0.56

39.53

47.03

50.79

53.10

54.17

53.40

0.56

39.70

46.45

50.65

52.57

54.00

54.60

52.84

46.79

0.56

39.87

46.60

51.57

53.68

53.96

53.60

52.05

47.74

28.67

0.56

39.81

46.88

51.16

53.21

54.35

54.53

53.29

49.80

24.68

10.30

0.56

39.90

46.46

50.62

53.21

54.75

53.92

52.82

49.77

39.35

21.89

0.56

39.54

46.86

51.33

53.41

54.25

55.18

53.58

50.84

44.55

23.65

1.39

39.67

46.54

51.22

53.85

54.74

54.22

54.32

51.82

48.17

37.95

2.00

39.53

46.64

51.83

53.99

54.91

55.10

53.92

53.32

50.12

44.12

20.91

2.54

WRN-16-4 SP XENT Adam TINY-IMAGENET

25

30

35

40

45

50

55

(a)

8.0e-05

1.6e-04

3.2e-04

6.4e-04

1.3e-03

2.6e-03

5.1e-03

1.0e-02

2.0e-02

4.1e-02

8.2e-02

1.6e-01

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

39.77

46.14

51.06

53.33

54.62

55.33

52.83

0.56

39.70

46.10

51.36

54.20

54.43

53.99

53.49

0.56

39.75

46.35

51.07

53.25

54.31

54.60

52.99

0.56

39.52

46.32

51.46

53.33

54.31

54.89

54.14

49.81

0.56

39.69

46.18

51.25

53.70

54.73

54.40

52.58

47.57

0.56

39.73

46.27

51.17

53.39

55.15

54.85

53.42

49.82

0.56

39.73

46.40

51.01

54.11

54.96

54.85

53.36

49.50

34.31

0.56

39.35

46.21

51.16

53.62

54.54

55.12

54.23

47.71

45.24

0.56

39.45

46.12

51.37

53.87

54.35

54.44

53.34

51.67

43.67

30.97

2.06

39.67

46.62

51.64

53.22

55.01

54.66

53.88

51.30

47.09

29.23

1.07

39.60

46.48

51.17

53.82

54.55

54.72

54.20

51.85

47.06

36.58

4.70

39.72

46.49

51.16

53.73

55.24

54.50

54.11

52.33

49.25

42.30

14.32

2.00

39.80

46.65

51.92

53.87

55.49

55.73

54.70

53.66

49.67

44.49

34.83

2.88

WRN-16-4 SP XENT GI-Adam TINY-IMAGENET

25

30

35

40

45

50

55

(b)

1.0e-04

2.0e-04

4.0e-04

8.0e-04

1.6e-03

3.2e-03

6.4e-03

1.3e-02

2.6e-02

5.1e-02

1.0e-01

2.0e-01

4.1e-01

8.2e-01

η

trgt

1

4

16

64

128

256

512

1024

2048

4096

T

wrm

4.47

4.25

4.08

3.99

3.93

3.88

3.84

4.48

5.44

5.72

5.83

5.94

5.89

6.25

4.47

4.25

4.08

3.99

3.93

3.88

3.85

4.46

5.36

5.74

5.77

5.75

5.84

6.15

4.47

4.25

4.08

3.99

3.93

3.88

3.85

3.85

4.61

5.57

5.76

5.86

5.94

6.09

4.48

4.26

4.08

3.99

3.93

3.88

3.85

3.85

4.36

5.56

5.81

5.81

5.90

6.19

4.48

4.26

4.09

3.99

3.93

3.87

3.85

3.84

4.22

5.06

5.78

5.71

5.94

6.14

4.48

4.26

4.09

3.99

3.93

3.87

3.85

3.84

4.12

4.47

5.68

5.87

6.02

6.17

4.49

4.27

4.10

3.99

3.93

3.88

3.84

3.83

3.84

4.46

4.59

5.77

6.06

6.06

4.50

4.28

4.10

3.99

3.93

3.88

3.84

3.83

3.83

4.38

4.50

5.70

5.91

6.08

4.52

4.29

4.11

4.00

3.94

3.89

3.84

3.83

3.82

4.07

4.46

4.50

5.79

5.98

4.57

4.33

4.15

4.02

3.96

3.90

3.86

3.83

3.83

3.88

4.14

4.46

4.50

5.79

LM-4-128 XENT Adam WikiText-2

4

.

00

4

.

25

4

.

50

4

.

75

5

.

00

5

.

25

5

.

50

5

.

75

6

.

00

(a)

1.0e-04

2.0e-04

4.0e-04

8.0e-04

1.6e-03

3.2e-03

6.4e-03

1.3e-02

2.6e-02

5.1e-02

1.0e-01

2.0e-01

4.1e-01

8.2e-01

η

trgt

1

4

16

64

128

256

512

1024

2048

4096

T

wrm

4.52

4.29

4.12

4.01

3.94

3.88

3.86

3.84

3.85

4.25

4.69

5.71

5.96

6.00

4.52

4.29

4.12

4.01

3.94

3.88

3.85

3.84

3.86

4.26

4.63

5.74

5.88

6.09

4.52

4.29

4.12

4.01

3.93

3.88

3.85

3.85

3.86

4.26

5.64

5.71

5.86

6.12

4.52

4.30

4.12

4.01

3.94

3.88

3.85

3.84

3.85

4.21

4.52

5.70

5.75

6.05

4.52

4.30

4.12

4.01

3.94

3.88

3.85

3.84

3.85

4.16

4.48

5.66

5.74

6.20

4.52

4.30

4.13

4.01

3.94

3.88

3.84

3.83

3.85

4.14

4.47

4.57

5.79

6.00

4.52

4.30

4.13

4.00

3.94

3.88

3.84

3.83

3.85

4.13

4.48

5.69

5.80

6.02

4.53

4.30

4.13

4.01

3.94

3.88

3.85

3.83

3.84

4.03

4.48

4.53

5.86

6.30

4.55

4.31

4.14

4.02

3.95

3.89

3.85

3.83

3.84

3.98

4.27

4.48

4.53

6.11

4.59

4.35

4.16

4.03

3.96

3.90

3.86

3.83

3.84

3.89

4.12

4.27

4.48

4.53

LM-4-128 XENT GI-Adam WikiText-2

4

.

00

4

.

25

4

.

50

4

.

75

5

.

00

5

.

25

5

.

50

5

.

75

6

.

00

(b)

instabilities, the model's performance is severely impacted, resulting in training failures rather than convergence to a reasonable minimum.

1.0e-04

3.0e-04

6.0e-04

8.0e-04

1.0e-03

3.0e-03

6.0e-03

1.0e-02

η

trgt

1

4

16

64

128

256

512

1024

2048

T

wrm

4.38

4.00

3.90

3.91

4.11

5.98

6.06

6.13

4.38

4.00

3.90

3.89

3.90

5.56

6.06

6.20

4.38

4.01

3.89

3.88

3.87

5.32

6.07

6.07

4.38

4.00

3.89

3.87

3.86

5.30

5.43

6.08

4.38

4.00

3.89

3.87

3.86

5.21

5.44

5.56

4.38

4.00

3.89

3.87

3.86

5.27

5.44

5.48

4.39

4.00

3.90

3.86

3.86

5.27

5.44

5.49

4.39

4.01

3.89

3.86

3.86

3.87

5.40

5.55

4.39

4.01

3.90

3.87

3.86

3.82

5.33

5.44

LM-4-768, wikitext-103, Crossentropy, Adamw

3

.

9

4

.

0

4

.

1

4

.

2

4

.

3

4

.

4

4

.

5

(a)

1.0e-04

3.0e-04

6.0e-04

8.0e-04

1.0e-03

3.0e-03

6.0e-03

1.0e-02

η

trgt

1

4

16

64

128

256

512

1024

2048

T

wrm

4.42

4.02

3.90

3.88

3.86

3.82

5.42

5.97

4.42

4.02

3.90

3.88

3.86

3.84

5.49

6.03

4.42

4.02

3.90

3.87

3.86

3.81

5.44

5.56

4.42

4.02

3.90

3.87

3.86

3.82

3.81

5.90

4.42

4.02

3.90

3.87

3.85

3.81

3.82

3.85

4.42

4.02

3.90

3.88

3.86

3.81

3.81

5.50

4.42

4.02

3.90

3.88

3.86

3.81

3.81

5.44

4.42

4.02

3.90

3.88

3.86

5.09

3.80

3.82

4.41

4.03

3.91

3.88

3.86

3.82

5.33

3.80

LM-4-768, wikitext-103, Crossentropy, GI-Adamw

3

.

9

4

.

0

4

.

1

4

.

2

4

.

3

4

.

4

4

.

5

(b)

1.0e-04

3.0e-04

6.0e-04

8.0e-04

1.0e-03

3.0e-03

6.0e-03

1.0e-02

η

trgt

1

4

16

64

128

256

512

1024

2048

T

wrm

4.42

4.02

3.90

3.87

3.86

3.84

5.37

5.53

4.42

4.02

3.90

3.87

3.86

3.82

5.37

5.57

4.42

4.02

3.90

3.88

3.86

3.82

5.34

5.57

4.42

4.02

3.90

3.87

3.86

3.82

3.84

5.67

4.42

4.02

3.90

3.88

3.85

3.81

3.82

3.85

4.42

4.02

3.91

3.88

3.86

3.81

3.82

3.86

4.42

4.02

3.90

3.87

3.86

3.81

3.81

3.84

4.42

4.02

3.90

3.88

3.86

5.17

3.80

3.81

4.41

4.03

3.91

3.88

3.86

3.82

5.38

3.80

LM-4-768, wikitext-103, Crossentropy, FlatAdamw

3

.

8

3

.

9

4

.

0

4

.

1

4

.

2

4

.

3

4

.

4

4

.

5

(c)

1.0e-04

3.0e-04

6.0e-04

8.0e-04

1.0e-03

3.0e-03

6.0e-03

1.0e-02

η

trgt

1

4

16

64

128

256

512

1024

2048

T

wrm

4.42

4.03

3.90

3.87

3.86

3.85

5.48

5.59

4.42

4.03

3.90

3.87

3.86

3.85

5.42

5.42

4.42

4.02

3.90

3.87

3.86

3.85

5.42

5.44

4.42

4.03

3.90

3.87

3.85

3.84

5.41

5.46

4.42

4.03

3.90

3.87

3.86

3.84

5.45

5.43

4.42

4.03

3.90

3.87

3.86

3.84

5.35

5.44

4.42

4.02

3.90

3.88

3.86

3.83

5.38

5.40

4.42

4.02

3.90

3.88

3.86

3.83

5.43

5.49

4.41

4.02

3.90

3.88

3.86

3.82

5.11

5.46

LM-4-768, wikitext-103, Crossentropy, RAdam

3

.

9

4

.

0

4

.

1

4

.

2

4

.

3

4

.

4

4

.

5

(d)

10

-

4

10

-

3

10

-

2

10

-

1

η

trgt

80

85

90

95

Test accuracy

Adam-1

GI-Adam-1

RAdam-1

Adam-4096

Adam-save-4096

GI-Adam-4096

RAdam-4096

These large loss catapults cause the gradients g to spike during early training, leading to a substantial increase in its second moment v . While the gradients return to a lower value after a few training steps, the second moment remains large in magnitude for a prolonged period. These large values of v result in a small effective learning rate, which hinders training to escape these high-loss regions. Consequently, the models remain stuck in a suboptimal state rather than converging. We refer to this as a training failure.

Upon closer examination of the individual layers during training failures, we found that certain layers or residual blocks output zero. This results in vanishing gradients except for the last layer bias and training halts. We defer the detailed analysis of Adam's failures to future work.

10 0

10 1

10 2

10 3

10 4

step

10 3

10 7

10 11

10 15

Training loss

η

trgt

30

60

90

120

150

(a)

10 0

10 1

10 2

10 3

10 4

step

0

20

40

60

80

Training accuracy

η

trgt

30

60

90

120

150

(b)

10 0

10 1

10 2

10 3

10 4

step

10 1

10 5

10 9

10 13

‖

gt

‖

η

trgt

30

60

90

120

150

(c)

10 0

10 1

10 2

10 3

10 4

step

10 1

10 6

10 11

10 16

‖

vt

‖

η

trgt

30

60

90

120

150

(d)

3.9e-01

7.8e-01

1.6e+00

3.1e+00

6.2e+00

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

T

reach

0

.

90

0

.

95

1

.

00

1

.

05

1

.

10

(a)

3.9e-01

7.8e-01

1.6e+00

3.1e+00

6.2e+00

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

0.0

0.0

0.0

0.0

0.0

-0.5

-0.5

-0.5

-0.5

-0.5

1.0

1.0

1.0

1.0

1.0

4.5

4.5

4.5

4.5

4.5

12.0

12.0

12.0

12.0

12.0

27.5

27.5

27.5

27.5

27.5

59.0

59.0

59.0

59.0

59.0

122.5

122.5

122.5

122.5

122.5

250.0

250.0

250.0

250.0

250.0

505.5

505.5

505.5

505.5

505.5

1017.0

1017.0

1017.0

1017.0

1017.0

2040.5

2040.5

2040.5

2040.5

2040.5

4088.0

4088.0

4088.0

4088.0

4088.0

T

save

0

1000

2000

3000

4000

(b)

H Additional results for the Initial Learning Rate Selection

This section provides additional results for the initial learning rate selection. Figure 29 shows the number of steps T reach required to reach the target learning rate and the effective number of steps saved T save for WRNs in µ P. Note that in such cases, η c ≈ η max. We observe that T reach = 1 for a wide range of learning rates, saving almost the entire warmup duration. Figures 30 and 31 show similar phase diagrams for WRNs in SP and µ P trained with Adam.

1.0e-04

2.0e-04

4.0e-04

8.0e-04

1.6e-03

3.2e-03

6.4e-03

1.3e-02

2.6e-02

5.1e-02

1.0e-01

2.0e-01

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

3

3

3

3

3

3

3

1

1

1

5

7

7

7

7

7

7

1

1

2

9

13

15

15

15

15

15

1

1

4

18

25

29

31

31

31

31

31

1

1

8

36

50

58

61

63

63

63

63

1

1

16

72

100

115

121

125

127

127

127

127

1

1

32

144

200

229

242

249

253

255

255

255

1

1

64

288

400

457

484

498

505

509

511

511

1

1

128

576

800

913

968

996

1010

1017

1021

1023

1

1

256

1152

1600

1825

1936

1992

2020

2034

2041

1

1

512

2304

3200

3649

3872

3984

4040

4068

4082

4089

T

reach

500

1000

1500

2000

2500

3000

3500

4000

(a)

1.0e-04

2.0e-04

4.0e-04

8.0e-04

1.6e-03

3.2e-03

6.4e-03

1.3e-02

2.6e-02

5.1e-02

1.0e-01

2.0e-01

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

0.0

0.0

0.0

0.0

0.0

0.0

0.0

0.0

0.0

-1.0

-0.5

0.0

0.5

0.5

0.5

0.5

0.5

0.5

0.5

1.0

1.5

0.0

0.5

0.5

0.5

0.5

0.5

0.5

4.0

4.5

4.0

1.5

0.0

0.5

0.5

0.5

0.5

0.5

11.5

12.0

10.5

4.0

1.5

0.0

0.5

0.5

0.5

0.5

27.0

27.5

23.0

10.5

4.0

1.5

0.0

0.5

0.5

0.5

0.5

58.5

59.0

48.5

23.0

10.5

4.0

1.5

0.0

0.5

0.5

0.5

122.0

122.5

100.0

48.5

23.0

10.5

4.0

1.5

0.0

0.5

0.5

0.5

249.5

250.0

204.5

100.0

48.5

23.0

10.5

4.0

1.5

0.0

0.5

0.5

505.0

505.5

413.0

204.5

100.0

48.5

23.0

10.5

4.0

1.5

0.0

0.5

1016.5

1017.0

830.5

413.0

204.5

100.0

48.5

23.0

10.5

4.0

1.5

0.0

2040.0

2040.5

1666.0

830.5

413.0

204.5

100.0

48.5

23.0

10.5

4.0

1.5

4087.5

4088.0

3337.5

1666.0

830.5

413.0

204.5

100.0

48.5

23.0

10.5

4.0

T

save

0

500

1000

1500

2000

2500

3000

3500

4000

(b)

1.0e-04

2.0e-04

4.0e-04

8.0e-04

1.6e-03

3.2e-03

6.4e-03

1.3e-02

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

1

T

reach

0

.

900

0

.

925

0

.

950

0

.

975

1

.

000

1

.

025

1

.

050

1

.

075

1

.

100

(a)

1.0e-04

2.0e-04

4.0e-04

8.0e-04

1.6e-03

3.2e-03

6.4e-03

1.3e-02

η

trgt

1

2

4

8

16

32

64

128

256

512

1024

2048

4096

T

wrm

0.0

0.0

0.0

0.0

0.0

0.0

0.0

-1.0

-1.0

-1.0

-1.0

-1.0

-1.0

-0.5

0.0

0.5

0.5

0.5

0.5

0.5

0.5

1.0

4.0

4.0

4.0

4.0

4.0

4.0

4.5

11.5

11.5

11.5

11.5

11.5

11.5

12.0

27.0

27.0

27.0

27.0

27.0

27.0

27.5

58.5

58.5

58.5

58.5

58.5

58.5

59.0

122.0

122.0

122.0

122.0

122.0

122.0

122.5

100.0

249.5

249.5

249.5

249.5

249.5

249.5

250.0

505.0

505.0

505.0

505.0

505.0

505.0

505.5

1016.5

1016.5

1016.5

1016.5

1016.5

1016.5

1017.0

2040.0

2040.0

2040.0

2040.0

2040.0

2040.0

2040.5

4087.5

4087.5

4087.5

4087.5

4087.5

4087.5

4088.0

T

save

0

500

1000

1500

2000

2500

3000

3500

4000

(b)

I Additional Results on GI-Adam

This section presents additional results for GI-Adam. We provide further insights into the mechanisms and interpretations of GI-Adam.

I.1 Warmup Mechanisms of GI-Adam

Figure 32 shows the training trajectories of FCNs with different parameterizations trained with GI-Adam. Notably, the pre-conditioned sharpness starts at significantly lower values than standard Adam. Specifically, for the µ P model, the initial pre-conditioned sharpness λ P -1 H is around 2000

(a)

10 0

10 1

10 2

10 3

step

0

.

015

0

.

020

0

.

025

0

.

030

0

.

035

Training loss

T

wrm

1

64

256

1024

(b)

10 0

10 1

10 2

10 3

step

10 3

10 4

10 5

10 6

10 7

λ P

-

1

H

t

T

wrm

1

64

256

1024

(c)

10 0

10 1

10 2

10 3

step

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

0

.

30

λ H

t

(a)

10 0

10 1

10 2

10 3

step

0

.

1

0

.

2

0

.

3

Training loss

T

wrm

1

64

256

1024

(b)

10 0

10 1

10 2

10 3

step

10 4

10 5

10 6

10 7

λ P

-

1

H

t

T

wrm

1

64

256

1024

(c)

10 0

10 1

10 2

10 3

step

10

15

20

25

30

λ H

t

instead of the value 10 5 observed for Adam (c.f. Figure 2). Remarkably, this almost eliminates initial sharpness reduction. Similarly, the pre-conditioned sharpness for the SP model starts around 10 4 instead of 10 6 . Notably, in the SP scenario, there is no initial spike in the T wrm = 1 (c.f. Figure 2), demonstrating that this simple modification effectively reduces instabilities during the early training.

I.2 GI-Adam as an Automated Warmup

In this section, we show that a bias correction is not required when the second moment is initialized with the gradients at initialization in GI-Adam. Therefore, employing a bias correction as in the original Adam algorithm in this case serves as an automated warmup given by η t = η trgt √ 1 β t .

The moving average of the second moment is given by:

where v 0 = g 2 0 . Following standard assumptions, we assume that the second moment of the gradient is constant during early training E [ g 2 t ] = σ 2 . Taking the expectation of the above equation over the gradient distribution yields

Simplifying the above equation, we have

This result demonstrates that when the second moment is initialized with the gradients at initialization, it does not require bias correction, as the expected value of the second moment is equal to the constant σ 2 . If we apply the usual bias correction on top of initializing the second moment with the gradients, we effectively downscale the second moment by a factor √ 1 -β t 2 . Assuming small enough ϵ , this can be viewed as a multiplicative factor to the learning rate. As a result, GI-Adam is equivalent to having a natural warmup given by η t = η trgt √ 1 -β t 2 .

-2

0

25k

50k

75k

100k

step

0

.

2

0

.

4

0

.

6

0

.

8

Test accuracy

Adam

Adam+warmup

GI-Adam

(a)

0

25k

50k

75k

100k

step

0

.

2

0

.

4

0

.

6

0

.

8

Test accuracy

Adam

Adam+warmup

GI-Adam

(b)

I.3 The Primary benefit of GI-Adam results from the magnitude of the second moment at initialization

To further assess if the primary cause of instability during early training is the large λ P -1 H , we randomly initialize v 0 but with the same norm as the gradients at initialization. We refer to this as Randomly Initialized Adam (RI-Adam). Like GI-Adam, this also results in improved performance as shown in Figure 33.

I.4 Warmup Mechanisms of Flat-Adam

Figure 34 shows the training trajectories of FCNs with different parameterizations trained with FlatAdam. We observe that further removing bias correction for the first moment completely removes the initial reduction in the pre-conditioned sharpness.

(a)

10 0

10 1

10 2

10 3

step

0

.

015

0

.

020

0

.

025

0

.

030

0

.

035

Training loss

T

wrm

1

64

256

1024

(b)

10 0

10 1

10 2

10 3

step

10 3

10 4

10 5

10 6

10 7

λ P

-

1

H

t

T

wrm

1

64

256

1024

(c)

10 0

10 1

10 2

10 3

step

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

0

.

30

λ H

t

(d)

10 0

10 1

10 2

10 3

step

0

.

1

0

.

2

0

.

3

Training loss

T

wrm

1

64

256

1024

(e)

10 0

10 1

10 2

10 3

step

10 3

10 4

10 5

10 6

10 7

λ P

-

1

H

t

T

wrm

1

64

256

1024

(f)

10 0

10 1

10 2

10 3

step

10

15

20

25

30

λ H

t

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We have summarized our main contributions in Section 1.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss the limitations in Section 7.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification: No theoretical results.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide extensive details of the experiments in Appendix D to reproduce the results. Furthermore, we provide additional details for each Figure in Appendix D.5.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The main results of the paper can be reproduced using the anonymous GitHub repository https://github.com/dayal-kalra/why-warmup .

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide extensive details of the experiments in Appendix D

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: The only sources of errors in our experiments are the variation arising from the model initialization and mini-batch sequence through training. Due to the extensive computational cost of generating the warmup phase diagrams ( 100 A100 hours), we do not provide error bars. Nevertheless, we have finely swept the learning rate and the warmup duration, and we do not observe a significant variation across the learning rate and warmup durations for various experiments. Hence, we do not expect the results to vary significantly with initialization and mini-batch sequence.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide the details in Appendix D.6.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We adhere to the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: This paper analyzes the underlying mechanisms of warmup. We do not think there are any societal impacts of the work that are worth mentioning here.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: We do not release data or models.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We have used standard image datasets and architectures and cited the original work.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: We do not release any assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: Human subjects were not involved in this study.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: Human subjects were not involved in this study.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: SUMMARY:
The paper analyses different aspects of warmup in the gradient based training focusing on SGD, ADAM and their variants under two types of parameterization. It shows through mostly empirical analysis that warmup facilitates training at higher learning rates and stabilizes the training dynamics by keeping it away from (what they call) divergence boundary (that results in failure). Furthermore, it suggests several improvements for hyperparameter initialization that shorten the training process and improves generalization.

STRENGTHS:
Choosing the learning rate  is critical for training large models. The paper proposes a nice analysis of the warmup procedure that agrees with some of the previous observations. And it suggests useful tips for practitioners.

WEAKNESSES:
The figures are not quite clear, especially because the captions do not describe the figures well enough. 
The observations are made mostly from empirical study.

QUESTIONS:
The study is done on convnets and resnets applied to images. Training LLMs is a more complex task. Would these ideas apply to attention based architectures? Or more widely, do the mechanisms described in the paper depend on the architecture?

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the mechanisms of the warmup technique. The authors experimentally demonstrate that the primary benefit of warmup is its ability to enable the network to handle larger learning rates.\n\nSTRENGTHS:\nWarmup is an essential trick for training modern deep neural networks, and understanding its role is a critical and open issue. \nThis paper makes a significant contribution to this exploration. \nFrom the perspective of training stability, the author highlights that the use of warmup enables network training to utilize larger learning rate.\nAdditionally, the author notes that different initializations correspond to different stability regimes at the beginning of training: sharpness reduction and progressive sharpening. \nWarmup is particularly important for maintaining training stability in the sharpness reduction regime.\n\nWEAKNESSES:\n- The experiments are primarily conducted on ResNet models on CIFAR, where warmup is not essential. \nHowever, for Transformer-based models, warmup appears to be indispensable, particularly in applications such as language model pretraining and Vision Transformer (ViT) training.\n\n- The observations in this paper are all based on experimental results. The findings would be more convincing if they could be theoretically validated in some settings (even in toy settings).\n\nQUESTIONS:\n- It is natural and insightful that extending warmup's time can tolerate the use of larger LR, but can the authors further clearly and quantitatively indicate this relationship?\n\n- Does the measure of sharpness, $\\lambda_{\\max}(H)$, apply to randomized algorithms? For example, for GD, the stability condition is typically $\\lambda_{\\max}(H)\\leq 2/\\eta$ ; however, for SGD, $\\lambda_{\\max}(H)$ may no longer be suitable sharpness measure about training stability, and it might be $||H||_{\\rm F}$ [1]. As the authors show in Figure 1(d) for GD, the results fully support the author's claim; however, in Figure 11(d) for SGD, the results do not completely align with the author's claim.\n\n- A closely related work [2] should be discussed. In Section 5 and Fig 3 in [2], the authors also discussed how initialization, warmup, and SGD noise influence progressive sharpening or sharpness reduction.\n\n\n[1] Wu et al. The alignment property of SGD noise and how it helps select flat minima: A stability analysis. (NeurIPS 2022)\n\n[2] Ziyin et al. Loss Symmetry and Noise Equilibrium of Stochastic Gradient Descent. (2024)\n\nLIMITATIONS:\nPlease refer to ``Weaknesses''.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

comments: SUMMARY:
The authors explain the mechanisms of the warmup technique showing that with warm up the loss of NN will go to a flatter space than direct optimization. Further, based on analysis, the authors propose a new optimization algorithm called GI-Adam.

STRENGTHS:
1. The authors explain why the warm-up technique can help networks converge better.

2. With the analysis, the authors show that the initialization of Adam is not "correct". Thus, the authors proposed a new initialization of Adam called GI-Adam.

WEAKNESSES:
1. The conclusions are from FCNs and WideResnet, which can be trained well without warm-up. Does the conclusion still hold for some "hard" models and datasets (e.g., Transformer)?

2. GI-Adam is used in [1].

[1] Zhang, Yushun, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhi-Quan Luo. "Adam can converge without any modification on update rules." Advances in neural information processing systems 35 (2022): 28386-28399.

QUESTIONS:
See Weaknesses.

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper examines the learning rate warmup technique from the perspective of its influence on the evolution of loss sharpness for different optimizers (GD, SGD(-M), Adam) and network parametrizations (Maximal Update Parameterization – $\\mu$P and Standard Parametrization – SP). The authors demonstrate that the warmup allows the network to tolerate larger learning rates by gradually reducing sharpness through a series of loss catapults and self-stabilizations. Experiments with different network parameterizations reveal how warmup influences training in progressive sharpening and sharpness reduction regimes and the lower importance of warmup for training $\\mu$P parameterized networks. Based on the empirical analysis, the paper proposes two practical training heuristics: (1) initializing the learning rate at an estimated critical value to eliminate unnecessary warmup steps and (2) introducing the GI-Adam optimizer, which initializes Adam's second moment with a squared gradient.\n\nSTRENGTHS:\n1. The paper provides solid experiments demonstrating that the gradual self-stabilization mechanism induced by the warmup is observed for different network parameterizations and optimizers.\n2. The paper discusses the specifics of the warmup effect on networks in different parameterizations and confirms that networks in $\\mu$P parametrization benefit less from it.\n3. The paper explains why warmup may be unstable for Adam and proposes a simple heuristic on how to deal with this instability. This heuristic potentially may be useful in practice.\n4. The paper includes an extensive study of warmup hyperparameters (warmup length and maximal learning rate) and suggests how to choose their optimal values. I specifically like the Persistent Catapult Warmup idea from the appendix and think it is promising.\n5. The paper is clearly written and easy to follow.\n\nWEAKNESSES:\nMy main concerns are related to the level of novelty of the empirical analysis of the warmup and the effectiveness of the proposed practical modifications: \n1. The novelty and significance of the empirical analysis in the first part of the paper seem limited. The study heavily relies on two previous works. Gilmer et al., 2021 (https://arxiv.org/pdf/2110.04369) investigate training instabilities from the perspective of sharpness, including a very similar analysis on how warmup decreases the sharpness and a discussion on how starting training from a flat initialization makes warmup much less important. Karla et al., 2023 (https://arxiv.org/pdf/2311.02076) demonstrate that $\\mu$P and SP exhibit different natural evolutions of sharpness. The paper combines these two perspectives, however, I am not sure if this combination leads to new insights. The warmup works similarly in both progressive sharpening and sharpness reduction regimes. The paper points out the instability in Adam warmup and the lower importance of warmup for $\\mu$P parametrized models (similar to the discussion in Gilmer et al., 2021), but both of these insights are not related to the different sharpness evolution regimes. \n2. The authors claim that starting warmup from the critical learning rate $\\eta_c$ is an effective strategy. However, this claim is not obvious, and adding empirical evaluation for it would improve the paper. For example, the experiment varying starting warmup learning rate can be added to demonstrate that starting from $\\eta_c$ value results in a shorter or more stable warmup than for lower and higher values.\n3. Test accuracy heatmaps are not provided for the experiments on the initial learning rate selection for warmup. Hence, it is not clear whether the proposed strategy results in high-quality solutions. Moreover, high values of $T_{\\text{save}}$ are observed for the small $\\eta_{\\text{trgt}}$ and large $T_{\\text{wrm}}$ (Figure 5b), but this configuration is clearly suboptimal since much shorter warmups work well for the low learning rates. For the most practically interesting hyperparameter regions associated with the shortest effective warmup for each learning rate, $T_{\\text{save}}$ is negligible. \n4. The maximum test accuracy achieved by GI-Adam appears indistinguishable from that of baseline Adam in most cases, so it is difficult to say whether the GI-Adam improves the training. Moreover, adding standard deviations to this comparison seems important due to the noisy behavior of the test accuracy. Also, GI-Adam is not the first Adam modification that increases stability at the beginning of training (see, e.g., RAdam from https://arxiv.org/pdf/1908.03265v4). A more accurate discussion of such methods and comparison with them would benefit the paper. \n5. The paper lacks experiments with Transformer architecture, while it is a primary case of using Adam optimizer with warmup. \n\nMinor comments\n1. Line 64: the idea that warmup is unnecessary if training is stable with the chosen learning rate is obvious and widely used in practice.\n2. Lines 185-186: I would not say that $\\mu$P does not benefit from warmup at all. As shown in Figure 3a, a longer warmup in the case of MSE training extends the range of converging learning rates.\n\nQUESTIONS:\nI would kindly ask the authors to address the main concerns from the Weaknesses section and focus on the following questions:\n1. Could you please summarize the main novel insights of the empirical analysis part of the paper compared to Gilmer et al., 2021 and Karla et al., 2023, and explain why analyzing the warmup behavior in different sharpness evolution regimes is important?\n2. Could you please provide any experiments demonstrating that starting warmup from the critical learning rate $\\eta_c$ is an effective strategy?\n3. What is the test accuracy with optimal warmup hyperparameters for baseline Adam, Adam with $\\eta_{\\text{init}}=\\eta_{c}$ and GI-Adam? Is there any statistically significant difference between them?\n4. Is the same self-stabilization mechanism observed when training Transformers?\n\nAdditional minor questions:\n1. Why do you use different target learning rates for $\\mu$P and SP in Figure 1? It is a bit confusing since the target learning rate and parameterization are changed between the two experiments at the same time, and it is not clear which change results in which effect. At the same time, Figure 2 uses an identical learning rate for both initializations.\n2. How do you measure the loss value for initial learning rate selection and the squared gradient for GI-Adam in the stochastic variants of the algorithms? Do you use a single batch or estimate these values over multiple batches? Using a single batch may result in higher variance in the estimates, which could be undesirable. On the other hand, estimating over several batches would incur additional computational costs.\n\nThere also exist several related works which the authors may find interesting:\n* Lobacheva et al., 2021 (https://arxiv.org/pdf/2106.15739) report an effect similar to warmup self-stabilization when training scale-invariant networks with weight decay and a constant learning rate. The decreasing weight norm increases the effective learning rate, which eventually leads to training instability and catapults. This periodic behavior allows the network to achieve flatter optima with higher test accuracy after several cycles.\n* A different cyclical behavior, the Slingshot effect, is observed in adaptive optimizers like Adam, as shown by Thilak et al., 2024 (https://openreview.net/pdf?id=OZbn8ULouY). This effect occurs in the terminal phase of training and involves a rapid growth of the last layer norm before catapulting, followed by an improvement in test performance.\n* In your experiments, optimal test performance is achieved with large learning rates close to the convergence boundary. However, Kodryan et al., 2022 (https://arxiv.org/pdf/2209.03695) show that training networks with weight decay and learning rates larger than optimal usually lead not to divergence but to a noisy stabilization of test error. Moreover, Andriushchenko et al., 2022 (https://arxiv.org/pdf/2210.05337) demonstrate that further reducing the learning rate from these stabilized solutions results in the model learning sparser features and achieving better final test performance.\n\nLIMITATIONS:\nThe authors adequately discuss the limitations of the paper in the conclusion.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}

comments: SUMMARY:
This paper finds that the learning rate warmup allow the network to tolerate larger learning rates. It gradually reduces the sharpness and forces the model to leave poorly conditioned areas of the lossscape and move toward flatter regions which can tolerate larger learning rates.

STRENGTHS:
1. This paper analyzes the mechanisms of warmup.
2. It also proposes a GI-Adam strategy, which is better than Adam.

WEAKNESSES:
1. Though the author claims to find that warmup allows for larger learning rates, it has been found in existing work such as "Gilmer, J., Ghorbani, B., Garg, A., Kudugunta, S., Neyshabur, B., Cardoze, D., Dahl, G.E., Nado, Z. and Firat, O., 2022, March. A loss curvature perspective on training instabilities of deep learning models. In International Conference on Learning Representations."  Further elaboration on the difference and novelty will make the paper more convincing.

2. They claimed to find "wasted time can be saved by making use of the catapult mechanism", but it seems this has been revealed in "Lewkowycz, A., Bahri, Y., Dyer, E., Sohl-Dickstein, J. and Gur-Ari, G., 2020. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.02218."

Not only the above two points, I can't well understand the novelty of this paper. I suggest further summarize the contribution part.

3. Could you give more explainations on the accuracy maps, such as Figures 3 and 4?

QUESTIONS:
I don't have questions.

LIMITATIONS:
Please refer to the Weaknesses
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-17