Noise-adaptive (Accelerated) Stochastic Heavy-Ball Momentum
Anh Dang Reza Babanezhad Sharan Vaswani
We analyze the convergence of stochastic heavy ball (SHB) momentum in the
smooth, strongly-convex setting. Kidambi et al. (2018) show that SHB (with
small mini-batches) cannot attain an accelerated rate of convergence even for
quadratics, and conjecture that the practical gain of SHB is a by-product of
mini-batching. We substantiate this claim by showing that SHB can obtain an
accelerated rate when the mini-batch size is larger than some threshold. In
particular, for strongly-convex quadratics with condition number $\kappa$, we
prove that SHB with the standard step-size and momentum parameters results in
an $O\left(\exp(-\frac{T}{\sqrt{\kappa}}) + \sigma \right)$ convergence rate,
where $T$ is the number of iterations and $\sigma^2$ is the variance in the
stochastic gradients. To ensure convergence to the minimizer, we propose a
multi-stage approach that results in a noise-adaptive
$O\left(\exp\left(-\frac{T}{\sqrt{\kappa}} \right) + \frac{\sigma}{T}\right)$
rate. For general strongly-convex functions, we use the averaging
interpretation of SHB along with exponential step-sizes to prove an
$O\left(\exp\left(-\frac{T}{\kappa} \right) + \frac{\sigma^2}{T} \right)$
convergence to the minimizer in a noise-adaptive manner. Finally, we
empirically demonstrate the effectiveness of the proposed algorithms.