Stochastic Gradient Descent with Momentum is Algorithmically Stable
It resolves the conjecture that momentum may degrade generalization by proving that SGDM generalizes well, offering theoretical guarantees for a widely used optimization algorithm in machine learning.
The paper provides a generalization analysis of stochastic gradient descent with momentum (SGDM) through algorithmic stability, establishing tight on-average model stability bounds for smooth convex problems without requiring Lipschitz loss functions, and deriving optimal excess population risk bounds for both Polyak's and Nesterov's momentum.
Stochastic gradient descent with momentum (SGDM) is one of the most widely used optimization algorithms in machine learning. While optimization properties of SGDM have been extensively studied in the literature, it remains insufficiently understood whether and when SGDM can generalize well to unseen data. In particular, it has been conjectured that while momentum accelerates training, it may degrade generalization. In this paper, we close this gap by developing a comprehensive generalization analysis of SGDM through the lens of algorithmic stability. More specifically, we introduce a generalized SGDM framework that encompasses both Polyak's and Nesterov's momentum schemes, and establish tight on-average model stability bounds for smooth and convex problems. Notably, the obtained bounds exploit small optimization error bounds along the trajectory, apply to any momentum parameter in the interval $[0, 1)$, and do not require the commonly assumed Lipschitzness of loss functions. We further derive optimization error bounds for the generalized SGDM, and combine them with our generalization analyses to obtain optimal excess population risk bounds for SGDM with both Polyak's and Nesterov's momentum.