Research map / Optimization / Implicit bias and generalization
Stochastic Gradient Flow Dynamics of Test Risk and its Exact Solution for Weak Features
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Derives via path integral and Laplace approximation a small-learning-rate formula for the test-risk correction of stochastic gradient flow relative to gradient flow, and computes it exactly for a weak features model with double descent.
Problem
Pure gradient flow analyses of double descent ignore the effect of SGD stochasticity on the test risk over the whole training dynamics; a general tractable description of the fluctuations of stochastic gradient flow around the deterministic trajectory is lacking.
Approach
Represent the Itô SDE transition probability as a path integral, apply a Laplace approximation for small learning rate to obtain a Gaussian covariance of fluctuations around the gradient-flow trajectory (C(t) via time-ordered exponential of the Hessian), and apply it to a weak features linear model to get the full time-dependent test-risk correction.
Key results
- Path-integral Laplace formula reproduces exactly the covariance of a solvable linear SDE.
- For the weak features model, the infinite-time stochastic correction is nonzero only in the underparametrized regime alpha < 1; it vanishes for alpha > 1.
- Asymptotic SGF-GF test-risk differences match discrete SGD simulations at small and large times but underestimate them for intermediate times.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add stochastic gradient-flow noise | weak features model · underparameterized regime, α < 1, t → ∞ | test risk | Hurt |
| add stochastic gradient-flow noise | weak features model · overparameterized regime, α > 1, t → ∞ | test risk | No clear effect |
| approximate SGD with the analytical SGF theory | weak features model · discrete-time SGD simulations | test risk | Mixed |
| add stochastic gradient-flow noise | weak features model, ψ = 2.5 · finite times across the interpolation threshold | test-risk difference between SGF and GF | Mixed |
Surprising findings
- The stochastic correction to test risk vanishes at infinite time above the interpolation threshold, so stochasticity has no lasting effect there.
- The asymptotic theory underestimates SGD test-risk differences at intermediate times.
Limitations
- Heuristic path integral and Laplace approximation, valid only for small learning rate; not rigorous for state-dependent covariance.
- Only a simple weak features linear model is treated; extension to more complex landscapes left open.
- Analytical theory underestimates SGD at intermediate times.
Experimental setup
Datasets: synthetic Gaussian data for weak features linear regression
Metrics: test risk difference between SGF and GF test risk
Baselines: pure gradient flow test risk discrete-time SGD simulations
Builds on: stochastic gradient flow approximation of SGD Freidlin-Wentzell large deviation theory Marchenko-Pastur law
Closest papers by meaning
- Learning Curves for SGD on Structured Features (ICLR 2022 · similarity 0.87)
- On the Overlooked Structure of Stochastic Gradients (NeurIPS 2023 · similarity 0.86)
- Implicit Bias of SGD for Diagonal Linear Networks: a Provable Benefit of Stochasticity (NeurIPS 2021 · similarity 0.86)
- Understanding the unstable convergence of gradient descent (ICML 2022 · similarity 0.86)
- The High Line: Exact Risk and Learning Rate Curves of Stochastic Adaptive Learning Rate Algorithms (NeurIPS 2024 · similarity 0.85)
- The Global Convergence Time of Stochastic Gradient Descent in Non-Convex Landscapes: Sharp Estimates via Large Deviations (ICML 2025 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.