Title: Analyzing & Reducing the Need for Learning Rate Warmup in GPT Training
Venue: NeurIPS 2024
Accepted: yes
Keywords: learning rate

warmup

schedule

AdamW

LION

rotational

GPT

LLAMA

transformer

deep learning

empirical

momentum

critical batch size

gradient diversity

signal-to-noise ratio
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

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 ∆ w t = η t u t early in training by using lower values for the learning rate η t . In this work we argue that warmup benefits training by keeping the overall size of ∆ w t limited, counteracting large initial values of u t . Focusing on small-scale GPT training with AdamW/Lion, we explore the following question: Why and by which criteria are early updates u t too large? We analyze different metrics for the update size including the ℓ 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 u t based on the aforementioned metrics.

1 Introduction

Neural networks are typically trained using variations of stochastic gradient descent. The weight updates ∆ w have the form ∆ w = η u , where η denotes the learning rate and u an unscaled update vector derived from the history of weight gradients. Throughout training, the learning rate is often adjusted over time t according to a learning rate schedule , η = η t . This schedule frequently includes an initial phase known as a learning rate warmup, where the learning rate starts low and is increased to a target value before being reduced according to 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 et al. [10] for large batch size ResNet [11] training, which is 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 ∆ w = η u , 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 η t values to counteract large u values early in training. How should we quantify what makes an update ∆ w large? Why do large updates adversely affect training?

In this work, we explore warmup from this perspective, focusing on GPT2 [29] training with adaptive optimizers like AdamW [24] and Lion [3]. We identify three key issues that necessitate warmup:

The way Adam handles momentum can lead to artificially large initial updates ∆ w .

Early updates ∆ w are large compared to the initial weight magnitude of w for matrices.

The gradients of early samples are highly correlated, limiting effective mini-batch sizes.

Wedemonstrate that simple modifications to the optimizer can mitigate the first two issues: eliminating the momentum bias correction in AdamW and scaling matrix updates to match their magnitude, akin to the Rotational Optimizers by Kosson et al. [19]. 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 signal-to-noise ratio of the gradient that can help mitigate this, functioning like an 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 [11], where a 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 et al. [6]. The practice of linear warmup in 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.1 we show that this warmup effect is insufficient in our setup. Relative optimizers like LARS [47] and LAMB [48] are also believed to reduce the necessity for warmup [21]. 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 et al. [41]. Xiong et al. [43] proposed the pre-LN normalization placement for transformers, showing it reduces the need for warmup. Huang et 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. [32] 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 as we discuss later.

The relation between stochastic gradient noise and learning rate has been studied in several works [46, 28, 49, 34, 22, 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 a 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 a 124M parameter GPT2 [29] model on the OpenWebText 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, as suggested by Chinchila [12]. The baselines use AdamW [24] (see algo. 1) with weight decay λ = 0 . 1 , momentum coefficient β 1 = 0 . 9 , smoothing coefficient β 2 = 0 . 95 , and ε = 10 -8 . The learning rate schedule consists of a

0

2500

5000

Iteration

0

1

Learning Rate Schedule

0%

2%

5%

10%

20%

10

3

10

2

Learning Rate

3.1

3.2

3.3

3.4

3.5

Validation Loss

0

0.05

0.50

1

Iteration (

scale)

4

6

8

10

Training Loss

0

0.05

0.50

1

Iteration (

scale)

0

1

2-update size / max

linear warmup followed by a constant phase and eventually linear cooldown spanning half of training (see examples in fig. 1). This schedule keeps the peak learning rate and decay phase identical for different warmup lengths. This differs 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. [19]. The hyperparameter values and base training configuration are adopted from NanoGPT. See appx. C.2 for experiments on additional architectures and datasets.

Figure 1 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 a given learning rate, but eventually falls behind leaving a permanent gap. Warmup not only stabilizes higher learning rates, but also prevents a 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 ℓ 2 -magnitude varies significantly throughout training with a large spike at the start of training.

4 The Interaction of Momentum and the ℓ 2 -Update Norm in AdamW

In this section we analyze the reason behind the large ℓ 2 -norm of early updates in our AdamW baseline seen in panel 4 of fig. 1. We find that this primarily stems from the β 1 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 ℓ 2 -norm of the update. Although we find this is to 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. 1) differ from simpler methods like SGD with momentum 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 in the gradient magnitude between layers. An important consequence of this is that the unscaled updates u are not large simply due to large initial gradients, unlike in plain SGD and other optimizers that don't normalize their

Algorithm 1 AdamW (PyTorch variant, differing from the original by Loshchilov and Hutter [24])

Require:

Learning rate η t , weight decay λ , momentum β 1 , magnitude smoothing β 2 , ε for numerical stability

1: Initialize: Time step t ← 0 , parameter vector θ 0 , momentum vector m 0 ← 0 , magnitude vector v 0 ← 0

2: while stopping criteria not met : 3: t ← t +1 4: g t ← Mini-batch gradient w.r.t. θ t -1 5: m t ← β 1 m t -1 +(1 -β 1 ) g t 6: v t ← β 2 v t -1 +(1 -β 2 ) g 2 t 7: ˆ m t ← m t / (1 -β t 1 ) 8: ˆ v t ← v t / (1 -β t 2 ) √

9: θ t ← (1 -η t λ ) θ t -1 -η t ˆ m t / ( ˆ v t + ε )

10

3

10

2

Learning Rate

3.1

3.2

3.3

3.4

3.5

Validation Loss

0%

2%

5%

10%

20%

0

0.05

0.50

1

Iteration (

scale)

4

6

8

10

Training Loss

0

0.05

0.50

1

Iteration (

scale)

0

1

2-update size / max

0

0.05

0.50

1

Iteration (

scale)

0

1

Avg

-update size / max

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 AdamW normalizes the update size based on the gradient, its magnitude can still vary throughout training as seen in fig. 1. This can be caused by changes in the gradient magnitude over time, especially when using different values of β 1 and β 2 . However, it can also be caused by momentum and especially the bias correction (algo. 1, line 7). The magnitude of m t depends on the alignment of subsequent gradients g 1 , . . . , g t whereas the normalization factor v t does not. For example, when each g t is an independent zero-mean random vector with a fixed second moment E [ g 2 t ] = σ 2 , we have (see appx. B.1 for details):

In this case the bias correction for β 1 is incorrect since it is derived for a constant gradient. With the bias correction the size becomes E [ ∥ ˆ m ∥ 2 ] = 1+ β t 1 1 -β t 1 1 -β 1 1+ β 1 σ 2 , amplifying the norm of early updates by √ (1 + β t 1 ) / (1 -β t 1 ) . 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 ℓ 2 -norm of AdamW updates can therefore vary significantly due to the initial bias correction, changes in the alignment 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-wise sign operation to normalize the update, giving +1 for positive values, -1 for negative values and 0 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 t instead of g t in line 6 while setting β 2 = 0 . It also uses a form of scaled Nesterov momentum instead of the traditional heavy-ball variant and a hyperparameter specification that differs significantly from AdamW. We propose a Lion variant, LionA (algo. 2), that keeps the hyperparameters more compatible with those of AdamW. The learning rate is kept comparable by scaling the ℓ 2 update size to match that of AdamW in the random-gradient scenario, see appx. B.1 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 as v being inaccurate from rapidly decreasing gradient magnitudes early in training, which can induce an additional warmup-like effect.

In fig. 2 we repeat the baseline sweep using LionA. Despite perfect control of the ℓ 2 update norm (as seen in panel 3), the benefit of warmup remains. This leads us to conclude that the ℓ 2 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 a 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 [3] optimizer for greater compatibility with AdamW (algo. 1). The sign operation replaces the magnitude smoothing, explicitly controlling the ℓ 2 -norm of each update. Additional scaling keeps the hyperparameters η, λ comparable to AdamW.

5 The Importance and Irregularity of the Angular Update Size

The effect of a weight vector w t ∈ R C used in a dot product with some vector x (e.g., in a neuron):

can be understood in terms of its magnitude ∥ w t ∥ and direction w t / ∥ w t ∥ . The magnitude acts like a gain, scaling the outputs, whereas the direction determines which input representations x the system responds to. The angular update size [38] of an update w t ↦→ w t +1 is defined as:

and measures how fast the direction of w t changes during training, and thus its 'preference' for x .

With BatchNorm [15] and similar operations [1, 14, 42, 4], a network can become invariant to the magnitude of weight vectors like w t , 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:

Note that although the system output is invariant to the magnitude ∥ w t ∥ , traditional optimizers are not. Scaling the value of a scale-invariant weight vector by a factor of c > 0 , results in a gradient that is scaled by c -1 and curvature that is scaled by c -2 (see appx. B.2). For SGD this scales the angular update by c -2 and for Adam-like optimizers it is scaled by c -1 . With weight decay the magnitude of scale-invariant vectors trends towards a certain stable equilibrium value over time which also results in a 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 be 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 scale-invariant, 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. 3. We normalize the angular update size to match the equilibrium value, replacing weight decay with projections similar to Kosson et 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 al. [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.

Algorithm 3 LionAR: A rotational version of algo. 2 inspired by Kosson et al. [19]. The parameter vector is divided into sub-vectors θ = [ θ (1) , . . . , θ ( P ) ] , each corresponding to either the weight vector of a neuron (e.g. a matrix row / a convolutional filter), or other parameters such as gains and biases. The updates of 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 as a scaling factor for the relative updates of neurons [19].

10

3

10

2

Learning Rate

3.1

3.2

3.3

3.4

3.5

Validation Loss

0%

2%

5%

10%

20%

0

0.05

0.50

1

Iteration (

scale)

4

6

8

10

Training Loss

0

0.05

0.50

1

Iteration (

scale)

0

1

Avg

-update size / max

0

0.05

0.50

1

Iteration (

scale)

0

1

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 it in this setting . Both the angular change and the ℓ 2 -norm are simple measures of the update magnitude in parameter space that do not account for the direction or other aspects of the update. In the next section we show how a 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. 3.

6 Early Gradient Alignment Results in Large Representation Changes

Our two approaches to measuring and controlling the update size in weight space failed to fully explain the need for warmup. As an alternative to the parameters, we can analyze changes in

the internal representations or activations of the neural network (feature learning). Although this is harder to 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 in 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.

A 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 a given magnitude on the representations will be larger than otherwise. We will analyze this for a dot product making up a single neuron:

where X = [ x 1 , . . . , x B ] ∈ R C × B are the C -dimensional representations of a random mini-batch of B inputs that is fed into the neuron, w ∈ R C is the weight vector, and y ∈ R B is a batch of outputs. For a weight update w ↦→ w +∆ w , we aim to quantify the size of the output change ∆ y = ∆ w ⊤ X computed on the same inputs. We focus on the Relative Representation Change (RRC) :

similar to the angular weight updates, as the sensitivity to the absolute change ∥ ∆ y ∥ can be unclear due to normalization or other scaling operations. Note that this is a measure of a 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 on the relatively tractable case of normalized gradient descent with updates:

where g b 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):

D1: We define g b =: ¯ g + ˜ g b where ¯ g = E [ g ] and ˜ g b is the difference with E [˜ g b ] = 0 .

D2: We define φ := E [ ∥ ¯ g ∥ 2 ] / E [ ∥ ˜ g b ∥ 2 ] as the Signal-to-Noise Ratio (SNR) of the gradient.

P1: For a neuron, g b ∥ x b , and hence x b = sign( ⟨ x b , g b ⟩ ) · ( ∥ x b ∥ / ∥ g b ∥ ) · (¯ g + ˜ g b ) .

L1: Consider two independent random vectors a ∈ R C and b ∈ R C , 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 a and b is given by E [ ⟨ a , b ⟩ 2 ] = E [ ∥ a ∥ 2 ] E [ ∥ b ∥ 2 ] /C .

̸

A1: We assume the following vector pairs satisfy L1: ( x i , ˜ g b ) when i = b , (¯ g , ˜ g b ) and ( w , x b ) .

This allows us to compute the expected square relative representation change (detailed in appx. B.4):

The expected relative change in the output for a given sample can be broken down into three sources, the contribution of the sample itself (first term), random interference from the 'noise' ˜ g i of other samples (second term), and finally amplification of the common mean component ¯ g (third term).

10

4

10 0

SNR

10

2

10

1

10 0

LR scaling for const RRC

B

=10 0

B

=10 2

B

=10 4

0

0.05

0.50

1

Iteration (

scale)

10

4

10

3

10

2

10

1

10 0

Signal-to-Noise Ratio

0%

5%

0

0.05

0.50

1

Iteration (

scale)

1.0

0.5

0.0

0.5

1.0

Alignment cos

(

gt

,

mt

1)

0

0.05

0.50

1

Iteration (

scale)

1.0

0.5

0.0

0.5

1.0

Velocity-Grad Cancellation

The RRC expression provides many interesting insights. In the case of large input dimension C →∞ and small SNR φ ≈ 0 , keeping the RRC constant for different batch sizes involves scaling the learning rate η ∝ √ B , 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 to break down around B = 1 /φ , matching the predicted critical batch size of e.g. McCandlish et al. [28]. The role of the dimension C in the expression is 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 µ -parameterization [44, 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 η ∝ √ 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.e. 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 micro-batches. Although this helps close the gap between warmup and no-warmup, the overall performance is 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 a factor of 32 × 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 a promising direction but require further work to be practical.

6.2 The Role of Momentum

Momentum is believed to be a key enabler of optimization with larger batch sizes [32, 39, 31]. However, it is unclear how it should change predictions for the critical batch size or relative representation change. Momentum spreads the application of a gradient out over multiple steps which tends to make each update smaller, especially for a random walk, which is 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 normalize the step size. In the last two panels of fig. 4 we observe that in our setup the gradient and momentum are negatively correlated, counteracting each other. We find momentum crucial for performant training, panel 2 of fig. 5 shows significant degradation without it.

10

3

10

2

Learning Rate

3.1

3.2

3.3

3.4

3.5

Validation Loss

LionAR

=0.9 +RRC comp

0%

5%

10%

10

3

10

2

Learning Rate

4.0

4.2

4.4

Validation Loss

LionAR  =0

0%

5%

20%

10

3

10

2

Learning Rate

3.1

3.2

3.3

3.4

3.5

Validation Loss

LionAR

=0.98 +N +MC

0%

2%

5%

10%

20%

10

3

10

2

Learning Rate

3.1

3.2

3.3

3.4

3.5

Validation Loss

LionA

=0.98 +N +MC

0%

2%

5%

10%

20%

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 β 1 bias correction in AdamW (algo. 1, line 7), or introducing an inverse bias correction in Lion like algorithms (see appx. B.1 for details). Higher β 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-ReLUs, 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 a 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, we 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 to exceed the critical batch size early in training. The second is 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 a 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 a single network which may not be 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 ∆ w = η u impacts the need for learning rate warmup. We showed that u can be large initially when measured in terms of its ℓ 2 -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-based measures of the update size, the ℓ 2 -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 can vary significantly 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 a 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-free 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.

Although we present new methods we consider promising, we still recommend the use of a short warmup in practice. Fully eliminating it seems to require significant modifications that also need further validation across additional settings. However, we hope to have provided the reader with a 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

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 , 2016. 5

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-21381, 2020. arXiv:2002.03432. 2

Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, and Quoc V Le. Symbolic discovery of optimization algorithms. In Thirty-seventh Conference on Neural Information Processing Systems , 2023. URL https://openreview.net/forum?id=ne6zeqLFCZ . arXiv:2302.06675. 1, 4, 5

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:1905.05894. 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 International Conference on Machine Learning , 2024. URL https://openreview.net/forum? id=0ksNeD1SJT . arXiv:2407.05872. 18

Jingwen Fu, Bohan Wang, Huishuai Zhang, Zhizheng Zhang, Wei Chen, and Nanning Zheng. When and why momentum accelerates sgd: An empirical study. arXiv preprint arXiv:2306.09000 , 2023. 2

Justin Gilmer, Behrooz Ghorbani, Ankush Garg, Sneha Kudugunta, Behnam Neyshabur, David Cardoze, George Edward Dahl, Zachary Nado, and Orhan Firat. A loss curvature perspective on training instabilities of deep learning models. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id=OcKMT-36vUs . arXiv:2110.04369. 2

Aaron Gokaslan and Vanya Cohen. Openwebtext corpus. http://Skylion007.github.io/ OpenWebTextCorpus , 2019. 2, 22

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. URL https://openreview.net/forum? id=r14EOsCqKX . arXiv:1810.13243. 2

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

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:1512.03385. 1, 2, 14

Jordan Hoffmann, Sebastian 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://arxiv.org/abs/2203.15556 . 2

Xiao Shi Huang, Felipe Perez, Jimmy Ba, and Maksims Volkovs. Improving transformer optimization through better initialization. In International Conference on Machine Learning , pages 4475-4483. PMLR, 2020. URL https://proceedings.mlr.press/v119/ huang20f.html . 2

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.06868. 5

Sergey Ioffe 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. arXiv:1502.03167. 5

Andrej Karpathy. nanogpt. https://github.com/karpathy/nanoGPT/ , 2023. 3

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 IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 24174-24184, 2024. arXiv:2312.02696. 5

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

Atli Kosson, Bettina Messmer, and Martin Jaggi. Rotational equilibrium: How weight decay balances learning across neural networks. In Forty-first International Conference on Machine Learning , 2024. URL https://openreview.net/forum?id=MQirNNU2pC . arXiv:2305.17212. 2, 3, 5, 6, 9, 14, 16

Alex Krizhevsky. Learning multiple layers of features from tiny images. self-published , 2009. URL https://www.cs.toronto.edu/~kriz/learning-features-2009-TR.pdf . 14

Zhiyuan Li, Kaifeng Lyu, and Sanjeev Arora. Reconciling modern deep learning with traditional optimization analyses: The intrinsic learning rate. Advances in Neural Information Processing Systems , 33:14544-14555, 2020. arXiv:2010.02916. 2

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. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , 2021. URL https://openreview.net/forum?id=goEdyJ_nVQI . arXiv:2102.12470. 2

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, 21

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations , 2019. URL https://openreview.net/forum? id=Bkg6RiCqY7 . arXiv:1711.05101. 1, 2, 3, 21

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.net/forum?id=xp5VOBxTxZ . arXiv:2206.07085. 16

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. 2, 21

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 https://openreview.net/forum?id=F2mhzjHkQP . arXiv:2205.10287. 2, 8

Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAI Dota Team. An empirical model of large-batch training. arXiv preprint arXiv:1812.06162 , 2018. 2, 8

Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. self-published , 2019. URL https://d4mucfpksywv.cloudfront.net/better-language-models/ language_models_are_unsupervised_multitask_learners.pdf . 1, 2

Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. Advances in neural information processing systems , 29, 2016. arXiv:1602.07868. 5

Christopher J Shallue, Jaehoon 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(112):1-49, 2019. arXiv:1811.03600. 8

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. arXiv:2006.15081. 2, 8

Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. Slimpajama: A 627b token cleaned and deduplicated version of redpajama. https://cerebras.ai/blog/ slimpajama-a-627b-token-cleaned-and-deduplicated-version-of-redpajama , 2023. 21, 22

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 4042-4050. PMLR, 2021. arXiv:2103.02351. 2

Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum in deep learning. In International conference on machine learning , pages 1139-1147. PMLR, 2013. URL https://proceedings.mlr.press/v28/sutskever13. html . 2

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 arXiv:2307.09288 , 2023. 22

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

Ruosi Wan, Zhanxing Zhu, Xiangyu Zhang, and Jian Sun. Spherical motion dynamics: Learning dynamics of normalized neural network using sgd and weight decay. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , volume 34, pages 6380-6391. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/paper/2021/file/ 326a8c055c0d04f5b06544665d8bb3ea-Paper.pdf . arXiv:2006.08419. 4, 5

Runzhe Wang, Sadhika Malladi, Tianhao Wang, Kaifeng Lyu, and Zhiyuan Li. The marginal value of momentum for small learning rate SGD. In The Twelfth International Conference on Learning Representations , 2024. URL https://openreview.net/forum?id=3JjJezzVkT . arXiv:2307.15196. 8

Ross Wightman. Pytorch image models. https://github.com/rwightman/ pytorch-image-models , 2019. 14

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://arxiv.org/abs/2309.14322 . 2

Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV) , pages 3-19, 2018. arXiv:1803.08494. 5

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. arXiv:2002.04745. 2

Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hyperparameter transfer. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , 2021. URL https: //openreview.net/forum?id=Bx6qKuBM2AD . arXiv:2203.03466. 8

Greg Yang, James B Simon, and Jeremy Bernstein. A spectral condition for feature learning. arXiv preprint arXiv:2310.17813 , 2023. 8

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. arXiv:1706.05699. 2

Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint arXiv:1708.03888 , 2017. 2

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. URL https://openreview.net/forum?id=Syx4wnEtvH . arXiv:1904.00962. 2

Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger B Grosse. Which algorithmic choices matter at which batch sizes? Insights from a noisy quadratic model. Advances in Neural Information Processing Systems , 32, 2019. arXiv:1907.04164. 2

100

200

300

400

500

600

700

800

Epochs

86

88

90

92

Test Accuracy

1 Standard

8 Standard

128 Standard

1 Leaky ReLU (0.1)

8 Leaky ReLU (0.1)

128 Leaky ReLU (0.1)

A The Detrimental Effects of Large Updates

To investigate the effects of large updates at the beginning of training, we conducted controlled experiments on a ResNet-20 model on CIFAR-10 [20] due to resource constraints. We controlled the average angular update throughout training using the rotational optimizer variant of AdamW proposed by Kosson et al. [19]. For the initial phase of training, 5 epochs, we either use a standard learning rate of 0 . 05 or amplify it by a factor of 8 or 128. This results in very large updates, scaling the rotation by either √ 8 or √ 128 . For all experiments, we used a weight decay of 0 . 01 , β 1 = 0 . 9 , β 2 = 0 . 999 , 5 epoch initial phase, and trained for 205 epochs in total with a cosine schedule unless otherwise specified. The data was pre-processed by normalizing it with a mean of (0 . 4914 , 0 . 4822 , 0 . 4465) and a standard deviation of (0 . 2023 , 0 . 1994 , 0 . 2010) and applying simple data augmentation techniques as described by He et al. [11]. To run the experiment, we used the codebase from Wightman [40] and extended the utilities from Kosson et al. [19].

As shown in fig. 6, the performance of 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 ReLU networks. This suggest that large, controlled, initial updates can result in a permanent performance degradation, similar to 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 to the performance degradation.

In fig. 7 we measure the fraction of dead ReLUs directly across different settings and scaling factors. We find that large initial updates do indeed lead to a 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, as 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 of samples at the beginning of training. The performance of 92 . 1 can be recovered in this case. Additionally, we explored stable rank measurements as a potential factor but did not find a significant relation, as detailed in fig. 8.

B Additional Mathematical and Technical Details

B.1 The magnitude of the Momentum Vector

Let's assume a scalar gradient g t (e.g. for some coordinate) that is a random variable that is independent across time and has a zero mean distribution that does not change across time, i.e. E [ g t ] = 0

0

200

400

600

800

1000

0.0

0.2

0.4

0.6

inactive activations

steps

S-fb 1

S-fb 8

S-fb 128

R-fb 1

R-fb 8

R-fb 128

R 1

R 8

R 128

S 1

S 8

S 128

0

200

400

600

800

1000

0.02

0.04

0.06

mean stable rank

steps

1 S (fb)

8 S (fb)

128 S (fb)

1 R (fb)

8 R (fb)

128 R (fb)

1 R

8 R

128 R

1 S

8 S

128 S

and E [ g 2 t ] = σ 2 . For standard heavyball-momentum m t , with m 0 = 0 and coefficient β (equivalent to β 1 for Adam) we have:

̸

̸

In the limit t → ∞ we have (1 -β 2 t ) → 1 . We can derive the size of the second-moment v t in AdamW in an analogous way, obtaining E [ v t ] = (1 -β t 2 ) σ 2 . For a random walk, the update size of Adam is scaled in a similar way. Since the update size of Lion is fixed and does not depend on β , we scale the update size to match that of AdamW for a random walk in a steady state, i.e. by γ = √ 1 -β 1+ β as seen in algo. 2.

For Nesterov momentum , the update is modified to use:

Note that m t -1 and g t are independent and zero-mean, allowing us to use the previous result for:

In the limit t →∞ this gives the Nesterov scaling factor used in LionA (algo. 2) to ensure that the update size corresponds to that of AdamW using an analogous Nesterov update.

Inverse bias correction for momentum. Adam uses a bias correction to attempt to 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 as a result, similar to applying the same update multiple times. Removing the β 1 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 so 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 ( γ in algo. 3) to vary over time:

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 c > 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 -2 . For normalized optimizers like Adam and Lion, where the update norm is not affected by the gradient magnitude, this factor is decreased to c -1 .

B.3 The Angular Update Size in LionAR

The scaling factor for the angular update size in algo. 3 is adopted directly from the AdamW value derived by Kosson et al. [19]. 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 as the RMS update norm we derived in appx. B.1.

B.4 Relative Representation Change for Normalized Gradient Descent

Property (P1): For a dot product y = ⟨ w , x ⟩ and loss L ( x b ) that depends on y , we have:

where ∂ L ( x b ) ∂y is a scalar, ensuring that g b := ∂ L ( x b ) ∂ w ∥ x b , assuming the vectors are not zero.

Lemma (L1): Consider two independent random vectors a ∈ R C and b ∈ R C , 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 a and b is given by:

Proof : Let a = ( a 1 , a 2 , . . . , a C ) and b = ( b 1 , b 2 , . . . , b C ) . The inner product ⟨ a , b ⟩ is given by:

We need to find E [ ⟨ a , b ⟩ 2 ] . Expanding the square of the inner product:

Taking the expectation, we get:

Since a and b are independent and their elements are IID, we have:

Given that at least one of a or b has a zero-mean distribution, say a without loss of generality, we have E [ a i ] = 0 . Thus:

Consider two cases:

When i = j :

̸

When i = j :

So, the expectation simplifies to:

Since a i and b i are IID, we have:

Therefore:

Recognizing that:

we have:

Thus:

This completes the proof.

Assumption (A1): We assume the following vector pairs satisfy L1: ( x i , ˜ g b ) when i = b , (¯ g , ˜ g b ) and ( w , x b ) .

̸

Vector pairs of the type ( x i , ˜ g b ) and (¯ g , ˜ g b ) should be independent and ˜ g b has a zero mean distribution. However, the elements of each vector are not necessarily IID. For ( w , x b ) , this is an even stronger assumption. Generally, neither w nor x b is guaranteed to be IID or zero mean, and their independence later in training does not necessarily hold. Applying weight standardization to w or batch normalization to x would suffice to make their mean zero. Overall, this assumption 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 , x b ⟩ throughout training and find that it is more complicated than assumed here. This will lead to additional factors that may affect the RRC but we do not attempt to analyze.

Deriving the Relative Representation Change: Applying L1 directly gives us the original expected square output :

For the expected square representation change we get:

̸

̸

(34)

where we have used the definitions from eq. (7) and D1. Using property P1, we can write:

Plugging this into the previous expression yields E [(∆ y b ) 2 ]

̸

Squaring the expression results in various cross but all remaining dot products except the sign one are zero in expectation (due to the noise ˜ g ) and independent from each other. The cross terms involving these thus all disappear under the expectation. We apply Lemma L1 to their squares and approximate the expected norms of x b and g b as being independent. This gives E [(∆ y b ) 2 ]

We can compute the expected magnitude of the batch gradient as:

and similarly

Combining this with the previous expression for E [ y 2 b ] and the definition (D2) of the signal-to-noise ratio φ := E [ ∥ ¯ g ∥ 2 ] / E [ ∥ ˜ g b ∥ 2 ] we obtain the expression in the main body:

B.5 Estimating the Signal-to-Noise Ratio

We use accumulation over the microbatches to estimate the SNR at a given time. Let's assume we have A microbatches of size M each, with the average gradient of a microbatch denoted g m and the average gradient of the whole batch denoted g = 1 A ∑ m g m .

We estimate the variance of the norm of a single gradient example, i.e. the noise power as:

The signal power is estimated as:

Our SNR estimate is then:

B.6 RRC Correction Factor

The RRC correction is done based on eq. (9) 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:

For numerical purposes, we clamp 1 ≤ ρ ≤ B which corresponds to φ = 0 and φ = ∞ for a large C → ∞ . The update scaling factor is the square root of an EMA of the inverse of this quantity. We use the same coefficient as 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, so scaling the gradient at each step size can not 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 AdamW run the run-to-run differences in the validation loss over different seeds are around 0.05. However, the relative ranking of different runs remained the same.

B.8 Computational Requirements

Our experiments are performed on A100 GPUs with either 40GB or 80GB of RAM. One training run for our GPT2 setup takes around 4h, running on a single GPU. Reproducing the GPT2 experiments reported in the main body should take on the order of 1000 GPU hours. Including our preliminary experiments brings this up to around 3x this amount.

C Additional Experiments

C.1 Comparison with RAdam

RAdamW combines the variance reduction technique of Liu et al. [23] 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 as a good comparison and contextualization of our work. Figure 9 shows that while RAdamW outperforms the AdamW baseline without warmup, it is unable to match longer warmups. Ma and Yarats [26] suggests that RAdamW is approximately equivalent to 4 steps of SGDMfollowed by AdamW with a special type of built-in warmup with an effective length of around 2 / (1 -β 2 ) = 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.

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 -norm of the update. This relates our approach of looking at the update size to the adaptive learning rate view of Liu et al. [23]. Secondly, we note that counteracting noisy estimates of the second moment can 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 at all, indicating there are additional factors that contribute to the need for warmup.

0

2500

5000

Iteration

0

1

Learning Rate Schedule

0%

2%

5%

10%

20%

10

3

10

2

Learning Rate

3.1

3.2

3.3

3.4

3.5

3.6

Final Validation Loss

AdamW

RAdamW

0

2500

5000

Iteration

0

1

Learning Rate Schedule

0%

2%

5%

10%

20%

10

3

10

2

Learning Rate

3.1

3.2

3.3

3.4

3.5

Final Validation Loss

AdamW

RAdamW

C.2 Model & Dataset Ablations

In this section we repeat some of our main experiments, varying the dataset and model architecture.

10

3

10

2

Learning Rate

3.0

3.1

3.2

3.3

Final Validation Loss

LionA

0%

5%

10%

10

3

10

2

Learning Rate

3.0

3.1

3.2

3.3

Final Validation Loss

LionAR

0%

5%

10%

10

3

10

2

Learning Rate

3.0

3.1

3.2

3.3

Final Validation Loss

LionAR

=0.98 +N +MC

0%

5%

10%

10

3

10

2

Learning Rate

3.0

3.1

3.2

3.3

Final Validation Loss

LionAR + RRC comp

0%

5%

Figure 10 shows the effects of changing the dataset used in our experiments from OpenWebText [8] to SlimPajama [33]. Overall the results are similar as before, when controlling the ℓ 2 -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.

10

3

10

2

Learning Rate

3.1

3.2

3.3

3.4

3.5

Final Validation Loss

LionA

0%

5%

10%

10

3

10

2

Learning Rate

3.1

3.2

3.3

3.4

3.5

Final Validation Loss

LionAR

0%

5%

10%

10

3

10

2

Learning Rate

3.1

3.2

3.3

3.4

3.5

Final Validation Loss

LionAR

=0.98 +N +MC

0%

5%

10%

Figure 11 shows the effects of changing the architecture from GPT2 to a 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 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. We expect that using larger batch sizes could necessitate these additional tricks for Llama, but do not explore this further here.

10

3

Learning Rate

2.8

2.9

3.0

3.1

Final Validation Loss

LionA

0%

5%

10

3

10

2

10

1

Learning Rate

2.8

2.9

3.0

3.1

Final Validation Loss

LionAR

0%

5%

10

3

10

2

10

1

Learning Rate

2.8

2.9

3.0

3.1

Final Validation Loss

LionAR

=0.98 +N +MC

0%

5%

Figure 12 uses a larger Llama model with twice the depth. It also trains for twice as many iterations to keep the ratio of tokens to parameters similar. The overall results resemble those of the smaller llama experiments in fig. 11.

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 to 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 in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

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

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

2. Limitations

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

Answer: [Yes]

Justification: We include a limitation section §9 which discusses the main overarching limitations from our perspective.

Guidelines:

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

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

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

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

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

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

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

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

3. Theory Assumptions and Proofs

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

Answer: [Yes]

Justification: Our work is mostly empirical, aside from the analysis of the relative representation 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 crossreferenced.

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

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

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

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

4. Experimental Result Reproducibility

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

Answer: [Yes]

Justification: We provide the main hyperparameters and other aspects of our training. We are not aware of 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, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

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

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

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

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

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

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

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

5. Open access to data and code

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

Answer: [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 NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

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

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

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

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

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

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

6. Experimental Setting/Details

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

Answer: [Yes]

Justification: We 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 of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

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

7. Experiment Statistical Significance

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

Answer: [No]

Justification: Our main results are qualitative and computationally expensive for us, so we do not estimate the variance for each point in 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 error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

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

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

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

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

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

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

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

8. Experiments Compute Resources

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

Answer: [Yes]

Justification: We 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 of compute required for each of the individual experimental runs as well as estimate the total compute.

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

9. Code Of Ethics

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

Answer: [Yes]

Justification: We read it and believe we fully conform to it.

Guidelines:

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

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

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

10. Broader Impacts

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

Answer: [NA]

Justification: This paper is focused on understanding deep learning in general, and does not have an specific societal impacts beyond those of the field overall.

Guidelines:

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

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

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

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

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

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

11. Safeguards

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

Answer: [NA]

Justification: We 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 a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

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

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

12. Licenses for existing assets

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

Answer: [Yes]

Justification: We 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 or dataset.

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

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

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

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

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

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

13. New Assets

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

Answer: [NA]

Justification: 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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

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

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

14. Crowdsourcing and Research with Human Subjects

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

Answer: [NA]

Justification: 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 in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

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

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

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

Answer: [NA]

Justification: No crowdsourcing / human subjects used

Guidelines:

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

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

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

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

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nSee Weakness.\n\nLIMITATIONS:\nI 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nSee weaknesses", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\nWarmup 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nCan 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.\n\nLIMITATIONS:\nPartially", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nTo 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nI 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-26"}