Analyzing & Reducing the Need for Learning Rate Warmup in GPT Training

Summary

This work explores the benefits of learning rate warmup in neural network training, focusing on the size of model updates via the GPT2 model. It finds that controlling update size in parameter space doesn't fully explain warmup's advantages, but quantifying updates in terms of neural representation changes shows promise. The study also highlights the role of high momentum in warmup and suggests potential methods for reducing the need for manual warmup configuration. Overall, the research provides insights into learning rate warmup's necessity and potential ways to eliminate it in practice.

Strengths

The paper addresses an intriguing topic, which aims to present a systematic understanding regarding the LR warmup heuristic from a novel perspective. However, I feel that the authors have attempted to cover too many aspects, which might be challenging to thoroughly demonstrate within the scope of a single conference paper.

Weaknesses

1. I noticed that the authors have not adequately discussed the highly relevant paper, "On the Variance of the Adaptive Learning Rate and Beyond," which addresses some of the questions raised by the authors. Please discuss the unique contributions of your work compared to the variance-based analysis presented in that paper. 2. 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. 3. The authors use linear transitions to analyze the representation changes, which seems too toy for me. 4. 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? 5. 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. 6. 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? 7. 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.

Questions

See Weakness.

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

I have not found any discussions about the limitations and potential negative societal impact. But in my opinion, this may not be a problem, since the work only focuses on analyzing the warmup heuristic in machine learning. Still, it is highly encouraged to add corresponding discussions.

Summary

The submission analyzes the underlying reason behind the need for a learning rate warmup in neural network training, focusing on GPT pre-training with AdamW and Lion optimizers. The authors identify three key reasons as to why the initial updates are large: 1. Momentum handling by AdamW, 2. Early updates not correlated with the initial weight magnitudes 3. Correlation between gradients of examples during early training The 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.

Strengths

* This paper analyzes various metrics that correlate with the benefit of learning rate. * The analysis of the normalized Gradient Descent is insightful and reproduces previously known scaling laws.

Weaknesses

* The authors begin by analyzing warmup for Adam and instead of directly modifying Adam, they modify the Lion optimizer. Direct modifications to the Adam optimizer would be more convincing and then moving to Lion would streamline the arguments. * 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. * 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.

Questions

* Which dataset is used for training the model? I don't think it is mentioned anywhere in the paper and its important for reproducing the results. * 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? * Did the authors try Adam with inverse bias correction for the momentum as suggested by equation 1?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

4

Limitations

See weaknesses

Summary

In this paper the authors investigate the performance benefits seen from the common practice of learning rate warmup and scheduling, attempt to understand the mechanistic underpinnings of those improvements, and engineer optimizers that mitigate the need for warmup. They conduct experiments using NanoGPT and consider controlling parameter updates, angular updates, and “relative representation” changes to close the gap between warmup and no warmup.

Strengths

Warmup length and peak learning rate are certainly some of the most important hyperparameters in large model training, and eliminating the need for a warmup phase would present a significant simplification to training. The paper is well motivated and systematic in its investigation of warmup and proposals to sidestep the necessity of warmup. The RRC is an interesting and promising angle on this question.

Weaknesses

The results do not suggest a clear prescription for learning rate scaling or straightforward changes that can be made to initializations or updates. In particular, the RRC is completely dependent on the inputs, but there does not seem to be any discussion or investigation of the effects of the input data. For the NanoGPT experiments there is no mention of what the data is. Presumably training was done with a cross entropy loss, but this is also not mentioned.

Questions

Can the authors provide some discussion about the sensitivity of RRC to variance in the inputs across and within batches? It is also not clear what direction the RRC results are suggesting. An “automatic warmup” that scales the update sizes according to online measurements of the signal-to-noise ratio is still a warmup phase, albeit a more principled way to arrive at what that schedule should look like. To be clear, I don’t think this is a bad thing, but it may be more representative of the results to propose an “adaptive” or “automatic” learning rate scheduler, rather than claim to make progress towards eliminating the need for scheduling.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

Partially

Summary

To train current deep neural network architectures, especially transformers, the learning rate of AdamW is usually first linearly increased to reach a peak before it's decreased to zero. The paper analyzes the impact of this so-called warming-up phase on GPT-2 models from the perspective of the update size. As a second contribution, the paper presents some small modifications for the Lion optimizer to mitigate some of the issues encountered in the experiments.

Strengths

- The paper is very well written. The experiments are nicely motivated and reasonable. - Warming up the learning rate is arguably common practice for training transformer models, but not well understood. The paper provides some interesting analysis of the matter, which could potentially lead to a more intuitive understanding of the problem and eventually better optimizers.

Weaknesses

- The results of the paper are somewhat inconclusive and after reading the paper, I am still not sure about the dynamics during the warming up phase. For example, while controlling angular updates seems to stabilize training to a certain degree, it eventually doesn't lead to better performance. Also, as the paper clearly states, the magnitude of the parameter updates doesn't really account for the gains of the warm-up phase. I am wondering if the paper approaches the problem actually from the right perspective. Having said that, I think the paper still provides some value and might help to stir future research. - 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.

Questions

- How sensitive are the results from Section 3 to the type of learning rate schedule? For example, how would Figure 1 look if you used, let's say, a cosine annealing schedule?

Rating

7

Confidence

2

Soundness

3

Presentation

4

Contribution

3

Limitations

I think the paper spells out all its limitations; however, for visibility, it might be better to move the corresponding paragraph from the appendix to the main text. I don't see any negative societal impacts of the work.

reply to authors

I thank the authors for addressing my comments. I will raise my score and vote for acceptance of the paper