On the Provable Suboptimality of Momentum SGD in Nonstationary Stochastic Optimization

TL;DR

The paper analyzes momentum SGD in nonstationary stochastic optimization, where the data distribution drifts and the minimizer moves over time. It develops finite-time, expectation, and high-probability bounds for SGD and momentum variants under uniformly strongly convex and smooth objectives, revealing a three-term error decomposition: initialization forgetting, irreducible gradient-noise floors, and drift-induced tracking lag. A key finding is the inertia window: as momentum $β$ approaches 1, momentum can amplify drift and delay adaptation, making SGD provably better in drift-dominated regimes, a fact supported by minimax lower bounds on dynamic regret. Experiments across drifting quadratics, linear/logistic regression, and teacher–student MLPs confirm the regime split: momentum helps in near-stationary, noise-dominated settings but harms tracking under genuine regime shifts, especially with ill-conditioning. These results provide a rigorous theoretical foundation for the empirical instability of momentum in dynamic environments and guide practical strategies like forgetting/restart mechanisms to mitigate drift effects.

Abstract

While momentum-based acceleration has been studied extensively in deterministic optimization problems, its behavior in nonstationary environments -- where the data distribution and optimal parameters drift over time -- remains underexplored. We analyze the tracking performance of Stochastic Gradient Descent (SGD) and its momentum variants (Polyak heavy-ball and Nesterov) under uniform strong convexity and smoothness in varying stepsize regimes. We derive finite-time bounds in expectation and with high probability for the tracking error, establishing a sharp decomposition into three components: a transient initialization term, a noise-induced variance term, and a drift-induced tracking lag. Crucially, our analysis uncovers a fundamental trade-off: while momentum can suppress gradient noise, it incurs an explicit penalty on the tracking capability. We show that momentum can substantially amplify drift-induced tracking error, with amplification that becomes unbounded as the momentum parameter approaches one, formalizing the intuition that using 'stale' gradients hinders adaptation to rapid regime shifts. Complementing these upper bounds, we establish minimax lower bounds for dynamic regret under gradient-variation constraints. These lower bounds prove that the inertia-induced penalty is not an artifact of analysis but an information-theoretic barrier: in drift-dominated regimes, momentum creates an unavoidable 'inertia window' that fundamentally degrades performance. Collectively, these results provide a definitive theoretical grounding for the empirical instability of momentum in dynamic environments and delineate the precise regime boundaries where SGD provably outperforms its accelerated counterparts.

Figures (5)

• Figure 1: Trade-off between drift tracking and stochastic noise under constant stepsize (expectation bounds). For SGD, the expected tracking error decomposes into an exponentially decaying initialization term plus two irreducible steady-state floors: a noise floor $\asymp \sigma^{2}\gamma/\mu$ and a drift-induced tracking floor $\asymp \Delta^{2}/(\mu^{2}\gamma^{2})$. For SGDM, momentum increases sensitivity to initialization by a factor $\asymp (1-\beta)^{-2}$ and inflates the noise floor by $\asymp (1-\beta)^{-1}$, while the drift floor retains the same $\gamma^{-2}$ scaling (up to constants). Moreover, stability requires $\gamma \le \mu(1-\beta)^{2}/(4L^{2})$, so large $\beta$ can indirectly worsen drift tracking by forcing smaller admissible stepsizes. Together, these effects formalize when momentum helps in stationary regimes yet becomes fragile under drift, and when SGD is provably more robust.
• Figure 2: Two minimax regimes and the inertia window.(a) After a regime change (distribution shift), momentum averages past gradients: this reduces noise but induces inertia, so the iterate lags behind the drifting minimizer for an "inertia window" whose length grows with the momentum parameter. (b) The minimax lower bound in \ref{['thm:minimax-lower-bound']} is the maximum of two contributions: a noise/variation-limited term which is bottlenecked by information rather than inertia, and an inertia-limited term that worsens with momentum and dominates when delay is the bottleneck. The crossover separates a statistical regime from an inertia regime, explaining when momentum is provably worse than SGD.
• Figure 3: Tracking a drifting minimizer under strong convexity. Mean squared tracking error $\mathbb{E}\!\left[\|\boldsymbol{\theta}_t-\boldsymbol{\theta}_t^\star\|_2^2\right]$ versus time for SGD, Heavy-Ball (HB), and Nesterov (NAG) on a strongly convex quadratic with noisy gradients. Parameters $(\gamma,\beta,\sigma^2)$ are shown above each panel with shaded bands denoting $\pm1$ std over random seeds. Across both regimes, momentum methods (HB/NAG) suppress short-term noise but exhibit a substantially larger steady-state tracking error—consistent with inertia-induced lag when the minimizer drifts—whereas SGD tracks the minimizer more closely but with higher variability.
• Figure 4: Non-stationary benchmark suite. Results for Linear and Logistic tasks. (Continued on next page.)
• Figure 5: Non-stationary benchmark suite (streaming) across linear, logistic, and teacher--student MLP tasks. We report mean $\pm$ std curves over seeds under two conditioning regimes (well-conditioned: $\kappa=10$; ill-conditioned: $\kappa=1000$). For linear/logistic tasks, tracking error is measured in parameter space as $\|\boldsymbol{\theta}_t-\boldsymbol{\theta}_t^\star\|$. For the teacher--student MLP regression task, tracking is measured in prediction space (function space), i.e., $\mathbb{E}_{x\sim\mathcal{D}}[\|f_{\boldsymbol{\theta}_t}(x)-f_{\boldsymbol{\theta}_t^\star}(x)\|^2]$ estimated on a fixed validation set (reported as tracking error and validation MSE). Methods compared: SGD, Heavy-Ball (HB), and Nesterov (NAG), with step sizes respecting the regime-dependent stability caps used in our analysis.

Theorems & Definitions (83)

• Definition 2.1: Conditional Orlicz norm
• Definition 2.2: Conditional Orlicz norm (vector)
• Theorem 3.1: Tracking error bound in expectation for SGD
• Theorem 3.2: Time to reach the asymptotic tracking error in expectation for SGD
• Lemma 3.1: Extended 2D recursion for SGD with momentum
• Theorem 3.3: Tracking error bound in expectation for SGD with momentum
• Theorem 3.4: Time to reach the asymptotic tracking error in expectation for SGD with momentum
• Theorem 3.5: High probability tracking error bound for SGD
• Theorem 3.6: High probability tracking error bound for SGDM
• Theorem 3.7: Minimax lower bound for strongly-convex function sequences using SGDM