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

Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-21 23:38:07

Abstract: Training Large Language Models (LLMs) incurs significant cost; hence, any strategy that accelerates model convergence is helpful. In this paper, we investigate the ability of a simple idea – checkpoint averaging along the trajectory of a training run – to improve both convergence and generalization quite early during training. Here we show that models trained with high learning rates observe higher gains due to checkpoint averaging. Furthermore, these gains are amplified when checkpoints are sampled with considerable spacing in training steps. Our training recipe outperforms conventional training and popular checkpoint averaging baselines such as exponential moving average (EMA) and stochastic moving average (SWA). We evaluate our training recipe by pre-training LLMs, where high learning rates are inherently preferred due to extremely large batch sizes. Specifically, we pre-trained nanoGPT-2 models of varying sizes—small (125M), medium (335M), and large (770M)—on the OpenWebText dataset, comprised of 9B tokens. Additionally, we present results for publicly available Pythia LLMs, ranging from 1B to 12B, which were trained on the PILE-deduped dataset containing 207B tokens.

Main paper: 
 introduction: 
 Large Language Models (LLMs) have made a significant leap from billion to trillion scale, both in terms of parameters (Chowdhery et al., 2022; Ren et al., 2023) and pre-training data size (Hoffmann et al., 2022) Touvron et al. (2023a; b) . This surge in both data and model size has rendered LLM pre-training increasingly slow and resource-intensive. For instance, a Llama 2 70B model trained with 2T tokens took 1720K GPU hours to train. To accelerate the training process, it is a popular practice in LLM pre-training (Biderman et al., 2023; Touvron et al., 2023a) to utilize exceptionally large batch sizes, thereby ensuring maximal GPU utilization and consequently necessitating the use of relatively high learning rates (Goyal et al., 2017) .In this paper, our goal is to improve the test generalization (log perplexity) of LLM pre-training while reducing the number of training steps, all without increasing the compute budget. To achieve this, we first demonstrate that: (a) models trained with higher learning rates exhibit greater improvements when averaged along the training trajectory, and (b) averaging distant model weights from a single training trajectory further amplifies these gains. We integrate these two insights to adapt LAWA (LAtest Weight Averaging) (Kaddour, 2022 )-a technique that performs checkpoint averaging throughout a training trajectory using a sliding window-for pre-training LLMs.We evaluate our methodology by pre-training nanoGPT-2 models of various scales, specifically 125M (small), 355M (medium), and 770M (large), using the OpenWebText dataset, which comprises 9B tokens. The experiments with nanoGPT-2 are conducted in a controlled environment to gain a deeper understanding of our training recipe. Furthermore, we extend our evaluation to publicly available Pythia LLMs (Biderman et al., 2023) , which include model sizes of 1B, 2.8B, 6.9B, and 12B, trained using 207B tokens. Our experiments with Pythia LLMs aim to demonstrate the impact of our work on real-world LLMs.Main Contributions. In summary, our findings are as follows, 2 ). This gain further amplifies when we sample distant checkpoints in the training run (Figure 10 ). We provide a intuitive explanation of this phenomenon in Section 2.2. We observe that the training trajectory of LAWA closely mimics that of a model being trained with a low LR. The primary advantage of LAWA is that it allows LLMs to be trained with high LR without compromising generalization (Section 4).3. We show that LAWA improves test generalization with fewer training steps compared to original training starting very early on during training; for both nanoGPT-2 and Pythia LLMs (Figure 1 and Figures 4 5 6 7 8 ). For Pythia checkpoints we perform zero-shot evaluation and observed similar gains, as shown in Table 2 .4. We compare our recipe with conventional training and popular baselines such as Exponential Moving Average -EMA (Szegedy et al., 2015) and Stochastic Weight Averaging -SWA (Izmailov et al., 2018) . These baselines were not originally proposed or evaluated for LLM pre-training, but we adapt them to set meaningful baselines. Our training recipe outperforms conventional pre-training, EMA, and SWA.5. Additionally, we perform a preliminary investigation of early weight averaging for a diffusion model for image generation (specifically, a 422M sized UNet model trained with the standard DDPM objective (Ho et al., 2020) ). We observe thematically similar improvements (evaluated by the FID metric) as shown in Figure 9 .Paper outline. The remainder of this paper is structured as follows: We start by expalining the problem through a toy example and intuition behind our approach in Section 2. Following that we detail our experimental set up in Section 3 and present our main results in Section 4. Thereafter we discuss our work relative to prior works in Section 5. Finally, we conclude the paper with a summary of our conclusions and potential directions for future research. 
toy setting: 
 We explain the setting using a simple toy problem of minimizing a two-dimensional loss function, represented as L(w 1 , w 2 ), where w 1 and w 2 are the parameters of the model. In this scenario, there exists an optimal batch size, B 0 , and an optimal learning rate, η 0 , that minimizes the loss function. rate, η, such that both are significantly larger than their optimal counterparts, i.e., B ≫ B 0 and η ≫ η 0 respectively. Suppose also that the loss function L(w 1 , w 2 ) exhibits a much higher curvature along w 1 , as compared to w 2 . It is widely known that the updates from the AdamW optimizer are mostly uniform across all weight dimensions (Liu et al., 2023) . When η ≫ η 0 the weight updates of w 2 will be accelerated, however this will cause oscillations along w 1 which in a long run adversely affects the convergence of w 1 . A naive approach to mitigate this problem is to use a smaller LR or decay LR to 0, which might hinder progress in flatter regions. It is conceivable that a LLM might exhibit extremely heterogeneous curvatures exacerbating this issue during pre-training. Optimization Viewpoint. We propose performing checkpoint averaging of model weights relatively early during training with high learning rates (η). The rationale behind this step stems from the fact that checkpoint averaging serves as a surrogate to LR decay, as demonstrated by Sandler et al. (2023) . However, this surrogate LR decay is decoupled from the weight update during optimization, as checkpoint averaging is conducted in a post-hoc manner.Employing this simple technique, we mitigate the oscillations in w 1 while swiftly traversing through w 2 , achieving enhanced generalization in fewer training steps as illustrated in Figure 3 .Diversity Viewpoint. The practice of averaging the weights of model checkpoints is broadly recognized as being functionally analogous to ensembling (Izmailov et al., 2018; Wortsman et al., 2022) . In model ensembling literature it is well established that diverse models improve the performance of the ensemble (Lakshminarayanan et al., 2016) . Therefore, it is fair to assume that this principle also applies to model averaging as well. In our context, we define the diversity of a model at two distinct training steps,1 N N i=1 1 y W1 i ̸ = y W2 iwhich calculates the number of disagreements between the two checkpoints. This equation computes the number of samples from the same held-out set where the check-point W 1 disagrees with checkpoint W 2 . A recent study by Athiwaratkun et al. has demonstrated that a higher LR can result in the generation of diverse model checkpoints. We observed (Figure 1 ) that this phenomenon can be further amplified by sampling far apart checkpoints in terms of training step. We combine both these insights to induce diversity in our checkpoints.We explain the Latest Weight Averaging (LAWA) algorithm below along with a python-style pseudo code (Algorithm 1). As shown in Figure 3 , LAWA maintains a first in first out (FIFO) queue of periodically sampled checkpoints with a large number of intervening steps (ν) in between two succesive samples. We adapt LAWA for our setting with minor modifications. Specifically, we introduce k stepsize (ν), and decoupled interval and k to effectively sample distant checkpoints in the training run. The original LAWA algorithm (Kaddour, 2022) assumes interval = k.LAWA runs a moving window at a predetermined interval to collect k latest checkpoints on sequence of saved checkpoints θ t . The LAWA derived checkpoints are computed as θ LAWA NanoGPT-2 Experiments. We conduct all our experiments utilizing autoregressive decoder-style Large Language Models (LLMs), specifically nanoGPT-2 and Pythia LLMs.We utilize three distinct sizes of nanoGPT-2 models: small (125M), medium (355M), and large (770M). We train nanoGPT-2 models from scratch using the OpenWebText dataset, which includes 9 billion training tokens and 4.4 million validation tokens. Throughout the experiments, we maintain a consistent sequence length of 1024 and a fixed batch size of 131K tokens per batch, the latter being the maximum batch size accommodated by our GPUs. The configurations for the model and pre-training were adapted from Sophia's (Liu et al., 2023) AdamW baseline, with adjustments made to the learning rate and batch size to align with our specific needs. Notably, we trained all the models with learning rates that were ten times higher and batch sizes that were twice as large compared to the configurations in Liu et al. (2023) , where the learning rate was tuned through a grid search. We compare LAWA with the original pre-training recipe, EMA (Szegedy et al., 2015) , and SWA (Izmailov et al., 2018) , which we adapt for LLMs. For EMA, we set the decay to 0.9 as per Kaddour (2022) and update the EMA model at every step, which is a standard practice. For SWA, we adhere to the original pre-training procedure until 75% completion, after which SWA training is initiated with a new SWA scheduler (cosine annealing). We compute the SWA uniform average every 10 steps.Pythia Experiments. The Pythia LLMs are publicly available in the Pythia suite (Biderman et al., 2023) . We report results on Pythia-1B, Pythia-2.8B, Pythia-6.9B, and Pythia-12B; Table .1 summarizes the details of these models. For our experiments, we use intermediate model checkpoints; saved after every 1000 update steps. The models are trained by Biderman et al. (2023) on the PILE dataset (Gao et al., 2020) , a publicly available, curated collection of English text corpus of size 800GB. The original PILE dataset is curated using 5 different genres of data namely, academia, internet, prose, Evaluation We evaluate the language modelling performance of nanoGPT-2 models pre-trained for 70K steps using log perplexity (perplexity and loss used interchangeably) on the held-out/val set. For nanoGPT-2 we use the moving window interval = 1K, k = 5 and we sample checkpoints interval = 200,1K apart for LAWA. Next we analyze the original training trajectories of Pythia LLMs and demonstrate the improvements achieved in test generalization using LAWA. We also present zero-shot evaluation results on Lambada OpenAI (Paperno et al., 2016) , SciQ (Welbl et al., 2017) , AI2 Reasoning Challenge-easy (ARC-e) (Clark et al., 2018) , and Wikitext (Merity et al., 2016) . We evaluate 4 Pythia LLMs using the intermediate model checkpoints on a subset of the test and validation dataset following the methodology prescribed by Xia et al. (2022) . For the purpose of evaluation we use the open source library lm-evaluation harness2 . We select representative subsets from the diverse genres encompassed by the full PILE validation and test dataset. This subset comprises PILE-philosophy papers, PILE-bookcorpus2, and PILE-YouTube subtitles datasets.We conduct zero-shot evaluation on the Pythia LLMs. For zero-shot evaluation we provide a natural language description of the downstream task, along with a textual example. The models then generate responses that are either open-ended or discriminatively select a proposed answer. This evaluation setup serves as a robust academic benchmark, as it assesses Pythia models of various scales on reasonably large PILE subsets and downstream datasets, both in terms of test performance and zero-shot. For both the test generalization and zero-shot experiments, we evaluate model checkpoints starting 21K steps to 141K steps (recall that subsequent checkpoints are 1K steps apart). Moreover, we choose to slide the averaging window at 3K steps (i.e. interval = 3K) and average last k intermediate checkpoints as discussed in LAWA algorithm (Algorithm 1). Our selection of LAWA parameters such as k = 5 and start step = 21K are based on the experiments discussed in Section A.1. 
exploring llm pre-training with nanogpt-2 at small scale: 
 LLMs trained with higher LRs observe higher gains with LAWA. We ran controlled experiments to better understand the correlation between LR and gains due to checkpoint averaging. Initially we train nanoGPT-2 small with two different LRs (6 × 10 -3 , 6 × 10 -4 ), keeping batch size and all relevant hyperparameters the same. 6 × 10 -4 is the assumed optimal LR computed using a grid search reported in Liu et al. (2023) . Subsequently, we pre-trained the same model using an LR of 6 × 10 -3 , which is tenfold higher in magnitude than the former. As shown in Figure 2 (a), models Figure 5 : LAWA speeds up convergence for Pythia-2.8B on subset of tasks from the original pretraining dataset i.e. PILE.. We present the original and the LAWA training trajectories for 3 different tasks from PILE namely philpapers, bookcorpus2 and youtube subtitles. trained with higher LR observe higher gains compared to its counterpart trained with lower LR due to post hoc checkpoint averaging in LAWA. From Figure 2 (b) we observe that the model trained with a higher LR converges faster but compromises on generalization, a phenomenon also observed by Kaur et al. (2022) . The gap in generalization is effectively mitigated by checkpoint averaging through LAWA, as shown in Figure 2(c ). Interestingly, we note that the training trajectory of LAWA approximates that of a model trained with a lower LR. This is an important insight: checkpoint averaging acts as a surrogate for LR decay, thereby enabling the model to be trained with a higher LR. In practical LLM pre-training scenarios, where conducting a grid search is challenging due to the model's size, adopting our proposed training recipe could be advantageous. One might select a higher LR (that doesn't cause divergence) and train an LLM faster without compromising much generalization compared to conventional pre-training strategy. LAWA improves test generalization in fewer training steps compared to original pre-training and relevant baselines. LAWA clearly outperforms the original pre-training run starting very early on during training, as shown in Figure 1 . Since we employed a reasonably high LRs for all the nanoGPT-2 LLMs, we observe higher gains in the early-mid stages of pre-training, and the gains start diminishing towards the final stages due to the LR scheduler that continuously decays the weight throughout the training cycle. Additionally LAWA also outperforms important baselines such as EMA and SWA. LAWA clearly has an edge over EMA throughout all training phases. Our experiments also reveal that applying SWA during the early stages of training leads to divergence (Figure 11 ). Consequently, LAWA outshines SWA in both performance and ease of implementation.The gains due to LAWA amplifies with far checkpoint averaging. As shown in Figure 1 , LAWA with higher k stepsize (ν) performs better particularly for larger models. Intuitively, we believe that the diversity between nearby checkpoints might be very low given that larger models learn faster (Li et al., 2020) . Hence, one needs to sample more distant checkpoints for larger models. This observation is also consistent with billion-parameter Pythia LLMs, as shown in Figure 10 .Figures 4 and 5 illustrate that the checkpoints derived using LAWA demonstrate better test generalization than the checkpoints saved during original training for the Pythia-1B and Pythia-2.8B models i.e. (moderate size LLMs). In Figures 7 and 8 , we observe significant improvements in test generalization during early-mid training regime and minor improvements towards the end for both Pythia-6.9B and Pythia-12B models. All the LAWA LLMs achieve lower perplexity with lesser training steps compared to the original training trajectory, thereby saving significant amount of GPU hours (Figure 6 ), subsequent training costs and ingested training data. The savings are computed based on Table 3 in the appendix. Additionally, LAWA proves beneficial in situations where training similar models from scratch is necessary but can only be conducted over a limited number of training steps due to strict compute budgets.To analyze the phenomenon wherein moderate size LLMs exibit higher gains compared to their larger counterparts in test performance across both early-mid and final training trajectories, we delve into the Pythia suite's training methodologies. The authors of Pythia suite (Biderman et al., 2023) report that they have used an exceptionally large batch size (2M tokens) and learning rate for Pythia-1B and Pythia-2.8B models to expedite the convergence. For the larger Pythia models, such as Pythia-6.9B and Pythia-12B, learning rates are reduced to 1.2 × 10 -4 while maintaining the batch sizes to 2M tokens, in line with prior work. Overall we observe thematically similar trends with the nanoGPT-2 experiments.Recent work (Kaplan et al., 2020; Zhang et al., 2022; Touvron et al., 2023a; Dey et al., 2023; Wortsman et al., 2023) report loss spikes-brief degradations in the performance along a training trajectory when scaling up the model size, batch size, and learning rate. In our evaluations, we observe two perplexity spikes (Figure 7 ). Interestingly, we find that LAWA mitigates the spikes during evaluation quite effectively. This can be intuitively explained by the smoothing effectaverage out the outliers to fit the trend -resulting from averaging checkpoint weights over a range Table 2 : LAWA improves zero shot performance on academic question answering and knowledge assessment downstream tasks starting very early on in the training. The checkpoints derived using LAWA requires less steps to reach higher zero-shot performance than the checkpoints derived using original training. We indicate the scores in bold font when the performance of LAWA surpasses the final score (at 141K steps) obtained using the original training or achieves comparable performance significantly earlier, specifically at 105K steps.of steps that are far apart. Note that since we have sampled the checkpoints at an interval of 3K for LAWA, we may have inadvertently overlooked some checkpoints exhibiting loss spikes in models other than Pythia-6.9B.LAWA improves the zero-shot performance in several ways. First, we observe that zero-shot performance of early-mid checkpoints (48K, 60K) achieves higher performance almost consistently, regardless of the scale as shown in Table 2 . For instance, the LAWA Pythia-1B checkpoint evaluated at 24K steps on the Lambada OpenAI task achieves higher accuracy than the original checkpoint evaluated at the 48K step.We note that the checkpoints derived using LAWA also exhibit improvements in the later stages of training (105K,141K) on the majority of tasks, highlighted by the bolded numbers in Table 2 . Moreover, we consistently witness gains until 105K steps across all models, which constitutes approximately 75% of the total training steps. Therefore our recipe proves to be beneficial in a compute optimal LLM training scenario where early stopping is employed at 75% of total training. Additionally, we find that our LAWA derived checkpoints of Pythia-6.9B reach the final accuracy/perplexity on multiple downstream tasks considerably earlier, specifically at the 105K step mark. Intuitively, we know that higher zero-shot performance on various different downstream tasks is associated with low perplexity in language modelling during training, a correlation that is mathematically substantiated (Saunshi et al., 2020) . Therefore, all the observations we made in Section 4.1 naturally apply to zero-shot performance as well.We also experiment with image diffusion models to gauge the effectiveness of LAWA on generative models beyond language. The underlying model is a 422M parameter UNet (Ronneberger et al., 2015; Ho et al., 2020) trained with ϵ-prediction objective and standard cosine schedule (Ho et al., 2020) optimizer using a learning rate of 1e-4. Figure 9 shows the FID on the evaluation set for the baseline checkpoints and LAWA averaging over the baseline checkpoints. Note that the baseline checkpoints themselves are obtained using the Exponential Moving Average (EMA) with decay rate of 0.9999 over the training trajectory, following the standard practice in training diffusion models. It is noteworthy that LAWA checkpoint averaging improves the FID over the already EMA'ed checkpoints. We defer a more thorough empirical investigation of LAWA for the family of diffusion models to future work. 
related work: 
 Weight Averaging (WA) has been studied and employed since the 1960s, predominantly in simple linear (Lakshminarayanan & Szepesvari, 2018) and convex settings (Polyak & Juditsky, 1992; Neu & Rosasco, 2018) . Recent WA approaches in deep learning can be broadly classified into two categories; First, approaches that simultaneously trains multiple models with different initialization and hyper-parameters (Wortsman et al., 2022; Rame et al., 2022; Jolicoeur-Martineau et al., 2023; Rame et al.) to later average them for better generalization. Second, approaches that focus on improving generalization of the final model or models close to convergence (Tarvainen & Valpola, 2017; Izmailov et al., 2018; Athiwaratkun et al., 2018b; Yang et al., 2019; Cha et al., 2021) . Stochastic WA (SWA) (Izmailov et al., 2018) employs a similar technique of averaging checkpoints along training trajectories but only works in the later stages of training (i.e. post 75% of the training run) with a new LR scheduler. This unsual halting and restarting the training with SWA with a new schedular limits its integration. We also show that SWA when applied early on during training diverges (Section A.2). Our recipe focuses on getting early gains through early averaging and can be generally applied to a wide range of training regimes. 
conclusion and future work: 
 In this paper we investigated a LLM pre-training setting where the LR is significantly higher than what is conventionally used. This scenario is particularly practical as LLMs are often trained using numerous GPUs in parallel, necessitating higher batch sizes for optimal GPU utilization. Here we introduce early weight averaging throughout the training trajectory utilizing LAWA. Our findings indicate that this strategy enables LLMs to generalize more effectively in fewer steps compared to the original pre-training scheme, and key baselines as demonstrated using nanoGPT-2 models. Subsequently, we applied LAWA to Pythia LLMs of varying scales-1B, 2.8B, 6.9B, and 12B parameters-to assess its effects. Remarkably, we observed consistent gains in both test and zeroshot performance across different scales. Similar performance enhancements were also noted when LAWA was applied to generative image diffusion models. Looking forward, several fascinating avenues for future research are evident. These encompass the application of proposed pre-training recipe in a fine-tuning scenario with a very restricted training budget and continual training of the intermediate averaged checkpoints.  
 
 Appendix: Averaging very initially i.e. before 8K steps during training may not always yield beneficial results (Figure 12 ). However, the technique does start showing efficacy fairly early in the training process. We highlight this phenomenon by presenting experimental results with the Pythia-1B model using a held-out set of PILE-bookcorpus2 and PILE-enron emails. We observe that LAWA trajectory undergoes a phase transition at the 8K training step. Beyond this transition, significant improvements in test performance can be seen. Such a phase transition may not occur for all Pythia LLMs. Following this phenomenon we presented our results starting 21K steps in Figures 4 5 6 7 8 . We further examine this phenomenon through the lens of linear mode connectivity.Note that W refers to the set of weights {w1, w2, ...}.https://github.com/EleutherAI/lm-evaluation-harness 


Decision: This paper makes an interesting twist to model averaging by suggesting that averaging be done with large learning rates and from spaced-out checkpoints. They show significant gains over previous methods in early stages of training. At convergence the gains are small but consistent. The core idea is thus quite simple, and the paper does not provide any theoretical insights. However, the paper may have practical impact for anyone trying to train a small/medium-scale model fast. All reviewers unanimously agree on this point. Experiments on large-scale LLMs is lacking, and it unclear how far these conclusions will generalize to larger models. However, in the small and medium scale, the paper provides exhaustive experiments. 

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: questions: - 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: (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: questions: Please see my "reasons to reject", and let me know if I misunderstand anything!
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: See above.
paper_summary: paper_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.
strength_weakness: strength_weakness: + 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: questions: 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).

Comments: title: Re: More clarification on Novelty
Ack
Invitation for further discussions
We appreciate your feedback and invite you for further discussions
We added additional experiments please provide us with some additional feedback
Let us know if we have addressed your concerns
comment: Thank you for engaging with us for further discussions. 

**Clear final loss improvements with far averaging not studied in original LAWA paper.**

Below we provide the final improvements in val loss with 1K far averaging compared to LAWA with not so far averaging. The table is created using New Fig. 1 from rebuttal and Fig. 2 b) in paper.

| Models                  | Step| LAWA (far=100)| LAWA (far=1K)| Original |
|----------------------|--------|----------------|------------|-------------|
| GPT-2 Small | 70K   |  3.28 | 3.26 | 3.28 |

| Models                  | Step| LAWA (far=200)| LAWA (far=1K)| Original |
|----------------------|--------|----------------|------------|-------------|
| GPT-2 Medium | 70K | 2.828 | 2.819 | 2.845 |

We show thematically similar results with pythia-1B (refer Fig. 10 b). 

**Major differences with original LAWA**

We emphasis that the key novelty of this work is not in showing that LAWA works , but to demonstrate that LAWA works better for early and far away averaging, this insight is neither trivial nor just an application (rather counter intuitive in optimization perspective as latest checkpoints are considered better than far away older checkpoints). Below we list some other major differences.
| Topic      | LAWA-Ous       | LAWA-Original|
|------------|------------|-------------|
| SWA baseline | ✔️         | ❌         |
| EMA baseline | ✔️         | ❌         |
| Large scale experiments | ✔️         | ❌         |
| Impact on models trained with large LR  | ✔️         | ❌         |
| Impact on downstream performance | ✔️         | ❌         |
| Additional analysis with linear model connectivity | ✔️         | ❌         |
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.
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.
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.
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.
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.