Hitting the High-Dimensional Notes: An ODE for SGD learning dynamics on GLMs and multi-index models

TL;DR

This paper develops a deterministic ODE framework to describe the high-dimensional dynamics of SGD on generalized linear and multi-index models, providing insights into stability, convergence, and statistical behavior.

Contribution

It introduces a novel ODE-based analysis of SGD dynamics in high dimensions, including stability thresholds and a simplified SDE model for statistics.

Findings

Deterministic ODE system accurately describes SGD dynamics.

Identifies learning rate thresholds for stability.

Numerical simulations match theoretical predictions.

Abstract

We analyze the dynamics of streaming stochastic gradient descent (SGD) in the high-dimensional limit when applied to generalized linear models and multi-index models (e.g. logistic regression, phase retrieval) with general data-covariance. In particular, we demonstrate a deterministic equivalent of SGD in the form of a system of ordinary differential equations that describes a wide class of statistics, such as the risk and other measures of sub-optimality. This equivalence holds with overwhelming probability when the model parameter count grows proportionally to the number of data. This framework allows us to obtain learning rate thresholds for stability of SGD as well as convergence guarantees. In addition to the deterministic equivalent, we introduce an SDE with a simplified diffusion coefficient (homogenized SGD) which allows us to analyze the dynamics of general statistics of SGD iterates. Finally, we illustrate this theory on some standard examples and show numerical simulations which give an excellent match to the theory.