Train faster, generalize better: Stability of stochastic gradient descent
This work addresses a foundational problem in machine learning by explaining generalization in SGD, which is incremental as it builds on existing stability theory.
The paper tackles the problem of understanding why stochastic gradient descent (SGD) generalizes well in practice, showing that models trained with few SGD iterations have vanishing generalization error due to algorithmic stability. It provides stability bounds for both convex and non-convex optimization, applying insights to neural networks and common training techniques.
We show that parametric models trained by a stochastic gradient method (SGM) with few iterations have vanishing generalization error. We prove our results by arguing that SGM is algorithmically stable in the sense of Bousquet and Elisseeff. Our analysis only employs elementary tools from convex and continuous optimization. We derive stability bounds for both convex and non-convex optimization under standard Lipschitz and smoothness assumptions. Applying our results to the convex case, we provide new insights for why multiple epochs of stochastic gradient methods generalize well in practice. In the non-convex case, we give a new interpretation of common practices in neural networks, and formally show that popular techniques for training large deep models are indeed stability-promoting. Our findings conceptually underscore the importance of reducing training time beyond its obvious benefit.