When does SGD favor ﬂat minima? A quantitative characterization via linear stability
Anonymous Author s)
AfﬁliationAddress
Abstract
The observation that stochastic gradient descent (SGD) tends to select ﬂat minimahas played a fundamental role in understanding implicit regularization of SGDand guiding the tuning of hyperparameters.
In this paper, we provide a quantitative explanation of this phenomenon by relating the particular noise structure of SGD to its linear stability (Wu et al., 2018).
Speciﬁcally, we consider trainingover-parameterized models 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, η denote the Frobenius norm of Hessian at θ ∗ , batch size, andlearning rate, respectively.
Otherwise, SGD will escape from that minimum exponentially fast.
Hence, for minima accessible to SGD, the ﬂatness—as measuredby the Frobenius norm of the Hessian—is bounded independently of the modelsize and sample size.
The key to obtaining these results is exploiting the particulargeometry awareness of SGD noise: 1) the noise magnitude is proportional to lossvalue; 2) the noise directions concentrate in the sharp directions of local landscape.
This property of SGD noise provably holds for linear networks and random featuremodels (RFMs) and is empirically veriﬁed for nonlinear networks.
Moreover, thevalidity and practical relevance of our theoretical ﬁndings are justiﬁed by extensivenumerical experiments.
Introduction
Modern machine learning (ML) models are often operated with far more unknown parametersthan training examples, a regime referred to as over-parameterization.
In this regime, there is alarge set of global minima, all of which have zero training loss but the test performance can besigniﬁcantly different [45].
Fortunately, SGD always converges to those good ones even withoutneeding any explicit regularizations [47], suggesting there must exist certain “implicit regularization”[31] mechanism at work.
More mysteriously, SGD solutions often generalize better than gradient descent (GD) solutions [20].
Therefore, the SGD noise must play a critical role in implicit regularization.
The most popularexplanation is that SGD converges to ﬂatter minima [20] than GD and ﬂatter minima generalizebetter [14].
This ﬂat-minima principle has been extensively and successfully adopted in practice totune the hyperparameters of SGD [42, 20] and to design new optimizers [16, 10, 41] 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 [42, 9, 46] show that SGD noise is highly anisotropic; [29, 39]
ﬁnd the magnitude of SGDnoise to be loss dependent.
Both structures are shown to be critical for SGD picking ﬂat minima.
However, these works all make unrealistic (even wrong) over-simpliﬁcations of SGD noise (see therelated work section for more details) in their analysis.
Moreover, 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 [24].
It is generallyunclear how much the SDE modeling is relevant for understanding SGD with a large LR—a regimepreferred in practice.
Consequently, these works only provide intuitive and empirical analyses,lacking a quantitative characterization of when and how SGD favors ﬂat minima.
Another line of works [44, 28] relate the selection bias of SGD to the dynamical stability of globalminima.
In over-parameterized case, all global minima are ﬁxed points of SGD but their dynamicalstabilities can be very different.
At unstable minima, a small perturbation will drive SGD to leave away,whereas, for stable minima, SGD can stay around and even converge back after initial perturbations.
Thus SGD prefers stable minima over unstable ones.
Speciﬁcally, [44, 28] analyze the linear stability[1] of SGD, showing that a linearly stable minimum should be ﬂat and uniform.
Different fromthe SDE-based analysis, this stability analysis is relevant for large-LR SGD and is even empiricallyaccurate in predicting the properties of minima selected by SGD [44, 19, 7].
In this work, we follow the linear stability analysis in [44, 28] 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, [44, 28] 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 55 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 57 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 the alignment strength.
We prove that there exists a size-independent positive constant µ 0 suchthat µ ( θ ) (cid:62) µ 0 for RFMs (Proposition 2.4) and linear networks (Theorem 2.1).• We then turn to analyze the linear stability of SGD with noise satisfying the above alignmentproperty.
We prove in Theorem 3.3 that if a global minimum θ ∗ is linearly stable, then(cid:107)
H ( θ ∗ ) (cid:107) F (cid:54) η −(cid:112) B/µ 0 .
Here the constant µ 0 quantiﬁes the alignment strength of SGDnoise.
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 quantitativecharacterization of when and how much SGD likes ﬂat minima.
• We provide extensive numerical experiments to justify the validity of our theoretical ﬁndingsfor a variety of nonlinear models.
In particular, the practical relevance is demonstrated inSection 4 by considering the classiﬁcation of full CIFAR-10 dataset with VGG networksand ResNets.
1.1 Related work
Noise structures.
[50, 18, 26] consider the Hessian-based approximation: Σ( θ )
≈
σ 2 H ( θ ) , whereσ is a small constant.
[51] proposes an improved version: Σ( θ )
≈ 2 L ( θ ) H ( θ ) .
But these approxima- tions in general cannot be accurate since Hessian is not semi-positive deﬁnite (SPD) in non-convexregion.
More recently, [29] and [39] study SGD by assuming Σ( θ ) =
2 L ( θ ) H ( θ ∗ ) ( θ ∗ is a minimumof interest) and Σ( θ )
=
σ 2 L ( θ ) I p , respectively.
These assumptions completely ignore the statedependence of noise direction.
In contrast, we only assume Σ( θ ) =
2 L ( θ ) C ( θ ) with C ( θ ) having anontrivial alignment with the Gram matrix G ( θ ) . Note that we do not impose any explicit structuralassumption on C ( θ ) .
As a result, our assumption is much weaker and can be rigorously justiﬁed formany ML models.
More importantly, we also show that this weak alignment property is sufﬁcient forstudying the linear stability of SGD.
Escape from sharp minima.
This escape behavior of SGD was initiated studied in [50, 44], asan indicator of how much SGD dislikes sharp minima.
It was shown that the escape happens inan unreasonably efﬁcient way.
However, the theoretical analysis there assumes the noise to bestate-independent, and consequently, the derived escape time depends polynomially on the lossbarrier.
Later [46, 29] attempt to study this issue using the classical diffusion-based framework [11].
However, the obtained escape rates depend on the loss barrier exponentially, which cannot explain theunreasonable escape efﬁciency at all.
See also [29] for an improved analysis.
[36] argues that the SGDnoise is heavy-tailed and thus SGD should be modeled as Lévy-SDE instead of Itô-SDE.
Moreover,[36, 49] shows that the heavy-tailedness can ensure the escape rate depends on the basin volumeinstead of the loss barrier.
However, the volume in high dimensions always scales exponentially withthe dimension.
Moreover, whether SGD noise is really heavy-tailed and whether the heavy-tailednessis really important for generalization are still debatable for neural networks [38, 25].
In contrast, weshow that the unreasonable escape efﬁciency comes from the particular geometry-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 asthe largest eigenvalue of Hessian [20], the trace of Hessian [8, 5], and many other variants [27, 37]that incorporate norm information into deﬁnitions.
However, it is unclear if these ﬂatnesses areperceivable to SGD? In our understanding, the “ﬂatness” perceived by SGD is related to the stability,which further depends on the noise structure. Hence, different noises induce different ﬂatnesses.
Speciﬁcally, we show that for many popular ML models, due to the particular noise structure, SGDcan perceive the Frobenius norm of Hessian—a ﬂatness that characterizes the linear stability of SGD.
As a comparison, the ﬂatness perceived by GD is only the largest eigenvalue of Hessian [44, 7].
Lastly, we particularly mention the work [32], 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 [32] 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 [29]:
[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 125 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 relativemagnitude of noise with respect to the loss value.
When β ( θ ) is bounded below, the inﬂuence of noise 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 (cid:80) ni
[Equation]
L i
[Equation]
θ ) T , Σ 2 ( θ ) = ∇ L ( θ ) ∇ L ( θ ) 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 141 of Σ 2 ( θ ) should be negligible compared to Σ 1 ( θ ) .
To disentangle the inﬂuences of them, we deﬁneµ
[Equation]
= [Equation] [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
[Figure]
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 [40, 32, 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]
( x ) = (cid:104) θ, Φ( 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 [17].
A general nonlinear model may locally behave like the linearized model:
[Equation]
with the tangent feature ∇ f ( x ; θ ∗ ) learned from data.
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 179 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 [Equation]
[-=
 g Ti Gg i-][Equation]
[-= (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 ) =-]Tr
[Equation]
G
[Equation]
F .
Thus µ 1 ( θ ) (cid:62) γ .
In addition, µ 2 ( θ ) =
[Equation]
, 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]
[Equation]
where ∇ f ( x ; θ ∗ ) is the tangent feature [17].
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 230 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) , 234 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., [44, 30]).
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 η
[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 µ 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]
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.
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 . when increasing the model size.
For the CNNs, one can seethat 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 complexities of solutions 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 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.
k H k F
FCN 4. 6.0 log 10 (modelsize) 7.
k H k F
CNN
Hence, 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 of the dimensionof parameter space.
This explains the unreasonable escape efﬁciency of SGD for training big models.
In contrast, the escape rates of existing works [46, 50, 49, 29] are either exponentially slow withrespect to 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]
L ( θ t )
RFM ˆ E [ L ( θ t )]
{+8 Step: t 10 −10 −
+}
L o ss : L ( θ t )
Linear net
ˆ E [ L ( θ t )]
{+8 Step: t 10 −10 −10 −
+}
L o ss : L ( θ t )
FCN ˆ E [ L ( θ t )]
{+8 Step: t 10 −10 −10 −
 L o ss : L ( θ t ) CNN ˆ E+}
[ L ( θ t )]
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 294 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) 297 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 σ is chosen
to ensure two types of noises have the same total variance for a fair comparison [50].
Note that p denotes the model size.
Analogous to Theorem 3.3, for the second isotropic noise,
[Equation]
Hence, the instability decreases with the parameter-space dimension and the resulting ﬂatnessconstraint is Tr( H ( θ ∗ ))
(cid:54) √ pB/η , which depends on the model size explicitly.
In contrast, for theﬁst noise, Theorem 3.3 implies (cid:107) H ( θ ∗ ) (cid:107) F (cid:54) √ B/η , independent of the model size.
This differencecan be intuitively explained as follows.
The isotropic noise wastes most energy in perturbing SGDalong ﬂat directions, which barely affects the instability.
In contrast, the geometry-aware noisefocuses more 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 relevance of our theory for practicalmodels and datasets.
Speciﬁcally, we consider classifying the CIFAR-10 dataset [21] with VGGnetworks [35] and ResNets [13] of a variety of depths.
In training, all explicit regularizations areremoved to keep consistent with our theoretical analysis.
We follow [15] to use square loss for theclassiﬁcation problem.
More details of the experimental 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 boundof different models.
We again observe that the ﬂatness is nearly independent of the model size forboth ResNets and VGG networks.
Moreover, we surprisingly discover that the upper bound obtainedfrom linear stability is rather tight in predicting the ﬂatness of convergent solution.
Lastly, we mention that the speciﬁc values of µ ( θ ) in Figure 5a are on the order of 0 .
01 , smaller thanthose in small-scale experiments. To understand the origin of this smallness, we conduct the sameexperiments but for classifying a two-class subset of CIFAR-10 and the result is reported in Figure 8of the appendix.
For this binary classiﬁcation, µ ( θ ) ’s are much larger, speciﬁcally on the order of0 .
1 , for all the models examined.
This comparison suggests that the alignment strength signiﬁcantlydepends on the intrinsic complexity of the problem, (nearly) independent of the model size.
[Figure]
Concluding remark
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—nondegenerate in the sharp directions of local landscape.
This geometry awareness together withlinear 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.
Obviously, many questions remains open.
For example, can we understand what roles the stabilityplays in the whole dynamic process of SGD instead of only around global minima?
Can we establishthe connection between the Hessian’s Frobenius norm and generalization? This is important forunderstanding implicit regularization of SGD.
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 areproprietary.
• 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’scontributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] Local analysis
[Table]
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
[Table] (d) Have you read the ethics review guidelines and ensured that your paper conforms to [Table]
them?
[Yes]2.
If you are including theoretical results...
[Table]
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
[Table]
(b) Did you include complete proofs of all theoretical results?
[Yes]
[Table]
If you ran experiments...
[Table]
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
[Table]
mental results (either in the supplemental material or as a URL)?
[Yes] See AppendixA.
(b) Did you specify all the training details (e
g., data splits, hyperparameters,
how they
[Table]
were chosen)?
[Yes] See Appendix A.
[Table]
(c) Did you report error bars (e.g., with respect to the random seed after running experiments 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]
[Table]
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
[Table]
(a) If your work uses existing assets, did you cite the creators?
[Yes]
[Table]
(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]
[Table]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
[Table]
using/curating?
[N/A]
[Table]
(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, ifapplicable?
[N/A]
[Table]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
[Table]
(c) Did you include the estimated hourly wage paid to participants and the total amount [Table] spent on participant compensation?
[N/A]
[Table]