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

Technical Summary: Scaling Laws with Training Steps and Batch Size

Problem Statement
Existing scaling laws for Large Language Models (LLMs), such as the Chinchilla law, model test loss as a function of model size ($N$) and total training data ($D$). While effective for determining the optimal allocation of compute between $N$ and $D$, these laws do not explicitly account for training hyperparameters, specifically the batch size ($b$) and the number of training steps ($K$). Previous attempts to model hyperparameter scaling often fail to integrate with loss-based scaling laws or rely on the concept of "critical batch size" which, in some theoretical formulations, implies an optimal batch size of one—a contradiction to empirical findings where optimal batch sizes scale with the token budget. Furthermore, fitting scaling laws for optimal hyperparameters typically requires exhaustive sweeps of batch sizes for various token budgets, incurring massive computational costs.

Methodology
The authors propose a three-term scaling law (3TL) that explicitly decomposes the training data budget $D$ into batch size in tokens ($M = b \times s$, where $s$ is sequence length) and training steps ($K$), such that $D = MK$. The proposed functional form for the loss $L$ is:

$$L(N, M, K) = E + \frac{A}{N^\alpha} + \frac{B}{M^\beta} + \frac{C}{K^\gamma}$$

where $(E, A, B, C, \alpha, \beta, \gamma)$ are parameters fitted from training runs.

The authors compare this 3TL against a two-term law (2TL) approach, where the model size $N$ is fixed, and the law is fitted independently for each $N$ using the form $L(M, K) = E + B/M^\beta + C/K^\gamma$.

Key methodological steps include:

1. Derivation of Optimal Batch Size: By minimizing the 3TL with respect to $M$ subject to the constraint $D=MK$, the authors derive an optimal batch size $M^\star$ that scales as a power law of $D$: $M^\star \propto D^{\frac{\gamma}{\beta+\gamma}}$.
2. Data Utilization: Unlike previous methods that require identifying the single optimal batch size for each configuration (via full sweeps), the 3TL is fitted using data from suboptimal batch sizes. This allows the law to be robustly fitted with significantly fewer training runs.
3. Suboptimal Regime Modeling: To address hardware constraints where the theoretical optimal batch size is infeasible, the authors propose a two-stage fitting procedure. First, they use the 3TL to inform a simplified functional form $L(b) = \tilde{E} + \tilde{A}b^{-\tilde{\alpha}} + \tilde{B}b^{\tilde{\alpha}}$ (assuming $\gamma \approx \beta$). Second, they fit a power law to the interval of $\epsilon$-suboptimal batch sizes (where $\epsilon$ represents a 5% compute waste threshold).
4. Critical Batch Size Analysis: The authors analyze the behavior of the number of steps $K_{\bar{L}}(b)$ required to reach a target loss $\bar{L}$ as a function of batch size, comparing the 3TL predictions against empirical findings regarding the critical batch size.

Key Results
The authors evaluated their approach on two datasets: Li (Li et al., 2025a) and OpenEuroLLM.

• Optimal Batch Size Scaling: The 3TL successfully recovers the scaling of the optimal batch size $M^\star$. On the Li dataset, the derived scaling $M^\star \approx 0.667 \cdot D^{0.566}$ aligns closely with the "Step-Law" ($M^\star \approx 0.58 \cdot D^{0.571}$) from Li et al. (2025a), which was derived from a full batch size sweep.
• Compute Efficiency: A primary finding is that the 3TL can be fitted using only 28% of the training runs required by traditional methods. By masking the dataset to include only 2 batch sizes per $(N, D)$ combination (instead of a full sweep), the 3TL still recovers the correct scaling law for $M^\star$. In contrast, directly fitting a power law to the few available "best" batch sizes (without modeling the loss surface) results in highly unstable and inaccurate scaling laws.
• Suboptimal Batch Sizes: The two-stage fitting procedure successfully models the interval of suboptimal batch sizes that incur at most 5% compute waste. The authors find this interval has a width of roughly $2^2$ (a factor of 4) in batch size, and its scaling behavior is consistent across different model sizes and datasets.
• Critical Batch Size: The 3TL correctly models the phenomenon where the number of steps to reach a target loss becomes less sensitive to batch size increases beyond a critical point. Crucially, the model shows that this critical behavior scales with the token budget $D$ but is largely invariant to the model size $N$, matching empirical observations by Zhang et al. (2025).
• Comparison to Chinchilla: When the 3TL is reduced to a standard Chinchilla form (by substituting the optimal $M^\star$ back into the loss equation), the implied compute-optimal scaling for model size $N$ differs from a direct fit of the Chinchilla law. The 3TL suggests a smaller scaling exponent for $D$, highlighting that the choice of scaling law formulation impacts the derived optimal configurations.

Significance and Claims
The paper claims that the three-term law offers a unified framework that bridges the gap between Chinchilla-style scaling laws (model size vs. data) and hyperparameter scaling laws (batch size). Its significance lies in three main areas:

1. Robustness with Sparse Data: The ability to fit the law using suboptimal batch sizes drastically reduces the computational cost of deriving scaling laws, making it feasible to determine optimal batch sizes without exhaustive sweeps.
2. Theoretical Consistency: The formulation is structurally similar to convergence bounds in stochastic optimization, providing a theoretical bridge between empirical scaling analysis and optimization theory.
3. Practical Applicability: Unlike previous models that imply an optimal batch size of one or fail to describe suboptimal regimes, the 3TL provides a non-trivial optimal batch size and explicitly models performance in suboptimal regimes, which is critical for practical training under hardware constraints.

The authors note limitations, including potential inconsistencies in parameter values across different datasets (e.g., the irreducible loss term $E$) and the fact that the law still requires optimal learning rates for each configuration. They also acknowledge that the 3TL's implied compute-optimal scaling differs from direct Chinchilla fits, suggesting it may not be the sole instrument for determining optimal model sizes but is superior for understanding the interplay of steps, batch size, and data.