On the Stability of Nonlinear Dynamics in GD and SGD: Beyond Quadratic Potentials
This work addresses a foundational issue in optimization theory for machine learning practitioners, providing new criteria for stability analysis beyond linear approximations.
The authors tackled the problem of whether linearized dynamics accurately capture the stability of gradient descent (GD) and stochastic gradient descent (SGD) near minima, showing that nonlinear terms can lead to stable oscillations in GD and that SGD stability can be dictated by a single unstable batch rather than average effects.
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.