Stochastic Differential Equations models for Least-Squares Stochastic Gradient Descent

TL;DR

This paper models the dynamics of stochastic gradient descent for least-squares problems using stochastic differential equations, providing convergence rates, distribution characterizations, and insights into heavy-tail phenomena.

Contribution

It introduces a continuous-time SDE framework for analyzing SGD in least-squares, extending previous work and offering detailed convergence and distribution results.

Findings

Non-asymptotic convergence rates to stationary distribution

Characterization of the asymptotic distribution including mean and deviations

Identification of heavy-tail emergence related to step-size

Abstract

We study the dynamics of a continuous-time model of the Stochastic Gradient Descent (SGD) for the least-square problem. Indeed, pursuing the work of Li et al. (2019), we analyze Stochastic Differential Equations (SDEs) that model SGD either in the case of the training loss (finite samples) or the population one (online setting). A key qualitative feature of the dynamics is the existence of a perfect interpolator of the data, irrespective of the sample size. In both scenarios, we provide precise, non-asymptotic rates of convergence to the (possibly degenerate) stationary distribution. Additionally, we describe this asymptotic distribution, offering estimates of its mean, deviations from it, and a proof of the emergence of heavy-tails related to the step-size magnitude. Numerical simulations supporting our findings are also presented.