How to Allocate Your Tokens? Scaling Laws with Training Steps and Batch Size

2026-07-01

AI Summary

This work addresses the critical challenge of optimally allocating computational resources in large-scale model training by balancing the number of training steps against batch size. The authors propose a novel three-factor scaling law that explicitly decomposes the total training data into training steps and batch size, thereby unifying the modeling of how model scale, training steps, and batch size jointly influence performance. This formulation enables robust estimation of scaling relationships using data from suboptimal batch sizes and is validated through extensive large-scale experiments. The proposed scaling law accurately reproduces optimal-batch-size behavior, empirically confirms the existence of a critical batch size, and substantially reduces the number of required training trials, offering principled guidance for efficient resource allocation in practice.

Abstract

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.

Problem

• scaling laws
• batch size
• training steps
• model scaling
• optimal allocation

Innovation

• scaling laws
• batch size
• training steps
• three-term law
• critical batch size