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

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 $\eta$, 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 ${\eta }^{-2}$ steps, in contrast to the local analysis of Blanc et al. (2020) that is only valid for ${\eta }^{-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(\kappa \ln d)$ samples for learning an $\kappa$-sparse overparametrized linear model in ${\mathbb{R}}^d$ (Woodworth et al., 2020), while GD initialized in the kernel regime requires $\Omega (d)$ samples. This upper bound is minimax optimal and improves the previous $\tilde{O}({\kappa }^2)$ upper bound (HaoChen et al., 2020).

arXiv comment56 pages, 2 figures; ICLR 2022

References in corpus (31)

• Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour
• Neural Tangent Kernel: Convergence and Generalization in Neural Networks
• One weird trick for parallelizing convolutional neural networks
• A Convergence Theory for Deep Learning via Over-Parameterization
• Train longer, generalize better: closing the generalization gap in large batch training of neural networks
• Gradient Descent Provably Optimizes Over-parameterized Neural Networks
• Learning Overparameterized Neural Networks via Stochastic Gradient Descent on Structured Data
• On Lazy Training in Differentiable Programming
• Three Factors Influencing Minima in SGD
• Essentially No Barriers in Neural Network Energy Landscape
• Characterizing Implicit Bias in Terms of Optimization Geometry
• Scaling Limits of Wide Neural Networks with Weight Sharing: Gaussian Process Behavior, Gradient Independence, and Neural Tangent Kernel Derivation
• Towards Explaining the Regularization Effect of Initial Large Learning Rate in Training Neural Networks
• Gradient Descent Maximizes the Margin of Homogeneous Neural Networks
• The Anisotropic Noise in Stochastic Gradient Descent: Its Behavior of Escaping from Sharp Minima and Regularization Effects
• First-order Methods Almost Always Avoid Saddle Points
• SGD Learns the Conjugate Kernel Class of the Network
• On the loss landscape of a class of deep neural networks with no bad local valleys
• Understanding the Loss Surface of Neural Networks for Binary Classification
• Spurious Valleys in Two-layer Neural Network Optimization Landscapes
• On the Noisy Gradient Descent that Generalizes as SGD
• Explaining Landscape Connectivity of Low-cost Solutions for Multilayer Nets
• The loss landscape of overparameterized neural networks
• Implicit Regularization for Optimal Sparse Recovery
• Why Are Convolutional Nets More Sample-Efficient than Fully-Connected Nets?
• On the Validity of Modeling SGD with Stochastic Differential Equations (SDEs)
• Shape Matters: Understanding the Implicit Bias of the Noise Covariance
• Stochastic Gradient and Langevin Processes
• Reconciling Modern Deep Learning with Traditional Optimization Analyses: The Intrinsic Learning Rate
• Label Noise SGD Provably Prefers Flat Global Minimizers
• The critical locus of overparameterized neural networks