Abstract: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 $O\left(\exp(-\frac{T}{\sqrt{\kappa}}) + \sigma \right)$ convergence, where $T$ is the number of iterations and $\sigma^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 $O\left(\exp\left(-\frac{T}{\sqrt{\kappa}}\right) + \frac{\sigma}{T}\right)$ rate. 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(-\frac{T}{\kappa}) + \frac{\sigma^2}{T})$ rate. We empirically demonstrate the effectiveness of the proposed algorithms.

Accelerated Over-Relaxation Heavy-Ball Methods with Provable Acceleration and Global Convergence

Global convergence and asymptotic optimality of the heavy ball method for a class of non-convex optimization problems

Accelerated Convergence of Stochastic Heavy Ball Method under Anisotropic Gradient Noise

Heavy Ball Momentum for Non-Strongly Convex Optimization

Exploring the Inefficiency of Heavy Ball As Momentum Parameter Approaches 1

On the fast convergence of minibatch heavy ball momentum

Normalized Stochastic Heavy Ball with Adaptive Momentum1

On the Convergence Analysis of Aggregated Heavy-Ball Method

An accelerated lyapunov function for Polyak's Heavy-ball on convex quadratics

A Globally Convergent Gradient Method with Momentum

Normalized Stochastic Heavy Ball with Adaptive Momentum<span Ref-Type="fn" Rid="faia230568_fn001" Style="display:None"> <sup>1</sup> &Lt;/span>

Non-Ergodic Convergence Analysis of Heavy-Ball Algorithms.

Stochastic Momentum Method with Double Acceleration for Regularized Empirical Risk Minimization

Unified Convergence Analysis of Stochastic Momentum Methods for Convex and Non-convex Optimization

Provable Acceleration of Nesterov's Accelerated Gradient Method over Heavy Ball Method in Training Over-Parameterized Neural Networks

Accelerated gradient methods for nonconvex optimization: Escape trajectories from strict saddle points and convergence to local minima

On the acceleration of the Barzilai-Borwein method

Averaged Heavy-Ball Method

(Accelerated) Noise-adaptive Stochastic Heavy-Ball Momentum

Acceleration with a Ball Optimization Oracle

On the Powerball Method for Optimization