On the Role of Batch Size in Stochastic Conditional Gradient Methods

TL;DR

This paper analyzes how batch size influences stochastic conditional gradient methods, revealing a regime-dependent effect where increasing batch size initially helps but can eventually hinder performance, and provides practical guidelines for training.

Contribution

It offers a new theoretical analysis of batch size effects in stochastic conditional gradient algorithms, including adaptive strategies for large-scale training.

Findings

Optimal batch size depends on the regime, with diminishing returns beyond a threshold.

Theoretical predictions align with empirical results on NanoGPT.

Guidelines for selecting batch size and stepsize in practice.

Abstract

We study the role of batch size in stochastic conditional gradient methods under a $\mu$-Kurdyka-{\L}ojasiewicz ($\mu$-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.