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

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.

• The paper introduces a three-term scaling law that models LLM loss as a function of model size, batch size, and training steps.
• It demonstrates that optimal batch size scales with dataset size and drastically reduces the need for exhaustive hyperparameter tuning.
• Empirical validation shows improved predictive accuracy even under suboptimal hardware constraints, offering practical guidance for compute allocation.

A Unified Scaling Law for Model Size, Training Steps, and Batch Size in LLM Training

Introduction

The extensive empirical investigation “How to Allocate Your Tokens? Scaling Laws with Training Steps and Batch Size” (2607.01487) proposes a unified power-law scaling framework for LLM pretraining that decomposes dataset size into training steps and batch size, effectively interpolating between classical Chinchilla-style scaling and batch size scaling analyses. This work substantially broadens the operational applicability of scaling laws, enabling robust loss prediction and compute allocation not only for optimal but also suboptimal hyperparameter choices, and reducing the empirical burden for scaling law estimation.

Three-Term Power Law and Its Theoretical Rationale

The principal contribution is the introduction of a three-term scaling law for the final loss $\mathcal{L}$ as a function of model size $N$ (parameter count), batch size $M$ (in tokens), and number of training steps $K$:

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

This approach decouples training data size $D$ as $D = MK$ and differentiates between sample efficiency (controlled by $M$ and $K$) and model efficiency ($N$) in loss reduction. The formulation is tightly linked with convergence results from stochastic optimization, where loss typically exhibits power-law decrease with $N$0 and $N$1 under optimal schedules.

The law naturally subsumes the canonical Chinchilla law when $N$2 is optimized, and plug-in optimal batch size $N$3 (with corresponding $N$4) yields a two-term law in $N$5 with reparametrized data scaling, aligning with Chinchilla functional form. The law thereby bridges empirical scaling analyses and theoretically-grounded optimization limits.

Empirical Validation and Comparison with Alternative Laws

A comprehensive study on two large-scale LLM training datasets (the Li dataset [Li2025] and OpenEuroLLM) establishes several empirical findings:

• The three-term law recapitulates the scaling of the optimal batch size $N$6 observed in prior work, demonstrating that $N$7 scales roughly as $N$8 and is independent of $N$9 in the regimes tested, in agreement with state-of-the-art hyperparameter scaling analyses.
• Fitting the three-term law is highly data-efficient: with only two batch size values per $M$0, the resulting estimate of the optimal batch size scaling is robust, significantly reducing the compute requirements compared to exhaustive batch size sweeps.

Figure 1: Estimates for $M$1-scaling coefficient $M$2 and implied scaling law for $M$3 compared against the empirical best.

The mean absolute deviation (MAD) analysis (Figure 2) shows that the three-term law achieves better out-of-sample error than model-size-specific two-term laws, even as the number of parameters grows.

Figure 2: MAD comparison of three-term and two-term laws on train (left) and validation (right) splits.

Even on datasets with incomplete batch size sweeps, the law retains its predictive power. Reducing the Li dataset to 42% and 28% of its runs does not substantially alter the scaling exponent of $M$4 (see Figure 3), in stark contrast to direct fitting of $M$5 which becomes highly unstable.

Figure 3: Scaling law fits for $M$6 with reduced dataset size (42% of batch size values per sweep).

Modeling Suboptimal and Critical Batch Size Regimes

A significant practical concern is characterization of loss behavior at suboptimal batch sizes induced by hardware constraints. The three-term law models loss for arbitrary $M$7 allocations, not just the optimum, enabling derivation of loss envelopes and the range of batch sizes that incur at most an $M$8-fraction compute regret.

Where higher accuracy is required, a two-stage fitting procedure is introduced:

1. For fixed $M$9, fit the cross-section loss as a function of $K$0 via a two-parameter power law.
2. Model the interval $K$1 where loss remains within a small excess of the minimum, showing this interval remains stable and scalable (see Figure 4). This provides a practical prescription for batch size selection when hardware is limiting.

Figure 5: For $K$2M, the three-term law predicts optimal batch size, but is less accurate in loss value at boundaries of token budget.

Figure 4: Batch size range $K$3 with $K$4-suboptimal loss derived for $K$5M.

An important theoretical consequence is the recovery of critical batch size phenomena: under the three-term law, the number of steps to achieve a fixed target loss is invariant to $K$6 but increases with $K$7, precisely mirroring empirical investigations of critical batch size scaling, and offering an explanation why earlier critical batch size theories which predict optimal batch size of one are inadequate for practical LLM training.

Figure 6: Function $K$8 at Chinchilla-optimal scaling, showing critical batch size varies with $K$9 but not with $$\mathcal{L}(N, M, K) = E + \frac{A}{N^\alpha} + \frac{B}{M^\beta} + \frac{C}{K^\gamma}$$0.

Optimal Compute Allocation and Limitations

While the three-term law reduces to a Chinchilla-style prediction for loss vs. $$\mathcal{L}(N, M, K) = E + \frac{A}{N^\alpha} + \frac{B}{M^\beta} + \frac{C}{K^\gamma}$$1 when $$\mathcal{L}(N, M, K) = E + \frac{A}{N^\alpha} + \frac{B}{M^\beta} + \frac{C}{K^\gamma}$$2 is optimal, inferred optimal model size scaling can differ from direct Chinchilla law fitting, especially in terms of data scaling exponents, highlighting that the choice of scaling law formulation affects optimal compute allocation. Consequently, accurate prediction of compute-optimal $$\mathcal{L}(N, M, K) = E + \frac{A}{N^\alpha} + \frac{B}{M^\beta} + \frac{C}{K^\gamma}$$3 pairs remains sensitive to model design.

Implications, Limitations and Future Work

The three-term law offers a practical tool for:

• Predicting LLM loss for arbitrary (even suboptimal) allocations of dataset into batch size and training duration, informing training under hardware constraints.
• Drastically reducing empirical effort needed for estimating hyperparameter scaling laws, making such estimation tractable for new model classes or architectures.
• Providing a unified framework for theory-grounded scaling laws that incorporate both model size and optimizer hyperparameters.

Limitations include:

• Quantitative discrepancies between datasets can arise due to varying training protocols; generalization to all architectures or regimes requires further study.
• The law does not explicitly model learning rate, which may modulate scaling in certain regimes.
• While optimal batch size is well-predicted, suboptimal performance envelopes require careful two-stage fitting, especially for high accuracy.

Extensions could incorporate learning rate scaling (as in [Shulgin2026], [Bjorck2024]) or modifications for different optimizer families (e.g., Muon [EssentialAI2025]), and analyze impacts in other modalities and tasks.

Conclusion

This work presents a three-term scaling law that subsumes and extends both the Chinchilla law and batch size scaling analyses, offering a unified, robust, and compute-efficient approach to managing data allocation in LLM training. It provides strong empirical support, theoretical justification, and actionable prescriptive guidance for practitioners and scaling law theorists alike, while highlighting the importance of law formulation in calculating truly compute-optimal training regimes.