Understanding the Generalization Beneﬁt of Normalization Layers: Sharpness Reduction
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Normalization layers (e.g., Batch Normalization, Layer Normalization) were introduced to help with optimization difﬁculties in very deep nets, but they clearly alsohelp generalization, even in not-so-deep nets.
Motivated by the long-held beliefthat ﬂatter minima lead to better generalization, this paper gives mathematicalanalysis and supporting experiments suggesting that normalization (together withaccompanying weight-decay) encourages GD to reduce the sharpness of loss surface.
Here “sharpness” is carefully deﬁned given that the loss is scale-invariant, aknown consequence of normalization.
Speciﬁcally, for a fairly broad class of neuralnets with normalization, our theory explains how GD with a ﬁnite learning rateenters the so-called Edge of Stability (EoS) regime and characterizes the trajectoryof GD in this regime via a continuous sharpness-reduction ﬂow.
Introduction
Training modern deep neural nets crucially relies on normalization layers to make the trainingprocess less sensitive to hyperparameters and initialization.
Popular normalization layers include
Batch Normalization (BN) [43], Layer Normalization [9], etc.
Normalization layers amount to areparametrization of the neural net so that the loss becomes invariant to the scale of most parameters(and with a minor change, to all parameters):
L ( c w ) =
L ( w ) for all scalings c > 0 [43, 7, 61].
Thecurrent paper assumes this scale-invariance for all parameters and analyzes the trajectory of gradientdescent with weight decay (WD):
[Equation]
Use of WD may appear nonsensical at ﬁrst sight because traditionally it is used to penalize largeparameter norm, which of course is inconsequential for scale-invariant loss — one can scale down theparameter norm arbitrarily without changing the loss value.
However, the scale of the parameter doesmatter for gradient and Hessian, and thus WD can affect the training dynamics.
In particular, simplecalculus shows ∇L ( w ) = 1 (cid:107) w (cid:107) 2 ∇L ( w (cid:107) w (cid:107) 2 ) and ∇ 2 L ( w ) = 1 (cid:107) w (cid:107) 22 ∇ 2 L ( w (cid:107) w (cid:107) 2 ) , so WD is in effecttrying to increase the norm of gradient and Hessian in training.
This makes the training dynamicsvery different from unnormalized nets, and requires revisiting classical convergence analyses, aswas done in various papers [7, 61, 64].
More notably, such a change in dynamics also yields a verydifferent solution at the end — in particular a solution that generalizes better [102].
The goal of the current paper is to improve mathematical understanding of how normalization togetherwith accompanying WD can improve generalization.
While this may arise from many places, we focuson sharpness-based generalization measures and exhibit settings where gradient descent persistentlyreduces sharpness in training normalized nets with WD, which we call the sharpness-reduction bias .
See Figures 1 and 2 for experiments on matrix completion (with BN) and CIFAR-10.
[Figure]
[Figure]
makes sense only if it is carefully deﬁned in consideration of various symmetries in neural nets.
Oneof the most straightforward measures of sharpness is the maximum eigenvalue of Hessian, namelyλ 1 ( ∇ 2 L ( w t )) .
But for normalized nets, this sharpness measure is vulnerable to weight rescaling,because one can scale the weight norm to make a minimizer arbitrarily ﬂat [22].
Also, this sharpnessmeasure may not decrease with the number of training steps: an empirical study by Cohen et al.
[16]shows that for various neural nets (including normalized nets), GD has an overwhelming tendencyto persistently increase λ 1 ( ∇ 2 L ( w t )) until it reaches the Edge of Stability (EoS) regime , a regimewhere λ 1 ( ∇ 2 L ( w t )) stabilizes around 2 / ˆ η ( ˆ η is the learning rate). See also Figure 2c.
The sharpness measure we use in this paper takes the scale-invariance property into account. We notethat techniques from previous works [75, 78, 27] can be easily adopted here to establish a PAC-Bayesbound on the test error, where our sharpness measure appears as an additive term (see Appendix C).
Deﬁnition 1.1 (Spherical Sharpness) .
For a scale-invariant loss L ( w ) (i.e., L ( c w ) =
L ( w ) for allc > 0 ), the spherical sharpness at w ∈ R D is deﬁned by λ 1 ( ∇ 2 L ( w (cid:107) w (cid:107) 2 )) , the maximum eigenvalueof the Hessian matrix after projecting w onto the unit sphere.
1.1 Our Contributions
In this work, we study the aforementioned sharpness-reduction bias of gradient descent with weightdecay (GD+WD), assuming the loss L is scale-invariant due to the presence of normalization.
For constant learning rate ˆ η and weight decay ˆ λ , we can rewrite gradient descent (1) on scale-invariantloss equivalently as Projected Gradient Descent (PGD) on the unit sphere with adaptive learningrates, θ t +1 ← Π( θ t − ˜ η t ∇L ( θ t )) , where θ t := w t (cid:107) w t (cid:107) 2 is the direction of w t , and ˜ η t is the “effective”learning rate at step t (deﬁned in Section 3).
We call ˜ η t adaptive because it can be shown that ˜ η tincreases when gradient is small and decreases when gradient is large, which resembles the behaviorsof adaptive gradient methods such as RMSprop [38].
Our main contributions are as follows:
We theoretically show that once θ t reaches a point near the manifold of minimizers of L , the
effective learning rate ˜ η t will keep increasing until 2 / ˜ η t roughly equals to the spherical sharpness,or in other words, gradient descent enters the EoS regime (Section 4.1).2.
In the EoS regime, we show that for gradient descent with a small (but ﬁnite) learning rate, θ toscillates around the manifold and moves approximately along a sharpness-reduction ﬂow, whichis a gradient ﬂow for minimizing spherical sharpness on the manifold (with gradient-dependent
learning rate) (Section 4.2).
As an application of our theory, we show that for linear regression with BN, GD+WD ﬁndsthe minimizer that corresponds to the linear model with minimum weight norm, which lookssurprisingly the same as the conventional effect of WD but is achieved through the completelydifferent sharpness-reduction mechanism (Section 5).
We experimentally veriﬁed the sharpness-reduction bias phenomena predicted by our theoremand its beneﬁts to generalization on CIFAR-10 with VGG-11 and ResNet-20, as well as matrixcompletion with BN (Appendix P).5.
We generalize our theoretical results of sharpness-reduction bias to a broader class of adaptivegradient methods, most notably a variant of RMSprop with scalar learning rate (Appendix B).
Our proof technique is novel and may have independent interest to the ML community.
Technical Contribution.
The main challenge to establish our theorem is that we need to analyzethe implicit bias of GD in the EoS regime.
In particular, we crucially rely on step size being ﬁnite —this is in sharp contrast to many previous works on implicit bias of GD [87, 86, 70, 47, 32, 31, 60,82 , 4, 14, 63, 71, 83, 88, 28] where the same bias exists at inﬁnitesimal LR.
Our analysis is inspiredby a previous line of works [12, 17, 65] showing that label noise can drive SGD to move on theminimizer manifold along the direction of minimizing the trace of Hessian, but the key differencehere is that we do not have any stochastic gradient noise.
Instead, we study the EoS regime andexhibit that GD oscillates around the minimizer manifold by connecting it to power method.
We showthat this oscillation is a driving power that pushes the parameter to move on the manifold, and analyzethe speed of this movement by modeling two key parameters of the dynamics as a 1-dimensional
Hamiltonian system.
To the best of our knowledge, our work is the ﬁrst one that theoretically provesa sharpness measure to decrease in standard training processes with gradient descent, without anyadditional regularization such as label noise [12, 17, 65], or using non-standard variants of gradientdescent update rule, e.g., using normalized GD or non-smooth wrappings on the loss function [8].
Related Works
Flatness and Generalization.
It has been long believed that minima locate in a ﬂat valley generalize better [39].
Li et al.
[58] ﬁrst visualized the loss landscape of neural networks and found“sharp” minima generalizes worse.
Keskar et al.
[49], Wu et al.
[94], Jastrz˛ebski et al.
[45] empiricallyveriﬁed the positive correlation between ﬂatness and generalization.
Neyshabur et al.
[78] gave atheoretical explanation for generalization beneﬁt of low sharpness using PAC-Bayes theory [75].
Inlight of this, Foret et al.
[27] proposed SAM algorithm to improve the generalization of SGD byminimizing the sharpness of the loss.
However, the deﬁnition of sharpness remains ambiguous inface of the positive homogeneity and invariance in ReLU networks [22], as networks with differentsharpness may represent the same function.
Towards resolving this challenge, multiple deﬁnitionsof scale-invariant sharpness have been proposed [99, 100, 91, 81].
Kwon et al.
[52] derived newalgorithms with better generalization by deﬁning new sharpness notions that are aware of positivehomogeneity and invariance.
He et al.
[34] argued that the local minima of modern deep networks aremore than being ﬂat or sharp and could be asymmetric valleys.
The theoretical implication of ﬂatnesshas also been explored speciﬁcally for two-layer nets [12, 77, 33, 65, 21] and deep linear nets [76].
Understanding Normalization Layers.
Ioffe and Szegedy [43] proposed BN with the originalmotivation to reduce Internal Covariate Shift (ICS), but Santurkar et al.
[84] challenged this viewby arguing that the effectiveness of BN comes from a smoothing effect on the training objective.
Acommon feature of normalization layers (including BN, LN [9], GN [96]) is that they make the lossinvariant to the scale of layer weights.
Based on this, several existing works have reported that trainingneural nets with normalization and weight decay can go out of the scope of the classical optimizationtheory.
Speciﬁcally, Li and Arora [61] showed that one can train the net to small loss even withlearning rates exponentially increasing; Li et al.
[62] showed that the dynamic of (full-batch) GD mayleave the stable regime; Lobacheva et al.
[67] empirically showed that the training loss can exhibit aperiodic behavior that sometimes improves the test accuracy.
However, despite a lot of effort devotedto understanding the optimization and generalization beneﬁts of normalization in various speciﬁcsettings, e.g. , [92, 11, 102, 69, 20, 66, 18, 19, 68, 53, 73], we still lack a complete and rigorousanalysis for the role of normalization in general, especially in terms of generalization.
Preliminaries
Let S D − 1 :
= { θ ∈ R D : (cid:107) θ (cid:107) 2 = 1 } be the unit sphere equipped with subspace topology.
We saya loss function
L ( w ) deﬁned on R D \ { 0 } is scale-invariant if L ( c w ) =
L ( w ) for all c > 0 .
In 119 other words, the loss value does not change with the parameter norm.
For a differentiable scaleinvariant function L ( w ) , the gradient is ( − 1) -homogeneous
and it is always perpendicular to w , i.e.,∇L ( c w ) =
c − 1 ∇L ( w ) for all c > 0 and (cid:104)∇L ( w ) , w (cid:105) = 0 .
The focus of this paper is the dynamicsof GD+WD on scale-invariant loss.
(1) gives the update rule for learning rate (LR) ˆ η and weightdecay (WD) ˆ λ .
We use θ t := w t (cid:107) w t (cid:107) 2 to denote the projection of w t onto the unit sphere at step t .
Anumber of works [92, 41, 7] propose to use the “effective learning rate” ˆ η (cid:107) w t (cid:107) 22 to measure the updateefﬁciency of θ t on S D − 1 .
Inspired by this notion, we write GD+WD on scale-invariant loss as aspeciﬁc kind of Projected Gradient Descent (PGD), and deﬁne ˜
η t := ˆ η (1 − ˆ η ˆ λ ) (cid:107) w t (cid:107) 22 to be the effectivelearning rate with slight abuse of terminology.
Lemma 3.1.
When the parameters w t are updated as (1) , θ t satisﬁes the following equation:
[Equation]
where ˜ η t := ˆ η (1 − ˆ η ˆ λ ) (cid:107) w t (cid:107) 22 is called the effective learning rate at step t , and Π : w (cid:55)→ w (cid:107) w (cid:107) 2 is theprojection operator that projects any vector onto the unit sphere.
The convergence rate of GD+WD has been analyzed by Li et al.
Here we present a variant oftheir theorem that bounds both the gradient and effective LR.
Theorem 3.2 (Variant of Theorem D.2 in [64]) .
Let L ( w ) be a scale-invariant loss function andρ 2 :
= sup {(cid:107)∇ 2 L ( w ) (cid:107) 2 : w ∈ S D − 1 } be the smoothness constant of L restricted on the unit sphere.
For GD+WD (1) with ˆ η ˆ λ ≤ 1 / 2 and ˜ η 0 ≤ [Equation] 2 ρ 2 (1 − ˆ η ˆ λ ) , let T 0 :
= (cid:108) 12ˆ η ˆ λ ln (cid:107) w 0 {+(cid:107) 
 ρ 
+}
π 2 ˆ η (cid:109) steps, there mustexist 0 ≤ t ≤ T 0 such that (cid:107)∇L ( θ t ) (cid:107) 22 ≤ 8 π 4 ρ 22 ˆ λ ˆ η and ˜ η t ≤ 2 π 2 ρ 2 (1 − ˆ η ˆ λ ) .
GD+WD on Scale-Invariant Loss Functions
This section analyzes GD+WD (1) on a scale-invariant loss L ( w ) , in particular what happens afterthe approximate convergence of Theorem 3.2.
We use w t to denote the trainable parameter at step tand θ t := w t (cid:107) w t (cid:107) 2 to denote its projection onto S D − 1 .
Section 4.1 analyzes the dynamics in the stableregime, where loss is guaranteed to decrease monotonically, and Theorem 3.2 suggests w t can getclose to a local minimizer at some time t 0 .
We show that the effective LR keeps increasing after t 0 ,causing GD+WD to eventually leave this stable regime and enter a new regime which we call the Edge of Stability (EoS).
In Section 4.2, we establish our main theorem, which connects the dynamicsof w t in the EoS regime to a sharpness-reduction ﬂow.
4.1 GD+WD Eventually Leaves the Stable Regime
A standard step of analyzing optimization methods is to do Taylor expansion locally for the lossfunction, and show that how the optimization method decreases the loss using a descent lemma .
In our case of scale-invariant loss functions, we use H ( w ) := ∇ 2 L ( w ) ∈
R D × D to denote the Hessianmatrix of L at w ∈ R D , and λ H1 ( w ) :
= λ 1 ( H ( w )) to denote the top eigenvalue of H ( w ) .
Lemma 4.1 (Descent Lemma) .
For scale-invariant loss L ( w ) , at step t of GD+WD we have
[Equation]
where λ ( t )max := sup α ∈ [0 , ˜ η t ] (cid:8) λ H1 ( θ t − α ∇L ( θ t ))
(cid:9) is an upper bound of spherical sharpness locally.
This descent lemma shows that the training loss L ( θ t ) keeps decreasing as long as the effective LR˜ η t is smaller than 2 /λ ( t )max , We call the regime of ˜ η t < 2 /λ ( t )max as the stable regime of GD+WD.
If˜ η t ≈ 2 /λ ( t )max with a small difference, then we call it as the Edge of Stability (EoS) regime .
This EoSregime is conceptually the same as that deﬁned by Cohen et al.
[16]; see Appendix G.3 for discussion.
Fix an initial point w 0 ∈ R D \ { 0 } .
Now we aim to characterize the dynamics of GD+WD when LRˆ η and WD ˆ λ are small enough.
Theorem 3.2 shows that for some t 0 ≤ T 0 , (cid:107)∇L ( θ t 0 ) (cid:107) 22 ≤ O (ˆ λ ˆ η ) 159 and ˜ η t 0 ≤ 1 π 2 ρ 2 < 2 ρ 2 , which means θ t 0 is an approximate ﬁrst-order stationary point of L on the unitsphere.
This does not guarantee that θ t 0 is close to any global minimizer, but in practice the trainingloss rarely gets stuck at a non-optimal value when the model is overparameterized [55, 79, 56, 101].
We are thus motivated to study the case where θ t 0 not only has small gradient (cid:107)∇L ( θ t 0 )
(cid:107) 2 2 ≤ O ( ˆ λ ˆ η )but also is close to a local minimizer θ ∗ ∈ S D − 1 of L in the sense that (cid:107) θ t 0 − θ ∗ (cid:107) 2 ≤ O ((ˆ λ ˆ η ) 1 / 2 )
(assuming smoothness, the latter implies the former).(a) (b) (c)
As the gradient is small near the local minimizer θ ∗ , startingfrom step t 0 , the norm of w t decreases due to the effectof WD.
See Figure 3a.
Since the effective LR is inverselyproportional to (cid:107) w t (cid:107) 22 , this leads to the effective LR to 169 increase.
Then Theorem 4.3 will show that the GD+WDdynamic eventually leaves the stable regime at some timet 1 > t 0 , and enters the EoS regime
where ˜ η t ≈ 2 /λ ( t )max .
To establish Theorem 4.3, we need to assume that L satisﬁes Polyak-Łojasiewicz (PL) condition locally, which is astandard regularity condition in the optimization literatureto ease theoretical analysis around a minimizer.
Intuitively,
PL condition guarantees that the gradient grows faster thana quadratic function as we move a parameter θ away fromθ ∗ .
Note that PL condition is strictly weaker than convexity as the function can still be non-convex under PLcondition (see, e.g., [48]).
Deﬁnition 4.2 (Polyak-Łojasiewicz Condition) .
For ascale-invariant loss L ( w ) and µ > 0 , we say that L satisﬁesµ -Polyak-Łojasiewicz condition (or µ -PL) locally around alocal minimizer θ ∗ on S D − 1 if for some neighborhood U ⊆ 185 S D − 1 of θ ∗ , [Equation]
Theorem 4.3.
Let L ( w ) be a C 2 -smooth scale-invariantloss that satisﬁes µ -PL around a local minimizer θ ∗ on theunit sphere, and
ρ 2 := sup {(cid:107)∇ 2 L ( w ) (cid:107) 2 : w ∈ S D − 1 } . 189 For GD+WD on
L ( w ) with learning rate ˆ η and weightdecay ˆ λ , if at some step t 0 , (cid:107) θ t 0 − θ ∗ (cid:107) 2 ≤ O ((ˆ λ ˆ η ) 1 / 2 ) and˜ η t 0 ≤ 2 ρ 2 < 2 λ H1 ( θ ∗ ) , and if ˆ λ ˆ η is small enough, then thereexists a time t 1 > t 0 such that (cid:107) θ t 1 − θ ∗ (cid:107)
= O ((ˆ λ ˆ η ) 1 / 2 )and ˜ η t 1 = 2 λ H1 ( θ ∗ )
+ O ((ˆ λ ˆ η ) 1 / 2 ) .
4.2 Dynamics at the Edge of Stability
From the analysis in the previous subsection, we know thatθ t can get close to a local minimizer θ ∗ and enter the EoSregime at some step t 1 .
But what happens after t 1 ?
Figure 3c gives a warm-up example on a 3D scale-invariant loss L : R 3 \ { 0 } → R , where the black 199 line is a manifold Γ consisting of all the minimizers.
In training with GD+WD, θ t ﬁrst goes closeto a local minimizer ζ 0 , then Theorem 4.3 suggests that WD causes the effective LR to steadilyincrease until the dynamic enters the EoS regime.
Now something interesting happens — θ t moves abit away from ζ 0 and starts to oscillate around the manifold Γ .
This oscillation is not completelyperpendicular to Γ but actually forms a small angle that pushes θ t to move downward persistentlyuntil θ t approaches the minimizer ζ ∗ denoted in the plot.
For a general scale-invariant loss L : R D \ { 0 } → R , which minimizer does θ t move towards?
In 206 this work, we consider the setting where there is a manifold Γ consisting only of local minimizers(but not necessarily all of them). We show that θ t always oscillates around the manifold once itapproaches the manifold and enters the EoS regime, and meanwhile θ t keeps moving in a directionof reducing spherical sharpness.
4.2.1 Assumptions
Now we formally introduce our main assumption on the local minimizer manifold Γ .
Assumption 4.4.
The loss function L : R D \ { 0 } → R is C 4 -smooth and scale-invariant.
Γ is a
C 2 -smooth, ( D Γ − 1) -dimensional submanifold of S D − 1 for some 0 ≤ D Γ < D , where every θ
∈
Γis a local minimizer of L on S D − 1 and rank( H ( θ )) =
D − D Γ .
Scale-invariance has become a standard assumption in studying neural nets with normalizationlayers [61, 62, 67].
For VGG and ResNet, the scale-invariance can be ensured after making minorchanges to the architectures (see Appendix Q.1).
The training loss L may not be smooth if theactivation is ReLU, but lately it has become clear that differentiable activations such as Swish [80], GeLU [37] can perform equally well.
Swish is indeed used in our VGG-11 experiments (Figure 2), but ResNet with ReLU activation also exhibits a sharpness-reduction bias empirically (see Appendix P.2).
For any local minimizer θ ∈ Γ , the eigenvalues λ H k ( θ ) must be non-negative.
And λ H k ( θ ) =
forall D − D Γ < k ≤ D , since Γ is of dimension D Γ − 1 .
The condition rank( H ( θ ))
=
D − D Γ 223 ensures that the Hessian is maximally non-degenerate on Γ , which also appears as a key assumptionin previous works [65, 8, 25].
This condition simpliﬁes the calculus on Γ in our analysis as it ensuresthat the null space of the matrix H ( θ ) equals to the tangent space of Γ at θ ∈ Γ .
It is also closelyrelated to PL condition (Deﬁnition 4.2) as Assumption 4.4 implies that L ( θ ) satisﬁes
µ -PL (for someµ > 0 ) locally around every θ ∈ Γ on the unit sphere (Arora et al.
[8], Lemma B.3).
To ease our analysis, we also need the following regularity condition to ensure that the largesteigenvalue is unique.
In our experiments, sharpness reduction happens even when the multiplicity ofthe top eigenvalue is more than 1 , but we leave the analysis of that case to future work.
Assumption 4.5.
For all θ ∈ Γ , λ H1 ( θ ) > λ H2 ( θ ) .
That is, the top eigenvalue of H ( θ ) is unique.
4.2.2 Main Theorem
First, we deﬁne η in :=
ˆ η ˆ λ as the intrinsic learning rate (name from Li et al. [62]) for convenience.
Assuggested in Theorems 3.2 and 4.3, θ t can get close to a local minimizer and be in the EoS regime atsome step t 1 : if ζ 0 is the local minimizer, then (cid:107) θ t 1 − ζ 0 (cid:107)
= O ( η 1 / 2 in ) and ˜ η t 1 = 2 λ H1 ( ζ 0 )
+ O ( η 1 / 2 in ) .
In our main theorem, we start our analysis from step t 1 while setting t 1 = 0 WLOG (otherwise wecan shift the step numbers).
We connect GD+WD in the EoS regime to the following gradient ﬂow (3)on the manifold Γ minimizing spherical sharpness (with gradient-dependent learning rate), and showthat one step of GD+WD tracks a time interval of η in in the gradient ﬂow.
[Equation]
Here we use the notation ∇ Γ R ( θ ) for any R : R D
→ R to denote the projection of ∇ R ( θ ) ontothe tangent space T θ ( Γ ) at θ ∈ Γ .
ζ ( τ ) reduces sharpness as it moves in direction of the negativegradient of log λ H1 ( ζ ( τ )) on Γ .
A simple chain rule shows how fast the spherical sharpness decreases:d
[Equation]
for small gradient
[Equation]
for large gradient .
Note that it is not enough to just assume that θ 0 is close to ζ 0 .
If θ 0 = ζ 0 holds exactly, then thesubsequent dynamic of w t is described by w t =
(1 − ˆ η ˆ λ ) t w 0 with direction unchanged.
There arealso some other bad initial directions of w 0 that may not lead to the sharpness-reduction bias.
Thismotivates us to do a smoothed analysis for the initial direction: the initial direction is ζ with tinyrandom perturbation, where the perturbation scale is allowed to vary from exp( − η − o (1) in ) to η 1 / 2 − o (1) in ,and we show that a good initial direction is met with high probability as η in → 0 .
Alternatively,one can regard it as a modeling of the tiny random noise in GD+WD due to the precision errors inﬂoating-point operations.
See Figure 4b; the training loss can never be exactly zero in practice.
Initialization Scheme. Given a local minimizer ζ 0 ∈ Γ , we initialize w 0 ∈ R D \ { 0 } as follows:draw ξ ∼ N ( 0 , σ 20 I /D ) from Gaussian and set the direction of w 0 to ζ 0 + ξ (cid:107) ζ 0 + ξ (cid:107) 2 , where σ 0 can takeany value in [exp( − η − o (1) in ) , η 1 / 2 − o (1) in ] ; then set the parameter norm (cid:107) w 0 (cid:107) 2 to be any value thatsatisﬁes (cid:12)(cid:12)(cid:12) ˜ [Equation] −
[Equation] 1) in , where ˜ η 0 := [Equation]
22 is the effective LR for the ﬁrst step.
Theorem 4.6. Under Assumptions 4.4 and 4.5, for GD+WD (1) with sufﬁciently small intrinsiclearning rate η in :=
ˆ η ˆ λ , if we follow the above initialization scheme for some ζ 0 ∈ Γ , then withprobability 1 − O ( η 1 / 2 − o (1) in ) , the trajectory of θ t :
=
w t (cid:107) w t (cid:107) 2 approximately tracks a sharpnessreduction ﬂow ζ : [0 , T ] → Γ that starts from ζ 0 and evolves as the ODE (3) up to time T (if solutionexists), in the sense that (cid:107) θ t − ζ ( tη in ) {+(cid:107) 
+} = O ( η 1 / 4 − o (1) in ) for all 0 ≤ t ≤ T/η in .
Remark 4.7 (Magnitude of Oscillation) .
As suggested by Figure 3c, θ t actually oscillates aroundthe manifold. But according to our analysis, the magnitude of oscillation is as small as O ( η 1 / 2 − o (1) in ) ,so it is absorbed into our ﬁnal bound O ( η 1 / 4 − o (1) in ) for the distance between θ t and ζ ( tη in ) .
4.2.3 Proof Idea
Throughout our proof, we view GD+WD for w t as a PGD for θ t with effective LR ˜ η t (Lemma 3.1).
To track θ t with ζ ( tη in ) , for each step t , we construct a local minimizer φ t ∈ Γ that serves asthe “projection” of θ t onto the manifold Γ , in the sense that the displacement x t := θ t − φ t isapproximately perpendicular to the tangent space of Γ at φ t .
Our entire proof works through induction.
According to the initial conditions, the dynamic is initially in the EoS regime: (cid:107) x t (cid:107) 2 ≤ η 1 / 2 − o (1) inand | ˜ η t − 2 /λ H1 ( φ t ) | ≤ η 1 / 2 − o (1) in at t = 0 .
In our induction, we maintain the induction hypothesisthat these two EoS conditions continue to hold for all t ≥ 0 .
Period-Two Oscillation.
A key insight in our proof is that after a few initial steps, θ t is oscillatingaround φ t along the ± v H1 ( θ ) directions, where v H1 ( θ ) is a unit top eigenvector of H ( θ ) and ischosen in a way that v H1 ( θ ) is continuous on Γ .
More speciﬁcally, x t = h t v H1 ( φ t ) + O ( (cid:107) x t (cid:107) 22 ) forh t := (cid:104) x t , v H1 ( φ t ) (cid:105) .
The oscillation is of period 2: h t > 0 when t is even and h t < 0 when t is odd.
See Figure 4d for an example.
This oscillation can be connected to a power method for the matrix I − ˜ η t H ( φ t ) .
In the EoS regime,we can approximate θ t +1 (when x t is small) as θ t +
=
Π( θ t − ˜ η t ∇L ( θ t ))
≈ Π( θ t − ˜ η t H ( φ t ) x t )
≈ 278 θ t − ˜ η t H ( φ t ) x t by Taylor expansions of ∇L and Π : R D \ { 0 } → S D − 1 .
We can further show that 279 φ t {++
+} ≈ φ t due to our choice of projections.
Then the connection to power method is shown below:
[Equation]
By simple linear algebra, v H1 ( φ t ) is an eigenvector of I − ˜ η t H ( φ t ) , associated with eigenvalue1 − ˜ η t λ H1 ( φ t )
≈ − 1 .
The remaining eigenvalues are { 1 − ˜
η t λ H i ( φ t ) } Di =2 , where λ H i ( φ t ) is the i -thlargest eigenvalue of H ( θ t ) , and they lie in the range ( − 1 , 1] since λ H i ( φ t ) ∈ [0 , λ H1 ( φ t )) .
Using asimilar analysis to power method, we show that x t quickly aligns to the direction of ± v H1 ( φ t ) after afew initial steps, as the corresponding eigenvalue has approximately the largest absolute value.
To formally establish the above result, we need a tiny initial alignment between x 0 and v H1 ( φ 0 ) , justas the initial condition in power method.
This is where we need the initial random perturbation.
[Figure]
Oscillation Drives φ t to Move.
This period-two oscillation is the driving power to push φ t tomove on the manifold.
The main idea here is to realize that the oscillation direction deviates slightlyfrom the direction of ± v H1 ( φ t ) by using a higher-order approximation.
We speciﬁcally use the
Taylor approximation to show that this deviation leads φ t to move slightly on Γ : after each cycle ofoscillation, φ t +2 ≈ φ t − 4 h 2 t ∇ Γ log λ H1 ( φ t )
+ O ( η 1 .
5 − o (1) in ) , which resembles two steps of gradientdescent on Γ to minimize the logarithm of spherical sharpness with learning rate 2 h 2 t ,
Periodic Behavior of h t and ˜ η t .
It remains to analyze the dynamics of h t so that we can knowhow fast the sharpness reduction is.
Our analysis is inspired by an empirical study from Lobachevaet al. [67], which reveals a periodic behavior of gradients and effective learning rates in trainingnormalized nets with weight decay.
In our theoretical setting, we capture this periodic behavior byshowing that h t and ˜ η t do evolve periodically.
See Figures 4c and 4d for an example.
The key is that ˜ η t changes as an adaptive gradient method: ˜ η t increases when gradient is small anddecreases when gradient is large (due to the effect of WD; see Figures 3a and 3b), and in our casethe gradient norm scales as | h t | since ∇L ( θ t )
≈ h t λ H1 ( φ t ) v H1 ( φ t ) . According to our power method 301 approximation, h t +2 ≈ (1 − ˜ η t λ H1 ( φ t ))
2 h t , so | h t | decreases when ˜ η t < 2 /λ H1 ( φ t ) .
But | h t | cannot 302 decrease forever, since ˜ η t increases when | h t | is sufﬁciently small.
When ˜ η t rises to over 2 /λ H1 ( φ t ) ,| h t | changes from decreasing to increasing according to our approximation.
But h t cannot increase 304 indeﬁnitely either, since ˜ η t decreases when | h t | is sufﬁciently large.
A period is ﬁnally ﬁnished when˜ η t drops below 2 /λ H1 ( φ t ) .
In our theoretical analysis, we connect this periodic behavior with a 1-dimensional Hamiltoniansystem (see Appendix H.2), and show that 2 h 2 t in each step can be approximated by its average valuein the period without incurring a large error. Further calculations show that this average value isapproximately [Equation] Γ log λ H1 ( ζ ( tη in ))
(cid:107) 2 , the learning rate in the ﬂow (3) multiplied with η in .
We cantherefore conclude that each step of φ t (or θ t ) tracks a time interval of η in in the ﬂow.
Extensions.
We note that this periodic behavior is not limited to GD+WD on scale-invariant loss,as the above intuitive argument holds as long as the effective LR changes adaptively with respectto gradient change.
Based on this intuition, an important notion called Quasi-RMSprop scheduleris proposed.
For a PGD method, a learning rate scheduler is a rule for changing the effective LRin each step, and Quasi-RMSprop is a speciﬁc class of schedulers we deﬁne, including the waythat the effective LR changes in GD+WD on scale-invariant loss (if viewed as PGD).
Our proofis done in a uniﬁed way that works as long as the effective LR changes in each step according toa Quasi-RMSprop scheduler.
As a by-product, a similar theorem can be proved for GD (withoutprojection) on non-scale-invariant loss if the LR changes as a Quasi-RMSprop in each step.
Forexample, we can extend our analysis to RMSprop with a scalar learning rate.
See Appendix B.
Case Study: Linear Regression with Batch Normalization
In this section, we analyze the GD+WD dynamics on linear regression with Batch Normalization(BN), as a simple application of our theory.
Let { ( x i , y i ) } ni
=1 be a dataset, where x i ∈ R d andy i ∈ R are inputs and regression targets.
We study the over-parameterized case where d (cid:29) n , andwe assume that the regression targets are generated by an unknown linear model.
A classic linear model is parameterized by ( w , b ) ∈ R d × R and outputs w (cid:62) x + b given input x ,but now we add a BN to the output.
More speciﬁcally, we consider a batch-normalized linear modelΦ( x ; w , γ, β ) :
= γ · w (cid:62) x − µ 1 σ+ β , where µ 1 , σ 1 are the mean and variance of { w (cid:62) x i } ni =1 over thewhole dataset 3 , and the bias term b is cancelled out due to BN.
Note that Φ( x ; w , γ, β ) is still a linearfunction with respect to x .
Let µ x ∈ R d and Σ x ∈ R d
×
d be the mean and covariance of the inputdata { x i } ni =1 .
Then Φ( x ; w , γ, β ) can be rewritten as:
[Equation]
where
[Equation]
No matter how w is set, the output mean and variance of Φ are always β and γ 2 .
To simplify ouranalysis, we ﬁx β, γ to be non-trainable constants so that the mean and variance of Φ ’s output matchwith those of { y i } n i =1 , that is, we set β = µ y and γ = σ y to be the mean and standard deviation of y iover the whole dataset.
Then the training loss is L ( w ) :
= 1 n (cid:80) i ∈ [ n ] (Φ( x i ; w , γ, β ) − y i ) 2 .
Theorem 5.1.
In our setting of linear regression with BN, the sharpness-reduction ﬂow ζ deﬁned in(3) converges to the solution w ∗ ∈ S d − 1 that minimizes sharpness λ H1 ( w ∗ ) on Γ , regardless of theinitialization.
Moreover, the coefﬁcients ( ˜ w , ˜ b ) associated with w ∗ (deﬁned in (4) ) are the optimalsolution of the following constrained optimization problem (M) :
[Equation]
At ﬁrst sight the result may appear trivial because the intent of WD is to regularize L 2 -norm.
But thisis deceptive because in scale-invariant nets WD changes from an explicit regularizer to an implicitone.
This also challenges conventional view of optimization.
GD is usually viewed as a discretizationof its continuous counterpart, gradient ﬂow (GF), and theoretical insight for the discrete updateincluding convergence rate and implicit bias is achieved by analyzing the continuous counterpart (See
Appendix A for a list).
However, GF does not have the same sharpness-reduction bias as GD.
Asdiscussed in [61], adding WD only performs a time-rescaling on the GF trajectory on scale-invariantloss, but does not change the point that GF converge to if we project the trajectory onto the unitsphere.
One can easily show that GF may converge to any zero-loss solution, but no matter how small
LR is, GD exhibits the sharpness-reduction bias towards the optimal solution of (M).
To our bestknowledge, this result is the ﬁrst concrete example where even with arbitrarily small LR, GD can stillgeneralize better than GF under natural settings.
Conclusions and Future Work
We exhibited settings where gradient descent has an implicit bias to reduce spherical sharpness intraining neural nets with normalization layers and weight decay, and we veriﬁed experimentally thissharpness-reduction bias predicted by our theorem as well as its generalization beneﬁt on CIFAR-10.
Our theoretical analysis applies to dynamics around a minimizer manifold and requires a small (butﬁnite) learning rate so that we can show that the parameter oscillates locally and approximately tracksa sharpness-reduction ﬂow.
We note that in practice a decrease in spherical sharpness is observedeven with moderate LR and even before getting close to a minimizer manifold.
Explaining thesephenomena is left for future work.
Now we list some other future directions.
The ﬁrst is to generalizeour results to SGD, where the sharpness measure may not be the spherical sharpness and coulddepend on the structure of gradient noise.
Second, to understand the beneﬁt of reducing sphericalsharpness on speciﬁc tasks, e.g., why does reducing spherical sharpness encourage low-rank on matrixcompletion with BN (Figure 1)?
Third, to study sharpness-reduction bias for neural net architecturesthat are not scale-invariant on all parameters (e.g., with certain unnormalized layers).
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
contributions and scope?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[N/A] We are
basically a theoretical work studying the generalization myestery in deep learning.
We do
not see any negative societal impact.
(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] The assumptionsof our main theorem for GD+WD on scale-invariant loss are stated in Section 4.2.1.
ForGD/PGD with Quasi-RMSprop schedulers in general, see Appendix B.3.
(b) Did you include complete proofs of all theoretical results?
[Yes] See Appendices G and Hfor the proofs for our main theorems, Appendix O for the proof for the linear example.3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimentalresults (either in the supplemental material or as a URL)?
[Yes] See our supplementarymaterial.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they werechosen)?
[Yes] See Appendix Q.
(c) Did you report error bars (e.g., with respect to the random seed after running experimentsmultiple times)?
[No] Most of our experiments have only run once due to computationalconstraints, but we veriﬁed the sharpness-reduction bias across various settings.
(d) Did you include the total amount of compute and the type of resources used (e.g., type ofGPUs, internal cluster, or cloud provider)?
[Yes] See Appendix Q.
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?
[N/A]
(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ﬁableinformation or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applica-
ble?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board
(IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]