(Accelerated) Noise-adaptive Stochastic Heavy-Ball Momentum

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.