On the Role of Batch Size in Stochastic Conditional Gradient Methods

arXiv:2603.21191 · cs.LG, math.OC, stat.ML · Submitted 2026-03-22 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Hi, I'm Tom.

Jane: And I'm Jane, here with Lu, Meng and our in-house large language model, Lalam.

Tom: We're looking at the paper: "On the Role of Batch Size in Stochastic Conditional Gradient Methods".

Jane: Here's the gist of the paper. The gist The study derives a new analysis that explicitly captures the interaction between stepsize, batch size, and stochastic noise in momentum-based stochastic conditional gradient methods under a µKurdyka–Łojasiewicz (µ-KL) condition.

Tom: Role of Batch Size in Optimization The paper investigates the role of batch size in stochastic conditional gradient methods under a µ-KL condition, focusing on momentum-based algorithms like Scion, and reveals a regime-dependent behavior where increasing the batch size initially improves optimization accuracy but saturates beyond a critical threshold beyond a critical threshold, the benefits saturate and can eventually degrade performance under a fixed token budget.

Jane: The theory predicts the magnitude of the optimal stepsize and aligns well with empirical practices observed in large-scale training.

Tom: Theoretical Framework and Convergence Guarantees The analysis is conducted for stochastic optimization under smoothness (A1), norm equivalence (A2), and a µ-KL error bound (A3).

Title and authors: Tom: So we're looking at the paper "On the Role of Batch Size in Stochastic Conditional Gradient Methods," and what it's about is really about how batch size interacts with stepsize and noise when you use momentum-based methods.

Jane: It sounds super technical, so if I translate that, are they trying to figure out if making your batch size bigger always helps you train better, or if there's a point where it actually starts hurting things?

Lu: Exactly. The core idea here is that they’ve created a new analysis that explicitly captures how those three things—stepsize, batch size, and stochastic noise—are all linked together under this mu-KL condition.

Meng: So they're moving beyond just saying "big batches are better" and showing us the specific mathematical reason why that benefit eventually hits a ceiling when you have a fixed token budget.

Lalam: They specialize their convergence bounds for the fixed-token setting, which means they've made it explicit how much your optimization error depends on the effective batch-sequence scale BS, which is what they callBS star zero S star zero T one over T zero mu one over mu zero rho one over rho zero L one over L0 two over three.

Tom: That specific scaling rule they derive is actually a huge piece of information because it gives us a concrete target for setting our batch size and sequence length, rather than just guessing.

Jane: So what this means for the listener is that we’re not just looking at batch size in isolation; we have to consider how it relates to the overall token budget of the training run.

The paper's summary: Tom: We've covered the context, so now let's get into what they actually found in "On the Role of Batch Size in Stochastic Conditional Gradient Methods." They show a regime-dependent behavior where increasing batch size initially helps accuracy, but then beyond a certain point, it just levels off or starts getting worse under a fixed token budget.

Jane: When you say "level off," what does that look like in practice? Is the model still learning effectively even if we keep feeding it more data?

Lu: It means they found three distinct regimes: first, a noise-dominated regime where bigger batches are good; second, an intermediate regime where the best error doesn't change much with batch size; and third, a large-batch regime where performance actually starts to degrade under a fixed token budget.

Meng: That "too-large batch hurts" idea is important because it moves it from just being an observation from training runs to something they’ve modeled mathematically.

Lalam: They show that this happens even when using momentum-based algorithms like Scion, and they align this theory with what we see happening in large-scale training practices.

Tom: So the big takeaway here is that performance isn't just a smooth climb; it has these distinct phases depending on how you set your batch size and sequence length relative to the total token budget.

Jane: It’s a bit like trying to optimize a recipe: sometimes adding more ingredients makes things taste better, but past a certain point, you just end up with too much salt or too much sugar for the dish.

Lu: Precisely. And they derive this BST scaling rule, which is BS equal to T two over three up to some problem-dependent factors that they figured out themselves <ref:2603.21191#pg1>.

The paper's improvements: Tom: Now we get into the actionable advice from "On the Role of Batch Size in Stochastic Conditional Gradient Methods." This is where they give us strategies for actually designing our training schedules.

Jane: Are we talking about a specific schedule, like a warm-up phase that tells us when to change our batch size dynamically as training goes on?

Lu: They derive principled guidelines for selecting and updating the momentum parameter alpha by deriving scaling relations seven and eight, which helps with this.

Meng: And they propose a two-stage protocol, which can be generalized into multi-stage protocols, that they tested empirically on NanoGPT. That suggests we can actually build these stages directly into our training pipeline for better performance when scaling up the model size.

Lalam: For increasing the model size, the strategy involves picking B and S using that specific formula involving BS star zero S star zero T one over T zero mu one over mu zero rho one over rho zero L one over L0 two over three.

Tom: And then they scale the momentum parameter alpha by a factor involving square roots of those ratios, which is exactly what they suggest for keeping things stable when you switch from smaller base models to bigger ones, so we don't have to manually retune it.

Jane: So these improvements are focused on giving us concrete scaling relations and protocols that we can use to dynamically adjust batch size and sequence length as training progresses based on the available token budget, which sounds very practical for our production pipelines.

Lu: The main improvement is deriving this BST scaling rule itself, which tells us we should pick the batch size and sequence length right at that crossover point between the second and third regimes <ref:2603.21191#pg1>.

Conclusion: Tom: So we've covered a lot about "On the Role of Batch Size in Stochastic Conditional Gradient Methods." We've talked about the theoretical framework for batch size scaling and the new protocols they suggest.

Jane: What does this all mean for us as a listener who doesn't read the paper? Can you give us that final summary of what this means in plain English?

Lu: It means we have a new way to think about batch size not as a fixed number, but something that should be adjusted dynamically based on where we are in the training stage and how many tokens we have left.

Meng: For production, it means we can build more robust schedules that handle that saturation effect where performance levels off and starts degrading under resource constraints if we don't follow the right scaling rules.

Lalam: This paper gives us a framework for global scaling laws that were missing before, providing a solid mathematical foundation for how those components should scale globally with the token budget.

Tom: To wrap up, "On the Role of Batch Size in Stochastic Conditional Gradient Methods" gives us a solid theoretical framework for understanding batch size scaling in stochastic conditional gradient methods and offers guidance for designing efficient training schedules.

Jane: And for us, it means moving towards more principled guidelines for selecting batch size and sequence length based on that BST scaling rule and adaptive strategies validated on NanoGPT.

Lu: It offers a way to understand the paper's core idea: that large batch sizes don't inherently degrade performance when you choose your parameters according to the BST scaling rule, suggesting setting BS at the crossover between the second and third regime.

Meng: I just want to add that they also point out some limitations, like they didn’t provide explicit optimization error laws as functions of B S T, nor did they characterize exactly when increasing batch size becomes provably detrimental under a fixed token budget.

Lalam: They also flagged that the transfer of the momentum parameter alpha might be shifting to lower values when increasing model size because they think that could be due to insufficient data points or perhaps we should incorporate parameters like the token budget into their scaling analysis.

Tom: So, "On the Role of Batch Size in Stochastic Conditional Gradient Methods" gives us a solid theoretical framework for understanding batch size scaling in stochastic conditional gradient methods and offers guidance for designing efficient training schedules.

Rustem Islamov, Roman Machacek, Aurelien Lucchi, Antonio Silveti-Falls, Eduard Gorbunov, Volkan Cevher

University of Basel, Switzerland · University of Bern, Switzerland · CentraleSupélec, France · MBZUAI, UAE · EPFL, Switzerland

cs.LG, math.OC, stat.ML

Submitted: 2026-03-22

Updated: 2026-10-18

Journal ref: Proceedings of the 43rd International Conference on Machine Learning, Seoul, South Korea. PMLR 306, 2026

Code: https://github.com/KellerJordan/modded-nanogpt

License: http://creativecommons.org/licenses/by-sa/4.0/

The gist: The gist The study derives a new analysis that explicitly captures the interaction between stepsize, batch size, and stochastic noise in momentum-based stochastic conditional gradient methods under a

Key concepts

$\mu$-Kurdyka–Łojasiewicz ($\mu$-KL) condition: This is a mathematical property that describes the shape of the loss function in optimization. It guarantees that the optimization process will converge reliably, even when dealing with non-smooth functions, by ensuring a certain level of regularity near the minimum.
Effective Batch–Sequence Scale (BS): This concept combines batch size ($B$) and sequence length ($S$) into a single measure. The paper shows that the optimal performance depends on this combined scale, not just batch size alone, guiding how to balance data utilization with computational constraints.
BST Scaling Rule: This is a derived formula suggesting an optimal relationship between the effective batch-sequence-token scale ($BS$) and the total token budget ($T$). It provides a concrete guideline for selecting $B$ and $S$ (e.g., $BS \sim T^{2/3}$) to maintain high efficiency across different model sizes.

Terminology

Accuracy — 1× in this paper · explained in 1 episode
Adaptive Strategy — 1× in this paper
Convergence Guarantees — 1× in this paper · explained in 3 episodes

Summary

The gist The study derives a new analysis that explicitly captures the interaction between stepsize, batch size, and stochastic noise in momentum-based stochastic conditional gradient methods under a µKurdyka–Łojasiewicz (µ-KL) condition.

Role of Batch Size in Optimization

The paper investigates the role of batch size in stochastic conditional gradient methods under a µ-KL condition, focusing on momentum-based algorithms like Scion, and reveals a regime-dependent behavior where increasing the batch size initially improves optimization accuracy but saturates beyond a critical threshold beyond a critical threshold, the benefits saturate and can eventually degrade performance under a fixed token budget. The theory predicts the magnitude of the optimal stepsize and aligns well with empirical practices observed in large-scale training.

Theoretical Framework and Convergence Guarantees

The analysis is conducted for stochastic optimization under smoothness (A1), norm equivalence (A2), and a µ-KL error bound (A3). Specializing the convergence bounds to the fixed-token setting T = KBS yields an explicit, non-monotone dependence of the achievable optimization error on the effective batch–sequence scale BS. Three regimes emerge: (i) a noise-dominated regime where increasing BS improves performance, (ii) an intermediate regime where the best achievable error is essentially independent of BS, and (iii) a large-batch regime where performance deteriorates as BS grows under a fixed token budget.

The BST Scaling Rule

Balancing the dominant terms yields a critical effective batch–sequence–token (BST) scale rule BS ≍ T 2/3 up to problem-dependent factors that we derive in this work, revealing how curvature, noise, geometry, and error-bound strength shift the optimal operating point. This suggests that large batch sizes do not inherently degrade performance: when batch size, sequence length, and learning rate are chosen according to our BST scaling rule, large-batch training remains effective and token-efficient. The BST scaling rule suggests the following choice of parameters in Algorithm 1: BS ∼ T 2/3, β ∼ 1/K, α = Const.

Actionable Scheduling Strategies

The paper derives principled guidelines for selecting and updating (β, B, S) during training under a fixed token budget. This includes deriving scaling relations (7)–(8) and a two-stage (and more generally multi-stage) protocol validated empirically on NanoGPT. For increasing model size, the strategy involves choosing B and S according to B1S1 = B⋆0S⋆0 T1/T0 µ1/µ0 ρ1/ρ0 L1/L0 2/3 and β = β⋆ 0 (√T0 √T1 µ1/µ0 ρ1/ρ0 L1/L0) 2/3.

Empirical Validation

Experiments on NanoGPT are consistent with the theoretical predictions and illustrate the emergence of the predicted scaling regimes. The empirical verification of Assumption 3.4 shows that after a short initial phase, the empirical variance stabilizes and fluctuates around the average, suggesting that the variance is fixed during most of the training. The joint evolution of the dual gradient norm and the training loss shows that once training loss falls below approximately 5, data points closely follow a linear relationship, empirically supporting the use of Assumption 3.3.

Conclusion

The work provides a theoretical framework for understanding batch size scaling in stochastic conditional gradient methods and offers guidance for designing efficient training schedules in large-scale optimization. The results complement classical large-batch heuristics such as linear learning rate scaling with warmup [Goyal et al., 2017] and adaptive batch size schedules [Smith et al., 2018], while offering a projection-free viewpoint rooted in conditional gradient geometry. The BST scaling rule suggests selecting the batch size and sequence length at the crossover between the second and third regimes.

Limitations and Future Work

The choice of radii η and variance initialization were adopted directly from Pethick et al. [2025a] without additional tuning. Another important question that requires further investigation is the transfer of the momentum parameter. The shift in the range of near-optimal values of the momentum parameter α to lower values (0.06–0.08) when increasing model size may be due to insufficient data points or because additional training parameters, such as the token budget, should be incorporated into the scaling analysis.

Acknowledgements

Rustem Islamov and Aurelien Lucchi acknowledge the financial support of the Swiss National Science Foundation, SNSF grant No 207392. Volkan Cevher acknowledges the financial support of the Swiss National Science Foundation, SNSF grant No 240094. This work was also supported under project ID 37 as part of the Swiss AI Initiative, through a grant from the ETH Domain and computational resources provided by the Swiss National Supercomputing Centre (CSCS) under the Alps infrastructure.

References

The references listed are: Foivos Alimisis, Rustem Islamov, and Aurelien Lucchi. Why do we need warm-up? a theoretical perspective. arXiv preprint arXiv:2510.03164, 2025. Shane Bergsma, Nolan Simran Dey, Gurpreet Gosal, Gavia Gray, Daria Soboleva, and Joel Hestness. Power lines: Scaling laws for weight decay and batch size in LLM pre-training. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du, Zhe Fu, et al. Deepseek llm: Scaling open-source language models with longtermism. arXiv preprint arXiv:2401.02954, 2024. Jérôme Bolte, Aris Daniilidis, and Adrian Lewis. The łojasiewicz inequality for nonsmooth subanalytic functions with applications to subgradient dynamical systems. SIAM Journal on Optimization, 2007. Jérôme Bolte, Shoham Sabach, and Marc Teboulle. Proximal alternating linearized minimization for nonconvex and nonsmooth problems. Mathematical Programming, 146(1):459–494, 2014. Enea Monzio Compagnoni, Rustem Islamov, Frank Norbert Proske, and Aurelien Lucchi. Unbiased and sign compression in distributed learning: Comparing noise resilience via sdes. arXiv preprint arXiv:2502.17009, 2025a. Enea Monzio Compagnoni, Tianlin Liu, Rustem Islamov, Frank Norbert Proske, Antonio Orvieto, and Aurelien Lucchi. Adaptive methods through the lens of SDEs: Theoretical insights on the role of noise. In The Thirteenth International Conference on Learning Representations, 2025b. Michael Crawshaw, Chirag Modi, Mingrui Liu, and Robert M Gower. An exploration of noneuclidean gradient descent: Muon and its many variants. arXiv preprint arXiv:2510.09827, 2025. Nolan Dey, Bin Claire Zhang, Lorenzo Noci, Mufan Li, Blake Bordelon, Shane Bergsma, Cengiz Pehlevan, Boris Hanin, and Joel Hestness. Don’t be lazy: Completep enables compute-efficient deep transformers. arXiv preprint arXiv:2505.01618, 2025. Ilyas Fatkhullin, Jalal Etesami, Niao He, and Negar Kiyavash.

Improvements for AI systems

Optimizing batch size and sequence length for large language models under a fixed token budget using a specific scaling rule ensures that large batch sizes do not degrade performance beyond a critical threshold by choosing B and S such that BS is proportional to T to the 2/3 power. This allows for maximizing hardware utilization while maintaining convergence guarantees.
Implementing an adaptive scheduling strategy where batch size and sequence length are increased during training based on the available token budget allows the system to leverage more data as it becomes available, leading to improved optimization efficiency across different training stages.
Utilizing a derived scaling rule for the Frank-Wolfe stepsize ensures that momentum parameter alpha transfers effectively when moving from smaller base models to larger ones, maintaining near-optimal performance without requiring retuning of this hyperparameter.

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.