# On the Generalization of Stochastic Gradient Descent with Momentum

arXiv: https://arxiv.org/abs/2102.13653 · HTML: https://mlsanity.com/paper/2102.13653/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 17 of its 38 references that are in the map.
- First posted on arXiv 26 Feb 2021.

## Builds on (most cited first, 17 of 17)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Train faster, generalize better: Stability of stochastic gradient descent](https://mlsanity.com/paper/1509.01240/) (2015)
- [A Differential Equation for Modeling Nesterov's Accelerated Gradient Method: Theory and Insights](https://mlsanity.com/paper/1503.01243/) (2015)
- [Stochastic Gradient Descent, Weighted Sampling, and the Randomized Kaczmarz algorithm](https://mlsanity.com/paper/1310.5715/) (2013)
- [Improving Generalization Performance by Switching from Adam to SGD](https://mlsanity.com/paper/1712.07628/) (2017)
- [Analysis and Design of Optimization Algorithms via Integral Quadratic Constraints](https://mlsanity.com/paper/1408.3595/) (2014)
- [Global convergence of the Heavy-ball method for convex optimization](https://mlsanity.com/paper/1412.7457/) (2014)
- [Stability and Convergence Trade-off of Iterative Optimization Algorithms](https://mlsanity.com/paper/1804.01619/) (2018)
- [Stochastic Heavy Ball](https://mlsanity.com/paper/1609.04228/) (2016)
- [Unified Convergence Analysis of Stochastic Momentum Methods for Convex and Non-convex Optimization](https://mlsanity.com/paper/1604.03257/) (2016)
- [iPiano: Inertial Proximal Algorithm for Non-Convex Optimization](https://mlsanity.com/paper/1404.4805/) (2014)
- [A Lyapunov Analysis of Momentum Methods in Optimization](https://mlsanity.com/paper/1611.02635/) (2016)
- [On the Convergence of Nesterov's Accelerated Gradient Method in Stochastic Settings](https://mlsanity.com/paper/2002.12414/) (2020)
- [Accelerated Linear Convergence of Stochastic Momentum Methods in Wasserstein Distances](https://mlsanity.com/paper/1901.07445/) (2019)
- [Nesterov's method with decreasing learning rate leads to accelerated stochastic gradient descent](https://mlsanity.com/paper/1908.07861/) (2019)
- [The Role of Memory in Stochastic Optimization](https://mlsanity.com/paper/1907.01678/) (2019)

## Abstract (arXiv)

While momentum-based methods, in conjunction with stochastic gradient descent (SGD), are widely used when training machine learning models, there is little theoretical understanding on the generalization error of such methods. In this work, we first show that there exists a convex loss function for which algorithmic stability fails to establish generalization guarantees when SGD with standard heavy-ball momentum (SGDM) is run for multiple epochs. Then, for smooth Lipschitz loss functions, we analyze a modified momentum-based update rule, i.e., SGD with early momentum (SGDEM), and show that it admits an upper-bound on the generalization error. Thus, our results show that machine learning models can be trained for multiple epochs of SGDEM with a guarantee for generalization. Finally, for the special case of strongly convex loss functions, we find a range of momentum such that multiple epochs of standard SGDM, as a special form of SGDEM, also generalizes. Extending our results on generalization, 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. Experimental evaluations verify the consistency between the numerical results and our theoretical bounds and the effectiveness of SGDEM for smooth Lipschitz loss functions.