On the Role of Batch Size in Stochastic Conditional Gradient Methods

Abstract: We study the role of batch size in stochastic conditional gradient methods under a $μ$-Kurdyka-Łojasiewicz ($μ$-KL) condition. Focusing on momentum-based stochastic conditional gradient algorithms (e.g., Scion), we derive a new analysis that explicitly captures the interaction between stepsize, batch size, and stochastic noise. Our study reveals a regime-dependent behavior: increasing the batch size initially improves optimization accuracy but, beyond a critical threshold, the benefits saturate and can eventually degrade performance under a fixed token budget. Notably, the theory predicts the magnitude of the optimal stepsize and aligns well with empirical practices observed in large-scale training. Leveraging these insights, we derive principled guidelines for selecting the batch size and stepsize, and propose an adaptive strategy that increases batch size and sequence length during training while preserving convergence guarantees. Experiments on NanoGPT are consistent with the theoretical predictions and illustrate the emergence of the predicted scaling regimes. Overall, our results provide a theoretical framework for understanding batch size scaling in stochastic conditional gradient methods and offer guidance for designing efficient training schedules in large-scale optimization.

• The paper establishes the BST scaling rule (BS ~ T^(2/3)) to determine optimal batch size in momentum-based SCG methods under a fixed token budget.
• The paper demonstrates non-monotonic convergence across noise-dominated, power-law, and iteration-starved regimes, validated through empirical NanoGPT experiments.
• The paper offers practical hyperparameter transfer guidelines that leverage token-complexity-aware analysis to improve efficiency in LLM training and projection-free optimization.

The Role of Batch Size in Stochastic Conditional Gradient Methods under a Fixed Token Budget

Introduction and Theoretical Insights

The paper "On the Role of Batch Size in Stochastic Conditional Gradient Methods" (2603.21191) establishes a unified theoretical and empirical framework for understanding scaling behaviors of batch size, sequence length, and optimizer stepsize in momentum-based stochastic conditional gradient (SCG) and projection-free algorithms. Modern large-scale LLM training is constrained primarily by token budget rather than iteration count, necessitating a global, token-complexity-aware analysis of hyperparameter selection. The work derives explicit non-monotonic error laws in terms of effective batch—sequence product ($BS$) under a fixed budget $T$ within the Kurdyka–Łojasiewicz ($\mu$-KL) landscape, capturing broad practical observations in a precise theoretical regime.

Concretely, SCG algorithms such as Scion or Muon (with weight decay decoupled from learning rate) are analyzed under structured nonconvexity and general geometries. The central result is the identification of regime-dependent scaling of convergence: there is a critical $BS\sim T^{2/3}$ (the “BST scaling rule”) beyond which further increases in effective batch cease to be beneficial and may become detrimental. This global scaling contrasts with local hyperparameter transfer frameworks (e.g., $\mu$P), which cannot capture batch size induced phase transitions under growing token budgets.

Empirical Verification of Key Assumptions

The analysis is grounded in empirically validated conditions. The $\mu$-KL property, which relates first-order stationarity to dual-norm optimality gaps, is directly validated on transformer-style LLMs, showing a robust linear scaling between loss and dual gradient norm once well into training.

Figure 1: Empirical verification of the $\mu$-KL condition in NanoGPT training; a linear relationship arises for sufficiently low loss, allowing for a robust estimate of $\mu$.

For gradient noise, the variance across batches/sequence lengths is shown to adhere to an inverse power law with $BS$, matching the presumed sub-Gaussian/bounded-variance noise model essential for the convergence proofs.

Figure 2: Empirical gradient variance versus batch size and sequence length. The variance scales almost inversely with $BS$, confirming the analytic noise model.

Scaling Laws and Phase Regimes

The explicit error law under a fixed token budget embodies three qualitative regimes:

1. Noise-Dominated Regime: For small $BS$, optimization error $\varepsilon$ decreases with increasing $BS$.
2. Batch-Independent Power Law Regime: For intermediate $BS$, the error saturates to a $T^{-1/3}$-like scaling—insensitive to $BS$ and driven by noise “floor”.
3. Iteration-Starved Regime: For $BS \gg T^{2/3}$, the optimization error grows linearly with $BS$, marking a sharp loss of token efficiency.

This non-monotonicity prescribes the BST scaling rule: $BS \sim T^{2/3}$, with stepsize $\beta\sim 1/K$ (total iterations). Larger $BS$ is not inherently harmful if learning rate/stepsize and batch size grow jointly according to theory. The region of optimality is further influenced by problem-dependent constants (geometry, curvature, dual-norm scaling), all of which are empirically estimated and provided.

Empirical Results and Ablation Studies

Systematic ablation confirms that the predicted $BS$ and stepsize values yield the lowest validation losses in NanoGPT experiments, with robust transfer across model scales and token budgets. Notably, methods proposed in prior literature (e.g., $\mu$P-style naive hyperparameter transfer) become strictly suboptimal as token budgets or models increase.

Figure 3: Batch size/sequence scheduling strategies with restarts—BST scaling and restarting leads to clear improvements and aligns with optimal scaling predictions; static $\mu$P tuning plateaus suboptimally.

Increasing $BS$ beyond the BST-predicted value results in degradation, confirming the theoretical “too large batch size” regime.

Figure 4: Increasing $BS$ beyond the BST rule (e.g., $B=2048,4096$) hurts performance; the BST-tuned configuration yields minimal validation loss.

Additional hyperparameter transfer studies reveal that momentum and stepsize predicted by the theory transfer well across both increasing token horizon and model scale.

Figure 5: Predicted optimal Frank-Wolfe stepsize ($\beta$) as a function of token budget matches the empirical minima over a range of token budgets for the 124M model.

Figure 6: Momentum ($\alpha$) transfer across budgets confirms the scaling law’s prediction—optimal values remain stable unless regime boundaries are crossed.

Figure 7: Scaling of optimal $\beta$ with model size and token budget for a 1B parameter model demonstrates transfer power of BST scaling.

Practical Hyperparameter Transfer and Adaptive Scheduling

Explicit guidelines are derived for practitioners: given tuned hyperparameters on a small model/token budget, scaling to larger settings is achieved by adjusting $BS$ and $\beta$ according to the BST scaling rule, possibly with stage-wise restarts as more data or compute becomes available. Empirical studies confirm that increasing batch size and sequence length progressively in two stages (with recomputed stepsize) provides superior convergence.

Implications and Future Directions

This work rigorously demonstrates that “large batch hurts” only emerges when scaling is misaligned with the problem’s budget and structure. Correct joint scaling of $(B,S,\beta)$ avoids the optimization slowdowns and instability typically associated with large $B$, while yielding maximal hardware efficiency.

Beyond LLM training, the results provide a template for scaling analysis in high-dimensional non-Euclidean or projection-free optimization. The structured analysis under the $\mu$-KL condition and global-tracking of noise and curvature constants supports robust out-of-sample transferability, motivating further development of hyperparameter-scaling transfer frameworks that account for token-complexity rather than simple model-width constraints.

Future theoretical investigations could extend to adaptively estimating problem constants online, or to more heterogeneous, non-i.i.d. data settings—especially those encountered in continual learning or federated environments. Practical future work should evaluate scaling laws outside of language modeling, such as in vision or multimodal settings, testing the global efficiency of SCG approaches vs. explicit projection-based SGD or adaptive methods.

Conclusion

This paper provides rigorous, data-informed scaling laws for batch size, sequence length, and optimizer stepsize in momentum-based SCG methods under global token budgets. Non-monotone dependence of convergence on batch size is theoretically and empirically established, with concrete predictive guidance for optimal batch scheduling. The results supersede local-approximation-based frameworks and remove much of the guesswork from large-scale optimizer configuration in LLM and related domains.