# Perfect Parallelization in Mini-Batch SGD with Classical Momentum Acceleration

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 39 of its 56 references that are in the map.
- First posted on arXiv 18 May 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [MobileNetV2: Inverted Residuals and Linear Bottlenecks](https://mlsanity.com/paper/1801.04381/) (2018)
- [OpenML: networked science in machine learning](https://mlsanity.com/paper/1407.7722/) (2014)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [SOAP: Improving and Stabilizing Shampoo using Adam](https://mlsanity.com/paper/2409.11321/) (2024)
- [The Power of Interpolation: Understanding the Effectiveness of SGD in Modern Over-parametrized Learning](https://mlsanity.com/paper/1712.06559/) (2017)
- [Katyusha: The First Direct Acceleration of Stochastic Gradient Methods](https://mlsanity.com/paper/1603.05953/) (2016)
- [Overfitting or perfect fitting? Risk bounds for classification and regression rules that interpolate](https://mlsanity.com/paper/1806.05161/) (2018)
- [Accelerated Gradient Descent Escapes Saddle Points Faster than Gradient Descent](https://mlsanity.com/paper/1711.10456/) (2017)
- [A Unified Analysis of Stochastic Momentum Methods for Deep Learning](https://mlsanity.com/paper/1808.10396/) (2018)
- [Momentum and Stochastic Momentum for Stochastic Gradient, Newton, Proximal Point and Subspace Descent Methods](https://mlsanity.com/paper/1712.09677/) (2017)
- [On the insufficiency of existing momentum schemes for Stochastic Optimization](https://mlsanity.com/paper/1803.05591/) (2018)
- [Stochastic Heavy Ball](https://mlsanity.com/paper/1609.04228/) (2016)
- [On the Computational Inefficiency of Large Batch Sizes for Stochastic Gradient Descent](https://mlsanity.com/paper/1811.12941/) (2018)
- [From Averaging to Acceleration, There is Only a Step-size](https://mlsanity.com/paper/1504.01577/) (2015)
- [A geometric alternative to Nesterov's accelerated gradient descent](https://mlsanity.com/paper/1506.08187/) (2015)
- [Trajectory of Mini-Batch Momentum: Batch Size Saturation and Convergence in High Dimensions](https://mlsanity.com/paper/2206.01029/) (2022)
- [Sharp Analysis of Sketch-and-Project Methods via a Connection to Randomized Singular Value Decomposition](https://mlsanity.com/paper/2208.09585/) (2022)
- [When and Why Momentum Accelerates SGD:An Empirical Study](https://mlsanity.com/paper/2306.09000/) (2023)
- [The Marginal Value of Momentum for Small Learning Rate SGD](https://mlsanity.com/paper/2307.15196/) (2023)
- [Accelerated Linear Convergence of Stochastic Momentum Methods in Wasserstein Distances](https://mlsanity.com/paper/1901.07445/) (2019)
- [Linearly convergent stochastic heavy ball method for minimizing generalization error](https://mlsanity.com/paper/1710.10737/) (2017)
- [Accelerating Stochastic Gradient Descent For Least Squares Regression](https://mlsanity.com/paper/1704.08227/) (2017)
- [On the fast convergence of minibatch heavy ball momentum](https://mlsanity.com/paper/2206.07553/) (2022)
- [An Accelerated Randomized Kaczmarz Algorithm](https://mlsanity.com/paper/1310.2887/) (2013)
- [Solving Dense Linear Systems Faster Than via Preconditioning](https://mlsanity.com/paper/2312.08893/) (2023)
- [Fast Last-Iterate Convergence of SGD in the Smooth Interpolation Regime](https://mlsanity.com/paper/2507.11274/) (2025)
- [Accelerated Stochastic Matrix Inversion: General Theory and Speeding up BFGS Rules for Faster Second-Order Optimization](https://mlsanity.com/paper/1802.04079/) (2018)

## Abstract (arXiv)

Accelerating stochastic gradient methods with classical momentum schemes, such as Polyak's heavy ball, has proven highly successful in training large-scale machine learning models, particularly when combined with the hardware acceleration of large mini-batch computations. Yet, the effect of classical momentum on stochastic mini-batch optimization has been poorly understood theoretically, with prior works requiring strong noise assumptions and extremely large mini-batches. In this work, we develop a general theory of stochastic momentum acceleration for optimizing over quadratics in the interpolation regime, a popular abstraction for studying deep learning dynamics which also includes classical methods such as randomized Kaczmarz and coordinate descent. Our framework encompasses both heavy ball and Nesterov-style momentum, allows for arbitrary mini-batch sizes, and makes minimal assumptions on the stochastic noise. In particular, we show that acceleration from classical momentum is directly proportional to the gradient mini-batch size (up to a natural saturation point), thereby enabling perfect parallelization of mini-batch computations. Our theory also provides a simple choice for the momentum parameter, which is shown to be effective empirically.