Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

What Happens after SGD Reaches Zero Loss? --A Mathematical Framework

Zhiyuan Li, Tianhao Wang, Sanjeev Arora

TL;DR

This work develops a global diffusion framework to analyze the implicit bias of SGD in heavily overparametrized models by projecting dynamics onto the manifold $\Gamma$ of local minimizers. Building on Katzenberger's theory, it derives a limiting SDE for $\Theta(\eta^{-2})$ steps with arbitrary noise covariance, revealing how tangent noise and Itô corrections shape the slow evolution along $\Gamma$ through tangent regularization and mixed/normal regularization terms. In particular, label-noise SGD is shown to induce a regularizer-driven flow on $\Gamma$ that can recover sparse groundtruth in overparametrized linear models with sample complexity $O(\kappa \log d)$, outperforming gradient descent in the kernel regime which requires $\Omega(d)$. The framework thus explains how stochasticity can promote generalization by guiding optimization toward low-complexity solutions, and it provides a principled avenue to analyze other stochastic optimizers through similar diffusion limits.

Abstract

Understanding the implicit bias of Stochastic Gradient Descent (SGD) is one of the key challenges in deep learning, especially for overparametrized models, where the local minimizers of the loss function $L$ can form a manifold. Intuitively, with a sufficiently small learning rate $η$, SGD tracks Gradient Descent (GD) until it gets close to such manifold, where the gradient noise prevents further convergence. In such a regime, Blanc et al. (2020) proved that SGD with label noise locally decreases a regularizer-like term, the sharpness of loss, $\mathrm{tr}[\nabla^2 L]$. The current paper gives a general framework for such analysis by adapting ideas from Katzenberger (1991). It allows in principle a complete characterization for the regularization effect of SGD around such manifold -- i.e., the "implicit bias" -- using a stochastic differential equation (SDE) describing the limiting dynamics of the parameters, which is determined jointly by the loss function and the noise covariance. This yields some new results: (1) a global analysis of the implicit bias valid for $η^{-2}$ steps, in contrast to the local analysis of Blanc et al. (2020) that is only valid for $η^{-1.6}$ steps and (2) allowing arbitrary noise covariance. As an application, we show with arbitrary large initialization, label noise SGD can always escape the kernel regime and only requires $O(κ\ln d)$ samples for learning an $κ$-sparse overparametrized linear model in $\mathbb{R}^d$ (Woodworth et al., 2020), while GD initialized in the kernel regime requires $Ω(d)$ samples. This upper bound is minimax optimal and improves the previous $\tilde{O}(κ^2)$ upper bound (HaoChen et al., 2020).

What Happens after SGD Reaches Zero Loss? --A Mathematical Framework

TL;DR

This work develops a global diffusion framework to analyze the implicit bias of SGD in heavily overparametrized models by projecting dynamics onto the manifold of local minimizers. Building on Katzenberger's theory, it derives a limiting SDE for steps with arbitrary noise covariance, revealing how tangent noise and Itô corrections shape the slow evolution along through tangent regularization and mixed/normal regularization terms. In particular, label-noise SGD is shown to induce a regularizer-driven flow on that can recover sparse groundtruth in overparametrized linear models with sample complexity , outperforming gradient descent in the kernel regime which requires . The framework thus explains how stochasticity can promote generalization by guiding optimization toward low-complexity solutions, and it provides a principled avenue to analyze other stochastic optimizers through similar diffusion limits.

Abstract

Understanding the implicit bias of Stochastic Gradient Descent (SGD) is one of the key challenges in deep learning, especially for overparametrized models, where the local minimizers of the loss function can form a manifold. Intuitively, with a sufficiently small learning rate , SGD tracks Gradient Descent (GD) until it gets close to such manifold, where the gradient noise prevents further convergence. In such a regime, Blanc et al. (2020) proved that SGD with label noise locally decreases a regularizer-like term, the sharpness of loss, . The current paper gives a general framework for such analysis by adapting ideas from Katzenberger (1991). It allows in principle a complete characterization for the regularization effect of SGD around such manifold -- i.e., the "implicit bias" -- using a stochastic differential equation (SDE) describing the limiting dynamics of the parameters, which is determined jointly by the loss function and the noise covariance. This yields some new results: (1) a global analysis of the implicit bias valid for steps, in contrast to the local analysis of Blanc et al. (2020) that is only valid for steps and (2) allowing arbitrary noise covariance. As an application, we show with arbitrary large initialization, label noise SGD can always escape the kernel regime and only requires samples for learning an -sparse overparametrized linear model in (Woodworth et al., 2020), while GD initialized in the kernel regime requires samples. This upper bound is minimax optimal and improves the previous upper bound (HaoChen et al., 2020).

Paper Structure

This paper contains 24 sections, 13 theorems, 25 equations, 1 figure.

Table of Contents

  1. Introduction
  2. Intuitive explanation of regularization effect due to SGD
  3. Our Approach: Separating the Slow from the Fast
  4. Related Works
  5. Loss Landscape of Overparametrized Models
  6. Implicit Bias in Overparametrized Models
  7. Modelling Stochastic First-Order Methods with Itô SDE
  8. Notation and Preliminaries
  9. Manifold of local minimizers
  10. Limiting Diffusion of SGD
  11. Recap of Katzenberger's Theorem
  12. Closed-Form expression of the limiting diffusion
  13. Main Result
  14. Implications and Examples
  15. Type I: Isotropic Noise.
  16. ...and 9 more sections

Key Result

Theorem 4.1

Suppose the loss $L$, manifold $\Gamma$ and neighborhood $U$ satisfies assump:manifoldassump:neighborhood. Let $\{X_n\}_{n\geq 1}$ be a sequence of Katzenberger process with $\{A_n\}_{n\geq1},\{Z_n\}_{n\geq 1}$. Let $Y_n(t)=X_n(t)-\phi(X(0),A_n(t))+\Phi(X_0)$. Under technical assumptions, it holds t

Figures (1)

  • Figure 1: Illustration for limiting flow in $\mathbb{R}^2$. $\Gamma$ is an 1D manifold of minimizers of loss $L$.

Theorems & Definitions (20)

  • Theorem 4.1: Informal version of \ref{['thm:limit_diffusion_main']}, katzenberger1991solutions
  • Lemma 4.1
  • Lemma 4.1
  • Definition 4.2: Lyapunov Operator
  • Lemma 4.3
  • Theorem 4.4
  • Remark 4.5
  • Corollary 5.0: Limiting Diffusion for Isotropic Noise
  • Corollary 5.0: Limiting Flow for Label Noise
  • Theorem 6.1
  • ...and 10 more