[Research Deep Dive] Scaling Law for Batch Size: Why µP Isn't Enough for Your Token Budget

TL;DR

Training large-scale models is a race against a fixed token budget (T). While frameworks like µP ensure your gradients don't explode at step zero, they don't tell you how to scale batch size ($B$) or sequence length ($S$) as you train longer. This paper provides a new theoretical compass: the BST Scaling Rule. By analyzing Stochastic Conditional Gradient (SCG) methods under the µ-KL condition, the authors prove that the optimal batch-sequence scale follows $BS\sim T^{2/3}$.

The Motivation: The "Too Large Batch" Trap

In modern LLM training, we often increase batch size to utilize hardware better. However, there is a known "critical batch size" beyond which optimization efficiency collapses under a fixed token budget.

Existing methods like µP (Maximal Update Parameterization) focus on local stability—ensuring updates stay $O(1)$ regardless of width. But optimization is a global journey. The authors argue that a learning rate that is optimal for a small budget becomes provably suboptimal as the token horizon expands. They shifted the perspective from "how many steps" to "how many tokens" ($T=K\cdot B\cdot S$).

Methodology: SCG and the µ-KL Geometry

The authors analyze Algorithm 1: Stochastic Conditional Gradient (SCG). Unlike standard SGD, SCG is projection-free and uses a Linear Minimization Oracle (LMO), making it a theoretical sibling to modern optimizers like Muon.

The Core Insight: The µ-KL Condition

The researchers validated that LLM training (specifically NanoGPT) follows the µ-Kurdyka–Łojasiewicz (µ-KL) condition. This condition effectively bridges the gap between the gradient's dual norm and the objective suboptimality: $$\Vert ablaf(x){\Vert }_{\ast }\ge \mu (f(x)-f^{\star })$$

The BST Scaling Rule

By solving the convergence bounds, they identified three distinct regimes for optimization error ($\epsilon$):

1. Noise-dominated: Performance improves as $BS$ increases.
2. Optimal/BST Regime: Error is independent of $B$ and $S$; this is the "sweet spot."
3. Iteration-starved: $B$ is so large that the model doesn't get enough updates ($K$ is too small), and error increases.

Figure 1: Empirical verification of the µ-KL condition. Note the linear fit between the dual gradient norm and training loss.

Practical Strategies: Hyperparameter Restart

The theory leads to a "Hyperparameter Restart" strategy. If you start with a small token budget $T_0$ and move to a larger $T_1$:

1. Increase Scale: $B_1{S}_1=B_0{S}_0(T_1/T_0{)}^{2/3}$.
2. Tune LR: Adjust $\beta$ according to the ratio of tokens and batch scale.
3. Keep Momentum: Interestingly, the momentum parameter $\alpha$ is a "constant of motion" that transfers directly.

Figure 2: Validation of gradient variance scaling, supporting the assumption that noise scales inverse-linearly with batch size.

Experimental Results

The authors tested their scaling laws by migrating from a 124M NanoGPT to a 1B parameter model.

• µP vs. BST: The µP baseline (fixed $B,S,\beta$) became increasingly suboptimal as the token budget grew.
• The Power of Restarts: Models that "restarted" their hyperparameters (increased $B$ mid-run) based on the BST rule outperformed all fixed large-batch baselines.
• Optimal Batching: For the 1B model, increasing $B$ beyond 1024 (the theory-predicted limit) resulted in a visible degradation in validation loss, confirming the "Iteration-starved" regime exists.

Figure 3: Restarting SCG (Yellow/Gray) consistently hits lower validation loss than fixed µP approaches.

Critical Analysis & Conclusion

Takeaway: This work provides a rigorous bridge between deep learning scaling laws and optimization theory. It proves that software-level hyperparameters (like batch size) must evolve with the data-level budget ($T$).

Limitations: The constants ($L,\mu ,\rho$) still require empirical estimation via power laws. While the authors provide these for NanoGPT, a practitioner on a completely different architecture (e.g., a Mixture-of-Experts) would need to re-estimate these.

Future Path: The interaction between these scaling rules and Schedule-Free optimizers or context-length extension techniques remains an open, high-value frontier.