Accelerated Linear Convergence of Stochastic Momentum Methods in Wasserstein Distances

Abstract

Momentum methods such as Polyak's heavy ball (HB) method, Nesterov's accelerated gradient (AG) as well as accelerated projected gradient (APG) method have been commonly used in machine learning practice, but their performance is quite sensitive to noise in the gradients. We study these methods under a first-order stochastic oracle model where noisy estimates of the gradients are available. For strongly convex problems, we show that the distribution of the iterates of AG converges with the accelerated $O(\sqrt{\kappa }\log (1/\epsilon ))$ linear rate to a ball of radius $\epsilon$ centered at a unique invariant distribution in the 1-Wasserstein metric where $\kappa$ is the condition number as long as the noise variance is smaller than an explicit upper bound we can provide. Our analysis also certifies linear convergence rates as a function of the stepsize, momentum parameter and the noise variance; recovering the accelerated rates in the noiseless case and quantifying the level of noise that can be tolerated to achieve a given performance. In the special case of strongly convex quadratic objectives, we can show accelerated linear rates in the $p$-Wasserstein metric for any $p\ge 1$ with improved sensitivity to noise for both AG and HB through a non-asymptotic analysis under some additional assumptions on the noise structure. Our analysis for HB and AG also leads to improved non-asymptotic convergence bounds in suboptimality for both deterministic and stochastic settings which is of independent interest. To the best of our knowledge, these are the first linear convergence results for stochastic momentum methods under the stochastic oracle model. We also extend our results to the APG method and weakly convex functions showing accelerated rates when the noise magnitude is sufficiently small.

72 pages

References in corpus (15)

• A Differential Equation for Modeling Nesterov's Accelerated Gradient Method: Theory and Insights
• Adding Gradient Noise Improves Learning for Very Deep Networks
• Non-convex learning via Stochastic Gradient Langevin Dynamics: a nonasymptotic analysis
• A Lyapunov Analysis of Momentum Methods in Optimization
• Unified Convergence Analysis of Stochastic Momentum Methods for Convex and Non-convex Optimization
• From Averaging to Acceleration, There is Only a Step-size
• Accelerating Stochastic Gradient Descent For Least Squares Regression
• Global Convergence of Stochastic Gradient Hamiltonian Monte Carlo for Non-Convex Stochastic Optimization: Non-Asymptotic Performance Bounds and Momentum-Based Acceleration
• A Tail-Index Analysis of Stochastic Gradient Noise in Deep Neural Networks
• Linearly convergent stochastic heavy ball method for minimizing generalization error
• A Universally Optimal Multistage Accelerated Stochastic Gradient Method
• Dissipativity Theory for Nesterov's Accelerated Method
• Breaking Reversibility Accelerates Langevin Dynamics for Global Non-Convex Optimization
• A single potential governing convergence of conjugate gradient, accelerated gradient and geometric descent
• Global convergence of the Heavy-ball method for convex optimization

Cited by in corpus (9)

• SlowMo: Improving Communication-Efficient Distributed SGD with Slow Momentum
• On the Convergence of Nesterov's Accelerated Gradient Method in Stochastic Settings
• Escaping Saddle Points Faster with Stochastic Momentum
• Bias-Variance Tradeoff in a Sliding Window Implementation of the Stochastic Gradient Algorithm
• Quickly Finding a Benign Region via Heavy Ball Momentum in Non-Convex Optimization
• Approximate Heavy Tails in Offline (Multi-Pass) Stochastic Gradient Descent
• A Modular Analysis of Provable Acceleration via Polyak's Momentum: Training a Wide ReLU Network and a Deep Linear Network
• Understanding Modern Techniques in Optimization: Frank-Wolfe, Nesterov's Momentum, and Polyak's Momentum
• Dynamics of Stochastic Momentum Methods on Large-scale, Quadratic Models