Early Weight Averaging meets High Learning Rates for LLM Pre-training
Authors: Sunny Sanyal
Atula Tejaswi Neerkaje
Jean Kaddour
Abhishek Kumar
sujay sanghavi
Venue: COLM (2024)
Keywords: LLM pre-training, High learning rates, Weight Averaging, Model merging
TL;DR: LLM pre-training requires large batch sizes paired with higher learning rates; we investigate a specific weight averaging scheme that improves language modelling in this pre-training setting.

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.

Published as a conference paper at COLM 2024
Early Weight Averaging meets High Learning Rates for LLM Pre-training
Sunny Sanyal*
Atula Neerkaje* UT Austin
Jean Kaddour University College London
Abhishek Kumar Google DeepMind
UT Austin
Sujay Sanghavi UT Austin
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 gener- alization quite early during training. Here we show that models trained with high learning rates observe higher gains due to checkpoint averag- ing. Furthermore, these gains are amplified when checkpoints are sampled with considerable spacing in training steps. Our training recipe outper- forms conventional training and popular checkpoint averaging baselines such as exponential moving average (EMA) and stochastic moving aver- age (SWA). We evaluate our training recipe by pre-training LLMs, where high learning rates are inherently preferred due extremely large batch sizes. Specifically, we pre-trained GPT2 models of varying (125M), medium (335M), and large (770M) on the OpenWebTexx 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. Code is available at https://gi thub. com/ sanyalsumy1111/Errr Weiigie
1 Introduction
Large Language Models (LLMs) have made 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 al. (2023a:b). This surge in both data and model size has rendered LLM pre-training increasingly slow and resource-intensive For instance, 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; Kaplan et al., 2020) to utilize exceptionally large batch sizes, thereby ensuring maximal GPU utilization. The usage of large batch sizes requires scaling the learning rates proportional to it's batchsize (Goyal et al., 2017), (Krizhevsky, 2014) for SGD or proportional to the square root of the multiplier of it's batch size for adaptive gradient methods (You et al., 2020; Malladi et al., 2022). Overall, high learning rates are preferred when utilizing large batch sizes.
In this paper, our goal is to improve the test generalization (log perplexity) also known as validation loss for L.L. 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 single training trajectory further amplifies these gains. We integrate these two insights to adapt LAWA (LAtest Weight Averaging) (Kaddour, 2022) technique that performs checkpoint averaging throughout training trajectory using a sliding window for pre-training LLMs. "These authors contributed equally this work. Correspondence sanyal wwwnsuuueens..eeeee. edu

---

Published as a conference paper at COLM 2024
325
16 22m22mml ama 0.2M
2s OILOPT22 matim CSSW
12 lane 173396
120
146
33
Drgat
oegial
235
30
280
LEGA 10
(ARK (00m
10000
20000
43030 5300E
5X0CE
7000
50090
30300
50000
73003
50009
20000
3000
5000e
7000
Tuining
Mapo
(a)
(b)
(c)
Figure 1: Across all model sizes, LAWA achieves faster convergence and generalizes better in comparison to original pretraining run and other baseline averaging schemes. Validation
loss on OpenWebTex Text with 70K training steps; (a) GPT2-small (125M) with Original is 2.963, EMA is 2.949, SWA is 2.952 and LAWA (ours-best) is 2.917, (b) GPT2-medium (355M) with Original is 2.855, EMA is 2.845, SWA is 2.837 and LAWA (ours-best) is 2.819, and (c) GPT2-large (770M) with Original is 2.977, EMA is 2.968, SWA is 2.961 and LAWA (ours-best) is 2.908.
We evaluate our methodology by pre-training GP'22 models of various scales, specifically 125M (small), 355M (medium), and 770M (large), using the Open Web'Text dataset, which comprises 9B tokens. The experiments with GPT2 are conducted in controlled environ- ment 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 real-world LLMss.
Main Contributions. In summary, our findings are as follows,
1. We empirically show that models trained with high learning rate (LR) show pro- nounced gains over original training on performing checkpointaveraginn very early on during training (Figure 2). This gain further amplifies when we sample distant checkpoints in the training run (Figure 1). We provide 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 com- pared to original training starting very early during training; for both GPT2 and Pythia LLMs (Figure 1 and Figures 4-7). LAWA also improves zero-shot perfor- mance for GPT2 and Pythia LLMs, 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 pro- posed 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 preliminary investigation of early weight averaging for 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 8.
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
2

---

Published as a conference paper at COLM 2024
34
01g Se-
34
Org
09%
32
32
23
18900
20006
33890
$9300
10006
23900
20009
48900
50000
33006
20069
30300
40090 52006
Thiling Strpe10)
Thariing crips
Thikiing derpe
(b)
(c)
Figure 2: We compare two independently trained GPT2 (125M) models with LR [6 x 10 3, 3,6 10 4] on OpenWeb Text data. For the sake of clarity we show results from the same training run in three parts (a) Pre-training curve with and without LAWA. LLMs trained with higher LR observes higher gain due to LAWA. (b) The model trained with high LR generalizes poorly compared to its counterpart trained with low / tuned LR. (c) The generalization gap caused by the high LR is effectively mitigated by LAWA.
models with different initialization and hyper-parameters (Wortsman et al., 2022; Rame et al., 2022; Jolicoeur-Martineal 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 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 new LR scheduler. This unsual halting and restarting the training with SWA with new schedular limits its integration. We also show that SWA when applied early during training diverges (Section C.2). Our recipe focuses on getting early gains through early averaging and can be generally applied to wide range of training regimes. To the best of our knowledge our paper is the first to show the improvements in LLM with weight averaging. We expand the other key differences of our work with SWA in Appendix B.
2
Intuition and Method
2.1 Toy Setting
We explain the setting using a simple toy problem of minimizing a two-dimensional loss function, represented as L(W1, (W), where Wi and W2 are the parameters of the model. In this scenario, there exists an optimal batch size, Bo, and an optimal learning rate, f/o, that minimizes the loss function. Assuming, for reasons outlined in Section 1, we are compelled to use a batch size, B, and a learning rate, !/, such that both are significantly larger than their optimal counterparts, i.e., B >> Bg and /J >> f/o respectively. Suppose also that the loss function L(w,, w2) exhibits a much higher curvature along W1, as compared to w2. It is widely known that the updates from the AdamW optimizer are mostly uniform across all weight dimensions (Liu et al., 2023). When 1j >> 1/o the weight updates of W2 will be accelerated, however this will cause oscillations along W1 which in a long run adversely affects the convergence of Wi A naive approach to mitigate this problem is to use smaller LR or decay LR to O, which might hinder progress in flatter regions. is conceivable that LLM might exhibit extremely heterogeneous curvatures exacerbating this issue during pre-training.
2.2 Intuition of Our Approach
Note that W refers to the set of weights {w), w2,
3

---

Optimization Viewpoint. We propose performing checkpoint averaging of model weights relatively early during training with high learning rates (f)) The rationale Published as a conference paper at COLM 2024
Interval
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 decoupled from the weight update during optimization, as checkpoint averag- ing is conducted in a post-hoc manner. Em- ploying this simple technique, we mitigate the oscillations in wi while swiftly travers- ing through W2, achieving enhanced gen- eralization in fewer training steps as illus- trated in Figure 3.
Wavg
Figure 3: LAWA illustration: Given weights Wi, W2, from high LR trajectory sepa- rated by k-stepsize (v, Algorithm 1), LAWA computes Wang at a given step.
Diversity Viewpoint. The practice of av- eraging the weights of model checkpoints broadly recognized as being functionally analogous to ensembling (Izmailov et al., 2018; Wortsman et al., 2022). model en- sembling literature it is well established that diverse models improve the performance of the ensemble (Lakshminarayanaa et al., 2016). Therefore, it is fair to assume that this principle also applies to model averag- ing as well. In our context, we define the diversity of a model two distinct training steps, N DN
Algorithm 1 LAWA: Pytorch-style pseudocode
def LAWA(ckpts, interval, k, k_stepsize):
ckpts: list of checkpoints
k: number of checkpoints average
k_stepsize: distance training steps
between checkpoints
everaged_ckpts = []
for i in range lencck)ts), interval):
start_idx =
W2] which calcu- lates the number of disagreements between the two checkpoints. This equation com- putes the number of samples from the same held-out set where the checkpoint W1 dis- agrees with checkpoint W2. A recent study by Athiwaratkun et al. has demonstrated that higher LR can result in the generation of diverse model checkpoints. We observed (Figure 1) that this phenomenon can be fur- ther amplified by sampling far apart checkpoints in terms of training step. We combine both these insights to induce diversity in our checkpoints.
end_idx -
select last k checkpoints
with k_stepsize between each checkpoint
selected_
[oots[eed__ix-kt]]]
end_idx+1 k_stpsitt
avg average averaged_ apppepttpttt.tdddd..eee.e append(avg) return averaged_ccptt
2.3 LAWA: LAtest Weight Averaging
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 first in first out (FIFO) queue of periodically sampled checkpoints with large number of intervening steps (v) in between two succesive samples. We adapt LAWA for our setting with minor modifications. Specifically, we introduce k_stepsize (v), and decoupled interval and to effectively sample distant checkpoints in the training run.
LAWA runs a moving window at predetermined interval collect latest check- points on sequence of saved checkpoints 0, The LAWA derived checkpoints are computed 0! AWA. & E k 0g where 0, the original checkpoints sampled several training steps apart in the training process.
4

---

Published as a conference paper at COLM 2024
3 Experimental Setup
GPT2 Experiments. We conduct all our experiments utilizing autoregressive decoder-style Large Language Models (LLMs), specifically GPT2 and Pythia LLMs. We utilize three distinct sizes of GPT2 models: small (125M), medium (355M), and large (770M). We train GPT2 models from scratch using the OpenWebText dataset, which includes 9 billion training tokens and 4.4 million validation tokens. Throughout the experiments, 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 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 a standard practice. For SWA, we adhere to the original pre-training procedure until 75% completion, after which SWA training is initiated with new SWA scheduler (cosine annealing). We compute the SWA uniform average every 10 steps.
maintain
Pythia Experiments. The Pythia LLMs are publicly available in the Pythia suite (Biderman al., 2023). We report results on Pythia-1B, Pythia-2.8B, Pythia-6.9B, and Pythia-12B; Table.3 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), publicly available, curated collection of English text corpus of size 800GB. The original PILE dataset is curated using different genres of data namely, academia, internet, prose, dialogue and miscellaneous. The PILE dataset contains 300 billion tokens prior to deduplication, and this number reduces to 207 billion tokens after the deduplication process. Our experiments use Pythia models trained with PILE-deduped dataset, as such models tend to memorize less (Lee et al., 2021). The batch size for all the Pythia models was set at 2.09 million tokens and the learning rate was scaled following Zhang et al. (2022).
Evaluation We evaluate the language modelling performance of GPT2 models pre-trained for 70K steps using log perplexity (perplexity and loss used interchangeably) on the held- out/val set. For GP112 we use the moving window interval = IK, k = 5 and we sample checkpoints interval = 200,1K apart for LAWA. Next we analyze the original training trajec- tories of Pythia LLMs and demonstrate the improvements achieved in test generalization using LAWA. We also present zero-shot evaluation results on Lambada OpenAl (Paperno et al., 2016), SciQ (Welbl et al., 2017), A12 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 prescribed by Xia et al. (2022). For the purpose of evaluation we use the open source library enrelaaatioon h We select representative subsets from the diverse genres encom- passed by the full PILE validation and test dataset. This subset comprises PILE-philosophy papers, PILE-bookcorpus2, and PILE- YoueTbbe subtitles datasets.
subset of the test and validation dataset following the methodology
We conduct zero-shot evaluation the Pythia LLMs following prior pre-training works (Zhang et al., 2022; Touvron et al., 2023a). 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 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
5

---

Published as a conference paper at COLM 2024
Models
Steps
Lambada
ARC-easy
SciQ
BoolQ
Average
LAWA
Original
LAWA
Original
LAWA
Original
LAWA
Original
LAWA
Original
C
50 K
33.77
27.65
44.65
44.07
74.10
71.60
52.08
48.26
51.15
47.89
70 K
35.57
33.57
44.65
44.36
75.8
74.6
54.31
54.43
52.5
51.74
C
50 K
35.88
30,95
44.74
44.28
75.8
72,9
50.76
50.49
49.72
51.72
70 K
36.95
33.59
45.96
45.33
75.2
74.2
52.94
51.62
52.76
51.18
GPT2(770M)
50 K
31.28
28.08
42.21
40.99
65.2
65.1
58.17
51.1
49.21
46.31
70 K
33.51
29.92
43.6
42.72
68.2
67.7
56.39
52.6
50.42
48.23
Table 1: The zero-shot performance of GPT2 LLMs on academic question answering and knowledge assessment downstream tasks is improved by LAWA. The checkpoints derived using our recipe require fewer steps to reach higher zero-shot performance than the check- points derived using original training.
LAWA
LAMA
LAMA
8
8
22K
120K
Trustop Stepact
Tratzag Suepetk)
(a) Philpapers
(b) Bookcorpus2
(c) Youtube Subtitles
Figure 4: LAWA speeds up convergence for Pythia-1B 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.
window at 3K steps (i.e. interval 3K) and average last 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 ??.
4 Results
4.1 Exploring LLM pre-training with GPT2 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 av- eraging. Initially we train GP12 small with two different LRs (6 x 10-3, 6 x 10-4), keeping batch size and all relevant hyperparameters the same. 6 x 10 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 x 10-3, which is tenfold higher in magnitude than the former. As shown in Figure 2(a), models 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, phenomenon also observed by Kaur et al. (2022). The gap in generalization is effectively mitigated by checkpoint averaging through LAWA, 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 's size, adopting our proposed training recipe could be advantageous. One might select 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 during training, shown in Figure 1. Since we employed

---

Published as a conference paper at COLM 2024
Models
Steps
Lambada
SciQ
WikiTexi())
ARC-easy
LAWA
Original
LAWA
Original
LAWA
Original
LAWA
Original
48 K
50.32
46.85
84.6
84.3
18.34
19.33
54.50
54.25
60 K
50.77
47.00
84.6
84.4
17.91
18.82
55.18
54.50
Pythia-1B
105 K
57.84
56.39
86.1
86.3
16.83
17.12
56.77
56.27
141 K
58.99
58.68
86.7
87.6
16.50
16.71
58.33
58.16
63.5
61.9
86.5
14.60
15.37
61.4
64.3
63.8
14.17
14.87
62.7
Pythia-2.8B
105
64.77
63.14
87.7
87.4
12.91
13.08
63.67
63.04
141 K
65.47
65.26
88.8
88.6
12.59
12.70
64.68
64.56
48
65.8
62.3
88.7
88.0
13.55
14.25
63.7
62.9
60
67.1
64.6
88.6
89.0
13.04
13.61
64.1
62.9
Pythia-6.9B
105 K
68.05
67.78
91.1
91.2
11.92
12.07
67.88
67.08
141 K
69.08
68.85
92.0
91.7
11.61
11.70
68.13
67.80
48 K
66.8
65.4
89.7
88.8
13.09
13.35
66.4
65.1
60 K
67.8
66.2
90.3
90.5
12.54
12.76
67.3
60.1
Pythia-12B
105 K
71.06
70.65
91.6
91.9
11.17
11.33
69.78
69.31
141 K
71.56
71.00
92.8
92.3
10.84
10.91
70.58
70.74
Table 2: LAWA improves zero shot performance of Pythia LLMs 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.
Origtnsl
Geigoal
Onigaal
Wed Pegiaan
8
15
200
608 60c
DOX
1208
1214
163K
28c
10k
6dC
826
1006
120k
140
256
60C
640C
346
1606
1106
1416
(a) Philpapers
(b) BookCorpus2
(c) Youtube Subtitles
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.
reasonably high LRs for all the GPT2 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 12). 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 (v) 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, needs to sample more distant checkpoints for larger models. This observation is also consistent with billion parameter Pythia LLMs, as shown in Figure
7

---

Published as a conference paper at COLM 2024
LAWA
45
LAMA
110
LAMA
(kigiia)
Meat Serjaaty
Truining Stepark)
Chuadag Stegatt
Tnstang Stepr
(a) Philpapers
(b) BookCorpus2
(c) Youtube Subtitles
Figure 6: LAWA speeds up convergence for Pythia-6.9B 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.
LAWA
90
Orglbal
Worl teeaate
20K
43K
61K
BIK
LOOK
130K
200K
3BK
61%
66X
123K
128K
100K
005
62%
100K
128K
140E
(a) Philpapers
(b) BookCorpus2
(c) Youtube Subtitles
Figure 7: LAWA speeds up convergence for Pythia-12B 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.
4.2 Scaling to Billion parameter Pythia LLMs
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 and 7, we observe significant improvements in test generalization during early-mid training regime and minor improvements towards the end for both Pythia-6.99 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 13), subsequent training costs and ingested training data. The savings are computed based on Table 4 in the appendix. Additionally, LAWA proves beneficial in situations where training similar models from scratch is necessary but can only be conducted over 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, delve into the Pythia suite's '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 x 10 while maintaining the batch sizes to 2M tokens, in line with prior work. Overall we observe thematically similar trends with the GPT2 experiments.
Recent work (Kaplan et al., 2020; Zhang et al., 2022; Touvron et al., 2023a; Dey et al., 2023; Wortsman al., 2023) report loss spikes- brief degradations in the performance along training trajectory when scaling up the model size, batch size, and learning rate. In our evaluations, we observe two perplexity spikes (Figure 6). Interestingly, we find that LAWA mitigates the spikes during evaluation quite effectively. This can be intuitively explained by the smoothing effect- average out the outliers to fit the trend - resulting from averaging checkpoint weights over range 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..
8

---

Published as a conference paper at COLM 2024
4.3 Improved Zero-shot performance
LAWA enhances the zero-shot performance of both GPT2 and Pythia LLMs (Table 1). In the GPT2 models, assess zero-shot performance at 50K and 70K steps, revealing that the checkpoints generated using LAWA at 50K consistently outperform original check- points-derived from conventional at 50K training steps (always) and also 70K training (in the majority of instances).
For Pythia LLMs 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 the scale as shown in Table 2. For instance, the LAWA Pythia-1B checkpoint evaluated at 24K steps on the Lambada OpenAl 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 Moreover, 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 compute optimal LLM training scenario where early stopping employed at 75% of total training. Additionally, we find that our LAWA derived checkpoints of Pythia-6.96 reach the final accuracy perplexxity 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 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 Section 4.1 naturally apply to zero-shot performance as well.
4.4 Diffusion models
We also experiment with image diffusion models to gauge the effectiveness of LAWA generative models beyond language. The underlying model is 422M parameter UNet (Ronneberger et al., 2015; Ho et al., 2020) trained with e-prediction objective and standard cosine schedule (Ho et al.,
Original
13-
LAWA K=2, v=20k
LAWA K=3, v=20k
FID score (eval)
31-
2020) on ImageNet 128x128 dataset. The
model was trained with the Adam (Kingma & Ba, 2014) optimizer using learning rate of 1e-4. Figure shows the FID on the eval- uation set for the baseline checkpoints and LAWA averaging over the baseline check- points. Note that the baseline checkpoints themselves are obtained using the Exponen- tial Moving Average (EMA) with decay rate of 0,9999 over the training trajectory, follow- ing the standard practice in training diffu- sion models. is noteworthy that LAWA checkpoint averaging improves the FID over the already EMA'ed checkpoints.
300K
400K
Training steps
Figure 8: LAWA speeds up the convergence for Image diffusion model, measured in terms of FID on the evaluation set for ImageNet- 128x128.
Conclusion
In this paper we investigated LLM pre-training setting where the LR significantly higher than what is conventionally used. This scenario is particularly practical as LLMs are often trained using numerous GPUs 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 GPT2 and Pythia models.

---

Published as a conference paper at COLM 2024
References
Ben Athiwaratkun, Marc Finzi, Pavel Izmailov, and A. Wilson. There are many consistent explanations of unlabeled data: Why you should average. International Conference on Learning Representations, 2018a.
Ben Athiwaratkun, Marc Finzi, Pavel Izmailov, and Andrew Gordon Wilson. Improv- ing consistency-based semi-supervised learning with weight averaging. ar preprint ar #xXit18866005538 2(9): 2018b.
Stella Biderman, Hailey Schoelkopf, Quentin Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar van der Wal. Pythia: A suite for analyzing large language models across training and scaling, ar arv preprint ar Xiv: Arxiv- 2304.01373, 2023.
Junbum Cha, Sanghyuk Chun, Kyungjae Lee, Han-Cheol Cho, Seunghyun Park, Yunsung Lee, and Sungrae Park. Swad: Domain generalization by seeking flat minima. Advances in Neural Information Processing Systems, 34:22405-22418, 2021.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Isvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Hen- ryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanu- malayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language modeling with pathways. GOOGLE, 2022.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. ArXiv, abs/ htss/1803.0545,, 2018.
Nolan Dey, Gurpreet Gosal, Zhiming, Chen, Hemant Khachane, William Marshall, Ribhu Pathria, Marvin Tom, and Joel Hestness. Cerebras-gpt: Open compute-optimal language models trained the cerebras wafer-scale cluster. Xrx preprint aix: mrxio-2304.03208, 2023.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The pile: An 800gb dataset of diverse text for language modeling. ar Xiv preprint ar Xiv: Arxio-2101.000027 2020.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangging Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in hour. ar Xiv preprint aix: Arxio-1706.02666 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851. 2020.
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. arv preprint ar arXi:o2203.15556, 2022.
10

---

Published as a conference paper at COLM 2024
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry P. Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. In Amir Globerson and Ricardo Silva (eds.), Proceedings of the Thirty-Fourth Conference on Uncer- tainty in Artificial Intelligence, UAI 2018, Monterey, California, USA, August 6-10, 2018, pp. 876 885. AÚAT Press, 2018. URL http://auat http:///aaaa..r/guaaaaa/a org/uai201 8/ proceeedings/pppprr///3 pdf.
Alexia JolicoeurMMartineau, Emy Gervais, Kilian Fatras, Yan Zhang, and Simon Lacoste- Julien. Population parameter averaging (papa). arXiv preprint arXiv: Arxio-2304.03094, 2023.
Jean Kaddour. Stop wasting my time! saving days of imagenet and bert training with latest weight averaging. arv preprint ar arXiv22209.14988 2022.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. Xr preprint Xiv: 008361, 2020.
Simran Kaur, Jérêmy E. Cohen, and Zachary Chase Lipton. On the maximum hessian eigenvalue and generalization. ICBINB, 2022. doi: 10.48550/ 10.4850/arXiv.2206.1666 arX.i.2206.106654
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint ar Xit::1412.6980, 2014.
Alex Krizhevsky. One weird trick for parallelizing convolutional neural networks. arv preprint Xiv: 1404.5997, 2014.
Balaji Lakshminarayanan, A. Pritzel, and C. Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. NIPS, 2016.
Chandrashekar Lakshminarayanan and Csaba Szepesvari. Linear stochastic approximation: How far does constant step-size and iterate averaging go? In International Conference Artificial Intelligence and Statistics, pp. 1347-1355. PMLR, 2018.
Katherine Lee, Daphne Ippolito, Andrew Nystrom, Chiyuan Zhang, Douglas Eck, Chris Callison-Burch Burch, and Nicholas Carlini. Deduplicating training data makes language models better. Xiv preprint ar Xiv: Arxie-2107.06499 2021.
Zhuohan Li, Eric Wallace, Sheng Shen, Kevin Lin, K. Keutzer, D. Klein, and Joseph Gonzalez. Train large, then compress: Rethinking model size for efficient training and inference of transformers. International Conference On Machine Learning, 2020.
Hong Liu, Zhiyuan Li, David Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training, arXiv preprint ar Xiv: 2305.14342, 2023.
Sadhika Malladi, Kaifeng Lyu, A. Panigrahi, and Sanjeev Arora. On the sdes and scaling rules for adaptive gradient algorithms. Neural Information Processing Systems, 2022. doi: 10.48550/ 10.48550/arXiv.220 arXiv.2205.102287
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016. URL https: //ap:v////xivvi///a wtw.aaawww...ggaa///aaaaaaaaaa 07843.
Gergely Neu and Lorenzo Rosasco. Iterate averaging regularization for stochastic gradient descent. In Conference On Learning Theory, pp. 3222-3242. PMLR, 2018
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Q. N. Pham, R. Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and R. Fernández. The lambada dataset: Word prediction requiring a broad discourse context. Annual Meeting Of The Association For Computational Linguistics, Linpulationallinuustiiisss 2016. doi: 10.1653/v1/PI1--1144
Boris Polyak and Anatoli Juditsky. Acceleration of stochastic approximation by averag- ing. STAM journal on control and optimization,, 30(4):838-855 1992.
11

---

Published as a conference paper at COLM 2024
Alexandre Rame, Kartik Ahuja, Jianyu Zhang, Matthieu Cord, Leon Bottou, and David Lopez-Paz. Model ratatouille: Recycling diverse models for out-of-distribution general- ization.
Alexandre Rame, Matthieu Kirchmeyer, Thibaud Rahier, Alain Rakotomamoniy, patrick gallinari, and Matthieu Cord. Diverse weight averaging for out-of-distribution general- ization. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL httpps://openreviers .eete net/ J.rumPidetq.JMqqB3...
Xiaozhe Ren, Pingyi Zhou, Xinfan Meng, Xinjing Huang, Yadao Wang, Weichao Wang, Pengfei Li, Xiaoda Zhang, A. V. Podolskiy, G. Arshinov, A. Bout, Trina Piontkovskaya, Jiansheng Wei, Xin Jiang, Teng Su, Qun Liu, and Jun Yao. Pangu-E: Towards trillion parameter language model with sparse heterogeneous computing. ARXIV.OG,, 2023. doi: 10.48550 2r11..033..22211555555.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Asee Assisted Interoention-MIICAA. 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pp. 234-241. Springer, 2015.
Mark Sandler, Andrey Zhmoginov, Max Vladymyrov, and Nolan Miller. Training trajectories, mini-batch losses and the curious role of the learning rate. ar aixiv preprint arXiv:2301 02312, 2023.
Nikunj Saunshi, Sadhika Malladi, and Sanjeev Arora. A mathematical exploration of why language models help solve downstream tasks. International Conference On Learning Representations, 2020.
Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jonathon Shlens, and Z. Wojna. Re- thinking the inception architecture for computer vision. Computer Vision And Pattern Recognition, 2015. doi: 10.1109
Antti Tarvainen and Harri Valpola. Mean teachers are better role models: Weight-averaged consistency targets improve semi-supervised deep learning results. Advances in neural information processing systems, 30, 2017.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, A Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models, ARXIV, 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton,, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, lliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models. Xiv preprint Xiv: 2307.09288, 2023b.
Johannes Welbl, Nelson F. Liu, and Matt Gardner. Crowdsourcing multiple choice science questions. ArXiv, abs/1707.0611 2017.
Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo- Lopes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith,
12

---

Published as a conference paper at COLM 2024
et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International Conference on Machine Learning, PP- 23965-23998. PMLR, 2022.
Mitchell Wortsman, Tim Dettmers, Luke Zettlemoyer, Ari Morcos, Ali Farhadi, and Ludwig Schmidt. Stable and (ow-precision training for large-scale vision-language models. arXiv preprint arvi Arxio-2304 13013, 2023.
M. Xia, Mikel Artetxe, Chunting Zhou, Xi Victoria Lin, Ramakanth Pasunuru, Danqi Chen, Luke Zettlemoyer, and V. Stoyanov. Training trajectories of language models across scales. 22.0903...
ARXIV.ORG, 2022. doi:
Guandao Yang, Tianyi Zhang, Polina Kirichenko, Junwen Bai, Andrew Gordon Wilson, and Chris De Sa. Swalp: Stochastic weight averaging in low precision training. In International Conference on Machine Learning, pp. 7015-7024 PMLR, 2019.
Yang You, Jing Li, Sashank J. 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 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. wos.:omwmmnidsyywtttttt...
URL https://openrevieee
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettlemoyer. Opt: Open pre-trained transformer language models. ARXIV.ORG, 2022.
13

---

Published as a conference paper at COLM 2024
Supplementary Materials: Appendix
Contents
A: Frequently Asked Questions.
B: Additional Related Works.
C: Ablations
D: Supplementary Experiments and Results.
A Frequently Asked Questions
A.1 Why your method seems to work well for small LLMs but not for the large ones?
The 6.9B and 12B pythia LLMs are trained with lesser LRs compared to their smaller counterparts hence one may argue that the gains due to LAWA diminished in Figure and 7. Although, the gains in terms of perplexity may have diminished; but shown in Figure 13, we observe that higher savings in compute can be achieved for larger models.
A.2
LAWA tends to improve the intermediate checkpoints than the final checkpoint. Could you clarify the significance of this work?
We humbly disagree with such an interpretation of our results as LAWA has improved the final model in terms of validation loss and downstream tasks for all the GPT2 models (refer Figure and Table 1). Next for pythia 1B and 2B models we have observed improvemente both in final model val loss and some downstream tasks. For Pythia-6.9D model we observe that the final checkpoint at 141K derived using LAWA clearly performs better than the one created using original pre-training on downstream tasks (refer Table 2. For the pythia 12B model the gains are limited and this could be due to related to the larger parameters and training with a reasonably low learning rate. Due to compute limitations we defer this study for future.
Some interesting implications. Our approach has no discernible downside (regarding the results we discussed in this paper with GP12 and large scale pythia LLMs) as incurs no extra compute or data to achieve better performance. Our approach advocates for far checkpoint averaging which typically leads to lesser requirements of disk space to save checkpoints. Today LLMs are pre-trained on trillions of tokens and trained for months. In a compute optimal setup where the training steps is bounded by the budget prior works Hoffmann et al. (2022); Kaplan et al. (2020) suggests early stopping and using an intermediate checkpoint as the final model. In such a scenario using our recipe one can sample intermediate checkpoints and make then perform better, thereby accelerating the entire model development pipeline.
A.3
Why this paper focuses too much on high LRs when the Pythia LLMs not trained with high LRs?
The paper focuses on high LRs as LAWA works best when applied to checkpoints trained with high LRs as shown in Figure 2. However the same Figure shows that LAWA integrates with well tuned LR (i.e. 6e-4) as well. Overall for the models under consideration it is safe to say that LAWA provides some gains for all LRs but works the best with high LRs. Moreover the idea of high LR also correlates well with our intuition as discussed in Section 2. The pythia LLMs (Biderman et al., 2023) are not trained by us and due to compute limitation cannot run LR ablations on billion size LLMss.
14

---

Published as a conference paper at COLM 2024
A.4 Why aren't existing solutions good enough?
In LLM pre-training literature to the best of our knowledge no work practically uses any form weight averaging particulary during pre-training. Some prior works like EMA and SWA are popularly used in vision setting that to for small models. We have adapted these two methods LLM pre-training and consistently outperform both of them as shown in Figure
B Additional Related Works
This section serves as supplement to the discussion about prior WA methods discussed in Section Here we further elaborate on prior work-SWA, and highlight the key differences of LAWA in the backdrop of training L.Mss. SWA Izmailov et al. (2018) is checkpoint averaging scheme similar to LAWA but is extensively used in vision settings. In this paper we adapt SWA (also EMA) for LLM pre-training setting and presented it as key baseline. The pseudo-code of SWA and EMA is provided as Algorithm 2. The major differences of our recipe and SWA are listed below.
In the SWA approach the checkpoints are averaged post 75% of regular training, i.e. close to convergence. This phase of training is termed as SWA training. Moreover the end of training.
the gain in generalization through SWA method achieved In contrast, LAWA averages checkpoints very early on during training (post 10-15% training) and achieves gains in test and zero-shot generalization with fewer training steps (less than 100% TB). SWA when applied early on during training performs poorly due to divergence of the training curve as shown in Figure 12.
SWA modifies the learning rate scheduler mid-training and requires different learn- ing rate schedulers for different architectures. Additionally, for models employing batch normalization, it necessitates a full pass over the entire dataset to update batch normalization statistics. LAWA free from these constraints and simpler to implement for large scale training.
C
Ablations
C.1 Pythia LLMs
48
Original
kmb
k=2. v=2K
Word Perplextyy
Word Perpleatty
k-2, U-5K
44
84
k=5. v-2X
42
42
k-5, U-5K
u-10K
40
40
38
Training Steps(K)
Training Stepstk)
(a)
(b)
Figure 9: Ablations studying test performance as function of (a) number of checkpoints to be averaged k = {2, 5, 10) at v = 1K, (b) distance of checkpoints to be averaged = {1K, 2K, 5K, 10K) at k (2. The checkpoints used for ablations is Pythia-1B.
We perform ablations to better understand the interplay between the performance of num- ber of checkpoints to be averaged gLLWA with varying k and distance between averaged checkpoints in training. We study the Pythia-1B model with held-out subset of PILE- bookcorpus2. Additionally, we also provide training trajectory Pythia-BB model on subset of PILE datasets much earlier than 21K steps (refer Section D).
15

---

Published as a conference paper at COLM 2024
GPT-2 small (125M)
4.2
Original, Batch Size=1 6K
4.0
LAWA (Ours), Batch Sizem16k
Original, Batch Size=131K
LAWA (Ours) Batch Size- Size-131K
Validation Loss
3.6
3.0
10000 20000 30000 40000 50000 60000 70000 Training steps (K)
Figure 10: Across different batch sizes (16K and 131K tokens) for GPT-2 small model, LAWA not only achieves faster convergence but also demonstrates better generalization compared to the original pre-training method. Training curves of GPT2-small model trained using 16K and 131K tokens per batch while using Ir=6.0 x 10-3, We present both the original training curve and gains achieved due to LAWA.
Test Performance with varying k and fixed 1K. We investigate the impact of varying k the model's test performance, while keeping constant at 1K. Our aim is to determine the optimal number of checkpoints to include in the average when selecting the latest checkpoints following the LAWA approach. As outlined in Section 3, the Pythia checkpoints are saved at frequency of 1K, so we have maintained 1 at 1K for this analysis. From Figure (a), it is apparent that smaller could be detrimental, but performance remains fairly stable for reasonably large values. Consequently, for our LLM experiments, we opted for k 5 since k 10 tends to occupy substantial amount of disk space, especially for larger models such as Pythia-12B.
Test Performance with varying v for k = (2,5). Memory requirements remain a key bottleneck in saving model checkpoints throughout the training trajectory, particularly for extremely large ballion-parameter models. Therefore, it is very interesting to know how far be averaged? We investigate this question As shown in Figure 9(b,, we observe for both recent checkpoints (keeping v small) works better than averaging stale weights (higher v). For instance LAWA using k = 5 and v {1K, 2K} consistently performed better than against LAWA with other parameters. Overall, we empirically find that moderate number of checkpoints (k = 5) saved in smaller frequencies (v 1K) works best.
away checkpoints in training trajectory
with far checkpoint averaging at = [1K, 2K, 5K, 10K} training steps apart for k = 5).
and k = 5, find that averaging more
C.2 GPT-2 models
Batchwise ablation with GPT2-small. We conducted ablation studies using GPT-2 models to assess the impact of LAWA on test generalization across different batch sizes while maintaining fixed learning rate of 6.0 x 10-3, In Figure 10 We trained GPTT-2 125M model with batch sizes of 16K and 131K tokens. Our findings revealed two key insights: first, LAWA consistently resulted in lower validation loss throughout the entire training run, even with varying batch sizes; second, increasing the batch size while using a relatively higher learning rate led to significant performance gains.
Baseline Comparison with Distant EMA and SWA averaging recipes. Averaging weights from checkpoints taken at distant intervals is a key focus of this paper. To ensure a fair comparison, have adapted the EMA and SWA baselines to also average checkpoints taken far apart during their training trajectories. It is important to note that such distant averaging was not considered in the original implementations of these baselines. In Figure
16

---

Published as a conference paper at COLM 2024
OPT22 small 1125V0 Batch Stee -BK tsters
2.8
OFT-2 mestum 1353M), Satch Sie- saxems
75% traiiing
4.2
DSS Lcainin
377
3.6
Wilkibn L33
3.5
Woldtitnn
2.4
3.6
Original
32
Onginal
2.4
EMA-Far (3K)
EMA-far (1K)
(1K)
34
SWAAfar (3K)
(Ours) v-1k
LAWA (Ours) v=1K
10000 20000 non0n 400n0 50000 60000 70000 Training sreps [K]
14000 20000 30000 40p00 50000 6cOOD 70004 Training steps
(a) GPT-2 small
(b) GPT-2 medium
Figure 11: Comparing our recipe with EMA and SWA baselines in training scenario where the checkpoints are averaged at 1K training steps apart. LAWA also outperforms all the baselines in this training regime. We sampled checkpoints from the original training run 1K apart for LAWA recipe. Similarly, we made adjustments to EMA and SWA implementations and also averaged checkpoints 1K step apart. We trained our models with 16K tokens per batch for all the baselines and our recipe.
uval (122
mad (25-M)
GPT2 large 127000
14
Oriind
2.8
axes axcS Tain dhepe d0
70202
Taiiing 1R)
(a)
(b)
(c)
Figure 12: Across all model sizes (125M, 355M, 770M), we observe that the early version of stochastic weight averaging diverges, in contrast to the originally proposed late version.
11, we train two models, GPT-2 Small and Medium, using LAWA along with the SWA and EMA baselines. For all approaches, including ours, model checkpoints were sampled at intervals of 1,000 steps. For both models, checkpoints generated using the LAWA approach show faster convergence and achieve a better final validation loss.
Early-SWA Experiments SWA has previously demonstrated benefits for smaller vision models, particularly in the later stages of training (typically after 75%). We explored the impact of applying SWA during the early stages of training, which we refer to as SWA Early. In Figure 12, we compare SWA applied from the very beginning of training with SWA applied after 75% of training, the latter being consistent with the original SWA implementation in Izmailov et al. (2018), which we refer to as SWA Late. Our findings indicate that SWA Early diverges quickly, suggesting that SWA does not offer benefits when applied early in the training process. In contrast, LAWA can be applied seamlessly throughout the training.
D Supplementary Experiments and Results
D.1 Compute Details and Savings
The LLM models pre-trained or used in this study are all decoder-style pre-trained on autoregressive loss. The key architectural details along with their respective best performing LRs used during pre-training of all the models are detailed in Table Note since the pythia models are not trained by us only report the LRs reported in Biderman et al. (2023). The
17

---

Published as a conference paper at COLM 2024
Pytio 1B
2
1530
2570
1920
2% 1%
perpleaty faad muddol
Figure 13: LAWA saves significant amount of GPU hours compared to original training. We compare the savings in GPU hours as a function of increase in final perplexity, i.e. perplexity achieved at 141K training step by the original checkpoint. This plot is created using held out set from the training subset PILE-phine -hil--hilsspphy papers.
Model Size
Layers
Hidden Size
Heads
Learning Rate
Equivalent Models
125M
12
12
768
6.0 10-3
GPT2 (small)
335M
24
1024
16
3,0 x 10-3
GP12 (medium)
770 M
1280
2.0 10
GPT2 (large)
1.0 B
16
2048
3.0 x 10-4
2.8
32
2560
32
1.6 x 10
GPT-Neo 2.7B, OPT-2.7B OPT-6.7B
6.9
32
4096
32
1.2 10
12 B
36
5120
40
1.2 x 10
Table 3: Overview of models and their architecture from the GPT2 suite and Pythia suite (Biderman et al., 2023) used in experiments. The model nomenclature for Pythia LLMs is pythia-xx with model size. Models marked as "equivalent" have the same architecture and number of non-embedding parameters.
compute details of the pythia LLMs for 1B-12B models are listed in Table 4. Checkpoints derived using LAWA saves lot of compute resources as shown in Figure 13. Here we plot the compute savings in GPU hours
function of tolerable increase in word perplexity metric over the held out set of subset of pre-training data. Moreover, we observe that significant compute savings can be obtained with mere 1% increase in perplexity for 1B-6.9B models. Finally this trend is interesting as this shows that larger the model becomes more savings be made with technique.
D.2 Phase Transition and Linear Mode Connectivity
Averaging very initially i.e before 8K steps during training may not always yield beneficial results (Figure 14). However, the technique does start showing efficacy fairly early the
Model Size
GPU Count
Total GPU hours required
1.0
64
4,830
2,8
64
14,240
6.9
128
33,500
12 B
256
72,300
Total
136,070
Table 4: Table from Biderman et (2023). Model sizes in the Pythia suite, number of GPUs used during training, and the total number of GPU hours, calculated (iteration time (s) x number of iterations X number of GPUs : 3600 s/ s/hour). All GPUs are A100s with 40GB of memory.
18

---

Published as a conference paper at COLM 2024
Algorithm 2 Pytorch-style pseudocode of EMA/SWA
def SWA_SWA(ckpt alpha, step_size, init_point):
ckpt: list of data points (could be checkpoints any data series) alpha: smoothing factor, alpha
If alpha 0, enables SWA. step_size: How often calculate average. Typically set to for EMA init_point: Step after which to start averaging. Typically 0 for EMA.
Initialize the series with the first data point series [ckpt[0]]
n_models 1
for i in (1, len(ch)t)):
# Calculate EMA/SWA
if i and i>init_point:
if alpha < 0:
factor (__odelss((__oooe + 1)
value (1 factor) series[i-1] + factor ckpt[i] n_models
else:
value (1 alpha) * series[i-1] alpha ckpt[i] series. append (value)
return series
200
Original
LAWA
180
Wort Perplesty
Phase transition
Word Perplestty
1690
140
Phase transition
60-
30-
80
60
10
20
30
40
50
60
10
20
4.0 Traming StepstK)
30
50
Training Stepe(K)
(a) BookCorpus2
(b) Enron Emails
Figure 14: Early weight averaging doesn't work at the very beginning of the training but works reasonably early during the training process. Here we compare original and LAWA early training trajectories for Pythia-1B model on different tasks namely bookcorpus2 and enron emails using held out set.
training process. We highlight this phenomenon by presenting experimental results with the Pythia-1B model using held-out set of PILE-bookcorpus2 and PILE-enron emails. We observe that LAWA trajectory undergoes 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 presented our results starting 21K steps in Figures 4-7. We further examine this phenomenon through the lens of linear mode connectivity.
To better comprehend the linear model connectivity of checkpoints, we perform a convex combination of model checkpoints at different training stages. For instance, a model checkpoint at 2K and 20K can be combined in this manner: & x 0x + (1 a) x 02ok In Fig. function of a using PILE-bookcorpus2. Here we observe
10, we plot word perplexity that initially the model checkpoints are not linear mode connected. However, based on the evaluated checkpoints shown in Figure 15, we posit that the model checkpoint attains linear mode connectivity (LMC) quite early and maintains this property until the end of training.
19

---

Published as a conference paper at COLM 2024
Wini pepiody
54
Weel perkels
41.3
113
66.7
103
48.
22
15
28.0
3odel302
(b)
(c)
Figure 15: The model checkpoints attains linear mode connectivity quite early but not at the very beginning of the training process. We plot word perplexity as a function of the model derived from the convex combination of different checkpoints i.e. 0yMC at = (0, 0.2, 0.4, 0.6, 0.8, 1}. In (a)
error barrier that means model checkpoint at 2K and 20K not linear mode connected, whereas both (b) and (c) shows the checkpoints under consideration are linear model connected.
20


Reviews:
[{"review_id": "dz5b9UcZpf", "reviewer": "Reviewer_GUxG", "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.", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the assessment.\n\n> (R1) The methods shows improvement in early stages…\n\nClarification: The gains are not just in the early stages, the diff. gets smaller but more significant in later stages. In later stages the LR gets fully annealed, hence this should hold for longer training runs where lr schedules will anneal much later in training. For all GPT2 and Pythia 1B-2.8B models (reasonably large) we observed consistent gains in final loss/perplexity and downstream performance. For extended discussion see A.1 A.2 in supplementary.\n\n\n>(R2) The experimental setup ..\n\nOur experiments are reasonably large such as GPT2 (770M and 9B tokens) and large scale results with Pythia (1B-12B with 207B tokens). GPT2-medium is trained with batch size(BS)=50K and GPT2-large with BS=16K due to GPU memory considerations hence GPT2-large has higher loss than medium.\n\n>(R3) The baseline setup is weak…\n\nLAWA observes higher gains with high LR but it also helps with well tuned LR and batch size (BS), for example [Pythia LLMs](https://arxiv.org/abs/2304.01373) are trained using optimized settings and still observes gains in final loss and downstream tasks for 1B-2.8B LLMs and more intermediate gains in 6.9B and 12B LLMs. We believe additional gains can be observed if llama 7b is trained with LAWA.\n\n> (R4) The EMA and SWA adaptation requires corresponding update to averaging interval similar to LAWA.\n\n**Additional results**  LAWA w/ far WA outperforms EMA and SWA with far WA. \n\nNew Fig1. https://i.postimg.cc/LX13SqPM/GPT2-small-far-comparison.png\n\nNew Fig2. https://i.postimg.cc/gjbq2YyL/GPT2-medium-far-comparison.png\n\n> (R5) The cost of …\n\nThe cost of hyper parameter search is well known to be extremely expensive as shown in several prior works such as [1][2]. Due to limited compute we didn’t show such a result and also for the fact is this is an intuitive and well established result in large scale pre-training.\n\n> Q1 How was ..\n\nBased on ablations shown in supplementary C.2 Ablations. Our ablations are done using the Pythia 1B model, no we didn’t vary model size and architecture due to compute limitations. The cost of compute and savings are provided in supplementary C.1 and Figure 9. \n\n>Q2 It’ ll be good …\n\nAll the GPT2 training conducted in this paper uses standard cosine LR scheduler following [1]. Pythia models are also trained using a similar scheduler. \n\n[1] [Liu et al.](https://arxiv.org/abs/2305.14342)\n[2] [Yang et al.](https://arxiv.org/abs/2203.03466)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The method of model weight averaging with higher learning rate is simple and shows improvement in validation loss and some downstream tasks.  \n- 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?  \n- 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.\n- 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\n- The EMA and SWA adaptation requires corresponding update to averaging interval similar to LAWA.\n- 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?\n- It'll be good to add learning rate schedule for different training methods.", "extra_scores": {}}, {"review_id": "57PKDqeKI1", "reviewer": "Reviewer_ZNy2", "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 \n\n(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.\n\n(2) For LAWA and high-LR training, using distant checkpoints performs better.\n\nThe 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.", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> (R1) Novelty ...\n\nLAWA [1] applies Polyak’s iterative weight averaging (WA) [2] to deep NNs. No prior work to our knowledge studied nature of iterates (ckpts) best suited for iterative WA. We show early, far and model trained with high lr ckpts are best suited for LAWA (ours) for LLM pre-training. Additionally we added EMA and SWA as baselines for LLMs and thoroughly analyzed LAWA. New Fig1. far away WA improves over LAWA [1]; https://i.postimg.cc/MpbR8vYh/lawa1-vs-lawa2.png\n\n> (R1 contd..) The use of LAWA …\n\nConventionally latest ckpts are thought to be better than far away older ckpts in large scale optimization specifically in early stages of training. The same goes with high lr. Our results are not so trivial and given its benefits to LLM pre-training it deserves to be studied.\n\n>(R2)  From the main results …\n\n**Clarification:** From Fig. 2 a) LAWA (ours) with lr=6e-4 (low lr) outperforms Original (no LAWA) with low lr at all intermediate steps and also at 50K but the gains are lower w.r.t intermediate steps as the lr schedule was fully annealed. Moreover for a well-tuned lr used by Pythia 1B-2B LLMs (Fig 4 and 5, Table 1) we see a clear final improvement with LAWA (ours).\n\n> (R2 contd …) \n\n**Consistent gains with LAWA(ours)** LR details are provided in Table 3. All GPT2-small, medium and large are trained with batch size (BS) 131K, 50K and 16K respectively. Pythia LLMs are trained with BS=2M tokens. LAWA achieves consistent gains with variable high lr and BS for GPT2 models in Fig 1 and fixed BS and variable LRs for Pythia LLMs.\n\n**Additional results**\n\na) variable BS and fixed LR-6e-3 for GPT2-small, New Fig3: https://i.postimg.cc/tCvLyNXx/GPT2-small-batchsize-comparison.png\n\nb) fixed BS for two GPT2 models + EMA and SWA far WA, New Fig4:\n\nhttps://i.postimg.cc/LX13SqPM/GPT2-small-far-comparison.png\n\nhttps://i.postimg.cc/gjbq2YyL/GPT2-medium-far-comparison.png\n\nWe emphasis that the gains are not limited to just the early stages. Next this is the first work to that shows improvements in LLM pre-training with WA. Our recipe has the potential to be effective when training with 1T tokens specifically for compute optimal training runs like [Chinchilla](https://arxiv.org/abs/2203.15556) which employs early stopping. For more implications please check A.2 FAQs in supplementary. \n\n[1] [Kaddour et al.](https://arxiv.org/abs/2209.14981)\n\n[2] [Polyak et al.](https://epubs.siam.org/doi/pdf/10.1137/0330046)\n\nWe will update the draft with new results.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "(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.\n\n(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.\n\n(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!", "extra_scores": {}}, {"review_id": "9Et126g6gw", "reviewer": "Reviewer_BhE4", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive assessment of our work. \n\n> R1 It would … weight average?\n\n**Additional results** we have discussed various weight averaging approaches in Sec. 1. We emphasis that conventionally no weight averaging (WA) approach is used during LLM pre-training. We have discussed differences with similar tail averaging schemes such as EMA and SWA in Sec 1 and in supplementary Sec. B. We have used EMA and SWA as our baselines shown results in Fig 1. \n\n(R1 contd..) For instance … stochastic weight average?\n\n**LAWA w/ far WA outperforms EMA and SWA with far WA** Thanks for this thoughtful remark. We provide additional experiments with GPT2-small and medium. \nNew Fig1. https://i.postimg.cc/LX13SqPM/GPT2-small-far-comparison.png\n\nNew Fig2. https://i.postimg.cc/gjbq2YyL/GPT2-medium-far-comparison.png\n\n> (R2) It seems that the …\n\nTrue, the improvement may look small but it is significant in val loss and downstream performance for all GPT2 models, Pythia 1B - 2.8B LLMs. Most importantly this gain is almost a free lunch.\nThe LR annealed fully towards the end of training hence we observe lesser gains compared to intermediate steps, But now one can sample intermediate checkpoints (say at 75%) with LAWA hence this recipe can help even in longer as well as compute optimal training runs. One can also design better LR schedulers for LAWA training runs for further gains.\n\n>(R3) The experimental … GPU utilization?\n\nThe LAWA is completely decoupled from the actual training run hence real throughput is not a valid metric for our setting. One can save checkpoints and perform LAWA even using CPUs both during or after training.\n\nIf you are convinced with our responses and additional results we hope you will re-consider your score.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1) The paper is very well written. The Section 2 is friendly to the readers who are not knowledgeable in weight averaging.\n2) 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?\n2) 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.\n3) 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.", "extra_scores": {}}, {"review_id": "yCxaOGfzNg", "reviewer": "Reviewer_rseH", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive assessment of our work.\n\n> (A1) Appendix A discusses …\n\nVery grateful for such a thoughtful remark. I will make a FAQ section in all my upcoming works.\n\n>(R1) The writing of some parts..\n\nNoted we will update the draft based on this remark.\n\n>(R2) Even though I love the intuition …\n\nWe agree the toy setting is intended to provide some intuition on why LAWA is working in a non convex setting but we make no rigorous claims there. We defer a more rigorous analysis of the optimization perspective as a future work.\n\nIn Section C.4, we discussed that LLMs achieve linear mode connectivity (LMC) not at the onset of training, but relatively early in the process. For example, for a Pythia 1B model, early weight averaging (EWA) is effective starting from 8,000 steps. This phase transition facilitates LMC, making weight ensembling feasible. Our understanding is shaped by these experimental observations.\n\n>(R3) Even though …\n\nLLMs are trained with (some sort of) large LRs due to the usage of large batch sizes. This idea is well studied for non-adaptive [1] and adaptive optimizers [2]. The Pythia LLMs are trained with well tuned LRs and we observe gains in both perplexity and downstream tasks. We emphasis that high LRs are preferred but our recipe also helps with well tuned LRs.\n\n> (R3 contd) As previous work…\n\nSophia is an optimizer and LAWA can work on top of any optimizer to amply gain. It’s a bit unfair to compare these two works as choice of LRs may vary based on the optimizer. \n\n>(R3 contd) Then … argument regarding learning rate.\n\nAs above LLMs are typically trained with larger batch sizes [3] for systemic efficiency concerns and typically large LRs are needed in such  scenarios and LAWA helps.\n\n>(R4) I noticed ...\n\nAs one can see in C.2 LAWA works better than the original training run for various different k and $\\nu$. One has to perform some ablations to select the best suited k and $\\nu$ for their usecase which may vary for different settings.\n\n>(Q1) As I said …\n\nThe averaged checkpoints to be used for further pre-training is something that we are considering as a future work.\n\n[1] [Goyal et al.](https://arxiv.org/abs/1706.02677)\n\n[2] [Malladi et al.](https://arxiv.org/abs/2205.10287)\n\n[3] [Birdman et al.](https://arxiv.org/abs/2304.01373)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "+ 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.\n+ 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.\n+ The intuition discussed in Section 2.2 makes a lot of sense to me.\n+ 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.\n+ 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.\n+ Even though I love the intuition in Section 2.2, they may not be convincing enough for some readers:\n  + The `Optimization Viewpoint` seems to hold for the toy setting, but how about the much more complicated setting of large language models?\n  + 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?\n+ 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.\n  + 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?\n  + Then, the improvement seems to be very marginal with large LR, e.g., Figure 2(a).\n  + 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.\n+ 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).", "extra_scores": {}}]

Author rebuttal:

Meta-review: 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. 
Decision: Accept