Analyzing & Reducing the Need for Learning Rate Warmup in GPT Training
Authors: Atli Kosson
Bettina Messmer
Martin Jaggi
Venue: NeurIPS 2024 poster (2024)
Keywords: learning rate
warmup
schedule
AdamW
LION
rotational
GPT
LLAMA
transformer
deep learning
empirical
momentum
critical batch size
gradient diversity
signal-to-noise ratio
Primary area: optimization_for_deep_networks
TL;DR: An empirical study on how learning rate warmup benefits GPT training by mitigating large initial updates, analyzed through various measures of update sizes

Abstract: Learning Rate Warmup is a popular heuristic for training neural networks, especially at larger batch sizes, despite limited understanding of its benefits. Warmup decreases the update size $\Delta \mathbf{w}_t = \eta_t \mathbf{u}_t$ early in training by using lower values for the learning rate $\eta_t$. In this work we argue that warmup benefits training by keeping the overall size of $\Delta \mathbf{w}_t$ limited, counteracting large initial values of $\mathbf{u}_t$. Focusing on small-scale GPT training with AdamW/Lion, we explore the following question: *Why and by which criteria are early updates $\mathbf{u}_t$ too large?* We analyze different metrics for the update size including the $\ell_2$-norm, resulting directional change, and impact on the representations of the network, providing a new perspective on warmup. In particular, we find that warmup helps counteract large angular updates as well as a limited critical batch size early in training. Finally, we show that the need for warmup can be significantly reduced or eliminated by modifying the optimizer to explicitly normalize $\mathbf{u}_t$ based on the aforementioned metrics.

Analyzing & Reducing the Need for Learning Rate Warmup in GPT Training
Atli Kosson
Bettina Messmer
Martin Jaggi
EPFL, Switzerland firstname www.aaameeeppllccmmmmmm ch
Abstract
Learning Rate Warmup is a popular hcuristic for training neural networks, espe- cially at larger batch sizes, despite limited understanding of its benefits. Warmup decreases the update size Aw; = ue early in training by using lower values for the learning rate 7/.- In this work we argue that warmup benefits training by keeping the overall size of Aw,, limited, counteracting large initial values of uy Focusing small-scale GPT training with AdamW/Lion, Why and by which criteria are early updates u, too large? ? We analyze different metrics for the update size including the f--norm, resulting directional change, and impact on the representations of the network, providing new perspective on warmup. In particular, we find that warmup helps counteract large angular updates as well as a limited critical batch size early in training. Finally, we show that the need for warmup can be significantly reduced or eliminated by modifying the optimizer to explicitly normalize un based on the aforementioned metrics.
explore the following question:
Introduction
Neural networks are typically trained using variations of stochastic gradient descent. The weight updates Aw have the form Au 17u, where 1) denotes the learning rate and unscaled update vector derived from the history of weight gradients. Throughout training, the learning rate is often adjusted over time according to a learning rate schedule, 1) = 7)- This schedule frequently includes an initial phase known as a learning rate warmup, where the learning rate starts low and is increased target value before being reduced according a decay schedule. Both the choice of warmup and decay strategy can significantly affect the final model performance. In this work, we focus on the linear warmup introduced by Goyal al. (10] for large batch size ResNet [I]] training, which also commonly used for transformers 37..
The length of the warmup is a hyperparameter that requires tuning, which is complicated by the fact that the reasons for its effectiveness are somewhat unclear. Warmup empirically helps stabilize training and allows the use of larger learning rates throughout the rest of training, which can speed up the process and provide beneficial regularization [10]. Since the learning rate simply scales the size of the updates Aw 7yu, warmup must achieve these effects by decreasing the size of early updates. However, it is not fully clear why this helps. Are the initial updates too large for some reason? For example, we might need small π) values counteract large u values early in training. How should we quantify what makes an update Aw large? Why do large updates adversely affect training?
In this work, we explore warmup from this perspective, focusing GPT2 29 optimizers like AdamW (24] and Lion We identify three key issues that necessitate warmup:
training with adaptive
1. The way Adam handles momentum can lead artificially large initial updates Aw.
2. Early updates Aw are large compared to the initial weight magnitude of w for matrices.
3. The gradients of early samples are highly correlated, limiting effective mini-batch sizes.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

the momentum bias correction in AdamW and scaling matrix updates to match their magnitude, akin to the Rotational Optimizers by Kosson et al. [9 We analyze the third issue in terms of the rate at which the internal neural representations of the network are changing (sometimes called feature learning). When the gradients of different samples are highly correlated these internal representations change too fast, which we conjecture can lead to issues with the non-linearities (e.g. activation functions) of the network. This can also be seen as the critical batch size [28] being too low early in training to enable the use of the peak learning rate. We derive a scaling factor based on the We demonstrate that simple modifications to the optimizer can mitigate the first two issues: eliminating
signal-to- ratio of the gradient that can help mitigate this, functioning like automatic learning rate warmup. Alternatively, we show that using high momentum values in conjunction to the first two methods may suffice to enable efficient training without warmup.
2 Related Work
Learning rate warmups have been used since at least ResNet [T]]. where lower constant learning rate was applied at the start of training. Earlier works may have employed similar concepts; for example, Sutskever et al. (35] utilized a momentum schedule that could induce a similar effect in the "effective learning rate" as defined by Fu al. [ The practice of linear warmup its current form was popularized by Goyal et al. [10 and Vaswani et al. (37.
Warmup has been studied indirectly in various neural network optimizer works. A notable example is RAdam 23 a modification of Adam [18] aimed at reducing the need for warmup. However, Ma and Yarats 26 demonstrated that RAdam essentially incorporates a fixed warmup schedule within the optimizer. In appx. C.T we show that this warmup effect is insufficient in our setup. Relative optimizers like LARS
and LAMB 48 are also believed to reduce the necessity for warmup 2I]. Bernstein et al. (2) propose a relative optimizer called Fromage and analyze how relative weight changes relate to representation changes, but differ from our approach in that they do not describe the effects of the gradient signal-to-noise ratio on this relationship. We build upon Kosson et al. 19 which showed that weight decay can make standard optimizers function as approximate relative optimizers and proposed variants that reduce the benefit of warmup without fully eliminating it.
The effect of warmup in transformers was empirically studied by Wortsman al. 41. Xiong et al. 43 proposed the pre-LN normalization placement for transformers, showing it reduces the need for warmup. Huang al. 13 studied initialization in transformers showing a link to warmup.
Finally, warmup has been studied directly on its own. Gotmare et al. [9] studied the effect of warmup. finding it helps avoid overly large updates to the weights of later layers which could be frozen to achieve a similar benefit. Gilmer et al. 7 study the need for warmup from a curvature perspective, showing it may help "push" the optimization trajectory towards flatter regions where higher learning rates are stable. Smith et al. [2] arrive at a similar conclusion, there is a stable learning rate that varies throughout training based on the curvature which limits the learning rate early on, necessitating warmup. These works focus on SGD with momentum, but it is less clear how curvature affects Adam-like or relative optimizers we discuss later.
The relation between stochastic gradient noise and learning rate has been studied in several works [4 28
49 34 27) They find that the update size can be increased roughly linearly with the batch size up to a certain critical batch size that depends on ratio of the mean and variance of the mini-batch gradient. We show how the signal-to-noise ratio (SNR) of the mini-batch gradient amplifies changes to the neural representations of network given a normalized update in weight space. We observe that the SNR starts out high but decreases over time, which translates to large early changes in the internal representations without warmup.
3 Baseline Experimental Setup & Results
Our main experiments focus on the training of 124M parameter GPT2 (29 model on the Open-
WebText corpus
[8]. The model has 12 transformer blocks with an embedding dimension of 768.
Our base training is performed at batch size 480 with a sequence length of 1024. We train for 5000
iterations which translates into roughly 20 tokens per parameter, suggested by Chinchila [12]
The baselines use Adam W 24 (see algo. D with weight decay X = 0.1, momentum coefficient
31 = 0.9, smoothing coefficient 3g = 0.95, and 10-8, The learning rate schedule consists of

---

Leaming Rate Schedule
0%
fr-upaaee size mmxx
5%
Validation Loss
3,4
Training Loss
10%
20%
3.3
2500
5000
10-3 10
0.50 Iteration (v scale)
0.05
0.50 Iteration (v scale)
0.05
Iteration
Leaming Rate
Figure Warmup sigmificantly benefits GPT2 training with Adam W. Panel 1: Trapezoidal learning rate schedules with different warmup lengths and 50% linear cooldown. Panel 2: Final validation loss for various learning rate and warmup configurations. Note the performance gap between no warmup (black) and other configurations. Panel 3: Training curves comparing the best no- warrmu to a 5% warmup with the same learning rate. The warmup run quickly surpasses the no-warmup run. Panel 4: Comparison of & update norms for these shows large initial updates without warmup.
linear warmup followed by constant phase and eventually linear cooldown spanning half of training (see examples in fig. D. This schedule keeps the peak learning rate and decay phase identical for different warmup lengths. This differr from other schedules, e.g. cosine, where the warmup length typically affects the whole shape. The learning rate value and the warmup length are swept for various configurations. Our code is based on NanoGPT [16] with additional utilities by Kosson et al. The hyperparameter values and base training configuration are adopted from NanoGPT. See for experiments additional architectures and datasets.
19
appx.
Figure shows the baseline performance for our setup. We observe that even short warmup can significantly improve performance. Not using warmup results in faster initial progress for given learning rate, but eventually falls behind leaving a permanent gap. Warmup not only stabilizes higher learning rates, but also prevents lasting degradation of the model that can not be mitigated by simply training for slightly longer. We notice that although Adam normalizes the update size, its le- fg-magnitude varies significantly throughout training with large spike at the start of training.
4 The Interaction of Momentum and the lg E2-Update Norm in AdamW
In this section we analyze the reason behind the large fz-norm of early updates in our AdamW baseline seen panel of fig. We find that this primarily stems from the 31 bias correction. We then explore to what extent these large initial updates contribute to the need for warmup by modifying the optimizer to directly control the f2-norm of the update. Although we find this is be insufficient to replace warmup on its own, these changes are an important component of our later methods.
Adam-like optimizers such as AdamW (algo. M differ from simpler methods like SGD with momen- tum in that they normalize the update size with the gradient magnitude. This makes them invariant to a rescaling of the loss function and helps counteract potential differences the gradient magnitude between layers. An important consequence of this is that the unscaled updates are not large simply due to large initial gradients, unlike plain SGD and other optimizers that don't t normalize their
Algorithm 1 Adam W (PyTorch variant, differing from the original by Loshchilov and Hutter (24))
Require:
Learning rate Mı, weight decay X, momentum 3,. magnitude smoothing 82, for numerical stability 1: Initialize: Time step +- 0, parameter vector 0,, momentum vector mo +- 0, magnitude vector vo + 0
2: while stopping criteria not met
3: t +
4:
g: +- Mini-batch gradient w.r.t. O
5:
m, < B,m,-1 1 + (1- B1 )g.
6:
$2u.-1 + (1- 82)9i
7:
m, m1 89
8:
v=1(1 (85)
9:
0, +- (1 1)1)).-11
- /prtts/
3

---

0%
AVe L. L-update size max
10
Validation Loss
3.4
5%
Training Loss
10%
20%
Ey-update size
3.3
3.1
10-3
10-2
0.50 Iteration (V scale)
0.05
0.50 Iteration (V' scalel
0.05
0.50 Iteration (v' scale)
0.05
Learning Rate
Figure 2: LionA (algo. 2) fails to significantly reduce the warmup advantage. Panel 1: Final validation loss across various learning rates and warmup percentages shows reduced but still significant no-warmup penalty compared to AdamW (fig. []. Panel 2: Training curves for 0% vs. 5% warmup at the highest stable learning rate for 0%, with warmup quickly overtaking no-warmup as before. Panel 3: LionA successfully controls the f2-upd norm. Panel 4: Early angular updates (see 5 are large without warmup and do not follow the learning rate schedule throughout training.
updates. Such un normalized optimizers might diverge to infinity if a high learning rate is combined with large initial gradients or large curvature, as the update size is unbounded. Preventing this could be an additional benefit of warmup for SGD on top of the effects discussed in this work.
Although Adam W normalizes the update size based on the gradient, its magnitude can still vary throughout training seen in fig. []] This be caused by changes in the gradient magnitude over time, especially when using different values of 31 and 33. However, it can also be caused by momentum and especially the bias correction (algo. [ line 7). The magnitude of m, depends on the alignment of subsequent gradients ....... 9r. whereas the normalization factor v. does not. For example, when each ge is an independent zero-mean random vector with fixed second moment Elg" = =², we have (see appx. B.1 for details):
E m/ (1 23") 1 - 81 0²
E(ey] (1 0-
(1)
1+31
In this case the bias correction for 31 is incorrect since it is derived for constant gradient. With the bias correction the size becomes ||| T- I+9
1+8
1-8, o², amplifying the norm of early updates by
/(1+ 81/(11111111 }. This factor is larger if the gradients between successive steps are negatively correlated, which we empirically observe happening in our setup (see $6.2).
The f2-norm of Adam updates can therefore vary significantly due to the initial bias correction, changes the aligmment of the gradients throughout training, and potential variations in the gradient norm over time. Lion [3] is a closely related optimizer that uses an element- sign operation to normalize the update, giving +1 for positive values, -1 for negative values and for zeros. Ignoring the possibility of zeros, this gives a constant update norm. Lion is closely related to Adam, and can be obtained by tracking the size of m, instead of g, in line while setting /3 0. also uses a form of scaled Nesterov momentum instead of the traditional heavy-ball variant and hyperparameter specification that differs significantly from AdamW. We propose Lion variant, LionA (algo.2) 2 that keeps the hyperparameters more compatible with those of Adam W. The learning rate is kept comparable by scaling the & update size to match that of Adam W in the -tradeet scenario, see appx. [B for the derivation of the scaling factors. Due to its ability to perfectly control the size of each update, we use Lion based methods for the remaining of this paper. This avoids confounding effects in Adam-like optimizers, such gradient magnitudes early in training, which can induce an additional warmup-like like effect.
being inaccurate from rapidly decreasing
In fig. 2 we repeat the baseline sweep using LionA. Despite perfect control of the & update norm (as seen in panel 3), the benefit of warmup remains. This leads us conclude that the &g update size is not sufficient to quantify the "effectively" large updates that we conjecture warmup mitigates. The final panel shows that the angular update size (see definition in the following section), proposed to be better measure of an effective step size by Wan et al. 38), still varies throughout training with a spike at the start of training. In the next section we explore the reasons for the large initial angular updates and how they significantly contribute to the need for warmup.

---

Algorithm 2 LionA: A modified version of the Lion [
optimizer for greater compatibility with
AdamW (algo. D. The sign operation replaces the magnitude smoothing, explicitly controlling the
g--norm of each update. Additional scaling keeps the hyperparameters 1), A comparable to Adam W.
Require: Learning rate yj. weight decay X, momentum 3. Nesterov flag r
I: Initialize: Time step <- 0. parameter vector Ov, momentum vector mho <-
2: while stopping criteria not met
3: t t-t+1
4:
g: + Mini-batch gradient w.r.t. 0
5:
mh < 8mm-1 + (1 8)ge
6:
if Nesterov flag is set
0. - (1-1111))0-11 - T/t V 32)2 + +.811110 17 sign( + (1 -)))))
7:
8:
else:
9:
O, c (1 n1)00-1 Th Sign(m))
5
The Importance and Irregularity of the Angular Update Size
The effect of weight vector w, E R used in a dot product with some vector (e.g.. in neuron):
(w,, x) = u [z| cos (<(w,, x))
(2)
can be understood terms of its magnitude | and direction /t/////zll |I. The magnitude acts like gain, scaling the outputs, whereas the direction determines which input representations æ the system responds to. The angular update size [38] of an update w. + wi+1 is defined as:
(w
(( w:) arccos
and measures how fast the direction of w, changes during training, and thus its "preference" for ac.
With BatchNorm [5] and similar operations I 14
a network can become invariant to the magnitude of weight vectors like w,, such that only the direction matters and the vector is said to be scale-invariant. Weight Normalization [30 provides a good example of this, changing the system to:
(w,// =) = || cos (L(w,, x))
(4)
Note that although the system output is invariant to the magnitude |we/], traditional optimizers are not. Scaling the value of a ocale-invariant weight vector by a factor of e > O, results in a gradient that is scaled by c- and curvature that is scaled by c (see appx. B.2). For SGD this scales the angular update by c2 and for Adam-like optimizers is scaled by e With weight decay the magnitude of scale-invariant vectors trends towards a certain stable equilibrium value over time which also results specific expected angular update size as described by Wan et al. [38 Kosson et al. 19].
This has several important implications. Changing the initialization magnitude of scale-invariant weights will scale the angular updates over time for standard optimizers, resulting in effects similar to modifying the learning rate schedule. For initial weight magnitudes that are small compared to the equilibrium magnitude, the early angular updates will large and these optimizers may benefit from learning rate warmup to counteract this. These effects also make the notion of "curvature" somewhat arbitrary as it can be scaled without changing the encoded function. Optimizers that specifically account for the weight magnitude would be invariant to these effects which may reduce the need for warmup from the traditional curvature perspective. Although standard transformers are not fully cale-invariiat, many of the angular update insights still approximately hold for un-normalized weights [19]. In light of this, we modify LionA to better control the angular update size by making the updates to weight matrices proportional to their weight magnitude, resulting in algo. B We normalize the angular update size to match the equilibrium value, replacing weight decay with projections similar to Kosson al. 19. However, unlike their RVs, we make the angular updates proportional to the learning rate schedule which we found was necessary for good performance in our case. We also do not rely on additional exponential moving averages to control the angular update size, instead utilizing the fixed update size from the LionA optimizer. This is similar to the Adam scheme used by Karras et [17] with good results for diffusion models. No additional normalization operations or scaling factors are introduced, which we still find to result in decent performance.
5

---

Algorithm 3 LionAR: A rotational version of algo. inspired by Kosson et al. [19]. The parameter vector is divided into sub-vectors 0 (0(), P)), each corresponding to either the weight vector of neuron (e.g. matrix row / convolutional filter), other parameters such as gains and biases. The updates neuronal weight vectors are scaled to be proportional to their magnitude which is kept constant through projections that replace weight decay. Additional hyperparameter adjustments are made for compatibility with AdamW. The weight decay hyperparameter remains, fulfilling its primary role scaling factor for the relative updates of neurons [19].
Require:
Learning rate 7t, weight decay A, momentum B, Nesterov flag parameter vector 0. momentum vector my 2: while stopping criteria not met
Initialize: Time step
3:
t +-t+1
4:
[gt
(1)
....... 91 (P)) Mini-batch gradient w.r.t 0:-1, divided into sub-vectors like 0
5:
for p E (1,..., p]:
6:
m. (2) +- 3m,_)
3)g(
7:
Nesterov flag set:
8:
4 8m (p)
+ (1 3)g
9:
82) + 84
# Nesterov momentim scaling factor
10:
else:
11:
(p)
m (p)
12:
# Heavy-ball momentum scaling factor
(())) E B c a neuronal weight vector:
-(2) -- ((4
max, (m A.. (100) I/VC) sigu(u, (*))
max (ir)
15:
0(2) + 0 110 1/10
# Reset the magnitude the initial value
else:
17:
0 + e0 m y sign (2!))
3.5
AVg L-update 1
2%
10
Avg Local RBC max
Valldation Loss
3.4
5%
Training Loos
10%
20%
3.3
3.2
3.1
10
10
0.05
0.50
0.05
0.50
1
0
0.05
0.50
Leaming Rate
Iteration scale)
Iteration
Iteration (V scale)
Figure 3: LionAR (algo. reduces but does not fully eliminate the benefit of warmup. Panel 1: LionAR is more stable across learning rates and shows a reduced but still significant performance gap without warmup. Panel 2: Comparing the 0% and 5% warmup for learning rate = 10 shows the warmup run overtaking early in training. Panel 3: LionAR precisely controls the angular update size throughout training. Panel 4: Despite fixed angular (and thus relative) updates in weight space, relative change of the internal representations (see $6) is large initially without warmup.
Figure 3 repeats the GPT2 training sweep with LionAR. Consistent with the findings of Kosson et al. 19 we find that controlling the angular updates stabilizes training and decreases the benefit from warmup, but does not entirely eliminate in this setting. Both the angular change and the 22-norm are simple measures of the update magnitude in parameter space that do not account for direction or other aspects of the update. the next section we show how fixed update size in parameter space can result in large changes to the internal representations of the network (a.k.a. features, activations etc), as shown in the final panel of fig.
Early Gradient Alignment Results in Large Representation Changes
Our two approaches measuring and controlling the update size in weight space failed to fully explain the need for warmup. As an alternative to the parameters, we analyze changes in

---

the internal representations activations of the neural network (feature learning). Although this harder analyze and control, it may ultimately be a better measure of the true impact of an update. A parameter update can only affect the network output, and hence the loss, by changing the representation of the network inputs at some layer. Large, sudden, changes the representations could significantly affect the non-linearities, potentially causing lasting issues such as dead ReLUS or vanishing gradients from saturated sigmoids. This could in turn explain the lasting performance degradation observed without warmup.
given parameter update will affect the representations of each distinct input sample differently. The gradients computed on these samples also generally differ, but can align to some extent. For a higher gradient alignment, the impact of a parameter update of given magnitude on the representations will be larger than otherwise. We will analyze this for a dot product making up single neuron:
= w
X (
= [(w æI)...... (w, xB)]"
(5)
[1 1I....... Fg E RC are the C-dimensional representations of random mini-batch of B imputs that is fed into the neuron, w E RC is the weight vector. and y E R is a batch of outputs. For a weight update +> w Aw,, we aim quantify the size of the output change Ay Aur X computed on the same inputs. We focus on the Relative Representation Change (RRC):
where X
A yyll
A T X| 1
(6)
x||
similar to the angular weight updates, as the sensitivity to the absolute change || | can be unclear due to normalization other scaling operations. Note that this is measure of local change, not accounting for changes in the inputs X from updates to preceding layers (global change).
6.1 Normalized Gradient Descent
We will focus our analysis the relatively tractable case of normalized gradient descent with updates:
B
9
Aw 7/
g
9- B
9b
b=1
where 9b is the gradient of some loss w.r.t. w for the b-th element of the mini-batch. We will use the following definitions, properties, lemmas and assumptions for this system (see appx. B.4 for details):
DI: We define =: ÿ + 9b where g = [E[g] and ÿb is the difference with E E(gn] = 0.
D2: We define > = || / [B [| as the Signal-to-Noise Ratio (SNR) of the gradient.
Pl: For a neuron, 9b 26. and hence ær, - sign( (ars, gb)) (1 /// ( (g gv).
Ll: Consider two independent random vectors E _ and b RC whose elements are independent and identically distributed (IID). If at least one of the vectors has a zero-mean distribution, then the expected value of the squared inner product of and b is given by E| (a, b) = E[|l a || |2]//.
A1: We assume the following vector pairs satisfy LI: (x,, @b) when + b, (g.gb) and (w, 2z).
This allows us to compute the expected square relative representation change (detailed in appx. B.4
E vb)
1²²C
B
Ellgg [13] +
Ellla
E
B2
1
12
C
(B 1122
||
ll
+ E E1
||ل||*
+ 2((-1))
(8)
C
C²C
1
(p+1) B-1 +
(B-1)&g
0 b) + +2(B-1)0)
B
w 2
p+1
The expected relative change in the output for given sample can be broken down into three sources, the contribution of the sample itself (first term), random interference from the "noise" ği of other samples (second term), and finally amplification of the common mean component (third term).
7

---

LR scaling for const RRC
100
Signal to Noise Ratio
100
0%
Alignment cas L(g..
Grad Cancellation lation
10-1
0.5
0.5
007
10
0.0
Velacity Soadctt--Grad
0.0
000
8-10º
B=10²
10
0.5
0.5
B-10
10
-1.0
10-4
10a
0.50 Iteration (v/ scale)
0.05
0
0.05
0.50
0.50 Iteration (V' scale)
0.05
SNR
Iteration (V scale)
Figure 4: Equation (9) predicts that the learning rate needs to downscaled for higher signal to noise ratios (f)) keep the relative representation change constant. Larger batch sizes are affected more, with scaling becoming significant when yp > B Panel 2: Mcasurements of the SNR for the two highlighted runs in fig. Note the SNR starts very high but is also remains large in comparison to our B = 480 for almost all of training. Panel 3: The gradient is strongly oppositely aligned with the momentum vector for most of training (shown for example layer). Panel 4: Projecting the momentum component of the updates onto the gradient component shows that this results in the momentum vector "cancelling" roughly half the gradient on average.
The RRC expression provides many interesting insights. In the case of large input dimension C oo and small SNR &p R 0, keeping the RRC constant for different batch sizes involves scaling the learning rate 1/ x vB, as suggested by Malladi et al. [27] for Adam. When the SNR is some finite and value and C is still large, this scaling rule instead starts break down around B 1/p. matching the predicted critical batch size of e.g. McCandlish et al. (28). The role of the dimension C' in the expression curious, suggesting that narrower layers experience larger changes due to random inference from other samples in a given batch. The C in the leading factor also suggests that the angular updates can be smaller for a larger input dimension, similar to what is proposed in parameteriiztion [4 45 Most importantly, this expression shows that if the SNR changes throughout training the learning rate needs to be adjusted to keep the RRC constant. In particular, with large batch sizes, a high initial SNR results in large representation changes which warmup can help prevent. The first panel of fig. 4 shows how eq. (9) predicts we should downscale the learning rate for different batch sizes and SNRs, assuming we originally scaled the learning rate 17 x / B and that C' is large. The second panel confirms that the SNR indeed starts out large, suggesting lower initial learning rates are needed, i.c. warmup.
In the first panel of fig. 5 we show the results of adding a term that scales the update size as predicted by eq. (9). This acts similar to an automatic warmup based on online measurements of the SNR which we obtain from the gradient accumulation of shicro-bacches. Although this helps close the gap between warmup and no-warmup. the overall performance slightly worse. One potential issue is that our batch size of 480 is quite large compared to the measured SNR, exceeding the critical batch size estimation throughout most of training. This results in a scaling of the step size throughout training, which distorts the decay phase. It also requires large learning rate values to counteract the scaling, which may destabilize the training of non-matrix weights like gains. We increase the weight decay by factor of 32x to try to increase the angular updates relative to gains in order to compensate, but this value was not tuned and is unlikely to be optimal. We believe approaches that aim to directly control the RRC are promising direction but require further work to be practical.
6.2 The Role of Momentum
Momentum is believed to be key enabler of optimization with larger batch sizes [32 39 31). How- ever, is unclear how it should change predictions for the critical batch size or relative representation change. Momentum spreads the application of gradient out over multiple steps which tends to make each update smaller, especially for random walk, which reflected in our scaling coefficients in algo. 2 and 3 The smaller updates are counteracted by an increased correlation in their direction, which can result in similar "long term" changes from each gradient sample, especially for simpler methods like SGD that don't 't normalize the step size. In the last two panels of fig. we observe that our setup the gradient and momentum are negatively correlated, counteracting each other. We find
momentum crucial for performant training, panel 2 of fig. shows significant degradation without it.
8

---

LionAR $=0.9 comp
LionAR $=0
LionAR $=0.98
LianA p
3.5
3.5
3,5
0%
10%
+ 0%
0%
Vallddtton Loss
3.4
5%
10%
Vallddatio Loss
4,4
5%
20%
Valldstioo Loss
3.4
2%
Validstioo Loss
3,4
2%
5%
10%
10%
3.3
4.2
3.3
20%
3.3
20%
3.2
3.2
3.2
4,0
3.1
111
3.1
3.1
10-3
10-2
10
10
10
10-2
10-3
10-2
Learning Rate
Learning Rate
Learning Rate
Learning Rate
Figure 5: Panel 1: LionAR with a correction factor for the RRC based on eq. 19.99 does not benefit from warmup. Panel 2: LionAR training without momentum results in drastically lower performance. Panel 3: In LionAR with increased momentum = 0.98, Nesterov momentum and an inverse bias correction for early momentum, warmup performs best. Panel 4: The same does not apply to LionA, suggesting that these changes are not sufficient without controlling the angular updates.
We believe the smaller update sizes for momentum combined with the potential for later gradients to counteract earlier gradients during their application over time, can help stabilize training. An otherwise large relative representation change is spread out over multiple steps and counteracted by later gradients. Higher values of momentum should amplify these effects. Looking at the total contribution of each gradient also implies that with momentum early updates should be smaller when measured in parameter space, otherwise the relative representation change for those samples is too large. This is equivalent to entirely removing the 31 bias correction in AdamW (algo. [ line 7), or introducing an inverse bias correction in Lion like algorithms (see appx. B.] for details). Higher 3 values should help amplify the stabilization effects of momentum. In fig. 5 we find that at higher momentum values LionAR no longer benefits from warmup unlike LionA which still needs it. These experiments use Nesterov momentum and the additional inverse bias correction, though these adjustments offer only minor improvements compared to higher momentum.
7 The Detrimental Effects of Large Updates
In appx. A we empirically investigate potential causes for the lasting performance degradation from large initial updates for a small ResNet model. We find that the performance impact best correlates with the number of dead ReLUs and is improved by the use of leaky-ReeUU which fits well with our perspective of large changes in the internal representations. We also investigated whether overfitting to initial training samples or the correlation between weight vectors of different neurons could explain the necessity for warmup, but these factors did not show significant impact.
8 The Role of Larger Batch Sizes
Warmup is often used with larger batch sizes in particular, for example in the setting where Goyal et al. (10 first proposed using linear warmup. Although this was for SGD, expect the need for warmup to be amplified at larger batch sizes for two of the reasons we identified. The first is that larger batch sizes are more likely exceed the critical batch size early in training. The second the size of early angular updates. As shown by Kosson et al. [19], the equilibrium weight magnitude depends on the learning rate and weight decay value. Common hyperparameter scaling rules for modified batch size only change the learning rate but not the weight decay, which shifts the equilibrium magnitude. The smaller the initialization magnitude is compared to the equilibrium magnitude, the larger the early angular updates will be relative to their steady state value, potentially necessitating warmup.
9
Limitations
Our main experiments focus on single network which may not broad enough to generalize to a wide range of networks. In appx. C.2 we experiments with an additional dataset and architecture but the scale of the experiments is still limited and they cover a limited range of hyperparameters. We believe we identify real factors that contribute to the need for warmup. but these may not be the only

---

ones across a broader range of settings. Similarly, the promising results for reducing or eliminating the warmup with higher momentum values or the relative representation correction would benefit from broader validation.
10 Conclusion
In this work we explored how the size of the updates Aw = ryu impacts the need for learning rate warmup. We showed that can be large initially when measured in terms of its 22-norm ($4). the resulting directional change in w (angular update, $5), as well as the resulting change in the internal representations of the network (relative representation change, $6). We argued that small initial values of the learning rate are beneficial to counteract large values of u, i.e. that a learning rate warmup simply keeps some notion of the overall "effective" update size reasonable. We showed this experimentally rather than theoretically by modifying the optimizers to normalize the size of u based on each metric and measuring how these changes affected the benefit of using learning rate warmup.
The two weight- measures of the update size, the f--norm and angular update did not fully account for the need for warmup. However, quantifying the update size in terms of the relative change in neural representations shows potential. This measure is closely linked to the angular update size but accounts for changes in the signal characteristics of the gradient, which throughout training. Effectively controlling neural representation changes is a challenging task we leave for future work, but our initial attempts show encouraging results in reducing the need for manually configured warmup. We also highlighted the importance of high momentum for warmup; when combined with angular update control and an inverse bias correction, it may enable efficient warmup--ree training. Overall, our work provides new insights into the benefit of learning rate warmup with modern optimizers beyond SGD and suggests potential directions for eliminating it.
vary significantly
Although we present new methods we consider promising, we still recommend the use of a short warmup in practice. Fully climinating it seems to require significant modifications that also need further validation across additional settings. However, we hope to have provided the reader with new perspective and simple intuition for why warmup is beneficial for training. We also hope our work inspires further exploration of how learning should be controlled and scheduled in neural network training. In particular, it seems that the learning rate in current optimizers does not really control the "rate of learning", making learning rate schedules and the use of warmup highly arbitrary.
References
[1] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint 1:X7v:6007004433 2016. 5
[2] Jeremy Bernstein, Arash Vahdat, Yisong Yue, and Ming- Yu Liu. On the distance between two neural networks and the stability of learning. Advances in Neural Information Processing Systems, 33:21370-21388 2020. arXiv:2002.03432 2
[3] Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, and Quoe V Le. Symbolic discovery of opti- mization algorithms. In Thirty-sevnntth Conference on Neural Information Processing Systems, 2023. URL htt.httts///ppereeiiee
arXiv:2302.06675
[4] Vitaliy Chiley, Ilya Sharapov, Atli Kosson, Urs Koster, Ryan Reece, Sofia Samaniego de la Fuente, Vishal Subbiah, and Michael James. Online normalization for training neural networks. Advances in Neural Information Processing Systems, 32, 2019. arXiv:
[5] Katie E Everett, Lechao Xiao, Mitchell Wortsman, Alexander A Alemi, Roman Novak, Peter J Liu. Izzeddin Gur, Jascha Sohl-Dickstein. Leslie Pack Kaelbling, Jaehoon Lee, and Jeffrey Pennington. Scaling exponents across parameterizations and optimizers. In Forty-first Inter- national Conference on Machine Learning, 2024. httlhhttss:///peeeeeeet wtwpeeeeeiiieeeeeeeeeeee nwt/
id=OksNeD1SSTT AiXiv:2407 200.02777 18
[6] Jingwen Fu, Bohan Wang, Huishuai Zhang, Zhizheng Zhang, Wei Chen, and Nanning Zheng. When and why momentum accelerates sgd: An empirical study. arXiv preprint 23Xiv:230600000 2023. 2
10

---

[7] 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 Learn- ing Representations, 2022. URL. http:////.paaarreeeeeeee///oaooa
arXiv:2110.04360
[8] Aaron Gokaslan and Vanya Cohen. Openwehtext corpus. hutp::/88kyiion077.gttth io/ OpenWebTextCorpus 2019,
22
[9] 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. httpp:ttts////peeeeee wwtp:/iwwwwweeettniiie.
id-r14EDsCqKX arXiv:1810.13243
[10] Priya Goyal, Piotr Dollár, Ross Girshick, Picter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangging Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in hour. arXiv preprint prrpiintaaxXiv: / 7706.02677 2017.
[11] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016. arXiv 112.:351
[12] Jordan Hoffmann, Scbastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556 2022. URL
https://arxii:
org/ abs/2203, 15556 2
[13] Xiao Shi Huang, Felipe Perez, Jimmy Ba, and Maksims Volkovs. Improving transformer optimization through better initialization. In International Conference on Machine Learn-
ing, pages 4475-4483, PMLR, 2020. URL httppa:/prroceddngs.nt mlr. www.mm....rr//v119... huang20f html
[14] Xun Huang and Serge Belongie. Arbitrary style transfer in real-time with adaptive instance normalization. In Proceedings of the IEEE international conference on computer vision, pages 1501-1510, 2017. arXiv:1703.0666
[15] Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pages 448-456. pmlr, 2015. MEXiv:1502.0311
5
[16] Andrej Karpathy. nanogpt. https: http://hhhh..ccmaaaraatttaaaa//a 2023.
[17] Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyzing and improving the training dynamics of diffusion models. In Proceedings of the IEEÉ/CVF Conference on Computer Vision and Pattern Recognition, pages 24174-24884, 2024.
arXiv:2312.02696 5
[18] Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), San Diega, CA. USA, 2015. arxiv: 141.6694226980
[19]
Atli Kosson, Bettina Messmer, and Martin Jaggi. Rotational equilibrium: How weight decay bal- ances learning across neural networks. Inemattto Conference on Machine Learn-
International
ing. 2024. URL http://p/ppaaeeeee.eeeaa/rrrrraa
arrp: a1Xiv:2305.17212 2
14 16
[20] Alex Krizhevsky. Learning multiple layers of features from tiny images. self-published, 2009.
URL httts:tp::: ptttps://wwrrr //www//wwwwccco/./ cs. toronto............... wewpontto..duu//e wect//krriz/eeannnngg.. featurrss-e 2009- TR. TR.pdf I4
[21] Zhiyuan Li, Kaifeng Lyu, and Sanjeev Arora. Reconciling modern deep learning with traditional
optimization analyses: The intrinsic learning rate, Advances Neural Information Processing
Systems, 33:14544-14553 2020. arXiv:2010.02916 2
11

---

[22] Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling SGD with stochastic differential equations (SDEs). In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wort- man Vaughan, editors, Advances in Neural Information Processing Systems, 2021. URL
nvQQ 1Xiv:2102.12470 2
[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. URL [tRppptts///ppeeeeiiee wwwpeeeeeeeeeeeeeeeeeeeeeee net/forum?
id=rkg22aKKDr 1XXv:1908.03266
[24] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019. URL httpp:://pperreiee wttp:/rreiiiieeeeefforrur net/forum?
id=Bkg6RiCqy7 arXiv:1711.05101
12
21
[25] Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normalization layers: Sharpness reduction. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL
https://openreview..n http://iiww...tetooooooooed/droo
ar arXiv:2206.0708
16
[26] Jerry Ma and Denis Yarats. On the adequacy of untuned warmup for adaptive optimization. In Proceedings of the AAA/ Conference on Artificial Intelligence, volume 35, pages 8828-8836,
2021. 1910:v:110..0422..
2 21
[27] Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the SDEs and scaling rules for adaptive gradient algorithms. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL httpltttts///oorreeete
of large-batch training. arXiv preprint ntt/foeumniddddeeeooooo.
arXiv:2205.1022 10287
[28] Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team. An empirical model
162, 2018. 2
[29] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei. and Ilya
Sutskever. Language models are unsupervised multitask learners. self-published.
2019. URL https:://4mmcfpkksyyat wwwudfront. www./eeeeeeerree htt.:aaw...eeeett...aaaaaaaogeoa language-models/
Language_modedt _rrellsssrr__e_e www.auuusssrrrsssesltlaaaaaaaaaa Learteesseeerrrrss.c pdf
[30] Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. Advances in neural prooessni
processing systems,
29, 2016. aarXiv:160.0000
5
[31] Christopher J Shallue, Jachoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E Dahl. Measuring the effects of data parallelism on neural network training. Journal of Machine Learning Research, 20(1 12): 2019. arXiv: w0w.:8811.3666 1.03600)
[32] Samuel Smith, Erich Elsen, and Soham De. On the generalization benefit of noise in stochastic gradient descent. In International Conference on Machine Learning, pages 9058-9067. PMLR, 2020. 1rXiv:2006.15081
2 g
[33] Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob
Steeves,
Joel Hestness, and Nolan Dey. Slimpajama: A 627b token cleaned and deduplicated version of redpajama. s1 ssimpajamas
ai/blog/ aappaajjaaaa-------toooe- token clanned-eeeddddde www.eeewwww.cceeeeeeeeeiioioooie redpa 2023. 21,221222
[34] Sebastian Stich, Amirkeivan Mohtashami, and Martin Jaggi. Critical parameters for scalable distributed learning with large batches and asynchronous updates. In International Conference on Artificial Intelligence and Statistics, pages
PMLR, 2021. arXiv:2103.02351
[35] Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initial- ization and momentum in deep learning. In International conference machine learning, pages 1 1139-1144 1147. PMLR, 2013. URL https::: htttps://pppt /ttpa/////oocceeeesee.oe mlr. www.ssreessvee aress//28/sutskeverre sutskever13
html
2
12

---

[36] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov. Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint urXiv:2307.09288, 2023. 22
[37] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. arXiv:1706.03700
[38] Ruosi Wan, Zhanxing Zhu, Xiangyu Zhang, and Jian Sun. Spherical motion dynam- ies: Learning dynamics of normalized neural network using sgd and weight decay. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 6380-6391. Cur- ran Associates, Inc., 2021, Jtt.htttss:/ppooeeeeee 22586550000600666666beee Paper pdf arXiv: 2006.02006.08419
nwwrips ccc htt.:c//paaaa pttpp/pappp/////eeeeeee
[39]
Runzhe Wang, Sadhika Malladi, Tianhao Wang, Kaifeng Lyu, and Zhiyuan Li. The marginal value of momentum for small learning rate SGD. The Twelfih International Conference on Learning Representations, 2024. URL
arXiv:2307.5119
[40] Ross Wightman. Pytorch image models.
pytorch-image image-models 2019. 14
[41] Mitchell Wortsman, Peter J Liu, Lechao Xiao, Katie Everett, Alex Alemi, Ben Adlam, John D Co-Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, et al. Small-scale proxies for large-scale transformer training instabilities. arXiv preprint arXiv:2309. 14322, 2023. URL
https://arrxv.orr//bb//000: 14322
[42] Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pages 19, 2018. arXiv:1803.08449 5
[43] 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 International Conference on Machine Learning, pages 10524-10533, PMLR, 2020. 2rX1v:2002.04745 2
[44] 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. URL https //openreview http://forummiddxxxxxxeeeeee
arXiv:2203.03466 g
[45] Greg Yang. James B Simon, and Jeremy Bernstein. A spectral condition for feature learning. arXiv preprint arxiv:23100178813 2023. 8
[46] Dong Yin, Ashwin Pananjady, Max Lam, Dimitris Papailiopoulos, Kannan Ramchandran, and Peter Bartlett Gradient diversity: a key ingredient for scalable distributed learning. In International Conference on Artificial Intelligence and Statistics, pages 1998-2007. PMLR,
2018. arx:: 1706.05699
2
[47] Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint arXiv: 03888, 2017. 2
[48] Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. In International Conference on Learning Representations. 2020. httL/htts////peereee http://ppppeeeeeeeeeeeeeetooooot wett:/forrmmdd-Sy///mmmt-t
arXiv: 104.:1990.009622 2
[49] Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger B Grosse. Which algorithmic choices matter which batch sizes? Insights from noisy quadratic model. Advances in Neural Information Processing Systems, 32, 2019. arXiv: 190..04141
13

---

Epach Sweep
92
Test Accuracy
Standard
90
Standard
128 Standard
88
1 Leaky ReLU (0.1)
8 Leaky ReLU (0.1)
86
128 Leaky ReLU (0.1)
100
200
300
400
500
600
700
800
Epochs
Figure 6: The performance gap caused by large initial updates persists despite extended training (800 epochs) in standard ResNet-20. We investigate the influence of network non-linearities by comparing two training methods while scaling update sizes during the 5 epoch initial phase by factors of 1, 8, and 128: Standard (S), which employs traditional ReLU activations, and Leaky ReLUU, which replaces ReL.Us with Leaky ReL.Us using a scaling factor of a = 0.1. We observe that training with Leaky ReLU results in smaller performance degradation from large initial updates, suggesting that the non-linearities in the network might substantially impact the observed performance degradation.
A The Detrimental Effects of Large Updates
To investigate the effects of large updates at the beginning of training, we conducted controlled experiments ResNet-20 model CIFAR-10 [20] due to resource constraints. We controlled the average angular update throughout training using the rotational optimizer variant of Adam W proposed by Kosson al. [19]. For the initial phase of training. epochs, either standard learning rate of 0.05 or amplify it by factor of 8 or 128. This results in very large updates, scaling the rotation by either v8 or /128. For all experiments, we used weight decay of 0.01, 31 = 0.9, A2 = 0.999, 5 epoch initial phase, and trained for 205 epochs in total with cosine schedule unless otherwise specified. The data was pre-processed by normalizing it with mean of (0.4914, 0.4822, 0.4465) and standard deviation of (0. 202202,, 0.1994, 0.2010) and applying simple data augmentation techniques as described by He et al. 1 To run the experiment, we used the codebase from Wightman 40 and extended the utilities from Kosson al.
As shown in fig. the performance standard training does not recover when large updates are used at the beginning of training, even when the training time is extended to four times the normal duration for ReI.U networks. This suggest that large, controlled, initial updates can result in a permanent performance degradation, similar what we observe without warmup in advanced settings. The impact is much smaller when replacing ReLUs with leaky ReLUs, suggesting that the non-linearities in the network might substantially contribute the performance degradation.
In fig. we measure the fraction of dead ReLUs directly across different settings and scaling factors. We find that large initial updates do indeed lead large number of permanently dead units and that the final accuracy suffers when this is the case. This effect can be mitigated by freezing the biases at the beginning of training, shown in the table in fig. 7 We also observe that replacing the actual gradients with random gradients has a much smaller impact, suggesting that the direction of the updates also matters for the degradation, not only their size.
Interestingly, we did not find a connection to overfitting to a small number samples the beginning of training. The performance of 92.1 can be recovered in this case. Additionally, we explored stable rank measurements as potential factor but did not find significant relation, as detailed fig. g
B Additional Mathematical and Technical Details
B.1
The magnitude of the Momentum Vector
Let's assume scalar gradient gt (e.g. for some coordinate) that a random variable that is indepen- dent across time and has zero mean distribution that does not change across time, i.c. E[go] 0
14

---

inactive activations
Method
Scule
Scale 8
Scale 128
0.4
Standard
92.13:0.2
89.48:1.22
88
Standard frozen hias
92.30+0.3
92.08+
92.30++022
0.0
Random
92.05+0.3
91.74+00.2
89.54-+0.3
Random frozen bias
92.27:0.2
92.12:044
92.20-0.1
steps
Leaky ReLU
92.1610.3
91.48:0.3
91.82404
S-fb
R-fb
R1
51
Leuky ReLU frozen bias
92,46+0,2
92.49::0.1
92.35=0.2
5-tb 128
R-fb 128
128
5 528
Figure Comparison the performance (final test accuracy) and fraction of dead ReLUs (inactive activations) across different settings. The learning rate in the initial phase of 5 epochs is scaled by factor of either or 128. Standard (S) denotes normal training, while Frozen Biases (fb) involves freezing the biases the onset of training. The Random (R) approach employs random gradient directions at the start of training, and Leaky ReLU replaces the ReLUs in standard models with Leaky ReLUs using scaling factor of 0.1. We observe notable correspondence between large initial updates, higher ratios of dead ReI.Us ResNet-20, and performance degradation.
0.06
0.04
0.02
200
400
600
800
1000
steps
(fb)
(fb)
(fb)
(fb)
(fb)
Figure 8: Impact of varying update sizes during the warmup phase on the stable rank of a standard ResNet-20. The learning rate in the initial phase of 5 epochs is scaled by factor of either or 128. We evaluate the effects across different training configurations: Standard (S) denotes normal training: Frozen Biases (fb) involves freezing the biases at the onset of training; and Random (R) employs random gradient directions at the start of training. The stable rank remains largely consistent across these methods, except when using extremely large updates specifically, scaling the learning rate by a factor of 128 without freezing biases-whie whhch leads to noticeable variations in the rank.
and E[gg] = o2. For standard heavyball-momentue my, with my and coefficient (equivalent to 31 for Adam) we have:
E[m?) = [[(mm 1 + (1 3)
(10)
27
(1 3)
B gt-
(11)
(1 8)2
325 2
9t
(1- 3)
82t-j g1 19g
(12)
j=0
1-1 Σ
1-1 t-1 Σ
-1 (1- 2)
E 2
+ (1- 2
320 1 3 E.
11
(13)
1=0
1-1
(1 3)2
32i 0 +0
(14)
1=0
15

---

(1 3)2 1 g
332 0²
(15)
1 g
(1 3)2 (1 3)(1 + B) o²
(16)
1 - 3
(1 8
(17)
8
In the limit t + 0o we have (1 - B2t)) - 1. We can derive the size of the second-moment in AdamW in an analogous way, obtaining E[[s]] (1 35)0². For a random walk, the update size of
Adam is scaled in similar way. Since the update size of Lion is fixed and does not depend on B, we
scale the update size to match that of AdamW for random walk in a steady state, i.e. by v T+: as seen in algo. 2
1-8
For Nesteroy momentum, the update is modified to use:
ut Bm, + - -8)gt
(18)
= 3 (8m--1 + (1 (1- 3)9() + (1-3)gr (1-
(19)
= mm- + - (1 + 3)9
(20)
Note that m and Ir are independent and zero-mean, allowing to use the previous result for:
Elu,] E (33m + (1 -1- ( (1++)))++ 8)9:)
(21)
33 + -1- (2)
(22)
-1- 1 1 1
(1 (1-33)) 82)80²
(23)
1
In the limit ++ oo this gives the Nesterov scaling factor used in LionA (algo. ensure that the update size corresponds to that of Adam using an analogous Nesterov update.
Inverse bias correction for momentum. Adam uses bias correction attempt fix the update size over time. This scales early updates resulting in the contributions of the corresponding gradients being amplified. The relative representation change for those samples is increased to applying the same update multiple times. Removing the 31 bias correction from AdamW removes this effect. LionA and LionAR similarly scale the update size, making it constant. We can counteract this by changing our scaling factors to use the time varying expressions based on the derivations above. Note however, that this assumed the gradients were uncorrelated it only approximately undoes the scaling effect for real values with arbitrary alignment of successive gradients. To summarize, the inverse bias correction for momentum changes the momentum scaling factors (y in algo. 3) to vary over time:
result, similar
3
Nesterov:
7t Va- -1- 3232 + (1 32t
2).4 1211--2-110000 -21-2)3
(24)
1 3
1- 3
Heavy-ball:
7 V (1- g2r)
(25)
1+B
B.2 Properties of Scale Invariance
Derivations for the gradient magnitude and curvature can be found in existing works, for example Lyu et al. [25]. When a scale invariant weight is scaled by a factor e >> 0, the gradient is scaled by c-1 which scales the ratio of the gradient norm and weight norm, and therefore the angular updates, by C For normalized optimizers like Adam and Lion, where the update norm is not affected by the gradient magnitude, this factor is decreased
B.3 The Angular Update Size in LionAR
The scaling factor for the angular update size in algo. 3 adopted directly from the AdamW value derived by Kosson et al. [1]] Since the Nesterov momentum does not change the total contribution of each gradient it does not affect the equilibrium magnitude. The expected angular updates are therefore scaled in the same way the RMS update norm we derived in appx. B.I
16

---

B.4
Relative Representation Change for Normalized Gradient Descent
Property (PI):
For dot product y (w, 2) and loss L(xs) that depends on y, we have:
(x)) (.((2)) dy (.)())))
2b
(26)
Sw
Oy
Sw Oy
where (.f) by
is scalar, ensuring that gh Dw
(ms)
||
Wh. assuming the vectors not zero,
Lemma (L1): Consider two independent random vectors a E RC are independent and identically distributed (IID). If at least one of the vectors has zero-mean distribution, then the expected value of the squared inner product of and is given by:
R whose elements
Ka, b) = El ||||||/E||10|
(27)
Proof: Let a (a1, ....... ac) and = (b), .... bc). The inner product (a, b) is given by:
(a, b) Σ a,by
We need to find (a, 2. Expanding the square of the inner product:
c
C
b). =
Σ
a;bi
abb
1=1
==
Taking the expectation, we get:
C c
C c
(a, b)3 = E
agbiazb;
= Σ
Bla,b,a,b,l
Since and are independent and their elements are IID, we have: E[a,b,ay6,] Ela.a,]B(.bbj)
Consider two cases:
1. When j:
Ela,baa, = [bb
2. When i + j:
Ela,blayyy] = E E,,
Given that at least one of a or has a zero-mean distribution, say a without loss of generality, we have Ela.] = 0. Thus:
Blabbaabby]= =
So, the expectation simplifies to:
E (a, b) Σ B|a
Since a, and b, are IID, have:
B|a = E[aa]] and Ebba1 = E(b?].
Therefore:
Bl(a,bb
CE
17

---

Recognizing that:
C
B||aa| = E
2
= CB(aj)
E
6 1
E
CB
we have:
Bllaa ||
E|a
and
E[
C
C
Thus:
E |[[|a|
BBI| ||
(a, b)
C
C
This completes the proof.
Assumption (A1): We assume the following vector pairs satisfy LI: (æt, when i o b, (g. 9b)
and (w, 2b).
Vector pairs of the type (x,, and (g.gv) should be independent and gb has zero mean dis- tribution. However, the elements of each vector are not necessarily IID. For (w, 2z). this is an I, is guaranteed to be IID or zero mean, and
stronger assumption. Generally, neither their independence later in training does not necessarily hold. Applying weight standardization to batch normalization can be viewed as a simplifying approximation to obtain reasonable predictions without additional information about these vectors. Everett et al. (5] explore the behavior of (w, 2s) throughout training and find that it is more complicated than assumed here. This will lead to additional factors that may affect the RRC but do not attempt analyze.
would suffice make their mean zero. Overall, this assumption
Deriving the Relative Representation Change: Applying I.1 directly gives us the original expected square output
|w
"E 0b 12
Elyil [w an)
(28)
For the expected square representation change we get:
E ( (Ay))²
(29)
[E[(-
g [2] x.))
(30)
27
1/2
E Elggit
gi, .b6
(31)
27
172
IE
sign (zo,, (b)) (6)))lg ||000|1 + Σ (gr, 206)
(32)
B
EL
[2]
i#
n72
E
sign((æs, (b)) |1 + (B - x1) + (g x6)
(33)
1
17
(34)
where we have used the definitions from eq. 7 and D1. Using property P1, write:
[
xs) Ø, sign( (xb, gb))
(g + gv)
(35)
[1gb]
x
sign( (xi,, 9b))
(lgg| 2 + (g. gb))
(36)
g
18

---

Plugging this into the previous expression yields B[(A8))
2
1/² _ E 1
1
[ (1l@) (0.00))) +
E
sign(( gbl) ||00|||<b||| + (B 1)
Σ (gi abb) i
1
(37)
Squaring the expression results various cross but all remaining dot products except the sign one are zero in expectation (due the noise 9) and independent from each other. The cross terms involving these thus all disappear under the expectation. We apply Lemma LI to their squares and approximate the expected norms of and Gb being independent. This gives ((()))
2 E a [?]
(B
1g
IE @6
[|g]]
(38)
B g 12]
[E 31
1) + B
+
[1
(39)
We can compute the expected magnitude of the batch gradient as:
B
B 9.)1' = B Σ
| [2] = [1]
Σ
g
9.)1
|g|l
(40)
1=1
1=1
and similarly
B[[lg<11]] | 2
[1
11. Using these facts we can further write (B((Aya)))
1/2
[[[[=r+|1]]
(B 1)2 |l
[B
[1
[|و] +1
| @b | ||g||2
B²
+
E 111g g ||
C
=11@@_1911 1) + C
(41)
Combining this with the previous expression for E[ys] and the definition (D2) the signal-to-noise we obtain the expression in the main body:
ratio yp : |
[ yb.) )1
C²2C
+
E
E [
B
w
1
Elgg
C
(B-
1)2
|
E 9bll
+
||g|| +
+ +2(B-1)lglP
(42)
B11lg8121
C
1²C
1
B-1
4 ((B-1)²
) + +2(B-1)))
(++1)
+
Ba
w
2
I B
C
x+1
(43)
B.5 Estimating the Signal-to-Noise Ratio
We use accumulation over the microbatches to estimate the SNR at given time. Let's assume we have A microbatches of size M each, with the average gradient of microbatch denoted gm and the average gradient of the whole batch denoted g = Em Im-
We estimate the variance the norm single gradient example, i.e. the noise power as:
PN
A
M I
1
1
Σ 9hh
g²
(44)
A
A
The signal power is estimated as:
P3 1T
PN
(45)
AM
Our SNR estimate is then:
6 Ps/PN
(46)
19

---

B.6
RRC Correction Factor
The RRC correction done based on eq. ( and the SNR estimation eq. (46 We assume the learning rate was originally scaled with the square root of the batch size. which is derived for an SNR of zero, and downscale the step size to compensate for the measured SNR and batch size. We define:
1 B(1+ + ())
(
B-1 (BB-112
(p+1 ( ( - +2(BB1)4)
(47)
For numerical purposes, clamp 1 S p S B which corresponds to 4p and 4 for large C- co. The update scaling factor is the square root of an EMA the inverse this quantity. We use the same coefficient for the momentum and compute this for the matrix of each linear layer independently. This form for the scaling factor is somewhat arbitrary, complicated by the fact that Lion-like algorithms fix the step size exactly, scaling the gradient at each step size can change the magnitude of the update. For Adam or SGD like algorithms we could scale the gradient contributions directly instead of scaling the update size.
B.7 Run-to-run Variance / Uncertainty Estimation
We do not quantify the uncertainty for every GPT2 configuration in our sweeps. This would require significantly more compute and our estimates of the uncertainty for select points indicate that this would not qualitatively change our results. For the baseline Adam W run the run-to-run differences in the validation loss over different seeds around 0.05. However, the relative ranking of different runs remained the same.
B.8 Computational Requirements
Our experiments are performed A A100 GPUs with either 40GB or 80GB of RAM. One training for our GPT2 setup takes around 4h, running on a single GPU. Reproducing the GPT2 experiments reported in the main body should take the order of 1000 GPU hours. Including our preliminary experiments brings this up to around 3x this amount.
20

---

C Additional Experiments
C.1
Comparison with RAdam
RAdamW combines the variance reduction technique of Liu al. (3 with the decoupled weight decay of Loshchilov and Hutter [24]. Since it is a well known technique for reducing the need for warmup it serves while RAdam W outperforms the Adam W baseline without warmup, it is unable to match longer warmups. Ma and Yarats (26) suggests that RAdamW is approximately equivalent to 4 steps of SGDM followed by Adam W with a special type of built-in warmup with an effective length of around 2/(1 - 32) 40, which is likely too short in our setting. For comparison the 2% warmup shown corresponds to 100 steps but is too short for optimal performance.
good comparison and contextualization of our work. Figure shows that
The analysis of Liu et al. [23] is based on the idea that early in training the second-moment estimates are not accurate (noisy) and can therefore not be trusted to scale the update properly. This could in turn contribute to the need for warmup, although Ma and Yarats (26) question this interpretation. We first note that without momentum, perfect estimates of the second moment at the current time step would control the expected 2--nrrm of the update. This relates our approach of looking the update size to the adaptive learning rate view Liu et al. [23 Secondly, we note that counteracting noisy estimates of the second moment not be the sole reason warmup is beneficial. This is supported by the fact that both SGD and Lion empirically need warmup in various settings but do not use the second moment all, indicating there additional factors that contribute to the need for warmup.
3.6
3.5
Learning Rate Schedule
AdamW RAdamW
Learning Rate Schedule
0%
AdamW
2%
5%
3.5
5%
Ainal Validation Loss
3.4
RAdamW
10%
Validdtot
3.4
10%
20%
3.3
3.3
3.2
3.2
3.1
3.1
2500
5000
10-3
10-2
25.00
5000
10-1
10-7
Learning Rate
Learning Rate
Figure 9: Comparison of the AdamW baseline with a cosine schedule and RAdamW. Panel 1: Cosine schedules with different warmup lengths. Note how the warmup shifts the curve, affecting the whole schedule including the decay portion. Panel 2: The final validation loss for GPT2- 124M training on Open WebText using the cosine schedules. Note that RAdamm helps but does not eliminate the need for warmup. Panel 3: The original trapezoidal schedules used in our experiments. Panel 4: Trapezoidal GPT2- 124M Open WebText results. The RAdamW results are similar those in panel 2.
C.2 Model & Dataset Ablations
In this section repeat some of our main experiments, varying the dataset and model architecture.
LianA
LionAR
LionAR 8-0.98 +N +MC
LionAR RAC camp
3.3
3.3
3.3
3.3
Final Validation Loss
0%
5%
Final Validation Loss
0%
5%
Fnnll validation Loss
0%
5%
Final Validation Los
+ 0%
5%
3.2
10%
3.2
10%
3.2
10%
3.2
3.1
3.1
3.1
3.1
10-3
10
10-3
10-2
10-3
10
10-3
10-2
Learning Rate
Learning Rate
Learning Rate
Learning Rate
Figure 10: Dataset ablation study, GPT2-124M, on SlimPajama 33 Four update control approaches
are shown (see titles). The overall results are similar to those for OpenWeeTTee in the manuscript.
21

---

Figure 10 shows the effects of changing the dataset used in our experiments from Open WebText g SlimPajama [3]. Overall the results are similar as before, when controlling the f--norm via LionA warmup is still beneficial, controlling the angular updates via LionAR decreases the gap significantly. The higher momentum LionAR with Nesterov momentum and our momentum correction eliminates the gap fully. The RRC also seems to eliminate the benefit of warmup but still has the same practical limitations as we describe in 6
LionA
LionAR
LionAR f-0.98 +N
3.5
3.5
3.5
Finll Validation Loss
5%
0%
Final Validation Logs
+ 0%
5%
0%
5%
3.4
3.4
3.4
10%
Rinal Validation
3.3
3.3
3.3
3.2
3.2
3.2
3.1
3.1
3.1
10
10
10-3
10-2
10-3
10
Learning Rate
Learning Rate
Learning Rate
Figure 11: Architecture ablation, Llama2- 124M on OpenWebText. Here LionAR is already sufficient to eliminate the need for warmup, additional RRC compensation is needed. This could be due to the critical batch size being larger (for unclear reasons).
Figure I1 shows the effects changing the architecture from GPT2 Llama style [36 while keeping the dataset and parameter count (124m) the same. This change consists of using SwiGLU activations, RoPE embeddings and RMSNorm. In this case LionAR able to fully eliminate need for warmup without any additional tricks like the RRC compensation momentum corrections. Based on our analysis these additional tricks are likely only needed when the critical batch size is very small initially. We expect that using larger batch sizes could necessitate these additional tricks for Llama, but do not explore this further here.
LionA
LionAR
LionAR 0=0.98 +N +MC
Final Final validation Loss
0%
0%
Final Validation Loss
0%
5%
3.0
Finll Validation
3.0
2.9
2.9
2.9
2.8
2.8
2.8
10
2
10-1
10
10-
Learinng Rate
Learning Rate
Learning Rate
Figure 12: Larger Ilama2-209M training on SlimPajama. This experiment changes the architecture, dataset, model size, and training length (proportional the model size). In this case LionAR suffices on its own again, no additional RRC correction needed as in the smaller llama experiments.
Figure 2 uses a larger Llama model with twice the depth. It also trains for twice as many iterations to keep the ratio of tokens parameters similar. The overall results resemble those of the smaller llaaa experiments in fig.
22

---

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 aimed to write our abstract and introduction accurately reflect the paper's contributions and scope.
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 the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation 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 include limitation section $ which discusses the main overarching limitations from our perspective.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
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?
23

---

Answer: [Yes]
Justification: Our work is mostly empirical, aside from the analysis of the relative represen- tation change where we state the assumptions upfront, with additional details deligated to the appendix.
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 in the statement of any theorems.
The proofs can either appear in 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 Reproducibility
Does the paper fully disclose all the information needed to reproduce the main perimental results the paper to the extent that 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 the main hyperparameters and other aspects of our training. We are not aware any missing details that would significantly hinder reproducibility.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer 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 a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access 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 submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) 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 describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
24

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Question:
Answer: [No]
Justification:
The datasets we use are freely available online but we have not released our code.
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
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https for more details.
/nips. wttpcc/ccccccccccccocooooop httpa//gggudeee/deuuuuubiiiiii wwweiisssonPoPoiiiioioo
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 supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We have attempted include all relevant information.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core 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, 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: Our main results are qualitative and computationally expensive for us, so do not estimate the variance for each point our sweeps. We give an estimate based on select points in.
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.
25

---

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 do this for our main experiments in appx. B.8
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 :ttpp::////uurree hww.nnuu...../. httpp/////..lletttccsaaieeeeeeee
Answer: [Yes]
Justification:
We read it and believe we fully conform to it.
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 is focused on understanding deep learning in general, and does not have specific socictal impacts beyond those of the field overall.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
26

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that 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 generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition 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 don't realease anything that falls under this category.
Guidelines:
The answer NA means that the paper poses no 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 to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We do this to the best of our ability.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URI.
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.
27

---

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-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: No new assets introduced
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:
No crowdsourcing human subjects used
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: No crowdsourcing human subjects used
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.
28

---

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.
29


Reviews:
[{"review_id": "Mar5dgqcSh", "reviewer": "Reviewer_VUaC", "summary": "To train current deep neural network architectures, especially transformers, the learning rate of AdamW is usually first linearly increased to reach a peak before it's decreased to zero. The paper analyzes the impact of this so-called warming-up phase on GPT-2 models from the perspective of the update size. As a second contribution, the paper presents some small modifications for the Lion optimizer to mitigate some of the issues encountered in the experiments.", "questions": "- How sensitive are the results from Section 3 to the type of learning rate schedule? For example, how would Figure 1 look if you used, let's say, a cosine annealing schedule?", "limitations": "I think the paper spells out all its limitations; however, for visibility, it might be better to move the corresponding paragraph from the appendix to the main text. I don't see any negative societal impacts of the work.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is very well written. The experiments are nicely motivated and reasonable.\n\n- Warming up the learning rate is arguably common practice for training transformer models, but not well understood. The paper provides some interesting analysis of the matter, which could potentially lead to a more intuitive understanding of the problem and eventually better optimizers.", "weaknesses": "- The results of the paper are somewhat inconclusive and after reading the paper, I am still not sure about the dynamics during the warming up phase. For example, while controlling angular updates seems to stabilize training to a certain degree, it eventually doesn't lead to better performance. Also, as the paper clearly states, the magnitude of the parameter updates doesn't really account for the gains of the warm-up phase. I am wondering if the paper approaches the problem actually from the right perspective. Having said that, I think the paper still provides some value and might help to stir future research.\n\n- While the empirical evaluation is insightful, it's limited to a single architecture and domain. This raises the question of how reliable the results actually are.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and feedback!\n\n**Generalizability of the results:** We have added experiments with a different GPT architecture (Llama2) and dataset (SlimPajama), see global response. We find that the results are similar, suggesting some transferability within GPT-style training. Originally we wanted to include more transformer tasks like vision transformers and translation but actually found that warmup did not have a significant impact in our target settings. This likely depends on the batch size among other factors, but instead of exploring this further we decided to narrow the scope to GPTs. We will change the title to reflect this, changing “neural network” to “GPT” as suggested by one of the other reviewers.\n\n**General Approach:** The main takeaway would be that the need for warmup largely arises from poor control of the update size during training. Modified optimizers, especially those that control the angular update size directly, can significantly reduce or eliminate the need. However, ultimately we believe that it is large changes in the internal representations that cause the need for warmup, which can not be fully captured by simple measures of the update size of the parameters. We show that the discrepancy between the parameter update size and representation changes can be linked to the noise in the gradient. We can compensate for this based on measurements of the gradient noise, leading to something like an \n“automatic warmup” in the parameter update size. The hope is that this will lead to an improved understanding of optimization dynamics and potentially eventual improvements in optimizer design. Controlling the angular updates improves performance without warmup, but we did not really attempt to improve the overall performance with warmup. We note that LionAR is at a bit of a disadvantage since the weight decay value and other settings are inherited from the baseline (not re-tuned) and we constrain the magnitude without any additional mechanism like learnable gains to compensate for this.\n\n**Cosine Schedule:** Thank you for bringing this up, this is something we will clarify further. We reran the AdamW baseline with a cosine schedule and added it to the global response (see left half of top figure). The trapezoidal schedule (aka warmup-stable-decay) we used in the manuscript has been becoming more popular for LLM training since it provides more flexibility in the training duration and can modify the data mixture in the cooldown phase while giving similar results [1, 2]. The reason we opted to use it is that it clearly separates the warmup phase from the rest of training, unlike common variants of the cosine schedule where the length of the warmup simultaneously affects the shape of the rest of the schedule. We wanted to eliminate this as a confounding factor, so that the apparent benefits of warmup were not coming from changes in the cooldown phase. **In terms of the gap between warmup and no warmup, the results are very similar.** However, in this case the cosine schedule performs marginally better and the learning rate transfers better between different warmup lengths. We believe the latter effect is because the area under the curve is roughly independent of the warmup length for the cosine (because the whole schedule shifts), but not the trapezoidal schedule.\n\n**Limitations:** Yes this is a fair point, we will move the limitations section to the main body in a future version of the manuscript.\n\nPlease let us know if you have any additional questions! We would also greatly appreciate it if you could inform us whether the additional experiments, proposed modifications, and clarifications at least partially mitigate your concerns.\n\n---\n\n[1]: Hu, Shengding, et al. \"Minicpm: Unveiling the potential of small language models with scalable training strategies.\" arXiv preprint arXiv:2404.06395 (2024).\n\n[2]: Hägele, Alexander, et al. \"Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations.\" arXiv preprint arXiv:2405.18392 (2024).", "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": "G8PUMG2b48", "reviewer": "Reviewer_pDv7", "summary": "In this paper the authors investigate the performance benefits seen from the common practice of learning rate warmup and scheduling, attempt to understand the mechanistic underpinnings of those improvements, and engineer optimizers that mitigate the need for warmup. They conduct experiments using NanoGPT and consider controlling parameter updates, angular updates, and “relative representation” changes to close the gap between warmup and no warmup.", "questions": "Can the authors provide some discussion about the sensitivity of RRC to variance in the inputs across and within batches? It is also not clear what direction the RRC results are suggesting. An “automatic warmup” that scales the update sizes according to online measurements of the signal-to-noise ratio is still a warmup phase, albeit a more principled way to arrive at what that schedule should look like. To be clear, I don’t think this is a bad thing, but it may be more representative of the results to propose an “adaptive” or “automatic” learning rate scheduler, rather than claim to make progress towards eliminating the need for scheduling.", "limitations": "Partially", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Warmup length and peak learning rate are certainly some of the most important hyperparameters in large model training, and eliminating the need for a warmup phase would present a significant simplification to training. The paper is well motivated and systematic in its investigation of warmup and proposals to sidestep the necessity of warmup. The RRC is an interesting and promising angle on this question.", "weaknesses": "The results do not suggest a clear prescription for learning rate scaling or straightforward changes that can be made to initializations or updates. In particular, the RRC is completely dependent on the inputs, but there does not seem to be any discussion or investigation of the effects of the input data. For the NanoGPT experiments there is no mention of what the data is. Presumably training was done with a cross entropy loss, but this is also not mentioned.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and feedback!\n\n**Missing Dataset Information:** Thank you for pointing this out! This was an unfortunate oversight on our behalf (and we geatly appreciate your attention to these details). The dataset we train on is the original OpenWebText dataset [1] used in the original GPT2 work. Training is indeed performed via next token prediction using cross entropy with teacher forcing. We will include this in future revisions of the manuscript.\n\n**RRC Dependency on Data:** We have expanded our experiments to include a separate dataset (SlimPajama [2]), obtaining similar results. You are correct that the RRC correction depends on the input data (by design), in a way that controlling the parameter update norm can not. This will vary depending on the data within a batch. The simplest example of this would be if the whole mini-batch consists of the same datapoint repeated. The RRC correction would indicate that we should use a lower learning rate (parameter update size) in this case than if the data has no similarity.\n\nHowever, the RRC correction does not measure the similarity of the data in the input space but rather through the gradient diversity. It therefore not only depends on the similarity of the data within each batch (which shouldn’t vary throughout training) but rather on the similarity of the “learning” to be made from each sequence. This can vary throughout time, for example early in training the model could largely learn syntax, word frequency, and potentially unlearn initialization biases (like maybe outputting the input sample). There is a strong overlap in the “lesson” to be learned from each input sequence in this context, resulting in similarity between the gradients which may be dominated by these simple concepts. Later in training the model could learn more advanced concepts that differ more depending on the semantic content of each sequence, resulting in greater gradient diversity. At this point we can perform larger parameter updates because the contributions of the input sequences do not all line up to change the representations in the same way. Such alignment would lead to large representation changes, which we hypothesized could lead to lasting issues e.g. with the non-linearities.\n\nIn its current form there is no dependency on the similarity between batches for the RRC correction. We hope that if the data is well shuffled the similarity within each mini-batch also reflects the similarity between batches, otherwise this could lead to issues. Let us know if you believe this kind of conceptual discussion of the RRC would be useful / informative we would be happy to incorporate it into the manuscript.\n\n**Eliminating Warmup vs Automatic Warmup:** This is a fair point and we do indeed refer to the RRC correction as an automatic warmup in lines 41 and 220. Overall we believe the definition of a warmup or a schedule is a bit subjective. For example Adam can be seen as a per-coordinate adaptive scheduler for the vanilla SGD learning rate, but what we refer to as the learning rate schedule does typically not account for this. In the same way an RRC corrected optimizer could eliminate the need for manually specified warmup, but would lead to a warmup like effect in the parameter update size. Using other metrics to measure the update size like the RRC, there would not necessarily be an observable warmup phase (and the learning rate could directly control this update size).\n\n**Clear Recommendations:** We do not touch upon initialization but we do show optimizer modifications (LionAR) and potentially further scaling via the RRC gradient noise correction is sufficient to significantly reduce or eliminate the need for warmup. People could attempt to directly apply these methods in other settings if warmup is undesirable for some reason. However we view the primary contribution of our work to be understanding of how the changes in the update size across time contribute to the need for warmup. We hope this will lead to a better understanding of optimization dynamics, that it could inform practitioners about the length of warmup (by just measuring the same metrics we do), and finally lead to better optimizer design in the future. We take the steps towards the optimizer design with LionAR and the RRC correction, but believe future work could improve upon them.\n\n**Additional Experiments:** Aside from repeating our experiments with a different dataset we have performed several other ablations for additional experimental evidence, see the global response.\n\nPlease let us know if you have any additional questions! We would also greatly appreciate it if you could inform us whether the additional experiments, proposed modifications, and clarifications at least partially mitigate your concerns.\n\n---\n\n[1]: Gokaslan, Aaron, and Vanya Cohen. \"OpenWebText Corpus.\" 2019, Skylion007.github.io/OpenWebTextCorpus.\n\n[2]: Soboleva, Daria, Faisal Al-Khateeb, Robert Myers, Jacob R. Steeves, Joel Hestness, and Nolan Dey. \"SlimPajama: A 627B Token Cleaned and Deduplicated Version of RedPajama.\" June 2023, www.cerebras.net/blog/slimpajama-a-627b-token-cleaned-and-deduplicated-version-of-redpajama.", "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": "uVbIU2bEAe", "reviewer": "Reviewer_gjgG", "summary": "The submission analyzes the underlying reason behind the need for a learning rate warmup in neural network training, focusing on GPT pre-training with AdamW and Lion optimizers. The authors identify three key reasons as to why the initial updates are large: \n1. Momentum handling by AdamW, \n2. Early updates not correlated with the initial weight magnitudes\n3. Correlation between gradients of examples during early training\n\nThe study introduces modifications to the Lion optimizer to mitigate the first two issues and proposes a method for controlling activation updates to address the third. Overall, I believe the paper's contributions are significant and hence I vote for acceptance.", "questions": "* Which dataset is used for training the model? I don't think it is mentioned anywhere in the paper and its important for reproducing the results.\n* Can the authors clarify the statement 'This factor is larger if the gradients are negatively correlated, which we empirically observe often happens early in training' on line 110?\n* Did the authors try Adam with inverse bias correction for the momentum as suggested by equation 1?", "limitations": "See weaknesses", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "* This paper analyzes various metrics that correlate with the benefit of learning rate. \n* The analysis of the normalized Gradient Descent is insightful and reproduces previously known scaling laws.", "weaknesses": "* The authors begin by analyzing warmup for Adam and instead of directly modifying Adam, they modify the Lion optimizer. Direct modifications to the Adam optimizer would be more convincing and then moving to Lion would streamline the arguments.\n* The LionAR algorithm is a much more complex solution than AdamW + warmup. Warmup duration is not a crucial hyperparameter, as a longer warmup duration does not hurt training.\n* The experiments are performed on a fixed setup: GPT-2 model with 100M parameters trained on a single dataset. To ensure the validity and generalizability of the results, it is crucial to extend the analysis to various model architectures, parameter sizes, and datasets.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and feedback!\n\n**Using Adam instead of Lion:** We actually experimented with direct modification to Adam originally before moving to Lion. This worked equally well or better than the Lion modifications. However, we realized that at the start of training the gradient norms can change rapidly, specifically decreasing in our setup. This causes the second moment in Adam to be larger than it would otherwise be, resulting in smaller update sizes and essentially giving an additional warmup-like effect. The update size is also affected by the alignment of the gradients in successive steps (through momentum), complicating the control of the update size. The main reason for moving to Lion was that it gives precise control of the update size at a given iteration, eliminating these confounding effects that we found hard to control for in Adam and allowing us to explore the effect of the update size directly. We will edit the manuscript to clarify this, but we acknowledge your point that the use of Lion complicates the story.\n\n**Complexity of LionAR:** This is a fair point, the algorithm for LionAR is more complicated than for AdamW. A significant portion of this complexity comes from trying to transfer the hyperparameters of Adam over to Lion, as well as handling the two types of momentum. This could be removed if we don’t need this which would simplify the algorithm considerably. Then the major difference is additional projection of the weight norm instead of using weight decay. In practice weight decay is also only applied to some parameters which could be expressed with a similar if/else branch. In terms of the optimization dynamics we expect the behavior of LionAR to be much more regular than when using weight decay which could matter more than the complexity of the code overall. That being said, we also agree that warmup is a perfectly fine solution in practice, we primarily want to offer insights into why it is needed.\n\n**Diversity of the Experiments:** We agree and would have liked to showcase a broader range of experiments. We are unfortunately somewhat restricted in our compute budget but have tried to add some simple ablations in the global response. We tried a different dataset (SlimPajama [1]) with the same GPT setup and found similar results. We also tried to change the architecture to a llama2 instead of the GPT2 style we used originally. In this case LionAR already suffices to eliminate the need for warmup without further tricks like the RRC or momentum changes. This is likely due to changes in the critical batch size although the exact reasons for why the architecture affects this are not clear to us. We want to further experiment with larger batch sizes and see if the llama behavior becomes more similar to that of GPT2.\n\n**Dataset:** Thank you for pointing this out. This was an oversight on our part, and we appreciate your thorough review. The dataset we used is OpenWebText [2], the same one referenced in the GPT2 paper and NanoGPT. We will update the manuscript to include this information.\n\n**Line 110 clarification:** We meant correlation over time / between optimizer steps, i.e. that when the gradients of successive steps point in opposite directions they cancel out in the momentum vector. This leads to smaller update sizes in Adam. However, the bias correction assumes that all future gradients will align with the current momentum vector and magnifies the update size accordingly, leading to much larger steps than we would see otherwise. We will edit this to clarify.\n\n**Adam with inverse bias correction** Yes, we did some exploratory experiments with modifications like these as well as removing the bias correction completely. This helps a bit, but overall the effect is too small to make a significant difference for momentum values like 0.9. At best this could result in warmup like effects of maybe 20 steps, which is too small to significantly decrease the no-warmup degradation in our setting. For comparison, see also the RAdamW results in the global response which might give effects similar to 40 steps of warmup which is not sufficient.\n\n---\n\n[1]: Soboleva, Daria, Faisal Al-Khateeb, Robert Myers, Jacob R. Steeves, Joel Hestness, and Nolan Dey. \"SlimPajama: A 627B Token Cleaned and Deduplicated Version of RedPajama.\" June 2023, www.cerebras.net/blog/slimpajama-a-627b-token-cleaned-and-deduplicated-version-of-redpajama.\n\n[2]: Gokaslan, Aaron, and Vanya Cohen. \"OpenWebText Corpus.\" 2019, Skylion007.github.io/OpenWebTextCorpus.", "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": "kLWahoV5H3", "reviewer": "Reviewer_ramW", "summary": "This work explores the benefits of learning rate warmup in neural network training, focusing on the size of model updates via the GPT2 model. It finds that controlling update size in parameter space doesn't fully explain warmup's advantages, but quantifying updates in terms of neural representation changes shows promise. The study also highlights the role of high momentum in warmup and suggests potential methods for reducing the need for manual warmup configuration. Overall, the research provides insights into learning rate warmup's necessity and potential ways to eliminate it in practice.", "questions": "See Weakness.", "limitations": "I have not found any discussions about the limitations and potential negative societal impact. But in my opinion, this may not be a problem, since the work only focuses on analyzing the warmup heuristic in machine learning. Still, it is highly encouraged to add corresponding discussions.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper addresses an intriguing topic, which aims to present a systematic understanding regarding the LR warmup heuristic from a novel perspective.  However, I feel that the authors have attempted to cover too many aspects, which might be challenging to thoroughly demonstrate within the scope of a single conference paper.", "weaknesses": "1. I noticed that the authors have not adequately discussed the highly relevant paper, \"On the Variance of the Adaptive Learning Rate and Beyond,\" which addresses some of the questions raised by the authors. Please discuss the unique contributions of your work compared to the variance-based analysis presented in that paper.\n\n2. Although the authors try hard to explain the need for warmup and how to potentially reduce it, I still did not find persuasive answers to the questions posed. The conclusions are primarily based on intuitive narrative explanations and a simple experiment involving GPT-2. Meanwhile, some of the conclusions seem to be evident. For instance, before I read the paper, I could understand the statement \"L2 update size is not sufficient to quantify the 'effectively' large updates\". The paper lacks convincing evidence to support its claims. Lastly, I recommend that the authors narrow down the scope of the title to accurately reflect the content presented in the paper.\n\n3. The authors use linear transitions to analyze the representation changes, which seems too toy for me.\n\n4. I think the gradient clipping operation may be quite related regarding the authors' idea, as it directly impacts the adaptive LR. Could the authors provide some research here?\n\n5. Regarding the writing of this paper, in my opinion, it is not particularly easy to follow. The organization feels somewhat messy. I think the authors should improve the clarity and structure. For example, including more detailed explanations and transitions between sections.\n\n6. In Figure 1, I observe that the performance may be quite similar when using a lower learning rate. Could the authors specify the lowest learning rate used in your experiments?\n\n7. I found the authors' use of the term \"update size\" to denote the step size in Adam somewhat confusing. I recommend that the authors use \"update step\", \"adaptive learning rate\" or \"effective learning rate\" instead, as these terms are clearer.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and feedback!\n\n**W7 Update Size:** We want to begin by clarifying this as it is a very important concept for our paper. By “update size” we literally mean “the size of the update”. This can be measured in different ways, for example through the L2 norm of the update, the angular change, or the representation change. We note that this is different from the learning rate, which scales the update size, but does generally not fully determine the update size. For example in SGD the L2 update size depends on the learning rate as well as the norm of the gradient. We specifically use “update size” or “update magnitude” to try to avoid confusion with these other terms that we do not feel capture our intended meaning.\n\n**W1 Comparison with RAdamW:** Thank you for this suggestion for improving our paper. Since RAdamW is a popular method for reducing the need for warmup it serves as a good comparison and contextualization of our work. Figure 1 in the global response shows that the RAdamW modifications are insufficient to eliminate the need for warmup in our setting. [2] suggest that RAdamW is approximately equivalent to 4 steps of SGDM followed by Adam with a special type of built-in warmup whose length is roughly 2/(1-beta2) = 40, which is likely too short in our setting.\n\nThe analysis in [1] is based on the idea that early in training the second-moment estimates are not accurate (noisy) and can therefore not be trusted to scale the update properly. This could in turn contribute to the need for warmup. We note that without momentum, perfect estimates of the second moment at the current time step would control the expected L2 norm of the update. This relates our approach of looking at the update size to the adaptive learning rate view you seem to favor. We note that although [1] focuses on this issue of counteracting noisy estimates of the second moment, this is not necessary the sole reason warmup is beneficial. This is supported by the fact that both SGD and Lion empirically need warmup in various settings but do not use the second moment at all, indicating there is more to the story.\n\nThe other aspects we explore relate to the momentum (e.g. the bias correction), weight decay and initialization (the angular updates), and how the gradient diversity affects how quickly the internal representations change (RRC compensation). We don’t believe [1] touches upon any of these to a significant extent. We also use Lion to control the L2 update size, avoiding the issues that [1] focuses on. Both works aim to reduce the need for learning rate warmup but focus on different aspects and take a completely different approach. Overall we therefore believe there is very little overlap in the contributions of these works, and if anything they could be complementary if we were to port our Lion modifications back to AdamW.\n\n**W2 Title:** Yes, we agree that just using GPT instead of neural networks would be more fitting. Thank you for this suggestion.\n\n**W2 L2 Norm is obvious:** We don’t believe this is the case, see both the discussion of gradient clipping and RAdamW which are closely related to this idea and not at all obvious shouldn’t work. We also note that clipping the L2 norm of the update will prevent a network from diverging to infinity, something that is often observed in unstable SGD training where warmup can help. Could you clarify further why you believe this is obvious or provide some reference to prior work that shows this?\n\n**W3 Linear Transformations:** This is of course a simplification compared to full neural networks. However, we believe this already gives interesting insights that are sufficient for our purposes. Specifically, this clearly shows that the parameter update size must be decreased when the gradient diversity is low in order to keep the representation changes small. We are also able to draw the same conclusions as existing works e.g. muP and hyperparameter scaling laws, showing this analysis provides useful results despite its simplicity. We believe a more complicated analysis would not fit within the scope of this paper, which is already on the broad side as you mention.\n\n**W4 Gradient Clipping:** Yes, gradient clipping can directly affect the update size. This is especially true when using SGD without momentum, where gradient clipping is similar to clipping the L2 norm of the update. However, controlling the L2 norm is not sufficient as we show. We believe the other ideas like the angular updates and representation changes are less related (but please clarify if you believe they are).\n\n**W5 Presentation:** Could you be more specific here? We note that some of the other reviewers found the paper to be “very well written” and “well motivated and systematic”, but we are happy to consider any concrete changes you recommend.\n\n**W6 Range of the sweep:** Yes, the performance is similar for smaller learning rates. This is consistent with our main hypothesis that large updates cause the degradation that warmup counteracts. Smaller learning rates result in smaller updates, thus avoiding this issue. The lowest learning rate in this sweep is 3e-4.\n\nPlease let us know if you have any further questions or suggestions! If you feel we have at least partially addressed your concerns, we would greatly appreciate it if you would consider reevaluating your review score.\n\n---\n\n[1] 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. URL https://openreview.net/forum? id=rkgz2aEKDr. arXiv:1908.03265.\n\n[2] Jerry Ma and Denis Yarats. On the adequacy of untuned warmup for adaptive optimization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 8828–8836, 2021. arXiv:1910.04209.", "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 would like to thank all the reviewers for their thoughtful reviews and feedback on our manuscript. We will try to address your specific concerns and questions in our individual responses.

Here we present additional experimental results with accompanying plots in the pdf:
* Figure 1 shows what the baseline plots would look like for a **cosine decay schedule** (specifically one-cycle cosine without momentum for these experiments). The reason we went with the trapezoidal schedule is that it fully separates the warmup phase from the rest of the schedule unlike in the cosine schedule where the warmup length affects the whole schedule. We wanted to avoid these confounding effects with our choice of the trapezoidal schedule.
* Figure 1 also shows the performance of **RAdamW** [1], a popular optimizer modification for reducing the need for warmup, in the baseline setup. We find that while it helps, it does not eliminate the need for warmup. The analysis of [2] suggest RAdamW functions similar to a 2/(1-beta2)=40 step warmup which seems to roughly match our findings (2% here would be 100 steps).
* Figure 2 shows the **effects of changing the dataset used in our experiments** from OpenWebText [3] to SlimPajama [4]. Overall the results are similar as before, when controlling the L2 norm via LionA warmup is still beneficial, controlling the angular updates via LionAR decreases the gap significantly. The higher momentum LionAR with Nesterov momentum and our momentum correction eliminates the gap fully. The RRC also seems to eliminate the benefit of warmup but still has the same practical limitations as we describe in section 6.1.
* Figure 3 shows the **effects of changing the architecture from GPT2 to the llama style [5]** while keeping the dataset and parameter count (~124m) the same. This includes using SwiGLU activations, RoPE embeddings and RMSNorm. In this case LionAR is able to fully eliminate the need for warmup without any additional tricks like the RRC compensation or momentum corrections. Based on our analysis these additional tricks are likely only needed when the critical batch size is very small initially. In the future we want to rerun these experiments using a larger batch size to verify this, but were not able to do it in time for this rebuttal.
* Figure 4 shows the results for a **larger 209m parameter llama2** trained on SlimPajama. Overall the results are similar to the smaller llama.

**We believe these additional experiments can increase the variety of our experimental setup, helping mitigate this limitation somewhat**, although an even broader range would of course be preferable.

Based on reviewer feedback, we have decided to limit our experimental scope to GPT variants and will change the title of the paper to reflect this, as suggested by one of the reviewers. Previously we wanted to include broader transformer experiments like DeiT and translation but found it hard to identify good setups where warmup has a significant impact but are computationally tractable for us. We found that many of the reference configurations we experimented with use warmup but don’t actually benefit significantly from it. This likely varies across batch sizes and other configuration aspects we found too expensive to tune.

---

[1] 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. URL https://openreview.net/forum? id=rkgz2aEKDr. arXiv:1908.03265.

[2] Jerry Ma and Denis Yarats. On the adequacy of untuned warmup for adaptive optimization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 8828–8836, 2021. arXiv:1910.04209.

[3]: Gokaslan, Aaron, and Vanya Cohen. "OpenWebText Corpus." 2019, Skylion007.github.io/OpenWebTextCorpus.

[4]: Soboleva, Daria, Faisal Al-Khateeb, Robert Myers, Jacob R. Steeves, Joel Hestness, and Nolan Dey. "SlimPajama: A 627B Token Cleaned and Deduplicated Version of RedPajama." June 2023, www.cerebras.net/blog/slimpajama-a-627b-token-cleaned-and-deduplicated-version-of-redpajama.

[5]: Touvron, Hugo, et al. "Llama 2: Open foundation and fine-tuned chat models." arXiv preprint arXiv:2307.09288 (2023).

Meta-review: The paper under consideration investigates the common but not well-understood practice of warming up learning rates while training transformer models.
Reviewers appreciated the importance of the topic and the clarity of the writing. In addition, reviewers expressed appreciation for the systematic investigation put forth in the paper.
While the reviewers noted some limitations, including the inconclusiveness of the results, their primary concern was the study's initial narrow focus on a single architecture and dataset. However, the authors' post-rebuttal additions, which included additional experiments, successfully addressed this concern by providing consistent results that reinforced the initial findings.
Overall, reviewers were overwhelmingly positive in their assessment and all unanimously recommended acceptance of the paper.
Decision: Accept (poster)