The Implicit Regularization of Dynamical Stability in Stochastic Gradient Descent

TL;DR

This paper investigates how stochastic gradient descent (SGD) implicitly regularizes model complexity through dynamical stability, leading to better generalization compared to gradient descent (GD), especially influenced by the learning rate.

Contribution

It establishes a theoretical link between stability metrics and generalization in SGD, contrasting it with GD, and highlights the role of learning rate in regularization strength.

Findings

Stable minima of SGD generalize well due to sharpness and norm equivalence.

GD's stability is too weak for effective regularization.

Larger learning rates enhance SGD's regularization effect.

Abstract

In this paper, we study the implicit regularization of stochastic gradient descent (SGD) through the lens of {\em dynamical stability} (Wu et al., 2018). We start by revising existing stability analyses of SGD, showing how the Frobenius norm and trace of Hessian relate to different notions of stability. Notably, if a global minimum is linearly stable for SGD, then the trace of Hessian must be less than or equal to $2/\eta$, where $\eta$ denotes the learning rate. By contrast, for gradient descent (GD), the stability imposes a similar constraint but only on the largest eigenvalue of Hessian. We then turn to analyze the generalization properties of these stable minima, focusing specifically on two-layer ReLU networks and diagonal linear networks. Notably, we establish the {\em equivalence} between these metrics of sharpness and certain parameter norms for the two models, which allows us to show that the stable minima of SGD provably generalize well. By contrast, the stability-induced regularization of GD is provably too weak to ensure satisfactory generalization. This discrepancy provides an explanation of why SGD often generalizes better than GD. Note that the learning rate (LR) plays a pivotal role in the strength of stability-induced regularization. As the LR increases, the regularization effect becomes more pronounced, elucidating why SGD with a larger LR consistently demonstrates superior generalization capabilities. Additionally, numerical experiments are provided to support our theoretical findings.