Perfect Parallelization in Mini-Batch SGD with Classical Momentum Acceleration

• URL: http://arxiv.org/abs/2605.18609v1
• Date: Mon, 18 May 2026 16:18:59 GMT
• Title: Perfect Parallelization in Mini-Batch SGD with Classical Momentum Acceleration
• Abstract summary: We develop a theory of momentum acceleration for optimizing over a quadratic saturation regime.<n>Our framework encompasses both heavy ball and Nesterov-style momentum, allowing for arbitrary mini-batch sizes.<n>Our theory also provides simple choice for the momentum, which is shown to be effective empirically.
• Score: 1.4323566945483497
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• From Atoms to Entropy: Optimal Noise Allocation for Diffusion Training in the Convex Regime [52.459513910242485]
  We develop a general statistical framework for studying optimal noise-level allocation in diffusion training.<n>We show that the optimized training schedule admits an atomic minimizer, concentrated on finitely many noise levels.<n>We then evaluate the entropic schedule in larger-scale experiments, where full schedule optimization is currently intractable.
  arXiv  Detail & Related papers  (2026-07-10T20:59:06Z)
• Deriving Hyperparameter Scaling Laws via Modern Optimization Theory [55.63126290312615]
  We study modern first-order boundss for methods based on the Linear Minimization Oracle (LMO)<n>Treating bounds in recent literature as a proxy and minimizing them across different tuning regimes yields closed-form power-law schedules for learning rate, momentum, and batch size.<n>Our results draw particular attention to the interaction between momentum and batch-size scaling, suggesting that optimal performance may be achieved with several scaling strategies.
  arXiv  Detail & Related papers  (2026-03-16T22:21:27Z)
• From Black Hole to Galaxy: Neural Operator: Framework for Accretion and Feedback Dynamics [70.27068115318681]
  We introduce a neural-based ''subgrid black hole'' that learns the small-scale local dynamics and embeds it within direct simulations.<n>Thanks to the great speedup in fine-scale evolution, our approach captures intrinsic variability in accretion-driven feedback, allowing dynamic coupling between the central black hole and galaxy-scale gas.
  arXiv  Detail & Related papers  (2025-12-01T11:47:49Z)
• Accelerated Convergence of Stochastic Heavy Ball Method under Anisotropic Gradient Noise [16.12834917344859]
  It is widely conjectured that heavy-ball momentum method can provide accelerated convergence and should work well in large batch settings. We show that heavy-ball momentum can provide $tildemathcalO(sqrtkappa)$ accelerated convergence of the bias term of SGD while still achieving near-optimal convergence rate. This means SGD with heavy-ball momentum is useful in the large-batch settings such as distributed machine learning or federated learning.
  arXiv  Detail & Related papers  (2023-12-22T09:58:39Z)
• The Marginal Value of Momentum for Small Learning Rate SGD [20.606430391298815]
  Momentum is known to accelerate the convergence of gradient descent in strongly convex settings without gradient noise regimes. Experiments show that momentum indeed has limited benefits for both optimization and generalization in practical training where the optimal learning rate is not very large.
  arXiv  Detail & Related papers  (2023-07-27T21:01:26Z)
• Losing momentum in continuous-time stochastic optimisation [42.617042045455506]
  momentum-based optimisation algorithms have become particularly widespread. In this work, we analyse a continuous-time model for gradient descent with momentum. We also train a convolutional neural network in an image classification problem.
  arXiv  Detail & Related papers  (2022-09-08T10:46:05Z)
• On the fast convergence of minibatch heavy ball momentum [5.4755933832880865]
  We show that heavy ball momentum retains the fast linear rate of (deterministic) heavy ball momentum on optimization problems.<n>The algorithm we study can be interpreted as an accelerated randomized Kaczmarz algorithm with minibatching and heavy ball momentum.
  arXiv  Detail & Related papers  (2022-06-15T14:12:45Z)
• Accelerate Distributed Stochastic Descent for Nonconvex Optimization with Momentum [12.324457683544132]
  We propose a momentum method for such model averaging approaches. We analyze the convergence and scaling properties of such momentum methods. Our experimental results show that block momentum not only accelerates training, but also achieves better results.
  arXiv  Detail & Related papers  (2021-10-01T19:23:18Z)
• Quantum algorithms for quantum dynamics: A performance study on the spin-boson model [68.8204255655161]
  Quantum algorithms for quantum dynamics simulations are traditionally based on implementing a Trotter-approximation of the time-evolution operator. variational quantum algorithms have become an indispensable alternative, enabling small-scale simulations on present-day hardware. We show that, despite providing a clear reduction of quantum gate cost, the variational method in its current implementation is unlikely to lead to a quantum advantage.
  arXiv  Detail & Related papers  (2021-08-09T18:00:05Z)
• Positive-Negative Momentum: Manipulating Stochastic Gradient Noise to Improve Generalization [89.7882166459412]
  gradient noise (SGN) acts as implicit regularization for deep learning. Some works attempted to artificially simulate SGN by injecting random noise to improve deep learning. For simulating SGN at low computational costs and without changing the learning rate or batch size, we propose the Positive-Negative Momentum (PNM) approach.
  arXiv  Detail & Related papers  (2021-03-31T16:08:06Z)
• Fast and differentiable simulation of driven quantum systems [58.720142291102135]
  We introduce a semi-analytic method based on the Dyson expansion that allows us to time-evolve driven quantum systems much faster than standard numerical methods. We show results of the optimization of a two-qubit gate using transmon qubits in the circuit QED architecture.
  arXiv  Detail & Related papers  (2020-12-16T21:43:38Z)
• On the Generalization Benefit of Noise in Stochastic Gradient Descent [34.127525925676416]
  It has long been argued that minibatch gradient descent can generalize better than large batch gradient descent in deep neural networks. We show that small or moderately large batch sizes can substantially outperform very large batches on the test set.
  arXiv  Detail & Related papers  (2020-06-26T16:18:54Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.