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: The authors propose an SDE approximation to study the implicit bias of SGD with infinitesimal learning rate. They show that given any noise covariance structure, after $\eta^{-2}$ steps, the SGD converges to an SGD on a certain manifold of local minimizer as the learning rate $\eta\rightarrow 0$. This recovers and strengthens some existing results in the literature. In particular, they show a sample complexity gap between label noise SGD and GD in the kernel regime for an overparametrized linear model, justifying the generalization benefit of SGD.