How to Allocate Your Tokens? Scaling Laws with Training Steps and Batch Size
For deep learning practitioners, this provides a more efficient and accurate method to predict optimal training configurations, though it is an incremental extension of existing scaling laws.
The paper proposes a three-term scaling law that separates training data into training steps and batch size, enabling robust fitting with fewer runs and correctly recovering optimal batch size scaling.
We propose a scaling law that takes into account model size and training data while explicitly splitting the latter into training steps and batch size (called three-term law). Fitting the proposed law on a large set of training runs, we find that it correctly recovers the scaling of the optimal batch size. Moreover, because it makes use of training runs with suboptimal batch size, our proposed law can be robustly fit with a significantly smaller amount of training runs. We further show that the three-term law can be used to derive scaling laws for suboptimal batch sizes, and that it matches previous empirical findings related to the critical batch size.