On the Role of Batch Size in Stochastic Conditional Gradient Methods

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

cs.LGMar 22, 2026

On the Role of Batch Size in Stochastic Conditional Gradient Methods

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

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

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.

Figures & tables

Panel of Figure 1 — Figure 1: Empirical verification of the validity of Assumption 3.3 during the training of a 124M NanoGPT model. The points with a loss below 5 fit a linear function well, with a slope equal to $\mu$ .

Panel of Figure 2 — Figure 2: Empirical gradient variance and fitted power-law models as functions of batch size $B$ with fixed sequence length $S=1024$ ( left ) and sequence length $S$ with fixed batch size $B=512$ ( right ) when training a 124M NanoGPT model on the FineWeb dataset under a fixed token budget $T=2.7$ B. For the left plot, the estimated scaling exponent is $\lambda\approx 0.9$ and $B_{\rm shift}\approx 90$ , while for the right plot they are $\lambda\approx 1.1$ and $S_{\rm shift}\approx 35$ . The fitted models support the validity of Assumption 3.4 .

Regime	Tuned at $\mathbf{T_{0}}$			Transfer Rule to $\mathbf{T_{1}}$			$\boldsymbol{\chi_{0\to 1}}$
	FW Stepsize / LR	Batch Size	Momentum	FW Stepsize / LR	Batch Size	Momentum
Increasing	$\beta_{0}^{\star}$	$B_{0}^{\star}S_{0}^{\star}$	$\alpha$	$\beta_{1}=\beta_{0}^{\star}{\color[rgb]{0.7,0,0}\left(\frac{T_{1}}{T_{0}}\right)^{-1/3}}\chi_{0\to 1}^{2/3}$	$B_{1}S_{1}=B_{0}^{\star}S_{0}^{\star}{\color[rgb]{0.7,0,0}\left(\frac{T_{1}}{T_{0}}\right)^{2/3}}\chi_{0\to 1}^{2/3}$	$\alpha$	$\frac{\rho_{1}}{\rho_{0}}$
Token Budget
Increasing							$\frac{\mu_{1}\rho_{1}L_{0}}{\mu_{0}\rho_{0}L_{1}}$
Token Budget and
Model Size

Table 1: Summary of BST scaling rules for different optimizers when transferring hyperparameters to larger token budget/model size.

$T=1.3$ B	Batch Size
$\boldsymbol{\beta},\times 10^{-4}$	64	128	256	512	1024	2048	4096
1.2	3.4258 ±0.0004	3.3889 ±0.0012	3.3857 ±0.0013	3.4074 ±0.0010	3.4587 ±0.0010	3.5598 ±0.0012	3.7715 ±0.0013
2.4	3.4394 ±0.0007	3.3880 ±0.0043	3.3706 ±0.0019	3.3801 ±0.0016	3.4144 ±0.0004	3.4940 ±0.0009	3.6694 ±0.0021
3.6	3.4554 ±0.0008	3.3945 ±0.0015	3.3717 ±0.0020	3.3765 ±0.0017	3.4065 ±0.0007	3.4799 ±0.0017	3.6472 ±0.0023
4.8	3.4766 ±0.0016	3.4072 ±0.0048	3.3790 ±0.0013	3.3807 ±0.0006	3.4115 ±0.0025	3.4945 ±0.0038	3.6611 ±0.0020
6.0	3.4967 ±0.049	3.4198 ±0.0002	3.3875 ±0.0019	3.3887 ±0.0024	3.4202 ±0.0022	3.5005 ±0.0038	3.7013 ±0.0160

Table 2: Final validation loss when training a 124M NanoGPT model varying the batch size while keeping the validation and train sequence lengths

1024

under the token budget

1.3

B (TPP 10.8). We report the average across

5

runs along with a standard deviation. Bold numbers indicate the best performance in the column. The runs in red indicate the best configuration of batch size, sequence length, and Frank–Wolfe stepsize across all runs for a given token budget.

$T=1.3$ B	Sequence Length
$\boldsymbol{\beta},\times 10^{-4}$	256	512	1024	2048	4096
1.2	3.7076 ±0.0084	3.4647 ±0.0240	3.4587 ±0.0010	3.4126 ±0.0014	3.4811 ±0.0025
2.4	3.9622 ±0.0585 ∗	3.4633 ±0.0091	3.3706 ±0.0019	3.3834 ±0.0026	3.4299 ±0.0021
3.6	4.0441 ±0.2055 ∗	3.4796 ±0.0131	3.3717 ±0.0020	3.3792 ±0.0022	3.4216 ±0.0013
4.8	3.9292 ±0.0852 ∗	3.5004 ±0.0063	3.3790 ±0.0013	3.3829 ±0.0020	3.4243 ±0.0029
6.0	3.9901 ±0.0170 ∗	3.5134 ±0.0059	3.3875 ±0.0019	3.3910 ±0.0024	3.4374 ±0.0037

Table 3: Final validation loss when training a 124M NanoGPT model varying the train sequence length while keeping the batch size

256

under the token budget

1.3

B. The validation sequence length is always

1024

. We report the average across

5

runs along with a standard deviation. ∗ indicates that not all runs had a stable decrease in validation loss. Bold numbers indicate the best performance in the column. The runs in red indicate the best configuration of batch size, sequence length, and Frank–Wolfe stepsize across all runs for a given token budget.

Model	$\boldsymbol{L}$	$\boldsymbol{\mu}$	$\boldsymbol{\rho}$	How to change $\boldsymbol{\beta}$ w.r.t. 124M model?	How to change $\boldsymbol{B}\boldsymbol{S}$ w.r.t. 124M model?
124M	$7.2$	$3.1$	$62.7$	$1\times$	$1\times$
1B	$10.6$	$2.9$	$111.9$	$\searrow 0.5\times\penalty\ {}^{(a)}$ $\searrow 0.54\times\penalty\ {}^{(b)}$	$\nearrow 4\times\penalty\ {}^{(a)}$ $\nearrow 4.37\times\penalty\ {}^{(b)}$

Table 4: Estimated problem-dependent constants, assuming that they change with the number of layers n_layer , embedding dimension n_embd , and batch size batch_size according to ( 15 ). The estimations of the change for

\beta

and

BS

are based on ( 7 ) and ( 8 ).

Panel of Figure 3 — Figure 3: Comparison of batch size and sequence length scheduling strategies when training a 1B model. The restarting schemes (in yellow and gray ) are compared against fixed schedules. The validation loss is evaluated with a smaller sequence length of $1024$ . The values of batch sizes $B_{0,1,2}$ , sequence lengths $S_{0,1}$ , and Frank–Wolfe stepsizes $\beta_{0,1}$ are given in the legends. The notation $(B_{0,1,2},S_{0,1,2},\beta_{0,1})$ characterizes which batch size, sequence length, and Frank–Wolfe stepsize are used for the particular setup. The notation $(B_{0},S_{0},\beta_{0})\to(B_{1,2},S_{1,2},\beta_{1})$ characterizes how parameters of Scion change after restart (e.g., batch size increases from $B_{0}$ to $B_{1,2}$ ), respectively. The notation $\mu$ P or BST indicates the rule used to select $B,S$ , and $\beta$ .

Panel of Figure 4 — Figure 4: Comparison of fixed large batch size strategies when training a 1B model. The validation loss is evaluated with a smaller sequence length $1024$ . Scion with a batch size of $1024$ suggested by our BST scaling rule achieves the best performance compared to other baselines with batch sizes $2048$ and $4096$ . The values of batch sizes $B_{1,2,3}$ , sequence lengths $S$ , and Frank–Wolfe stepsizes $\beta_{0,1}$ are given in the legends. The notation $(B_{1,2,3},S,\beta_{0,1})$ characterizes which batch size, sequence length, and Frank–Wolfe stepsize are used for the particular setup, respectively. The notation BST indicates the rule used to select the $B,S$ , and $\beta$ .

Panel of Figure 5 — Figure 5: The final performance of the 124M model when varying the Frank–Wolfe stepsize $\beta$ under different token budgets ( left: 2.7B, center: 5.3B, right: 8.0B). We average the train loss over 3 random seeds and report the moving average in a window of size 500. We observe that the BST scaling rule predicts a good estimate for the optimal $\beta$ when increasing the token budget. Moreover, the difference in performance between BST and $\mu$ P baselines grows with a token budget.

Panel of Figure 6 — Figure 6: The final performance of the 124M model when varying the momentum $\alpha$ under different token budgets ( left: 1.3B, center left: 2.7B, center right: 5.3B, right: 8.0B). We average the train loss over 3 random seeds and report the moving average in the window of size 500. We observe that rule momentum parameter $\alpha$ transfers under BST scaling.

Panel of Figure 7 — Figure 7: The final performance of the 1B model when varying the momentum parameter $\alpha$ ( left ) and Frank–Wolfe stepsize $\beta$ ( right ) under different token budget 10.6 TPP. We report the final train loss, smoothed in the window of size $500$ . We observe that the BST scaling rule predicts a good estimate for both optimal $\alpha$ and $\beta$ when transferring from a smaller 124M model to a larger 1B model.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	124M Model	775M Model	1B Model
Layers	12	36	18
Heads	6	20	16
Embedding Size	768	1280	2048
Weight Tying	Yes
Activation Function	ReLU 2
Vocabulary Size	50304

Appendix

Table B.1 : The model configurations and training details used in Section 6 .

Panel of Figure 13 — Figure B.1 : Evolution of the empirical gradient variance varying batch size $B$ with fixed sequence length $S=1024$ ( left ) and sequence length $S$ with fixed batch size $B=512$ ( right ) when training a 124M NanoGPT model on the FineWeb dataset. We observe that the variance quickly stabilizes after a short initial phase.

n_embd $\,\backslash\,$ n_layer	3	6	9	12	15	18	21	24	27	30
384	–	5.3	6.2	6.5	6.2	18	7.6	–	27	7.84
576	–	6.4	7.7	7.0	7.4	–	7.5	8.5	7.9	8.2
768	6.1	6.5	6.9	7.6	8.3	18	9.9	8.5	8.7	10.0
1152	–	–	–	8.8	9.4	10.8	9.5	–	–	–
1536	3	7.9	9.2	12	9.9	–	–	–	–	–
2304	–	9.9	–	–	–	–	–	–	–	–

Appendix

Table B.2 : Estimated

L

constant for various model configurations.

n_embd $\,\backslash\,$ n_layer	3	6	9	12	15	18	21	24	27	30
384	–	35.5	41.3	48.2	48.7	–	50.7	–	–	58.0
576	–	42.1	53.8	61.1	62.9	–	64.2	64.6	66.2	68.6
768	31.2	52.6	64.1	67.1	72.4	–	80.2	87.0	86.3	89.2
1152	–	–	–	76.6	81.1	87.3	—	–	–	–
1536	–	67.5	77.4	–	–	–	–	–	–	–
2304	–	–	–	–	–	–	–	–	–	–

Appendix

Table B.3 : Estimated

\rho

constant for various model configurations.

batch_size	256	512	1024	2048	4096
$\boldsymbol{\rho}$	48.9	52.6	55.0	56.4	57.9

Appendix

Table B.4 : Estimated

\rho

constant for a configuration with

6

layers and

768

embedding dimension when varying the batch size.

n_embd $\,\backslash\,$ n_layer	3	6	9	12	15	18	21	24	27	30
384	–	3.4	3.1	3.1	3.0	–	2.9	–	–	2.7
576	–	3.3	3.2	3.0	2.9	–	2.7	2.6	2.5	2.4
768	3.7	3.2	3.0	2.9	2.8	–	2.6	2.5	2.3	2.4
1152	–	–	–	2.7	2.7	2.8	2.6	–	2.5	–
1536	–	3.2	2.9	–	2.9	3.0	–	–	–	–
2304	–	3.6	–	–	–	–	–	–	–	–

Appendix

Table B.5 : Estimated Kurdyka–Łojasiewicz constant

\mu

for various model configurations.

Panel of Figure 18 — Figure C.1 : Comparison of two strategies: BST scaling rule, where $BS\sim T^{2/3}$ (fixed $1024$ batch size with $\beta_{1}=1.8\cdot 10^{-4}$ in light green and restarted version in gray ) and square-root rule ( 16 ), where $BS\sim T^{1/2}$ (fixed $736$ batch size with $\beta_{2}=1.27\cdot 10^{-4}$ in orange and restarted version in pink ). The validation and train sequence lengths are fixed to $1024$ . ( large batch size strategies when training a 1B model. Scion with a batch size of $1024$ (fixed from the beginning or after a restart) achieves slightly better performance compared to the baselines, where the batch size is set according to the square-root rule ( 16 ). The notation $(B_{1,2},S,\beta_{0,1})$ characterizes which batch size, sequence length, and Frank–Wolfe stepsize are used for the particular setup, respectively. The notation $(B_{0},S,\beta_{0})\to(B_{1,2},S,\beta_{1,2})$ characterizes how parameters of Scion change after restart (e.g., batch size increases from $B_{0}$ to $B_{1,2}$ ), respectively.

Panel of Figure 19 — Figure C.2 : Final performance of 124M model with $B=256,S=1024,\beta=3.6\cdot 10^{-4}$ and varying the momentum parameter $\alpha$ and device batch size.

Stochastic Optimization

Stochastic optimization studies how to find good decisions or model parameters when objectives, data, or feedback are random or noisy.

See all 210 papers

Adaptive Batch Sizes Using Non-Euclidean Gradient Noise Scales for Stochastic Sign and Spectral Descent
Feb 3, 2026Hiroki Naganuma, Shagun Gupta, Youssef Briki +4Stochastic Optimization Sign-Based Optimization
Mila, Montreal, Canada · Université de Montréal, Montreal, Canada · Meta Platforms, Menlo Park, California, USA
Deriving Hyperparameter Scaling Laws via Modern Optimization Theory
Mar 16, 2026Egor Shulgin, Dimitri von Rütte, Tianyue H. Zhang +3LMO Methods Scaling Laws
King Abdullah University of Science and Technology (KAUST), Thuwal, Saudi Arabia · ETH Zürich, Zürich, Switzerland · Mila, Quebec AI Institute & Université de Montréal, Quebec, Canada +1
How to Allocate Your Tokens? Scaling Laws with Training Steps and Batch Size
Jul 1, 2026Fabian SchaippLanguage Model Scaling Laws Scaling Laws
Inria, ´Ecole Normale Sup´erieure, PSL Research University, Paris

Feb 3, 2026J/K move · Enter open · S save

Adaptive Batch Sizes Using Non-Euclidean Gradient Noise Scales for Stochastic Sign and Spectral Descent

Hiroki Naganuma, Shagun Gupta, Youssef Briki, Ioannis Mitliagkas+3

Mila, Montreal, Canada · Université de Montréal, Montreal, Canada · Meta Platforms, Menlo Park, California, USA

Stochastic Optimization Sign-Based Optimization

Open paper AskPDF arXiv

To maximize hardware utilization, modern machine learning systems typically employ large constant or manually tuned batch size schedules, relying on heuristics that are brittle and costly to tune. Existing adaptive strategies based on gradient noise scale (GNS) offer a principled alternative. However, their assumption of SGD's Euclidean geometry creates a fundamental mismatch with popular optimizers based on generalized norms, such as signSGD / Signum (

\ell_\infty

) and stochastic spectral descent (specSGD) / Muon (

\mathcal{S}_\infty

). In this work, we derive gradient noise scales for signSGD and specSGD that naturally emerge from the geometry of their respective dual norms. To practically estimate these non-Euclidean metrics, we propose an efficient variance estimation procedure that leverages the local mini-batch gradients on different ranks in distributed data-parallel systems. Our experiments demonstrate that adaptive batch size strategies using non-Euclidean GNS enable us to match the validation loss of constant-batch baselines while reducing training steps by up to 66% for Signum and Muon on a 160 million parameter Llama model.

Abstract

Stochastic Optimization

Explore similar work

Adaptive Batch Sizes Using Non-Euclidean Gradient Noise Scales for Stochastic Sign and Spectral Descent

Deriving Hyperparameter Scaling Laws via Modern Optimization Theory

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