Why Warmup the Learning Rate? Underlying Mechanisms and Improvements
Authors: Dayal Singh Kalra
Maissam Barkeshli
Venue: NeurIPS 2024 poster (2024)
Keywords: Learning rate Warmup
Optimization Dynamics
Primary area: optimization_for_deep_networks

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
Dayal Singh Kalra
Maissam Barkeshli
Abstract
It is common in deep learning to warm up the learning rate 17. often by a linear schedule between Thnit = O and a predetermined target fugt- In this paper, we show through systematic experiments using SGD and Adam that the overwhelming bene- fit of warmup arises from allowing the network to tolerate larger '/rrt by forcing the network to more well-conditioned areas of the loss landscape. The ability handle larger V/ug 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 progressive sharpening or sharpness reduction phase, which in turn depends on the initialization and parameterization. Using these insights, we show how /iii 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 initialization for the variance in Adam which provides benefits similar to warmup.
Introduction
One of the most important choices to make in gradient-based based optimization is the learning rate (step size) 1). 1) is too small, then learning may take place too slowly or the model might get stuck in unfavorable regions of the loss landscape. If 17 is too large, training will typically diverge. In practice, it is common to pick a dynamical learning rate schedule 7) 22 40 26]. Modern learning rate schedules for deep learning typically consist of a warmup period where Th is increased linearly from
zero to a target value T/rr over a warmup time T Tem to eventually decay the learning rate, for example via cosine decay schedule (34 26 40
13 34
After the warmup period, is common
Given that warmup standard in the practitioner's toolkit, it is important to understand it deeply and identify improvements. In modern settings, perhaps the carliest work to use warmup was [14), which used a small constant learning rate for the first few epochs of training and then switched learning rate. A linear warmup schedule was later introduced in [13]. The intuition given was that to scale the minibatch size SGD by a factor of is natural to also scale the learning rate by factor of k, provided the model is not changing too rapidly and successive gradients are roughly aligned. However at the beginning training, the model changing rapidly, so natural to start with a lower learning rate and gradually increase it to the target value after the network has stabilized.
larger
Other explanations suggest that since the network is initialized randomly, the gradient steps 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 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,
Department of Physics and Joint Quantum Institute, University of Maryland, College Park Institute for Physical Science and Technology, University of Maryland, College Park
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

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 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 [3] 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, TinylmageNet, WikiTex(), 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 | 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 learming 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 withstand larger target learning rate in turn makes hyperparameter tuning of the target learning rate more robust, since the network responds well to 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 (33 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 rreccondiiooeed sharpness typically starts off at high values (even the large batch case). causing considerable instabilities at high learning rates. Such instabilities, which may be retained in Adam's '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 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 Tinit- 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.
Notations and Preliminaries
SGD(-M): Given gradients g. :- (o,(O,) step Stochastic Gradient Descent with momentum
updates the parameters O, using learning rate 7h and momentum m, with coefficient B. The update
equations are: m4+1 9: + 3,,,, and 0c+1 O. m 1. corresponds to SGD,
Adam: Adam 20 updates the parameters 0 according to the equations: m441 = 8im, + (1-31)g.
0 = B2 + (1 B2 ),? and 0 = 0, - 17p Vou-tte
m441 = 0₂ P:+1 where m
m.
1-09
2

---

0..6
(a)
(b)
T
RIG
0100
117
10
(c)
(d)
Twasing los
256
11
Figure I: Training loss and sharpness trajectories of FCNs trained on 5k subset of CIFAR- 10 with MSE loss using GD., the sharpness plot, the dashed lines represent the 2/1, curves, and when X above these curves, training exceeds the instability threshold (17 > ne). (top) jP with Tggt = 1/>5' (bottom) SP with Vg = 32/x". Similar mechanisms are observed across different architectures, loss functions, and mini-batch sizes, shown in Appendix B
and 0, = Pi are the bias birs-cor moments, c is small scalar used for numerical stability and Pt = (1 B1) [diag (v + eI] is the preconditioner.
This is defined by the schedule γ = Timit + (7/rg "Tii //Tm: The warmup rate = (Me -hal Tal/TTu.. Twrm corresponds to constant learning rate. Unless otherwise specified, we set Vinit = when referring to linear warmup. We propose strategies for selecting Tinit Section
Linear Warmup:
Sharpness:
The sharpness is defined as the maximum eigenvalue of the Hessian of the loss X Amax (V)L), with subscript indexing the training step. Adaptive optimizers, such as Adam, effectively perform gradient descent transformed space determined by Φ = p1/20 (for details, see Appendix B.I). Hence, their stability is determined by the largest eigenvalue of the pre-conditioned Hessian, denoted by AF
Amax V2L), rather than the sharpness itself.
Parameterizations in Neural Networks: The mechanism of warmup and its effectiveness heavily influenced by the network parameterization (see Sections and 5). Standard Parameterization (SP) 3 staple in common libraries 28 Another notable parameterization is the Neural Tangent Parameterization (NTP) [77,, which along with SP resides in the kernel learning class at infinite width. Ref. (37) proposed Maximal Update Parameterization (nP) 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 warmup is intimately tied to training instabilities. These training instabilities, often referred to 'catapults' [22 5). arise when the learning rate exceeds critical threshold Me. where both 17 and 7)c generally change with time. The critical learning rate 7]c is influenced by variety of factors, including the choice of optimizer 6), mini-batch size 35 6), and
model properties such as depth, width, parameterization, and initialization [8 19 overview of instability thresholds, see Appendix
For detailed
When the instability threshold is exceeded (n) > 7c), two cases arise: (i) if the learning rate is higher than the instability threshold but smaller than maximum stable learning rate (which varies with time), ie., T) < 7 < Tmax training stabilizes through a self-stabilization process and training continues, (ii) if the learning rate exceeds this maximum stable learning rate 17 > Tmmax 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 training failure, where the loss fails to improve significantly over its initial value, as we demonstrate in Section
The self eeff-stabilization mechanism of GD can be understood through both empirical observations (Figure [T and a theoretical model. We first describe model derived by Ref. [7] that effectively captures this phenomenon. The model assumes that the top eigenvector changes slowly through
3

---

training and can be treated as constant and considers a cubic approximation of the GD dynamics around a reference point 0- The dynamics along the projection 2+ := u' u²(0, ) is given by two coupled non-linear equations:
I++1
(1 - T Jar:
H = X + ( Bxz)
(I)
where a V quaatiffs the instantaneous change in sharpness and 8 := V 11 controls the non-linear change in sharpness. In this model, an instability arises when 1/: > 7) = 2/x"" Ref.
7 considered a constant learning rate 7 and assumed progressive sharpening (a 0). contrast, we consider a time-dependent learning rate and allow a to attain both positive and negative values in order analyze different warmup mechanisms.
The seff-stabilization mechanism manifests as four-step process 22 four steps of the self- ceff-stabliizttio mechanism using the above model and the Twm = 64 trajectories illustrated in Figure [c. d):
7 Below, we describe the
(1) Approaching instability:
Due to increasing learning rate and/or progressive sharpening, training approaches the instability threshold 1/t = 1/c 2/x Figure d), this occurs within the first 10 steps due to increasing learning rate.
(2) Blow up: On exceeding the instability threshold (n > ne), Equation predicts exponential growth in 2t, empirically manifesting as a sharp increase in loss, as observed in Figure c).
(3) Sharpness reduction: For small enough learning rates, |xy] (and the loss) continues to increase until the higher-order term in the sharpness update equation causes a decrease in sharpness (I, > %/9).
This is observed an abrupt decrease in sharpness in Figure dd). If the sharpness fails to decrease over extended steps, it may result in training divergence (e.g., see Twrm = case in the same figure).
(4) Return to stability: Once the sharpness has decreased appreciably so that A < 2, stability is restored and the loss eventually decreases.
While the self-stabilization process for more complex optimizers remains poorly understood, qualitatively similar mechanism 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 ne. As the learning rate is increased during warmup, training instabilities can be triggered. Assuming the warmup rate is not too high, these instabilities induce 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 the increased learning rate. In other words, 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 (1) < n)e). training
naturally experiences either a progressive increase or decrease in sharpness, as observed in Figure [ which is unrelated to warmup. For instance, consider the sharpness trajectories with Twm 1024 in the above figure. In Figure [bb, sharpness has a natural preference for increasing, whereas in Figure dd) it tends to decrease on its own. This natural sharpness evolution can defined as the sharpness evolution under gradient flow, corresponding to a in Equation (T). 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 (I) and the experiments in Figure to describe these distinct dynamics.
The combined effect of naturally increasing sharpness while the learning rate is also being increased results in a "head-on
(C1) Natural Progressive Sharpening (ax > 0; top row of Figure D

---

collision" at which the instability threshold is exceeded (71 > nc). This causes the loss to increase, leading to a decrease in sharpness. Once the sharpness has decreased appreciably, the stability is restored (1) <. me). 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 "h R2 2/2f x fle for the remainder of the warmup period, as seen in Figure bb).
(C2) Natural Sharpness Reduction (a < O; bottom row of Figure I]: 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 dramatically more pronounced decrease in sharpness than would naturally occur, ultimately restoring stability (7 < ne). 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 [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., Twm 64 in Figure ())). Upon reaching such region, the dynamics aligns with the natural progressive sharpening scenario.
When natural sharpness reduction is sigmificant (large negative a), 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, cooperates with warmup in decreasing sharpness. Conversely, if the natural trajectory of training is towards increasing sharpness, opposes the warmup's effort, leading a persistent cycle of catapults.
The Effect of Warmup Duration:
Given a fixed target learning rate Mugt. increasing the warmup duration Twrm delays the point at which training exceeds the instability threshold nec 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 decrease sharpness actively. Consequently, increasing Twrm results in catapults that are both delayed and smaller in magnitude, seen in Figure D 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 Twm allows the sharpness to naturally increase. As 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 lower warmup rate Muxt/Twwe.. As result, the fluctuations observed on exceeding the instability threshold are much smaller in magnitude, as seen in Figure [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 (JPP (37 in the large width limit appropriately using normalizing layers (c.g. standard Transformer architectures, see Figure [7 in Appendix E.5). are characterized by small initial network output. Such initializations start in flat regions where gradients point toward increasing sharpness [1], placing them in the progressive sharpening category (CI). As we will see in Section such initializations may not significantly benefit from warmup they already start in flat region. In contrast, large initializations, such as FCNS, CNNs, ResNets with Standard Parameteriza- tion (SP) initialized at criticality 22 31 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] [9] and adhere to the second mechanism, especially for prolonged warmups.
5

---

Tere
7am
(a)
Traitgg
10
107
3t1p
(d)
1000
(c)
107
(f)
Tranieba
10
0.30
0100
107
step
Figure 2: Training loss and sharpness trajectories of FCNs trained on the entire CIFAR- 10 dataset with MSE loss using full batch Adam. (top) simple-;PP (for details, Appendix D.2.1) with Tugt = 0.003 and (bottom) SP with learning rate Magt = 0.001. The dashed lines the sharpness figures illustrate the instability thresholds (2+20:)/ (1- Similar mechanisms are observed for different architectures, loss functions, and smaller batch sizes as detailed in Appendix
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
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 Appendix [E.2
During early training, the loss may decrease non-monotonically on incorporating momentum, even small learning rates. Such oscillations are also observed when quadratic loss functions optimized using GD with momentum [ These oscillations make it challenging to differentiate between warmup-induced catapults and fluctuations 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 itself evolves differently during training. It changes from 2/>"" initialization to (2+23/>X" later in training. Moreover, the late-time instability threshold is significantly influenced by the batch size, exhibiting much smaller value than SGD for the same batch size. These properties make more challenging analyze the training dynamics of SGD with momentum. Nonetheless, the fundamental warmup mechanisms closely mirror the vanilla SGD case. We leave 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 & = pl/20, we show in Appendix [.. This analysis suggests that their local stability should be determined by the largest eigenvalue of V p V.L,, which we refer to as the pre-conditioned Hessian. Indeed, this is what we observe in Figure which shows the training loss, pre-conditioned sharpness X 17 and sharpness trajectories for full batch Adam. In these figures. sharpness is significantly smaller than instability threshold (2+29.)/n 4000, indicating that sharpness does not determine stability. Instead, loss catapults are associated with P "H exceeding its instability threshold.
The pre-conditioned sharpness starts high for both progressive sharpening (simple-uP) and sharpness reduction (SP) scenarios considered the previous section. For simplicity. we considered a simpler version of μP, detailed in Appendix D.2.1 In particular, for jPP models, X ~ 10 despite being initialized in flat region as measured sharpness, while for SP models, 2o 1 109. These large initial values arise because P (1 B1) [diag(g]) + el] and go has components that are near zero.

---

WRN-16-4
MSE
SGD CIFAR-10
WRN-16-4
SGD CIFAR-10
(a)
2046
4500
MLUE
2048
4096
ICB
KCM
1004
85
85
532
512
10:0
MMN
MALL
256
126
25
175
M
SUW
70
K
37
M
KCM
16
15
40
10
WWW 91
2
M
1.7e--00
5.4e-00
1.de: 6.
de
6: 42
00
6
07
5 -0
Be-02
1.6e
2.5kv
MSE CIFAR-DD
WRN
CIFAR-10
2040
4206
6777
60:00
7.00
11
54)
2043
N0
HV
11
50
85
$12
250
PM
UTTE
517-
60
120
0.09
26:4
125
75
6/
20
64
NJ
32
11
60
73:0
65
13
PLY
45
40
NG
7.0603
33x-0
E4e-12
Ldedil
Sdedf
1.0e
2.0s.00
4
1.
3.7ed3
7.5md
L2
L
Mellll
3
Figure 3: Test accuracy heatmaps of WRNs trained on CIFAR-10 using different and parameteriza- tions loss functions using SGD: (a) μP and MSE loss, (b) μP and cross-entropy loss, (c) SP and MSE loss, and (d) SP and cross-entropy loss. Empty cells correspond to training divergences. Similar phase diagrams are generically observed for different architectures and datasets, shown in Appendix F
The large Xa- X H 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 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, seen in Figure 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- pre-conditioned sharpness. Similar to the momentum case, Adam's 's stability threshold at late training times significantly decreases for smaller batch sizes 6). also shown in Appendix EE.4
2 dd), potentially leading training failures. Increasing the
Impact of Warmup on Training and Generalization
Here we investigate the impact of warmup on training efficacy and generalization by disentangling the role of Trr and Twm- Our key findings are that generalization capability primarily determined by 7ggt and that Adam is particularly sensitive to large catapults. The role of increasing Twin is to (i) allow the network to tolerate larger Flgs and (ii) move training further away from the divergence (failure) boundary, leading to marginal improvement in generalization.
We consider Weds (WRNs) and Transformers (LM) parameterized in either SP or jaP. WRNs are trained on CIFAR-10, 10, CIFAR-100 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 tross-entropy (xent) loss functions using SGD or Adam optimizers for a fixed training budget of T 10s steps unless otherwise specified. Training begins with linear warmup phase from Timit = 0 to Tar over Twrm steps. After warmup, training continues at Tug for the remaining training budget. In some cases, following the warmup period, we decrease the learning rate using cosine decay 224. Further details are provided in Appendix D
Experimental Setup:
7

---

LMM--755 wrktect 108
Adamw
LM-4-768 welctsot-C5
(a)
Ki
s+
1853
44
16A
4.4
S12
80
130
4.
S12
200
128
64
1.0
10
46
1.0
1.0
400
14
140
59
Ldwal
60edd
3.0-0M
3G
6
Figure 4: Test loss heatmaps of Pre-LN Transformers in SP trained
WikiText- 103 with cross- entropy loss for a single epoch using (a) Adam, and (b) GI-Adam (introduced in Section 66. Additional results are presented in Appendix F33
5.1 Stochastic Gradient Descent (SGD)
Figure presents heatmaps that show the best test accuracy achieved during training, plotted in the Thg-:TTrm plane for different parameterizations and loss functions, These phase diagrams of warmup also show the convergence-divergenee boundary, with empty cells indicating training divergences, illustrating the interplay between warmup duration and the maximum trainable 7/rgg- 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- eross-entropy loss and smaller initializations (like μP). The diminished benefit for μPP is likely due to its initialization in a relatively flat region of the loss landscape, which can already facilitate training at higher learning rates initialization. This consistent increase in maximum Vlrgt with warmup durations can be understood through the lens of warmup mechanisms 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.
described in the previous section. As observed in Figure T
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 (c.g.. Figure 3 C)). and (ii) models that attain optimal performance without warmup (e.g.. Figure 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 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, demonstrate that these results hold on incorporating momentum and employing cosine learning rate decay.
5.2 Adaptive Gradient Methods (Adam)
Figure 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
8

---

Frod
(a)
4546
4000
1024
1004
S13
3000
512
5000
255
17H
28A
64
64
1000
M4
200
1-eatd
Le
2.4
1e
7.4vll
1.4 no
1460
2
Line
2.3422
7ow
1:4-110
2.5
2.8
Figure 5: Heatmaps showing (a) the steps to reach Bugt (Treach) and (b) the effective steps saved (Tsave)
on setting Tmma = 7)c for WRNs in SP trained on CIFAR-10 using SGD with cross-entropy loss.
training at higher rates similar to SGD, warmup further improves Adam's '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 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 SGD where optimal learning rates observed near the divergence boundary. This discrepancy stems from Adam's '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 e remains large for an extended period, leading to small effective learning rate p/ As result, training struggles to escape high-loss regions. Therefore, longer warmup is more beneficial for Adam compared to SGD, as is crucial to stay away from the failure boundary.
Improved Hyperparameter Initialization Schemes for Optimizers
Initial Learning Rate Selection for Warmup: Setting the initial learning rate to Tinit = 0 is
common practice in warmup (27 g 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 Vini = appears suboptimal, as can significantly delay the learning rate from exceeding the instability threshold, thus delaying the primary effect of warmup.
An effective strategy involves setting Thii Te 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 Ne.. The method consists of two stages: (i) an exponential search, starting from an initial guess Too iteratively multiplies 7/ by factor k > 1 until the loss increases. This identifies an interval [Mwr. Nuppe] containing Te. (ii) a binary search further narrows down [Yhwr: Tuppo by evaluating the loss at the midpoint 7/mia Nerr2/. If the loss increases, T/uppr is updated to Mmiia otherwise, 7ew is set to 7/mid- This process repeated until the loss in the next step T.(0,) ) satisfies the condition (0)) < ((0o)(1 + 8). for some hyperparameter 5 >>000 O. For details, see Appendix B.3
By setting Tinic = Ne. training can achieve the target learning rate earlier. Consider the modified warmup schedule: πソ = Tuii + (t/TT(/ which attains Tug in Treach T [w (1-9o/nes) steps, saving Twm (Te/vme) steps. Incorporating the computational cost of additional forward passes Tip required for estimating T)c (~ 10 in number), and noting that one training step approximately equates two forward passes, the net computational savings is Tawve Twm (M</Twy) Te/2. Figure demonstrates how T Tresch and Taave vary with the Twr and T/gr-- For Thr < TJo, the target learning rate is reached in single step, nearly saving the entire duration of the warmup. whereas for (Yirg > ne). starting Thni Z 2/ can save up to half of the allocated warmup duration, although this saving diminishes on approaching the divergent/failure boundary.
9

---

(a)
0.3
107
25
Tranre tm
10
101
18
9.7
10
117
100
10 127
10
10
2000
stop
st0p
Figure 6: Training loss and sharpness trajectories of FCNs in SP. The experimental setup is identical Figure but with GI-Adam instead of standard Adam.
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 fraction of the maximum stable learning rate, such as T/nit
1os/100
GI-Adam: Improving Adam by Initializing The Second Moment using Gradients
In Section we observed that the pre-conditioned sharpness for Adam starts at a high value, even for low sharpness initializations like JPP, and can lead training failures at large learning rates. We propose Gradient Initialized Adam (GI-Adam), which initializes the second moment using the gradient squared, vn- = -.. In Appendix 1.2 we show that a bias correction not required when the second moment is initialized using the gradients. As result, GI-Adam can be viewed as standard Adam with an automated warmup given by f) =uug V $5-
This simple trick reduces the initial pre-conditioned sharpness by around two orders magnitude (more precisely by factor of vT- Figure (c.f. Figure d-)). Moreover, it consistently shows improvement over standard Adam across datasets and prevents training failures by pushing the training failure boundary to higher aggt as shown in Figure 4(). We provide additional results for different datasets in Appendix F3 In
-1-8)) (2) at initialization, preventing large catapults, as illustrated
Appendix we show that Gl-Adam consistently performs par or better than RAdam 23 while offering a simple modification Adam.
To further assess that the primary cause of instability during early training is the large pre-conditioned sharpness, we randomly initialize vo but with the same norn as the gradients at initialization. Like GI GI-Adam, this also results in improved performance as shown in Appendix 1.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 [.4] this modification eliminates the initial decrease in pre-conditioned sharpness. We leave 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 away from the failure (divergence) boundary. Looking under the hood, we found variety of underlying mechanisms, which also suggested several improvements for hyperparameter initialization. In Appendix A we provide
practical guidance for practitioners choosing the warmup duration.
Our analysis also motivates a potential parameter-fres 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 and defer further development to future work.
The maximum learning rate be written 1mmax /0m//>-1 P-1 Here showed how warmup effectively decreases X P which is local measure of sharpness. There is also another possible effect of warmup, that it can cause increase in Cmax. which can be viewed more non-local measure of sharpness. Further analysis is required to understand how warmup helps increase Cmax-
Limitations: Our experiments were conducted 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 31, 82.
10

---

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 the draft. The authors acknowledge the University 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
[1] Tiny imagenet challenge. pdf
https://cs231n. stw.iis.ssaanooddddooo.oo. http://hhsaaaaooteeee///////////
[2] Stephen Boyd and Lieven Vandenberghe. Convex optimization. Cambridge university press, 2004.
[3] 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.
[4] Sébastien Bubeck et al. Convex optimization: Algorithms and complexity. Foundations and Trends@ in Machine Learning, ((3-4):231-357,201
[5] 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.
[6] 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.
[7] 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.
[8] 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ć, Wojciech Stokowiee, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020.
[9] Emily Dinan, Sho Yaida, and Susan Zhang. Effective theory of transformers at initialization, 2023.
[10] 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.
[11] Gabriel Goh. Why momentum really works. Distill, 2017.
[12] 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.
[13] 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 I hour. arXiv preprint arXiv:1706.02677, 2017.

---

[14] 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.
[15] 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.
[16] Dan Hendrycks and Kevin Gimpel. Bridging nonlinearities and stochastic regularizers with gaussian error linear units. CoRR, abs/ (ass/1606.08415, 2016.
[17] 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. Fesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018
[18] 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.
[19] Dayal Singh Kalra, Tianyu He, and Maissam Barkeshli. Universal sharpness dynamics in neural network training: Fixed point analysis, edge stability, and route to chaos. arXiv preprint MxKiv:3311.02077 2023.
[20] Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), San Diego, CA, USA, 2015.
[21] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[22] 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:2003002/18 2020.
[23] 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.
[24] Ilya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts. In International Conference Learning Representations, 2017.
[25] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In International Conference on Learning Representations, 2017.
[26] Kevin P Murphy. Probabilistic machine learning: introduction. MIT press, 2022.
[27]
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, high- performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc, E. Fox, and R. Gareett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[28] Adam Paszke. Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Ginelshein, Luca Antiga, al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[29] Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. In NIPS, 2016.
[30] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models unsupervised multitask learners. 2019.
12

---

[31] Daniel A. Roberts, Sho Yaida, and Boris Hanin. The Principles of Deep Learning Theory. Cambridge University Press, 2022. httpe///depleeaminghhorrr com
[32] 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.
[33] 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.0733 2020.
[34] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[35] Lei Wu, Chao Ma, et How sgd selects the global minima in over-paranne poar-prammterized learning: dynamical stability perspective. Advances in Neural Information Processing Systems, 31, 2018.
[36] 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 Proceedings of Machine Learning Research, pages 10524-1058 10533. PMLR, 13-18 Jul 2020.
[37] Greg Yang and Edward J. Hu. Tensor programs iv: Feature learning in infinite-widtt 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.
[38] Greg Yang. Edward 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.
[39] Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. 2017.
[40] Aston Zhang, Zachary C Lipton, Mu Li, and Alexander Smola. Dive into deep learning. Cambridge University Press, 2023.
[41] Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. In International Conference on Learning Representations, 2018.
13

---

A Practical Guidance for Practitioners
How to Select the Warmup Duration? Given a target learning rate Tggr if the training loss during the warmup period exhibits large instabilities (loss spikes), the warmup duration Tmm should be increased until such instabilities are sufficiently small. This effectively moves training away from the divergent/ failure boundary, as illustrated Figure 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? 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 least greater than the instability threshold initialization.
When to Decay the Learning Rate? Figure suggests that employing learning rate decay at small learning rates result performance degradation for fixed training budget. Therefore, the learning rate should be decayed at large target learning rates only. The underlying intuition is that we large target learning rates to train in 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 Efffcient 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 JuP, is advantageous. These initializations might allow for achieving optimal performance without the need for warmup, as observed in Figure
B Instability Thresholds
B.1
Instability Thresholds for Adaptive Optimizers
For adaptive optimizers, such as Adam, a pre-conditioner P.+1 multiplies the gradients in the update equation
0 = 0: Pi-+l (O,(
(2)
Next, we define Φ. = P1 then the gradient and Hessian are V and V2L p-1 VZL. As result, the update equations for & is given by
VoL(O.)
(3)
We observe that the update equation of an adaptive optimizer mirrors that of SGD under change of variables. Using the stability arguments of SGD, conclude that Amax (241 (()) 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 2/xØ" early training. Further investigation by Kalra and Barkeshli demonstrated that sharpness reduction during early training causes Ne increase with depth and 1/width. In such scenarios, ne can be as large as 40/20- Cohen al. [5] demonstrated that sharpness late training times for GD with momentum coefficient 3 oscillates above (2+28)/n, suggesting πe Z (2+20)/x² at late training times. Expanding on this, Cohen et al. 6 analyzed adaptive optimizers and found that for Adam, the pre-conditioned sharpness P oscillates around (2t2 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).
18
14

---

B.3 Estimating the Instability Threshold
This section describes the method for estimating the instability threshold de initialization (or generically, any point 8τ) using only forward passes. The method consists of two stages:
Exponential Search: An exponential search, starting from initial guess flo iteratively multiplies Yo by a factor = until the loss increases. This identifies an interval (/wr. Muper containing fJe- The
detailed algorithm described in Algorithm | Unless specified, we use fo 10 as our initial guess. If the loss already increases at the initial guess njo, set Tait 7ju-
Algorithm Exponential Search
Input: (Initial weights: do. Initial guess: 170)
2: Output: Interval [Twr. Nuppr] containing 1ec
3: Evaluate initial loss I(Ob))
4: Initialize 7 +
5: 0 t Optimizer(n) 0o)
6: Evaluate L(O1)
7: while L(O,) L(Oa) do
8: if 2 'trrtt then
9:
17 f- Tirgt
10: break end if
11:
12:
1) 2m
13:
0 < Optimizer(n, 0n)
14:
Evaluate L(02)
15: end while
Thpper f 1) 17: Thwr t- 1
18: return (/wr, Tuppr
Binary search: A binary search further narrows down [Twr: Nuper by evaluating the loss at the midpoint 'mmid Rups/// the loss increases, Nuprr is updated to Jama: otherwise, f/ww is set to 7mid This process is repeated until the loss in the next step L(Ø,) satisfies the condition L(,)) < L( + 6), for some 8 The algorithm is detailed in Algorithm 2 In all our experiments, we set 0.1.
Algorithm 2 Binary Search
Input: (Initial weights: Oo, Tolerance: Initial search interval (Yter Mapper
2: Output: Estimate of T/c
3: Evaluate L(Og)
4: 0, --tiimizer(Papppn Nupprr 0o)
5: Evaluate Luppe t- L
6: while Luppe (1 8) do
7: Amid & (7wr + Mupprr 1/2
8: Omid +- Optimizer (Tmim 0u)
9: Evaluate Lmnid <- L(emid)
10:
if Tomid < I(ff)) then
11:
Tiwr t- Mnid
12:
else
13:
Tappe f Mmmid
Luippe ((Omia))
15:
end if
end while
17:
return 17c + Nuppr
While both no and 8 are additional hyperparameters, the method does not heavily depend on these choices. A poor initial guess of no would only take few more iterations to find an interval [[we, Thppr]
15

---

Tam
64
256
1024
Training
10"
10
10P
10 10
10
10
10
10"
step
(a)
(b)
Figure 7: Comparison of persistent catapult warmup (in black) with linear warmup with different durations. The experimental setup is the same in Figure [ but the model is trained on the entire CIFAR-10 10 dataset using SGD with a batch size = 512.
Meanwhile, any small value of E (0, 1] is effective in finding Te. as small initial loss spikes have minimal impact the overall dynamies. Note that for Adam, small 5 (~> 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 rre-condiioonedd sharpness), thereby facilitating training at higher learning rates. Given a target learning rate 7/gg,, the strategy consists of the following steps:
1. Start with 'stable' reference point 0*, defined as point where the loss decreases in the next step and estimate the interval (Thrr, Thoppe] containing as described in Appendix B.3
2. Induce catapult by increasing the learning rate to
T/pppr-
3. Continue training and wait until the loss falls below the reference point, i.e., L(O,) L(8*). This new point now becomes the stable reference point.
4. Repeat the above steps until the target learning is achieved, i.e., Yhigt-
Here, the initial stable reference point the model's initialization. The detailed algorithm described in Algorithm 3 in Appendix C
Figure 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 T]e 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 more efficient and effective warmup.
Although persistent catapult warmup 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-fi 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: hhtpps
//gi gi thub. com/ dawallaaalaaaaaaaaaaaaaa warmup
16

---

Algorithm 3 Persistent Catapult Warmup
1: Input: (Initial weights: Oo. Target learning rate: Vlgt, Tolerance: 8)
2:
0* Bg Reference point
3: while 17 < Tuugt do
4: if' L(Oz) L(0r) then
5:
Estimate [upppr Thwrl containing 7e using Algorithms and with tolerance
6:
!7 x 1/uppr 8" + 0
7:
8:
else
9:
continue
10:
end if
11:
end while
D.1 Datasets Details
D.1.1 Image Classification Tasks
We consider standard image classification datasets such as CIFAR-IO, CIFAR-100 (21), and Tiny- ImageNet [.. 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 [4]]
D.1.2 Language Modeling Tasks
We consider the next token prediction task the Wikitext-2 and Wikitext- 103 datasets 225 The Wikitext-2 dataset consists of ~ 2M tokens, whereas the Vikitext-103 dataset has ~ 0.1B tokens. We use Byte Pair Encoding (BPE) tokenizer (32) with 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- non-linearittee and initialize the weights with truncated normal distribution | with variance o = 2.0 in appropriate parameterizations (details below), except for the last layer, which has weight variance of o 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 (0,, (.ttam,))) and the biases are initialized to zero.
Maximal Update Parameterization (μP): For pP, different schemes are employed for the inter- mediate and last layers. The intermediate layers are initialized using A0 ri/tan..) and the layer outputs are scaled by the factor / /Fanaa/fmm. In comparison, the layer weights are initialized with NO
(0, -²//an.,, and the final output is rescaled by the factor /i/fanm.. Conveniently, for SGD, the learning rate does 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 9//Emmer 1/vTum, and 1/fam.. Since we are utilizing pP only to obtain flat initializations, we omit the additional scaling of the learning rate for Adam in some experiments (e.g.. Figure 2). As result, the instability threshold only dependent on the target learning rate Thr during late training, rather than on the largest learning rate across layers. We refer to this parameterization as 'simple-uP" for Adam.
for details, see http:////....aaaaaoooo/////////s
nn. initializers
cruncated_nooeee www._ooraallllmllmmooo html
17

---

D.2.2 Architectures
Fully Connected Networks (FCNs): We consider fully connected networks with constant width of n and a depth of layers. These networks are denoted by FCN-d-n. Unless specified, we considered d = 4 layer FCNs with width 512.
WideResNets (WRNs): We consider WedeResNets [39 with d layers, 5 stages, and a widening factor of k denoted by WRN-d-k: The number of channels in each stage [0, S) is given by 2 x 16 x k, with the input layer having 16 channels. For example, WRN-16-4 consists of S = stages, each with [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 e.//tam, 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 Transformer with blocks and an embedding dimension as LM-d-n. Unless specified, the model has d = 4 blocks, embedding dimension n 128, context length Tenn = 64 and are trained for 10 steps.
D.3
Optimization Details
D.3.1
Optimizers
SGD(-M): Given gradients 9, step Stochastic Gradient Descent with momentum (SGD-M) updates the parameters O, using learning rate and momentum m, with coefficient B. The update
equations are:
m = 9: + 9m,,
(4)
0tt1 0, m
(5)
Here, 3 corresponds to SGD. In all experiments incorporating momentum, the default value of the coefficient set to 3 0.9.
Adam: Given gradients g: at step f, Adam 20 updates the parameters O, using learning rate n
and the first two moments of the gradient m, and v; with their coefficients }1 and f8,, respectively. The equations governing the updates are:
m++1 8,m, (1)99 )gr,
(6)
Ut+1 -32u; + (1 --22
(7)
m
0 0, 7) v 1z- +
(8)
where m. =
1y4 T-3 and 0, = 2/ T are the bias-corrected moments, and a small scalar used for numerical stability. The pre-conditioner for Adam is given by:
P = (1 -1-18 -11-3(t) +1+)
diag
1i+1
+ +el
8 a++i
In all experiments, the default values are set to 31 0.9, 33 0.999, and = 10-8, unless otherwise specified.
GI-Adam: Given gradients g: at step I, GI-Adam updates the parameters like Adam but with the second moment initialized using the gradients at initialization, i.e., vn = - The equations governing the updates are:
18

---

m&+1 = 3im; + (1 1--11)9 31)gt.
(10)
1/+1 = f22; + +1 -B2)g/
(11)
M
0.+1 = 0, 1 +
(12)
where 8N
01 is the bias-correction second moment and is small scalar used for numerical 1-33
stability.
Flat-Adam:
Given gradients 9; at step f, Flat-Adam updates the parameters 6, using learning rate Ye and the first moments of the gradients m, and v, with their coefficients 31 and 82, 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 Bo g6 The equations governing the updates are:
m++1
m, (1 B1 )gr.
(13)
4t+1 320; + (1 32)g
(14)
Mp++1
0++1 0: T)t 0 +
(15)
where 0, 1-33
v. is the bias-correction second moment and is small scalar used for numerical
stability. The pre-conditioner for Flat-Adam is given by:
Pr+1
diag
Ut++ 82
+el
(16)
V
D.3.2 Linear Warmup
Warmup linearly increases the learning from initial value Thii target value Arg over T Twm training steps. The learning rate at step is given by:
T = Vinit + (Thrg
Vinit
(17)
Twnn
is referred to the rate of warmup. Under the above definition, constant learning rate training corresponds to T
Here, a
(Dage That)
Tang
Twrm corresponds constant learning rate. Unless otherwise specified, we set Muii = 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 reduce the learning rate small value. Cosine decay is a commonly used method for decaying the learning rate from an initial value of Turt down value 7/min over Teos steps, according the rule:
1
T/. T
)
?) Vnrrgt + (Jmin Tugg)
1
cos
(18)
2
where / governs the rate of decay, with p = 1 being the standard. Note that with p= 0, the learning rate is not decayed and instead maintained at Tug- In the above expression, / counts the steps from the initiation of cosine decay and the current training step. As per standard practice, we consider 1 and decay the learning rate Timin 1h//10
19

---

D.3.4 Target Learning Rate Sampling for Phase Diagrams
For SGD, target learning rates Bggt are exponentially sampled using the initial sharpness X Starting with Vhrgt = 1/x0* subsequent rates are sampled until divergence as 2 2"/x" for values of æ increased in integer steps starting from zero. For WRNs trained with Adam, we sample target learning rates exponentially as TArgt 2 x 10-5, where is incremented integer steps starting from zero until training failure. For Transformers, we sample the learning rate in similar fashion but starting from 10 and increment æ in steps 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 Thtrr = 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 with a new initial guess for the eigenvector. We allowed for at most 10 restarts with the maximum number of iterations set to Witter 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 5k subset of CIFAR- 10 using MSE loss and GD in (top) μP with Tlg = 1/25", where X as 0.05, and (bottom) SP with Vrrgt 32/25' where X Rs 50.
Training loss and sharpness trajectories of 4 layer FCNs with width 512, in (top) pP with learning rate Thrg = 0.003 and (bottom) SP with Tegtt = 0.001 trained the CIFAR- 10 dataset with MSE loss using full batch Adam with 3 = 0.9, .3 0.999 and 10-8, these experiments, we use data augmentation as described in Appendix D.I.I.
Figure 2
Test accuracy heatmaps of WRN-16-4 trained on CIFAR-IO using different parameter- izations and loss functions using SGD with 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 steps. In these experiments,
Figure 3
data augmentation as described in Appendix D.I.I
Test loss heatmaps of Pre-LN Transformers in SP trained on WikiText- 103 with cross- entropy loss using (a) Adam, and (b) Flat-Adam (introduced in Section 6 over Adam. The Trans- former models have d = 4 blocks, embedding dimension n 768, a context length of Tenet = 64, and a batch size B 64. These experiments also employ cosine decay and weight decay with ==1 10-4
Figure 4
Heatmaps showing (a) Treach- number of steps to reach Tugt and (b) Teaee. the effective number of steps saved on setting 7un ne for WRN-16-4 in SP trained on CIFAR-10 with cross- entropy loss using SGD with B 128 for 10a steps. For a fair comparison with linear warmup, we Te/T.. as our initial guess.
Figure 5
choose 7o0
Figure 6 Training loss and sharpness trajectories of FCNs in SP. The experimental setup is identical to Figure 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 100 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
20

---

0.65
T
0.
D.C4
0,8
Trining
1024
2um
0.2
0.1
0.02
102
step
s01p
(a)
(b)
10
T
319
109
Training
10
10
100 107
10
10"
(c)
(d)
Figure 8: Training loss and sharpness trajectories of FCNs trained on CIFAR-I0 with MSE loss using SGD with a batch size B = 512. The dashed lines in the sharpness figures illustrate the instability thresholds 2/n. (top) pP with learning rate 1/25" (bottom) SP with learning rate 32/2//
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 A 100 hours. This brings the total computational time amounted to approximately 4500 A100 hours. Preliminary experiments took about 1000 A A100 hours. Hence, we estimate the total computational cost to be around 5500 A 1100 hours.
E
Additional Results for Mechanisms of Warmup
This section presents additional trajectories for warmup mechanisms discussed in Section covering various architectures, loss functions, and optimizers.
E.1 Stochastic Gradient Descent
Figure shows that the warmup mechanisms for full batch GD are also observed in the SGD with 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 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 To illustrate this, consider the full-batch GD with momentum. The middle row Figure demonstrates the sharpness reduction case with the learning rates well below the stability threshold (7 << ne). 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 loss due to the intrinsic effects of momentum. Nevertheless, we can still observe loss spikes correlated with an abrupt decrease in sharpness large learning rates, seen in the bottom row 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 7) for SGD with momentum evolves during training. For simplicity of explanations, again consider the full batch GD case. The stability threshold ne for SGD with
21

---

10-1
10
Trining
10m0
7
Zi
10-1
(a)
(b)
Mi
Trining los
0.44
11
117
10
(d)
101
Traning
110
200
10
step
stap
(e)
(f)
Figure 9: Training loss and sharpness trajectories of FCNs trained on 5k subset of CIFAR- using MSE loss and full batch GD with momentum 8 0.9: (top) μP with learning rate 1/x" (middle) SP with learning rate 1/x8" and (bottom) SP with learning rate 32/25. The dotted lines in the sharpness figures correspond to the (2+28/n, curves, while dashed lines show the 2/1, for reference.
momentum changes from 2/25" initialization (2+28)/>, late in training. At initialization, the momentum vector set to zero my = O, and the stability is given by vanilla GD threshold 2/x;" As training progresses, the momentum m, increases in magnitude, and the instability threshold at late training time becomes (2+28)/>, The bottom row of Figure 9 show the sharpness trajectories with both 2/1, and (2+28)/ curves. For Twm 64, the learning rate curve 2/7, causes abrupt decrease in sharpness, which is coupled with a loss spike, For longer warmup durations, the sharpness decreases before training exceeds the 2/x,5
Finally, the instability threshold T/c for SGD with momentum significantly decreases for smaller batch sizes. Figure 10 shows the training trajectories under the same setup in Figure 8 but with momentum coefficient (2+20)/n, (and even below 2/x(), whereas the vanilla SGD counterpart oscillates on the 2/7, curve. This indicates a strong dependence of the instability threshold on batch size.
0.9. The late-time sharpness trajectories oscillate well below the
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 those observed with MSE loss with one crucial difference. Near convergence, sharpness first increases and then abruptly decreases. The decrease in sharpness towards end of training observed in previous studies analyzing SGD with fixed learning rate (5). Additionally, we observe higher
22

---

0,
0.4
255
Tesking
RLCO
1024
0.00
0.2
0.1
201
11F
step
step
(h)
1234
Traning
10-1
Taza
11/
10
1024
112
10
(d)
Figure 10: Training loss and sharpness trajectories FCNs trained on CIFAR-10 with MSE loss using SGD with batch size B = 512 and momentum 3 0.9: (top) μP with learning rate 1/x5. and (bottom) SP with learning rate s²//a" The dotted lines in the sharpness figures correspond to the (2+28)/n, curves, while dashed lines show the ?/4. for reference. Similar mechanisms are observed for cross-entropy loss with a decrease in sharpness at late training times, as detailed in Appendix E.3
2.54
2.0
Traning bas
Tora
Tem
11:44
0.5m
stop
5060
(b)
101
318
Trining
35
Texs
10
101
step
(c)
(d)
Figure 11: Training loss and sharpness trajectories of FCNs trained on CIFAR-10 with cross-entropy loss using SGD with batch size B 512. (Top row) jP with learning rate 1/28" (Bottom row) SP with learning rate 32/2g-
fluctuations compared the MSE loss case. Figure shows trajectories 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 SP trained on CIFAR-10 with cross-entropy loss using full batch GD with and without momentum.
23

---

Trining
2
Pizza
10
10"4
110
stop
100
(a)
(b)
107
7-
210
10F
Training
10
10
100
10P
stop
(c)
(d)
Figure 12: Training loss and sharpness trajectories of FCN-4-512 in SP trained 5k subset of CIFAR-10 with cross-entropy loss using full batch GD with learning rate 32/2 with momentum coefficient (top) 3 0.0 and (bottom) 0.9.
0.4
T.
Tux
10
0.04
0.3-
10
MPN
Tur
0.022
-55
10
1330
10
13
10
10
37
10
107
317
(a)
(b)
(c)
T....
T
10
60
256
Training
1004
1234
16
0.0
10-
0.0
11
10 107 30
10'
107
10
10
107
(d)
(f)
Figure 13: Training loss and sharpness trajectories FCN-4-512 in (top) pP and (bottom) SP trained on CIFAR-10 with MSE loss using Adam with learning rate 0.001, batch size B = 512. 31 = 0.9 and 82 0.999. The dashed lines the sharpness figures illustrate the instability thresholds
(2+20L)/o.11-
E.4 Warmup Mechanisms of Adam
As discussed in Section 4.3 the instability threshold for Adam is determined by the pre-conditioned sharpness X 3 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 the same setting as in Figure 2 but with batch size of B 512. Similar to the SGD with momentum case, the late

---

30
Tems
200
100
1.34
16
10
0.0
0-
10
102
step
(b)
(c)
2.0
T
167
T
2,0
107
100
M
1.5
10
1.0
50
0.5
102
0.0
10
10
107
step
step
(d)
(e)
(f)
Figure 14: Training loss and sharpness trajectories of FCNs in (top) μPP and (bottom) SP trained on CIFAR- 10 with cross-entropy loss using full-batch Adam with learning rate 7 0.001, 31 0.9 and B2 0.999. The dashed lines in the sharpness figures illustrate the instability thresholds (2+20L)/nr(11
time sharpness oscillates far below the instability threshold (22+2.)/n.(1 suggesting that the instability threshold heavily decreases with 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
0,05
Turrs
10F
Tim
751
Training
0.04
1234
101
10P
step
stap
(b)
Tum
Tranng
1234
35 10P
31
10
10-1
201
11F
13P
109
(c)
(d)
Figure 15: WRN-16-1 trained on CIFAR-10 with MSE loss using vanilla SGD with batch size
B 512: (top) μP with Tag 1/20" and (bottom) SP with Thrt 32/2".
25

---

2.5
10
m
2.2
290
Taining
2024
T
2.0
101
1.0
10
11F
107
step
107
22.
1024
2024
step
stap
(c)
(d)
Figure 16: WRN-16-1 trained on CIFAR-10 10 with cross-entropy loss using vanilla SGD with batch size 512: (top) /P with 'hrgtt = 1/25" and (bottom) SP with T/rgt x//x
In the previous sections, we confined our analysis to FCNs to thoroughly explore the effects of different optimizers and loss functions. This section expands those results by demonstrating that the observed warmup mechanisms apply to ResNets and Transformers as well. The Resnet experiments also employ data augmentation detailed in Appendix [D.I]
Figures 15 and T6] 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 significant sharpness spike 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 the first LayerNorm parameters, which also resolves automatically by 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 (2 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 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 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.
Additional Phase Diagrams
This section presents further results related to the phase diagrams of warmup shown in Section 5
26

---

254
007
15
107
10
30
10
10
101
(a)
(b)
107
94
206
Traiing
14
L0G4
200
L024
12F
10'
102
10
17
(c)
(d)
Figure 17: LM-4-128 trained on the WikiText-2 dataset with cross eoss-ey loss using SGD with batch size B 512 and context length Tensx 64. The top shows the warmup mechanisms of a Pre-LN Transformer with Thg 5.65/26" while the bottom row shows the results for the same Pre-LN Transformer but with the last LayerNorm removed and learning rate of 1/ogt = 8/x0"
54
Tetta
Traning les
6024
715
:56
10
10f
200p
stop
(b)
(c)
Figure 18: Pre-LN LM-4-128 trained on the WikiText-2 dataset with cross-entropy loss using Adam with target learning rate Tug 0.003, a batch size 512 and context length Temix 64.
F.1 Phase Diagrams for different Models and Datasets
Figure 19 shows the test accuracy heatmaps of WRN-16-4 trained CIFAR-100 and -may- 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 F3
Figure 20 shows the test loss heatmaps of Pre-LN Transformer trained on the WikiText-2 dataset using SGD with a batch size B = 64. Figure 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 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 DD3.3 minimally affects the warmup phase diagrams. While the conclusions regarding warmup presented the main text remain unaffected, note few interesting observations.
27

---

WRN-16-4 SP XENT
SGD CIFAR-102
WRN-16-4 SP XENT SGD TINY- -MY-IMMEEETT
4096
2048
67.5
2048
3343
X2
41A0
SURS
AUA
1024
512
13
84
05
61
08
600
SUAT
65.0
1024
512
3007
4144
6:21
20.6
45
256
50.54
60
92
40
G1
02.5
128
110
126
33.99
41
50.5
40
00.0
64
3349
36
4143
4.
41
40
+0:00
50.00
4300
0327
58.00
474
6.2
97.5
16
35
55.0
4144
W
4982
3311
39
4134
14
41.0
90
2 438
10:00
0.30
4063
6.30
x
7255
Se-18
T.
14-00
2.h-
5.76400
1.
233--1
Me
9.
1.
3
50.0
2
1.
2.
2
4.4
1.76-8
3 3.3-01
6.76-81
1.3k--00
2.7e-00
25
(8)
(b)
Figure 19: Test accuracy heatmaps of WideResNets (WRNs) in SP trained on (a) CIFAR-100 and (b)
Tiny ImageNet with cross-entropy loss using SGD with batch size B = 128.
4896
LM-4-122 MM-4-128 XENT SGD WIKITEXT-2
6.00
LM-4-128 XENT SGD WIKITEXT-2
8.0
2848
F4
135
+20
421
418
4:1
400
6.75
284B
5.8
1004
TACA
5.50
5.6
S12
256
G4
47
455
4:2
1.00
6.25
512
504
1.
5:05
4:53
126
5.00
5.75
64
4.75
5.0
16
4.00
L8
U
455
4.25
k:
16
2-4e-11
3-4e81
4.560f
6 4.9--2
9.7e.0
day D0
Abx 00)
2 (0
3
5.4600
7.6e33
1.
1.5
2.2-02
300-02
4.3e
8.6
2e-01
17700
2.
Figure 20: Test loss heatmaps of LM-4-128 in SP trained on WikiText-2 with cross-entropy loss
using SGD with batch size B 64: (a) Pre-LN Transformer and (b) Pre-LN Transformer without the last LayerNorm.
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 further enhancement, indicating a synergistic interaction between the two.
F.3 Phase Diagrams of Adam and GI-Adam
25 compare the warmup phase diagrams of Adam and GI-Adam of WRNs trained on CIFAR- 100, Imageeete and of Transformers trained on WikiText-2 and Wikitext-109 103 dataset. Similar to the results shown in the main text, GI-Adam enhances performance over standard Adam by pushing the failure boundary.
Figures 23
F.4
Comparison with RAdam
Figures 26 and 27 and table T compare the performance of GI-Adam with RAdam. We observe that GI-Adam performs par better than RAdam while offering simpler modification to Adam.
28

---

WRN-1644 SP XENT SGD CIFAR-10
WRN-1664 SP XENT SGD-M CIFAR-10 10
4096.
90
2046
N-23
30
10.00
SQAL
2048
ALST
VLUA
ALOB
1024
RCAA
BCM
85
1024
512
20.00
80
512
22
23
3039
128
126
76
04 2.31
81
KLO
SLA
BLSC
$0.00
42
64
0862
944
40
22-40
10:30
16
63
M
53
1.8003
7.3ed3
125
1e
Se-01
7e-01
9 R-4-01
1.3k-
60
1.8k-3
3.76-0
7.36.3
1.5
2.0
5.96-02
1.26-11
2.3e-1
4.76-8
9:4e-01
1k
50
(a)
(b)
WRN-1664
XENT SGD CIFAR-10
WRN-16-4 SP XENT SGD-M CIFAR-10
4096 X0.
21:2
4096
AUM
MILL
WE
MUDI
MAN
2046
2048
00
1004. 11.4
12
1074
ALUS
85
512 70
10
08.00
REAT
$0.30
512
WMT
12
21
90-3
25.15
33.9
256- 7030
21
90
75
04 70.1
REAT
64
9030
S155
32 71.77
32
70
7030
63
60
40
00
71.1
55
53
1.8633
3.76433
7.3e03
1.5e.22
129
$.0-02
1.2-01
2.9600
7e-01
R46-8
50
50
1.800
3.76.03
7.3600
1e000
2.0002
6.9652
1.2501
2.3601
4.7001
9.:00
1 00
50
1.9ks
(d)
Figure 21: Test accuracy heatmaps of WideResNets (WRNs) in SP trained on CIFAR-10 with cross- entropy loss using SGD with batch size B 128: (top row) no cosine decay (a) no momentum, (b) momentum with 3 = 0.9, and (bottom row) with cosine decay (c) no momentum, and (d) momentum with 3 = 0.9. The setting of (a) is the same as in Figure 3cc but with different mini-batch sequence.
WRRN-1-4
SP
XENT Adam CIFAR-10
WRN-16-M
12-N-16-55
SP XENT GlAddm CIFAR-10
4098
11
NG
4096
PLTS
WENI
KUN
2048-511
4354
2045
0211
00
1004
1074
55
256
GEN
32
V0.9
5353
1325
250
5345
VUSD M-
54.4
20.02
95.55
50
128
31
70
126
1253
81
M:30
1250
70
64
DUDE
70
16
655
16
0G
8
24
60
65
56
Ge:tS
04
d.
44
63
2a-02
02
02
Ge:0t
60
8.0c00
14
d4
1.:00
1.0002
1e-02
01
Figure 22: Test accuracy heatmaps of WRN-16-4 trained on CIFAR-100 with cross-entropy loss
using (left) standard Adam, and (right) GI-Adam with batch size B 128.
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
29

---

WRIN-1644 SP XENT Adam CIFAR-100
WRN-16-4
SP
XENT GI-Adam CIFAR-100
4028-4229
2048
2048
70
1024 4040
1.99
1024
512 4230
512
12
15.00
3.33
1.35
256
65
0066
64
1.55
1.40
1140
60
37
000
16
16
500
2.44
80ed8
16x04
3.0ed
6.4604
26e03
200d2
A.
8.2e-22
3.5e-11
8.6e-01
8.0e05
16604
0.2e04
6.16
10-12
20--22
8.2e022
feedl
3.3-01
6.60-1
(a)
(b)
Figure 23: Test accuracy heatmaps of WRN-16-4 trained on CIFAR-100 with cross-entropy loss
using (left) standard Adam, and (right) GI-Adam with batch size B = 128.
WRN-16-4
16-4 SP
XENT Adam TINY- IMAENEET
WRN-16-4
4
SP
XENT GLAdem TINY-IMAGENET
4028-1331
55
2048
2048
1004.
1074
0545
1000
50
512 1303
500
20
11.00
3:00
512 3:
454
11
56
03.0
50
EAT
256 1381
3430
65
64
00
32 SR53
SUTO
50
325
32 X-7
32
05
16 4332
16 413
35
4,4311
4.
2 4.133
2. 0313
M
1,00
1.30
80
cm
1 4834
933
M0
25
8.0edt
1.fe-04
b4
6.4--0
1.3600
2.66
5.1e:3
2.0e22
A.1e-2
0.2e-2
1 1.6el
8.0e:55
1.5e00 04
2.66:3
5eled
1.0e-2
A.le002
à 4.3022
1.6e01
(a)
(b)
Figure 24: Test accuracy heatmaps of WRN-16-4 trained on ImageNet with cross-entropy loss using (left) standard Adam, and (right) GI-Adam with batch size B = 128.
LM-4-128 XENT
Adam Wiki Text-2
LM-4-128 128 XENT Gl-Adiam WikiText-2
4166
6.50
4364
0.75
28A2
200
200
33
5.50
VIDA
A50
4.00
412
430
880
343
204
0.30
0.99
5.25
S12
156
256
AX2
4XR
5
412
4:0
6
4.00
4.25
41
5
400
4.00
1.
240e-04
4.0e-00
8.
51e-12
1.0em
4.Jedll
6.2681
1.0cti
2.
4.0e.l
16-00
6.e-13
02
O2
21
(a)
(b)
Figure 25: Test loss heatmaps of LM-4-128 in SP trained WikiText-2 with cross-entropy loss
using (a) standard Adam, and (right) GI-Adam with batch size B = 64.
instabilities, the model's performance is severely impacted, resulting in training failures rather than convergence reasonable minimum.
30

---

2544
MM-47668 wikitec-603 163, Empeentmppy Adamw
45
2048
LM-4-768. wikitct-163 Creeentrap.. Gi-Adamw
12
345
12
335
1%
X
Set
30
1034
18
EET
13
1
512
10
512
F1
43
265
1
256
442
4
133
Ten
129
4.2
64
110
16
4.7
113
4.9
F18
39
1x
1M
308
422
24
Be-84
1.0
1.06.4
3.Bc-04
B.
6.0c-0
wikitect-
Crosentropy
FlatAdame
MM-4768
wkiteet
Grouentrapy
2048
2048
1024
4.4
1524
4.4
512
43
512
Tan
128
4.2
39
1.0e-0
3.0.01 00
100-0
1.0.0t
3.0e
B.Be-ff
1.0e00 00
9.0e 0
6:0e-0
1.0e D2
Rg
Figure 26: Test loss heatmaps of LM-4-768 in SP trained on WikiText- 103 with cross-entropy loss
using (a) standard Adam, (b) GI-Adam, (c) Flat-Adam, and (d) RAdam with batch size B = 64.
36
Tee: 00
Atsmm
B
Adam44066
Als tas 4206
G1-Adm--430
10
10
10
Figure 27: Comparison of optimizer performance on WRNs trained on CIFAR-10. The value next to the optimizer name corresponds the warmup duration. Each value corresponds to average over five initializations. Here, Adam-save corresponds to Adam with warmup 7)t Thoit + Thr T
These large loss catapults cause the gradients g to spike during early training, leading to substantial increase in its second moment While the gradients return lower value after few training steps, the second moment remains large in magnitude for prolonged period. These large values of result small effective learning rate, which hinders training escape these high-loss regions. Consequently, the models remain stuck in a suboptimal state rather than converging. We refer to this 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 's failures to future work.
31

---

Optimizer
Test Accuracy (mean std)
Adam-1
94.1315 = 0.0980
GI-Adam-1
94.6123 + 0.1273
Radam-1
94.4244 + 0.1052
Adam-4096
94.6163 = 0.2528
Gl-Adam-4096
94.6499 + 0.1112
Adam-save-4096
94.6084 +: 0.1664
Radam-4096
94.6460 + 0.2532
Table 1: Performance comparison of different optimizers with varying warmup durations.
80
Traning Sco/Dg
10:1
Trainin
150
1.0
12
10
10P
107
L(P
10
slep
step
(a)
1013
109
10
100
100
10
10
16-
10F
100
10
(c)
(d)
Figure 28: Training trajectories of WRNs trained on CIFAR-10 10 using Adam with cross-entropy loss and varying learning rates. The setup is identical to the Twrm = 1 row of Figure 4 but without employing cosine learning rate decay. The first training failure is observed at a learning rate of 'rrgt 0.02048. To investigate the behavior beyond the training failure boundary, learning rates are sampled from V/rrtt = 0.01024 (just below the failure boundary) up to 1/rgt R 150.
T
T
4096
1.10
4096
0000
2048
2083
2413
26.3
296
226
1024
KIX
X174
15113
1333
1003
512
1.05
258
256
128
128
403
125
122
1.00
64
648
513
963
860
2000
16-
16
0.95
44
44
1000
31
7-8--l
3.aa-00 0
6.2e-100 0
0.00
7.Brt
3.1e-90
6.7e-
Figure 29: Heatmaps illustrating the computational savings in reaching 7ug for WRN-16-4 pP
trained on CIFAR-10, using SGD and cross-entrope loss. Colored cells show the warmup steps T (left) and the total steps saved Tave (right), while the empty cells correspond to training divergence.
32

---

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 Teach 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, Te ~ TJmax. We observe that T reach for
wide range of learning rates, saving almost the entire warmup duration. Figures 30 and show similar phase diagrams for WRNs in SP and μP trained with Adam.
4096
4003
4006
4000
7048
3300
2048
VISE
3000 23R
1024
1024
A230
DAL
ASS
305
512
3000
512
3000
756
128
25003
128
2500
64
2000
64
EN
32
32
OW
15003
1500
500
45
66
25
36
05
05
500
2.0004
Ae:044
6.Ged4
3.2b-13
44400
1.3e02
226
5.1e.02
1.:01
2.0e01
1.06:04
2.0e04
4.0044
8 0e:04
1e:03
6.4
1.5--22
le:02
1.0e01
2
Figure 30: Heatmaps illustrating the computational savings in reaching Vrr for WRN-16-4 SP
trained on CIFAR-10, using Adam and cross-entropy loss. Colored cells show the warmup steps Tranch (left) and the total steps saved T save (right), while the empty cells correspond to training divergence.
Tynsh
Jane
4096
1.303
4000
2048
3500
1024
LERS
512.
1,00
512
XG-
16.0
3000
256-
256
201
26.
2016
128.
11,000
64
2000
32
31
250
16.
0:995
16
455
11.
115
20
00363
1000
0.925
500
1.0eH
2.06-44
Ge-04
8.0e.04
3.26-33
6 Ae-03
0.903
2.0e.0
9.0e
Li 03
G 4e400
L.3e-0
Figure 31: Heatmaps illustrating the computational savings in reaching Yurt for WRN-16-4 in uP
trained CIFAR-IO, using Adam and cross-entropy loss. Colored cells show the warmup steps Tranc (left) and the total steps saved Tave (right), while the empty cells correspond to training divergence.
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 X P-*H is around 2000
33

---

0.005
107
01:0
ILEGO
10
01:0
Tanig
01.00
10
0.20
10.00
60
01
300
10
102
(a)
Tranting
1224
10
(c)
10
Figure 32: Training loss and sharpness trajectories of FCNs (top) jPP and (bottom) SP. The experimental setup is identical to Figure but with GI-Adam instead of standard Adam.
Remarkably, this almost eliminates initial sharpness reduction. Similarly, the pre-conditioned sharpness for the SP model starts around 10 instead of 10% Notably, in the SP scenario, there is no initial spike in the Twm 1 (c.f. Figure 2 demonstrating that this simple modification effectively reduces instabilities during the early training.
instead of the value 10t observed for Adam (c.f. Figure 2)
1.2 GI-Adam as an Automated Warmup
In this section, we show that bias correction is not required when the second moment initialized with the gradients at initialization in GI-Adam. Therefore, employing a bias correction as in the original Adam algorithm in this case serves an automated warmup given by т): Argt / 3
The moving average of the second moment is given by:
= (1 - (2)
Bzgi- + +330.
(19)
1=0
where vh = -6 Following standard assumptions, we assume that the second moment of the gradient
constant during early training Elg]] = a². Taking the expectation of the above equation over the gradient distribution yields
E(ex)
(2)
B5E(vo).
(20)
Simplifying the above equation, we have
32)0 1- 31 1- 3300² 0².
E [..
(21)
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 a2. 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 8% Assuming small enough c, this can be viewed as multiplicative factor to the learning rate. As result, GI-Adan is equivalent to having natural warmup given by 1); Tug / 35.
34

---

0.5
Adn Adars Aboe
Test socuraty
0.6
Mar Adan Alti
00.4
Ridder
(b)
Figure 33: Comparison of test accuracy trajectories of WRNs trained with different Adam variants for two target different learning rates: (a) Thrgt 0.020480, and (b) TArtt 0.040960. For warmuu.
the warmup duration is set to T 1024.
1.3
The Primary benefit 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 XP we randomly initialize vo but with the same norn the gradients 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
1.4 Warmup Mechanisms of Flat-Adam
Figure shows the training trajectories FCNs with different parameterizations trained with Flat- Adam. We observe that further removing bias correction for the first moment completely removes the initial reduction in the pre-conditioned sharpness.
01005
Tu
167
0.55
100
(a)
Trattte
SCAA
10
10
07
18
161
102
100
0.3
167
000
250
102
(d)
Tomig
0.2
Sg 10
2034
15
10
10
Figure 34: Training loss and sharpness trajectories of FCNs in (top) μP and (bottom) SP. The experimental setup is identical to Figure 2 but with Flat- Adam instead of standard Adam.
35

---

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
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 perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how be expected to generalize other settings.
much the results
is fine to include aspirational goals as motivation as long 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
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create 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 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might 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 address problems of privacy and faimess.
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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant 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 complete (and correct) proof?
Answer: [NA]
36

---

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 cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide 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 supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibiliyy
Does the paper fully disclose all the information needed to reproduce the main ex- perimental 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)?
Question:
Answer: [Yes]
Justification: We provide extensive details of the experiments Appendix D to reproduce the results. Furthermore, 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, 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 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 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 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 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 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 submis- sions 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 reproduce that algorithm. (b) If the contribution is primarily
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 way to access this model for reproducing the results or 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 deseribe 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., 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 instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
37

---

Answer: [Yes]
Justification: The main results of the paper can be reproduced using the anonymous GitHub repository bttps://githube com http://aaaal--aaaaara/aaaaaaaaaa
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https for more details.
//ipps.c//ppi wttpccpppppppppeeeeeeeeeeeeeeeee
The authors should provide instructions on data access and preparation, including how 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 subset of experiments 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) 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, hyper- parameters, 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 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, 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 generating the warmup phase diagrams (100 A100 hours). we do not provide eπror bars. Nevertheless, duration, and we do not observe significant variation across the learning rate and warmup durations for various experiments. Hence, do not expect the results to vary significantly with initialization and mini-batch sequence.
have finely swept the learning rate and the warmup
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
38

---

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 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 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
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 compute required for each of the individual experimental runs as well 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 into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https::ps http://urripp..ccppcbe///eeeeeee hc./puubicc/tthcccc.cdee.....
Answer: [Yes]
Justification:
We adhere to the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- cration due to laws or regulations 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 warmup. We do not think there are any socictal impacts of the work that are worth mentioning here.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
39

---

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), faimess 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 tied to particular applications, let alone deployments. However, if there is 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 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 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 system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We do release data models.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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 challenging, and many papers do not require this, but encourage authors to take this into account and make 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 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 used and, 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.
40

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssit
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if has changed) should be provided. If this information not available online, the authors are encouraged to reach out to the asset's 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/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset 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. Crowdsoureing 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 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 the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then 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 equivalent approval/review based the requirements of your country 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 research with human subjects.
Depending 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.
41

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere 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.
42


Reviews:
[{"review_id": "wTR37r5gGL", "reviewer": "Reviewer_ZUpa", "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.", "questions": "I don't have questions.", "limitations": "Please refer to the Weaknesses", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. This paper analyzes the mechanisms of warmup.\n2. 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.\n\n2. 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.\"\n\nNot only the above two points, I can't well understand the novelty of this paper. I suggest further summarize the contribution part.\n\n3. Could you give more explainations on the accuracy maps, such as Figures 3 and 4?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Though the author claims to find that warmup allows for larger learning rates, it has been found in existing work such as Gilmer 2022. Further elaboration on the difference and novelty will make the paper more convincing....I suggest further summarize the contribution part.\n\nWhille Gilmer 2022 indeed demonstrated that warmup allows for larger learning rates primarily for models trained with SGD, our work extends and differentiates itself in several key ways that were not studied in prior work:\n\n* **Warmup mechanisms of SGD:** Gilmer 2022 demonstrated that warmup gradually reduces sharpness to facilitate training at higher learning rates. However, we go even further and ''look under the hood\" by showing that generically there are two distinct underlying mechanisms. We show that the interplay between natural sharpness evolution and increasing learning rate during warmup leads to either (1) a persistent catapult cycle in progressive sharpening cases, or (2) separated loss catapults in sharpness reduction cases.\n\n* **Warmup Mechanisms of Adam:** We provide a comprehensive analysis of Adam's underlying warmup mechanisms, especially during early training times, which was not analyzed in Ref. [1]. In particular, we find that the pre-conditioned sharpness, which determines stability, decreases during early training, *regardless of the natural evolution of sharpness.* This initial reduction in pre conditioned sharpness, before it eventually increases, suggests the existence of flatter initializations (wrt pre-conditioned sharpness) for Adam, which can enable training at higher learning rates from the start. Based on this insight, we propose a simple alternative initialization method called GI-Adam, which provides benefits similar to warmup and consistently improves over standard Adam by pushing the training failure boundary to higher target learning rates.\n\n* **The Phase Diagrams of Warmup:** The test accuracy heatmaps (Figures 3 and 4) demonstrate how the maximum target learning rate changes with warmup duration and disentangles the effect of warmup time and target learning rate. These phase diagrams reveal that the final performance primarily depends on the target learning rate, with longer warmup durations mainly helping to avoid the convergent-divergent (failure) boundary.\n\n* **Dual Advantage of Warmup:** Our results not only show that warmup improves model performance by allowing larger target learning rates, but also that warmup gives rise to a wider range of target learning rates that yield optimal results, which makes learning rate tuning more robust. This additional benefit of warmup depends on the phase diagram results and were not discussed in prior work. \n    \n* **Initial Learning Rate Selection:** As the primary effect of warmup is to facilitate training at higher learning rates by annealing sharpness (or pre-conditioned sharpness for Adam), setting the initial learning rate to $\\eta_c$ induces loss increase at initialization and thereby sharpness decrease right from initialization, saving warmup training steps. We provide a simple and practical method to pick $\\eta_c$ based on the loss catapult mechanism. \n\n* **Persistent Catapult Warmup:** We have also introduced 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 without specifying warmup duration. We demonstrate encouraging preliminary experiments in Appendix C of the submission. We intend to move these results into the main text.\n\n> They claimed to find \"wasted time can be saved by making use of the catapult mechanism\", but it seems this has been revealed in\n\nWe respectfully disagree with the reviewer. Lewkowycz et al 2020 introduced the catapult mechanism, however, it never suggested utilizing it as a way to estimate $\\eta_c$ and then setting the initial learning rate in warmup to be equal to $\\eta_c$. The development of this idea is unique to our work.\n\n> Could you give more explainations on the accuracy maps, such as Figures 3 and 4\n\nFigures 3 and 4 show the best test accuracy achieved during training as a function of target learning rate $\\eta_{\\text{trgt}}$ and warmup duration $T_{\\text{wrm}}$. These phase diagrams of warmup also show the convergence-divergence boundary, indicated by empty cells, illustrating the interplay between warmup duration and the maximum trainable learning rate. These results reveal: (i) flat initializations such as $\\mu$P benefit less from warmup, whereas large initializations such as SP may require long warmup to attain the optimal performance, (ii) The final performance primarily depends on the target learning rate and the improvement from increasing warmup duration comes from keeping training away from the convergent-divergent (failure) boundary, (iii) warmup makes learning rate tuning more robust, as mentioned above.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IhYzgsO3d9", "reviewer": "Reviewer_6XDG", "summary": "The 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.", "questions": "I 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.", "limitations": "The authors adequately discuss the limitations of the paper in the conclusion.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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.", "weaknesses": "My 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and effort in reviewing our paper and providing comments.\n\n> 1. Could you summarize the main novel insights ...\n\nWhile Gilmer et. al. 2021 is a key reference that our paper builds on, there are several novel insights in the empirical analysis part of our paper. First, our work emphasizes qualitatively distinct phenomena that arise at early training times, depending on whether the network starts in progressive sharpening or sharpness reduction phases. While Kalra et. al. 2023 also discuss sharpness reduction and progressive sharpening at early training times, the implications for learning rate warmup are not discussed. Our work combines ideas from both papers with novel empirical analysis to demonstrate for the first time the two separate underlying regimes of operation for learning rate warmup. \n\nA key result of our analysis is that models that experience sharpness reduction require more warmup compared to models that experience progressive sharpening at initialization. \nThis result forms the basis for examining Adam. Models trained with Adam only experience a reduction in pre-conditioned sharpness, regardless of natural sharpness evolution. It thus becomes evident why Adam generally requires warmup to perform well. This early reduction also suggests the possibility of flatter initializations for Adam, like GI-Adam, which we propose. Without a characterization of underlying warmup mechanisms in terms of early-time sharpness evolution, it is not obvious why Adam would generally require warmup.\nTherefore, the different sharpness evolution regimes are indeed related to the instability in Adam warmup. \n        \nAs for $\\mu$P, while it has been demonstrated in Kalra et al., 2023 that $\\mu$P is a flat initialization and one can speculate that it may not require warmup, our results show that models in $\\mu$P can still benefit from warmup (Fig 3), as also pointed by the reviewer. \n\nMoreover, our analysis in Sec 5 disentangles the role of $T_{wrm}$ and $\\eta_{trgt}$, which has not been performed in prior work. These results reveal that the final performance primarily depends on $\\eta_{trgt}$, with longer warmup durations mainly helping to avoid the convergent-divergent (failure) boundary. This also leads us to the point that warmup has another advantage: it makes learning rate tuning more robust. This point was not, to our knowledge, made in prior work. \n\nWe have also introduced persistent catapult warmup and shown some encouraging preliminary experiments in App. C. We intend to move these results to the main text.\n\n> 2. Could you provide experiments ...\n\nWe have conducted experiments to demonstrate the effectiveness of starting warmup from the critical learning rate $\\eta_c$. These results are shown in Fig 4 of the attachment to the global response. When training WRNs on CIFAR-10 using Adam, setting the initial learning rate to $\\eta_c$ (referred to as Adam-save in these results) yields solutions of similar quality to Adam. In this case, we save 576 training steps. In the updated manuscript, we will include test accuracy heatmaps for initial learning rate selection to further illustrate this point.\n\nAs described in Sec 6.1, the total number of steps saved is $T_{save} = T_{wrm} \\eta_c / \\eta_{trgt}$. For flat initializations, we have observed that $\\eta_c$ is close to the optimal target learning rates. This means that for the optimal target learning rates, flat initializations can save the full warmup time, $T_{save} \\approx T_{wrm}$. The benefit is less significant for large initializations where $\\eta_c$ is small. Fig 22 in App H shows heatmaps of $T_{save}$ for WRNs in $\\mu$P, where $\\eta_c$ is close to the maximum trainable learning rate. In these cases, we observe that $T_{save} \\approx T_{wrm}$ for most learning rates and warmup durations.\n\nOur analysis in Sec 4 demonstrates that before the `collision,' warmup does not actively reduce sharpness. Therefore, setting $\\eta_{init} = \\eta_c$ should not be expected to negatively impact the training dynamics. Furthermore, a priori, we cannot predict which part of the heatmap our chosen hyperparameters correspond to. Thus, adopting this strategy of starting from $\\eta_c$ does not negatively impact training and potentially saves a significant number of training steps in favorable cases.\n    \n> 3. What is the test accuracy with ...\n\nWe have consistently observed that for small warmup durations, GI-Adam outperforms Adam. As shown in Table of Fig 4 of the attachment, when training WRNs on CIFAR-10, GI-Adam improves test accuracy by $0.5\\%$ over Adam without warmup. This improvement is greater than the standard deviation in test accuracy across different initializations, indicating statistical significance. Moreover, in our language modeling experiments (detailed in the global response), we observed a loss improvement of $0.1$ with GI-Adam. For long warmups, we do not observe a significant advantage of GI-Adam over Adam. This is perhaps not surprising as GI-Adam performs a natural warmup, as described in the global response.\n \nAnother benefit of GI-Adam is that it widens the range of optimal learning rates by pushing the failure boundary further, as demonstrated in Fig 4(b) of our submission. This makes $\\eta_{trgt}$ easier to tune. Note that, a apriori, it is hard to predict where the experimental hyperparameters lie in the heatmap. Thus, it is always beneficial to use GI-Adam as it does not degrade performance, and yet it is a minor modification to Adam.\n    \nWe have also included Radam in our experiments for reference. GI-Adam performs similarly to Radam while being a significantly simpler modification to Adam. \n\n> 4. Is the same self-stabilization..\n\nWe have extended our experiments to include Transformers trained on language datasets with SGD and Adam. Our results readily extend to this setting. These results are detailed in the global response.\n\nWe provide replies to the minor questions in the comment below.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bavOSLCuxE", "reviewer": "Reviewer_LAU1", "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.", "questions": "See Weaknesses.", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The authors explain why the warm-up technique can help networks converge better.\n\n2. 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)?\n\n2. GI-Adam is used in [1].\n\n[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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and effort for reviewing our paper and providing comments.\n\n> 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)?\n\nWe have extended our experiments to include Transformer models trained on language modeling tasks with SGD and Adam. Our findings demonstrate that the conclusions drawn from FCNs and WideResNets generalize to Transformers trained on language modeling tasks. The additional results are presented in the global response.\n    \n> GI-Adam is used in [1].\n\nWe have reviewed Ref. [1] and respectfully disagree. Below we would like to take the opportunity to clarify the distinction between the Adam initialization used in Ref. [1] and our proposed GI-Adam. In their theoretical analysis, Ref. [1] initializes both first and second moments using gradients at initialization as a replacement for bias correction (Algorithm 1). This approach is primarily used for simplifying their theoretical analysis and does not include empirical evaluation of their proposed initialization.\n    \nIn comparison, GI-Adam initializes the second-moment using gradients ($v_0 = g_0^2$), while initializing the momentum to zero ($m_0 = 0$) and keeping the bias corrections. As shown in the derivation below, under standard assumptions for deriving the bias correction, initializing the second moment with the gradients at initialization does not require bias correction (also mentioned in Ref. [1]). Hence, for small $\\epsilon$, the bias correction on top of setting $v = g_0^2$ can be viewed as a multiplicative factor to the learning rate. As a result, GI-Adam is equivalent to having natural warmup given by $\\eta_t = \\eta_{\\text{trgt}} \\sqrt{1 - \\beta_2^t}$. \n     \nFurthermore, the referenced study has not performed any empirical analysis of their Adam initialization and it appears to be used as a replacement for bias correction for simplifying the mathematical analysis. In comparison, we show that GI-Adam has a smaller pre-conditioned sharpness at initialization compared to Adam and requires less warmup. \n\nThese distinctions highlight that GI-Adam is a novel approach, different from the Adam initialization in Ref. [1]. Nevertheless, we acknowledge the relevance of the reference and will cite it in the related works section of our updated manuscript.\n    \n\n**Derivation:** The moving average of the second moment is given by:\n\n\\begin{align}\nv_t = (1 - \\beta_2) \\sum_{i=0}^{t-1}    \\beta_2^i g_{t-i}^2 + \\beta_2^t v_0,\n\\end{align}\n\nwhere $v_0 = g_0^2$. Following standard assumptions, we assume that the second moment of the gradient is constant during early training $\\mathbb{E}[g_{t}^2] = \\sigma^2$. Taking the expectation of the above equation over the gradient distribution yields\n\n\\begin{align}\n    \\mathbb{E} [ v_t ] = (1 - \\beta_2) \\sum_{i=0}^{t-1} \\beta_2^i \\mathbb{E} [ g_{t-i}^2] + \\beta_2^t \\mathbb{E}[ v_0].\n\\end{align}\n\nSimplifying the above equation, we have\n\n\\begin{align}\n    \\mathbb{E}[v_t] = (1 - \\beta_2) \\sigma^2 \\frac{1 - \\beta_2^t}{1 - \\beta_2} + \\beta_2^t \\sigma^2 = \\sigma^2.\n\\end{align}\n\nTherefore, initializing the second moment with gradients at initialization does not require the corresponding bias correction.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sHI3M6dlHJ", "reviewer": "Reviewer_YVuK", "summary": "This 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.", "questions": "- 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)", "limitations": "Please refer to ``Weaknesses''.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Warmup 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.", "weaknesses": "- 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for reviewing our paper and they encouraging feedback.\n\n> The experiments are primarily ....\n\nWe have extended our experiments to include Transformers trained on language modeling tasks and found that our results extend to these models. The results are detailed in the global response. Furthermore, we would like to draw your attention to Appx F of our submission, where we show phase diagrams for CIFAR-100 and TinyImageNet. \n\n> observations in this paper are all based on..\n\nWe can use a toy model to understand the two warmup mechanisms. Following Ref. [3], we can understand the self-stabilization mechanism through a toy model resulting from a third-order approximation of the loss function.\n\nConsider a loss function $L(\\theta)$ with parameters $\\theta$. Let $\\lambda^H_t$ and $u$ denote the sharpness and top eigenvector. The model assumes that the top eigenvector $u$ changes slowly through training and can be treated as constant. Next, consider a cubic approximation of the dynamics along a reference point $\\theta^*$. The dynamics along the projection $x_t:= u^T (\\theta_t - \\theta^*)$ is given by two coupled equations\n\n$$ x_{t+1} = (1 - \\eta_t \\lambda_t^H)x_t, $$   $$ \\lambda_{t+1}^H = \\lambda_t^H + \\eta_t (\\alpha - \\beta x^2_t), $$\n\nwhere $\\alpha := - \\nabla \\lambda^H \\cdot \\nabla L(\\theta)$ quantifies the instantaneous change in sharpness and $\\beta:= \\|\\nabla \\lambda^H  \\|^2$ controls to the non-linear change. Ref. [3] considered a constant learning rate $\\eta$ and $\\alpha > 0$. Here, in contrast we consider a time-dependent learning rate and allow $\\alpha$ to attain both positive and negative values.\n\nIn this model, an instability arises when $\\eta_t \\lambda_t^H  > 2$. During instability, $x_t$ continues to grow until the higher order term in the sharpness equation causes a significant decrease in sharpness. Once the sharpness has decreased sufficiently, the stability is restored ($\\eta_t \\lambda_t^2 < 2$), and training continues.\n    \nNext, we consider the two natural sharpness evolution scenarios considered in our work:\n    \n1. **Natural progressive sharpening ($\\alpha > 0$):** The combined effect of naturally increasing sharpness ($\\alpha > 0$) and the increasing learning rate from warmup leads to instability ($\\eta_t \\lambda_t^H > 2$). Resultantly, $x_t$ increases until the higher order term in the sharpness update cause a decrease in sharpness ($x_t^2 > \\frac{\\alpha}{\\beta}$). Once the sharpness has decreased appreciably so that $\\eta_t \\lambda^H_t < 2$, stability is restored and the training continues. As training proceeds, progressive sharpening and increasing learning rate cause instability, resulting in a persistent catapult cycle characterized by $\\eta_t \\lambda_t^H \\approx 2$.\n         \n2. **Natural Sharpness Reduction ($\\alpha < 0$):** In this case, sharpness is naturally decreasing during training ($\\alpha < 0$). If the learning rate is increased quick enough relative to decreasing sharpness, an instability occurs $(\\eta_t \\lambda_t^H > 2)$. The increase in $x_t$ causes a more pronounced decrease in sharpness than it would have occurred naturally, restoring instability. To exceed the instability threshold again, the learning rate must significantly increase to account for the decreased sharpness. This results in one or more separated catapults.\n\nWe will include this toy model in the updated version of our paper, which will complement our experimental results in Section 4.\n    \n[3] Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability, ICLR 2023\n\n> It is natural and insightful that ....\n\nWe have already qualitatively demonstrated that increasing warmup duration facilitates training at higher target learning rates. Figs 1 and 2 illustrate that increasing the warmup duration results in smaller loss catapults, indicating improved stability at higher learning rates. This observation can also be argued through the toy model presented in the prior response. Figs 3 and 4 specifically show how the maximum target learning rate scales with warmup duration. To further address the reviewer's request for a more quantitative analysis, we can provide fitted curves to the maximum learning rate - warmup duration trends in the final version of the manuscript.\n\n> Does the measure of sharpness...\n\nOur extensive empirical analysis shows that even for randomized algorithms like SGD, the relevant stability measure is $\\lambda_{max}(H)$, not $||H||_F$. However, it is true that the stability threshold can change significantly for small batch sizes [1].\nPrior empirical work [4] has indeed shown that $\\lambda^H$ oscillates at a lower threshold than $\\frac{2}{\\eta}$ at late training times. Nevertheless, our results in Appendix E show that similar warmup mechanisms are observed for SGD, albeit saturating at a lower threshold (Fig 10 & 13).\n\nThe discrepancy in Fig 11(d) can be primarily attributed to loss function rather than the use of SGD. Fig 8(d) shows an experiment with the same batch size but with MSE loss. At late training times, $\\lambda^H$ oscillates slightly below $\\frac{2}{\\eta}$, suggesting a minimal deviation from GD. In comparison, for cross-entropy loss (Fig 11d), we observe that (i) sharpness oscillates slightly above $\\frac{2}{\\eta}$ during training and (ii) sharpness dramatically decreases towards the end. These phenomena align with the findings from prior studies. \nRef. [4] showed that sharpness decreases at the end of training for cross-entropy loss. Ref. [5] observed that for cross-entropy, the loss starts to catapult around $\\eta \\approx \\frac{4}{\\lambda^H}$. \n\n[4] Gradient descent on neural networks typically occurs at the edge of stability, ICLR, 2021\n\n[5] Phase diagram of early training dynamics in deep neural networks, NeurIPS 2023.\n \n> A closely related work [2] ....\n\nWe thank the reviewer for bringing Ref. [2] to our attention. We will incorporate a discussion of this paper in the related works section of our updated manuscript.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "guhb8fYigt", "reviewer": "Reviewer_kKLj", "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.", "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", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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. \nThe observations are made mostly from empirical study.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking the time to review our paper and for their encouraging comments. Below are our responses for the questions and comments.\n\n> The figures are not quite clear...\n\nWe noticed some errors in the Figure 3 caption in the submission and have revised it as follows:\n''Test accuracy heatmaps of WRNs trained on CIFAR-10 using different and parameterizations loss functions using SGD: (a) $\\mu$P and MSE loss, (b) $\\mu$P and cross-entropy loss, (c) SP and MSE loss, and (d) SP and cross-entropy loss. Empty cells correspond to training divergences. Similar phase diagrams are generically observed for different architectures and datasets, as shown in Appendix F.''\n\n> The observations are made mostly from empirical study,,\n\nWe can use a toy model to understand the two warmup mechanisms. Following Ref. [1], we can understand the self-stabilization mechanism through a toy model resulting from a third-order approximation of the loss function.\n    \nConsider a loss function $L(\\theta)$ with parameters $\\theta$. Let $\\lambda^H_t$ and $u$ denote the sharpness and its corresponding eigenvector. The model assumes that the top eigenvector $u$ changes slowly through training and can be treated as constant. Next, consider a cubic approximation of the dynamics along a reference point $\\theta^*$. The dynamics along the projection $x_t:= u^T (\\theta_t - \\theta^*)$ is given by two coupled non-linear equations:\n\n\\begin{align}\n        & x_{t+1} = (1 - \\eta_t \\lambda_t^H)x_t,\\\\\n        & \\lambda_{t+1}^H = \\lambda_t^H + \\eta_t (\\alpha - \\beta x^2_t),\n\\end{align}\n\nwhere $\\alpha := - \\nabla \\lambda^H \\cdot \\nabla L(\\theta)$ quantifies the instantaneous change in sharpness and $\\beta:= \\|\\nabla \\lambda^H  \\|^2$ controls to the non-linear change in sharpness. Ref. [1] considered a constant learning rate $\\eta$ and $\\alpha > 0$. Here, in contrast, we consider a time-dependent learning rate and allow $\\alpha$ to attain both positive and negative values.\n\nIn this model, an instability arises when $\\eta_t \\lambda_t^H  > 2$. During instability, $x_t$ continues to increase until the higher order term in the sharpness update equation causes a significant decrease in sharpness. Once the sharpness has decreased sufficiently, the stability is restored ($\\eta_t \\lambda_t^2 < 2$), and training continues.\n    \nNext, we consider the two natural sharpness evolution scenarios considered in our work:\n    \n 1. **Natural progressive sharpening ($\\alpha > 0$):** The combined effect of naturally increasing sharpness ($\\alpha > 0$) and the increasing learning rate from warmup leads to instability ($\\eta_t \\lambda_t^H > 2$). Resultantly, $x_t$ increases until the higher order term in the sharpness update cause a decrease in sharpness ($x_t^2 > \\frac{\\alpha}{\\beta}$). Once the sharpness has decreased appreciably so that $\\eta_t \\lambda^H_t < 2$ , stability is restored and the training continues. As training proceeds, both progressive sharpening and increasing learning rate cause instability, resulting in a persistent catapult cycle characterized by $\\eta_t \\lambda_t^H \\approx 2$.\n         \n2. **Natural Sharpness Reduction ($\\alpha < 0$):** \nIn this case, sharpness is naturally decreasing during training ($\\alpha < 0$). If the learning rate is increased quick enough relative to decreasing sharpness, an instability occurs $(\\eta_t \\lambda_t^H > 2)$. The increase in $x_t$ causes a more pronounced decrease in sharpness than it would have occurred naturally, restoring instability. To exceed the instability threshold again, the learning rate must significantly increase to account for the decreased sharpness. This results in one or more separated catapults.\n    \nWe will include this toy model in the updated version of our paper, which will complement our experimental results in Section 4.\n    \n[1] `Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability', Alex Damian, Eshaan Nichani, Jason D. Lee, ICLR 2023\n\n> The study is done on convnets and resnets applied to images...\n\nWe have extended our experiments to include Transformers trained on language tasks and found that our results also apply to attention-based methods. These results are detailed in the global response section.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their time and effort in reviewing our paper.  Based on reviewers feedback, we have added the following results:

1. **Language Modeling Experiments:** We have extended our experiments to include Transformer models trained on language modeling tasks using SGD and Adam. In these experiments, we trained 4 layer Pre-LN transformers on the WikiText-2 dataset. The results are shown in the PDF attached.
    
    * Figure 1 shows the warmup mechanisms for Transformers trained on Wikitext-2 using SGD. We observe that the initial sharpness of the Pre-LN Transformers in Standard Parameterization (SP) is surprisingly small (\~5) and exhibits progressive sharpening right from the onset. This low initial sharpness is due to the last LayerNorm, and removing this layer results in a large sharpness (\~200), and reveals early sharpness reduction behavior observed in FCNs and ResNets in SP. 

    * Figure 2 shows the warmup mechanisms of Pre-LN Transformers in SP trained on Wikitext-2 using Adam. Consistent with our findings in image classification tasks, we observe a reduction in pre-conditioned sharpness during early training, even for flat initializations that show an increase in sharpness from the onset.

    * Figure 3 presents the phase diagrams of warmup for Transformers trained on Wikitext-2 using both Adam and our proposed GI-Adam. In line with our image classification results, we observe:
        
        1. The final performance is primarily determined by the target learning rate and increasing the warmup duration keeps training further away from the divergence (failure) boundary.
            
        2. GI-Adam exhibits a wider range of target learning rates that achieve optimal performance, compared to standard Adam.
        
    These findings support the generalization of our conclusions to Transformers trained on language modeling tasks.

2. **Toy model for the two mechanisms of warmup:** We can use a toy model to understand the two warmup mechanisms. Following Ref. [1], we can analyze the self-stabilization mechanism through a model derived from a third-order approximation of the loss function.
    
    Consider a loss function $L(\theta)$ with parameters $\theta$. Let $\lambda^H_t$ and $u$ denote the sharpness and its corresponding eigenvector. The model assumes that the top eigenvector $u$ changes slowly through training and can be treated as constant. Next, consider a cubic approximation of the dynamics along a reference point $\theta^*$. The dynamics along the projection $x_t:= u^T (\theta_t - \theta^*)$ is given by two coupled non-linear equations: $$ x_{t+1} = (1 - \eta_t \lambda_t^H)x_t, $$ $$ \lambda_{t+1}^H = \lambda_t^H + \eta_t (\alpha - \beta x^2_t),$$
    where $\alpha := - \nabla \lambda^H \cdot \nabla L(\theta)$ quantifies the instantaneous change in sharpness and $\beta:= \|\nabla \lambda^H  \|^2$ controls to the non-linear change in sharpness. Ref. [1] considered a constant learning rate $\eta$ and $\alpha > 0$. Here, in contrast, we consider a time-dependent learning rate and allow $\alpha$ to attain both positive and negative values.

    In this model, an instability arises when $\eta_t \lambda_t^H  > 2$. During instability, $x_t$ continues to increase until the higher order term in the sharpness update equation causes a significant decrease in sharpness. Once the sharpness has decreased sufficiently, the stability is restored ($\eta_t \lambda_t^2 < 2$), and training continues.

    Next, we consider the two natural sharpness evolution scenarios:
 
    * **Natural progressive sharpening ($\alpha > 0$):** The combined effect of naturally increasing sharpness ($\alpha > 0$) and the increasing learning rate from warmup leads to instability ($\eta_t \lambda_t^H > 2$). Resultantly, $x_t$ increases until the higher order term in the sharpness update cause a decrease in sharpness ($x_t^2 > \frac{\alpha}{\beta}$). Once the sharpness has decreased appreciably so that $\eta_t \lambda^H_t < 2$, stability is restored and the training continues. As training proceeds, both progressive sharpening and increasing learning rate cause instability, resulting in a persistent catapult cycle characterized by $\eta_t \lambda_t^H \approx 2$.
     
    * **Natural Sharpness Reduction ($\alpha < 0$):**  In this case, sharpness is naturally decreasing during training ($\alpha < 0$). If the learning rate is increased quickly enough relative to decreasing sharpness, an instability occurs $(\eta_t \lambda_t^H > 2)$. The increase in $x_t$ causes a more pronounced decrease in sharpness than it would have occurred naturally, restoring instability. To exceed the instability threshold again, the learning rate must significantly increase to account for the decreased sharpness. This results in one or more separated catapults.
    
    We will include this analysis in the updated version of our paper, which will complement our empirical results in Section 4.
    
    [1] Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability, ICLR 2023

3. **Improved understanding of GI-Adam:** Since the submission, we have improved our understanding of GI-Adam. Initializing the second moment with initial gradients eliminates the need for bias correction (we can provide a derivation). Hence, for small $\epsilon$, a bias correction, in addition to setting $v = g_0^2$, can be viewed as a learning rate multiplier. As a result, GI-Adam can be viewed as Adam with a natural warmup given by $\eta_t = \eta_{\text{trgt}} \sqrt{1 - \beta_2^t}$.

4. **Persistent Catapult Warmup:** We have also introduced a parameter-free warmup strategy `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 without specifying warmup duration. We demonstrate encouraging preliminary experiments in Appendix C of the submission.

Meta-review: This paper provides a detailed empirical analysis of the warmup mechanism in neural network training, focusing on its role in improving training stability and enabling the use of larger learning rates. The authors also introduce the GI-Adam optimization method, which aims to enhance the initialization of the Adam optimizer. Strengths of the paper include the solid experimental validation of the warmup process, particularly its influence in different sharpness evolution regimes (as noted by 6XDG and kKLj), and the extension of these findings to Transformer models, which addressed concerns about the generality of the results (raised by LAU1 and YVuK). Additionally, the paper’s insights into the interplay between warmup duration and learning rate are appreciated for their potential practical applications in training neural networks (mentioned by kKLj and YVuK).

However, there are some weaknesses. ZUpa expressed concerns about the novelty of the paper, specifically the lack of a clear distinction from previous research, which was not fully alleviated by the authors’ response. Additionally, 6XDG pointed out that while the proposed GI-Adam method is reasonable, it does not offer significant improvements over existing alternatives like RAdam, which limits its practical impact.

While concerns remain regarding the novelty and the practical effectiveness of GI-Adam when compared to existing methods, I believe these issues can be adequately addressed if the manuscript clearly discusses these points in the main text. The numerical analysis of the warmup mechanism, focusing on different regimes during training, provides significant insights that are valuable to the NeurIPS community. Therefore, I recommend accepting this paper.
Decision: Accept (poster)