[Paper Review] Stochastic Modified Flows, Mean-Field Limits and Dynamics of Stochastic Gradient Descent

Benjamin Gess, Sebastian Kassing+1 more|arXiv (Cornell University)|Feb 14, 2023

Markov Chains and Monte Carlo Methods4 citations

TL;DR

This paper introduces stochastic modified flows (SMFs), a new class of SDEs that model stochastic gradient descent (SGD) in the small learning rate regime with regular diffusion coefficients, avoiding the irregularity issues of prior stochastic modified equations (SMEs). It further establishes distribution-dependent SMFs as the fluctuating mean-field limit of SGD in the infinite-width, small-learning-rate scaling regime, matching multi-point statistics and enabling rigorous analysis of overparameterized networks.

ABSTRACT

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.

Motivation & Objective

• To resolve the irregularity of diffusion coefficients in stochastic modified equations (SMEs), which hinders analysis in overparameterized deep learning.
• To develop a limiting SDE model that matches the multi-point statistics of SGD, improving upon prior continuum limits.
• To establish a fluctuating mean-field limit for SGD in the infinite-width, small-learning-rate regime using distribution-dependent stochastic modified flows.
• To provide a unified framework for analyzing the dynamics of SGD that is robust to architectural and data-dependent variations.

Proposed method

• Proposes stochastic modified flows (SMFs) as SDEs driven by cylindrical Brownian motion, with diffusion coefficients derived directly from individual loss gradients rather than the degenerate covariance matrix Σ.
• Defines the SMF SDE as dXₜ^η(x) = -∇(R(Xₜ^η(x)) + η/4 |∇R(Xₜ^η(x))|²) dt + √η ∫_Θ G(Xₜ^η(x),θ) W(dθ,dt), where G(x,θ) = ∇R̃(x,θ) - ∇R(x).
• Uses a martingale problem formulation to show that SMFs reproduce the same statistical behavior as SMEs but with improved regularity.
• Introduces distribution-dependent SMFs to model the empirical measure dynamics of SGD in the infinite-width limit, with coefficients depending on the law of the solution.
• Applies Wasserstein-2 distance and propagation of chaos arguments to compare empirical measures of SMF solutions with finite SGD dynamics.
• Employs a combination of SDE theory, Malliavin calculus, and propagation of chaos to prove convergence rates in expectation and distributional sense.

Experimental results

Research questions

• RQ1 Can a new SDE-based limiting dynamics be constructed for SGD that avoids the irregularity of diffusion coefficients in classical stochastic modified equations?
• RQ2 Does the proposed stochastic modified flow model preserve the multi-point statistics of SGD, enabling a more accurate continuum approximation?
• RQ3 What is the fluctuating mean-field limit of SGD in the infinite-width, small-learning-rate regime, and how can it be described using distribution-dependent SDEs?
• RQ4 Can the convergence of finite SGD dynamics to the proposed limiting SDEs be quantified in terms of distributional and moment-based errors?

Key findings

• The proposed stochastic modified flow (SMF) SDE has regular diffusion coefficients derived from individual gradient differences, avoiding the need for the ill-behaved square root of the covariance matrix Σ.
• The SMF matches the multi-point statistics of SGD and satisfies the same martingale problem as the classical SME, ensuring statistical consistency.
• For every f ∈ C⁴_b(ℝᵈ), the expected value of f under the SMF and SGD dynamics are within O(η²) over finite time horizons, establishing second-order accuracy.
• In the infinite-width, small-learning-rate regime, the empirical measure of SGD converges to a distribution-dependent SMF, which captures the fluctuating mean-field dynamics.
• The convergence rate of the empirical measure to the limiting SMF is O(η²) in expectation, with additional error bounds depending on the initial measure's moments and sample size M.
• Under suitable moment and regularity conditions, the convergence rate is O(η²) even when the initial measure is approximated by a finite empirical measure, with explicit dependence on M and η.