Exponential convergence rates for momentum stochastic gradient descent in the overparametrized setting
This work provides theoretical guarantees for MSGD convergence, which is important for practitioners in machine learning dealing with non-convex optimization problems, though it is incremental as it builds on existing momentum and SGD analyses.
The paper proves explicit exponential convergence bounds for momentum stochastic gradient descent (MSGD) in non-convex optimization, showing faster convergence than plain SGD in scenarios like small step-sizes, flat minima, or high noise, with results applicable to overparametrized settings and functions satisfying a local Polyak-Lojasiewicz inequality.
We prove explicit bounds on the exponential rate of convergence for the momentum stochastic gradient descent scheme (MSGD) for arbitrary, fixed hyperparameters (learning rate, friction parameter) and its continuous-in-time counterpart in the context of non-convex optimization. In the small step-size regime and in the case of flat minima or large noise intensities, these bounds prove faster convergence of MSGD compared to plain stochastic gradient descent (SGD). The results are shown for objective functions satisfying a local Polyak-Lojasiewicz inequality and under assumptions on the variance of MSGD that are satisfied in overparametrized settings. Moreover, we analyze the optimal choice of the friction parameter and show that the MSGD process almost surely converges to a local minimum.