Accelerated Convergence of Stochastic Heavy Ball Method under Anisotropic Gradient Noise
This provides rigorous theoretical justification for using heavy-ball momentum in large-batch settings like distributed or federated learning to reduce communication rounds, though it is incremental as it builds on existing momentum methods.
The paper tackles the lack of theoretical understanding of heavy-ball momentum in SGD under anisotropic gradient noise for quadratic regression, showing it achieves $ ilde{\mathcal{O}}(\sqrt{\kappa})$ accelerated convergence for bias and near-optimal variance rates, implying overall convergence close to the statistical minimax rate.
Heavy-ball momentum with decaying learning rates is widely used with SGD for optimizing deep learning models. In contrast to its empirical popularity, the understanding of its theoretical property is still quite limited, especially under the standard anisotropic gradient noise condition for quadratic regression problems. Although it is widely conjectured that heavy-ball momentum method can provide accelerated convergence and should work well in large batch settings, there is no rigorous theoretical analysis. In this paper, we fill this theoretical gap by establishing a non-asymptotic convergence bound for stochastic heavy-ball methods with step decay scheduler on quadratic objectives, under the anisotropic gradient noise condition. As a direct implication, we show that heavy-ball momentum can provide $\tilde{\mathcal{O}}(\sqrtκ)$ accelerated convergence of the bias term of SGD while still achieving near-optimal convergence rate with respect to the stochastic variance term. The combined effect implies an overall convergence rate within log factors from the statistical minimax rate. This means SGD with heavy-ball momentum is useful in the large-batch settings such as distributed machine learning or federated learning, where a smaller number of iterations can significantly reduce the number of communication rounds, leading to acceleration in practice.