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.

New logarithmic step size for stochastic gradient descent

Large Stepsize Gradient Descent for Logistic Loss: Non-Monotonicity of the Loss Improves Optimization Efficiency

Barzilai-Borwein Step Size for Stochastic Gradient Descent

Towards Noise-adaptive, Problem-adaptive (Accelerated) Stochastic Gradient Descent

Provably Faster Gradient Descent via Long Steps

Step-size Optimization for Continual Learning

Adaptive Step Sizes for Preconditioned Stochastic Gradient Descent

Bandwidth-based Step-Sizes for Non-Convex Stochastic Optimization

Second-order step-size tuning of SGD for non-convex optimization

Learning the Step-size Policy for the Limited-Memory Broyden-Fletcher-Goldfarb-Shanno Algorithm

Accelerated Gradient Descent via Long Steps

Step Size Matters in Deep Learning

Painless step size adaptation for SGD

Accelerated Gradient Descent by Concatenation of Stepsize Schedules

Relationship between Batch Size and Number of Steps Needed for Nonconvex Optimization of Stochastic Gradient Descent using Armijo Line Search

Stochastic Ratios Tracking Algorithm for Large Scale Machine Learning Problems

Large Stepsize Gradient Descent for Non-Homogeneous Two-Layer Networks: Margin Improvement and Fast Optimization

Scaling transition from momentum stochastic gradient descent to plain stochastic gradient descent

Exact Mean Square Linear Stability Analysis for SGD

Speed learning on the fly