# Towards understanding how momentum improves generalization in deep learning

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

## Key facts

- Cited by 40 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 51 references that are in the map.
- First posted on arXiv 13 Jul 2022.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Wide Residual Networks](https://mlsanity.com/paper/1605.07146/) (2016)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [The Implicit Bias of Gradient Descent on Separable Data](https://mlsanity.com/paper/1710.10345/) (2017)
- [Implicit Regularization in Deep Matrix Factorization](https://mlsanity.com/paper/1905.13655/) (2019)
- [Implicit Regularization in Matrix Factorization](https://mlsanity.com/paper/1705.09280/) (2017)
- [Don't Decay the Learning Rate, Increase the Batch Size](https://mlsanity.com/paper/1711.00489/) (2017)
- [Gradient Descent Maximizes the Margin of Homogeneous Neural Networks](https://mlsanity.com/paper/1906.05890/) (2019)
- [Towards Understanding Ensemble, Knowledge Distillation and Self-Distillation in Deep Learning](https://mlsanity.com/paper/2012.09816/) (2020)
- [Train longer, generalize better: closing the generalization gap in large batch training of neural networks](https://mlsanity.com/paper/1705.08741/) (2017)
- [Implicit Bias of Gradient Descent for Wide Two-layer Neural Networks Trained with the Logistic Loss](https://mlsanity.com/paper/2002.04486/) (2020)
- [The Marginal Value of Adaptive Gradient Methods in Machine Learning](https://mlsanity.com/paper/1705.08292/) (2017)
- [Towards Explaining the Regularization Effect of Initial Large Learning Rate in Training Neural Networks](https://mlsanity.com/paper/1907.04595/) (2019)
- [Momentum Improves Normalized SGD](https://mlsanity.com/paper/2002.03305/) (2020)
- [Path-SGD: Path-Normalized Optimization in Deep Neural Networks](https://mlsanity.com/paper/1506.02617/) (2015)
- [An Improved Analysis of Stochastic Gradient Descent with Momentum](https://mlsanity.com/paper/2007.07989/) (2020)
- [Analysis and Design of Optimization Algorithms via Integral Quadratic Constraints](https://mlsanity.com/paper/1408.3595/) (2014)
- [Shape Matters: Understanding the Implicit Bias of the Noise Covariance](https://mlsanity.com/paper/2006.08680/) (2020)
- [The Implicit and Explicit Regularization Effects of Dropout](https://mlsanity.com/paper/2002.12915/) (2020)
- [Theoretical Analysis of Auto Rate-Tuning by Batch Normalization](https://mlsanity.com/paper/1812.03981/) (2018)
- [Norm matters: efficient and accurate normalization schemes in deep networks](https://mlsanity.com/paper/1803.01814/) (2018)
- [On the insufficiency of existing momentum schemes for Stochastic Optimization](https://mlsanity.com/paper/1803.05591/) (2018)
- [Convergence of a Stochastic Gradient Method with Momentum for Non-Smooth Non-Convex Optimization](https://mlsanity.com/paper/2002.05466/) (2020)
- [The Two Regimes of Deep Network Training](https://mlsanity.com/paper/2002.10376/) (2020)

## Built on by (most cited first, 30 of 40)

- [Flora: Low-Rank Adapters Are Secretly Gradient Compressors](https://mlsanity.com/paper/2402.03293/) (2024)
- [Vision Transformers provably learn spatial structure](https://mlsanity.com/paper/2210.09221/) (2022)
- [Solving Regularized Exp, Cosh and Sinh Regression Problems](https://mlsanity.com/paper/2303.15725/) (2023)
- [Robust Learning with Progressive Data Expansion Against Spurious Correlation](https://mlsanity.com/paper/2306.04949/) (2023)
- [Momentum Provably Improves Error Feedback!](https://mlsanity.com/paper/2305.15155/) (2023)
- [Why Does Sharpness-Aware Minimization Generalize Better Than SGD?](https://mlsanity.com/paper/2310.07269/) (2023)
- [Implicit regularization in Heavy-ball momentum accelerated stochastic gradient descent](https://mlsanity.com/paper/2302.00849/) (2023)
- [The Benefits of Mixup for Feature Learning](https://mlsanity.com/paper/2303.08433/) (2023)
- [The Marginal Value of Momentum for Small Learning Rate SGD](https://mlsanity.com/paper/2307.15196/) (2023)
- [Momentum Further Constrains Sharpness at the Edge of Stochastic Stability](https://mlsanity.com/paper/2604.14108/) (2026)
- [Non-convex Stochastic Composite Optimization with Polyak Momentum](https://mlsanity.com/paper/2403.02967/) (2024)
- [Leveraging Continuous Time to Understand Momentum When Training Diagonal Linear Networks](https://mlsanity.com/paper/2403.05293/) (2024)
- [DASH: Warm-Starting Neural Network Training in Stationary Settings without Loss of Plasticity](https://mlsanity.com/paper/2410.23495/) (2024)
- [Understanding the Generalization Benefits of Late Learning Rate Decay](https://mlsanity.com/paper/2401.11600/) (2024)
- [On the Clean Generalization and Robust Overfitting in Adversarial Training from Two Theoretical Views: Representation Complexity and Training Dynamics](https://mlsanity.com/paper/2306.01271/) (2023)
- [SLIM: Spuriousness Mitigation with Minimal Human Annotations](https://mlsanity.com/paper/2407.05594/) (2024)
- [Entropy-Guided Token Dropout: Training Autoregressive Language Models with Limited Domain Data](https://mlsanity.com/paper/2512.23422/) (2025)
- [Single-Timescale Multi-Sequence Stochastic Approximation Without Fixed Point Smoothness: Theories and Applications](https://mlsanity.com/paper/2410.13743/) (2024)
- [Provable Benefit of Cutout and CutMix for Feature Learning](https://mlsanity.com/paper/2410.23672/) (2024)
- [DeepKD: A Deeply Decoupled and Denoised Knowledge Distillation Trainer](https://mlsanity.com/paper/2505.15133/) (2025)
- [On the Role of Label Noise in the Feature Learning Process](https://mlsanity.com/paper/2505.18909/) (2025)
- [How Does Label Noise Gradient Descent Improve Generalization in the Low SNR Regime?](https://mlsanity.com/paper/2510.17526/) (2025)
- [Mean-field analysis for heavy ball methods: Dropout-stability, connectivity, and global convergence](https://mlsanity.com/paper/2210.06819/) (2022)
- [Gradient Descent with Polyak's Momentum Finds Flatter Minima via Large Catapults](https://mlsanity.com/paper/2311.15051/) (2023)
- [Changing the Training Data Distribution to Reduce Simplicity Bias Improves In-distribution Generalization](https://mlsanity.com/paper/2404.17768/) (2024)
- [Understanding Private Learning From Feature Perspective](https://mlsanity.com/paper/2511.18006/) (2025)
- [Hallmarks of Optimization Trajectories in Neural Networks: Directional Exploration and Redundancy](https://mlsanity.com/paper/2403.07379/) (2024)
- [From Learning to Optimize to Learning Optimization Algorithms](https://mlsanity.com/paper/2405.18222/) (2024)
- [Understanding the Benefits of SimCLR Pre-Training in Two-Layer Convolutional Neural Networks](https://mlsanity.com/paper/2409.18685/) (2024)
- [Rethinking Benign Overfitting in Two-Layer Neural Networks](https://mlsanity.com/paper/2502.11893/) (2025)

## Abstract (arXiv)

Stochastic gradient descent (SGD) with momentum is widely used for training modern deep learning architectures. While it is well-understood that using momentum can lead to faster convergence rate in various settings, it has also been observed that momentum yields higher generalization. Prior work argue that momentum stabilizes the SGD noise during training and this leads to higher generalization. In this paper, we adopt another perspective and first empirically show that gradient descent with momentum (GD+M) significantly improves generalization compared to gradient descent (GD) in some deep learning problems. From this observation, we formally study how momentum improves generalization. We devise a binary classification setting where a one-hidden layer (over-parameterized) convolutional neural network trained with GD+M provably generalizes better than the same network trained with GD, when both algorithms are similarly initialized. The key insight in our analysis is that momentum is beneficial in datasets where the examples share some feature but differ in their margin. Contrary to GD that memorizes the small margin data, GD+M still learns the feature in these data thanks to its historical gradients. Lastly, we empirically validate our theoretical findings.