Stochastic Gradient Descent with Momentum is Algorithmically Stable

Abstract: 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.

• The paper establishes the on-average model stability of SGDM, quantifying how momentum increases instability by a factor of O(1/(1-β)^(3/2)).
• It introduces a unified SGDM framework that covers both Polyak's and Nesterov's momentum variants under non-Lipschitz smooth convex conditions.
• Empirical and theoretical results confirm that with controlled step sizes, SGDM achieves optimal excess risk rates of O(1/√n), validating its generalization properties.

Algorithmic Stability and Generalization of Stochastic Gradient Descent with Momentum

Introduction and Context

Stochastic Gradient Descent with Momentum (SGDM), particularly in the forms of Polyak's (heavy-ball) and Nesterov's accelerated variants, is central to large-scale optimization in machine learning. While their convergence behavior has been extensively analyzed, a persistent open question concerns their generalization properties—specifically, whether the use of momentum fundamentally impairs algorithmic stability and thus generalization, as often conjectured. Previous attempts at analyzing the stability of momentum-based variants have yielded results under strong convexity, Lipschitzness, or highly restrictive momentum parameters, leaving practical settings largely uncharacterized.

Generalized SGDM Framework

This work introduces an encompassing SGDM algorithmic framework that unifies both Polyak’s and Nesterov’s momentum variants through flexible parameterization:

• Polyak's momentum: $\gamma = 0$, weight update $w_{t+1} = w_t - \eta m_t$.
• Nesterov's momentum: $\eta = \beta \gamma$, resulting in the canonical Nesterov update via step reorganization.

The analysis operates under nonnegative, $\alpha$-smooth, convex loss functions, without requiring Lipschitz continuity—a relaxation compared to much of the prior literature.

On-Average Model Stability Analysis

The core theoretical contribution is a comprehensive on-average model stability characterization for SGDM. The analysis demonstrates that, under suitable step size choices, both Polyak and Nesterov momentum variants retain algorithmic stability, with the following key findings:

• Stability Bounds: The on-average model stability is upper bounded by a function that grows at most linearly with the number of optimization steps and inversely with the training set size, but is adversely affected by the momentum parameter. Specifically, the bound increases by a factor of $\mathcal{O}(1/(1-\beta)^{3/2})$ as $\beta \rightarrow 1$.
• Tightness: The derived bounds align, up to this constant factor, with known optimal stability results for vanilla SGD without momentum, confirming their tightness.
• No Lipschitz Requirement: The analysis exploits the self-bounding property of smooth functions, circumventing the standard need for Lipschitz continuity by directly bounding gradient norms in terms of function values.

Empirical experiments with binary logistic regression on diverse datasets demonstrate that both higher momentum and larger step sizes lead to increased instability in practice, in concordance with the theoretical bounds.

Figure 1: Tests on different step sizes for Polyak’s momentum variant demonstrate monotonic increases in output divergence with step size increments.

Figure 2: The effect of increasing momentum parameter $\beta$ on the distance between models produced on neighboring datasets confirms that higher momentum amplifies instability in both Polyak and Nesterov SGDM.

Optimization Error and Excess Population Risk Bounds

An auxiliary sequence, evolving by a reduced SGD recursion, is utilized to facilitate tight optimization error bounds for SGDM. These, together with the stability analysis, yield high-probability excess risk upper bounds:

• Excess Population Risk: For both momentum schemes, with chosen step sizes, the excess population risk $E[L(\bar{w}_t)] - L(w^*)$ attains the optimal rate $\mathcal{O}(1/\sqrt{n})$, matching minimax lower bounds.
• Low-Noise Regime: When the Bayes risk is exceptionally low ($L(w^*) < 1/n$), the risk further improves to $w_{t+1} = w_t - \eta m_t$0, matching optimistic results for pure SGD.

The step size and momentum parameter selection are precisely specified to maintain these rates, and the bounds are shown to be both theoretically and empirically robust.

Implications and Future Directions

The analytic results rigorously refute conjectures that momentum necessarily destroys generalization for convex and smooth problems. Instead, it is shown that momentum degrades stability and thus generalization only by a constant factor, scalable with $w_{t+1} = w_t - \eta m_t$1, provided step sizes are controlled. This insight is immediately relevant for practitioners deploying SGDM in large-scale risk minimization, as it justifies widespread empirical observations of negligible generalization degradation in deep learning.

Open directions include:

• Nonconvex Settings: Extension of the analysis to general nonconvex regimes, including deep neural networks, remains a natural target.
• Beyond On-Average Stability: Uniform stability analysis for momentum methods under weaker or stochastic conditions, potentially leveraging recent developments in high-dimensional concentration.
• Generalization in the Presence of Strong Momentum: Understanding the practical tradeoff between accelerated optimization and generalization, particularly as $w_{t+1} = w_t - \eta m_t$2, in deep or implicit regularization scenarios.

Conclusion

This work establishes, via rigorous on-average model stability and optimization error analyses, that both Polyak and Nesterov SGDM enjoy provable, tight generalization guarantees for smooth convex risk minimization problems. The stability degradation induced by momentum is precisely quantified, and optimal excess risk rates are shown to be preserved under momentum. The theoretical developments, empirical verification, and introduced analytic techniques substantially advance the understanding of algorithmic stability in momentum-driven stochastic optimization.