# (Accelerated) Noise-adaptive Stochastic Heavy-Ball Momentum

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 33 references that are in the map.
- First posted on arXiv 12 Jan 2024.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [SGD: General Analysis and Improved Rates](https://mlsanity.com/paper/1901.09401/) (2019)
- [Just Interpolate: Kernel "Ridgeless" Regression Can Generalize](https://mlsanity.com/paper/1808.00387/) (2018)
- [The Power of Interpolation: Understanding the Effectiveness of SGD in Modern Over-parametrized Learning](https://mlsanity.com/paper/1712.06559/) (2017)
- [Fast and Faster Convergence of SGD for Over-Parameterized Models and an Accelerated Perceptron](https://mlsanity.com/paper/1810.07288/) (2018)
- [Stochastic Polyak Step-size for SGD: An Adaptive Learning Rate for Fast Convergence](https://mlsanity.com/paper/2002.10542/) (2020)
- [An Improved Analysis of Stochastic Gradient Descent with Momentum](https://mlsanity.com/paper/2007.07989/) (2020)
- [Does data interpolation contradict statistical optimality?](https://mlsanity.com/paper/1806.09471/) (2018)
- [Global convergence of the Heavy-ball method for convex optimization](https://mlsanity.com/paper/1412.7457/) (2014)
- [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)
- [A Second look at Exponential and Cosine Step Sizes: Simplicity, Adaptivity, and Performance](https://mlsanity.com/paper/2002.05273/) (2020)
- [Momentum via Primal Averaging: Theoretical Insights and Learning Rate Schedules for Non-Convex Optimization](https://mlsanity.com/paper/2010.00406/) (2020)
- [A Modular Analysis of Provable Acceleration via Polyak's Momentum: Training a Wide ReLU Network and a Deep Linear Network](https://mlsanity.com/paper/2010.01618/) (2020)
- [Trajectory of Mini-Batch Momentum: Batch Size Saturation and Convergence in High Dimensions](https://mlsanity.com/paper/2206.01029/) (2022)
- [Tight Dimension Independent Lower Bound on the Expected Convergence Rate for Diminishing Step Sizes in SGD](https://mlsanity.com/paper/1810.04723/) (2018)
- [The Marginal Value of Momentum for Small Learning Rate SGD](https://mlsanity.com/paper/2307.15196/) (2023)
- [A Universally Optimal Multistage Accelerated Stochastic Gradient Method](https://mlsanity.com/paper/1901.08022/) (2019)
- [Dynamics of Stochastic Momentum Methods on Large-scale, Quadratic Models](https://mlsanity.com/paper/2106.03696/) (2021)
- [Towards Noise-adaptive, Problem-adaptive (Accelerated) Stochastic Gradient Descent](https://mlsanity.com/paper/2110.11442/) (2021)
- [On the fast convergence of minibatch heavy ball momentum](https://mlsanity.com/paper/2206.07553/) (2022)
- [Provable non-accelerations of the heavy-ball method](https://mlsanity.com/paper/2307.11291/) (2023)
- [Provable Acceleration of Heavy Ball beyond Quadratics for a Class of Polyak-Lojasiewicz Functions when the Non-Convexity is Averaged-Out](https://mlsanity.com/paper/2206.11872/) (2022)
- [Accelerated Convergence of Stochastic Heavy Ball Method under Anisotropic Gradient Noise](https://mlsanity.com/paper/2312.14567/) (2023)
- [Almost sure convergence rates for Stochastic Gradient Descent and Stochastic Heavy Ball](https://mlsanity.com/paper/2006.07867/) (2020)

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

- [A Nonlinear Time-Series Theory for Constant-step SGD with Momentum](https://mlsanity.com/paper/2610.16034/) (2026)

## Abstract (arXiv)

Stochastic heavy ball momentum (SHB) is commonly used to train machine learning models, and often provides empirical improvements over stochastic gradient descent. By primarily focusing on strongly-convex quadratics, we aim to better understand the theoretical advantage of SHB and subsequently improve the method. For strongly-convex quadratics, Kidambi et al. (2018) show that SHB (with a mini-batch of size 1) cannot attain accelerated convergence, and hence has no theoretical benefit over SGD. They conjecture that the practical gain of SHB is a by-product of using larger mini-batches. We first substantiate this claim by showing that SHB can attain an accelerated rate when the mini-batch size is larger than a threshold b^* that depends on the condition number kappa. Specifically, we prove that with the same step-size and momentum parameters as in the deterministic setting, SHB with a sufficiently large mini-batch size results in an Oleft(exp(-fracT√kappa) + σ right) convergence when measuring the distance to the optimal solution in the ℓ2 norm, where T is the number of iterations and σ2 is the variance in the stochastic gradients. We prove a lower-bound which demonstrates that a kappa dependence in b^* is necessary. To ensure convergence to the minimizer, we design a noise-adaptive multi-stage algorithm that results in an Oleft(expleft(-fracT√kapparight) + fracσ√Tright) rate when measuring the distance to the optimal solution in the ℓ2 norm. We also consider the general smooth, strongly-convex setting and propose the first noise-adaptive SHB variant that converges to the minimizer at an O(exp(-fracTkappa) + fracσ2T) rate when measuring the distance to the optimal solution in the squared ℓ2 norm. We empirically demonstrate the effectiveness of the proposed algorithms.