On the Stability of Nonlinear Dynamics in GD and SGD: Beyond Quadratic Potentials

TL;DR

This paper investigates the nonlinear stability of gradient descent and stochastic gradient descent, revealing that nonlinear effects can cause divergence even when linear analysis suggests stability, and providing exact criteria for stability.

Contribution

It derives an exact nonlinear stability criterion for GD and extends the analysis to SGD, highlighting the limitations of linear stability analysis in nonlinear settings.

Findings

Nonlinear terms can cause divergence in SGD even if linear analysis indicates stability.

GD can stably oscillate near linearly unstable minima, challenging linear stability assumptions.

If all batches are linearly stable, SGD's nonlinear dynamics are stable in expectation.

Abstract

The dynamical stability of the iterates during training plays a key role in determining the minima obtained by optimization algorithms. For example, stable solutions of gradient descent (GD) correspond to flat minima, which have been associated with favorable features. While prior work often relies on linearization to determine stability, it remains unclear whether linearized dynamics faithfully capture the full nonlinear behavior. Recent work has shown that GD may stably oscillate near a linearly unstable minimum and still converge once the step size decays, indicating that linear analysis can be misleading. In this work, we explicitly study the effect of nonlinear terms. Specifically, we derive an exact criterion for stable oscillations of GD near minima in the multivariate setting. Our condition depends on high-order derivatives, generalizing existing results. Extending the analysis to stochastic gradient descent (SGD), we show that nonlinear dynamics can diverge in expectation even if a single batch is unstable. This implies that stability can be dictated by a single batch that oscillates unstably, rather than an average effect, as linear analysis suggests. Finally, we prove that if all batches are linearly stable, the nonlinear dynamics of SGD are stable in expectation.