Stability of Stochastic Gradient Method with Momentum for Strongly Convex Loss Functions

Ali Ramezani-Kebrya, Ashish Khisti, and Ben Liang

Abstract: While momentum-based methods, in conjunction with the stochastic gradient descent, are widely used when training machine learning models, there is little theoretical understanding on the generalization error of such methods. In practice, the momentum parameter is often chosen in a heuristic fashion with little theoretical guidance. In this work, we use the framework of algorithmic stability to provide an upper-bound on the generalization error for the class of strongly convex loss functions, under mild technical assumptions. Our bound decays to zero inversely with the size of the training set, and increases as the momentum parameter is increased. We also develop an upper-bound on the expected true risk, in terms of the number of training steps, the size of the training set, and the momentum parameter.

Keywords: Generalization Error, Stochastic Gradient Descent, Uniform Stability

TL;DR: Stochastic gradient method with momentum generalizes.