When does SGD favor ﬂat minima? A quantitative characterization via linear stability
s
AfﬁliationAddress
Abstract
The observation that stochastic gradient descent (SGD) favors ﬂat minima hasplayed a fundamental role in understanding implicit regularization of SGD andguiding the tuning of hyperparameters.
In this paper, we provide an explanation ofthis striking phenomenon by relating the particular noise structure of SGD to its linear stability (Wu et al., 2018).
Speciﬁcally, we consider training over-parameterizedmodels with square loss.
We prove that if a global minimum θ ∗ is linearly stable for SGD, then it must satisfy (cid:107) H ( θ ∗ ) (cid:107) F (cid:54) O ( √ B/η ) , where (cid:107) H ( θ ∗ ) (cid:107) F , B, η denotethe Frobenius norm of Hessian at θ ∗ , batch size, and learning rate, respectively.
Otherwise, SGD will escape from that minimum exponentially fast.
Hence, forminima accessible to SGD, the ﬂatness—as measured by the Frobenius norm of the Hessian—is bounded independently of the model size and sample size.
The key toobtaining these results is exploiting the geometry awareness of SGD noise: 1) thenoise magnitude is proportional to loss value; 2) the noise directions concentratein the sharp directions of local landscape.
This property of SGD noise provablyholds for linear networks and random feature models (RFMs) and is empiricallyveriﬁed for nonlinear networks.
Moreover, the validity and practical relevance ofour theoretical ﬁndings are justiﬁed by extensive numerical experiments.
Introduction
Modern machine learning (ML) models are often operated with far more unknown parameters thantraining examples, a regime referred to as over-parameterization.
In this regime, there are manyglobal minima, all of which have zero training loss but their test performance can be signiﬁcantlydifferent [44].
Fortunately, it is often observed that SGD converges to those generalizable ones, evenwithout needing any explicit regularizations [46]. This suggests there must exist certain “implicitregularization” [30] mechanism at work.
More mysteriously, SGD solutions often generalize better than gradient descent (GD) solutions[19].
Therefore, the SGD noise must play a critical role in implicit regularization.
The mostpopular explanation is that SGD favors ﬂatter minima [19] and ﬂatter minima generalize better[14].
This ﬂat-minima principle has been extensively and successfully adopted in practice to tunethe hyperparameters of SGD [41, 19] and to design new optimizers [15, 10, 40] for improvinggeneralization.
Therefore, understanding how SGD noise biases SGD towards ﬂatter minima is ofparamount importance, which is the main focus of this paper.
The works [41, 9, 45] show that SGD noise is highly anisotropic; [28, 38]
ﬁnd the magnitude of SGD noise to be loss dependent.
Both structures are shown to be critical for SGD picking ﬂatminima.
However, these works all make unrealistic (even wrong) over-simpliﬁcations of SGD noise(see the related work section for more details) in their analysis.
In addition, instead of studying SGD, they all consider the continuous-time stochastic differential equation (SDE), which is a good modeling of SGD only in ﬁnite time and when the learning rate (LR) is sufﬁciently small [23].
Itis generally unclear how much the SDE modeling is relevant for understanding SGD with a large LR—a regime preferred in practice.
Consequently, these works only provide intuitive and empiricalanalyses, lacking a quantitative characterization of when and how SGD favors ﬂat minima.
Another line of works [43, 27] relate the selection bias of SGD to the dynamical stability .
In overparameterized case, all global minima are ﬁxed points of SGD but their dynamical stabilities can bevery different.
At unstable minima, a small perturbation will drive SGD to leave away, whereas, forstable minima, SGD can stay around and even converge back after initial perturbations.
Thus SGDprefers stable minima over unstable ones.
Speciﬁcally, [43, 27] analyze the linear stability [1] of SGD, showing that a linearly stable minimum should be ﬂat and uniform.
Different from SDE-basedanalysis, this stability-based analysis is relevant for large-LR SGD and is even empirically accuratein predicting the properties of minima selected by SGD [43, 18, 7].
In this work, we follow the linear stability analysis in [43, 27] but take the particular geometry-awarestructure of SGD noise into consideration.
We establish a direct connection between the linearstability and the ﬂatness, which allows us to obtain a quantitative characterization of the ﬂatness ofminima accessible to SGD.
In contrast, [43, 27] have to introduce the another quantity: non-uniformitytogether with ﬂatness to characterize the linear stability because of neglecting the noise structure.
Setup.
Let { ( x i , y i ) } ni =1 with x i ∈ R d , y i ∈
R be the training set and f ( · ; θ ) with θ ∈ R p be 54 our model.
Here p denotes the model size.
We focus on training the model with square loss .
Let
L i ( θ ) = 12 | f ( x i ; θ ) − y i | 2 be the ﬁtting error at the i -th sample and L ( θ ) =
1 n (cid:80) ni
=
L i ( θ ) be the 56 empirical risk.
To minimize L ( · ) , we consider the mini-batch SGD:
[Equation]
where η and B are the learning rate and batch size, respectively.
This SGD can be rewritten asθ t +
=
θ t − η ( ∇ L ( θ t ) + ξ t ) , where the noise ξ t satisﬁes E [ ξ t ] =
and E [ ξ t ξ Tt ]
= Σ( θ t ) /Bwith
Σ( θ ) =
1 n (cid:80) ni =1 ∇ L
i ( θ ) ∇ L i ( θ ) T − ∇ L ( θ ) L ( θ ) T .
To characterize the local geometry of losslandscape, we use the Gram matrix: G ( θ ) =
1 n (cid:80) ni
=1 ∇ f ( x i ; θ ) ∇ f ( x i ; θ ) T and the Hessian matrix:
H ( θ ) =
G ( θ ) + 1 n (cid:80) ni
=
( f ( x i ; θ ) − y i ) ∇ 2 f ( x i ; θ ) .
We focus on the over-parameterized regime, forwhich if θ ∗ is a global minimum, L ( θ ∗ )
=
and H ( θ ∗ )
= G ( θ ∗ ) .
In addition, throughout this paper,we denote by { λ 1 ( A ) } i (cid:62) 1 the eigenvalues of matrix A in the non-decreasing order.
Our contributions are summarized as follows.
• We ﬁrst show that for many ML models, the SGD noise is geometry aware: 1) the noisemagnitude is proportional to the loss value; 2) the noise covariance aligns with the Grammatrix.
Hence SGD noise is non-degenerate and concentrates along the sharp directions oflocal landscape.
Speciﬁcally, a loss-scaled alignment factor µ ( θ ) is deﬁned to quantify thealignment strength.
We prove that there exists a size-independent positive constant µ 0 suchthat µ ( θ ) (cid:62) µ 0 for linear networks (Theorem 2.1) and RFMs (Proposition 2.4). • We then turn to analyze the linear stability of SGD with noise satisfying the above alignment property.
We prove in Theorem 3.3 that if a global minimum θ ∗ is linearly stable, then(cid:107)
[Equation] (cid:112) B/µ 0 .
Here the constant µ 0 quantiﬁes the alignment strength of SGD noise.
Hence, for minima accessible to SGD, the Hessian’s Frobenius norm—the ﬂatnessperceived by SGD—is bounded independently of the model size and sample size.
On the
other hand, if a minimum is too sharp, violating the preceding stability condition, SGD
will escape from it exponentially fast (Theorem 3.6).
Together, we obtain a quantitative characterization of when and how much SGD likes ﬂat minima.
• We provide extensive numerical experiments to justify the validity of our theoretical ﬁndings for a variety of nonlinear models.
In particular, the practical relevance is demonstrated in Section 4 by considering the classiﬁcation of full CIFAR-10 dataset with VGG networksand ResNets.
1.1 Related work
Noise structures.
[49, 17, 25] consider the Hessian-based approximation: Σ( θ )
≈
σ 2 H ( θ ) , whereσ is a small constant.
[50] proposes an improved version: Σ( θ )
≈ 2 L ( θ ) H ( θ ) .
But these approximations in general cannot be accurate since Hessian is not semi-positive deﬁnite (SPD) in non-convex region.
More recently, [28] and [38] study SGD by assuming Σ( θ ) =
2 L ( θ ) H ( θ ∗ ) , where θ ∗ is aminimum of interest, and Σ( θ )
=
σ 2 L ( θ ) I p , respectively.
These assumptions completely ignorethe state-dependence of noise direction.
In contrast, we only assume Σ( θ ) =
2 L ( θ ) C ( θ ) with C ( θ ) having a nontrivial alignment with the Gram matrix G ( θ ) , which does not rely on any explicitstructural assumption on C ( θ ) .
As a result, our assumption is much weaker and can be rigorouslyjustiﬁed for many ML models both theoretically and empirically.
More importantly, we show thatthis weak alignment property is sufﬁcient for studying the linear stability of SGD.
We anticipate thatour alignment assumption can be also adopted to analyze other properties of SGD.
Escape from sharp minima.
The escape behavior of SGD was ﬁrst studied in [49, 43], as anindicator of how much SGD dislikes sharp minima.
One of the most mysterious observation is thatthe escape happens in an unreasonably efﬁcient way.
However, the theoretical analysis there assumesthe noise to be state-independent, and consequently, the derived escape time depends polynomiallyon the loss barrier.
Later [45, 28] attempt to study this issue using the classical diffusion-basedframework [11].
However, the obtained escape rates depend on the loss barrier exponentially, whichcannot explain the unreasonable escape efﬁciency at all.
See also [28] for an improved analysis.
[35]argues that the SGD noise is heavy-tailed and thus SGD should be modeled as Lévy-SDE instead of Itô-SDE.
Moreover, [35, 48] shows that the heavy-tailedness can ensure the escape rate depends onthe basin volume instead of the loss barrier.
However, the volume in high dimensions always scalesexponentially with the dimension.
Moreover, whether SGD noise is really heavy-tailed and whetherthe heavy-tailedness is really important for generalization are still debatable for neural networks[37, 24].
In contrast, we show that the unreasonable escape efﬁciency comes from the particulargeometry-aware structure of SGD noise, regardless of whether the noise is heavy- or light-tailed.
Notions of ﬂatness.
In the literature, a variety of ﬂatness notions have been proposed, such as thelargest eigenvalue of Hessian [19], the trace of Hessian [8, 5], the basin volume [48], and the onesscaled by parameter norms [26, 36] in order to achive the scaling-invariance for ReLU networks.
These ﬂatnesses are proposed for either computation easiness or bounding generalization gaps. However, it is unclear if these ﬂatnesses are perceivable to SGD, let alone how the boundedness ofthem depends on the batch size, learning rate, as well as the model size and sample size.
We showthat for SGD solutions, the Frobenius norm of Hessian—a ﬂatness perceived by SGD through thelinear stability—is bounded by a size-independent quantity.
Note that a similar stability argumentalso applys to GD, but only yielding the boundedness of the largest eigenvalue of Hessian [43, 7].
Lastly, we particularly mention the work [31], which provides a rigorous analysis of the implicitregularization of training two-layer diagonal linear networks.
This work is related to ours since weboth consider the magnitude and direction structure of SGD noise simultaneously.
However, theanalysis in [31] is limited to the speciﬁc toy model but ours is relevant for general models.
The geometry-aware structure of SGD noise
To gain some intuition, we ﬁrst consider the decoupling approximation [28]:
[Equation]
The second step assumes that the inﬂuence of ∇ L ( θ ) can be neglected; the third step assumes that{∇ f ( x i ; θ ) } i and { L i ( θ ) } i are nearly decoupled.
This approximation is very crude (cannot be 126 accurate in general) but tells us two critical properties of SGD noise: 1) The noise magnitude isproportional to the loss value; 2) the noise covariance aligns with the Gram matrix.
To utilize the above structure in a rigorous and quantitative way, we deﬁne
[Equation]
Here α ( θ ) quantiﬁes the strength of standard alignment between the noise covariance and Grammatrix, where the latter characterizes the local geometry of landscape.
The more α ( θ ) is close to 1 ,the more the noise concentrates in sharp directions of local landscape.
β ( θ ) characterizes the relative magnitude of noise with respect to the loss value.
When β ( θ ) is bounded below, the inﬂuence ofnoise is non-degenerate even if its absolute magnitude is small.
Lastly, µ ( θ ) =
α ( θ ) β
( θ ) can beviewed as a loss-scaled alignment factor, combining the direction and magnitude structure of noisetogether.
When the decoupling approximation holds, we have α ( θ ) =
β ( θ )
= µ ( θ )
= 1 .
The speciﬁc form of µ ( θ ) deﬁned above is inspired by our linear stability analysis in Section 3.
Inthis paper, the noise is said to satisfy the ( µ -) alignment if there exist a positive constant µ 0 such that
[Equation]
This alignment property ensures that in sharp directions, SGD noise is non-degenerate.
A relaxed alignment.
Let Σ 1 ( θ ) = 1 n [Equation] =
[Equation]
∇ L
[Equation]
θ ) T .
ThenΣ( θ ) =
Σ 1 ( θ ) − Σ 2 ( θ ) .
It is often believed that during the late phase of training, the full-batchgradient ∇ L is relatively small compared to the sample gradients {∇ L i } i .
As a result, the inﬂuence 142 of Σ 2 ( θ ) should be negligible compared to Σ 1 ( θ ) .
To disentangle the inﬂuences of them, we deﬁne
[Equation]
Then µ ( θ ) = µ 1 ( θ ) − µ 2 ( θ ) .
Our linear stability analysis in Section 3 show that µ 1 ( θ ) (cid:62) µ 1 > 0 , acondition we refer to as µ 1 -alignment , is sufﬁcient to ensure that SGD only selects
ﬂat minima.
2.1 Over-parameterized linear models
Consider an over-parameterized linear model (OLM): f ( x ; θ ) =
F ( θ ) T x , where F : Ω (cid:55)→ R d denotesa general re-parameterization function.
f ( · ; θ ) only represents linear functions but the correspondinglandscape can be highly non-convex.
Typical examples include the (deep) linear network: F ( θ ) =
W 1 W 2 · · · W L and the quadratically parameterized linear model: F ( θ ) =
( α 21 − β 21 , . . .
, α 2 d − β 2 d ) T .
The latter has attracted a lot of attention recently in analyzing the implicit regularization of GD and
SGD [39, 31, 12, 3].
The following theorem provides a precise characterization of (online)-SGDnoise for OLM models, whose proof is deferred to Appendix C.
Theorem 2.1.
Denote by N (0 , S ) the Gaussian distribution with mean zero and covariance matrix S .
Suppose f ( · ; θ ) is a general OLM and x ∼ N (0 , S ) .
Consider the online SGD setting, i.e., n = ∞ .
Then, Σ(
[Equation]
This theorem shows that the alignments hold in the entire parameter space.
Moreover, the alignmentstrengths are independent of model size.
More importantly, this theorem provides a closed-formexpression of the noise covariance, which might be useful for analyzing other properties of SGDbeyond the linear stability considered in Section 3.
It should be stressed that the alignment isonly proved for the inﬁnite-sample case.
Similar results should hold for ﬁnite-sample cases byconcentration inequalities as long as n is relatively large but this straightforward extension does notbring new insights.
It is more interesting to consider the low-sample regime (i.e., n < d ) and Figure 1in Section 2 shows that the alignments indeed hold in this regime, at least in typical regions exploredby SGD.
We leave the comprehensive analysis of this regime to further work.
2.2 Feature-based models
We now turn to the feature-based model f ( x ; θ ) =
[Equation]
ϕ
[Equation]
[Equation] θ, Φ( x ) (cid:105) .
Here Φ( x ) ∈ R mdenotes the feature of x .
In this case, the Hessian and Gram matrix are both constant but the SGDnoise is still state-dependent.
Let g i =
Φ( x i ) be the sample feature of x i .
Then,
[Equation]
Our aim is to identify conditions on the features { g i } ni =1 such that the alignments hold for the noise.
Remark 2.2.
Although feature-based models are linear, the analysis provided below is also relevantfor understanding nonlinear models.
For instance, wide neural networks trained by SGD areclose to RFMs [16].
A general nonlinear model may locally behave like the linearized model:
f lin ( x ; θ ) : = f
( x ; θ ∗ ) + (cid:104) θ − θ ∗ , ∇ f ( x ; θ ∗ ) (cid:105) with the tangent feature ∇ f ( x ; θ ∗ ) learned from data.
174 Hence, the analysis here at least can explain the alignment property of SGD noise in a local region,which is sufﬁcient for characterizing the linear stability of SGD (see Section 3).
Assumption 1.
Let χ i
= g Ti Gg
i , ¯ χ = 1 n (cid:80) ni =1 χ i .
Assume γ :=
min i ∈ [ n ] χ i / ¯ χ > 0 .
This assumption means that the feature norms (weighted by G ) are uniformly bounded below, whichcan ensure that the corresponding SGD noise satisﬁes the alignments (at least the µ 1 -alignment).
Lemma 2.3.
Suppose Assumption 1 holds.
Let λ 1 ( G ) (cid:62) λ 2 ( G ) (cid:62) · · · (cid:62) λ n ( G ) be the eigenvalues 180 of
G .
Then µ 1 ( θ ) (cid:62) γ, µ 2 ( θ ) (cid:54) τ ( G ) :
= λ 21 ( G ) / (cid:80) j λ 2 j ( G ) , and µ ( θ ) (cid:62) γ − τ ( G ) .
Proof.
Noticing ¯ χ = 1 n (cid:80) ni
{+=
 g Ti Gg i+}
{+= (cid:107) G (cid:107) 2 F , then Tr(Σ 1 ( θ ) G ) = 1 n (cid:80) ni =1 | g Ti θ | 
 Tr( g i g Ti G ) =+}
[Equation]
ni
[Equation]
θ ) (cid:107) G (cid:107) 2 F .
Thus µ 1 ( θ ) (cid:62) γ .
In addition, µ 2 ( θ ) = [Equation] =
[Equation]
2 F , where the last step follows from Lemma E.1.
We prove below that the uniform boundedness of feature norms can be rigorously veriﬁed for the
ReLU random feature model.
As a result, the µ 1 - and µ -alignment provably hold with a highprobability as stated in the following proposition.
The proof is presented in Appendix D.3, wheresimilar results for general RFMs are provided (see Proposition D.7).
Proposition 2.4.
Let ϕ j
( x ) = ReLU( w Tj x ) with w j iid ∼ Unif( √ d S d − 1 ) and x ∼ Unif( S d − 1 ) .
Forany δ ∈ (0 , 1) , assume n (cid:38) d 5 log(1 /δ )
, then w.p. at least 1 − δ , µ 1 ( θ ) (cid:62) 1 and µ ( θ ) (cid:38) d − 1 .
2.3 Empirical validations
Reported in Figure 1a are the values of α ( θ t ) , β ( θ t ) , µ ( θ t ) during the SGD training of four types ofmodels, including the RFM and linear networks analyzed above.
First, one can see that t α ( θ t ) ’sare quite close to 1 during the whole training, suggesting the concentration of SGD noise on sharpdirections of local landscape.
Second, β ( θ t ) ’s are bounded below, implying the noise magnitudesare sufﬁciently large with respect to the training loss.
As a result, we see clearly that the loss-scaledalignment factor µ ( θ t ) is signiﬁcantly positive for all models examined.
Figure 1b further examines how the extent of over-paramaterization affects the alignment strength.
One can see clearly that for linear networks and RFMs, µ ( θ ) ’s are independent of the model size,which conﬁrms our theoretical results proved above.
In addition, we also observe that for nonlinearnetworks the alignment strength is also (nearly) independent of the model size.
[Figure]
In Figure
[Figure]
we plot the (relative) feature norms of three nonlinear models.
Here we compute χ i withg i = ∇ f
( x i ; θ ) , which corresponds to the features of linearized model f lin ( · ; θ ) at θ .
It is shownthat Assumption 1, i.e., the uniform boundedness of feature norms, holds during the SGD training.
According to Lemma 2.3, this implies that the SGD noise of the local linearized model satisﬁes theµ 1 -alignment, which is sufﬁcient for the linear stability analysis in Section 3.
[Figure]
We proved that the alignment of SGD noise holds in the whole parameter space for the simpliﬁedmodels, but this does not necessarily generalize to arbitrarily models.
The experiments provided aboveonly suggest that for nonlinear networks, the alignment property is satisﬁed in typical regions exploredby SGD, including the random initialization and the convergent solutions.
Similar observations forlarger-scale models can be found in Section 4.
The linear stability analysis
Let θ ∗ be a global minimum of L ( · ) .
When θ t is close to θ ∗ , the local dynamical behavior of SGDcan be characterized by linearizing the dynamics around θ ∗ :
[Equation]
(5) where ∇ 2 L i ( θ ∗ ) =
∇ f
( x i ; θ ∗ ) ∇ f ( x i ; θ ∗ ) T .
This corresponds to the local quadratic approximationof L ( · ) or the local linearization of the model around θ ∗ :
[Equation]
where ∇ f ( x ; θ ∗ ) is the tangent feature [16].
Note that (5) is exactly the SGD of training (6).
Deﬁnition 3.1 (Linear stability) .
A global minimum θ ∗ is said to be linearly stable if there exists a
C > 0 such that it holds for the linearized dynamics (5) that E [ L (˜ θ t )
] (cid:54) C E
[ L (˜ θ 0 )] , ∀ t (cid:62) 0 .
If the local quadratic approximation is non-degenerate around θ ∗ , then θ t is close to ˜ θ t when θ t isclose to θ ∗ .
In other words, the local behavior of the original SGD around θ ∗ can be characterizedby the linearized SGD.
However, in the over-parameterized case, the local quadratic approximationis degenerate in ﬂat directions.
Consequently, one may be concerned about the relevance of linearstability.
Fortunately, the stability is measured with the change of loss value.
Thus, the instabilitymostly comes from noise perturbations in sharp directions and the noise also concentrates in thesedirections (guaranteed by the alignment property).
In sharp directions, the local quadratic approximation is always validated.
This explains why linear stability is still relevant in over-parameterized case.
The rigorous formulation of this intuition is left to future work.
In current paper, we instead resort tonumerical experiments to demonstrate the validity.
For simplicity, we will use θ t to denote ˜ θ t ; let θ ∗ = 0 and g i = ∇ f ( x i ; θ ∗ ) .
For the linearized modelf lin ( · ; θ ) , we have L ( θ ) = 12 n (cid:80)
[Equation]
ni =
[Equation]
g
[Equation]
g Ti , where we omit 231 the dependence on θ ∗ for simplicity.
Note that the Gram and Hessian matrix are constant but thenoise covariance Σ( θ )
= 1 n (cid:80) ni
=1 | g Ti θ |
g
i g Ti − Hθθ T H is still state-dependent.
Before considering the speciﬁc linearized SGD (5), we ﬁrst have a general result.
Lemma 3.2.
Consider a general SGD: θ t +1 =
θ t − η ( ∇ L ( θ t ) + ξ t ) for the linearized model (6) , 235 where ( ξ t ) t (cid:62) 1 are any noises satisfying E [ ξ t ] = 0 , E [ ξ t ξ Tt ]
= S ( θ t ) .
Then we have
[Equation]
7) where ν ( θ ) =
Tr( HS ( θ )) / 2 and r ( θ ) (cid:62) 0 .
Moreover, if η (cid:54) 2 /λ 1 ( H ) , then r ( θ ) (cid:54) 1 .
Proof.
The update of expected loss can be written as
[Equation]
where we use E [ ξ t ] = 0 and E [ ξ t ξ Tt ]
= S ( θ t ) .
Substituting ∇ L ( θ ) =
Hθ yields r ( θ ) = 1 −2 η θ T H
2 θ θ T Hθ
+ η 2 θ T H
3 θ θ T Hθ .
By Lemma E.2, r ( θ ) (cid:62)
and if η (cid:54) 2 /λ 1 ( H ) , then r ( θ ) (cid:54) 1 .
The two terms r ( θ t ) L ( θ t ) and η 2 v ( θ t ) denote the contributions from the full-batch gradient ∇ L ( θ t )and the noise ξ t , respectively.
The stability is affected by both terms simultaneously.
It is well-knownthat if θ ∗ is linearly stable for GD, then λ 1 ( H ( θ ∗ ))
(cid:54) 2 /η
(see, e.g., [43, 29]).
This also holds for
SGD but SGD imposes a stricter condition because of the extra η 2 ν ( θ t ) term.
Speciﬁcally,
[Equation]
It is implied that the more S ( θ ) aligns with H , the more unstable that minimum is .
3.1 The linear stability imposes size-independent ﬂatness constraints
For the mini-batch SGD, the following theorem characterizes how the batch size and learning rateaffect the ﬂatness—as measured by the Frobenius norm of Hessian—of minima accessible to SGD.
Theorem 3.3.
Let θ ∗ be a global minimum that is linearly stable.
Denote by µ ( θ ) the alignmentfactors for the linearized SGD (5) and model (6) .
If µ ( θ ) (cid:62) µ 0 , then (cid:107) H ( θ ∗ ) (cid:107) F (cid:54) 1 η (cid:113)
[Equation]
Proof.
By (9) and the deﬁnition of µ ( θ ) , we have
[Equation]
To ensure the stability, we must have µ 0 η 2 (cid:107) H (cid:107) 2 F /B (cid:54) 1 , leading to (cid:107) H (cid:107) F (cid:54) (cid:112)
B/µ 0 /η .
We have shown in Section 2 that µ 0 is (nearly) size-independent, and thus the obtained upper bound ofﬂatness is also (nearly) size-independent.
As a comparison, for GD, the linear stability can only ensureλ 1 ( H ( θ ∗ ))
(cid:54) 2 /η .
This gives a bound of the Hessian’s Frobenius norm: (cid:107) H ( θ ∗ ) (cid:107) F (cid:54) 2 √ p/η ,depending on the model size explicitly.
The comparison of two bounds partially explains why SGDtends to select ﬂatter minima than GD.
We show below that the µ -alignment can be further relaxed to the µ 1 -alignment.
The proof is similarto the one of Theorem 3.3 and deferred to Appendix E.
Proposition 3.4.
Under the setting of Theorem 3.3, if the noise of linearized SGD satisﬁes µ 1 ( θ ) (cid:62) µ 1 ,then (cid:107) H ( θ ∗ ) (cid:107) F (cid:54) min
[Equation]
When B (cid:29) 1 , the bound becomes B/ ( η
[Equation]
) µ 1 ) ≈ (cid:112)
B/µ 1 /η , which is the same as the caseof µ -alignment.
Thus the inﬂuence of ∇ L is indeed negligible compared with {∇ L i } i .
Remark 3.5.
The linear stability is local in nature and hence our analysis essentially only needs thealignment properties of SGD noise to hold locally around minima of interest.
It worth noting that inthe derivation, we only consider the instability caused by the noise component, with the full-batchcomponent completely ignored.
To obtain a tigher bound of ﬂatness, one needs to consider bothcomponents simutaneously, which is much more complicated and left to future work.
Numerical validations.
Figure 3 numerically shows that the actual ﬂatness (not only the upperbound) is also (nearly) independent of the extent of over-parameterization.
For instance, for the FCNs, the ﬂatness only oscillates around 2 . 4 when increasing the model size.
For the CNNs, one cansee that the ﬂatness only increases roughly from 7 . 6 to 8 . 6 as increasing the model size from 10 4 to10 6 , an increase of two orders of magnitude.
Similar results for larger-scale models can be found in
Section 4.
These further justify our ﬁndings: the learning rate, batch size, and the alignment strengthtogether control the ﬂatness of minima reachable to SGD, regardless of the model size.
3.2 SGD escapes from sharp minima exponentially fast
The following theorem shows that the pure noise-driven escape from a sharp minimum is exponentiallyfast , whose proof follows trivially from (10).
Theorem 3.6.
Under the setting of Theorem 3.3, if (cid:107) H ( θ ∗ ) (cid:107) F > 1 η (cid:113)
Bµ 0 , then the linearized SGDsatisﬁes E [ L ( θ t )] (cid:62) γ t 0 E [ L ( θ 0 )] with γ 0 =
η 2 µ 0 B (cid:107) H ( θ ∗ ) (cid:107) 2 F > 1 .
0 log 10 (model size) 2 .2 .2 .2 .k
H k F
Linear net 4.
6.0 log 10 (modelsize) 1.8 2.
k H k F
FCN 4.0 4.5 5.0 5.5 6.0 log 10 (modelsize) 7.2 7.
k H k F
CNN
Hence, linearized SGD takes roughly log γ 0 (1 /ε ) steps to escape from a O ( ε ) -loss region to a O (1) -loss region.
The escape time depends on the loss barrier only logarithmically and is independent ofthe parameter space dimension.
Due to the local closeness between linearized SGD and the original
SGD, this partically explains the unreasonable escape efﬁciency of SGD for training big models.
Incontrast, the escape rates of existing works [45, 49, 48, 28] are either exponentially slow with respectto the loss barrier or suffer from the curse of dimensionality.
Figure 4 shows the trajectories of SGD escaping from sharp minima.
It is demonstrated that the escapeis indeed exponentially fast and speciﬁcally, 10 steps are enough for SGD escaping to a high-lossregion for all the models examined.
In addition, we observe that the escape is still exponentially fastin the high-loss region, although our analysis only applies to a local region.
How can we explain thisnonlocal escape behavior?
We leave the study of this interesting phenomenon to future work.
[Figure]
[Figure]
[Figure]
( θ [Figure] [Figure] θ
[Figure]
[Figure]
[Figure] [Figure] − [Figure]
[-8 Step: t 10 −10 −
-]− [Figure]
−
[Figure] ( θ t )
[-8 Step: t 10 −10 −10 −
-][Figure]
[Figure] ( θ
[Figure] [Figure]
[-8 Step: t 10 −10 −10 −
 L o ss : L ( θ t ) CNN ˆ E-]−
[Figure]
3.3 The importance of the noise structure
The magnitude structure.
Theorem 3.6 together with its proof suggests that the loss dependenceof noise magnitude is critical for obtaining the exponentially fast escape.
The intuition is as follows.
When θ t is perturbed by noise to θ t +1 where L ( θ t +1 ) >
L ( θ t ) , the noise magnitude becomes largerthere and thus θ t +1 is easier to be perturbed to a larger-loss region.
This positive feedback drives
SGD to leave exponentially fast.
On the contrary, the following lemma shows that if the noise isuniformly bounded, the noise-driven escape is at most linear in time.
Lemma 3.7.
Under the setting of Lemma 3.2, assume η (cid:54) 2 /λ 1 ( H ) and E [ HS ( θ )]
(cid:54) 2 σ 2 .
Then 299 E [ L ( θ t ) − L ( θ 0 )] (cid:54) η 2 σ 2 t.
Here we set η (cid:54) 2 /λ 1 ( H ) to avoid the exponential escape caused by the GD part.
Proof.
By Lemma E.2, when η (cid:54) 2 /λ 1 ( H ) , r ( θ ) (cid:54) 1 .
Thus Lemma 3.2 implies E [ L ( θ t +1 )] (cid:54) 302 E [ L ( θ t )]
+ η 2 σ 2 , which implies E [ L ( θ t )]
(cid:54) E
[ L ( θ 0 )]
+ η 2 σ 2 t .
The direction structure.
We now turn to consider the impact of direction structure.
Considergeneral SGDs: θ t +1 =
θ t − η ( ∇ L ( θ t ) + ξ t ) with E [ ξ t ξ t ]
= S ( θ t ) /B for the linearized model (6).
We compare two type of (unrealistic) noises:
• Geometry-aware noise: S 1 ( θ ) =
2 L ( θ ) H .
• Isotropic noise: S 2 ( θ ) = 2 σ 2 L ( θ )
I p with σ 2 =
Tr( H ) /p .
Here,
the value of σ 2 is chosen
to ensure two types of noises have the same total variance for a fair comparison [49].
Note that p denotes the model size.
Analogous to Theorem 3.3, for the second isotropic SGD,
[Equation]
Hence, the instability decreases with the parameter-space dimension and the resulting ﬂatnessconstraint is Tr( H ( θ ∗ ))
(cid:54) √ pB/η , depending on the model size explicitly.
In contrast, for the ﬁstnoise, Theorem 3.3 implies (cid:107) H ( θ ∗ ) (cid:107) F (cid:54) √ B/η , independent of the model size.
This difference canbe intuitively explained as follows.
The isotropic noise wastes most energy in perturbing SGD alongﬂat directions, which barely affects the instability.
In contrast, the geometry-aware noise focusesmost energy on perturbing SGD along sharp directions, causing much more instability.
Larger-scale experiments
We have provided small-scale experiments to justify the validity of our theoretical ﬁndings for avariety of nonlinear models.
Here we turn to demonstrate the practical relevance by classifyingthe CIFAR-10 dataset [20] with VGG networks [34] and ResNets [13].
In training, all explicitregularizations are removed to keep consistent with our theoretical analysis.
More details of theexperimental setup can be found in Appendix A.1.
Figure 5a shows that the alignment factors are signiﬁcantly positive during the training.
We againobserve that the alignment strength is (nearly) independent of the model size.
In addition, Figure 9 inthe appendix shows that the escape from sharp minimum is exponential for both ResNets and VGGnetworks.
These further conﬁrm our theoretical ﬁndings.
Figure 5b reports the ﬂatness of convergent solution and the corresponding theoretical upper bound ofdifferent models.
We again observe that the ﬂatness is nearly independent of the model size for both ResNets and VGG networks.
Moreover, we surprisingly discover that the upper bound obtained fromlinear stability analysis is rather tight in predicting the ﬂatness of convergent solution.
This tightnessseems to suggest that SGD runs nearly at the edge of stability (nEoS) [43, 7].
In addition, Figure5c shows that the ﬂatness bound is much tigher in the small-batch regime.
Note that our bound isobtained by inspecting only the noise-driven instability, with the instability caused by the full-batchcomponent completely ignored. When decreasing the batch size, the noise component becomes largerand the full-batch component keeps unchanged. Hence, in the small-batch regime, it is not surprisingto observe the phenomenon in Figure 5c.
[Figure] [Figure] Lastly, we point out that the speciﬁc values of µ ( θ ) on the order of 0 .
01 for 337 CIFAR-10 experiments (see Figure 5a), clearly smaller than those in smallscale experiments (see Figure 1). In the right ﬁgure, we examine classifyingsubsets of CIFAR-10 with different number of classes. We can see a trendthat the alignments become stronger when decreasing the number of classes. More experiments with similar results can be found in the appendix.
Theseexperiments seem to suggest that the alignment strengths (at least in the SGD trajectory) depends moreon the intrinsic complexity of the problem instead of the model size.
A systematical investigationof how the alignment strengths depend on model size, sample size, and model architecture is animportant questions, which we leave to future work.
Conclusion
We provide a stability-based explanation of why SGD tends to select ﬂat minima.
Our currentunderstanding is as follows.
1) For popular ML models , the SGD noise is geometry
aware—non- degenerate
in the sharp directions of local landscape. This geometry awareness together with linear stability
ensures that the ﬂatness of minima
accessible to SGD is size -independent Moreover
,this understanding is made rigorous and quantitative by introducing a loss-scaled alignment factor tocharacterize the particular noise structure.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information onhow to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .
You are strongly encouraged to include a justiﬁcation to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] See Section ?? .
• Did you include the license to the code and datasets?
[No] The code and the data are proprietary.
• Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] Local analysis
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes] See Appendix A.
(b) Did you specify all the training details (e
g., data splits, hyperparameters,
how they
were chosen)?
[Yes] See Appendix A.
(c) Did you report error bars (e g., with respect to the random seed after running experi- ments multiple times)?
[Yes] Estimated over 5 runs
(d) Did you include the total amount of compute and the type of resources used (e g., typeof GPUs, internal cluster, or cloud provider)?
[No]
If you are using existing assets (e g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e ) Did you discuss whether the data you are using/curating contains personally identiﬁable
[Table]
information or offensive content?
[N/A]
[Table]
If you used crowdsourcing or conducted research with human subjects...
[Table]
(a) Did you include the full text of instructions given to participants and screenshots, if [Table] applicable?
[N/A]
[Table]
(b) Did you describe any potential participant risks, with links to Institutional Review
[Table]
Board (IRB) approvals, if applicable?
[N/A]
[Table]
(c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?
[N/A]
[Table]