Data-Dependent Stability of Stochastic Gradient Descent

TL;DR

This paper introduces a data-dependent stability framework for SGD, leading to new generalization bounds that depend on initial risk or curvature, offering practical stabilization strategies.

Contribution

It develops a data-dependent stability analysis for SGD, providing novel generalization bounds based on initial risk and curvature, unlike previous worst-case results.

Findings

Generalization bounds depend on initial risk in convex cases.

Curvature around initialization influences non-convex generalization.

Data-driven stabilization strategies improve SGD performance.

Abstract

We establish a data-dependent notion of algorithmic stability for Stochastic Gradient Descent (SGD), and employ it to develop novel generalization bounds. This is in contrast to previous distribution-free algorithmic stability results for SGD which depend on the worst-case constants. By virtue of the data-dependent argument, our bounds provide new insights into learning with SGD on convex and non-convex problems. In the convex case, we show that the bound on the generalization error depends on the risk at the initialization point. In the non-convex case, we prove that the expected curvature of the objective function around the initialization point has crucial influence on the generalization error. In both cases, our results suggest a simple data-driven strategy to stabilize SGD by pre-screening its initialization. As a corollary, our results allow us to show optimistic generalization bounds that exhibit fast convergence rates for SGD subject to a vanishing empirical risk and low noise of stochastic gradient.