Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations

🌈 Abstract

The article discusses the importance of understanding a model's scaling properties for effectively designing training setups and future architectures. It argues that the commonly used cosine learning rate schedule has been overly complex, and investigates a simpler alternative of constant learning rate with a cooldown phase. The article demonstrates that this alternative matches the performance of cosine, provides additional advantages, and enables significant reductions in compute and GPU hours for scaling law experiments.

🙋 Q&A

[01] Scale and Training Research

1. What are the key issues with the cosine learning rate schedule discussed in the article?

• The cosine schedule achieves optimal loss only when the cycle length matches the training duration, but underestimates the model's performance during training.
• This means that when performing experiments, one must train multiple models for different lengths from scratch to have reliable estimates of the quality of training and scaling behavior, which is much more expensive.
• The cosine schedule also complicates continuation of training, as the improvement in loss precisely happens because of the learning rate decay, and extrapolating the loss curve beyond the end of the cycle is generally not possible.

2. What is the alternative learning rate schedule proposed in the article? The article proposes an alternative of using a constant learning rate for the majority of training, followed by a cooldown phase where the learning rate decreases, typically in a linear or (1-sqrt) form.

3. What are the key advantages of the constant learning rate with cooldown schedule?

• It does not require specifying the number of training steps in advance, which is convenient for large runs.
• It allows for continual learning by default, as training can be resumed from a checkpoint prior to the cooldown phase.
• It can enable changing the data mixture during the cooldown phase as a form of finetuning.

4. How does the performance of the constant learning rate with cooldown schedule compare to the cosine schedule? The article shows that the constant learning rate with cooldown schedule matches the performance of the optimally tuned cosine schedule, and in some cases even outperforms it, especially with the (1-sqrt) cooldown form.

[02] Stochastic Weight Averaging (SWA) and Schedule-Free Optimizer (SFO)

1. How does stochastic weight averaging (SWA) perform compared to the cooldown schedule? SWA provides a significant performance boost over a constant learning rate, but does not fully match the performance of the cooldown schedule. However, SWA can be applied on top of any schedule, including cosine, to improve performance during training.

2. How does the schedule-free optimizer (SFO) perform compared to the cooldown schedule? The article finds that SFO performs well but is still matched or outperformed by the cooldown schedule, particularly when comparing the same configuration of momentum parameters.

[03] Implications for Scaling Law Research

1. How do the alternative learning rate schedules impact the cost of scaling law experiments? The article demonstrates that the constant learning rate with cooldown schedule, as well as SWA, enable scaling law experiments to be performed with significantly reduced compute and GPU hours compared to the traditional approach of retraining models from scratch with the cosine schedule.

2. What are the estimated savings for the Chinchilla model suite if the cooldown schedule had been used? The article estimates that the Chinchilla model suite could have been trained with less than half the compute that was originally used, by utilizing a single run for each model size with a 10% cooldown.

3. Why are the presented results important for the future of scaling law research? The reduced computational cost enabled by the alternative schedules makes scaling law research more accessible and allows for more frequent updates, which is important given recent findings on the data-dependency of scaling laws.

</output_format>