Stochastic Modified Flows, Mean-Field Limits and Dynamics of Stochastic Gradient Descent
This work provides more precise theoretical tools for analyzing SGD fluctuations, which is incremental for researchers in optimization and machine learning theory.
The authors tackled the problem of accurately modeling stochastic gradient descent (SGD) dynamics by proposing stochastic modified flows (SDEs) that improve upon existing stochastic modified equations, offering regular diffusion coefficients and matching multi-point statistics, and extended this to distribution-dependent flows for SGD in a small learning rate-infinite width regime.
We propose new limiting dynamics for stochastic gradient descent in the small learning rate regime called stochastic modified flows. These SDEs are driven by a cylindrical Brownian motion and improve the so-called stochastic modified equations by having regular diffusion coefficients and by matching the multi-point statistics. As a second contribution, we introduce distribution dependent stochastic modified flows which we prove to describe the fluctuating limiting dynamics of stochastic gradient descent in the small learning rate - infinite width scaling regime.