Early Weight Averaging meets High Learning Rates for LLM Pre-training

Summary

This paper proposed a method to do weight averaging along the LLM pretraining trajectory. The experiment results show that the obtained checkpoints are better than the original checkpoint (with the fixed training steps) in terms of validation loss and zero-shot performance on downstream tasks.

Rating

7

Confidence

4

Ethics flag

1

Reasons to accept

+ The proposed method is simple and easy to implement, and it costs relatively little extra computation. It can be used for any LLM's pretraining. + To the best of my knowledge, this work is the first to use the technique of weight averaging **during pretraining**, which IMO is a kind of novelty in some sense. + The intuition discussed in Section 2.2 makes a lot of sense to me. + The experiment results are rich. The proposed method seems to be effective as it improves both the perplexity and the zero-shot performance of LLMs, especially for models trained with large learning rates. Its effect in mitigating loss spikes is also interesting. + Appendix A discusses some FAQs, which were helpful when I finished reading the main text.

Reasons to reject

+ The writing of some parts could be improved. As a person who is not familiar with LAWA/EMA/SWA, I wrongly assumed the proposed method was to add some checkpoints obtained by weight averaging during pretraining when I just finished reading Section 2.3. And finally, it took me some time to understand the proposed method does not influence the checkpoints used for training. + Even though I love the intuition in Section 2.2, they may not be convincing enough for some readers: + The `Optimization Viewpoint` seems to hold for the toy setting, but how about the much more complicated setting of large language models? + The `Diversity Viewpoint` talks about ensembling, but when the model training is at the early stage the models are not strong enough (i.e., they may not be models that have some could-be-essential capabilities, which are much more common in previous ensembing works). In this case, how can the intuition of ensembling explain that? + Even though the authors argue that (1) ` where conducting a grid search is challenging due to the model’s size, adopting our proposed training recipe could be advantageous` and (2) `it is safe to say that LAWA provides some gains for all LRs but works the best with high LRs`, I still have the concern that the proposed method relies too much on the assumption that the model is pretrained with a large learning rate. + As previous works already pointed out that we can use a small LR (e.g., Sophia), why should we assume that people in the future have to do an extensive grid search for the pretraining learning rate? + Then, the improvement seems to be very marginal with large LR, e.g., Figure 2(a). + Fortunately your proposed method does not bring (too much) degradation in all the experiments conducted, and then you may argue that "we can always adopt the method in any situation". But the mentioned facts show that the effectiveness is kind of limited, so you may have to get a stronger argument regarding the learning rate. + I noticed some discussions on hyper-parameters in Appendix C.2 but they are just too empirical. I hope to see a more intuitive discussion or explanation so it can guide people who use the proposed method in the future. However, it is not a major concern because people could be able to tune the hyperparameters with such a lightweight method.

Questions to authors

As I said in Reasons To Reject, I had a misunderstanding about your proposed method. I am curious about how the authors think about the method that I was wrongly thinking of - the checkpoints obtained by weight averaging are also added to the training trajectory (i.e., such checkpoints continue to be trained).

Summary

This paper studies early weight merging to speedup the LLM pretraining. Authors found that models trained with high learning rates observe higher gains due to checkpoint averaging. Authors proposed one simple but effective strategy and surpasses popular checkpoint averaging baselines such as exponential moving average (EMA) and stochastic moving average (SWA) when training nanoGPT-2 models from 125M to 770M.

Rating

7

Confidence

4

Ethics flag

1

Reasons to accept

1) The paper is very well written. The Section 2 is friendly to the readers who are not knowledgeable in weight averaging. 2) The results are promising under authors' current setting. It looks like there is a clear convergence speedup at the early stage of training.

Reasons to reject

1) It would be great to discuss more about the difference between this approach and existing approaches. For instance, can I say it simply change the distribution of EMA and stochastic weight average? 2) It seems that the improvement is very small at the late stage of training. Maybe the reason is that the averaging actually provides a smaller learning rate in practice, which learns the things that will be learned by smaller LR. If so, I have a concern about whether the approach will really work when training with a very large dataset like 15T tokens from llama3. 3) The experimental results about the real throughput are missing. Loading/merging the checkpoints would also be time-consuming if we train very large models on the GPU cluster. Would the speedup benefits diminish under this type of setting because of lower GPU utlization?

Questions to authors

See above.

Summary

This paper studies the use of LAWA (latest weight averaging; Kaddour, 2022) -- a type of checkpoint averaging technique -- in the context of pre-training language models. The key insights discussed by the paper is that (1) Models trained with high learning rates particularly benefit from checkpoint averaging at the early stage of training. The loss trajectories of such models mimic those of models trained with low learning rates. Thus, checkpoint averaging can be used as a technique to improve high-LR training's generalization. (2) For LAWA and high-LR training, using distant checkpoints performs better. The authors conducted thorough experiments, studying models with different LRs, different scales, and different training data. The evaluation is done on both perplexity and zero-shot downstream task performance. The experiments show that LAWA is better than other methods such as EMA (exponential moving average; Szegedy et al., 2015) and SWA (stochastic weight averaging; Izmailov et al., 2018). LAWA also leads to faster convergence at the early stage of training.

Rating

4

Confidence

4

Ethics flag

1

Reasons to accept

(1) The experiments in the paper are very thorough, spanning across different models, different scales, different data, and different evaluations. The proposed training paradigm leads to faster convergence in high-LR trainings at the early training stage. (2) The two insights -- high-LR training + checkpoint averaging helps improve generalization, and LAWA should use larger intervals -- are effective and are verified by thorough experiments.

Reasons to reject

(1) Novelty: the main method, LAWA, is not new. The use of LAWA in high-LR training and using larger interval for LAWA (the two key insights in the paper) seem to me a natural applications (and hyperparameter tuning) of the original LAWA method. (2) From the main results, it seems that the proposed training method at most can match the standard low learning rate training (Figure 2). Please let me know if I missed something here (I assume in most experiments, you control the learning rate to be the high one, except in Figure 2). Now the authors can argue that high learning rates are useful in (a) large batch size scenario and (b) very short training. However, there are no experiments on (a) in the paper (batch sizes are all fixed?) and for (b), it seems that the gain is only significant at very very early training stages (from Figure 2). Note that compared to nowadays pre-training scale (1T tokens), the experiments are already quite toy (10B, 207B tokens), thus "an early stage training" of this setting may have little empirical benefit implications.

Questions to authors

Please see my "reasons to reject", and let me know if I misunderstand anything!

Summary

The paper investigates the effect of checkpoint averaging along the trajectory of a model training. It adapts a previously proposed method - LAtest Weight Averaging (LAWA), for decoder-only LLMs. The authors were able to train models with higher learning rate to observe gains in validation loss and some downstream tasks. They also adapt exponential moving average (EMA) and stochastic moving average (SWA) for their setup and compare the methods.

Rating

4

Confidence

4

Ethics flag

1

Reasons to accept

- The method of model weight averaging with higher learning rate is simple and shows improvement in validation loss and some downstream tasks. - The paper is easy to follow, and compares with some relevant methods.

Reasons to reject

- The methods shows improvement in early stages of training, and the difference gets increasingly smaller at later stages. It is seen in both loss value and downstream task performance. This raises the question: will the conclusions still hold for longer training regime such as llama family of models? - The experimental setup is limited, and there might be some issue in the GPT2 training experiment. Generally, we expect the loss to get smaller as we increase the model size. However, GPT2-large has worse val loss than GPT2-medium. - The baseline setup is weak. Training without any averaging can improve further with better optimized batch size and learning, and they're correlated. Higher learning rate along with larger batch size can lead to better performance. e.g. llama 7b was trained with 4M batch size and LR=3e-4 - The EMA and SWA adaptation requires corresponding update to averaging interval similar to LAWA. - The cost of hyperparameter search for LAWA should be discussed when comparing with the baseline.

Questions to authors

- How was the moving window interval and number of model checkpoints to average decided? what is the corresponding cost? does it vary with model size and architecture? - It'll be good to add learning rate schedule for different training methods.

Let us know if we have addressed your concerns

Dear Reviewer GUxG, Since we are very close to the end of the discussion period, we are eager to know if our responses have satisfactorily addressed your concerns regarding our paper. We have also provided additional results to address your concerns about EMA and SWA baselines. In case we you feel we have sufficiently addressed your concerns please reconsider your score.

We added additional experiments please provide us with some additional feedback

Dear Reviewer ZNy2, Since we are very close to the end of the discussion period, it is a gentle reminder that please engage with us for discussing our comments and newly added results. We worked very hard to generate 4 new figures of additional experiments. Let us know your thoughts and if you feel we have sufficiently addressed your concerns we invite you to improve your scores.

We appreciate your feedback and invite you for further discussions

Dear Reviewer rseH, Thanks for taking some time to review our work. We feel your thoughtful remarks have improved the quality of our work. Since we are close to end of discussion period we would love to get some additional feedback on the comments we made during the rebuttal. We also invite you to update your scores only if you believe we have sufficiently addressed your concerns.

Invitation for further discussions

Dear Reviewer BhE4, Thank you for your thoughtful remarks. We have added additional experiments to address your remark on distribution of EMA and SWA. We would be very grateful if you may engage in further discussions with us. Meanwhile we if you belive we have addressed your concerns well enough please consider improving your scores.

Ack

I acknowledge that I have read the author's rebuttal; however, I still stand for my original evaluation and thus will keep the score unchanged. My main concern is still how novel the proposed method is compared to the original LAWA (even with the updated result; the gain is pretty small) and also the application of the finding.