From Gradient Flow on Population Loss to Learning with Stochastic Gradient Descent

TL;DR

This paper establishes a theoretical framework connecting the convergence of Gradient Flow on population loss to the convergence of Stochastic Gradient Descent, providing new insights into when and why SGD succeeds in training complex models.

Contribution

It introduces a general converse Lyapunov theorem linking GF convergence to SGD convergence, applicable to a wide range of non-convex problems including phase retrieval and matrix square-root.

Findings

Unified analysis for GD/SGD across classical and complex problems

Conditions under which SGD converges based on GF convergence

Extension of results to non-convex and structured objectives

Abstract

Stochastic Gradient Descent (SGD) has been the method of choice for learning large-scale non-convex models. While a general analysis of when SGD works has been elusive, there has been a lot of recent progress in understanding the convergence of Gradient Flow (GF) on the population loss, partly due to the simplicity that a continuous-time analysis buys us. An overarching theme of our paper is providing general conditions under which SGD converges, assuming that GF on the population loss converges. Our main tool to establish this connection is a general converse Lyapunov like theorem, which implies the existence of a Lyapunov potential under mild assumptions on the rates of convergence of GF. In fact, using these potentials, we show a one-to-one correspondence between rates of convergence of GF and geometrical properties of the underlying objective. When these potentials further satisfy certain self-bounding properties, we show that they can be used to provide a convergence guarantee for Gradient Descent (GD) and SGD (even when the paths of GF and GD/SGD are quite far apart). It turns out that these self-bounding assumptions are in a sense also necessary for GD/SGD to work. Using our framework, we provide a unified analysis for GD/SGD not only for classical settings like convex losses, or objectives that satisfy PL / KL properties, but also for more complex problems including Phase Retrieval and Matrix sq-root, and extending the results in the recent work of Chatterjee 2022.