Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations — part 1/3
Cleaned-text code points [0, 15840). Bibliography omitted here; source.md includes it.
# Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations
###### 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/ .
## 1 Introduction
Training large language models is expensive — in time, energy, and compute. Moreover, it requires a complex algorithmic recipe of model architecture and training data to obtain high-quality models. Therefore, the workflow of training large models consists of iterating over small experiments to verify success before extrapolating to larger scales. This is then either done by computing specialized scaling laws ( OpenAI, 2023 ; Bi et al., 2024 ; Hu et al., 2024 ; Team et al., 2023 ) , relying on established laws ( Hoffmann et al., 2022 ; Anil et al., 2023 ) or training past compute-optimality to save cost at inference ( Touvron et al., 2023a ; Touvron et al., 2023b ) .
Despite large advances across data and training recipes, one 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., 2021 ) . Importantly, the Chinchilla project ( Hoffmann et al., 2022 ) 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, from scratch , to have reliable estimates of the quality of training and the scaling behavior. This is much more expensive than training a suite of models just once. Even more, it 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 a multitude of training runs, we demonstrate how a simple alternative of performing a cooldown after a constant learning rate — which was already suggested in the literature ( Zhai et al., 2022 ) and recently used by released models ( Hu et al., 2024 ; Shen et 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 as reliable as cosine, outperforming it for sufficiently long 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 a 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 to the need to retrain models from scratch. We demonstrate this empirically by performing a 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 et al., 2024 ) .
## 2 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 et al., 2017 ; Loshchilov & Hutter, 2016 ) .
In the context of LLMs, the most commonly used cosine strategy presents a particular trade-off by which the LR reaches its maximum early after the warm-up stage and then gradually decreases, typically to 10% of the maximum LR (see Figure 1 , right). Since the seminal works of GPT ( Radford et al., 2018 ; Radford et al., 2019 ; Brown et al., 2020 ) and large models like Gopher ( Rae et al., 2021 ) , PaLM2 ( Anil et al., 2023 ) or LLaMA ( Touvron et al., 2023a ; Touvron et al., 2023b ) , cosine has stayed the de-facto standard schedule.
Experimental setup. Our first goal is 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 to the LLaMa ( Touvron et al., 2023a ; Touvron et al., 2023b ) or Noam architecture ( Ormazabal et al., 2024 ) . Throughout this paper, we use the AdamW optimizer with weight decay ( Kingma & Ba, 2014 ; Loshchilov & Hutter, 2017 ) with common LLM training parameters. We train on a subset of SlimPajama ( Soboleva et al., 2023 ) with 6B tokens 1 1 1 https://huggingface.co/datasets/DKYoon/SlimPajama-6B , a 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 A.1 .
The pitfalls of cosine. In the results of Figure 1 , 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 is 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 ( Ibrahim 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 a good trade-off between a high learning rate and cooling down the model sufficiently, which is expanded proportionally to the training steps.
The alternative — constant + cooldown. The tradeoff of cosine can also be achieved by a different schedule. Here, the learning rate (LR) is kept constant for the majority of the training and only decreases in a short final phase, known as the cooldown (or decay/annealing phase). This schedule has previously been referred to as a trapezoidal ( Zhai et al., 2022 ) , and later as the warmup-stable-decay schedule (WSD) ( Hu et al., 2024 ) . To avoid overloading of terms (e.g., weight decay), we refer to this approach as 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 2 ). Formally, we can define it as η ( n ) = { n N warmup ⋅ η max if n < N warmup η max if N warmup < n ≤ N − N decay f ( n , N , N decay ) ⋅ η max if n > N − N decay \eta(n)=\begin{cases}\frac{n}{N_{\text{warmup}}}\cdot\eta_{\text{max}}&\text{if }n<{N}_{\text{warmup}}\\ \eta_{\text{max}}&\text{if }N_{\text{warmup}}<n\leq N-N_{\text{decay}}\\ f(n,N,N_{\text{decay}})\cdot\eta_{\text{max}}&\text{if }n>N-N_{\text{decay}}\end{cases} (1) with the peak learning rate η max \eta_{\text{max}} , the total steps N N with warmup and cooldown steps N warmup N_{\text{warmup}} and N decay N_{\text{decay}} , and a monotonically decreasing function f ( n , N , N decay ) f(n,N,N_{\text{decay}}) that handles the cooldown.
### 3.1 The Advantages
The main advantage of the constant schedule is 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 a high LR; this avoids loss spikes and brittle training when rewarming the learning rate (cf. Figure 1 with cosine). Moreover, rewarming has been reported to hurt performance and introduce forgetting compared to single annealing training ( Ibrahim et 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.
Data mixtures. During the cooling phase, the data mixture can be changed ( Hu et al., 2024 ) as a form of finetuning; beyond aiming for specific downstream tasks, this allows one to assess the quality of specific mixes, e.g., as recently done for Llama 3 ( Dubey et al., 2024 ) . Although we focus on the same mixture, understanding the curriculum aspect of the separate phases is explored in the concurrent work of Blakeney et al. (2024) and remains open for further research.
Scaling studies. Since a cooldown can also be initiated retrospectively from a 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 5 .
### 3.2 Experimental Comparison
We follow our experimental setup from the previous section (details in A.1 ) and train a 210M parameter model on SlimPajama with constant LR and the cooldown schedule defined in ( 1 ). That is, 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, we 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.1 . We identify a new function ( 1-sqrt ) that outperforms the linear cooldown. This improvement is maintained in a smaller number of decay steps, different learning rates, and various timestamps (see Appendix B.1 ). In Figure 4 , we train a 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 ( 1-sqrt ) cooldown, which we define as: f ( n , N , N decay ) = ( 1 − n − ( N − N decay ) N decay ) f(n,N,N_{\text{decay}})=\textstyle\left(1-\sqrt{\frac{n-(N-N_{\text{decay}})}{N_{\text{decay}}}}\,\right) (1-sqrt)
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, in Fig. 6 we demonstrate that using only 5% decay with the ( 1-sqrt ) cooldown can nearly match the performance of the cosine schedule on a 20B token run (much beyond Chinchilla optimal). This is practically important, as we ideally keep the cooldown as short as possible but get strong performance.
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 15 ). 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 a 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.e., 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 a connected minimum of the loss.
Decaying cosine to less than 10%. Throughout our experiments, we 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.1 . We find that, perhaps unsurprisingly, the influence of the final LR is non-negligible: decaying to lower values improves performance to match the best (1-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.
### 3.3 Scaling Up: 1B and 8B Models
We expand our ablations to realistic scenarios of 1B 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 schedule with the high-quality FineWeb corpus ( Penedo et al., 2024 ) . Following DeepSeek scaling laws ( Bi et al., 2024 ) , we set an approximately optimal batch size and LR of 1.8M and 8 e- 4 8\text{e-}4 for 100B and 460B token runs; evaluation is done throughout training on the most common benchmarks like MMLU ( Hendrycks et al., 2021 ) . The setup is described in detail in Appx. A.2 .