# Online Learning and Information Exponents: On The Importance of Batch size, and Time/Complexity Tradeoffs

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 43 references that are in the map.
- First posted on arXiv 4 Jun 2024.

## Builds on (most cited first, 30 of 33)

- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [A Mean Field View of the Landscape of Two-Layers Neural Networks](https://mlsanity.com/paper/1804.06561/) (2018)
- [On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport](https://mlsanity.com/paper/1805.09545/) (2018)
- [Neural Networks can Learn Representations with Gradient Descent](https://mlsanity.com/paper/2206.15144/) (2022)
- [High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation](https://mlsanity.com/paper/2205.01445/) (2022)
- [SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics](https://mlsanity.com/paper/2302.11055/) (2023)
- [The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks](https://mlsanity.com/paper/2202.08658/) (2022)
- [Online stochastic gradient descent on non-convex losses from high-dimensional inference](https://mlsanity.com/paper/2003.10409/) (2020)
- [Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup](https://mlsanity.com/paper/1906.08632/) (2019)
- [Mean Field Analysis of Neural Networks: A Central Limit Theorem](https://mlsanity.com/paper/1808.09372/) (2018)
- [Trainability and Accuracy of Neural Networks: An Interacting Particle System Approach](https://mlsanity.com/paper/1805.00915/) (2018)
- [Saddle-to-Saddle Dynamics in Deep Linear Networks: Small Initialization Training, Symmetry, and Sparsity](https://mlsanity.com/paper/2106.15933/) (2021)
- [High-dimensional limit theorems for SGD: Effective dynamics and critical scaling](https://mlsanity.com/paper/2206.04030/) (2022)
- [How Two-Layer Neural Networks Learn, One (Giant) Step at a Time](https://mlsanity.com/paper/2305.18270/) (2023)
- [The staircase property: How hierarchical structure can guide deep learning](https://mlsanity.com/paper/2108.10573/) (2021)
- [Smoothing the Landscape Boosts the Signal for SGD: Optimal Sample Complexity for Learning Single Index Models](https://mlsanity.com/paper/2305.10633/) (2023)
- [Learning time-scales in two-layers neural networks](https://mlsanity.com/paper/2303.00055/) (2023)
- [On Learning Gaussian Multi-index Models with Gradient Flow](https://mlsanity.com/paper/2310.19793/) (2023)
- [Phase diagram of Stochastic Gradient Descent in high-dimensional two-layer neural networks](https://mlsanity.com/paper/2202.00293/) (2022)
- [A Theory of Non-Linear Feature Learning with One Gradient Step in Two-Layer Neural Networks](https://mlsanity.com/paper/2310.07891/) (2023)
- [From high-dimensional & mean-field dynamics to dimensionless ODEs: A unifying approach to SGD in two-layers networks](https://mlsanity.com/paper/2302.05882/) (2023)
- [Online Stochastic Gradient Descent with Arbitrary Initialization Solves Non-smooth, Non-convex Phase Retrieval](https://mlsanity.com/paper/1910.12837/) (2019)
- [Gradient-Based Feature Learning under Structured Data](https://mlsanity.com/paper/2309.03843/) (2023)
- [Homogenization of SGD in high-dimensions: Exact dynamics and generalization properties](https://mlsanity.com/paper/2205.07069/) (2022)
- [Hitting the High-Dimensional Notes: An ODE for SGD learning dynamics on GLMs and multi-index models](https://mlsanity.com/paper/2308.08977/) (2023)
- [Global Convergence of EM Algorithm for Mixtures of Two Component Linear Regression](https://mlsanity.com/paper/1810.05752/) (2018)
- [On the Impact of Overparameterization on the Training of a Shallow Neural Network in High Dimensions](https://mlsanity.com/paper/2311.03794/) (2023)
- [Fast matrix completion without the condition number](https://mlsanity.com/paper/1407.4070/) (2014)
- [Fast Exact Matrix Completion with Finite Samples](https://mlsanity.com/paper/1411.1087/) (2014)
- [On the different regimes of Stochastic Gradient Descent](https://mlsanity.com/paper/2309.10688/) (2023)

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

- [From Information to Generative Exponent: Learning Rate Induces Phase Transitions in SGD](https://mlsanity.com/paper/2510.21020/) (2025)
- [Specialization of softmax attention heads: insights from the high-dimensional single-location model](https://mlsanity.com/paper/2603.03993/) (2026)
- [Average Gradient Outer Product in kernel regression provably recovers the central subspace for multi-index models](https://mlsanity.com/paper/2605.15082/) (2026)

## Abstract (arXiv)

We study the impact of the batch size nb on the iteration time T of training two-layer neural networks with one-pass stochastic gradient descent (SGD) on multi-index target functions of isotropic covariates. We characterize the optimal batch size minimizing the iteration time as a function of the hardness of the target, as characterized by the information exponents. We show that performing gradient updates with large batches nb lesssim d^fracℓ2 minimizes the training time without changing the total sample complexity, where ℓ is the information exponent of the target to be learned citeparous2021online and d is the input dimension. However, larger batch sizes than nb gg d^fracℓ2 are detrimental for improving the time complexity of SGD. We provably overcome this fundamental limitation via a different training protocol, textitCorrelation loss SGD, which suppresses the auto-correlation terms in the loss function. We show that one can track the training progress by a system of low-dimensional ordinary differential equations (ODEs). Finally, we validate our theoretical results with numerical experiments.
