Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations
Authors: Alexander Hägele
Elie Bakouch
Atli Kosson
Loubna Ben allal
Leandro Von Werra
Martin Jaggi
Venue: NeurIPS 2024 spotlight (2024)
Keywords: Scaling Laws
Large Language Models
Learning Rate Schedules
Weight Averaging
Primary area: optimization_for_deep_networks
TL;DR: We show reliable scaling behavior of an alternative LR schedule as well as stochastic weight averaging for LLM training, thereby making scaling law experiments more accessible.

Abstract: Scale has become a main ingredient in obtaining strong machine learning models. As a result, understanding a model's scaling properties is key to effectively designing both the right training setup as well as future generations of architectures. In this work, we argue that scale and training research has been needlessly complex due to reliance on the cosine schedule, which prevents training across different lengths for the same model size. We investigate the training behavior of a direct alternative --- constant learning rate and cooldowns --- and find that it scales predictably and reliably similar to cosine. Additionally, we show that stochastic weight averaging yields improved performance along the training trajectory, without additional training costs, across different scales. Importantly, with these findings we demonstrate that scaling experiments can be performed with significantly reduced compute and GPU hours by utilizing fewer but reusable training runs. Our code is available at https://github.com/epfml/schedules-and-scaling/.

Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations
Alexander Hägele !
Elie Bakouch²
Atli Kosson
Loubna Ben Allal?
Leandro Von Werra"
Martin Jaggi'
EPFL
Hugging Face
aleeanderr hagele@epff..h
Abstract
Scale has become main ingredient in obtaining strong machine learning models. As a result, understanding a model's scaling properties is key to effectively designing both the right training setup as well future generations of architectures. In this work, we argue that scale and training research has been needlessly complex due to reliance on the cosine schedule, which prevents training across different lengths for the same model size. We investigate the training behavior of direct alternative - constant learning rate and cooldowns and find that it scales pre- dictably and reliably similar to cosine. Additionally, we show that stochastic weight averaging yields improved performance along the training trajectory, without ad- ditional training costs, across different scales. Importantly, with these findings we demonstrate that scaling experiments can be performed with significantly reduced compute and GPU hours by utilizing fewer but reusable training runs. Our code is available whttps:/// thub. com/ wwt/ http///pmm//ccddlleses/ and-scaling/
Introduction
Training large language models is expensive - time, energy, and compute. Moreover, it requires complex algorithmic recipe of model architecture and training data obtain high-quality models. Therefore, the workflow of training large models consists of iterating over small experiments to verify success before extrapolating larger scales. This is then either done by computing specialized
scaling laws (OpenAl 2023 Bi et al., 2024 Hu et al. 20244] Team et al. 2023). relying on established laws (Hoffmann et al.) 2022 Anil et 2023) or training past compute compute-optimant optimality to save cost at inference (Touvron et al. 2023a b).
Despite large advances across data and training recipes, aspect of large language model (LLM) pretraining has remained surprisingly prevalent: the cosine learning rate schedule (Loshchilov & Hutter 2016,] Radford et al., 2018 Rae et al. 202T). Importantly, the Chinchilla project (Hoffmann et al. (022) showed that the cosine schedule achieves optimal loss only when the cycle length matches the training duration, but underestimates the model performance during training. This means that when performing experiments - e.g., for architectural changes or data mixtures one must train multiple models for different lengths, scratch, have reliable estimates of the quality of training and the scaling behavior. This is much more expensive than training suite of models just once. Even more, is restrictive for the final model for which the training length must be decided in advance.
In this work, our goal is to revisit and question the necessity of the cosine learning rate schedule for large model training. Through multitude of training runs, demonstrate how simple alternative of performing a cooldown after a constant learning rate which was already suggested in the literature (Zhai al. 2022) and recently used by released models (Hu al. 2024 (hhn al. 2024) matches the performance of cosine. We expand on this and provide and analyze different recipes for the decay form and length, which scale reliable as cosine, outperforming for sufficiently long
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Perplexity at Different Lengths
Cosine Scheduler
35
Cycle
0.0020
Oyde modk
bm8k
2-20
Learning Rate
Perplexity
30-
0.0015
20G
0.0010
25-
0.0005
optimality
0.0000
5000
10000
15000
20000
0
5000
10000
15000
20000
Steps
Steps
Figure 1: Revisiting cosine optimality for language models. We revisit the observation from Chin- chilla (Hoffmann et al. (2022) that in order to achieve the best model after certain training length (to-
kens), the cosine schedule must match the total duration training. This comes at the cost of neither
being able to stop before or going beyond the cycle an issue we show how to alleviate in Section 3
cooldowns. Going beyond, we investigate stochastic weight averaging (Izmailov et al. 2018 and
a schedule-free optimizer (Defazio et al.) 2024), which give strong (but not optimal) performance at any point during training and can act as replacement for the learning rate decay, if the performance gap is acceptable and avoiding separate cooldowns is preferred.
These findings suggest that research on training recipes and scaling laws has been needlessly complex due the need to retrain models from scratch. We demonstrate this empirically by performing small-scale experiment of scaling laws which only uses a fraction of the compute and GPU hours that were previously needed. Following, we discuss that this makes scaling research more accessible and enables more frequent computation of laws for data mixtures (Bi et al.
2024 Goyal et al. 2024
Aghajanyan et al. 2023) or novel architectures (Gu & Dao 2023
De al. 2024).
Background: Cosine Learning Rate Schedule for LLMs
Revisiting the optimality of the cosine schedule. We start our argument by revisiting the use of the cosine schedule in large language model (LLM) training. For any machine learning model, the learning rate value (LR) and schedule both are crucial choices for training. From optimization theory, our understanding is that a slow annealing of the learning rate is essential to find good minima in the loss landscape particularly for deep networks, whereas higher values help exploration (Smith jet al. 2017 Loshchilov & Hutter 2016).
In the context of LLMs, the most commonly used cosine strategy presents particular trade-off by which the L. reaches its maximum early after the warm-up stage and then gradually decreases, typically to 10% of the maximum LR (see Figure [ right). Since the seminal works of GPT (Radford cet al. 2018 2019 Brown et al. 2020) and large models like Gopher (Rac et al., 2021). PaLM2 (Anil ct al. 2023 or LLaMA Touvron et al. 2023a b), cosine has stayed the de-facto standard schedule.
Experimental setup. Our first goal to understand the importance of the length of the schedule for performance of the model. To this end, we implement the common decoder-only transformer (Vaswani et al.] 2017) identical the LLaMa (Touvron et al., 2023a b) or Noam architecture Ormazabal et al. (2024). Throughout this paper, we use the Adam W optimizer with weight decay (Kingma & Bal 2014; Loshchilov Hutter 2017) with common L.M training parameters. We train on a subset of SlimPajama (Soboleva et al. 2023) with 6B tokens]] cleaned and deduplicated corpus for LLM pretraining, which we split into train and validation sequences and report validation loss (perplexity). We provide all the details in Appendix [.I]
The pitfalls of cosine. In the results of Figure we see that the key parameter for the cosine schedule is the length of training: At specific step counts, the best perplexity is always achieved by the cosine schedule that matches the length. This is the main observation of Hoffmann et al. (2022) in order to achieve the best model for a specific token count, the training duration must be known in advance and match the cosine length. However, this brings particular issues. First, cosine suboptimal during training and underestimates the model's performance for the same token count. At the same time, cosine strongly complicates continuation of training. For example, one could easily be mistaken

---

to extrapolate a loss curve of a cosine schedule beyond the end of the cycle. The improvement in loss precisely happens because of the LR decay; afterwards, the final learning rate will generally be too low to continue making large progress. In contrast, rewarming leads to spikes of which the training only slowly recovers, similarly reported in the continual learning literature (lbrahim et al. 2024).
3 A Different Route: Constant Learning Rate with Cooldown
Why does the cosine schedule with a single cycle work well for LLM training? Arguably, it provides good trade-off between high learning rate and cool- ing down the model sufficiently, which is expanded proportionally to the training steps.
1.00
Learning rate
0.75
10%
0.50
Casine
1 Sapt
The alternative constant cooldown. The trade- off of cosine can also be achieved by a different sched- ule. Here, the learning rate (LR) is kept constant for the majority of the training and only decreases in a short final phase, known the cooldown (or decay/annealing phase). This schedule has previ- ously been referred as trapezoidal (Zhai et al. schedule
0.25
Cooldww
0.00
200
400
600
1000
Steps
Illustration of schedules. Cosine (red) follows a slow decrease in learning rate, typically 10% of the maximum for LLMs. The alternative is characterized by an aggres- sive decrease in learning rate, e.g., via a linear (blue) or square root (yellow) cooldown.
Figure 2:
2022). and later as the warmup-stable-dess (WSD) (Hu et al. 2024). To avoid overloading of terms (e.g., weight decay), we refer to this approach constant LR cooldown.
The cooldown phase typically has the LR go to zero, mirroring the warmup phase, which gives rise to an overall trapezoidal shape (see Figure Formally, we can define it as
Naawn
7max
if n N Nwarmup
nn()
mmax
N warmup < n S Naecay
fnn, N, Ndscay) 7max
if > N - Naceay
with the peak learning rate Tmaxx the total steps N with warmup and cooldown steps Nammup and Ndccay. and monotonically decreasing function [n, N, Naccay) that handles the cooldown.
3.1 The Advantages
The main advantage of the constant schedule that it does not require one to specify the number of training steps in advance. This is particularly convenient for large runs, as the cooldown can be initiated at any time to observe model behavior and decide whether to stop. We want to highlight the following in more detail.
Continual learning. The constant + cooldown schedule allows for continual learning by default. Here, the natural approach is to use checkpoints before the cooldown to continue training with high LR: this avoids loss spikes and brittle training when rewarming the learning rate (cf. Figure Tith cosine). Moreover, rewarming has been reported to hurt performance and introduce forgetting compared single annealing training (Ibrahim al. 2024), though careful strategies can alleviate such issues. It remains an interesting question if a single cooldown schedule is absolutely optimal given a total compute budget for LLM training.
During the cooling phase, the data mixture can be changed (Hu et al. 2024) as form of finetuning; beyond aiming for specific downstream tasks, this allows to assess the quality of specific mixes, e.g., as recently done for Llama 3 (Dubey et al.] (024).. Although we focus on the same mixture, understanding the currieulum aspect of the separate phases is explored in the concurrent work of Blakeney et al. (2024) and remains open for further research.
Data mixtures.
Scaling studies. Since cooldown can also be initiated retrospectively from checkpoint, the schedule allows to see the same model at different scales of compute. This enables much cheaper scaling studies than commonly done; we show this in depth in Section
3.2 Experimental Comparison
We follow our experimental setup from the previous section (details in A. and train 210M parameter model on SlimPajama with constant LR and the cooldown schedule defined in (.. That is,
3

---

Validation (210M)
LR Sensitivity (Cosine vs, Cooldown)
35
Coe 22k
Cooldown 22k
24
Couldvens (238) E3%
Cos lilk
Cooldown so 3ik
Cuttze 33k
Colltow (2175) 23%
Perplexity
30
Cos 4ck
Cooldown to 1ik
Perplexity
22
Chise -4k
Caniiinn (119)) 44k
25
20
20
18
15
0
10000
20000
30000
40000
0.0001
0.0002
0.0005
0.001
0.002
Steps
Learning Rate
Figure 3: The difference in loss curves of cosine vs. constant learning rate with cooldown. The cooldown phase initiates a sharp decrease in loss (left) to match cosine; the training perplexity follows the same behavior (Fig. [5). We find the LR sensitivity (right) to be similar for both schedules, albeit less for cooldown, where the optimum lies slightly below at half of the optimal cosine maximum LR.
(1 Sqrt.) vs Linear decay
Cosine (1 Sqrt) & Linear decay
22
24
Linear
Caita
1 Syp
Parplexity
22
Perplexity
20
1- Sett Cuclloen
20
18
18-
16-
50000
100000 150000 200000
140000
160000
180000
200000
Steps
Steps
Figure 4: A different cooldown schedule can improve performance. Perhaps surprisingly, find that a different decay phase in the functional form of (I-sqrt) can consistently outperform the standard linear decay, where both are better than the chosen (untuned) cosine for long lengths.
we compare the same length of warmup and training, but replace the cosine decay after warmup with a constant LR and a cooldown for 20% of steps linearly going to zero. We additionally sweep the maximum learning rate for both approaches. In the results shown in Figure 3 perhaps suprisingly, observe an almost perfect match between the performance of the best cosine and cooldown schedule even for different training durations, all while exhibiting slightly less sensitivity to variations in the LR.
Different cooldown schedules. We investigate various functions describing the cooldown shape, including cosine, square, and square root shapes in Appendix B. We identify a new function (I-qqri) that outperforms the linear cooldown. This improvement is maintained in a smaller number of decay steps, different learning rates, and various timestamps (see Appendix [. I). In Figure 4 we train model for 200K steps (approximately 20B tokens), applying cooldown every 20K steps for 20%. The results indicate that the longer the training duration, the more linear cooldown is outperformed by the [I-qrt) cooldown, which we define as:
f(n, N, Naccoy = (1
(N Namn
(1-sqrt)
Navay
Takeaway 1:
The constant LR- + cooldown schedule offers significant convenience by not requiring the number of training steps to be specified in advance, and provides similar performance compared to a well tuned cosine schedule.
Takeaway 2: We find cooldown form -sqrtt) that consistently performs better than linear decay.
How long do you need to cooldown? To effectively utilize the schedule, it is essential to determine the optimal number of decay steps. Our study of the relative number of cooldown steps, as shown in Fig. [5] reveals that the benefits of extended cooldown periods plateau at around 20%, which we select for the remaining experiments. Additionally, decay with the (I-sqrt) cooldown can nearly match the performance of the cosine schedule on 20B token run (much beyond Chinchilla optimal). This practically important, as we ideally keep the cooldown as short as possible but get strong performance.
Fig. 6 we demonstrate that using only 5%

---

Fraction Cooldown Length (124M)
Fraction Cooldown Leugth (124M, 25k)
24
24
1ia Stepa
L.000000G
Perplexity
22-
23
Guskse
Steps
Stepe
Stugo
Perplexity
22
23
LRR: 0.002
Cioisec
0.031
21
20
20-
19
19
0.0
0.2
0.4
0.6
0.8
0.0
0.2
0.4
0.6
0.8
Fraction of Cooldown Steps
Fraction of Couldown Steps
Figure 5: Longer cooldown helps to achieve lower loss. We investigate the effect of the cooldown length fraction of the total steps for a 124M model. We find that the cooldown surpasses cosine between 10-20% of steps (left), but largely stops improving when done over a majority of training. This also holds when sweeping the LR (right). Additional ablations are provided in Appendix B.1
10k decay steps (5%)
10k decay steps (5%)
24
22
Costne
Cotlie
22
Lisear Cooddown
Linoe Couldown
Perplexity
1 Sqrt Cooldown
Perplexity
20
Sert Cookiown
20
18
18-
16
50000
100000
150000 200000
185000
190000
195000 200000 205000
Steps
Steps
Figure 6: A long training run suggests that small number of cooldown steps can match cosine for long training. From Fig. 5 and Fig. 20 we find that the required duration of cooldown to match the cosine loss decreases with longer training, we validate with a long training run (200k steps. same LRs), and find that just 10k cooldown steps almost perfectly match cosine in performance.
Takeaway 3: The constant learning rate with a short cooldown (< 20% of total training steps) can achieve the same final loss as cosine, and only outperforms for longer fraction of cooldown steps. For long training runs, our results suggest that the cooldown length can be less than 20% to match cosine if it is enough in absolute steps see Figure 6
What happens during the cooldown? It is remarkable how the sudden drop in loss is consistent for both train and validation loss and aligns closely with the decay in LR (Figure [5) Hu et al. (2024) investigate the cooldown phase and find that the first-order directional derivative diminishes with each step, whereas the curvature of the loss function increases; they attribute this to proximity to local optimum.
We expand upon these results and aim to understand the optimization landscape around the trajectory of the cooldown. For that, we evaluate the loss along the straight line trajectory when moving between a checkpoint before and after cooldown, i.c., a linear interpolation of the weights of the two models. Perhaps surprisingly, we find that the smooth drop in loss also occurs for this interpolation, as visualized in Figure) 7 This aligns with the findings of Hu et al. (2024). These results suggest that upon decaying the LR, the model immediately descends into connected minimum of the loss.
Takeaway 4: The cooldown phase is smooth transition basin in the loss landscape.
Decaying cosine to less than 10%. Throughout our experiments, follow the common heuristic of setting the final LR of cosine to 10% of the maximum. We also provide experiments in which we ablate annealing to zero (or a very small value) in Fig 22 in Appendix B.] We find that, perhaps unsurprisingly, the influence of the final LR is negligbgee decaying to lower values improves performance to match the best (I-sqrt) cooldown. However, when evaluating on downstream benchmarks (see next Section 3.3), we see that annealing to zero can hurt metrics; we posit this comes from too early saturation. Combining these two arguments implies that the maximum and final LR should be set (and ideally swept over) independently.
5

---

Train Loss Connectivity (100 batches)
Val Loss Connectivity
22k Hee
22k iter
24
24
33k iter
33k iter
Perplexityy
44k iter
Perplexity
44k iee
22
22
20
20-
18
18
0.0
0.2
0.4
0.6 0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Interpolation Weight
Interpolation Weight
Figure 7: The smooth drop in loss also occurs when moving linearly in weight space between checkpoints before and after the cooldown. This suggests that in the cooldown phase, the model directly moves within connected basin in the loss landscape.
Aggr. Score (1B model, 100B tokens)
Aggregated Score
0.45
Score
Cosine to 10%
46.26
0.40
Cosine to 0
45.88
Cudnr to 10% Liuar Cookdown (205))
Lincar (20%)
46.20
0.35
(1-Sqrt) (20%)
46.23
0
20
40
60
80
100
Tokens (B)
(a) Metrics throughout training.
(b) Final scores (100B tokens).
Figure 8: The performances of both schedules also match for downstream tasks. We run a realistic setting of IB model trained for 100B tokens of Web and establish matching performance of both schedules for common LLM benchmarks. Interestingly, there is similar boost in performance with the cooldown. Detailed numbers over the course of training and 460B token runs are in Appx. B.5
3.3 Scaling Up: 1B and 8B Models
We expand our ablations to realistic scenarios of IB and 8B parameter models, with an additional focus on model evaluations. Our goal in this section is therefore twofold: first, we validate the dynamics of the schedule at practical scale; second, we establish the connection between the pure loss values and downstream benchmarks. Although tasks often scale reliably with loss (Du et al. 2024 Gadre et al. 2024]. this connection is crucial, as downstream abilities are ultimately the main metric of interest.
1B with downstream evals. We train a 1B parameter model with cosine and a cooldown sched- ule with the high-quality FineWeb corpus (Penedo al., 2024). Following DeepSeek scaling laws (Bi et al. 2024). we set an approximately optimal batch size and LR of 1.8M and 8e-4 for 100B and 460B token runs; evaluation is done throughout training on the most common bench- marks like MMLU (Hendrycks et al. 2021). The setup is described in detail in Appx. A.2
With the aggregated results Fig. we find a clear agreement of the downstream performance for both schedules. Interestingly, there is an uptick for some
8B Model Run
30
Training Perplexity
25
evaluations, similar to the drop in loss, with the start of the LR cooldown; see Fig. 29 However, not all
benchmarks follow this trend, which requires further
20
research. The results for the much longer 460B run, where performance also matches, and detailed benchmark numbers are provided in Appx. B..
Codtse
15
1-Seet (20%)
0
10
12
Tokens (B)
8B. While a full 8B training run is beyond our limits, we provide a first investigation in Fig. (9) In line with all our experiments, we see matching loss values and, promisingly, no instability at such a large scale, though the two runs are clearly short (20k steps). Moreover, instabilities arising from a high LR for longer training can be alleviated by methods like QK norm (Wortsman et al. 2024
Figure Results at 8B scale, We validate the behavior with much larger model (architecture of Llama 3) for a single short run (12B tokens of FineWeb-Edu), where the cooldown matches the cosine schedule again.

---

SWA A of Long Constant LR (210M)
SWA of Long Cosine LR (210M)
30
30
Constant
Cusine LR
SWA Conseant LR
SWA of Longest Cusine LR
Perplexity
Perplexity
25
25-
20
20
0
20000 40000
60000
0
20000
40000
60000
Steps
Steps
Figure 10: SWA improves generalization and simulates decayed learning rates. Using SWA for the constant LR phase (left) strongly boosts the loss, but a gap to the cooldown remains, SWA also improves the generalization of cosine schedule (right), where the intermediate checkpoints of SWA largely overlap with optimal loss trajectory of shorter cosine runs.
4 Do We Even Need to Cooldown?
In Section 3 we show that a constant LR with a short cooldown phase can replace the cosine decay. Ideally, however, we would not need a decay at all, but aim to obtain an optimal model any point in training, even when prolonged. This would save even more computational resources and time. In this section, we investigate two potential approaches: weight averaging and schedule-free optimizer.
4.1 Stochastic Weight Averaging (SWA)
Motivation. While slow annealing of the learning rate can be essential to find good minima (Smith et al. 2017: Loshchilov & Hutter 2016). (Sandle et al. (2023) show theoretical and empirical equivalence between stochastic weight averaging (SWA) and cooldowns in training vision models. There, intuitively and naturally averaging reduces noise and thus improves generalization (Wortsman et al. 2022). Motivated by these results, we aim to answer the same question in LLM training: Can weight averaging replace the cooldown phase?
Method. We opt for a form of SWA (Izmailov et al. 2018) that splits the training in fixed windows and averages within window, allowing us to keep the average as single additional copy of the model parameters. In our experiments, we set the window to h = 500 steps and save the window averages as checkpoints every h steps, which allows ad hoc evaluation of longer windows to see potential improvements akin to latest weight averaging (LAWA, Kaddour 2022] Sanyal et al.) 2023). For all our experiments, we find that windows below at 2500 steps (256M tokens) are optimal. We also experimented with an exponential moving average (EMA), which performed worse than SWA, and therefore we do not report EMA.
SFO vs Linear Cooldown
Experimental results. We evaluate SWA with the same 210M model and show the results in Fig. 10 for
24
Choldons (0) -1.0000-0000
320 (8, - 0.5LS - 6.0)
both a constant LR (left) and cosine (right). Notably,
Perplexity
22
Cooldown (s) 630.9 (350)
590 0.95)
we find significant performance boost for SWA on top of a constant LR. Yet, it does not reach the loss values of explicit cooldowns. On the other hand.
20
in line with previous work (Kaddour 2022 Sanyal see a
ee al. 2023 Andriushchenko et al.. (023), similar boost for SWA on top of cosine. This suggests that, regardless of schedule, SWA is a compelling approach to achieving strong models along the data- scale axis that can serve as a replacement for models trained with fewer steps, if the performance gap is ac- ceptable and one wishes to avoid cooldowns. This particularly ad advantageous as it can arguably be done free on top of existing well-tuned optimizers.
16
0
25000
50000
75000
100000
Steps
Figure 11: The cooldown schedule outper- forms SFO even when tuning momentum parameters. We find that SFO is sensitive to the choice of (3), 32) momentum parameters. It gives strong performance for well-tuned momentum, but falls short of cooldown.
Takeaway 5:
Irrespective of the schedule and without additional overhead, SWA improves performance along the training trajectory and provides better models during training. While it does not match the cooldown, it reduces the gap without the need for a separate decay phase.
7

---

4.2 Schedule-Free Optimizer (SFO)
Very recently, Defazio et al. (2024) introduced schedule-free optimizer (SFO) which uses an interpolation between standard averaging and accelerated method (Nesterov) (983). As such, the optimizer does not require a decreasing learning rate schedule, making it relevant for continual training. We seek to investigate if it can outperform the presented cooldown schedule and provide comparison in the same LLM setting.
averaging, inspired by Nesterov's
Results. We compare the results of a long 210M model run with cooldown vs. SFO with Adam W in Figure [1 Although the optimizer does not require learning rate schedule, the authors point out that training is more sensitive to the choice of the momentum parameters (2), (which are not identical to the momentum in Adam). They note that the optimal parameters may depend on the length of training. making it not fully schedule-free. We observe this sensitivity in our experiments, where the choice of (0.9, 0.95) performs significantly worse and even increases the loss toward the end of training. For (31 = 0.95, B2 = 0.99), SFO performs remarkably well. Nevertheless, both settings matched or outperformed by the cooldown schedule, in particular when comparing the same momentum configuration. We did not perform any further hyperparameter tuning for either method.
5 The Implications for Scaling Law Research
Our investigation shows the efficacy of a constant LR for flexible training of LLMs by introducing short cooldown phase at any point during training. In this section, we go beyond a single model, and focus on the reliability of the alternative LR schedule compared to cosine across model sizes and scales. Ultimately, we discuss the importance for the future of scaling law experiments.
Importance of scaling laws. At their core, scaling laws aim to establish a functional form of a model's performance, commonly modelled as the loss L, as function of the parameters N or training tokens D; for LLMs, this is usually expressed as the power law
A B
L(N, D)
+
Na
D3
where (A, 3, E} are variables be estimated (Kaplan et al. 2020). Such scaling laws serve multitude of critical purposes- - from optimally spending fixed amount of compute (FLOPs) for achieving the lowest loss, to trading off data-sourees of different quality (Bi al. [024]] Goyal et al.)
2024) or modalities (Aghajanyan et al. (2023), to comparing the efficiency of different architectures
(Gu & Dao 2023 De 2024).
Crucially, Hoffmann et al. (2022) demonstrated that is necessary to vary the number of training
tokens for fixed family of models and fully decay the LR (see also Kaddour al. 2023). At the time,
their results suggested that LLMs were over-sized, leading to substantial increase in data collection
and training for much longer, beyond the Chinchilla optimal point (N, D) (Touvron et al. 2023a b).
Why do the presented results matter for scaling? Following the results of Chinchilla, scaling laws require family of models each trained from scratch with a cosine schedule that is fit to different training lengths (cf. Section 2). In contrast, the cooldown schedule as well as weight averaging allow much cheaper alternative in two phases: first, model sweep with single sufficiently long training run for each model size in the family; then, using the model checkpoints to perform a cooldown or averaging. This reduces scaling law experiments to only the model scaling axis, effectively dividing the number of necessary training runs by one order of magnitude. At the same time, it allows for flexible continual training beyond any predetermined number of steps.
Experimental setup. We mimic a small-scale experimental setup for scaling laws: We train a range of model sizes (33M-360M) across different token scales (0.3B-10B) on the same SlimPajama 6B dataset. For each model, we choose exactly three token counts (around the Chinchilla optimal ratio of D/N=20) in increments of 10, 20 and 30 tokens per parameter. With the cosine schedule, each model is trained from scratch three times. In contrast, for averaging, we train each model just once for the longest cycle and then use the averages at the same token count as cosine; for cooldown, we similarly take checkpoints along the constant LR trajectory and perform three annealing periods to match the token counts. We adjust the LR to be higher for smaller models. In line with the findings of Sect. 3 and for a fair comparison, we set the constant LR to be half the maximum LR for cosine and perform 20% (linear) cooldown steps. More details are given in Appendix A.
Results. We show the validation loss envelopes as a function of the training FLOPs in Fig. 12 (left) and compare each obtained model (i.c., same parameter & tokens) with cosine and its alternatives
8

---

Cosine Schedule (Chinchills)
Cooldowa Schedille
Cooldown ve
Constaat ve
6.0
6.0
5.0
181
5.0
151M
Conknn Perplesity
Libear
10
1300
12402
Loos
4.0-
6354
30-
300
5334
3.344
10
10
10
10
PLOPS
FLOPS
PPLX Cooldown
PPLX SWA Constant
Figure 12: Cooldown LR schedule SWA scale reliably. We train range of models (33M-360M),
each for three different cosine lengths. We then run each model with a constant LR just once and take the snapshots at the same token count as cosine (for SWA) or perform post-train cooldowns to the same length. Each final model is represented by a dot. Left: The loss curve envelopes. Right: The perplexity of cosine (y-axis) vs. cooldowns and SWA. Points on the diagonal indicate the same loss for both methods; above outperforming cosine, below fall short. We see alignment between both methods.
Total FLOPS spent
Total GPU Hours for all Models
1x
Goske L.
Gare S.
GPU Hours
PLOPS
200-
FLOPs
1'000"
1000
Chinchilla
5.59 x 1023
w/ Cooldown
2.36 x 1023
(a) FL.PPs and GPU hours for our models.
(b) Estimated FL.PPs savings for Chinchilla.
Figure 13: Scaling laws for fraction of the cost. The reliable behavior of both the cooldown schedule and SWA allows scaling experiments with a drastic reduce in both compute and GPU hours; in both our experiments (left) and the original Chinchilla (right) factor of or less. The more training runs are performed per model size (e.g. 4 for Chinchilla), the larger the difference becomes.
(right). We find that the cooldown learning rate schedule scales reliably, much like cosine and SWA, to achieve optimal losses with the recipe we establish in previous sections. This is particularly visible in Fig. 12 (right), where each model's performance lies almost perfectly on the diagonal, which signals that cosine (y-axis) and cooldown (x-axis) reach the same loss after the same amount of tokens. A similar result is visible for SWA, though the reciprocal (y-offset) is negative, which agrees with our findings from Sect. 4.. that the averaging does not fully close the gap to a LR cooldown.
Compute savings. Crucially, the alternatives enable scaling laws for just a fraction of the cost. In Fig. [3a we report FL.PPs and GPU hours (real wall-clock time) for all model runs for both methods. They substantially reduce both compute and GPU hours. In our experiments, where we space the runs to token ratios of 10, 20 and 30, it saves half the time and FLOPs, enabling scaling laws for only a fraction of the previous cost. We report the detailed savings for all models in Appendix B.2
Estimating savings for Chinchilla. We take our analysis further and estimate how much cheaper the Chinchilla model suite would have been if performed with 10% cooldowns after a single run for each model size using the model configurations reported in Table A9 (Hoffmann et al. 2022). Since the authors do not report exact training configurations. we consider a sequence length of 1024, batch size of 0.5M tokens and token ratios M D/N E {10, 15, 20, 25} for each model. With this, we arrive at roughly 5.59 x 202 total FL.OPs originally vs. 2.36 x 1023 (Figure 13b). This means that less than half the compute could have been used.
Takeaway 6: Scaling experiments can be done with significantly reduced compute and GPU hours by utilizing fewer but reusable training runs with a constant learning rate and ad-hoc cooldowns.
Additional results. We plot the training curves of all models in Appx. B.3 In addition, we show the same findings with experiments on Open WebText2 in Appx. B.4
6 Limitations
We conduct our experiments on models of up to 8B parameters, with long training runs of a 1B model on multiple hundred tokens. The trends we find are consistent across all scales, but training

---

behavior can be more brittle at modern scales and extremely long training (Wei et al. 2022 Tay al. 202Tp. However, instabilities arising from a high learning rate for a large part of training can be alleviated (Wortsman et al. 2024).
7 Related work
Cosine Schedules and Alternatives for Transformers. The cosine decay was originally introduced by Loshchilov & Hutter (2016) for cyclic schedules in vision tasks, where it is common practice to
have stepwise or cyclic LRs to escape bad minima when training multiple epochs (Smith et al.) 2017).
For language models where data more vast, the cosine schedule with a single cycle is currently the de-facto standard for training, with few exceptions of T5 (Raffel et al. (2200) and PaLM1 (Chowdhery
al.
2023) that used form of inverse square root. In line with our work, recently released models opt
For alternatives such as stepwise schedules (Bi et al.) (024) or the presented constant cooldown (Shen
2024] Hu et al. 2024). These alternatives were previously also explored for vision transformers by /Zhai et al. (2022), who find reciprocal square-root with cooldown to perform best, and in the
al.
context of continual learning (Ibrahim et al.) 2024] Gupta et al. 2023). Defazio al. (2023) investigate
the gap between theory and practice of LR schedules and suggest that a linear decay is optimal, also
for LLM training. We similarly find that a long linear decay can slightly outperform cosine.
Weight Averaging, Weight averaging over past iterates (Polyak & Juditsky 1992) has long been known to be beneficial for convergence. [Imailov et al. (2018) introduce stochastic weight averaging for better generalization in deep learning models. Similarly, the exponential moving average is
commonly used in vision (Morales-Brotowe et al., 2024). Importantly, (andler et al. (2023) show the equivalence of WA with decaying learning rate schedules. In the context of LLM training, close to our work is |Sanyal et al.] (2023) which showed that form of latest averaging (Kaddour 2022) can be used to improve the performance of models early in training. However, they did not investigate the relation of weight averaging to compute optimality and its implications for scaling experiments.
Scaling Law Experiments for Neural Language Models. Kaplan et al. (2020) were the first to estab- lish scaling laws for language models. Important ings and demonstrate specific methods to establish scaling laws, notably training a family of models for different cosine lengths. The subsequent models like 1/amal/2 (Touvron et al. 2023a b) further improve performance of smaller models by training beyond the Chinchilla optimal point, motivated
Hofffmnn et al. (2022) revise these find-
by lower inference costs (Gadre et al. 2024 De Vries 2023 Sardana & Frankle 2023). Recent works (Muennighoff et al. 2023; Bi et al. [2024]] Goyal et al. 2024) highlight how data repetition and quality affect the scaling behavior, which suggests that scaling laws should be updated more frequently. How- ever, these works do not consider efficient experiments for scaling laws, which is the focus of our work.
In concurrent work, Porian al. (2024) and PParce & Song (0024) find that the discrepancy between the Kaplan and Chinchilla scaling laws is not attributed to the learning rate decay, but they establish that the optimal poren//arameter ratio can be obtained with constant learning rate without any cooldown; however, the actual performance is then suboptimal, and cooldown schedule as suggested in our work is needed to properly estimate model performance, in particular for downstream tasks.
8 Conclusion
We have demonstrated the reliability of an alternative learning rate schedule to replace cosine for LLM training, which uses a constant rate with a cooldown phase. Across a multitude of experiments, analyze different recipes for the decay form and length. Importantly, established the best learning rate schedule - instead, we investigate and demonstrate how an arguably simple recipe can match the performance of the current best practice of cosine, and discuss how it provides compelling advantages such as continual training and strong reduction in costs for scaling law research. In addition, we find that SWA can give reliable (strong, but not optimal) estimates of models during runs, without additional overhead or training.
do not claim to have
We believe the results are of great importance to the present and future of LLM training: The presented methods facilitate research for the current post- cost-Chinchilla era, where models are trained much beyond compute-optimall by allowing more flexibility to continue training whenever needed. At the same time, recent results that suggest data dependency in scaling (Bi et al. 2024 Goyal et al. 2024 Aghajanyya et al.) 2023 Pandey 2024) imply the need to frequently update scaling laws, which is economically more feasible with reduced costs. We therefore hope that our work will make scaling research more accessible researchers and practitioners alike.
10

---

Acknowledgements
We thank Anastasia Koloskova, Amirkeivan Mohtashami and Thomas Wolf for helpful discussions regarding the paper and its implications. We also thank Stella Biderman, whose open discussion ? inspired parts of our experiments.
References
Aghajanyan, A., Yu, L., Conneau, A., Hsu, W.-N., Hambardzumyan, K., Zhang, S., Roller, S., Goyal,
N., Levy, O., and Zettlemoyer, L. Scaling laws for generative mixed-modal language models. In International Conference on Machine Learning, pp. 265 2669 PMLR, 2023.
Andriushchenko, M., D' Angelo, F., Varre, A., and Flammarion, N. Why do we need weight decay in modern deep learning? Oct 2023. URL http:tttpt//i..oooo org ab1/ 04415v1
Anil, R., Dai, A. M., Firat, O., Johnson, M., Lepikhin, D., Passos, A., Shakeri, S., Taropa, E., Bailey, P., Chen, Z., et al. Palm 2 technical report. arXiv preprint arXiv:2305. (00:3055000000 10403, 2023.
Bi, X., Chen, D., Chen, G., Chen, S., Dai, D., Deng, C., Ding, H., Dong, K., Du. Q.. Fu, Z.. et al. Deepseek I1m: Scaling open-source language models with longtermism. arXiv preprint arXiv:2401.02954, 2024.
Bisk, Y., Zellers, R., Bras, R. L., Gao, J., and Choi, Y. Piqa: Reasoning about physical commonsense in natural language. Thirty-Four Fourth AAAI Conference on Artificial Intelligence, 2020.
Blakeney, C., Paul, M., Larsen, B. W., Owen, S., and Frankle, J. Does your data spark joy? perfor- mance gains from domain upsampling at the end of training. arXiv preprint arXiv:2406.0344 2024.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P.. Sastry, G., Askell, A., et Language models are few-shot learners. Advances neural information processing systems, 33:187 13:1877-1901, 2020.
Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-113, 2023.
Clark, P., Cowhey, I,, Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.000 05457,, 2018.
Dao, T., Fu, D., Ermon, S., Rudra, A., and Ré, C. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 15:16344-16359 2022.
De. S., Smith, S. L., Fernando, A., Botev, A., Cristian-Muaary G., Gu, A., Haroun, R., Berrada, L., Chen, Y.. Srinivasan, S.. Desjardins, G., Doucet, A., Budden, D., Teh, Y. W., Pascanu, R., Freitas, N. D., and Gulcehre, C. Griffin: Mixing gated linear recurrences with local attention for efficient language models. Feb 2024. URL http:///wxxxv..ooa////////a hbtp:bbb22202222927/0 httt:://aaaa...r//bb///22/0222
De Vries, H. Go smol or go home, 2023, URL https:/www.hardderrtttt model-size-vee compute- overtead/
Defazio, A., Cutkosky, A., Mehta, H., and Mishchenko, K. When, why and how much? adaptive learn- httg//bb/220 2ttpp//00.0.000.11221111
ing rate scheduling by refinement. Oct 2023. URL httltt/pp///xiii
Defazio, A., Yang, X., Mehta, H., Mishchenko, K., Khaled, A., and Cutkosky, A. The Road Less Scheduled. May 2024. URL
15682v1
Du, Z., Zeng, A., Dong, Y., and Tang, J. Understanding emergent emergnt abilities of language models from the loss perspective, arXiv preprint 1776, 2024.

---

Dubey, A., Jauhri, A., Pandey. A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A.,
Yang, A., Fan, A., et al. The llama herd of models. arXiv preprint arXiv:2407.2178,, 2024.
Fourrier, C., Habib, N., Wolf, T., and Tunstall, L. Lighteval: A lightweight framework for Ilm evaluation, 2023. URL https:t/t:/// github...t....
Gadre, S. Y., Smyrnis, G., Shankar, V., Gururangan, S., Wortsman, M., Shao, R., Mercat, J., Fang,
A. Li, J., Keh, S., Xin, R., Nezhurina, M., Vasiljevic, I,, Jitsev, J., Dimakis, A. G., Ilharco, G., Song, S., Kollar, T., Carmon, Y., Dave, A., Heckel, R., Muennighoff, N., and Schmidt, Language models scale reliably with over-training and downstream tasks. Mar 2024. URL attp://arxiig wwwpsiiii.rrr.eeooooe hat///b//20/ 2403 08540v1
Gao, L., Biderman. S., Black, S., Golding, L.,, Hoppe, T., Foster, C., Phang, J., He, H., Thite, A., Nabeshima, N., Presser, S., and Leahy, C. The Pile: An 800GB Dataset of Diverse Text for Language Modeling, arXiv preprint arXiv:2101.00027, 2020.
Goyal, S., Maini, P., Lipton, Z. C., Raghumathan, A., and Kolter, J. Z. Scaling Laws for Data Filtering-Data Curation cannot be Compute Agnostic. arXiv preprint arXiv:2404.07177, 2024.
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces, Dec 2023. URL http://arxiv. http://...rr/abbb21..0.0000 0072vv
Gupta, K., Thérien, B., Ibrahim, A., Richter, M. L., Anthony, Q., Belilovsky, E., Rish, I., and Lesort, T. Continual pre-training of large language models: How (re)warm your model? Aug 2023.
URL htt:ttp: /taxxiv//xxxv... arxiv.iiivvv... wrg/abs/////eoooogggo....g...a.r abs 2abs/23088 2308. 04014v2.
Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, J. Mea- suring Massive Multitask Language Understanding. In International Conference on Learning Representations, 2021.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, d. L., Hendricks, L. A., Welbl, J., Clark, A., et al. Training compute-optinas large language models. Advances in Neural Information Processing Systems, 35:30016-30030, 2022.
Hu, S., Tu, Y., Han, X., He, C., Cui, G., Long, X., Zheng, Z., Fang. Y., Huang, Y.. Zhao, W., Zhang, X., Thai, Z. L., Zhang, K., Wang, C., Yao, Y., Zhao, C., Zhou, J., Cai, Zhai Z., Ding, N., Jia, C., Zeng, G., Li, D., Liu, Z., and Sun, M. Minicpm: Unveiling the potential of small language models with scalable training strategies. Apr 2024. URL (tLtattt:://aat www.aariiii.ooggar.aaaaa/ wbwprrg////aaaaeeeeeeeee ab//a/b/22444443550 06395v2
Ibrahim, A., Thérien, B., Gupta, K., Richter, M. L., Anthony, Q., Lesort, T., Belilovsky, E., and Rish, I. Simple and scalable strategies continually pre-train large language models. Mar 2024. URL https://arxiv.. org/ abs////s/// abtpa/bs.../////////// 2403 08763v3
Izmailov, P., Podoprikhin, D., Garipov, T. Vetrov, D., and Wilson, A. G. Averaging weights leads to wider optima and better generalization. Mar 2018. URL fttt:://axxii....t///// org/ /asb//113 05407v3
Kaddour, Stop wasting my time! saving days of imagenet and bert training with latest weight averaging. Sep 2022. httptttpp/arxii.rrggeee a20p: .299//9999.9191
Kaddour, J., Key, O., Nawrot, P., Minervini, P., and Kusner, M. J. No train no gain: Revisiting efficient training algorithms for transformer-ban language models. Advances in Neural Information Processing Systems, 36, 2023.
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford,
A. Wu, J., and Amodei, D. Scaling laws for neural language models. Jan 2020. URL U
7/arxiv. org/ 2000 08361v1
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint 1412.644 2014.
Loshchilov, and Hutter, F. Sgdr: Stochastic gradient descent with warm restarts. Xrviv preprint 038.088. 2016.
12

---

Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. arXiv preprint arX1v:711.000000011111 2017.
Micikevicius, P., Narang, S., Alben, J., Diamos, G., Elsen, E., Garcia, D., Ginsburg, B., Houston, M.,
Kuchaiev, O., Venkatesh, G., et al. Mixed precision training. arXiv preprint arXiv:
2017.
Mihaylov, T., Clark, P., Khot, T., and Sabharwal, A. Can a suit of armor conduct electricity? new dataset for open book question answering. EMNLP, 2018.
Morales-Brotons, D., Vogels, T., and Hendrikx, H. Exponential moving average of weights in deep learning: Dynamics and benefits. Transactions Machine Learning Research, 2024.
Muennighoff, N., Rush, A. M., Barak, B., Scao, T. I.., Piktus, A., Tazi, N., Pyysalo, S., Wolf, T., and Raffel, C. Scaling data-constrained language models. May 2023. URL sattp://arxxxe org/ abs 2 16264v4
Nesterov, Y. A method of solving a convex programming problem with convergence rate o (1/k** 2). volume 269, pp. 543. Russian Academy of Sciences, 1983.
OpenAl. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, abs/2303.08774, 2023. URL atttps:/arxxx org ab////bb//// 200../2333.0000 08774
Ormazabal, A., Zheng, C., d' Autume, C. d. M., Yogatama, D., Fu, D., Ong, D., Chen, E., Lamprecht, E., Pham, H., Ong, I., al. Reka core, flash, and edge: A series of powerful multimodal language models. arXiv preprint wxXi::240.11287,, 2024.
Pandey, R. gzip predicts data-dependent scaling laws. May 2024. URL Utthhtt://aaxiv.orr// 2405. 16684v1
Paszke, A., Gross, S., Chintala, S., Chanan, G., Yang, E., Detito, Z., Lin, Z., Desmaison, A., Antiga, L., and Lerer, A. Automatic differentiation in pytorch. 2017.
Pearce, I. and Song. J. Reconciling kaplan and chinchilla scaling laws. Jun 2024. URL http //arxiv. org/ abs/2466. 12907
Penedo, G., Kydlfček, H., allal. L. B., Lozhkov, A., Mitchell, M., Raffel, C., Werra, L. V., and
Wolf. T. The fineweb datasets: Decanting the web for the finest text data at scale, 2024. URL
atttps://aree wwgpxivoorrr//eo
Polyak, B. T. and Juditsky, A. B. Acceleration of stochastic approximation by averaging. SIAM Journal on Control and Optimization, 30(4))838-855 1992. doi: 10.1 10.1337/033004.
Porian, T., Wortsman, M., Jitsev, J., Schmidt, L.. and Carmon, Y. Resolving discrepancies in compute- optimal scaling of language models. Jun 2024. URL httptpp:///axxx.or////bab///// /20es.2246...e 19146
Radford, A., Narasimhan, K., Salimans, T., Sutskever, I., et al. Improving language under-
standing by generative pre-training 2018. URL https://s3-us-vestt amazonaus
com/openai-oeeee htw.:oponal-assstt/reeeeeeeeee http///wwwrrrooooooo//oggaaaaaaa Lnderstaaniing_pppe. www.:iiinn.ppp...p.......
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Rae. J. W., Borgeaud, S., Cai, T., Millican, K., Hoffmann. J., Song, F., Aslanides, J., Henderson, S., Ring, R., Young, S., et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:21 12. 1446, 2021.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S.. Matena, M., Zhou, Y., Li, W. and Liu. P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140): 2020. http: http::://aaxvv.eee wwg.ziivoor///ba http//bs/19999 10683v4
Sakaguchi, K., Bras, R. I.,, Bhagavatula, C., and Choi, Y. Winogrande: An adversarial winograd schema challenge scale. Communications of the ACM, 64(9):99-106. 2021.
13

---

Sandler, M., Zhmoginov, A., Vladymyrov, M., and Miller, N. Training trajectories, mini-batch losses and the curious role of the learning rate. Jan 2023. URL http://aarxii..orraa/////// 02312v2
Sanyal, S., Neerkaje, A., Kaddour, J., Kumar, A., and Sanghavi, S. Early weight averaging meets high learning rates for IIm pre-training. Jun 2023. URL httt:://aaxi..or///ababb.aa 1rtg0//babb00000002111222
Sap. M.. Rashkin, H., Chen, D., LeBras, R., and Choi, Y. Socialiqa: Commonsense reasoning about social interactions. arXiv preprint arXiv:1904009723 2019.
Sardana, N. and Frankle, J. Beyond chinchilla-optimal: Accounting for inference in language model scaling laws. arXiv preprint arXiv:2401.00448, 2023. URL https///aaxxi..or///bbppp 00448
Shazeer, N. Glu variants improve transformer. Feb 2020. URL 05202v1
Shen, Y., Guo, Z., Cai, T., and Qin, Z. Jetmoe: Reaching Ilama2 performance with 0.1m dollars. Apr 2024. URL http:://axxx wtgg 07413v1
Smith, S. L., Kindermans, P.-.,, Ying, C., and Le. Q. V. Don't decay the learning rate, increase the batch size. arXiv preprini arXiv: 1711:017 2017.
Soboleva, D., Al-Khateeb, F., Myers, R., Steeves, R., Hestness, and Dey, N. SlimPajama: 627B token cleaned and deduplicated version of RedPajama. http://iwwn.t.tppe///////eerrrrr www.bras........................ net/bblgg/ slimpa ama-a- 627b-
June 2023. URL http:////w/gggaggreee////a/a//aa am26677B
Su, J., Ahmed, M., Lu, Y., Pan, S., Bo, W., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063 2024.
Talmor, A., Herzig, J., Lourie, N., and Berant, J. CommonsenseQA: A Question Answering Challenge Targeting Commonsense Knowledge. Burstein, J., Doran, C., and Solorio, (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pp. 4149-4158, Minneapolis, Minnesota, June 2019, Association for Computational Linguistics. doi:
URL htthssps::/// htttt:::://aclaathooogg.g org NN/ 1421
Tay, Y., Dehghani, M., Rao, J., Fedus, W., Abnar, S., Chung, H. W., Narang, S., Yogatama, D., Vaswani, A., and Metzler, D. Scale efficiently: Insights from pre-training and fine-tuning trans- formers. Sep 2021. URL http: org/ aws/ 10686v2
Team, G., Anil, R., Borgeaud, S., Wu, Y., Alayrac, J.-B., Yu, J., Soricut, R., Schalkwyk, J., Dai, A. M., Hauth, A., et al. Gemini: a family of highly capable multimodal models. arXiv preprint 1.Xi::3312.11005 2023.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023a.
Touvron, H., Martin, L.. Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine-tuned chat models, arXiv preprint arXiv:2307.00228
2023b.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, E., and Polosukhin, I. Attention is all you need, Advances in neural information processing systems, 30, 2017.
Wei, J., Tay, Y., Bommmsani, R., Raffel, C., Zoph, B., Borgeaud, S., Yogatama, D., Bosma, M., Zhou, D., Metzler, D., Chi, H., Hashimoto, T., Vinyals, O., Liang, P., Dean, J., and Fedus, W. Emergent abilities of large language models. Jun 2022. URL http: / arxiv org/ abs/ 2206 07682v2
14

---

Wortsman, M., Ilharco, G., Gadre, S. Y., Roclofs, R., Gontijo-Lopes, R., Morcos, A. S., Namkoong, H., Farhadi, A., Carmon, Y. Korbbiihh, S., al. Model soups: averaging weights of multiple fine- tuned models improves accuracy without increasing inference time. In International conference on machine learning, pp. 11963-23996
PMLR, 2022.
Wortsman, M., Liu, J., Xiao, L., Everett, K. E., Alemi, A. A., Adlam, B., Co-Reyes, J. D., Gur, I., J., Xu, K., Lee, J., Gilmer, J., and Kornblith, S. Small-scale proxies for large-scale transformer training instabilities. The Twelfth International Conference on Learning Representations, 2024. URL httppp://ppeereeeeee net/ fettforummid d8w0pmvXbZ
Kumar, A., Novak, R., Pennington, J.,
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. HellaSwag: Can a machine really finish your sentence? In Korhonen, A., Traum, D., and Mårquez, L. (eds.), Proceedings of the 57th Anmual Meeting of the Association for Computational Linguistics, pp. 4791-4800, Florence, Italy, July 2019. Association for Computational Linguisties. doi: 10.18653/v URL 1472
Zhai, X., Kolesnikov, A., Houlsby, N., and Beyer, L. Scaling vision transformers. pp. 12104-12113, 2022. /tt.ttttp://arxiy http://rrg/bbb///eeee 04560v2
Zhang, B. and Sennrich, R. Root mean square layer normalization. Advances Neural Information Processing Systems, 32, 2019.
15

---

A Experimental Details
A.1 Overview
Architecture and training parameters. We implement the most commonly used decoder-only architecture with SwiGLU activations (Shazeer 2020), RoPE embeddings (Su et al. 2024), RM- SNorm (Zhang Sennrich 2019) and alternating attention and MLP blocks. Unless otherwise noted, we follow standard practices in LLM training and use the A optimizer with beta parameters B2) (0.9, 0.95), decoupled weight decay of 0.1 (Kingma & Ba 2014 Loshchilov Hutter 2017) and gradient clipping with 1.0. For warmup steps, we use a short warmup of 300 steps for majority of runs and 1000 3000 for longer runs (above 100k total steps). The cosine schedule decays the learning rate to 10% of the maximum learning rate. For most of our experiments, we use batch size of 200, i..., roughly O.IM tokens for sequence length of 512. The vocabulary is based on the GPT-2 tokenizer (Radford et al. 2019) and contains 50304 tokens.
Dataset and evaluation. Our main body focuses on results using SlimPajama (Soboleva et al. 2023 a cleaned and deduplicated corpus that includes webcrawl, code, papers and other sources, which is commonly used for pretraining LLMss We use subset of the full corpus that comprises roughly 6B tokens and randomly sample validation set of roughly 3M tokens. During training, we evaluate the models with fixed set of 32 batches of sequence length 512 (the same context length as training) to establish validation loss curves. At the end of training, we compute the full validation set perplexity. We perform further experiments that verify our main findings with OpenWebTexts WebText2 (Gao al. 2020
in Appx. B.4
Implementation and infrastructure. Our code is based on an extension of NanoGPT and uses Py PyTorch (Paszke et al. 2017) as well as FlashAttention (Dao et al. 2022). We incorporate bfloat 16 for memory and throughput, trained with mixed precision float32 parameters and 6 activations (Micikevicius et al. 2017). All experiments (aside from IB and 8B, see see A.2) were performed using cluster A100 GPUs (both 40GB/80GB RAM) with 2 data-paralle (i.e. 2 GPUs per run). Some selected runs used a single node of 8 H100s. We estimate that the full cost of all experiments for this project (including prototyping) to amount to roughly 2500-3000 GPU hours.
Model configurations. We provide an overview of the model sizes and configurations in Table [ and the parameters for training and length in Table All models for the scaling law experiments are trained with 300 warmup steps and a sequence length of 512.
Model Size
d_model
n layers
ffw fizssiee
kv_size
n heads
33M
384
1024
64
53M
512
1536
64
60M
512
10
1536
64
93M
640
12
1792
64
10
124M
768
12
2048
64
12
M
768
16
2048
64
12
210M
768
24
2048
64
12
360M
1024
24
2816
64
16
Table Model configurations for scaling law experiments. We provide an overview of the model sizes and hyperparameters for the different models in the scaling experiments.
A.2 Setup for the 1B and 8B Runs
Setup. In order to efficiently scale to larger models, we use the nanotron library that enables all forms of training parallelism while being performant and flexible. The configuration and hyperpa- rameters of the two models are described in Table 3 Similarly to our other experiments, we use the most common setup of the AdamW optimizer with (31, B2) = (0.9, 0.95). weight decay of 0.1,
http:://gttabb.cccaaaapptttaaarp
16

---

and gradient clipping with 1.0. For the IB model, we use a custom BPE tokenizer (that is very similar GPT-2) because the FineWeb dataset (Penedo et 2024) was already pre-tokenized on our cluster. Each run for the IB model was performed on 4xH1(0ss GPUs. For the 8B model, we reuse the original Llama3 tokenizer (Dubey al. (024) for the similar reason that Web-Edu was pre-tokenized our cluster; here, we use 12 nodes, each composed of 4xGH200 GPUs. The model is split within one node with tensor parallel (TP) set to
Evaluation. For evaluation, we use the lighteval library (Fourrier et al. with nanotron checkpoints. We run the most common LLM benchmarks inc uding
2023) that works directly
MMLU (Hendrycks)
et al. 2021), ARC (Clark et al., 2018), OpenBookQA (Mihaylov et al.
2018). PIQA Bisk et al.
2020
HellaSwag (Zellers et al. 2019), CommonSense@u (Talmor al., 2019)) SIQA (Sap et al.
2019) Winogrande (Sakaguchi et 202T), truncating large benchmarks to 1000 samples so that we
could efficiently evaluate over course of training. This is custom setup that is equal to that used for Web ablations; the setup to reproduce the eval is described here We report the accuracy
normalized by sequence length (acc_norm).
Model
LR (Cos/Const)
BS
Steps
Tokens
Token/Params Ratio
33M
(2e-3, le-3)
0.1M
[3k, 7k, 10k]
[0.3B, 0.7B, 1.0B]
[9.2, 21.4, 30.6]
53M
(2e-3, le-3)
0.IM
[3.7K. 7.5K, 11.2K]
[0.4B, 0.8B, 1.2B]
[7.2, 14.5, 21.7]
(2e-3, le-3)
0.IM
[7.5K, 12.5K, 17.5K]
[0.8B, 1.3B, 1.8]]
[12.8, 21.4, 30.0]
93M
(2e-3, le-3)
0.IM
[10K, 17.5K, 25K]
[1.0B. 1.8B, 2.6B]
[11.0, 19.2, 27.5]
124M
(le-3,5e-4)
0.IM
[15K, 25K, 35K]
[1.5B, 2.6B,
[12.4, 20.7,
151M
(le-3, 5e-4)
0.1M
[25K, 37.5K, 50K]
[2.6B. S.IB]
[16.9, 33.7]
210M
(Ie-3, 5e-4)
0.IM
[37.5K, 50K, 62.5K]
[3.8B. 5.1B. 6.4B]
[18.4, 24.6, 30.7]
360M
(Ie-3,
0.2M
[25K, 37.5K. S0K]
[5.1B, 7.7B, 10.2B]
[14.2, 21,3,
Table Training parameters for scaling experiments. We describe the learning rates, training lengths and ratios for the different models in the scaling experiments.
Parameter
1B
8B
d model
1792
4096
layers
24
32
ffw size
4864
14336
ky size
128
128
n heads
14
32
n kv heads
14
8
vocab_sia _ize
49152
128256
seq
2048
4096
warmup_sted
2000
1000
batch_ste size
(1.8M, 2M)
0.6M
Total Steps
(55k, 220k)
20k
Peak LR
8c-4
3e-4
Table 3: Model configurations for larger runs.
The batch size and learning rate for the IB model tokens were estimated using DeepSeek scaling laws for 100B tokens. The two values for BS and the total steps of the 1B model distinguish the runs (100B,460B). For the 8B model. the architecture is identical to Llama3 and the batch size was set according to the available GPU limits on our cluster at the time of running experiments.
https: http:///hh..honaacoccocooooooooa
eva 
httpg://H...ocoeee/gee/eeeeeeeee py#L12
17

---

A.3 FLOPs computation
We do not rely on the heuristic 6=ND for computing the FLOPs of Transformer model, but use more thorough computation that involves the embeddings, attention and MLP operations directly. For reproducibility and other researchers to use, we provide our Python code in Figure [4
def embedding (seq_len, pocab_size, d_model): return * seq_len vocab_size d_model
def Ittentionsseq_1ee d_model, key_size, num_heads):
projections logits seq_len seq_len (key_size num_heads
seq_len d_model (key_size num_ aum_heads)
softmax = 3 num_heads seq_len seq_len deftmax.ueer_redducee num_heads)
2 len seq_len (key_size
seq_len (key_size num num_heads) d_model return projections logits softmax final_layer
final_layer
10
12 def dense (seq_len, d_model, ffu_size, swiglu-False).
13
if sviglu:
14
return 2 seq_len (3 d_model ffw_size)
else:
seq_len (2 d_model ffw_size)
17
18 def final_logits (seq_len, d_model, vocab_size):
19
return * seq_len d_model vocab_size
def flops(
22
n_layers,
23
seq_len,
vocab_size,
d_model,
26
key_size,
27
num_heads,
28
ffv_size,
swiglu-True,
30
31
flops_single
32
embedding vocab_size, d_model)
33
+ n.layers
35
attention (seq d_model, sizesize, num_heads) dense (seg d_model, ffu_size, svigln-swiglu)
36
37
+ final fonal_logits (seq_len, d_model, vocab_size)
40
assume backward pass twice the FLOPs of the forward pass return flops_single
41
Figure 14: FLOPs computation. Instead of the common approximation of 6=ND, use more detailed calculations for the FLOPs estimation based on the Transformer model configuration. We provide the Python code above.
18

---

Train (210M)
Validation (210M)
Cos 22k
Cooldown to 222
Cos 22k
Cooldown to 22L
Coollown
Coolhonn
30
Clo dik
Cooldown to 4%
30
Coe, 4ik
Cooldown to 46k
Perplexits
Perplexity
25-
-2
15
10000
30000
40000
10000
30000
40000
Steps
Steps
Figure 15: The difference in loss curves of cosine constant learning rate with cooldown. The cooldown phase initiates a sharp decrease in loss for both training (left) and validation (right). The training perplexity is averaged out over window to achieve smoother curve. For validation, use a fixed set of 32 validation batches to report the loss each step, which creates a smooth curve by design.
B Additional Results
B.1 More Results on Cooldown
Ablation of functional form of cooldown. We provide additional experiments to show that results are consistent and robust to different changes. For simplicity, we only cooldown to 50K steps here, but we found these results hold regardless of the step which decay. In Fig. anl Fig. 17 we test different functions for the cooldown phase:
Linear: Classical linear decay.
Sqrt: Defined in Eq. (T-sqrt).
Cosine: Similar to the classical cosine decay, but applied only during the decay phase.
Mirror Cosine: The symmetric counterpart of the cosine function with respect to the linear decay.
1 Square:
The square root function in Eq. [-sqrt) is replaced by a square function.
Note that the order might change for substantially different cooldown lengths; we focus on 10% and 20% because they practically relevant.
In addition, we perform comparison of different exponents for the square-root function: since the (I-Sqrt) can be expressed as (1 x²) where 0.5, we sweep the parameter a < 0.5. The results are shown in Figure [8] Apart from 0.1 and 0.2, which perform noticeably worse because the learning rate is too low for many steps, the other exponents only show a marginal difference to = 0.5, which still comes out on top.
We aim to further investigate the impact of the functional form of the cooldown in future work.
Length of cooldown. We repeat our ablation from Sect. 3 with a 210M model in Figure 19 (fractional x-axis) and with the absolute number of decay steps (not fractional) in Figure 20
LR sensitivity. The optimal learning rate also transfers to different cooldown lengths in our experi- ments, see the results in Figure 21
Annealing cosine to less than 10%. We ablate the choice of the final learning rate for the cosine an- nealing in Fig. 22 where find that the final value should be set lower than simply 10% of the maxi- mum. However, when evaluating on downstream benchmarks (see Section(B.5), we see that annealing zero can hurt metrics; we posit this comes from early saturation. Combining these two argu- ments implies that the maximum and final LR should be set (and ideally swept over) independently.
B.2 Additional Results and Compute Savings
We give the savings for all models used in the scaling experiments in terms of FLOPSs in Figure 24 and GPU hours in Figure 25
19

---

LR with different cooldnwn finettion
Vilidisiãon with differett cooldown function
22
1.00
Mire:
Tearning rate
Perplexity
20-
Sggtt
0.75
0.50
18
0.25
16
0.00
40000
42500
45000
47500
50000
40000
42500
45000
47500
50000
Steps
Steps
Figure 16: Different cooldown functions. We test various cooldown schedule functions and consistently observe the same order of performance (left), with (I-sqrt) being the most effecttiee Remarkably, the drop in loss closely follows the learning rate (right) even for different functions.
Cooldown functions for 1) 5e
Cooldown functions 10% decay steps
22
22
1- Squar
Mese
Perplexity
20
Perplexity
20
18-
16
40000
42500
45000
47500
50000
46000
48000
50000
Steps
Steps
Figure 17: The order and results of different cooldowns hold across settings. Left: We change the learning rate to 5e - instead of 1e- as previous experiments. Right: The performance for 10% cooldown steps instead of 20%. Note that the order might change for substantially different cooldown lengths; we focus on 10% and 20% they are practically relevant.
Cooldown Form (1 (1=x") Perplexity
Cooldown Form (1 - 2²) LR
20
1.00
a-0.5 (vest)
L-0.4
Perplexity
a-1.1
am00.
Learning
0.50
a-111
18
0.25
10000
42000
441000
16000
18000
50000
40000
42000
44000
48000
48000
50000
Steps
Steps
Figure 18: Different exponents for functions. Since the negative square root performs remarkably well, experiment with varying the exponent of the functional form (1 x²) for 0.5. Besides 0.1 and 0.2 which perform noticeably worse, the other exponents only show marginal difference in this experiment, with 0.5 (the square root) still coming out on top.
20

---

Fraction Cooldown Length (210M)
Fraction Cooldown Length (210M)
22
Stage
Stage
Stops
5Sc Stupe
6%% Stupe
Ill Stegs
671: Stops
4lk Stoge
Oodine
Custne
Perplexity
Perplexiit
20
20
18
18
0.05
0.10
0.15
0.20
0.0
0.4
Fraction Cooldown Steps
Fraction of Cooldown Steps
Figure 19: Parabola shape of the relationship between cooldown length and final perplexity. We repeat the experiment from Fig. 5 with a 210M parameter model and the zoomed-in in view on the left.
Cooldown (210M, Absolute Steps)
Cooldown (210M, Absolute Steps)
22
Steps
Stops
20%
56k
SSk Stepe
67k Strpp
Stepe
Stope
Perplexity
Perplexity
442 Stape
Cosine
20
18
5000
10000
15000
10000
20000
30000
40000
Number of Cooldown Steps
Number of Cooldown Steps
Figure 20: The effect of the cooldown length in terms of absolute steps. We repeat the plots from Fig. 19 with the absolute number of steps on the x-axis.
LR. Sensitivity (Cosine Cooldown)
LR. Sensitivity (Cooldown Length)
24
Onlar
Coollown (317))
22
Uookddowa (19%)
Choldown (2075)
Cooktonn (aF)) 33k
Dookiooo (106) 33
Cholioon (19%) 33k
Cuenne 4tk
Cooktowa (307))
Usokdovia (109%)
Cooldown (2075)
Perplexity
22
Perplexity
20
20
19-
18
0.0001
0.0002
0.0005
0.001
0.002
0.0001
0.0002
0.0005
0.001
0.002
Learning Rate
Learning Rate
Figure 21: Learning rate sensitivity for cosine and cooldown with different lengths. The optimal learning rate also transfers to different cooldown lengths.
21

---

Cosine Decay to 1e-7 (124M)
LR Sensitivity (124M, Cosine 1e-7, 25k)
30.0
27.5
Cosise 1e-7
25k
26
Cosine To-7
Perplexity
(20%
Perplexity
I Sept Deray (20% steps)
21
22.5
20.0
10000
20000
30000
10000
50000
0.0001
0.0005
0.001
0.002
Steps
Learning Rate
Figure 22:
The cosine schedule should be decayed to lower values. While setting the minimum to 10% of the maximum LR is the most common heuristic, the final LR should be set lower and inde- pendently of the maximum. our experiments, cosine then matches the cooldown with (I-Sqrt) (left; optimal LR shown), albeit the LR sensitivity remains (right; larger maximum LR led to divergence).
Equality Across Model Sizes Training Scales
Cosine LR
Cooldown
SWA Cosine
SWA Constant
Val Perplexity
30
33M
63M
60M
93M
124M
151M
210M
360M
Model sizes
Figure 23: Final validation perplexity of all models in scaling experiments across different runs.
FLOP Savings for Model Sweeps
Coeine LR
Cooldown
SWA Cosine
SWA Constant
FLOPS
1029
10a8_
33M
53M
60M
93M
121M
151M
210M
360M
Model sizes
Figure 24: FLOPs savings for all models in scaling experiments across different runs (log scale). The savings amount to a factor of across all models. more runs across different lengths were to be performed (compared to 3 per model in our experiments), the difference would be even more significant.
22

---

GPU Hour Savings Model Sweeps
100
Cosine LR
Cooldown
SWA Cosine
SWA Constant
GPU Hours
75
50
25
33M
53M
60M
93M
124M
151M
210M
360M
Model sizes
Figure 25: GPU hours savings for all models in scaling experiments across different runs, We report the actual runtimes summed up over the sweeps for all models. The savings terms runtime become especially more prominent for bigger models and longer training runs. Please note that the hours for SWA of the 360M model are slightly off because of congestion in our cluster during the runs.
50
40
Perplerity
50
a0V SVIA Crecraa
38z SMA Cuaza I3
Porplestty
40-
STEA
NW. Cauakenk
Perplesite
1494 DWA Crnaken
<d5 12:
60-
135
30-
30-
25/0)
5000)
7500
10000
10000 Steps
Steps
Steps
93M
124M
151M
KM Caier ER
35-
[S Cribe 1R
Porplccty
055
J0M Courazk L2
Cooao 12t
Perplesitt
90
12:l SWA Cular
Perplecity
13150 Constat 1I1
SIFA Cheisr
13-M SMA Conkeen
30
20
20-
10000 Sseps
20000
10000
20003
30000
20000 Steps
40000
Stegs
Perplocity
25
23004 SWAA Costae
20103 SWA Consttnt 141
Parplocity
20-
39:0:
39:32
Chmmita
Chicae
20000
Steps
Steps
Figure 26: Validation Loss Curves (Perplexity) of all Models in Scaling Experiments. We visualize all training runs for the models used in the scaling experiments Section 5
B.3 Learning Curves for All Models of Scaling Experiments
We provide the learning curves for all models the scaling experiments in Figure 26 validation perplexity for all models and methods is given in Figure
The final
B.4 Additional Experiments on OpenWebText2
In addition all results on SlimPajama, we perform experiments on the commonly used benchmark of Oeen2e (Gao et al. 2020) with models of sizes 60M, 93M and 166M. As shown in
23

---

45
60M OWT2
00
93M OWT2
168M OWT2
E00 Codee
KIM Coline
Porplesiyy
E000 SWA Contas 2.1
100C Conalani
Purplcity
35
103 Conoana
2354 SIAAA Costee 1at
Perpleciss ledittyy
30
55
6000 SWA Conat on 13
50
2034 SIAAA Conatast 122
25
30-
5000
10000
15000
10000
20000
10000
20000
30000
Stops
Steps
Steps
Figure 27: Results Transfer to Open Text2. We see the same behavior for cooldown schedules and SWA, verifying the reliability of our findings.
A ggr. Score (1B model, 100B tokens)
Aggr. Score (1B model, 460B tokens)
0.45
0.45
Score
Score
0.40
0.40
Contte Sn
1:get Cscldown (2275)
Cheise Cuildumn
0.35
Cuuldran (2FF)
0.35
0
25
50
75
100
0
200
400
Tokens (B)
Tokens (B)
Figure 28: Aggregate metrics throughout training of the 1B model on 100B and 460B tokens. We train IB model on 100B (left) and 460B tokens (right) of FineWeb (Penedo et al. 2024), and find that the performance cosine and cooldown matches. Though cosine to zero improves the loss (Figure 22 it leads to saturation before the end of training, hurting overall performance.
Figure 27 our findings from previous experiments succesfully transfer, where the cooldown recipe matches the performance of cosine and SWA boosts performance during training.
B.5 Full Results of Large Model Runs
In this section, we report the detailed results the 1B runs with the downstream benchmarks.
In Figure 28 we compare the aggregate metrics throughout training for both the 100B (left) and the 460B token run (right). We equally plot individual benehmark curves in Figure 29 for the 100B run with a zoomed view after 80B tokens in Figure 30, interestingly. we observe similar uptick in performance for some metrics (e.g., MMLUU HellaSwag) with the cooldown, while others do not benefit as clearly (e.g.. Open This observation is an interesting direction for further research.
The final numbers are given Table (100B) and Table (460B). Notably, longer cooldowns do necessarily improve the metrics (e.g. going from 5% to 20%).
24

---

XIMLU
ARC
OpenBo
PIQA
0.70
Scate
0.28
Score
10.25
0.123
Scme
0.65
0,300
0.50
0.270
Conter 1o
Ulser 201t
0.55
Challonn (1900
(33)
Tokema (B3)
HellaSSoug
QA
Winogrande
0.42
0.5
0.41
0.52
0.60
0.3
0.25
Lian Cookliin 2116
20
60
100
20
60
50
100
20
-0
60
100
0
100
60
80
100
Tolenns (1)
Tolaras (B)
Tobems (D)
Toboms (D)
Figure 29: Detailed benchmarks throughout training of the 1B model on 100B tokens. For the cooldown (starting 80B tokens), observe similar uptick in performance for some metrics (e.g., MMLU, HellaSwag) akin to the observed drop in loss. Other metrics do not benefit as clearly from the cooldown (e.g., OpenBookQA).
Aggregated Seore
MMLU
ARC
0.42
Score
0.16-
Score
0.306
0.45
Singet
0,300
Hight
2010
90
Tokens (B)
Tokens (B)
Tokens (B)
OpenBookQA
PIQA
Hella ad!
0.36
0.76
0.56
Score
Score
0.34
0.74
0.54
Liaes Choklamn 1299)
Cooeld parto
Lauane Coottons gelto
80
90
100
30
90
100
80
100
Tokens (13)
(11)
Tokens (13)
CommonSenseeAA
SIQA
Winogrande
0.38
Score
Litee Cuitoons deics
Score
0.41
0.54
Chedive carl
LES
0.36
0.40
0.52
Tokens (TB)
Toloess
(T)
Tokems (B3
Figure 30: Zoomed-in view of the 1B model benchmarks after 80B tokens. We repeat Figure
29 with focus on the cooldown phase, which starts 80B tokens.
25

---

Metric
Cosine 10%
Cosine 0
1-Sqrt 20%
Linear 20%
Aggregated Score
46.26
45.88
46.23
46.20
MMLU
30.21
30.70
30.84
30.71
ARC
41.25
42.35
41.80
41.25
OpenBookQA
35.40
33.80
36.00
35.80
PIQA
76.10
75.50
76.50
77.00
HellaSwag
56.30
55.60
55.80
55.90
CommonSenseQA
36.10
36.20
37.00
36.70
SIQA
40.70
40.20
39.60
39.50
Winogrande
54.00
52.70
52.30
52.70
Table 4: Final evaluation results after 100B tokens. Both cosine and the cooldown schedules have comparable final numbers, with only slight differences for certain benchmarks.
Metric
Cosine to
1-Sqrt 5%
Linear 5%
Linear 10%
Linear 20 %
Aggregated Score
48.03
47.91
47.84
47.98
47.92
MMLU
31.25
31.71
31.65
31.78
31.84
ARC
45.20
44.10
44.05
44.15
45.05
OpenBookQA
37.60
37.80
38.00
38.20
37.20
PIQA
78.10
77.40
77.40
77.30
77.80
HellaSwag
59.90
59.60
59.50
60.00
59.20
CommonSense(A
37.70
37.30
36.70
36.90
36.90
SIQA
39.90
39.50
39.40
39.50
39.70
Winogrande
54.60
55.90
56.00
56.00
55.70
Table 5: Final evaluation results after 460B tokens. The findings of Table transfer to much
longer training runs with 460B tokens, where the performances of cosine and cooldowns match well. Notably, longer cooldowns do not necessarily improve the metrics (e.g. going from 5% to 20%).
26

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We provide an in-depth analysis of the proposed methods across different settings and the empirical results match the claims made.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We explicitly discuss limitations in Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
27

---

Justification:
We do not provide theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide experimental details to the best of our efforts both in the main paper well as the Appendix and have released our source code.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
28

---

Answer: [Yes]
Justification: The main dataset used in our experiments is publicly available on HuggingFace with the specified link. We also released code repository htttp:///iithu.... com epfml/ scheduless andedulleeadd--eeee aca-scaling
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ ptapps..c/publlc/gguueeeee////// /etec////uiieee//eeuuuubi Cate///oeeeSeeaaaaaaeeee ssi onPolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide all experimental details, including hyperparameters, optimizers and data splits, to the best of our efforts both in the main paper as well as Appendix A.1
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Training LLMs requires extensive time and compute, thereby making it generally impractical do runs across different seeds. Nonetheless, our results are consistent across runs and scales and are highly transferrable as shown by neural scaling laws.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
29

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma enror bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide the description of compute resources Appendix A.I
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http://aaarrrp.cccppbb.c.cccccla
Answer: [Yes]
Justification:
We have read and conform to the code of ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work on fundamental machine learning research. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
30

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We do not release pretrained models.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: To the best of our knowledge, adequately cite and mention all used assets in this paper.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
31

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssit
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not provide new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
We do not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper doess not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification:
We do involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
32

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
33


Reviews:
[{"review_id": "NdPoDUBjeL", "reviewer": "Reviewer_J233", "summary": "The paper focuses on learning rate schedules in scaling large language models (LLMs). Traditionally, LLMs are trained with a cosine learning rate schedule, which requires a pre-specified training duration for learning rate decay. This makes it difficult to dynamically adjust the training duration, as early stopping and continued training after the decay result in suboptimal model performance. To address this, the authors study an alternative schedule of constant learning rate with cooldown, which enables one to get the optimal model at any step during the training process. The authors systematically compare different choices of cool down duration and cool down schedule, and demonstrate that constant learning rate with 1 - sqrt cooldown schedule matches the performance of cosine schedule tuned for that duration.\n\nBeyond cooldown, the authors compare two different approaches of removing the need of learning rate schedule: stochastic weight averaging and schedule-free optimizer. The authors show that while they exhibit promising performance, the performance of models obtained in the middle of the training still cannot match the performance of the proposed constant learning rate with cooldown.\n\nFinally, the authors validate the proposed learning rate schedule in the scaling law study of compute optimal models, and demonstrate that the constant learning rate with cooldown schedule leads to the same results as cosine learning rate schedule, while saving half of the compute flops.", "questions": "Another highly relevant use of training LLMs without fixed duration is continued pre-training, which is quite common in practice as people often continue the pre-training on specialized data for different use cases (such as turning general models into code specialized models). It would be great if the authors could include some discussions and experiments on how the proposed constant learning rate with cooldown can be applied in continued pre-training. For example, one question to ask would be would it be better to continue from the checkpoint before or after cooldown.", "limitations": "The authors have sufficiently addressed the limitations and societal impacts.", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper presents a simple approach to address a major limitation of cosine learning rate schedule, and demonstrates its effectiveness with solid experimentations. Overall I think the paper is of very high quality.\n\n\nThe proposed method is simple and highly practical, which makes it relevant for a broad range of today’s scaling studies. It is widely known that scaling is the key to the success of LLMs, and it is important to establish compute optimal scaling laws as large scaling training is really expensive. The proposed learning rate schedule is able to halve the flops of the chinchilla compute optimal horizon study without adding extra implementation or tuning complexity, which I believe is a major contribution.\n\nThe proposed method is supported with solid experimentations. The authors include comprehensive comparisons of different schedules of cooldown and hyperparameter choices, which convincingly establish the advantage of the proposed schedule and can also serve as tuning guide when adopting it. The authors also compare the proposed method to popular alternative approaches.\n\nThe paper is well written. The main idea and the experiments are laid out in a logical way that makes it easy to follow. The experiment results are also presented in an intuitive way, making it easy to compare methods.", "weaknesses": "While the experiment results are very solid, I do hope that the authors could expand the scale of experiments to larger models. Larger models are usually more sensitive to suboptimal hyperparameters, and therefore it would be great if the proposed method can also be verified in larger models, such as a 3B or 7B model which are common sizes for practical LLMs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you very much for your valuable comments and kind words!\n\n> **Expansion to larger models**\n\nPlease see the general response above for the results on both a 8B and a 1B model, which we were able to train without issues out of the box. This makes us confident that the match of performance of the schedules also holds for larger scales, and we aim to investigate even larger models in the future. We also look forward to more adoption and ablations in different settings of the community.\n\n> **Discussion on continued pretraining**\n\nWe fully agree on the relevance to continued pretraining and have expanded our discussion on these points in the updated version of the paper.\nThe natural approach is to use checkpoints before the cooldown to continue training with a high LR; this avoids loss spikes and brittle training when rewarming the learning rate (cf. Figure 1 in our paper with cosine). Moreover, rewarming has been reported to hurt performance and introduce forgetting compared to single annealing training [1], though careful strategies can alleviate such issues. We think it remains an interesting question if a single cooldown schedule is absolutely optimal given a total compute budget for LLM training.\n\n---\n\nPlease let us know if you would like us to elaborate any additional points, we’d be happy to discuss.\n\n[1] Ibrahim, A., Thérien, B., Gupta, K., Richter, M. L., Anthony, Q., Lesort, T., Belilovsky, E., and Rish, I. Simple and scalable strategies to continually pre-train large language models. Mar 2024. URL https://arxiv.org/abs/2403.08763v3", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "saWhBny20y", "reviewer": "Reviewer_5fmf", "summary": "Scaling Laws and Compute-Optimal Training without Fixed Training Duration\n\nSummary\nA major weakness of the cosine annealing learning rate schedule, one of the most prevalent learning rate schedules in LLM training, is that for optimal performance the cycle length must be adjusted based on the training duration. Intermediate checkpoints generally do not achieve optimal loss. This is inconvenient since achieving the optimal loss for various training durations requires individual runs for each duration. The authors investigative alternative scheduling methods for alleviating this difficulty. Namely, the authors seek solutions which allow provide near optimal intermediate checkpoints along one long training run, potentially with a relatively small overhead from applying cooldown or averaging. Such solutions are especially useful for computing scaling laws and hyperparameter tuning by reducing the computation cost significantly. Through careful experimentation the authors provide several insights about the following solutions (which have been proposed previously in the literature).\n\n1.\tTrapezoidal LR schedule with warmup + constant lr + cooldown\n2.\tStochastic Weight Averaging\n3.\tSchedule Free Optimizer from Defazio et al. [1]\n\nIn particular it is demonstrated that the trapezoidal schedule is a simple and effective solution for provide near optimal checkpoints along a potentially infinitely long trajectory. \n\n[1] Defazio, A., Mehta, H., Mishchenko, K., Khaled, A., Cutkosky, A., & others. (2024). The road less scheduled. arXiv preprint arXiv:2405.15682.", "questions": "1. In figure 1 why are there no spikes? For example the n=8k curve is rewarmed every 8k iterations. Are the curves smoothed?\n2. Can the SWA be clarified a bit? Is every evaluation done on the iterate average over all checkpoints in the last 500 steps? How many checkpoints does that include?\n3. What about a cooldown function like (1-x^a) where a < 0.5? All the other cooldown functions tried seem to lie \"above\" 1 - sqrt(x).", "limitations": "Yes.", "rating": 9, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper discusses an important and practical topic of LLM training. Hyperparameter tuning has been established as an important aspect of LLM training but result in prohibitive costs especially when scaling. The paper discusses solutions for trying to \"eliminate\" the hyperparameter of training duration which is a worthy goal with a large benefit. The paper is also written and presented clearly with many experiments and interesting results.", "weaknesses": "The technical novelty is a bit limited because all the approaches evaluated have been proposed before in the literature (modulo some small tweaks). However, I do not view this as a major weakness since a thorough, well-presented review of these approaches is still a strong contribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for the positive and detailed feedback to our work!\n\nWe aim to clarify your individual questions below:\n\n> **Spikes in Figure 1**\n\nFor all the figures in our paper, we plot the validation perplexity of a fixed set of sequences in the dataset; this therefore entails very smooth curves as there is no batch sampling. You are right that the curve  for n=8k (dashed black) is rewarmed cyclically, which results in a strong jump in loss in the left part of the Figure. Does this answer your question?\n\n> **Clarification of SWA**\n\nApologies that the description of SWA was not clear enough -- we have updated the paper to reflect this feedback. We perform SWA within a window of 500 steps. For example, at step 1125, we uniformly average the 125 previous steps down to 1000; at step 1500, the average contains all steps down to 1000. This is done by keeping a single running average.  We then store these averages for each window as checkpoints (i.e., at steps 1000, 1500, 2000, …). \nStoring these checkpoints allows investigating the optimal horizon by retrospectively averaging the averages and finding the minimal validation loss. This is what we use to plot the evaluation curves in Figure 8. We found that starting in the middle stages of training, the optimal horizon stabilizes around 2500 steps (i.e. 5 checkpoints). \n\n> **A cooldown of (1-x^a)**\n\nThis is a very good point and an interesting comparison! We will include such an ablation in the final version of our paper. With that, we might be able to better answer the question of what the ‘optimal’ cooldown looks like.\n\n---\n\nPlease let us know if you have additional questions, we are happy to discuss more.", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ps59Pnt1Rj", "reviewer": "Reviewer_yzXN", "summary": "The authors present an empirical study on using constant learning rates (plus a short cooldown) instead of cosine schedules for training LLMs. The authors show that:\n\n- Constant LR + cooldown roughly matches cosine schedules. (Fig 3,4)\n- SWA of a long schedule almost matches the performance of shorter schedule given the same number of steps. (Fig 8).\n- Chinchilla-type scaling laws can be derived with constant LR schedules + cooldowns which will save compute. (Fig 10)", "questions": "1. Why plot perplexity instead of val loss?\n2. Could you have a subsection called experimental setup? On line 104 you mention “. We follow our experimental setup and train a 210M model”, but I don’t think the details of this model is given earlier.\n3. Based upon Figure 2, it seems like your schedule goes to 0 whereas cosine typically does not go to 0. Is this a confounder in your experiments?\n4. Fig 4 shows constant LR beating cosine while Fig 3 shows them matching. Why this discrepancy?\n5. Fig 6 (right) seems to show that there is still a small gap between constant (1-sqrt) and cosine – is this correct?", "limitations": "na", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The potential impact of constant LR instead of cosine is high. \n- The paper is rather well written.", "weaknesses": "- The experimental setup is unclear. But this should be very easy to fix in an updated version.\n- It is still unclear to me exactly how well the match between constant LR and cosine is, especially for longer training. \n- The results on SWA and constant LR seems almost orthogonal. \n- Higher performance is not reached, so the only upside is in doing scaling studies. This might only be relevant for a few well-funded labs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for taking the time to read our work and the detailed questions. You can find the replies below:\n\n> **On the experimental setup and perplexity (weaknesses and Questions 1-2):**\n\nThank you for the feedback -- we have made the setup and pointers to the exact details more clear in the updated version of the paper. In essence, we follow the most common setup of LLM pretraining with a decoder-only Llama architecture and a large web scale dataset. In the submitted version, we describe this setup in the background Section 2 and follow it for the rest of the paper, with exact details given in Appendix A.1. \nRegarding the metrics, we opt for the perplexity instead of the loss simply because it makes interpreting results easier, as the values are spread further apart.\n\n> **The match between cooldown and cosine (from weaknesses):**\n\nAll our experimental results point towards a strong match between cosine and the constant LR when the cooldown is performed between 10-20% of training duration. In particular, this is the goal of Section 5 and Figure 10, where each schedule leads to the same result for the same model and training length.\nNotably, we have been able to verify this alignment with our larger scale runs of a 1B model using downstream evaluations; see our general response above.\n\n\n> **Results on SWA and constant LR (from weaknesses):**\n\nThere is a strong theoretical and empirical connection between weight averaging and learning rate decay (see [1]). We were particularly motivated to potentially replace the cooldown with SWA entirely; while it does not fully match the performance, we found a strong improvement that has practical relevance and merits a mention in the paper. Moreover, we believe these insights might inspire future work in this direction, e.g., cleverly combining SWA and cooldowns.\n\n> **Relevance only for large labs (from weaknesses):**\n\nWe agree that scaling studies are compute intensive and therefore most accessible to well-funded labs. However, the aim of our work is precisely to reduce the GPU hours required to train models or run scaling studies. This is possible by making it easier to 1) continue training a model for more tokens without restart, and 2) see a model at different scales of compute via checkpoints and ad-hoc cooldowns. We therefore hope that our findings enable (and encourage) smaller labs to perform scaling studies.\n\n\n> **The scale of the cooldown (Question 3):**\n\nYou are right that cosine does not typically go to zero, whereas the LR cooldown does. We have therefore added new ablations to investigate this. \nWhile going to zero helps cosine achieve lower loss values, we find that it overall hurts benchmark metrics. We posit that annealing to zero with cosine leads to a saturation before the end of training, as large portions will be spent with very low LR values.\nPlease find the detailed results in our general response to all reviewers above. \n\n> **Difference between Figure 3, 4 and 6 (Questions 4-5):**\n\nWe apologize for the confusion. Figure 3 shows the match between cosine and a linear cooldown for 20%, where we sweep the peak LR and plot the optimal one for both approaches.\nFigure 4, on the other hand, shows the curves with the same LR (not necessarily optimal) that was found in the previous experiment for much longer runs (200k steps instead of 45k), where we find that the cooldown (again for 20% of steps) can outperform cosine.\nIn contrast, Figure 6 shows a cooldown for only 5% of training steps. This was motivated by our ablation on the cooldown length in Figure 5. Importantly, we find a negligible gap between cosine and such a short cooldown (compared to the 20% used before).\nFollowing your feedback, we have adapted the captions to make this distinction clearer.\n\n---\n\nPlease let us know if this response clarifies your questions or if you would like to discuss more points.\n\n\n[1] Sandler, M., Zhmoginov, A., Vladymyrov, M., and Miller, N. Training trajectories, mini-batch losses and the curious role of the learning rate. Jan 2023. URL http://arxiv.org/abs/2301.416 02312v2.", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank everyone for taking the time to read our work and the positive & constructive feedback. We are very encouraged by the positive comments of the reviewers to our findings and the writing.

We also thank all reviewers for the useful comments with detailed questions, to which we reply in individual comments. We look forward to engaging in further discussion, either for questions or potential improvements.

Since the submission date and with the points raised, we have been working diligently on expanding and generalizing our results. In particular, we want to highlight **larger scale runs and investigation of evaluation metrics** to all reviewers. In summary, we have results for 1) a 8B model run, 2) evaluation benchmarks (beyond loss) on a 1B model for 100B and 460B tokens, and 3) ablations of cosine to zero at both small and large scale. 

## 1) Validation on a 8B run
We have been able to validate the LR schedule for a **8B model (Llama 3.1 architecture)**, showcasing the exact same behavior as we have demonstrated on smaller scales. We rely on the same peak LR as Llama3 (0.0003) with a smaller batch size (~0.3M tokens) and train on FineWeb-Edu [3]. In Figure 1 in the attached PDF, we show the loss curves for both cosine and the constant LR (20% cooldown with 1-Sqrt) and find that the **training is stable** and **performance matches**. Albeit it is a short run (6B tokens total), we believe these are promising results on a much larger scale.

## 2) Results for 1B model with benchmarks
Beyond just loss values, we have been aiming to include **downstream evaluation benchmarks** -- which are ultimately the main interest -- and have obtained results for a **1B model on long runs of 100B and 460B tokens**.

To be more specific, we used scaling rules by Bi et al. [1] to determine an estimate of the optimal peak learning rate and batch size for two dataset sizes of 100B and 460B tokens of the FineWeb [2] dataset. Notably, the training was stable for both cosine and the constant LR, thus we were able to run the experiments **without issues or tuning, out of the box**.  With the results in the tables below, we find a **clear agreement between the performance of both cosine and the cooldown schedule**, which validates the method beyond just loss values.

## 3) Annealing cosine to 0
In addition, we have added ablations for annealing the cosine schedule to 0, the same final learning rate as for the cooldown schedule. While it helps achieve lower loss values (Figure 2), we find that overall it hurts benchmark metrics (Figure 3 and Tables below). We posit that when annealing to 0, the shape of the cosine curve results in a large portion of training spent with a LR too close to 0; the final LR should therefore also be tuned, which adds complexity to cosine.

All results, including a detailed description of the experimental setup, the evaluation metrics throughout training, and ablations of cosine to 0, will also be available in the updated version of the paper.

---

## Table 1: Results on 100B tokens
* Batch size: 1.86M tokens
* Total steps: 55’000
* Peak Learning Rate: 0.0008

| Schedule              | Aggregated Score | MMLU  | ARC   | OpenBookQA | PIQA  | HellaSwag | CommonSenseQA | SIQA  | Winogrande |
| --------------------- | ---------------- | ----- | ----- | ---------- | ----- | --------- | ------------- | ----- | ---------- |
| Cosine to 10%         | 46.26            | 30.21 | 41.25 | 35.40      | 76.10 | 56.30     | 36.10         | 40.70 | 54.00      |
| Cosine to 0           | 45.88            | 30.70 | 42.35 | 33.80      | 75.50 | 55.60     | 36.20         | 40.20 | 52.70      |
| 1-Sqrt Cooldown (20%) | 46.23            | 30.84 | 41.80 | 36.00      | 76.50 | 55.80     | 37.00         | 39.60 | 52.30      |
| Linear Cooldown (20%) | 46.20            | 30.71 | 41.25 | 35.80      | 77.00 | 55.90     | 36.70         | 39.50 | 52.70      |

## Table 2: Results on 460B tokens
* Batch size: 2.1M tokens
* Total steps: 220’000
* Peak Learning Rate: 0.0008

| Schedule            | Aggregated Score | MMLU  | ARC   | OpenBookQA | PIQA  | HellaSwag | CommonSenseQA | SIQA  | Winogrande |
| ------------------- | ---------------- | ----- | ----- | ---------- | ----- | --------- | ------------- | ----- | ---------- |
| Cosine to 0         | 48.03            | 31.25 | 45.20 | 37.60      | 78.10 | 59.90     | 37.70         | 39.90 | 54.60      |
| 1-Sqrt Cooldown 5%  | 47.91            | 31.71 | 44.10 | 37.80      | 77.40 | 59.60     | 37.30         | 39.50 | 55.90      |
| Linear Cooldown 5%   | 47.84            | 31.65 | 44.05 | 38.00      | 77.40 | 59.50     | 36.70         | 39.40 | 56.00      |
| Linear Cooldown 10% | 47.98            | 31.78 | 44.15 | 38.20      | 77.30 | 60.00     | 36.90         | 39.50 | 56.00      |
| Linear Cooldown 20% | 47.92            | 31.84 | 45.05 | 37.20      | 77.80 | 59.20     | 36.90         | 39.70 | 55.70      |


---

## References
[1] Bi, X., Chen, D., Chen, G., Chen, S., Dai, D., Deng, C., Ding, H., Dong, K., Du, Q., Fu, Z., et al. Deepseek LLM: Scaling Open-Source Language Models with Longtermism. arXiv preprint arXiv:2401.02954, 2024.

[2] Penedo, G., Kydlíček, H., Allal, L. B., Lozhkov, A., Mitchell, M., Raffel, C., Werra, L. V., and Wolf, T. The fineweb datasets: Decanting the web for the finest text data at scale, 2024. URL https://arxiv.org/abs/2406.17557.

[3]  Lozhkov, A., Allal, L. B., Werra, L. V., and Wolf, T. FineWeb-Edu, 2024. URL https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu.

Meta-review: The paper offers a practical and well-validated approach to optimizing training schedules for LLMs, with significant implications for continual learning and reducing computational costs in scaling law studies. While the technical novelty is modest, the thorough experimental analysis and potential impact make this work a strong contribution to the field. All reviewers are in favor of this paper. This is a clear accept.
Decision: Accept (spotlight)