[Research Deep Dive] Scaling Law for Batch Size: Why µP Isn't Enough for Your Token Budget
On the Role of Batch Size in Stochastic Conditional Gradient Methods
This paper introduces a token-budget-aware theoretical framework for momentum-based Stochastic Conditional Gradient (SCG) methods, such as Scion, under a µ-Kurdyka–Łojasiewicz (µ-KL) condition. It identifies a "BST scaling rule" that governs the optimal relationship between batch size (B), sequence length (S), and token budget (T), achieving state-of-the-art results in hyperparameter transfer for NanoGPT.
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 () or sequence length () 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 .
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 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" ().
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:
abla f(x)\|_* \geq \mu(f(x) - f^\star)$$ ### The BST Scaling Rule By solving the convergence bounds, they identified three distinct regimes for optimization error ($\varepsilon$): 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.