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

1 INTRODUCTION . The implicit bias underlies the generalization ability of machine learning models trained by stochastic gradient descent ( SGD ) . But it still remains a mystery to mathematically characterize such bias . We study SGD in the following formulation xη ( k + 1 ) = xη ( k ) − η ( ∇L ( xη ( k ) ) + σξk ( xη ( k ) ) ) ( 1 ) where η is the learning rate ( LR ) , L : RD → R is the training loss and σ ( x ) = [ σ1 ( x ) , σ2 ( x ) , . . . , σΞ ( x ) ] ∈ RD×Ξ is a deterministic noise function . Here ξk is sampled uniformly from { 1 , 2 , . . . , Ξ } and it satisfies Eξk [ σξk ( x ) ] = 0 , ∀x ∈ Rd and k. It is widely believed that large LR ( or equivalently , small batch size ) helps SGD find better minima . For instance , some previous works argued that large noise enables SGD to select a flatter attraction basin of the loss landscape which potentiallyx benefits generalization ( Li et al. , 2019c ; Jastrzebski et al. , 2017 ) . However , there is also experimental evidence ( Li et al. , 2020b ) that small LR also has equally good implicit bias ( albeit with higher training time ) , and that is the case studied here . Presumably low LR precludes SGD jumping between different basins since under general conditions this should require Ω ( exp ( 1/η ) ) steps ( Shi et al. , 2020 ) . In other words , there should be a mechanism to reach better generalization while staying within a single basin . For deterministic GD similar mechanisms have been demonstrated in simple cases ( Soudry et al. , 2018 ; Lyu & Li , 2019 ) for GD and referred to as implicit bias of gradient descent . The current paper can be seen as study of implicit bias of Stochastic GD , which turns out to be quite different , mathematically . Recent work ( Blanc et al. , 2020 ) shed light on this direction by analyzing effects of stochasticity in the gradient . For sufficiently small LR , SGD will reach and be trapped around some manifold of local minimizers , denoted by Γ ( see Figure 2 ) . The effect is shown to be an implicit deterministic drift in a direction corresponding to lowering a regularizer-like term along the manifold . They showed SGD with label noise locally decreases the sharpness of loss , tr [ ∇2L ] , by Θ ( η0.4 ) in η−1.6 steps . However , such an analysis is actually local , since the natural time scale of analysis should be η−2 , not η−1.6 . The contribution of the current paper is a more general and global analysis of this type . We introduce a more powerful framework inspired by the classic paper Katzenberger ( 1991 ) . 1.1 INTUITIVE EXPLANATION OF REGULARIZATION EFFECT DUE TO SGD . We start with an intuitive description of the implicit regularization effect described in Blanc et al . ( 2020 ) . For simplification , we show it for the following canonical SDE approximation ( See Section B.1 for more details ) of SGD ( 1 ) ( Li et al. , 2017 ; Cheng et al. , 2020 ) . Here W ( t ) is the standard Ξ-dimensional Brownian motion . The only property about label noise SGD we will use is that the noise covariance σσ > ( x ) = ∇2L ( x ) for every x in the manifold Γ . ( See derivation in Section 5 ) dX̃η ( t ) = −η∇L ( X̃η ( t ) ) dt+ η · σ ( X̃η ( t ) ) dW ( t ) . ( 2 ) Suppose X̃η ( 0 ) is already close to some local minimizer point X∗ ∈ Γ . The goal is to show X̃η ( t ) will move in the tangent space and steadily decrease tr [ ∇2L ] . At first glance , this seems impossible as the gradient∇L vanishes around Γ , and the noise has zero mean , implying SGD should be like random walk instead of a deterministic drift . The key observation of Blanc et al . ( 2020 ) is that the local dynamics of X̃η ( t ) is completely different in tangent space and normal space — the fast random walk in normal space causes X̃η ( t ) to move slowly ( with velocity Θ ( η2 ) ) but deterministically in certain direction . To explain this , letting ∆ ( t ) = X̃η ( t ) − X∗ , Taylor expansion of ( 2 ) gives d∆ ( t ) ≈ −η∇2L ( X∗ ) ∆dt+ ησ ( X∗ ) dW ( t ) , meaning ∆ is behaving like an Ornstein-Uhlenbeck ( OU ) process locally in the normal space . Its mixing time is Θ ( η−1 ) and the stationary distribution is the standard multivariate gaussian in the normal space scaled by √ η ( see Figure 1b ) , because noise covariance σσ > = ∇2L . Though this OU process itself doesn ’ t form any regularization , it activates the second order Taylor expansion of ∇L ( X∗ + ∆ ( t ) ) , i.e. , − 12∂ 2 ( ∇L ) ( X∗ ) [ ∆ ( t ) , ∆ ( t ) ] , creating a Θ ( η2 ) velocity in the tangent space . Since there is no push back force in the tangent space , the small velocity accumulates over time , and in a longer time scale of Ω ( η−1 ) , the time average of the stochastic velocity is roughly the same as the expected velocity when ∆ is sampled from its stationary distribution . This simplifies the expression of the velocity in tangent space to η 2 2 ∇T tr [ ∇ 2L ] ( see Figure 1c ) , where ∇T means the gradient is only taken in the tangent space . However , the above approach only gives a local analysis forO ( η−1.6 ) time , where the total movement due to implicit regularization is O ( η2−1.6 ) = O ( η0.4 ) and thus is negligible when η → 0 . In order to get a non-trivial limiting dynamics when η → 0 , a global analysis for Ω ( η−2 ) steps is necessary and it can not be done by Taylor expansion with a single reference point . Recent work by Damian et al . ( 2021 ) glues analyzes of multiple local phases into a global guarantee that SGD finds a ( , γ ) stationary point for the regularized loss . This technique works for constant size LR , but this still doesn ’ t show convergence for trajectory when η → 0 . The main technical difficulty here is that it ’ s not clear how to separate the slow and fast dynamics in different spaces and how to only take limit for the slow dynamics , especially when shifting to a new reference point in the Taylor series calculation . 1.2 OUR APPROACH : SEPARATING THE SLOW FROM THE FAST . In this work , we tackle this problem via a different angle . First , since the anticipated limiting dynamics is of speed Θ ( η2 ) , we change the time scaling to accelerate ( 2 ) by η−2 times , which yields dXη ( t ) = −η−1∇L ( Xη ( t ) ) dt+ σ ( Xη ( t ) ) dW ( t ) . ( 3 ) The key idea here is that we only need to track the slow dynamic , or equivalently , some projection of X onto the manifold Γ , Φ ( X ) . Here Φ : RD → Γ is some function to be specified and hopefully we can simplify the dynamics Equation ( 3 ) via choosing suitable Φ . To track the dynamics of Φ ( Xη ) , we apply Ito ’ s lemma ( a.k.a . stochastic chain rule , see Lemma A.8 ) to Equation ( 3 ) , which yields dΦ ( Xη ( t ) ) = −η−1∂Φ ( Xη ( t ) ) ∇L ( Xη ( t ) ) dt+ ∂Φ ( Xη ( t ) ) σ ( Xη ( t ) ) dW ( t ) + 1 2 ∑D i , j=1 ∂ijΦ ( Xη ( t ) ) ( σ ( Xη ( t ) ) σ ( Xη ( t ) ) > ) ijdt . Note the first term −η−1∂Φ ( Xη ) ∇L ( Xη ) is going to diverge to∞ when η → 0 , so a natural choice for Φ is to kill the first term . Further note −∂Φ ( X ) ∇L ( X ) is indeed the directional derivative of Φ at X towards −∇L , killing the first term becomes equivalent to making Φ invariant under Gradient Flow of −∇L ( X ) ! Thus it suffices to take Φ ( X ) to be the limit of gradient flow starting at X . ( Formally defined in Section 3 ; see Lemma C.2 for a proof of ∂Φ ( X ) ∇L ( X ) ≡ 0 . ) Also intuitively Xη will be infinitely close to Γ , i.e. , d ( Xη ( t ) , Γ ) → 0 for any t > 0 as η → 0 , so we have Φ ( Xη ) ≈ Xη . Thus we can rewrite the above equation as dXη ( t ) ≈ ∂Φ ( Xη ( t ) ) σ ( Xη ( t ) ) dW ( t ) + 1 2 ∑D i , j=1 ∂ijΦ ( Xη ( t ) ) ( σ ( Xη ( t ) ) σ ( Xη ( t ) ) > ) ijdt , ( 4 ) and the solution of ( 4 ) shall converge to that of the following ( in an intuitive sense ) : dX ( t ) = ∂Φ ( X ( t ) ) σ ( X ( t ) ) dW ( t ) + 1 2 ∑D i , j=1 ∂ijΦ ( X ( t ) ) ( σ ( X ( t ) ) σ ( X ( t ) ) > ) ijdt , ( 5 ) The above argument for SDE was first formalized and rigorously proved by Katzenberger ( 1991 ) . It included an extension of the analysis to the case of asymptotic continuous dynamics ( Theorem 4.1 ) including SGD with infinitesimal LR , but the result is weaker in this case and no convergence is shown . Another obstacle for applying this analysis is that 2nd order partial derivatives of Φ are unknown . We solve these issues in Section 4 and our main result Theorem 4.6 gives a clean and complete characterization for the implicit bias of SGD with infinitesimal LR in Θ ( η−2 ) steps . Finally , our Corollary 5.2 shows ( 5 ) gives exactly the same regularization as tr [ ∇2L ] for label noise SGD . The main contributions of this paper are summarized as follows . 1 . In Section 4 , we propose a mathematical framework to study the implicit bias of SGD with infinitesimal LR . Our main theorem ( Theorem 4.6 ) gives the limiting diffusion of SGD with LR η for Θ ( η−2 ) steps as η → 0 and allows any covariance structure . 2 . In Section 5 , we give limiting dynamics of SGD with isotropic noise and label noise . 3 . In Section 6 , we show for any initialization , SGD with label noise achieves O ( κ ln d ) sample complexity for learning a κ-sparse overparametrized linear model ( Woodworth et al. , 2020 ) . In this case , the implicit regularizer is a data-dependent weighted ` 1 regularizer , meaning noise can help reduce the norm and even escape the kernel regime . The O ( κ ln d ) rate is minimax optimal ( Raskutti et al. , 2012 ) and improves over Õ ( κ2 ) upper bound by HaoChen et al . ( 2020 ) . In contrast , vanilla GD requires Ω ( d ) samples to generalize in the kernel regime . For technical contributions , we rigorously prove the convergence of GF for OLM ( Lemma 6.3 ) , unlike many existing implicit bias analyses which have to assume the convergence . We also prove the convergence of limiting flow to the global minimizer of the regularizer ( Lemma 6.5 ) by a trajectory analysis via our framework . It can not be proved by previous results ( Blanc et al. , 2020 ; Damian et al. , 2021 ) , as they only assert convergence to stationary point in the best case .

Summary: This paper proposes a new mathematical framework for explaining implicit regularization of SGD after attaining zero training loss. Technical novelty includes a different time scaling of the SGD that leads to the Katzenberger process, and then applying weak convergence of the Katzenberger process to obtain an elegant characterization of the SGD dynamics on zero loss manifold. This new technical device leads to a new proof of the sample complexity bound on the overparametrized linear model with label noise regularization.