# Unlocking optimal batch size schedules using continuous-time control and perturbation theory

arXiv: https://arxiv.org/abs/2312.01898 · HTML: https://mlsanity.com/paper/2312.01898/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 8 of its 22 references that are in the map.
- First posted on arXiv 4 Dec 2023.

## Builds on (most cited first, 8 of 8)

- [Optimization Methods for Large-Scale Machine Learning](https://mlsanity.com/paper/1606.04838/) (2016)
- [Don't Decay the Learning Rate, Increase the Batch Size](https://mlsanity.com/paper/1711.00489/) (2017)
- [Coupling Adaptive Batch Sizes with Learning Rates](https://mlsanity.com/paper/1612.05086/) (2016)
- [Causal optimal transport and its links to enlargement of filtrations and continuous-time stochastic optimization](https://mlsanity.com/paper/1611.02610/) (2016)
- [A Diffusion Theory For Deep Learning Dynamics: Stochastic Gradient Descent Exponentially Favors Flat Minima](https://mlsanity.com/paper/2002.03495/) (2020)
- [Stochastic Modified Flows, Mean-Field Limits and Dynamics of Stochastic Gradient Descent](https://mlsanity.com/paper/2302.07125/) (2023)
- [A continuous-time analysis of distributed stochastic gradient](https://mlsanity.com/paper/1812.10995/) (2018)
- [Adversarial Training for Gradient Descent: Analysis Through its Continuous-time Approximation](https://mlsanity.com/paper/2105.08037/) (2021)

## Built on by (most cited first, 3 of 3)

- [Predictable Scale: Part I, Step Law -- Optimal Hyperparameter Scaling Law in Large Language Model Pretraining](https://mlsanity.com/paper/2503.04715/) (2025)
- [Fast Catch-Up, Late Switching: Optimal Batch Size Scheduling via Functional Scaling Laws](https://mlsanity.com/paper/2602.14208/) (2026)
- [Stochastic Modified Flows for Riemannian Stochastic Gradient Descent](https://mlsanity.com/paper/2402.03467/) (2024)

## Abstract (arXiv)

Stochastic Gradient Descent (SGD) and its variants are almost universally used to train neural networks and to fit a variety of other parametric models. An important hyperparameter in this context is the batch size, which determines how many samples are processed before an update of the parameters occurs. Previous studies have demonstrated the benefits of using variable batch sizes. In this work, we will theoretically derive optimal batch size schedules for SGD and similar algorithms, up to an error that is quadratic in the learning rate. To achieve this, we approximate the discrete process of parameter updates using a family of stochastic differential equations indexed by the learning rate. To better handle the state-dependent diffusion coefficient, we further expand the solution of this family into a series with respect to the learning rate. Using this setup, we derive a continuous-time optimal batch size schedule for a large family of diffusion coefficients and then apply the results in the setting of linear regression.