Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations — part 2/3

Cleaned-text code points [15840, 31820). Bibliography omitted here; source.md includes it.

With the aggregated results in Fig. 8 , we find a clear agreement of the downstream performance for both schedules. Interestingly, there is an uptick for some 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 research. The results for the much longer 460B run, where performance also matches, and detailed benchmark numbers are provided in Appx. B.5 .

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 ) .

## 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 at 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 a 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 ) , Sandler 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 a window , allowing us to keep the average as a single additional copy of the model parameters. In our experiments, we set the window to h = 500 h=500 steps and save the window averages as checkpoints every h h steps, which allows ad hoc evaluation of longer windows to see potential improvements akin to latest weight averaging ( Kaddour, 2022 ; Sanyal et al., 2023 , LAWA,) . For all our experiments, we find that windows below or 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.

Experimental results. We evaluate SWA with the same 210M model and show the results in Fig. 10 for both a constant LR (left) and cosine (right). Notably, we find a 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, in line with previous work ( Kaddour, 2022 ; Sanyal et al., 2023 ; Andriushchenko et al., 2023 ) , we see a 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 acceptable and one wishes to avoid cooldowns. This is particularly advantageous as it can arguably be done for free on top of existing well-tuned optimizers.

### 4.2 Schedule-Free Optimizer (SFO)

Very recently, Defazio et al. (2024) introduced a schedule-free optimizer (SFO) which uses an interpolation between standard averaging and Polyak-Ruppert averaging, inspired by Nesterov’s accelerated method ( Nesterov, 1983 ) . 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 a comparison in the same LLM setting.

Results. We compare the results of a long 210M model run with cooldown vs. SFO with AdamW in Figure 11 . Although the optimizer does not require a learning rate schedule, the authors point out that training is more sensitive to the choice of the momentum parameters ( β 1 , β 2 ) (\beta_{1},\beta_{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 ) (0.9,0.95) performs significantly worse and even increases the loss toward the end of training. For ( β 1 = 0.95 , β 2 = 0.99 ) (\beta_{1}=0.95,\beta_{2}=0.99) , SFO performs remarkably well. Nevertheless, both settings are 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 a 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 L , as a function of the parameters N N or training tokens D D ; for LLMs, this is usually expressed as the power law L ⁡ ( N , D ) = A N α + B D β + E , L(N,D)=\frac{A}{N^{\alpha}}+\frac{B}{D^{\beta}}+E\;, where { A , α , B , β , E } \{A,\alpha,B,\beta,E\} are variables to be estimated ( Kaplan et al., 2020 ) . Such scaling laws serve a multitude of critical purposes — from optimally spending a fixed amount of compute (FLOPs) for achieving the lowest loss, to trading off data-sources of different quality ( Bi et al., 2024 ; Goyal et al., 2024 ) or modalities ( Aghajanyan et al., 2023 ) , to comparing the efficiency of different architectures ( Gu & Dao, 2023 ; De et al., 2024 ) .

Crucially, Hoffmann et al. (2022) demonstrated that it is necessary to vary the number of training tokens for a fixed family of models and fully decay the LR ( Kaddour et al., 2023 , see also) . At the time, their results suggested that LLMs were over-sized, leading to a substantial increase in data collection and training for much longer, beyond the Chinchilla optimal point ( N , D ) (N,D) ( Touvron et al., 2023a ; Touvron et al., 2023b ) .

Why do the presented results matter for scaling? Following the results of Chinchilla, scaling laws require a 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 a much cheaper alternative in two phases: first, a model sweep with a 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.1 .

Results. We show the validation loss envelopes as a function of the training FLOPs in Fig. 12 (left) and compare each obtained model (i.e., same parameter & tokens) with cosine and its alternatives (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.1 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. 13(a) , we report FLOPs 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 use token ratios of 10 , 20 10,20 and 30 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 as reported in Table A9 ( Hoffmann et al., 2022 ) . Since the authors do not report exact training configurations, we consider a sequence length of 1024, a batch size of 0.5M tokens and token ratios M = D / N ∈ { 10 , 15 , 20 , 25 } M=D/N\in\{10,15,20,25\} for each model. With this, we arrive at roughly 5.59 × 10 23 5.59\times 10^{23} total FLOPs originally vs. 2.36 × 10 23 2.36\times 10^{23} (Figure 2(a) ). This means that less than half the compute could have been used.

Additional results. We plot the training curves of all models in Appx. B.3 . In addition, we show the same findings with experiments on OpenWebText2 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 et al., 2021 ) . 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 is 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., 2020 ) and PaLM1 ( Chowdhery et al., 2023 ) that used a form of inverse square root. In line with our work, recently released models opt for alternatives such as stepwise schedules ( Bi et al., 2024 ) or the presented constant + cooldown ( Shen et al., 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 context of continual learning ( Ibrahim et al., 2024 ; Gupta et al., 2023 ) . Defazio et 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. Izmailov 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-Brotons et al., 2024 ) . Importantly, Sandler 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 a 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 establish scaling laws for language models. Important to our work, Hoffmann et al. (2022) revise these findings and demonstrate specific methods to establish scaling laws, notably training a family of models for different cosine lengths. The subsequent models like LLama1/2 ( Touvron et al., 2023a ; Touvron et al., 2023b ) further improve performance of smaller models by training beyond the Chinchilla optimal point, motivated 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. However, these works do not consider efficient experiments for scaling laws, which is the focus of our work.

In concurrent work, Porian et al. (2024) and Pearce & Song (2024) 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 token/parameter ratio can be obtained with a constant learning rate without any cooldown; however, the actual performance is then suboptimal, and a 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, we analyze different recipes for the decay form and length. Importantly, we do not claim to have 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 a 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.

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-Chinchilla era, where models are trained much beyond compute-optimal, 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 ; Aghajanyan 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 to researchers and practitioners alike.