Exact Mean Square Linear Stability Analysis for SGD

TL;DR

This work derives an explicit mean-square stability threshold for SGD near minima, showing a necessary-and-sufficient condition on the learning rate $\eta$ that depends on the batch size $B$ through a pair of Kronecker-structured matrices ${\boldsymbol{C}}$ and ${\boldsymbol{D}}$, and proves that the threshold is nondecreasing in $B$ so smaller batches can reduce stability. It reveals an equivalent mixture process ALG$(p)$ with $p=(n-B)/(B(n-1))\approx 1/B$ that shares SGD's stability properties, linking mini-batch SGD closely to full-batch GD for typical batch sizes. The analysis extends to non-interpolating (regular) minima and provides the asymptotic covariance of the dynamics, along with simple necessary stability criteria, all corroborated by MNIST experiments. These results offer a precise, actionable stability criterion and illuminate how batch size and learning rate jointly shape SGD dynamics in deep learning contexts.

Abstract

The dynamical stability of optimization methods at the vicinity of minima of the loss has recently attracted significant attention. For gradient descent (GD), stable convergence is possible only to minima that are sufficiently flat w.r.t. the step size, and those have been linked with favorable properties of the trained model. However, while the stability threshold of GD is well-known, to date, no explicit expression has been derived for the exact threshold of stochastic GD (SGD). In this paper, we derive such a closed-form expression. Specifically, we provide an explicit condition on the step size that is both necessary and sufficient for the linear stability of SGD in the mean square sense. Our analysis sheds light on the precise role of the batch size $B$. In particular, we show that the stability threshold is monotonically non-decreasing in the batch size, which means that reducing the batch size can only decrease stability. Furthermore, we show that SGD's stability threshold is equivalent to that of a mixture process which takes in each iteration a full batch gradient step w.p. $1-p$, and a single sample gradient step w.p. $p$, where $p \approx 1/B $. This indicates that even with moderate batch sizes, SGD's stability threshold is very close to that of GD's. We also prove simple necessary conditions for linear stability, which depend on the batch size, and are easier to compute than the precise threshold. Finally, we derive the asymptotic covariance of the dynamics around the minimum, and discuss its dependence on the learning rate. We validate our theoretical findings through experiments on the MNIST dataset.

Figures (3)

• Figure 1: Sharpness vs. step size and batch size. We trained single hidden-layer ReLU networks using varying step sizes and batch sizes on a subset of MNIST. Panel subfig:Sharpness vs. step size visualizes the sharpness of the converged minima versus learning rate for different batch sizes. For small batch sizes, $\lambda_{\max}( {\boldsymbol{H} } )$ deviates significantly from $2/\eta$. Yet as the batch size increases to a moderate value, these curves coincide, indicating that in terms of stability, SGD behaves similarly to GD. Panel subfig:Sharpness vs. batch size plots the sharpness against the batch size for three different learning rates $\eta_1 = 0.043, \eta_2 = 0.012, \eta_3 = 0.002$. Here we see a similar trend where SGD behaves like GD for $B \geq 32$.
• Figure 2: (Generalized) Sharpness vs. step size. We trained single hidden-layer ReLU networks using varying step sizes and batch sizes on MNIST dataset. For each pair of hyper-parameters $(\eta, B)$, we measured the sharpness of the minimum (yellow), our necessary condition for stability (blue), and the optimized bound (purple), which their relations are given in \ref{['eq:relationship']}. We see that for small batch sizes $B = 1$ and $B = 2$, the optimized bound \ref{['eq:optimized bound']} coincides with $2/\eta$, confirming that SGD converged at the edge of stability ($\eta = \eta^*_{\mathrm{var}}$). For additional insights and detail, see Sec. \ref{['sec:Experiments']}.
• Figure 3: Sharpness vs. learning rate. Additional results for the experiment in Sec. \ref{['sec:Experiments']}. These two figures complete the results of Fig. \ref{['Fig:A']}. Here we see that SGD with big batch sizes behaves like GD.

Theorems & Definitions (19)

• definition 1: Linearized dynamics
• theorem 1: Stability of the mean
• definition 2: Interpolating minima
• theorem 2: ma2021on, Thm. 1 + Cor. 3
• theorem 3: Mean square stability for interpolating minima
• proposition 1: Monotonicity of the stability threshold
• proposition 2: Equivalent mixture process
• proposition 3: Stability gap
• proposition 4: Necessary conditions for stability
• definition 3: Regular minima