Understanding the Generalization Beneﬁt of Normalization Layers: Sharpness Reduction
Abstract
Normalization layers (e.g., Batch Normalization, Layer Normalization) were introduced to help with optimization difﬁculties in very deep nets, but they clearly also help generalization, even in not-so-deep nets.
Motivated by the long-held belief that ﬂatter minima lead to better generalization, this paper gives mathematical analysis and supporting experiments suggesting that normalization (together with accompanying weight-decay) encourages GD to reduce the sharpness of loss surface.
Here “sharpness” is carefully deﬁned given that the loss is scale-invariant, a known consequence of normalization.
Speciﬁcally, for a fairly broad class of neural nets with normalization, our theory explains how GD with a ﬁnite learning rate enters the so-called Edge of Stability (EoS) regime, and characterizes the trajectory of GD in this regime via a continuous sharpness-reduction ﬂow.
Introduction
Training modern deep neural nets crucially relies on normalization layers to make the training process less sensitive to hyperparameters and initialization.
The two of the most popular normalization layers are Batch Normalization (BN) [55] for vision tasks and Layer Normalization (LN)
[9] for language tasks.
Recent works also proposed other normalization layers aiming for better performance, most notably including Group Normalization (GN) [120], Weight Normalization (WN) [102], Scaled Weight
Standardization (SWS) [97, 53, 14] etc.
Most normalization layers amount to a reparametrization 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 [55, 7, 77].
The current paper assumes this scale-invariance for all parameters and analyzes the trajectory of gradient descent with weight decay (WD):
[Equation]
The use of WD is a common practice that has been adopted in training state-of-the-art neural nets, such as ResNets [46, 47] and Transformers [29, 15].
Previous ablation studies showed that adding
WD to normalized nets indeed leads to better generalization [126, 72, 125].
More notably, Liu et al.
[83] conducted experiments of training ResNets initialized from global minima with poor test accuracy, and showed that SGD with WD escapes from those bad global minima and attains good test accuracy.
In contrast, training with vanilla SGD yields signiﬁcant generalization degradation.
In the traditional view, WD regularizes the model by penalizing the parameter norm, but this may appear nonsensical for scale-invariant loss because one can scale down the norm arbitrarily without changing the loss value.
However, the scale of the parameter does matter in backward propagation, and thus WD can affect the training dynamics.
In particular, simple calculus shows ∇L
[Equation]
=
[Equation] [Equation]
WD is in effect trying to enlarge the gradient and Hessian in training.
This makes the training dynamics very different from unnormalized nets and requires revisiting classical convergence analyses [77, 78, 84, 80].
[Figure]
[Figure]
The current paper aims to improve mathematical understanding of how normalization improves generalization.
While this may arise from many places, we focus on studying the dynamics of (fullbatch) GD (1), which is a necessary ﬁrst step towards understanding SGD. We show that the interplay between normalization and WD provably induces an implicit bias to persistently reduce the sharpness of the local loss landscape during the training process, which we call the sharpness-reduction bias .
It is long believed that ﬂatter minima generalize better [50, 63, 95], but the notion of sharpness/ﬂatness makes sense only if it is carefully deﬁned in consideration of various symmetries in neural nets.
One of 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 [31].
Also, this sharpness measure may not decrease with the number of training steps: an empirical study by Cohen et al.
[24] shows that for various neural nets (including normalized nets), GD has an overwhelming tendency to persistently increase λ 1 ( ∇ 2 L ( w t )) until it reaches the Edge of Stability (EoS) regime , a regime where λ 1 ( ∇ 2 L ( w t )) stays around 2 / ˆ η
( ˆ η is the learning rate).
See also Section 6 and Figure 2c.
1.1 Our Contributions
The sharpness measure we use in this paper takes care of the scale-invariance in normalized nets.
We are motivated by our experiments on matrix completion (with BN) and CIFAR-10, where our sharpness measure decreases as the training proceeds, and the generalization improves accordingly; see Figures 1 and 2.
We note that techniques from previous works [92, 95, 37] can be easily adopted here to establish a PAC-Bayes bound 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 all c > 0 ), the spherical sharpness at w ∈ R D is deﬁned by λ 1 ( ∇ 2 L ( w (cid:107) w (cid:107) 2 )) , the maximum eigenvalue of the Hessian matrix after projecting w onto the unit sphere.
Based on Deﬁnition 1.1, we study the aforementioned sharpness-reduction bias in training normalized nets with GD+WD (deﬁned in (1)).
For constant learning rate ˆ η and weight decay ˆ λ , we can rewrite this rule equivalently as Projected Gradient Descent (PGD) on the unit sphere with adaptive learning rates, θ 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 (see Lemma 3.1).
We call ˜ η t adaptive because it can be shown to resemble the behaviors of adaptive gradient methods (e.g., RMSprop [49]): ˜ η t increases when gradient is small and decreases when gradient is large (Figure 3).
Our main contributions are as follows:
After θ t reaches a point near the manifold of minimizers of L , we
theoretically
show that the effective learning rate ˜ η t increases until GD enters a regime where 2 / ˜ η t roughly equals to the spherical sharpness (or equivalently 2 / ˆ η ≈ λ 1 ( ∇ 2 L ( w t )) ), namely the EoS regime (Section 4.1).
In the EoS regime, we show that for GD with a small (but ﬁnite) learning rate, θ t oscillates around the manifold and moves approximately along a sharpness-reduction ﬂow, which is 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 ﬁnds the minimizer that corresponds to the linear model with minimum weight norm, which looks surprisingly the same as the conventional effect of WD but is achieved through the completely different sharpness-reduction mechanism (Section 5).
We experimentally veriﬁed the sharpness-reduction phenomenon predicted by our theorem and its beneﬁts to generalization on CIFAR-10 with VGG-11 and ResNet-20, as well as matrix completion with BN (Appendix P).
We generalize our theoretical results of sharpness-reduction bias to a broader class of adaptive gradient methods, most notably a variant of RMSprop with scalar learning rate (Appendix B).
Technical Contribution.
Our proof technique is novel and may have independent interest to the
ML community.
The main challenge is that we need to analyze the implicit bias of GD in the EoS regime which crucially relies on step size being ﬁnite — this is in sharp contrast to many previous works on implicit bias of GD [107, 106, 87, 59, 43, 42, 76, 100, 4, 22, 79, 88, 101, 108, 38] where the same bias exists at inﬁnitesimal LR.
Our analysis is inspired by a previous line of works [13, 25, 81] showing that label noise can drive SGD to move on the minimizer manifold along the direction of minimizing the trace of Hessian.
We borrow a few lemmas from those analyses, but the overall proof strategy is very different because our setting does not even have any stochastic gradient noise.
Instead, we connect the dynamics in the EoS regime to power methods and show that GD oscillates around the minimizer manifold.
This oscillation then becomes a driving power that pushes the parameter to move on the manifold.
Finally, we analyze the speed of this movement by modeling two key parameters of the dynamics as a 1-dimensional Hamiltonian system (Figure 6).
To the best of our knowledge, we are the ﬁrst to provide theoretical proof for a sharpness measure to decrease during the standard GD training, without any additional regularization (e.g., label noise [13, 25, 81]) and without involving uncommon variants of GD (e.g., normalized GD or non-smooth wrappings on the loss function [8]).
Related Works
Sharpness and Generalization.
It has been long believed that ﬂat minima generalize better [50].
Several empirical studies [63, 74, 117, 57] veriﬁed the positive correlation between ﬂatness and generalization.
Neyshabur et al.
[95] justiﬁed this via PAC-Bayes theory [92].
Several other theoretical papers explored the generalization properties of ﬂat minima speciﬁcally for two-layer nets [13, 94, 44, 81, 30] and deep linear nets [93].
Jiang et al.
[60] conducted extensive experiments for all existing generalization measures to evaluate their correlation and causal relationships with generalization error, concluding that sharpness-based measures perform the best overall.
In light of this, Foret et al.
[37] proposed SAM algorithm to improve the generalization by minimizing the sharpness.
Despite so many positive results on sharpness-based measures, a common issue of many works is that the measures may suffer from sensitivity to rescaling of parameters in deep nets [31].
Another issue is that the minima could lie in asymmetric valleys that are ﬂat on one side and sharp on the other [45].
Understanding Normalization Layers.
The beneﬁts of normalization layers can be shown in various aspects.
A series of works studied the forward propagation of deep nets at random initialization, showing that normalization layers stabilize the growth of intermediate layer outputs with depth [14, 10, 28], provably avoid rank collapse [26] and orthogonalize representations [27].
Although these works mainly focused on BN [55], Lubana et al.
[85], Labatie et al.
[67] provided thorough discussions on the applicability of these arguments to other normalization layers.
It is also believed that BN has a unique regularization effect through the noise in batch statistics [86, 111, 104].
Several other works argued that normalization layers lead to a smoothening or preconditioning effect of the loss landscape [103, 12, 39, 61, 82, 68], which may help optimization.
By analyzing the training dynamics, Arora et al.
[7] rigorously proved that normalization yields an auto-tuning effect of the effective learning rate ˜ η t , which makes the asymptotic speed of optimization much less sensitive to the learning rate and initialization.
In linear regression settings, Cai et al.
[16], Kohler et al.
[65] showed that training with BN leads to a faster convergence rate; Wu et al.
[119] studied the implicit regularization effect of WN [102].
For two-layer nets with normalization, Ma and Ying [90] derived a mean-ﬁeld formulation of the training dynamics; Dukler et al.
[33] proved a convergence rate via
NTK-based analysis.
The current paper focuses on the interplay between normalization and WD during training, whereas all the above works either do not analyze the dynamics or assume no WD.
Interplay Between Normalization and WD.
A common feature of normalization layers (including but not limited to BN, WN, LN, GN, SWS) is that they make the loss invariant to the scale of layer weights.
In presence of both scale-invariance and WD, training dynamics can go out of the scope of the classical optimization theory, e.g., one can train the net to small loss even with learning rates exponentially increasing [77].
A series of works investigated into the interplay between normalization and WD and argued that the training dynamic with SGD eventually reaches an “equilibrium” state, where the parameter norm [78, 113, 21] and the size of angular update [114] become stable.
Li et al.
[78], Wang and Wang [115] provided empirical and theoretical evidence that the function represented by the net also equilibrates to a stationary distribution that is independent of initialization.
This could be related to Liu et al.
[83]’s experiments on the ability of SGD with WD to escape from bad initialization, but it remains unclear why the generalization should be good at the equilibrium state.
In this paper, we focus on (full-batch) GD, which is the most basic and important special case of SGD.
Preliminaries
Let S D − 1 :
= { θ ∈ R D : (cid:107) θ (cid:107) 2 = 1 } be the unit sphere equipped with subspace topology.
We say a loss function L ( w ) deﬁned on R D \ { 0 } is scale-invariant if L ( c w ) =
L ( w ) for all c > 0 .
In 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) =
(see Lemma D.1).
The focus of this paper is the dynamics of GD+WD on scale-invariant loss.
(1) gives the update rule for learning rate (LR) ˆ η and weight decay (WD) ˆ λ .
We use θ t := w t (cid:107) w t (cid:107) 2 to denote the projection of w t onto S D − 1 at step t .
We write GD+WD on scale-invariant loss as a speciﬁc kind of Projected
Gradient Descent (PGD) and deﬁne the effective learning rate to be the LR ˜ η t :=
[Equation] [Equation]
(cid:107) w t (cid:107) 22 that appears in the update rule of PGD.
This notion is slightly different from the effective learning rate ˆ η (cid:107) w t (cid:107) 22 deﬁned in previous works [113, 52, 7], but ours is more convenient for our analysis.
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
[Equation]
w
[Equation]
2 is the projection operator that projects any vector onto the unit sphere.
GD+WD on Scale-Invariant Loss Functions
This section analyzes GD+WD (1) on a scale-invariant loss L ( w ) , in particular what happens after approaching a manifold of local minimizers.
Section 4.1 analyzes the dynamics in the stable regime, where loss is guaranteed to decrease monotonically, and Theorem 4.2 suggests w t can get close 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 dynamics of 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 loss function, 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 Hessian matrix 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 .
We remark that this condition for EoS regime is essentially the same as ˆ η ≈ 2 /λ H1 ( w ) in Cohen et al.
[24]’s deﬁnition because ˜ η t · λ ( t )max ≈ ˆ η · λ H1 ( w ) ; see Appendix G.3.
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.
The convergence rate of GD+WD has been analyzed by Li et al.
Here we present a variant of their theorem that bounds both the gradient and effective LR.
Theorem 4.2 (Variant of Theorem D.2, Li et al.
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 ≤ 1 π
[Equation]
[Equation] [Equation]
η ˆ λ ) , let T 0 :=
[Equation]
steps, there must exist 0 ≤ t ≤ T 0 such that (cid:107)∇L ( θ t ) (cid:107) 22 ≤ 8 π 4 ρ 22 ˆ λ ˆ η and ˜ η t ≤
[Equation]
Theorem 4.2 shows that for some t 0 ≤ T 0 , (cid:107)∇L ( θ t 0 ) (cid:107) 22 ≤ O (ˆ λ ˆ η ) and ˜ η t 0 ≤ 1 π 2 ρ 2 < 2 ρ 2 , which means θ t 0 is an approximate ﬁrst-order stationary point of L on the unit sphere.
This does not guarantee that θ t 0 is close to any global minimizer, but in practice the training loss rarely gets stuck at a non-optimal value when the model is overparameterized [70, 96, 71, 125].
We are thus motivated to study the case where θ t 0 not only has small gradient (cid:107)∇L ( θ t 0 )
(cid:107) 22 ≤ 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).
[Figure]
As the gradient is small near the local minimizer θ ∗ , starting from step t 0 , the norm of w t decreases due to the effect of WD.
See Figure 3a.
Since the effective LR is inversely proportional to (cid:107) w t (cid:107) 22 , this leads to the effective LR to increase.
Then Theorem 4.4 will show that the GD+WD dynamic eventually leaves the stable regime at some time t 1 > t 0 , and enters the EoS regime where ˜ η t ≈ 2 /λ ( t )max .
To establish Theorem 4.4, we need to assume that L satisﬁes Polyak-Łojasiewicz (PL) condition locally, which is a standard regularity condition in the optimization literature to ease theoretical analysis around a minimizer.
Intuitively,
PL condition guarantees that the gradient grows faster than a 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 PL condition (see, e.g., [62]).
Deﬁnition 4.3 (Polyak-Łojasiewicz Condition) .
For a scale-invariant loss L ( w ) and µ > 0 , we say that L satisﬁes µ -Polyak-Łojasiewicz condition (or µ -PL) locally around a local minimizer θ ∗ on S D − 1 if for some neighborhood U ⊆
S D − 1 of θ ∗ , ∀ θ ∈ U : 12 (cid:107)∇L ( θ ) (cid:107) 22 ≥ µ · ( L ( θ ) − L ( θ ∗ )) .
Theorem 4.4.
Let L ( w ) be a C 2 -smooth scale-invariant loss that satisﬁes µ -PL around a local minimizer θ ∗ on the unit sphere, and
ρ 2 :=
sup {(cid:107)∇ 2 L ( w ) (cid:107) 2 : w ∈ S D − 1 } .
For GD+WD on
L ( w ) with learning rate ˆ η and weight decay ˆ λ , 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 there exists a time t 1 > t 0 such that (cid:107) θ
[Equation]
− θ ∗ (cid:107)
= O ((ˆ
[Equation]
ˆ
[Equation]
2 ) and ˜
[Equation]
=
[Equation]
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
EoS regime at some step t 1 .
But what happens after t 1 ?
Figure 4 gives a warm-up example on a 3D scale-invariant loss L : R 3 \ { 0 } → R , where the black line is a mani- fold Γ consisting of all the minimizers.
In training with
GD+WD, θ t ﬁrst goes close to a local minimizer ζ 0 , then
Theorem 4.4 suggests that WD causes the effective LR to steadily increase until the dynamic enters the EoS regime.
Now something interesting happens — θ t moves a bit away from ζ 0 and starts to oscillate around the manifold Γ .
This oscillation is not completely perpendicular to Γ but actually forms a small angle that pushes θ t to move downward persistently until θ 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 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 it approaches the manifold and enters the EoS regime, and meanwhile θ t keeps moving in a direction of reducing spherical sharpness.
4.2.1 Assumptions
Now we formally introduce our main assumption on the local minimizer manifold Γ .
Assumption 4.5.
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 normalization layers [77, 78, 84].
For VGG and ResNet, the scale-invariance can be ensured after making minor changes to the architectures (see Appendix Q.1).
The training loss L may not be smooth if the activation is ReLU, but lately it has become clear that differentiable activations such as Swish [98],
GeLU [48] 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 ( θ ) =
for all D − D Γ < k ≤ D , since Γ is of dimension D Γ − 1 .
The condition rank( H ( θ ))
=
D − D Γ ensures that the Hessian is maximally non-degenerate on Γ , which also appears as a key assumption in previous works [81, 8, 35].
This condition simpliﬁes the calculus on Γ in our analysis as it ensures that the null space of the matrix H ( θ ) equals to the tangent space of Γ at θ ∈ Γ .
It is also closely related to PL condition (Deﬁnition 4.3) as Assumption 4.5 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 largest eigenvalue is unique.
In our experiments, sharpness reduction happens even when the multiplicity of the top eigenvalue is more than 1 , but we leave the analysis of that case to future work.
Assumption 4.6.
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.
[78]) for convenience.
As suggested in Theorems 4.2 and 4.4, θ t can get close to a local minimizer and be in the EoS regime at some 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 we can 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 show that one step of GD+WD tracks a time interval of length η in in the gradient ﬂow.
[Equation]
Here we use the notation ∇ Γ R ( θ ) for any R : R D
→ R to denote the projection of ∇ R ( θ ) onto the tangent space T θ ( Γ ) at θ ∈ Γ .
ζ ( τ ) reduces sharpness as it moves in direction of the negative gradient 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 the subsequent dynamic of w t is described by w t =
(1 − ˆ η ˆ λ ) t w 0 with direction unchanged.
There are also some other bad initial directions of w 0 that may not lead to the sharpness-reduction bias.
This motivates us to do a smoothed analysis for the initial direction: the initial direction is ζ with tiny random 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 5b; 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 take any 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 that satisﬁes (cid:12)(cid:12)(cid:12)
˜
η 0 − 2 λ H1 (
[Equation]
0 ) (cid:12)(cid:12)(cid:12) ≤ η 1 / 2 − o (1) in , where ˜ η 0 :
= ˆ η (1 − ˆ η ˆ λ ) (cid:107) w 0 (cid:107) 22 is the effective LR for the ﬁrst step.
Theorem 4.7.
Under Assumptions 4.5 and 4.6, for GD+WD (1) with sufﬁciently small intrinsic learning rate η in :=
ˆ η ˆ λ , if we follow the above initialization scheme for some ζ 0 ∈ Γ , then with probability 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 solution exists), 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.8 (Magnitude of Oscillation) .
As suggested by Figure 4, θ t actually oscillates around the 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 as the “projection” of θ t onto the manifold Γ , in the sense that the displacement x t := θ t − φ t is approximately 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) in and | ˜ η t − 2 /λ H1 ( φ t ) | ≤ η 1 / 2 − o (1) in at t = 0 .
In our induction, we maintain the induction hypothesis that these two EoS conditions continue to hold for all t ≥ 0 .
[Equation]
[Figure]
Period-Two Oscillation.
A key insight in our proof is that after a few initial steps, θ t is oscillating around φ t along the ± v H1 ( θ ) directions, where v H1 ( θ ) is a unit top eigenvector of H ( θ ) and is chosen 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 ) for h 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 5d 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 )
≈ θ t − ˜ η t H ( φ t ) x t by Taylor expansions of ∇L and Π : R D \ { 0 } → S D − 1 .
We can further show that φ 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 eigenvalue 1 − ˜ η t λ H1 ( φ t )
≈ − 1 .
The remaining eigenvalues are { 1 − ˜
η t λ H i ( φ t ) } Di =2 , where λ H i ( φ t ) is the i -th largest eigenvalue of H ( θ t ) , and they lie in the range ( − 1 , 1] since λ H i ( φ t ) ∈ [0 , λ H1 ( φ t )) .
Using a similar analysis to power method, we show that x t quickly aligns to the direction of ± v H1 ( φ t ) after a few 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 ) , just as the initial condition in power method.
This is where we need the initial random perturbation.
Oscillation Drives φ t to Move.
This period-two oscillation is the driving power to push φ t to move on the manifold.
The main idea here is to realize that the oscillation direction deviates slightly from 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 of oscillation, φ t +2 ≈
φ t − 4 h 2 t ∇ Γ log λ H1 ( φ t )
+ O ( η 1 .
5 − o (1) in ) , which resembles two steps of gradient descent 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 know how fast the sharpness reduction is.
Our analysis is inspired by an empirical study from Lobacheva et al.
[84], which reveals a periodic behavior of gradients and effective learning rates in training normalized nets with weight decay.
In our theoretical setting, we capture this periodic behavior by showing that h t and ˜ η t do evolve periodically.
See Figures 5c and 5d for an example.
The key is that ˜ η t changes as an adaptive gradient method: ˜ η t increases when gradient is small and decreases when gradient is large (due to the effect of WD; see Figures 3a and 3b), and in our case the gradient norm scales as | h t | since ∇L ( θ t )
≈ h t λ H1 ( φ t ) v H1 ( φ t ) .
By the power method approximation, h t +2 ≈ (1 − ˜ η t λ H1 ( φ t ))
2 h t , so | h t | decreases when ˜ η t < 2 /λ H1 ( φ t ) .
But | h t | cannot 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 indeﬁnitely either, since ˜ η t decreases when | h t | is sufﬁciently large.
A period ﬁnishes when ˜ η t < 2 /λ H1 ( φ t ) holds again.
In our theoretical analysis, we connect this periodic behavior with a 1-dimensional Hamiltonian system (see Appendix H.2), and show that 2 h 2 t in each step can be approximated by its average value in the period without incurring a large error.
Further calculations show that this average value is approximately
[Equation]
the learning rate in the ﬂow
[Equation]
3) multiplied with η in .
We can therefore 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, since the above intuitive argument holds as long as the effective LR changes adaptively with respect to gradient change.
Based on this intuition, an important notion called Quasi-RMSprop scheduler is proposed.
For a PGD method, a learning rate scheduler is a rule for changing the effective LR in each step, and Quasi-RMSprop is a speciﬁc class of schedulers we deﬁne, including the way that the effective LR changes in GD+WD on scale-invariant loss (if viewed as PGD).
Our proof is done in a uniﬁed way that works as long as the effective LR changes in each step according to a Quasi-RMSprop scheduler.
As a by-product, a similar theorem can be proved for GD (without projection) on non-scale-invariant loss if the LR changes as a Quasi-RMSprop in each step.
For example, 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 and y i ∈ R are inputs and regression targets.
We study the over-parameterized case where d (cid:29) n , and we 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 Φ(
[Equation]
; w , γ, β ) :=
[Equation]
+ β , where µ 1 , σ 1 are the mean and variance of { w (cid:62)
x i } ni =1 over the whole dataset 3 , and the bias term b is cancelled out due to BN.
Note that Φ( x ; w , γ, β ) is still a linear function with respect to x .
Let µ x ∈ R d and Σ x ∈ R d
×
d be the mean and covariance of the input data { 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 our analysis, we ﬁx β, γ to be non-trainable constants so that the mean and variance of Φ ’s output match with those of { y i } ni =1 , that is, we set β = µ y and γ = σ y to be the mean and standard deviation of y i over 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 the initialization.
Moreover, the coefﬁcients ( ˜ w , ˜ b ) associated with w ∗ (deﬁned in (4) ) are the optimal solution 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 this is deceptive because in scale-invariant nets the regularization effect of WD is not explicit.
This result also challenges conventional view of optimization.
GD is usually viewed as a discretization of its continuous counterpart, gradient ﬂow (GF), and theoretical insight for the discrete update including 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.
As discussed in [77], adding WD only performs a time-rescaling on the GF trajectory on scale-invariant loss, but does not change the point that GF converge to if we project the trajectory onto the unit sphere.
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 best knowledge, this result is the ﬁrst concrete example where even with arbitrarily small LR, GD can still generalize better than GF under natural settings.
Discussion
Experimental Veriﬁcation of Sharpness Reduction.
Besides Figures 1 and 2, Appendix P. provides additional matrix completion experiments with different data size, and Appendix P. provides CIFAR-10 experiments with ResNet-20.
In all these experiments, we observed that GD continues to improve the test accuracy even after ﬁtting the training set, and this phenomenon is correlated with the decreasing trend of spherical sharpness.
See also Appendix P.3 for the validation for the periodic behavior we analyze in theory.
Ablation Studies on Normalization and Weight Decay.
Our theoretical analysis crucially relies on the interplay between normalization and WD to establish the sharpness-reduction ﬂow.
We also conducted ablation studies on normalization and WD to highlight the importance of this interplay.
First, if normalization is removed, the spherical sharpness becomes undeﬁned, and we do not know if
GD implicitly minimizes any sharpness measure.
But even if a similar measure does exist, it cannot be strongly related to generalization, because we can verify that the test accuracy becomes very bad without normalization ( 56 . 8% on CIFAR-10, Figure 14), and continuing training after ﬁtting the training set no longer improves test accuracy.
Second, if WD is removed, the analysis in Arora et al.
[7] guarantees convergence in the stable regime, and we can verify that the spherical sharpness and test accuracy stop changing when the loss is small.
The ﬁnal test accuracy is stuck at 66 .
4% (Figure 15), whereas training with WD leads to 84 .
Explaining the Progressive Sharpening and EoS Phenomena.
Cohen et al.
[24] conducted extensive empirical studies on the dynamics of GD in deep learning (without weight decay),
formally w t +1 ← w t − ˆ η ∇ ˜ L ( w t ) .
They observed the progressive sharpening phenomenon: λ 1 ( ∇ 2 ˜ L ( w t )) tends to increase so long as it is less than 2 / ˆ η .
Then they observed that the training typically enters the EoS regime, which they deﬁne as a regime that (1) λ 1 ( ∇ 2 ˜ L ( w t )) hovers right at, or just above 2 / ˆ η ; and (2) the training loss ˜ L ( w t ) goes up and down over short timescales, yet still decreases in the long-term run.
A recent research trend focuses on explaining the progressive sharpening and
EoS phenomena [1, 91, 8, 18].
Our work corresponds to an important special case where ˜ L ( w ) is a scale-invariant loss with L 2 -regularization, namely L ( w ) + ˆ λ 2 (cid:107) w (cid:107) 22 .
By analyzing the interplay between normalization and WD, the ﬁrst part of our results (Section 4.1) attributes progressive sharpening to norm change, and the second part (Section 4.2) justiﬁes in theory that the training can make progress in the EoS regime.
See Appendix G.3 for more discussion.
Conclusions and Future Work
We exhibited settings where gradient descent has an implicit bias to reduce spherical sharpness in training neural nets with normalization layers and weight decay, and we veriﬁed experimentally this sharpness-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 tracks a sharpness-reduction ﬂow.
We note that in practice a decrease in spherical sharpness is observed even with moderate LR and even before getting close to a minimizer manifold.
Explaining these phenomena is left for future work.
Now we list some other future directions.
The ﬁrst is to generalize our results to SGD, where the sharpness measure may not be the spherical sharpness and could depend on the structure of gradient noise.
Second, to understand the beneﬁt of reducing spherical sharpness on speciﬁc tasks, e.g., why does reducing spherical sharpness encourage low-rank on matrix completion with BN (Figure 1)?
Third, to study sharpness-reduction bias for neural net architectures that are not scale-invariant on all parameters (e.g., with certain unnormalized layers).
Acknowledgements
This work is funded by NSF, ONR, Simons Foundation, DARPA and SRC.
ZL is also supported by
Microsoft Research PhD Fellowship.
B A General Theory for a Broader Class of Adaptive Gradient
Methods
B.
RMSprop
and
Quasi-RMSprop Scheduler
B
B
B.3.
Spherical Optimization . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.3.
Full
Space
Optimization . . . . . . . . . . . . . . . . . . . . . . . . . . .
C PAC-Bayes Bounds Based on Spherical Sharpness
D Additional Preliminaries
D.1 Additional Notations . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.2 Scale-Invariant Functions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.3 Polyak-Łojasiewicz Condition . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.3.
Full Space
Optimization . . . . . . . . . . . . . . . . . . . . . . . . . . .
D.3.
Spherical
Optimization . . . . . . . . . . . . . . . . . . . . . . . . . . . .
E Supplementary Material for Appendix B.
E.1 Proof for Theorem B.7 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
E.2 Proof for Theorem B.9 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
F Details of the 3D Example
G.1 Proof for Descent Lemma . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
G.2 Proof for Theorem 4.4: GD Eventually Enters the EoS Regime . . . . . . . . . . .
G.3 Connection to the EoS Regime in Cohen et al.’s
Deﬁnition . .
H Proof Outlines of Our Theorems on Sharpness Reduction
H.1 Additional Notations . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
H.2 RMS-drift Process: Introduction . . . . . . . . . . . . . . . . . . . . . . . . . . .
H.3 Reduction to RMS-drift Process: The Case of Full Space Optimization . . . . . . .
H.3.1 Construction of Working Zones . . . . . . . . . . . . . . . . . . . . . . .
H.3.2 Good Initialization . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
H.3.3 Alignment Phase . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
H.3.4 Drifting Phase . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
H.4 Reduction to RMS-drift Process: The Case of Spherical Optimization . . . . . . .
H.4.1 Construction of Working Zones . . . . . . . . . . . . . . . . . . . . . . .
H.4.2 Good Initialization . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
H.4.3 Alignment Phase . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
H.4.4 Drifting Phase . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
H.5 RMS-drift Process: Analysis . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
H.6 Finalizing Proofs . .
I
Lemmas for Working Zones
I.
Construction of Working Zones . .
I.
Gradient Flow Projection . . . . . . . . . . . . . . . . . . . . . . . . . . .
I.2.
Full Space Optimization . . . . . . . . . . . . . . . . . . . . . . . . . . .
I.2.
Spherical Optimization . . . . . . . . . . . . . . . . . . . . . . . . . . . .
J
Lemmas
for
Gradient Descent
J.
Full Space Optimization . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
J.
Spherical Optimization . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
K Lemmas for Quasi-RMSprop Schedulers
L
Reduction
to
RMS-drift
Process: The Case of Full Space Optimization
L.1 Good Initialization . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
L.2 Alignment Phase . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
L.3 Drifting Phase . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
M Reduction to RMS-drift Process: The Case of Spherical Optimization
M.1 Good Initialization . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
M.2 Alignment Phase . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
M.3 Drifting Phase . .
N.1 Conservation of Energy . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
N.2 Flow Approximation . .
O Proofs
for Linear Regression with Batch Normalization
P Experiments
P.
Validation of Sharpness Reduction on Matrix Completion . .
P.
Validation of Sharpness Reduction on CIFAR-10 . . . . . . . . . . . . . . . . . .
P.
Periodic Behaviors . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
P.
Ablation Study: Normalization . . . . . . . . . . . . . . . . . . . . . . . . . . . .
P.
Ablation Study: Weight Decay . . . . . . . . . . . . . . . . . . . . . . . . . . . .
P.
Ablation Study: Initial Effective Learning Rate . . . . . . . . . . . . . . . . . . .
Q Experiment Details
Q.1 Additional Details of CIFAR-10 Experiments . . . . . . . . . . . . . . . . . . . .
Q.2 Additional Details of Linear Regression Experiments . . . . . . . . . . . . . . . .
Q.3 Computing Spherical Sharpness . .
R Discussion on the Afﬁne Parameters of the Final BN
R.1 Squared Loss . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
R.2 Crossentropy Loss . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
A Additional Related Works
Sharpness Measures and Parameter Rescaling.
To best capture the generalization performance, the measure of sharpness should give the same value whenever the function represented by the neural net is the same.
As mentioned in the introduction, λ 1 ( ∇ 2 L ( w )) does not satisfy this property for normalized nets because it is sensitive to weight rescaling.
The spherical sharpness takes care of scale-invariance w.r.t.
all parameters, but it is certainly not the only rescaling symmetry in deep nets.
However, to the best of our knowledge, the spherical sharpness is the only measure that provably decreases in training normalized nets.
Although many other sharpness measures may take care of more symmetries [123, 124, 112, 99, 37], it is unclear whether GD/SGD is implicitly reducing them.
SGD Noise Helps to Escape Sharp Minima.
It has been a folklore that noise in stochastic gradient helps escapes sharp local minima.
With the simpliﬁcation of assuming the loss is quadratic and treating SGD as its canonical continuous SDE approximation, Zhu et al.
[127], Xie et al.
[121] showed that anisotropic noise ( e.g. , the noise covariance is equal to Hessian) has a better escape efﬁciency of SGD out of a sharp minimizer, in comparison with isotropic noise.
Under the same assumptions, Ibayashi and Imaizumi [54] proved the exponential escape efﬁciency without assuming
SDE reaches the stationary distribution.
Kleinberg et al.
[64] showed that SGD can escape sharp local minima assuming one-point convexity.
While all previous escaping analysis of SGD are based on the continuous approximation, another approach called stability analysis is able to show that SGD cannot converge to sharp local minima when learning rate is larger than some threshold [117, 118, 89].
Implicit Bias of GD.
There are mainly two types of implicit bias result for GD, where the ﬁrst type of result applies essentially to the continuous limit of gradient descent, namely gradient ﬂow, and tolerates error discretization and stochasticity when the learning rate is sufﬁciently small [107, 106, 87, 59, 43, 42, 76, 100, 4, 22, 79, 88, 101, 108, 38].
The analyses of GD/SGD based on Neural
Tangent Kernel (NTK) also essentially belong to this type, because though the analysis in NTK regime tolerates stochasticity and ﬁnite learning rate, GD/ SGD do not learn different solutions compared to gradient ﬂow.
Such works includes (but are not limited to)
This type of result typically relates the generalization quality to the initialization of GD, and GD in such regimes cannot escape from bad local minima once reaching there.
In contrast, the second type of results, to which the current paper belongs, fundamentally relies on the discrete nature of gradient descent.
For example, Barrett and Dherin [11] showed that for small LR, gradient descent is approximately equal to gradient ﬂow minimizing a new objective, i.e. , the original objective plus η times squared norm of gradients.
Kong and Tao [66] studied a special class of “multiscale” loss functions, and they showed that large learning rate introduces chaos to the dynamics of GD and provides a mechanism to escape local minima.
Wang et al.
[116] proved that for the matrix factorization problem, GD with a large learning rate has an implicit bias towards a solution with balanced matrix norms.
Stability analysis [117, 118, 94, 89] also belongs to this type.
Comparison with Arora et al.
The paper by Arora et al.
[8] is probably the most related work among the second type of the implicit bias results.
They assume that there is a smooth function L which satisﬁes certain regularity conditions around the minimizer manifold (including that Hessian is maximally non-degenerate on the manifold), and show that running normalized GD on L or GD on √ L with sufﬁciently small LR tracks a deterministic ﬂow on the minimizer manifold and decreases the largest eigenvalue of ∇ 2 L , which has a similar ﬂavor to our result.
However, our setting (scale invariant loss + WD) is more natural and their result and technique do not apply to our setting because there is no minimizer under their deﬁnition, not to mention manifolds.
To show the spherical sharpness decreases, we have to develop new proof techniques, including connecting the dynamics to a 1-dimensional Hamiltonian system.
Another difference is that our analysis applies to gradient descent directly, but the analysis by Arora et al.
[8] requires injecting stochastic noise to gradients.
Edge of Stability.
Cohen et al.
[24] provided an extensive empirical study showing that GD typically occurs at the Edge of Stability (EoS), where the top eigenvalue of Hessian is approximately 2 / LR and the descent lemma does not guarantee the loss to decrease.
Ahn et al.
[1] explored the dynamics of GD in the EoS regime through insightful experiments, and they attribute the EoS phenomenon to the lack of ﬂat stationary points near GD trajectory and the existence of a subset near minima that is forward invariant under GD update.
Ma et al.
[91] proved the EoS phenomenon for a class of loss functions that are decomposable as a sum of 1-dimensional functions with subquadratic growth.
Chen and Bruna [18] provided detailed analyses for the EoS phenomenon on two-layer single-neuron net and matrix factorization.
The aforementioned work by Arora et al.
[8] analyzes normalized GD on L or GD on √ L also in the EoS regime, where the latter case corresponds to a class of loss functions that grow approximately linearly near minima (e.g., the absolute value function | x | ).
The evolution of the eigenvalues of Hessian has also been studied for SGD [40, 58, 73].
Our work can be seen as a theoretical explanation of the EoS phenomenon for scale-invariant loss functions with L 2 -regularization (see also Appendix G.3), which are a more broad and natural class of training loss in deep learning compared with those being studied in previous theoretical analyses.
B A General Theory for a Broader Class of Adaptive Gradient Methods
Our main theorem for GD+WD on scale-invariant loss, Theorem 4.7, is actually a corollary of a more general theorem that holds for any PGD on S D − 1 with effective learning rates changing adaptively according a speciﬁc kind of update rules.
We name this kind of update rules as Quasi-RMSprop
Scheduler , where the name is due to its similarity to RMSprop [49].
One of the main reasons that we deﬁne this concept is that proving a theorem for quasi-RMSprop schedulers in general is no harder than proving that only for GD+WD on scale-invariant loss, and sometimes the math involved is more simple and elegant when we analyze the dynamics at a higher level through quasi-RMSprop.
In this section, we introduce this key notion, Quasi-RMSprop scheduler.
To motivate it, we ﬁrst recall the update rule of RMSprop and present one of its variants with a single scalar learning rate.
Then we formally introduce the notion of quasi-RMSprop schedulers, which is a class of rules for setting effective learning rates similarly as RMSprop.
Then we prove a recursive formula for the effective learning rates when GD+WD optimizes a scale-invariant loss, and we categorize it as an instance of quasi-RMSprop scheduler.
After introducing this key notion, we then present our general theorem that holds for any PGD, or even
GD, as long as the effective learning rate is set by a quasi-RMSprop scheduler in each step.
Examples include Scalar RMSprop (Corollary B.14) and GD+WD on scale-invariant loss (Theorem 4.7).
B.1 Scalar RMSprop and Quasi-RMSprop Scheduler
The usual RMSprop algorithm maintains a vector v t storing the moving average of the squared gradients for every coordinate, i.e., v t +1 ← β v t + (1 − β )( ∇L ( θ t ))
(cid:12) 2 , where g (cid:12) 2 stands for the vector obtained by squaring g coordinatewise.
When updating the training parameter, RMSprop divides the usual GD update by the square root of the moving average coordinatewise, i.e., θ t +1 ← θ t − η √ v t +
+ (cid:15) (cid:12) ∇L ( θ t ) , where (cid:15) is a small constant to avoid division by zero.
Here all the addition, division, square root operations are coordinatewise, and (cid:12) stands for coordinatewise multiplication.
Now we consider a variant of RMSprop, which we call Scalar RMSprop , where the moving average v t is maintained as a scalar storing the moving average of the squared norm of gradients, rather than a vector that stores the moving averages separately for each coordinate.
Deﬁnition B.1 (Scalar RMSprop, Standard Form) .
Scalar RMSprop is an iterative method with the following update rule:
[Equation]
Besides that v t is changed from a vector to a scalar, another difference is that the gradient is divided by √ v t in (5), while in the usual RMSprop it is √ v t +1 .
In fact, our later analysis also applies if √ v t is changed to √ v t +1 , but the version with √ v t leads to simpler math.
An alternative view of Scalar RMSprop is to regard it as GD with time-varying learning rate ˜ η t , that is, θ t +1 ← θ t − ˜ η t ∇L ( θ t ) , where ˜ η t is the learning rate being used at the t -th step, which we call the effective learning rate at step t .
We view the effective learning rate here as the output of a learning rate scheduler .
In this view, we call the learning rate scheduler for Scalar RMSprop as RMSprop scheduler .
Deﬁnition B.2 (Gradient-Based Learning Rate Scheduler) .
A gradient-based learning rate scheduler is an algorithm H that reads from a stream of vectors g 0 , g 1 , g 2 , . . .
, For all t ≥ 0 , H outputs a real number ˜
η t :=
H ( g 0 , g 1 , . .
. , g t ) as soon as H ﬁnishes reading the ﬁrst t + 1 vectors.
At each step of gradient descent equipped with a gradient-based learning rate scheduler H , the gradients in training are revealed one by one as an input stream to H , and ˜ η t produced by H is used as the effective learning rate at step t .
Deﬁnition B.3 (RMSprop Scheduler) .
Given a constant ˜ v 0 > 0 , a base learning rate η and a decay rate β as hyperparameters, an RMSprop scheduler H RMS is a gradient-based learning rate scheduler that reads from a stream of vectors g 0 , g 1 , . . . , and generates effective learning rates ˜
η t := H RMS ( g 0 , . . . , g t ) for all t ≥
according to the following recursion:
[Equation]
where
[Equation]
Deﬁnition B.4 (Scalar RMSprop, Alternative Form) .
Scalar RMSprop is gradient descent with a
RMSprop scheduler H RMS .
[Equation]
where
[Equation]
Note that we use ˜ v t as an internal state of H RMS in Deﬁnition B.3.
It is easy to verify that ˜ v t is nothing but a reparameterization of v t : setting ˜ v t = v t /η 2 in Deﬁnition B.3 recovers the update rule in Deﬁnition B.1.
Now we introduce the notion of quasi-RMSprop scheduler, which is a class of gradient-based learning rate schedulers that have update rules similar to RMSprop scheduler.
At ﬁrst reading, one can just ignore the details and regard quasi-RMSprop scheduler as an RMSprop scheduler with negligible perturbations when β → 1 .
Deﬁnition B.5 (Quasi-RMSprop Scheduler) .
A quasi-RMSprop scheduler H is a gradient-based learning rate scheduler parameterized by a base learning rate η and a decay rate β , satisfying the following properties:
( η, β ) is allowed to take value from a hyperparameter space P H ⊆ (0 , + ∞ ) ×
There exist thresholds η max , β min , a continuous function δ : (0 , + ∞ ) → R , and a polynomial P : R → R such that the following holds.
If ( η, β ) ∈ P H and η
< η max , β > β min , for any input stream g 0 , g 1 , g 2 , . .
. , there exists a sequence of positive real numbers ˜ v 0 , ˜ v 1 , ˜ v 2 , . . .
such that ˜ η t :=
H ( g 0 , . .
. , g t ) satisﬁes the following two inequalities for all t ≥ 0 :
[Equation]
where ¯
[Equation]
The real number ˜ v t is called the moment estimate at step t associated with the input stream and effective learning rates.
It is clear that a RMSprop scheduler can be seen as a quasi-RMSprop scheduler with the same hyperparameters η, β , and δ ≡ 0 .
B.2 Reformulation of GD+WD on Scale-Invariant Loss via Quasi-RMSprop Scheduler
Now we reformulate GD+WD on scale-invariant loss as PGD with a quasi-RMSprop scheduler.
Recall that we say a loss function L ( w ) is scale-invariant if L ( c w ) =
L ( w ) for all c > 0 , and
Lemma 3.1 converts GD+WD on scale-invariant functions to Projected Gradient Descent (PGD) on unit sphere, i.e., θ t +1 ← Π( θ t − ˜ η t ∇L ( θ t )) , where ˜ η t :=
ˆ η
[Equation]
− ˆ η ˆ λ )
[Equation]
w
[Equation]
22 is the effective learning rate at step t .
However, the evolution of ˜ η t over time is unclear unless we know how the parameter norm (cid:107) w t (cid:107) 2 changes.
Similar as the above analysis for Scalar RMSprop, where we decompose Scalar RMSprop as the GD method and an RMSprop scheduler, here we abstract the evolution of ˜ η t as a learning rate scheduler, which we call GWSI scheduler (name picked from the initials of G D+ W D on S cale-
I nvariant loss) and view GD+WD on scale-invariant loss as PGD on S D − 1 with effective learning rates being set by a GWSI scheduler.
Deﬁnition B.6 (GWSI Scheduler) .
Given a constant ˜ v 0 > 0 , a base learning rate η and a decay rate β as hyperparameters, a GWSI scheduler H GWSI is a gradient-based learning rate scheduler that reads from a stream of vectors g 0 , g 1 , . . . , and generates effective learning rates ˜
η t :=
H GWSI ( g 0 , . . . , g t ) for all t ≥
according to the following recursion:
[Equation]
where
[Equation]
Theorem B.7.
For gradient descent (1) with learning rate ˆ η > 0 and weight decay ˆ λ > 0 on scale-invariant function L ( w ) , let θ t :=
w t (cid:107) w t (cid:107) 2 be the direction of w t and ˜ η t :
=
ˆ η
[Equation]
1 −
[Equation]
η
[Equation]
w t
[Equation]
22 be the effective learning rate at time t .
Then ˜ η t evolves exactly the same as the GWSI scheduler with hyperparameters ˜ v 0 :
= (1 − ˆ η ˆ λ ) 2 (cid:107) w 0 (cid:107) 42 ˆ η, β := (1 − ˆ η ˆ λ ) 4 , η := (cid:112)
[Equation]
β
[Equation]
− 1) / 2 , and we can write the dynamics of θ t as:
[Equation]
where
[Equation]
Remark B.8.
To the best of our knowledge, this particular form of the recursion formula (Deﬁnition B.6) for the effective learning rates of GD+WD on scale-invariant loss does not appear in prior works, but some variants have been studied before.
Arora et al.
[7] derived a similar formula when the weight decay is zero.
Li et al.
[78] obtained a Stochastic Differential Equation (SDE) for SGD+WD on scale-invariant functions, which is essentially a continuous approximation of our formulation when ˆ η ˆ λ → 0 .
Tanaka and Kunin [110] studied the continuous-time approximation of the momentum method with WD on scale-invariant loss and establish a connection to adaptive gradient methods.
We defer the proof for Theorem B.7 to Appendix E.1.
The formula of GWSI scheduler clearly resembles RMSprop scheduler: the only difference is that GWSI scheduler has an extra 1 4 β ˜ v t (1 − β ) 2 ¯ g 4 t term, which is negligible when (1 − β ) 2 is small.
This intuition is formalized through the deﬁnition of quasi-RMSprop.
It can be easily seen that GWSI scheduler is a quasi-RMSprop scheduler with the same hyperparameters η, β .
In fact, we can use Theorem B.7 as a basis to obtain a better way to write GD+WD on scaleinvariant loss as PGD with a quasi-RMSprop scheduler, in which η, β are expressed more simply in terms of the intrinsic learning rate η in :
=
ˆ η ˆ λ .
The main idea is that β = 1 − 4 η
in + O ( η 2in ) and η = √ 2 η in · (1 + O ( η in )) when η in is small, and we can absorb these approximation errors into the error bounds in Deﬁnition B.5.
We defer the details to Appendix E.2.
Theorem B.9.
There exists a quasi-RMSprop scheduler H QRMS with hyperparameter space { ( η, β ) : β
= 1 − 2 η 2 , η ∈ (0 , 1 √ 2 ) } such that the following holds for GD+WD on scale-invariant loss.
If the intrinsic learning rate η in :=
ˆ η ˆ λ lies in the range (0 , 1 / 4) , then we can set the hyperparameters of H QRMS to be η = √ 2 η in , β = 1 − 4 η in so that the dynamics of θ t can be written as
[Equation]
where
[Equation]
B.3 Main Results for GD/PGD with Quasi-RMSprop Scheduler
As mentioned in Section 4.2.3, a key step in our analysis is to show the oscillation and periodic behaviors of the effective learning rates.
But these behaviors can show up in other algorithms besides
GD+WD on scale-invariant loss.
In the following, we ﬁrst present Theorem 4.7, which holds for any PGD on S D − 1 with effective learning rates being set by a quasi-RMSprop scheduler.
We focus on the case where η > 0 is small and β =
C b
η 2 + O ( η 4 ) for some constant C b > 0 .
Our theorem shows that if θ 0 is initially near a local minimizer manifold Γ and is in the EoS regime, then θ t approximately tracks a sharpness-reduction ﬂow deﬁned as follows:
[Equation]
Besides, we also generalize our theorem to any GD with quasi-RMSprop scheduler, where Scalar
RMSprop serves as an important example.
The proof outlines of these theorems are given in Appendix H, and the proof details are spread over
Appendices I to N.
B.3.1 Spherical Optimization
Consider Projected Gradient Descent (PGD) on the unit sphere S D − 1 :
[Equation]
where Π : w (cid:55)→ w (cid:107) w (cid:107) 2 is the projection operator, and the effective learning rate ˜ η t is set by a quasi-RMSprop scheduler with base learning rate η > 0 and decay rate β
= 1 − C b
η 2 + O ( η 4 ) .
Recall that the dynamic is in the EoS regime if ˜ η t ≈ 2 /λ ( t )max (see Lemma 4.1).
When θ 0 is around a local minimizer ζ 0 and ˜ η 0 is generated by a quasi-RMSprop scheduler, the condition of being in
EoS is essentially 1 √ ˜ v 0 ≈ 2 λ H1 ( ζ 0 ) .
In the following theorem, we show that PGD with quasi-RMSprop scheduler evolves as (8) if it is initially in the EoS regime.
Initialization Scheme.
Given a local minimizer ζ 0 ∈ Γ and a hyperparameter α 0 , we initialize the initial direction θ 0 and initial moment estimate ˜ v 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 take any value in [exp( − α 20 ) η, α 0 η ] ; then set ˜ v 0 to be any positive value that satisﬁes (cid:12)(cid:12)(cid:12) 1 √ ˜ v 0 −
[Equation]
Theorem B.10.
Under Assumptions 4.5 and 4.6, for PGD described as (9) and initialized as above scheme for some ζ 0 ∈ Γ and some 1 ≤ α 0 ≤ η − o (1) , with probability 1 − O ( α 0 η (cid:112) log(1 /η )) , the trajectory of θ t approximately tracks a sharpness-reduction ﬂow ζ : [0 , T ] → Γ that starts from ζand evolves as the ODE (8) (if solution exists), in the sense that (cid:107) θ t − ζ ( tη 2 ) (cid:107)
[Equation]
for all 0 ≤ t ≤ T/η 2 .
Theorem 4.7 is a direct corollary of Theorem B.10, as GD+WD on scale-invariant loss can be seen as
PGD with GWSI scheduler.
Proof for Theorem 4.7.
Theorem B.9 implies that GD+WD on scale-invariant loss with LR ˆ η and
WD ˆ λ can be seen as a PGD on S D − 1 with quasi-RMSprop scheduler, where the base learning rate and decay rate of this scheduler is η = √ 2 η in and β = 1 − 4 η in .
Therefore, β = 1 − C b
η 2 for
C b = 2 .
Now we apply Theorem B.10 to prove Theorem 4.7.
It is easy to translate the initial conditions of
Theorem 4.7 to Theorem B.10 with α 0 =
η − o (1) =
η − o (1) in .
Let ˆ ζ ( t ) be the sharpness-reduction ﬂow deﬁned as in (8) with C b = 2 and horizon ˆ T := 2 T .
[Equation]
[Equation]
Then with probability 1 − O ( α 0 η (cid:112) 1 /δ )
= 1 − O ( η 1 − o (1) ) , θ t approximately tracks ˆ ζ ( τ ) in the sense that (cid:107) θ t − ˆ ζ ( tη 2 ) (cid:107)
=
O ( η 1 / 2 − o (1) )
for all 0 ≤ t ≤ ˆ T/η 2 .
Replacing η with √ 2 η in gives (cid:107) θ t − ˆ ζ (2 tη in ) (cid:107)
= O ( η 1 / 4 − o (1) in ) for all 0 ≤ t ≤ T/η in , and the success probability becomes 1 − O ( η 1 / 2 − o (1) in ) .
We can ﬁnish the proof by noting that ˆ ζ (2 τ ) is just ζ ( τ ) deﬁned in (3).
B.3.2 Full Space Optimization
Now we present our general theorem for GD on R D with quasi-RMSprop scheduler.
Here GD with quasi-RMSprop scheduler can be written as:
[Equation]
where the effective learning rate ˜ η t is set by a quasi-RMSprop scheduler with base learning rate η >and decay rate β = 1 − C b η 2 + O ( η 4 ) .
Similar to the spherical case, we assume that there is a manifold Γ consisting of local minimizers, but now we are not assuming scale-invariance.
Assumption B.11.
The loss function L : R D
→ R is C 4 -smooth.
Γ is a C 2 -smooth, D Γ -dimensional submanifold of S D − 1 for some 0 ≤ D Γ ≤ D , where every θ ∈ Γ is a local minimizer of L on R D and rank( H ( θ ))
=
D − D Γ .
The following assumption is essentially the same as Assumption 4.6 except that now Γ is deﬁned differently.
Assumption B.12.
For all θ ∈ Γ , λ H1 ( θ ) > λ H2 ( θ ) .
That is, the top eigenvalue of H ( θ ) is unique.
Based on Assumptions B.11 and B.12 above, we study GD starting in the EoS regime.
Similar to the spherical case, we focus on the case where θ 0 is around a local minimizer ζ ∈ Γ , and EoS is the regime
in which ˜ η t ≈ 2 λ H1 ( ζ 0 ) , which essentially means 1 √ ˜ v 0 if the effective LR is set by a quasi-RMSprop scheduler.
In the theorem below, we show that GD with quasi-RMSprop scheduler tracks the sharpness-reduction ﬂow deﬁned in (8).
Note that the ODE here is the same as the spherical case, but Γ is deﬁned differently.
Initialization Scheme.
Given a local minimizer ζ 0 ∈ Γ and a hyperparameter α 0 , we initialize the initial parameter θ 0 and initial moment estimate ˜ v 0 as follows: draw ξ ∼ N ( 0 , σ 20 I /D ) from
Gaussian and set w 0 ← ζ 0 + ξ , where σ 0 can take any value in [exp( − α 20 ) η, α 0 η ] ; then set ˜ v 0 to be any positive value that satisﬁes (cid:12)(cid:12)(cid:12) 1 √ ˜ v 0 −
[Equation]
Theorem B.13.
Under Assumptions B.11 and B.12, for GD described as (11) and initialized as above scheme for some ζ 0 ∈ Γ and some 1 ≤ α 0 ≤ η − o (1) , with probability 1 − O ( α 0 η (cid:112) log(1 /η )) , the trajectory of θ t approximately tracks a sharpness-reduction ﬂow ζ : [0 , T ] → Γ that starts from ζand evolves as the ODE (8) (if solution exists), in the sense that (cid:107) θ t − ζ ( tη 2 ) (cid:107) 2 = O ( α 20 η 1 / 2 log(1 /η )) for all 0 ≤ t ≤ T/η 2 .
A direct corollary is that Scalar RMSprop follows the sharpness-reduction ﬂow (8), since the RMSprop scheduler is a quasi-RMSprop scheduler.
Corollary B.14.
The statement of Theorem B.13 holds for Scalar RMSprop if the decay rate is set to β
= 1 − C b
η 2 + O ( η 2 ) for some constant C b > 0 .
C PAC-Bayes Bounds Based on Spherical Sharpness
In this section we give the PAC-Bayes bound for generalization error using spherical sharpness (Deﬁnition 1.1).
We will start with our setting and then recap the classic PAC-Bayes theorem in [92].
The main result in this section is Theorem C.2.
Setting.
Let (cid:96) ( w , z ) be the loss of parameter w on data point z and assume (cid:96) max = sup z ∈Z , w ∈ S D − 1 (cid:96) ( w , z ) < ∞ .
Let S :
= { z i } ni =1 where z i are sampled independently.
Different to the previous notation, we use
L S to denote the empirical loss on training dataset S where we run op- timization algorithms and
L to denote the population loss.
Let ρ
( L S
[Equation]
Since S D − 1 is compact and (cid:107)∇ 3 L S ( · ) (cid:107) 2 is continuous, ρ 3 ( L S ) is ﬁnite.
Theorem C.1 (PAC-Bayes theorem [92]) .
Given any distribution P on R D , with at least 1 − δ probability over the randomness of the dataset S , for all distribution Q on R D , it holds that
[Equation]
Now we are ready to state the main theorem in this section, Theorem C.2, which shows that small spherical sharpness λ 1 ( ∇ 2 L S ( θ )) leads to small generalization error.
Theorem C.2.
For any σ ≤ 1 2+2 √ (ln n ) /D , with at least 1 − δ probability over the randomness of the dataset S , where S := { z i } ni =1 and every z i is sampled independently, for
any θ ∈ S D − 1 ,
E (cid:15) ∼N (
[Equation]
Thus with the standard assumption in [37] that E (cid:15) ∼N ( 0 ,σ 2 I D /D )
L ( θ + (cid:15) )
≥ L ( θ ) , we have the same upper bound for L ( θ ) − L S ( θ ) .
To prove Theorem C.2, we will ﬁrst need the following lemma.
The proof is standard so it is omitted.
Lemma C.3.
Let Q = N ( µ Q , σ 2 Q I D ) and P = N ( µ P , σ 2 P I D ) , we have that
[Equation]
Proof of Theorem C.2.
Let Q :=
N ( θ , σ 2 I D /D ) , P := N ( 0 , σ 2 I D /D ) .
Then D KL ( Q (cid:107) P ) =
D 2 σ 2 .
By Theorem C.1 we have
[Equation]
Let h ( σ ) :
=
(1 + (cid:112) (ln n ) /D ) σ .
By assumption it holds that h ( σ ) ≤ 12 .
Thus by Lemma 1 in [69], we have for any positive t :
[Equation]
Therefore, with probability 1 − 1 / √ n , we have
[Equation]
Thus for any θ with (cid:107) θ (cid:107) 2 =
[Equation]
we have
[Equation]
By Taylor expansion,
[Equation]
Note that (cid:107) θ + λ (cid:15) (cid:107) 2 ≥ (cid:107) θ (cid:107) 2 − λ (cid:107) (cid:15) (cid:107) 2 ≥ 12 for all λ ∈ [0 , 1] , it holds that (cid:13)(cid:13) ∇ 3 L S ( θ + λ (cid:15) )
(cid:13)(cid:13) =
[Equation]
θ + λ
[Equation]
L
[Equation]
≤ 8 ρ 3 ( L S ) , we have
[Equation]
Thus we conclude that
[Equation]
which completes the proof.
D Additional Preliminaries D.1 Additional Notations
We use cl( M ) to denote the closure of a set M .
For θ ∈ R D and M ⊆ R D , we use d 2 ( θ , M ) :
= inf {(cid:107) θ − φ (cid:107) 2 : φ ∈ M} to denote the L 2 -distance from θ to M .
For θ ∈ R D and (cid:15) ≥ 0 , B (cid:15) ( θ ) :
=
{ θ (cid:
) ∈ R D : (cid:107) θ − θ (cid:48) (cid:107) 2 < (cid:15) } is the open (cid:15) -ball centered at θ .
For a set M ⊆ R D ,
M (cid:15) :
= (cid:83) θ ∈M B (cid:15) ( θ ) is the (open) (cid:15) -neighborhood of M .
All the manifolds in our paper refer to manifolds without boundary.
For a manifold Γ , we use T p ( Γ ) and N p ( Γ ) to denote the tangent and normal space of Γ at a point p ∈ Γ .
Given a function f : R D → R n and vectors x , v ∈ R D , we use ∂f x [ v ] to denote the directional derivative ∂f x [ v ] :=
lim t → 0 1 t f ( x + t v ) , which also equals to the Jacobian of f at x multiplied with v .
We use ∂ 2 f x
[ v , u ] to denote the second-order derivative ∂ ( ∂f x [ v ])
x [ u ] .
For a real-valued function L : R D → R , we use ∇L ( θ ) for gradient, ∇ 2 L ( θ ) for Hessian.
For third-order derivatives of L , we deﬁne ∂ 3 L θ [ v , u ] := ∂ 2 ( ∇L ) θ [ v , u ] ∈ R D .
If L is C 4 -smooth, then the following Taylor expansion holds for ∇L :
[Equation]
D.2 Scale-Invariant Functions
The following lemma summarizes a few important properties of scale-invariant loss that have been exploited in previous works [7, 77].
For completeness, we include a proof here.
Lemma D.1.
The following hold for a twice-differentiable scale-invariant 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 ; 2.
The Hessian matrix is ( − 2) -homogeneous
i
e., ∇ 2 L ( c w ) =
c − 2 ∇ 2 L ( w ) for all c > 0 . 3.
∇ 2 L ( w ) w = −∇L ( w ) .
Proof.
Taking gradients with respect to w on both sides of L ( c w )
=
L ( w ) gives ∇L ( c w ) =
c − 1 ∇L ( w ) .
Taking gradients again proves that ∇ 2 L ( c w ) =
c − 2 ∇ 2 L ( w ) .
Taking derivative with respect to c on both sides of L ( c w )
=
L ( w ) gives (cid:104)∇L ( c w ) , w (cid:105) = 0 .
Taking c = 1 gives (cid:104)∇L ( w ) , w (cid:105) = 0 .
Finally, we take gradients with respect to w , then ∇ 2 L ( w ) w + ∇L ( w ) = 0 .
Now we provide proofs for Lemma 3.1 and Theorem 4.2 in Section 3.
Proof for Lemma 3.1.
By (1) and deﬁnition of θ t ,
[Equation]
Since Π( x ) =
Π( c x ) for all c > 0 , we can divide w t +1 by (1 − ˆ η ˆ λ ) (cid:107) w t (cid:107) 2 and obtain θ
[Equation]
Note that ∇L ( w t ) =
[Equation]
t
[Equation]
2 L ( θ t ) by Lemma D.1.
So we can further rewrite the above formula:
[Equation]
which proves the lemma by deﬁnition of ˜ η t .
Proof of Theorem 4.2.
The ﬁrst claim is directly from the Lemma D.2 of [80].
The second claim about ˆ η holds by scrutinizing their proof.
D.3 Polyak-Łojasiewicz Condition
Deﬁnition D.2 (Polyak-Łojasiewicz) .
For a loss function L ( θ ) and a constant µ > 0 , we say that L satisﬁes µ -Polyak-Łojasiewicz condition (or µ -PL for brevity) on a set U if
[Equation]
for all θ ∈ U .
D.3.1 Full Space Optimization
Theorem D.3.
Let L : R D → R be a C 3 -smooth function, and Γ be a C 1 -smooth, D Γ -dimensional submanifold of R D , where every θ ∈ Γ is a local minimizer of L and rank( ∇ 2 L ( θ ))
=
D − D Γ .
If
Z is a compact subset of Γ , then there exist (cid:15) > 0 , µ > 0 such that cl( Z (cid:15) )
∩ Γ is compact and L satisﬁes µ -PL on Z (cid:15) .
Proof.
Since Z is compact and Γ is a submanifold of R D , we can choose a small δ > 0 such that
N := Γ ∩ cl( Z δ ) is compact.
It can be shown that there exists an open neighborhood U of the compact submanifold N such that for every θ ∈ U , the nearest point on N , P ( θ ) :
= arg min {(cid:107) θ − φ (cid:107) 2 : φ ∈ N} , exists and is unique [36].
We choose (cid:15) <
δ/ 2 to be small enough so that (cid:15) <
δ/ 2 and Z (cid:15) ⊆
U .
For θ ∈ Z (cid:15) , P ( θ ) lies in the interior of the manifold N , since d 2 ( P ( θ ) , Z ) ≤ (cid:107) P ( θ ) − θ (cid:107) 2 + d ( θ , Z ) <
δ .
Then it must hold that θ − P ( θ ) ∈ N P ( θ ) ( N ) ; otherwise the differential of (cid:107) θ − φ (cid:107) 22 on N is non-zero at φ = P ( θ ) , which contradicts to the fact that P ( θ ) is the nearest point to θ on N .
Since L ∈ C 3 , and ∇ 2 L ( θ ) is of constant rank D − D Γ on the compact manifold N , there exist λ min > 0 , λ max > 0 such that λ D − D Γ ( ∇ 2 L ( θ )) ≥ λ min and λ 1 ( ∇ 2 L ( θ )) ≤ λ max for all θ ∈ N .
Also by
L ∈ C 3 and compactness of N , there exists C 3 > 0 such that the following Taylor expansions hold for all θ ∈ Z (cid:15) ,
[Equation]
Then L satisﬁes µ -PL on Z (cid:15) for µ :=
[Equation] [Equation]
+ C
[Equation]
) , which is positive if we choose (cid:15) to be small enough in the beginning.
Theorem D.4.
If θ ∗ is a local minimizer of a C 2 -smooth function L : R D → R and L is µ -PL on an open neighborhood U of θ ∗ , then for any θ 0 sufﬁciently close to θ ∗ , a gradient ﬂow d θ d t = −∇L ( θ ) starting with θ 0 converges to a point θ ∞ as t → + ∞ and (cid:107) θ ∞ − θ ∗ (cid:107)
= O ( (cid:107) θ 0 − θ ∗ (cid:107) 2 ) .
Proof.
Let T :=
inf { t : θ ( t ) / ∈ U } .
For all t < T ,
[Equation]
Since ( L ( θ ) − L ( θ ∗ )) 1 / 2 ≤ 1 √ 2 µ (cid:107)∇L ( θ ) (cid:107) 2 , we have
[Equation]
Integrating on both sides proves the following
[Equation]
So if (cid:107) θ 0 − θ ∗ (cid:107) 2 is small enough, then T = + ∞ and θ ( t ) converges to a point in U as t → + ∞ .
Moreover,
[Equation]
D.3.2 Spherical Optimization
Theorem D.5.
Let L : R D → R be a C 3 -smooth scale-invariant function, and Γ be a C 1 -smooth, ( D Γ − 1) -dimensional submanifold of S D − 1 , where every θ ∈ Γ is a local minimizer of L on S D −and rank( ∇ 2 L ( θ ))
=
D − D Γ .
If Z is a compact subset of Γ , then there exist (cid:15) > 0 , µ > 0 such that cl( Z (cid:15) )
∩ Γ is compact and L satisﬁes µ -PL on Z (cid:15) ∩ S D − 1 .
Proof.
Let Γ (cid:48) :
= { ν θ : θ ∈ Γ , ν > 0 } .
Then Γ (cid:48) is a C 1 -smooth, D Γ -dimensional submanifold of R D , where every θ ∈ Γ is a local minimizer of L on R D and rank( ∇ 2 L ( θ ))
=
D − D Γ .
By
Theorem D.5, there exist (cid:15) > 0 , µ > 0 such that L satisﬁes µ -PL on Z (cid:15) , so it satisﬁes µ -PL on
Z (cid:15) ∩ S D − 1 .
Theorem D.6.
If θ ∗ is a local minimizer of a C 2 -smooth and scale-invariant function L : R D \{ 0 } → R and
L is µ -PL on an open neighborhood U of θ ∗ on S D − 1 , then for any θ 0 ∈ S D − 1 sufﬁciently close to θ ∗ , a gradient ﬂow d θ d t = −∇L ( θ ) starting with θ 0 converges to a point θ ∞ as t → + ∞ and (cid:107) θ ∞ − θ ∗ (cid:107)
= O ( (cid:107) θ 0 − θ ∗ (cid:107) 2 ) .
Proof.
Since L is µ -PL on U ⊆ S D − 1 and scale-invariant, we know L is µ 2 -PL on an open set in R D ,
U (cid:48)
= { w : w (cid:107) w (cid:107) 2 ∈ U, (cid:107) w (cid:107) 2 ∈
[ 1 √ 2 , √ 2] }
The proof is completed by applying Theorem D.4.
E Supplementary Material for Appendix B. E.1 Proof for Theorem B.
Lemma E.1.
In the setting of Theorem B.7,
[Equation]
Proof.
Recall that w t +
=
(1 − ˆ η ˆ λ ) w t − ˆ η ∇L ( w t ) .
By scale-invariance, (cid:104)∇L ( w t ) , w t (cid:105) = 0 .
Then by Pythagorean theorem (or Gougu Theorem),
[Equation]
Since ∇L ( w t ) = 1 (cid:107) w t (cid:107) 2 ∇L ( θ t ) by scale-invariance, we can rewrite the last term ˆ η 2 (cid:107)∇L ( w t ) (cid:107) 22 as ˆ η 2 (cid:107) w t (cid:107) 22 (cid:107)∇L ( θ t ) (cid:107) 22 , which implies (13).
Proof for Theorem B.7.
Squaring both sides of (13), we have
[Equation]
Let β :=
(1 − ˆ η ˆ λ ) 4 , ˜ v t :=
1˜ η 2 t
=
(1 − ˆ η ˆ λ ) 2 ˆ η(cid:107) w t (cid:107) 42 .
Then ˜ v t +
[Equation]
= (1 − ˆ
[Equation]
ˆ
[Equation] [Equation]
where the last equality uses the deﬁnition of β and ˜ v t .
Let η := (cid:113) 1 − β 2 β
and ¯ g t :=
(cid:107)∇L ( θ t ) (cid:107) 2 /η .
Then
[Equation]
which is exactly the update rule of GWSI scheduler.
E.2 Proof for Theorem B.
Proof.
We specify the quasi-RMSprop scheduler H QRMS as follows.
Let P H QRMS :
= { ( η, β ) :
β
= 1 − 2 η 2 , η ∈ (0 , 1 √ 2 ) } .
Given hyperparameters η, β , we deﬁne β (cid:48) :
= (1 − 14 (1 − β )) 4 and η (cid:48) :
= (cid:112) ( β (cid:48) − 1) / 2 .
Then H QRMS produces the effective LRs as a GWSI scheduler with ( η (cid:48) , β (cid:48) ) :
˜
[Equation]
where
[Equation]
When η = √ 2 η in and β = 1 − 4 η in , it is easy to see that H QRMS produces the same effective LRs as GD+WD on scale-invariant functions (Theorem B.7).
Now we only need to verify that H QRMS is indeed a quasi-RMSprop scheduler.
When β is close enough to 1 , we have β (cid:48)
=
β + O ((1 − β ) 2 ) , β (cid:48) ≥ 1 / 2 , η (cid:48)
=
η · (1+ O (1 − β )) .
Let
C 0 be a constant such that η/η (cid:48) ≤ C 0 , |
[Equation] [Equation]
β
[Equation]
− β | ≤ C 0 (1 − β ) 2 .
Let ¯ g t := (cid:107) g t (cid:107) 2 /η .
Then ˆ g t =
( η/η (cid:48) )¯ g t , and thus ˆ g
[Equation]
g
[Equation] [Equation]
We only need to verify that (cid:12) (cid:12) ˜ v t +1 − (cid:0)
β ˜ v t +
[Equation]
1 − β
[Equation]
¯ g 2 t (cid:1)(cid:12) (cid:12) ≤ δ (˜ v t ) ·
[Equation]
1 − β ) 2 · P (¯ g t ) for some continuous function δ and some polynomial P .
[Equation]
For the ﬁrst term, we have
[Equation]
For the second and third terms, we have
[Equation]
Finally we can conclude
[Equation]
which veriﬁes that H QRMS is indeed a quasi-RMSprop scheduler.
F Details of the 3D Example
In this section, we give more details for Figure 4.
The loss L ( w ) is constructed as follows.
First, we deﬁne the following scale-invariant function:
[Equation]
By taking gradient on S 2 , one can easily see that the minimum is attained when ( x, y ) points to (1 , 1) in direction.
In other words, the minimizer manifold of F is Γ := { ( x, y, z ) ∈ S 2 : x = y > 0 } .
Then we ﬁx an orthogonal matrix U (generated randomly) and deﬁne L : R 3 \ { 0 } → R , w (cid:55)→ F ( Uw ) , i.e., the function F after an orthogonal transformation.
For plotting the ﬁgure, we transform the coordinates back to the domain of F .
The initial point is w 0 = (0 . 3 , 1 . 3 , 1 . 2) in the domain of F .
We run gradient descent on L with LR ˆ η = 0 .
5 and WD ˆ λ = 0 .
It can be seen from the ﬁgure that θ t does not stop moving after reaching eventually
One can check that the Hessian matrix of F at ( x, y, z ) ∈ Γ is
H ( x, y, z
[Equation]
= x 2 + y(cid:34)− 1−     (cid:35)
= 1 − z(cid:34)− 1−
(cid:35) .
Therefore, the spherical sharpness is controlled by | z | .
The smaller the absolute value of z , the ﬂatter the minimizer.
And the ﬂattest one is ζ ∗ , which has z -coordinate being zero.
This matches with our theory of sharpness-reduction bias as GD+WD moves along Γ and oscillates near ζ ∗ in the end.
G Supplementary Material for Section 4. G.1 Proof for Descent Lemma
Proof for Lemma 4.1.
By Taylor expansion,
[Equation]
which proves the lemma.
G.2 Proof for Theorem 4.4: GD Eventually Enters the EoS Regime
Let R :
=
C 0 (ˆ λ ˆ η ) 1 / 2 be a radius so that θ ∗ is a minimizer of L on U :=
B R ( θ ) ∩ S D − 1 and µ -PL holds within U , where C 0 is a large constant to be speciﬁed later.
Let λ max := sup { λ H1 ( θ ) : θ ∈ U } .
Then we know that λ max = λ H1 ( θ ∗ )
+ O
(ˆ λ ˆ η ) .
Let T 0 be the largest number so that θ t ∈ U for all t 0 ≤ t ≤ T 0 , We deﬁne a potential function Ψ( θ ) :
= (cid:112) L ( θ ) − L ( θ ∗ ) .
Lemma G.1.
If θ t ∈ U and ˜ η t < 2 λ max for some t 0 ≤ t < T 0 , then
[Equation]
Proof.
By descent lemma,
[Equation]
Then
[Equation]
By µ -PL, (cid:107)∇L ( θ t ) (cid:107) 2 ≥ √ 2 µ · Ψ( θ t ) .
Combining these together proves the lemma.
Lemma G.2.
There exists C 2 =
O (1) such that (cid:107)∇L ( θ ) (cid:107) 22 ≤ C 2 ( L ( θ ) − L ( θ ∗ )) for all θ ∈ U .
Proof.
It is equivalent to give an upper bound for sup θ ∈ U { G ( θ ) } , where G ( θ ) :=
(cid:107)∇L ( θ ) (cid:107) 22 L ( θ ) −L ( θ ∗ ) .
Since L ∈ C 2 , G is continuous in its domain.
So it sufﬁces to upper bound G ( θ ) around every singular point, i.e., around every minimizer of L on U .
And for every minimizer θ (cid:48) ∈ U , we can do
Taylor expansions for L and ∇L to show that G ( θ ) is indeed bounded by O (1) around θ (cid:48) .
Proof for Theorem 4.4.
Recall that ˜ η t 0 ≤ 2 ρ 2 < 2 λ H1 ( θ ∗ ) .
Let δ :=
2 − ˜ η t 0 λ max ∈ (0 , 2) and T 1 be the largest number so that ˜ η t ≤ 2 − δ/ 4 λ max for all t 0 ≤ t ≤ T 1 .
By Lemma G.1, for all t 0 ≤ t < min { T 0 , T 1 } ,
[Equation]
Telescoping the sum we have √ 2 µ 16 δ (cid:80) t − 1 τ = t 0 ˜ η τ (cid:107)∇L ( θ τ ) (cid:107) 2 ≤ Ψ( θ 0 ) .
By smoothness of L , Ψ( θ t 0 )
=
O ( (cid:107) θ t 0 − θ ∗ (cid:107) 2 ) =
O (( ˆ λ ˆ η ) 1 / 2 ) .
So for all t 0 ≤ t ≤ min { T 0 , T 1 } ,
[Equation]
which implies that T 0 > T 1 or T
= T 1 = + ∞ if we choose C 0 to be large enough.
By Theorem B.7, ˜ η t can be seen as the output of a GWSI scheduler with β =
(1 − ˆ λ ˆ η ) 4 = 1 − Θ(ˆ λ ˆ η ) and η = (cid:112) ( β − 1 − 1) / 2 =
O ((ˆ λ ˆ η ) 1 / 2 ) .
Then by the update rule,
[Equation]
where
[Equation]
For θ ∈ U , we have (cid:107)∇L ( θ ) (cid:107) 22 ≤ C 2 ( L ( θ ) − L ( θ ∗ )) by Lemma G.2.
So for all t 0 ≤ t ≤ T 0 , we have
[Equation]
O
[Equation] [Equation]
Then ˜ η − 2 t +1 = β ˜ η − 2 t + (1 − β ) · O (1) , which implies ˜ η t ≥ Ω(1) for some t =
t
+ O ( 1 1 − β log(˜ η − 2 0 )) .
Therefore, we can infer that it must hold for some steps t that ˜ η t ∈ (cid:104)
c min , 2 − δ/ 4 λ max (cid:105) , where c min is some constant.
As ¯ g t ≥ 0 , the update rule (14) also implies ˜ η − 2 t +1 ≥ β ˜ η − 2 t , or equivalently ˜
η t +1 ≤ (1 − ˆ η ˆ λ ) − 2 ˜ η t .
This suggests that the number of steps such that ˜ η t ∈ (cid:104) 2 − c λ max , 2 − δ/ 4 λ max (cid:105) , is at least Ω( 1 / (ˆ η ˆ λ ) ) .
When ˜ η t does lie in this range, by Lemma 4.1 we have
[Equation]
Combining with µ -PL gives
[Equation]
Thus the loss decays by a constant factor in every step (the factor is in (0 , 1) as we can choose c min as small as we want).
As this process lasts for at least Ω( 1 / (ˆ η ˆ λ ) ) steps, the loss ﬁrst decreases to
L ( θ ∗ )
+ O
((ˆ η ˆ λ ) 10 ) after O (log 1ˆ η ˆ λ ) steps, then it stays small until T 1 .
Now we show that T 1 is ﬁnite.
By Lemma G.2 and (14), when the loss is L ( θ ∗ )
+ O
((ˆ η ˆ λ ) 10 ) the effective LR steadily grows as ˜ η t +
=
(1 − ˆ η ˆ λ ) − 2 ˜ η t + o (1) .
So at some step t , it must hold that ˜ η t > 2 − δ/ 4 λ max , which proves T 1 < + ∞ .
Let C 1 be a large constant, and T 2 be the largest number so that ˜ η t < 2 − 2 C 1 (ˆ η ˆ λ ) 1 / 2 λ max for all t 0 ≤ t ≤ T 2 .
By Lemma G.1, for all T 1 ≤ t < min { T 2 , T 0 } ,
[Equation]
(cid:107)∇L
[Equation]
Telescoping the sum gives (ˆ η ˆ λ ) 1 / 2 (cid:80) t − 1 τ = T 1 ˜ η τ (cid:107)∇L ( θ τ ) (cid:107) 2 ≤ O (Ψ( θ T 1 ))
≤ O ((ˆ η ˆ λ ) 5 ) , where the last inequality is due to L ( θ T 1 ) =
L ( θ ∗ )
+ O
((ˆ η ˆ λ ) 10 ) .
This shows that (cid:107) θ t − θ T 1 (cid:107)
= O ((ˆ η ˆ λ ) 4 . 5 ) , and thus (cid:107) θ t − θ ∗ (cid:107)
= O ((ˆ η ˆ λ ) 1 / 2 ) by triangle inequality.
Now we have T 0 > T 2 or T 0 =
T 2 = + ∞ when C 0 is chosen to be large enough.
We can ﬁnish the proof with a similar argument as for T 1 to show that T 2 cannot be inﬁnite either.
G.3 Connection to the EoS Regime in Cohen et al.’s Deﬁnition
Now we elaborate how our deﬁnition of EoS ˜
η t ≈ 2 /λ ( t )max is related to the original deﬁnition of EoS in Cohen et al.
In their work, they studied the dynamics of GD (without weight decay).
When the loss is ˜ L , the update rule is given by w t +1 ← w t − ˆ η ∇ ˜ L ( w t ) .
They deﬁne the EoS regime as a regime in which (1) λ 1 ( ∇ 2 ˜ L ( w t )) hovers right at, or just above 2 / ˆ η ; and (2) the training loss ˜ L ( w t ) goes up and down over short timescales, yet still decreases in the long-term run.
View I: Rewriting as GD.
We can write GD+WD on scale-invariant loss as GD on scale-invariant loss with L 2 -regularization, i.e., GD on ˜ L ( w ) :
= L ( w ) + ˆ λ 2 (cid:107) w (cid:107) 22 .
Now we show that ˜ η t ≈ 2 λ ( t )max is essentially the same as ˆ η ≈
[Equation]
It sufﬁces to show ˜ η t · λ ( t )max ≈ ˆ η · λ 1 ( ∇ 2 ˜ L ( w t )) .
When the gradient is small, θ t does not move far in one step, then λ ( t )max ≈ λ 1 ( ∇ 2 L ( θ t )) .
By scale-invariance, λ 1 ( ∇ 2 L ( θ t ))
= (cid:107) w t (cid:107) 22 · λ 1 ( ∇ 2 L ( w t )) (Lemma D.1).
Recall that ˜ η t := ˆ η (1 − ˆ η ˆ λ ) (cid:107) w t (cid:107) 22 .
Then we have ˜ η t · λ ( t )max ≈
[Equation]
≈
[Equation]
Note that ˆ η · λ 1 ( ∇ 2 ˜ L ( w t ))
=
ˆ η · λ 1 ( ∇ 2 L ( w t ))
+ ˆ η ˆ λ .
When ˆ η ˆ λ is small, we can then conclude that ˜ η t · λ ( t )max ≈ ˆ η · λ 1 ( ∇ 2 ˜ L ( w t )) .
Now we show below that our main theorem on sharpness-reduction bias implies the second condition in Cohen et al.
[24]’s deﬁnition, that is, ˜ L decreases in the long-term run.
For the regularizer ˆ λ 2 (cid:107) w t (cid:107) 22 , note that (cid:107) w t (cid:107) 22 ≈ ˆ η/ ˜
η t ≈ 12 ˆ ηλ 1 ( ∇ 2 L ( θ t )) in the EoS regime.
So (cid:107) w t (cid:107) 2
as well as the regularizer is decreasing due to the sharpness-reduction bias (Theorem 4.7).
The scale-invariant part L ( θ t ) is not always decreasing, but now we show that its time average can be upper bounded by the time average of norm squared.
By Lemma E.1 and Lemma D.1, we have
[Equation]
Since (cid:107) w t (cid:107) 22 decreases in the long run, we know that for any long enough time window T 0 to T 1 − 1 , (cid:80) T 1 −
[Equation] [Equation]
=
[Equation]
2ˆ
[Equation]
w t (cid:107) 22 .
Further due to the alignment between the gradient and the top eigenvalue of the Hessian in the EoS regime, we have L ( w t ) ≈
[Equation] [Equation]
≈
[Equation]
w
[Equation] [Equation]
Therefore, we conclude that the average loss over a long enough time window is always upper bounded by the average of squared weight norm, that is,
[Equation]
ˆ η (cid:107) w t (cid:107) 22 ≈
T 1 − T
T 1 − 1 (cid:88) t =
Tˆ λ 2 (cid:107) w t (cid:107) 22 , where the last step uses the fact that GD operates in EoS. So L ( w t ) decreases in the long-term run.
Combining the above two parts, we can conclude that the regularized loss ˜ L ( w t ) has a tendency to decrease in the long-term run.
View II: Generalizing EoS to PGD.
For a gradient-based method in general, 2 / ˆ η should be replaced to the maximum sharpness bound that the loss function is guaranteed to decrease through
Taylor expansions, e.g., Cohen et al.
[24] derived the bounds exactly for Polyak and Nesterov momentum in Appendix B of their paper.
In our deﬁnition, we view GD+WD on scale-invariant loss as PGD on S D − 1 , and thus we deﬁne the EoS regime for PGD as the regime where ˜ η t ≈ 2 /λ ( t )max , where λ ( t )max is the local upper bound of spherical sharpness (Lemma 4.1).
This captures the ﬁrst condition of Cohen et al.
[24]’s deﬁnition.
Repeating our argument in View I, we can show the second condition, namely the condition that the loss L ( θ t ) decreases in the long-term run.
Progressive Sharpening.
All the above discussion is about the EoS phenomenon.
Another phenomenon identiﬁed by Cohen et al.
[24] is progressive sharpening, which is the phenomenon that λ 1 ( ∇ 2 ˜ L ( w t )) tends to increase so long as it is less than 2 / ˆ η .
Theorem 4.4 in our paper justiﬁes this phenomenon in View II, i.e., if ˜ η t is less than λ H1 ( θ ∗ ) , then ˜ η t increases until it reaches 2 λ H1 ( θ ∗ ) .
The key insight in our analysis is that WD decreases the norm when gradient is small, and smaller norm leads larger λ 1 ( ∇ 2 ˜ L ( w t ))
This shows that the progressive sharpening phenomenon in our case can be well explained by the interplay between normalization and WD
H Proof Outlines of Our Theorems on Sharpness Reduction
In this section, we give proof outlines of Theorems B.10 and B.13.
The main proof idea for the spherical case is stated in Section 4.2.3, but technically it is easier to state and prove the lemmas for full space optimization.
Therefore, we present the full details for the full space case and omit certain details for the spherical case if they are similar to the full space case.
As mentioned in Section 4.2.3, a key ingredient in our proof is to show that the period-2 oscillation drives the parameter to move along the manifold.
For both full space and spherical cases, we project θ t onto the manifold Γ with a carefully-deﬁned projection function, φ t :=
Φ( θ t ) .
Then we show that after a period of oscillation ( 2 steps), the projection drifts from φ t to a new position φ t + approximately along the direction of ∇ Γ log λ H1 ( φ t ) .
To analyze the speed of each drift, we show that the oscillation can be tracked with two variables h ∈ R , u ∈ R , where h is related to the displacement of θ t from the manifold, u is related to the closeness of the current dynamic to the edge of stability.
We formally deﬁne a discrete process called RMS-drift process and show that the oscillations in both full space and spherical cases can be regarded as RMS-drift processes.
The rest of the section is organized as follows.
In Appendix H.1 we introduce some additional notations.
In Appendix H.2 we formally introduce the concept of RMS-drift process.
In Appendix H. we show how to reduce our analysis in the full space case into studying an RMS-drift process.
In
Appendix H.4 we show how to reduce our analysis in the spherical case into studying an RMS-drift process.
Finally, in Appendix H.5 we analyze the RMS-drift process and show that the projections of parameters in training can be tracked with the sharpness-reduction ﬂow deﬁned in (8).
H.1 Additional Notations
We follow the notations in Appendix D. We also need some additional notations in this section.
It is implied by the uniqueness of the top eigenvalue (Assumption B.12 or Assumption 4.6) and L ∈ Cthat λ H1 ( θ ) is C 2 -smooth on Γ , and we can construct a C 2 -smooth function v H1 ( θ ) such that v H1 ( θ ) is a unit top eigenvector of H ( θ ) on Γ .
Let P H0 ( θ ) be the projection matrix onto the null space of H t , and P H NZT ( θ ) be the projection matrix onto the space spanned by the eigenvectors of H t corresponding to non-zero and non-top eigenvalues ( λ (cid:54) = 0 , λ H1 ( φ t ) ).
For any local minimizer θ ∈ Γ , we deﬁne γ ( θ ) :
= 1 λ H1 ( θ ) min (cid:8) λ H1 ( θ ) − λ H2 ( θ ) , λ H D − D Γ ( θ ) (cid:9) to be the relative eigenvalue gap (between λ H and λ H2 , or between λ H D − D Γ and 0 ).
Then P H0 ( θ ) , P H NZT ( θ ) , γ ( θ ) are all C 2 -smooth on Γ .
Finally, we deﬁne µ ( θ ) :
= 2 λ H1 ( θ ) .
We deﬁne Φ( θ ) to be the convergence point of a gradient ﬂow of L starting from θ .
That is, the limit of ˜ θ ( t ) as t → + ∞ when ˜ θ ( t ) is described by the following ODE:
[Equation]
where ˜
[Equation]
We leave Φ( θ ) undeﬁned if the ODE does not converge to any point.
For GD/PGD with quasi-RMSprop scheduler (Deﬁnition B.5), the state at step t can be written as ( θ t , ˜ v t ) , where θ t is the trainable parameter in GD/PGD and ˜ v t is the moment estimate.
Whenever Φ( θ t ) exists and is in Γ , we deﬁne
[Equation]
Most importantly, we deﬁne two hidden variables ( h t , u t ) as follows:
[Equation]
Note that the h t deﬁned here differs with that in Section 4.2.3 by a factor of 1 /η .
We introduce this factor for the sake of convenience, and we will use the deﬁnition with the factor 1 /η only in our theoretical analysis.
H
Introduction
Deﬁnition H.1.
A drift state is described by a tuple S
= ( h, u, φ ) in the drift state space S :
= R × R × Γ .
We say that S is α -bounded if max {| h | , | u |} ≤ α .
Deﬁnition H.2.
Given two drift states S t = ( h t , u t , φ t ) , S t +2 =
( h t +2 , u t +2 , φ t +2 ) in the drift state space S , for learning rate η > 0 and hyperparameter C b > 0 , we say that the transition S t → S t +2 is a C 0 -RMS-drift transition if for all α ≥ 1 , as long as S t is α -bounded, S t +2 is close to an auxiliary state S (cid:48)
t +2 :
=
( h (cid:48)
t +2 , u (cid:48) t +2 , φ (cid:48) t +2 ) in the following sense:
[Equation]
For a sequence of states S 0 , S 2 , S 4 , . . .
, S 2 M , we say it is a C 0 - RMS-drift process if S t → S t +2 is a
C 0 -RMS-drift transition for all even numbers 0 ≤ t < 2 M .
In our analysis of GD/PGD with quasi-RMSprop scheduler, we can rewrite the dynamics as RMS-drift processes (after a few warm-up steps), where φ t :=
Φ( θ t ) is the gradient ﬂow projection of the parameter at step t onto Γ , and h t , u t are two hidden variables deﬁned in (15).
This RMS-drift process serves as an abstraction of the original dynamics that contains the minimal but sufﬁcient amount of information so that we can compute the continuous approximation for the trajectory of φ t .
In RMS-drift process, φ t evolves as gradient descent on Γ for minimizing log λ H1 ( θ ) , and the corresponding learning rate is changing with h t per step.
To obtain the ﬁnal ﬂow approximation (8), we need to sum up the learning rates over time.
An intuitive way to understand RMS-drift process is to use the following ﬁrst-order continuous approximation with time scaling ( h ( τ ) , u ( τ ) , φ ( τ ))
≈ ( h τ/η , u τ/η , φ τ/η ) and ignore all the second order terms O ( η 2 ) :
[Equation]
This approximation gives an important insight: h t and u t are changing much faster than φ t when η is small.
Therefore, we can analyze this ﬁrst-order approximation to obtain an average value of 2 η 2 h 2 t , and use this average value as the “effective” learning rate in the ﬂow approximation.
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
can be associated with a physical system in which a unit-mass particle moves in a potential well U ( x )
=
K 2 e 2 x − C b x without any energy loss; see (a).
log | h | can be seen as the position of the particle, and − u can be seen as the velocity.
This system must be periodic because of the conservation of energy, and it can be seen clearly from its phase portrait (b).
In fact, ( h, u ) forms a 1 -dimensional Hamiltonian system after a coordinate transformation and evolves periodically in the above ODE.
To see this, we can compute the time derivatives of (log | h | , − u ) while letting K :=
(cid:112)
[Equation]
b + (cid:107)∇
[Equation]
[Equation]
Consider a unit-mass particle in the system with position x ( t ) , velocity v ( t ) =
x (cid:48) ( t ) and acceleration a ( t ) = v (cid:48)
( t ) .
Suppose that this system has the potential energy U ( x ) :
=
K 2 e 2 x − C b x .
Then we can see that the position and velocity of this particle evolve exactly the same as (log | h | , − u ) !
This also shows that (log | h | , − u ) evolves periodically because 1 -dimensional Hamiltonian system with unimodal potential energy must be periodic.
See also Figure 6.
In Appendix H.5 below, we will use this observation to obtain the ﬂow approximation nicely.
But now we ﬁrst outline how to reduce the original dynamics to an RMS-drift process.
H.3 Reduction to RMS-drift Process: The Case of Full Space Optimization
Now we outline how to reduce the dynamics to an RMS-drift process in the setting of Theorem B.13, the main theorem for the case of full-space optimization.
H.3.1 Construction of Working Zones
Let Z := { ζ ( t ) : t ∈ [0 , T ] } ⊆
Γ be the set of points passed by the sharpness-reduction ﬂow (8).
Inspired by Arora et al.
[8], we construct a two-level nested working zone (
Z (cid:15) 0 , Z (cid:15) 1 ) , where Z (cid:15) 0 , Z (cid:15)are essentially the (cid:15) 0 - and (cid:15) 1 -neighborhoods of Z with (cid:15) 0 , (cid:15) 1 carefully chosen.
In our later analysis, we will ensure that θ t is always in Z (cid:15) 0 , and its gradient ﬂow projection φ t is always well-deﬁned and lies in Z (cid:15) 1 ∩ Γ .
The following lemma shows some important properties of the working zone in our construction.
We defer the proof to Appendix I.1.
Lemma H.3 (Working Zone Lemma) .
There exist 0 < (cid:15) 0 < (cid:15) 1 such that Z (cid:15) 0 , Z (cid:15) 1 satisfy the following:
cl( Z (cid:15) 1 ) ∩ Γ is compact; 2.
L satisﬁes µ PL -PL on Z (cid:15) 1 for some µ PL > 0 ; 3.
Φ is well-deﬁned on Z (cid:15) 0 , and Φ( θ ) ∈ Z (cid:15) 1 ∩ Γ for all θ ∈ Z (cid:15) 0 ; 4.
Φ is C 3 -smooth on Z (cid:15) 0 ; 5.
γ ( θ ) ≥ γ min holds uniformly on Z (cid:15) 1 ∩ Γ for some γ min > 0 .
It can be seen from Lemma H.3 that φ t , H t , µ t , U t , x t , h t , u t are all well-deﬁned as long as θ t ∈ Z (cid:15) 0 .
Now we deﬁne some useful notions in the EoS regime.
A state is α -bounded if θ t is close to Γ and the dynamic is in the EoS regime due to ˜ v t ≈ µ − 2 t .
A state is α -deviated if θ t is not too close to Γ .
Deﬁnition H.4 ( α -Bounded State) .
We say that the state ( θ t , ˜ v t ) at some step t is α -bounded if θ t ∈ Z (cid:15) 0 , (cid:107) x t (cid:107) 2 ≤ αη and | u t | ≤ α .
Deﬁnition H.5 ( α -Deviated State) .
We say that the state ( θ t , ˜ v t ) at some step t is α -deviated (from the manifold Γ ) if (cid:107) x t (cid:107) 2 ≥ η exp( − α 2 ) or θ t / ∈ Z (cid:15) 0 .
Next, we deﬁne a quantitative measurement for how much x t aligns to the top eigenvector of
H t .
Note that x t can be decomposed into the projections onto the top eigenspace, the null space, and the space spanned by non-zero and non-top eigenvectors; we can write this decomposition as x t = h t η v H 1 ( φ t ) + P H 0 ( φ t ) x t + P H NZT ( φ t ) x t .
The following lemma shows that P H NZT ( φ t ) is always negligible, so we characterize the alignment to the top eigenvector only through comparing ηh t v H1 ( φ t ) and P H NZT ( φ t ) x t .
Lemma H.6.
At any step t , if θ t ∈ Z (cid:15) 0 , then (cid:107) P H0 ( φ t ) x t (cid:107) 2 =
O ( (cid:107) x t (cid:107) 22 ) .
Proof.
Direct consequence of Lemma I.3.
Deﬁnition H.7 ( p -Misaligned State) .
We say that the state ( θ t , ˜ v t ) at some step t is at most pmisaligned (to the top eigenvector) if θ t ∈ Z (cid:15) 0 , (cid:107) P H NZT ( φ t ) x t (cid:107) 2 ≤ p · h t η .
H.3.2 Good Initialization
First, we show that the initialization satisﬁes some desirable properties with high probability.
The proof is deferred to Appendix L.1.
Lemma H.8.
There exists δ
= O ( α 0 η (cid:112) log(1 /η ))
such that the following holds.
With probability 1 − δ , the initial state is O
( α 0 (cid:112) log(1 /δ ))
-bounded, O ( α 0 + (cid:112) log(1 /δ ))
-deviated, at most O (1 /δ ) - misaligned, and satisﬁes (cid:107) φ 0 − ζ 0 (cid:107) 2 ≤ O ( α 0 η (cid:112) log(1 /δ )) .
H.3.3 Alignment Phase
At initialization, the state is not well-aligned to the top eigenvector.
But in the following, we show that it becomes at most O ( | h t | η ) -misaligned after only η − o (1) steps.
We defer the proofs to Appendix L.2.
The key lemma is the following, which gives good approximations for various important variables.
Lemma H.9.
For small enough base learning rate η , at any step t , if the state ( θ t , ˜ v t ) is α -bounded for some 1 ≤ α ≤ η − o (1) , then θ t +1 ∈ Z (cid:15) 0 , and
[Equation]
[Equation] [Equation]
Applying the above lemma through an induction proves the following theorem.
Theorem H.10.
There exists δ
= O ( α 0 η (cid:112) log(1 /η ))
and T 1 =
O (log 1 η + α 0 (cid:112) log(1 /δ ))
such that the following holds.
Let α max := α 0 (cid:112) log(1 /δ ) .
If the initial state is O ( α max ) -bounded, O ( α max ) - deviated, at most O (1 /δ )
-misaligned, and satisﬁes (cid:107) φ 0 − ζ 0 (cid:107) 2 ≤ O ( α max η ) , then at step t
= T 1 , the state is at most O ( | h t | η ) -misaligned while still being O ( α max ) -bounded, O ( α max ) -deviated, and satisfying (cid:107) φ t − ζ 0 (cid:107) 2 ≤ O ( α max η ) .
H.3.4 Drifting Phase
After the alignment phase, the state is now at most O ( | h t | η ) -misaligned.
Then we have the following lemma showing that ( h t , u t , φ t ) evolves as an O (1) -RMS-drift process.
We defer the proof to
Appendix L.3.
Lemma H.11.
For small enough base learning rate η , at any step t , if for some 1 ≤ α ≤ η − o (1) , the state ( θ t , ˜ v t ) is α -bounded and at most O
( | h t | η ) -misaligned, then θ t +2 ∈ Z (cid:15) 0 , and h t +2 =
[Equation]
In other words, ( h t , u t ) → ( h t +2 , u t +2 ) is an O (1) -RMS-drift transition if the state at step t is α -bounded and at most O
( | h t | η ) -misaligned.
H.4 Reduction to RMS-drift Process: The Case of Spherical Optimization
Now we outline how to reduce the dynamics to an RMS-drift process in the setting of Theorem B.13, the main theorem for the case of spherical optimization.
The basic logic is the same as the case of full space optimization, so we only list the new lemma and theorem statements here.
The proofs in this section are deferred to Appendices I.1 and M.
H.4.1 Construction of Working Zones
We still deﬁne Z := { ζ ( t ) : t ∈ [0 , T ] } ⊆ Γ .
The construction of working zone becomes the following.
We defer the proof to Appendix I.1.
Lemma H.12 (Working Zone Lemma) .
There exist 0 < (cid:15) 0 < (cid:15) 1 such that Z (cid:15) 0 , Z (cid:15) 1 satisfy the following:
cl( Z (cid:15) 1 ) ∩ Γ is compact; 2.
L satisﬁes µ PL -PL on Z (cid:15) 1 for some µ PL > 0 ; 3.
Φ is well-deﬁned on Z (cid:15)
∩ S D − 1 and Φ
( θ ) ∈
Z (cid:15) 1 ∩ Γ for all θ ∈ Z (cid:15)
∩ S D − 1 ; 4.
Φ is C 3 -smooth on Z (cid:15) 1 ∩ S D − 1 ; 5.
γ ( θ ) ≥ γ min holds uniformly on Z (cid:15) 1 ∩ Γ for some γ min > 0 .
In the working zone, we continue to deﬁne α -bounded states, α -deviated states, at most p -misaligned states following the same deﬁnitions as the full space case except that the deﬁnition of working zone is changed.
Deﬁnition H.13 ( α -Bounded State) .
We say that the state ( θ t , ˜ v t ) at some step t is α -bounded if θ t ∈ Z (cid:15)
∩ S D − 1 , (cid:107) x t (cid:107) 2 ≤ αη and | u t | ≤ α .
Deﬁnition H.14 ( α -Deviated State) .
We say that the state ( θ t , ˜ v t ) at some step t is α -deviated (from the manifold Γ ) if (cid:107) x t (cid:107) 2 ≥ η exp( − α 2 ) or θ t / ∈ Z (cid:15)
∩ S D − 1 .
Lemma H.15.
At any step t , if θ t ∈ Z (cid:15) 0 ∩ S D − 1 , then (cid:107) P H0 ( φ t ) x t (cid:107) 2 =
O ( (cid:107) x t (cid:107) 22 ) .
Proof.
Direct consequence of Lemma I.6.
Deﬁnition H.16 ( p -Misaligned State) .
We say that the state ( θ t , ˜ v t ) at some step t is at most pmisaligned (to the top eigenvector) if θ t ∈ Z (cid:15)
∩ S D − 1 , (cid:107) P H NZT ( φ t ) x t (cid:107) 2 ≤ p · h t η .
H.4.2 Good Initialization
Lemma H.17.
The same statement as Lemma H.8 holds for the spherical case.
H.4.3 Alignment Phase
Lemma H.18.
For small enough base learning rate η , at any step t , if the state ( θ t , ˜ v t ) is α -bounded for some 1 ≤ α ≤ η − o (1) , then θ t +1 ∈ Z (cid:15)
∩ S D − 1 , and (17) to (21) in the full space case continue to hold in the spherical case.
Theorem H.19.
The same statement as Theorem H.10 holds for the spherical case.
H.4.4 Drifting Phase
Lemma H.20.
For small enough base learning rate η , at any step t , if for some 1 ≤ α ≤ η − o (1) , the state ( θ t , ˜ v t ) is α -bounded and at most O
( | h t | η ) -misaligned, then θ t +2 ∈ Z (cid:15)
∩ S D − 1 , and (22) to (25) in the full space case continue to hold in the spherical case.
H.5 RMS-drift Process: Analysis
Now we outline how to obtain the ﬁnal ﬂow approximation (8) from the RMS-drift process.
We say that an RMS-drift process S 0 , . . .
, S 2 M is in the working zone if φ t ∈ Z (cid:15) 1 ∩ Γ for all even numbers 0 ≤ t ≤ 2 M .
We focus on RMS-drift processes in the working zone, and later we will show that the
RMS-drift processes of interest are indeed in the working zone.
First, we deﬁne a potential function that resembles the total energy (or Hamiltonian) in physics.
Deﬁnition H.21 (Energy) .
For a drift state S
= ( h, u, φ ) , we deﬁne the energy E ( S ) as follows:
[Equation]
If the energy is bounded by α 2 at some step t , then it is easy to see that the state at step t is
O ( α ) -bounded
and O ( α ) -deviated.
The ﬁrst key lemma is the conservation of energy in RMS-drift process, which shows that the energy is preserved for O (1 /η 2 ) steps.
The proof is deferred to Appendix N.1.
Theorem H.22.
For an O (1) -RMSdrift process S 0 , . . .
, S 2 M in the working zone, if E ( S 0 ) ≤ αfor some parameter 1 ≤ α ≤ η − o (1) and M =
O (1 /η 2 ) , then E ( S t ) =
O ( α 2 ) for all even numbers 0 ≤ t ≤ 2 M .
Next, we show that as long as the states are η − o (1) -bounded, the RMS-drift process tracks the sharpness-reduction ﬂow nicely.
We defer the proof to Appendix N.2.
Theorem H.23.
Let ζ : [0 , T ] (cid:55)→ Z (cid:15) 1 ∩ Γ be a sharpness-reduction ﬂow deﬁned in (8) .
For an
O (1) -RMS-drift process S 0 , . .
, S 2 M in the working zone, where M := (cid:98) T 2 η 2 (cid:99) , if (cid:107) φ 0 − ζ (0) (cid:107) 2 ≤
O ( α 2 η 1 / 2 ) , and S t is O ( α ) -bounded for all even numbers 0 ≤ t ≤ 2 M , then (cid:107) φ t − ζ ( tη 2 ) (cid:107) 2 ≤
O ( α 2 η 1 / 2 ) for all even numbers 0 ≤ t ≤ 2 M .
H.6 Finalizing Proofs
Proof for Theorem B.13.
First, we use Lemma H.8 to ensure a good initialization.
Then we apply
Theorem H.10 to show that the state at t =
T 1 is at most O ( | h t | η ) -misaligned, O ( α max ) -bounded,
O ( α max ) -deviated, and satisﬁes (cid:107) φ t − ζ 0 (cid:107) 2 ≤ O ( α max η ) .
By smoothness of ζ , we also have (cid:107) φ t − ζ ( tη 2 ) (cid:107) 2 .
Then we consider the even-indexed and odd-indexed steps separately and do an induction for each of them to show that the later dynamics (1) has bounded energy O ( α 2max )
(Theorem H.22); (2) stays at most O ( | h t | η ) -misaligned
(Lemma H.11); (3) follows the ﬂow (Theorem H.23).
The proof is done when the induction proceeds to T/η 2 .
Proof for Theorem B.10.
Same as above
but we invoke the spherical version of lemmas and theorems.
I Lemmas for Working Zones I.1 Construction of Working Zones
Proof for Lemma H.3.
By Assumption B.11 and Theorem D.3, there exists (cid:15) 1 such that Items 1 and 2 hold.
By Theorem D.4, we can choose (cid:15) (cid:48)
> 0 small enough so that any gradient ﬂow starting in
Z (cid:15) (cid:48) 0 moves at most O ( (cid:15) (cid:48) 0 ) in distance and converges in Z (cid:15) 1 .
We can further combine this with the results by Falconer [34] to show that Φ is C 3 -smooth in a neighborhood of Z .
Then we can choose (cid:15) 0 > 0 small enough so that (cid:15) 0 < (cid:15) (cid:48)
and Z (cid:15) 0 is a subset of that neighborhood, which ensures Itemsand 4.
Finally, Item 5 is directly implied by the compactness of cl( Z (cid:15) 1 ) ∩ Γ .
Proof for Lemma H.12.
We can use a similar argument as above but with Theorems D.5 and D.6.
I.2 Gradient Flow Projection
In the working zone, Φ( θ ) is well-deﬁned and C 3 -smooth.
Now we highlight some useful properties.
I.2.1 Full Space Optimization
The following two lemmas are from Li et al.
Lemma I.1.
For φ ∈ Z (cid:15)
∩ Γ , ∂ Φ φ [ x ] =
P H0 ( φ ) x , which also equals to the projection of x onto the tangent space T φ ( Γ ) at φ .
Lemma I.2.
For φ ∈ Z (cid:15)
∩ Γ and x ∈ N φ ( Γ ) , ∂ 2 Φ φ [ x , x ] = 0 .
The following lemma can be proved by Taylor expansion.
Lemma I.3.
For θ ∈ Z (cid:15) 0 and φ = Φ( θ ) , then (cid:107) ∂
Φ φ [ θ − φ ] (cid:107) 2 ≤ O ( (cid:107) θ − φ (cid:107) 22 ) .
Proof.
When (cid:107) θ − φ (cid:107) 2 is small enough, the linear interpolation of θ and φ lies in Z (cid:15) 1 .
By Taylor expansion, Φ( θ ) =
Φ( φ ) + ∂ Φ φ [ θ − φ ] + O ( (cid:107) θ − φ (cid:107) 22 ) .
In fact, Φ( θ ) =
φ =
Φ( φ ) .
So we can conclude ∂ Φ φ [ θ − φ ]
= O ( (cid:107) θ − φ (cid:107) 22 ) .
I.2.2 Spherical Optimization
In the spherical case, we have the following lemmas that highly resemble those in the full space case.
All these lemmas can be proved in the same manner: we can ﬁrst deﬁne Γ (cid:48) :
= { ν θ : θ ∈ Γ , ν ∈ (1 / 2 , 2) } and apply the counterparts of these lemmas in the full space case; then translate the results back to the spherical case.
Lemma I.4.
For φ ∈ Z (cid:15)
∩ Γ , ( I − φφ (cid:62) )
∂
Φ φ [ x ] =
( I − φφ (cid:62) )
P H0 x , which also equals to the projection of x onto the tangent space T φ ( Γ ) at φ .
Lemma I.5.
For φ ∈ Z (cid:15)
∩ Γ and x ∈ N φ ( Γ ) , if (cid:104) x , φ (cid:105) = 0 , then ∂
Φ φ [ x , x ] = 0 .
Lemma I.6.
For θ ∈ Z (cid:15) 0 ∩ S D − 1 and φ =
Φ( θ ) , then (cid:107) ∂
Φ φ [ θ − φ ] (cid:107) 2 ≤ O ( (cid:107) θ − φ (cid:107) 22 ) .
J Lemmas for Gradient Descent J.1 Full Space Optimization
Lemma J.1.
Under Assumption B.11, consider one step of gradient descent θ t +
=
θ t − ˜ η t ∇L ( θ t ) with effective learning rate ˜ η t .
Let φ ∈ Z (cid:15)
∩ Γ be a local minimizer.
If for some parameters α = η − o (1) and r = O ( α ) , (cid:107) θ t − φ (cid:107) 2 = O ( rη ) and ˜ η t = 2 λ H1 ( φ ) + O ( αη ) , then after one step, we have the following approximations for θ t +1 :
Zeroth-order approximation:
[Equation]
First-order approximation:
[Equation]
Second-order approximation:
[Equation]
Proof.
For the zeroth-order approximation, we expand L ( θ t ) around φ using Taylor expansion:
[Equation]
For the ﬁrst-order approximation, we expand L ( θ t ) around φ using Taylor expansion:
[Equation]
So
[Equation]
For the second-order approximation, we again use Taylor expansion to expand L ( θ t ) :
[Equation]
So we have
[Equation]
where the last equality uses r
= O ( α )
[Equation]
Lemma J
[Equation]
Under Assumption B.
[Equation]
consider two steps of gradient descent with effective learning rates ˜
η t and ˜ η t +1 :
[Equation]
Let φ ∈ Z (cid:15)
∩ Γ be a local minimizer.
If for some parameters α = η − o (1) and r = O ( α ) ,
[Equation] [Equation]
and ˜
[Equation]
) , then after two steps,
[Equation]
where ψ φ (ˆ x ) :=
[Equation]
Proof.
Let ˆ x τ := θ τ − φ for all τ ∈ { t, t + 1 , t + 2 } .
By Lemma J.1, we have ˆ x t +
= O ( rη ) and the following ﬁrst-order and second-order approximations for ˆ x t +1 :
[Equation]
Note that ˜ η t +1 =
˜ η t + O (
[Equation] [Equation]
r
[Equation]
η
[Equation]
=
[Equation] [Equation]
O ( αη ) .
By Lemma J.1 again, we have the following second-order approximation for ˆ x t +2 :
[Equation]
Now we combine the two steps together.
[Equation]
where the second equality uses I − ˜ η t +1 H ( φ )
=
U φ + O ( αη ) .
The following lemma characterizes the function ψ φ (ˆ x ) in Lemma J.2 when ˆ x is the top eigenvector of H ( φ ) .
We will use this property later in Appendix L.3.
Lemma J.3.
Under Assumptions B.11 and B.12, for φ ∈ Γ ,
[Equation]
where ψ φ is deﬁned as in Lemma J.2.
Moreover,
[Equation]
Proof.
Note that U φ v H1 ( φ ) = − v H1 ( φ ) .
Then we can rewrite ψ φ ( v H1 ( φ )) as follows:
[Equation]
Since (2 I −
[Equation]
−
[Equation]
φ
[Equation]
= 0 , we have (cid:104) ψ φ ( v H1 ( φ
[Equation]
, v H1 ( φ ) (cid:105) = 0 .
The projection matrix onto the tangent space of Γ at φ equals to P H0 ( φ ) by Lemma I.1.
Then
P
[Equation]
which proves the lemma.
J.2 Spherical Optimization
Lemma J.4.
Under Assumption 4.5, consider one step of projected gradient descent on S D − 1 , θ t
+1 =
Π( θ t − ˜ η t ∇L ( θ t )) with effective learning rate ˜ η t .
Let φ ∈ Z (cid:15)
∩ Γ be a local minimizer.
If for some parameters α = η − o (1) and r = O ( α ) ,
[Equation]
and ˜ η
[Equation]
then after one step, we have the following approximations for θ t +1 :
Zeroth-order approximation:
[Equation]
First-order approximation:
[Equation]
Second-order approximation:
[Equation]
Proof.
Let ˆ θ t
+
=
θ t − ˜ η t ∇L ( θ t ) .
Then θ t +1 =
[Equation]
Since ∇L ( θ t ) is perpendicular to θ t , (cid:107) ˆ θ t +1 (cid:107)
= (cid:112) (cid:107) θ t (cid:107) 22 + ˜ η 2 t (cid:107)∇L ( θ t ) (cid:107) 22 .
By Taylor expansion, ∇L ( θ t ) =
O ( rη ) .
Then the norm of ˆ θ t +1 can be estimated by
[Equation]
Then θ t +1 =
ˆ θ t +1 (cid:107) ˆ θ t +1 (cid:107)
=
ˆ θ t +
+ O ( r 2 η 2 ) .
Combining with Lemma J.1 proves the zeroth- and ﬁrst-order approximations.
To prove the second-order approximation, we need a tighter estimate for the norm of ˆ θ t +1 .
By
Taylor expansion, we have ∇L ( θ t ) =
H ( φ )( θ t − φ ) + O ( r 2 η 2 ) .
Squaring both sides gives
[Equation]
( θ t )
[Equation]
22 =
[Equation]
H
[Equation]
φ )
[Equation]
θ
[Equation]
− φ )
[Equation] [Equation]
O
[Equation]
r 3 η 3 ) .
Then we have
[Equation]
So θ t +1 can be estimated by
[Equation]
Then θ t +1 − φ can be estimated by
[Equation]
where the last equality uses the second-order approximation in Lemma J.1.
Lemma J.5.
Under Assumption 4.5, consider two steps of projectioned gradient descent on S D −with effective learning rates ˜
η t and ˜ η t +1 :
[Equation]
Let φ ∈ Z (cid:15)
∩ Γ be a local minimizer.
If for some parameters α = η − o (1) and r = O ( α ) , (cid:107) θ t − φ (cid:107) 2 = O ( rη ) , ˜ η t = 2 λ H1 ( φ ) + O ( αη ) , and ˜ η t +1 = 2 λ H1 ( φ ) + O ( αη ) , then after two steps,
[Equation]
where
[Equation]
Proof.
Let ˆ x τ := θ τ − φ for all τ ∈ { t, t + 1 , t + 2 } .
Let (cid:98) ψ φ ( x ) be the following function:
[Equation]
By Lemma J.
[Equation]
we have ˆ x t +
= O ( rη ) and the following ﬁrst-order and second-order approximations for ˆ x t +1 :
[Equation]
Note that ˜ η t +1 =
˜ η t + O ((1+ r 2 ) η 2 )
= 2 λ H1 ( φ ) + O ( αη ) .
By Lemma J.4 again, we have the following second-order approximation for ˆ x t +2 :
[Equation]
Now we combine the two steps together.
[Equation]
where the second equality uses I − ˜ η t +1 H ( φ )
=
U φ + O ( αη ) .
Finally, we note that
[Equation]
where the second equality uses the fact that U φ
= ( I −
[Equation]
(Lemma D.1)
[Equation]
Lemma J.6.
Under Assumptions 4.5 and 4.6, for φ ∈ Γ ,
[Equation]
where ψ φ is deﬁned as in Lemma J
[Equation] [Equation]
Moreover,
[Equation]
Proof.
Let V φ :=
I − 2 λ H1 ( φ ) H ( φ ) .
Using a similar argument as in Lemma J.3,
[Equation]
Also notice that
[Equation]
Combining these together proves ψ φ ( v H1 ( φ ))
=
V φ ∇ log λ H1 ( φ ) + 4 φ .
To obtain the last two equations (26), (27), we ﬁrst note that φ is a 0 -eigenvector of H ( φ ) since
H ( φ ) φ = −∇L ( φ ) = 0 by Lemma D.1.
So we have (cid:104) φ , v H1 ( φ ) (cid:105) = 0 and P H0 ( φ ) φ = φ .
To prove (26), we note that V φ v H1 ( φ )
= 0 .
Then
[Equation]
To prove (27), ﬁrst we note that P H0 ( φ ) V φ = 2 I and P H0 ( φ ) φ = φ , which implies
[Equation]
Then by Lemma I.4, we can decompose P H0 ( φ ) ∇ log λ H1 ( φ ) as a component parallel to φ and a component perpendicular to φ :
[Equation]
For the second term, we note that λ H1 ( c φ ) =
c − 2 λ H1 ( φ ) by scale-invariance, and thus we have ∂∂c log λ H1 ( c φ ) = − 2 /c and ∂∂c log λ H1 ( c φ
[Equation]
c =1 = − 2 .
Combining all these together gives
[Equation]
which proves (27).
K Lemmas for Quasi-RMSprop Schedulers
Lemma K.1.
Given gradients { g t } t ≥ 0 , let ˜ η 0 , ˜ η 1 , ˜ η 2 , . .
. be the effective learning rates produced by a quasi-RMSprop scheduler with base learning rate η and decay rate β , and let { ˜ v t } t ≥ 0 be the corresponding moment estimating sequence.
Consider the case of η = o (1) , β =
[Equation]
O
[Equation]
η
[Equation]
) for some C b = Θ(1) .
For µ = Θ(1) and some t ≥ 0 , deﬁne ˆ u τ := 1 η ( µ 2 ˜ v τ − 1) for τ ∈ { t, t + 1 } .
If (cid:107) g t (cid:107) 2 ≤ αη and | µ 2 ˜ v t − 1 | ≤ αη for some
α = ( 1 η ) o
(1) at step t , then the following holds
[Equation] [Equation]
Proof.
By deﬁnition of RMSprop-like learning rate scheduler,
[Equation]
By deﬁnition of ˆ u t , we can express ˜ v t as ˜ v t = 1 µ 2 (1 + η ˆ u t ) .
Combining this with (30) proves (28):
[Equation]
By substituting 1 µ 2 (1 + η ˆ u t ) for ˜ v t in (31) we have
[Equation]
Multiplying µ 2 and subtracting 1 on both sides gives
[Equation]
Now we divide η on both sides.
Then we have
[Equation]
which proves (29).
Lemma K.2.
In the setting of Lemma K.1, for µ = Θ(1) and some t ≥ 0 , deﬁne ˆ u τ := 1 η ( µ 2 ˜ v τ − 1) for τ ∈ { t, t + 1 , t + 2 } .
If (cid:107) g t (cid:107) 2 ≤ αη , (cid:107) g t
+1 (cid:107) 2 ≤ αη , and | µ 2 ˜ v t − 1 | ≤ αη for some
α = ( 1 η ) o
(1) at step t , then the following holds
[Equation]
Proof.
Note that (29) implies ˆ u t
+
= O ( α ) .
Applying the inequalities (28) and (29) to step t +proves (32) and (33).
L Reduction to RMS-drift Process: The Case of Full Space Optimization
In this section, we let { ( θ t , v t ) } t ≥ 0 be a trajectory of gradient descent with quasi-RMSprop scheduler, and let η, β be the base learning rate and decay rate.
We follow the notations and terminologies in
Appendix H.
L.1 Good Initialization
Proof for Lemma H.8.
Let ˆ x 0 :
= θ 0 − ζ 0 and r = (cid:107) ˆ x 0 (cid:107) 2 for short.
By Gaussian concentration and anti-concentration, with probability− δ , the following holds:
[Equation]
By Taylor expansion of Φ and Lemma I.1,
[Equation]
Thus we can approximate x 0 by
[Equation]
Now we give a lower bound for | h 0 | η :
[Equation]
Since r = O ( σ 0 (cid:112) log(1 /δ )) , we can choose δ :=
C 0 α 0 η (cid:112) log(1 /η ) with a large enough C 0 such that the above inequality gives | h 0 | η ≥ Ω( δr ) .
Now we verify the conditions claimed in the lemma statement.
First, we can see from the following that the initial state is O ( α 0 (cid:112) log(1 /δ ))
-bounded:
[Equation]
It is also O ( α 0 + (cid:112) log(1 /δ ))
-deviated since
[Equation]
Next, we verify that the initial state is O (1 /δ )
-misaligned:
[Equation]
Finally, (cid:107) φ 0 − ζ 0 (cid:107) 2 ≤ O ( r ) ≤ O ( α 0 η (cid:112) log(1 /δ )) .
L.2 Alignment Phase
We deﬁne r t = (cid:107) x t (cid:107) 2 /η and the following notations for this subsection.
[Equation]
Lemma L.1.
If θ t ∈ Z (cid:15) 0 at step t , then (cid:107) P H NZT ( φ t )
( I − µ t H t ) (cid:107) 2 ≤ 1 − 2 γ min .
Proof.
If λ H1 ( φ t ) , . . . , λ H D ( φ t ) are the eigenvalues of H ( φ t ) , then { 1 − 2 λ H i ( φ t ) / λ H1 ( φ t ) }
Di =1 are the eigenvalues of I − µ t H t .
By deﬁnition of γ min , we have 1 − 2 λ
H
i ( φ t ) / λ H1 ( φ t ) ∈ [ − 1+2 γ min , 1 − 2 γ min ] as long as λ H i ( φ t ) (cid:54) = 0 , λ H1 ( φ t ) .
Therefore we have (cid:107) P H NZT ( φ t )
( I − µ t H t ) (cid:107) 2 ≤ 1 − 2 γ min .
Lemma L.2.
In the setting of Lemma H.9,
[Equation]
Proof for Lemma L.2.
By Lemma J.1, we have the following ﬁrst-order approximation for ˆ x t +1 :
[Equation]
which proves (34)
Then we can prove (35) as follows:
[Equation]
To prove (36) from (34), we note that if η is sufﬁciently small, then ˜ η t is sufﬁciently close to µ t .
In this case, by Lemma L.1 we have (cid:107) ( I − ˜ η t H t ) x t (cid:107) 2 ≤ (1 − 1 .
9 γ min ) (cid:107) P H NZT ( φ t ) x t (cid:107) 2 .
Combining this with (34) proves (36).
Finally, we prove (37) By Taylor expansion of L around φ t , (cid:107)∇L ( θ t ) (cid:107)
=
O ( (cid:107) x t (cid:107) 2 )
=
O ( r t η ) .
By Lemma K.1, we can approximate ˆ u t +1 by
[Equation]
which proves (37).
Proof for Lemma H.9.
Combining (34) with a Taylor expansion of Φ gives
[Equation]
By Lemma I.1 and Lemma I.3, ∂ Φ φ t [( I − µ t H t ) x t ] = ∂
Φ φ t [ x t ] =
O ( r 2 t η 2 ) .
So we have φ t +1 =
φ t + O ( r 2 t η 2 ) , which proves (17).
Then x t +1 =
ˆ x t +1 + φ t − φ t +1 =
ˆ x t +
+ O ( r 2 t η 2 ) , which proves (18).
Finally, by C 1 -smoothness of µ ( φ ) , v H1 ( φ ) , P H NZT ( φ ) on Γ , (35), (36), (37) imply (19), (20), (21) respectively.
Proof for Theorem H.10.
Let δ :=
C 0 α 0 η (cid:112) log(1 /η ) , where C
=
O (1) is a sufﬁciently large constant so that our proof can work.
Let α max := α 0 (cid:112) log(1 /δ ) .
Note that the initial state is O (1 /δ )
-misaligned.
So (cid:107) x 0 (cid:107)
=
O
( h 0 η/δ ) .
By Lemma H.9, for all t ≤ T 0 :
= (cid:108)
[Equation]
log
[Equation] [Equation]
(cid:109) , we can prove by induction that the states are O ( α max ) -bounded, and
[Equation]
At t = T 0 , (cid:107) x t (cid:107) 2 =
O ( h 0 η ) , φ t = φ 0 + O ( (cid:107) x 0 (cid:107) 22 ) , | h t − ( − 1) t h 0 | ≤ O ( α max h 0 η/δ ) , (cid:107) P H NZT ( φ t ) x t (cid:107)
= O ( h 0 η ) .
When C 0 is sufﬁciently large, it also holds that | h t − ( − 1) t
h 0 | ≤ | h 0 | / 4 .
After that, for all T 0 ≤ t ≤ T 1 := (cid:108)
[Equation]
= O ( α max ) , we can prove by induction that the states are O ( α max ) -bounded and
[Equation]
At t =
T 1 , (cid:107) P H NZT ( φ t ) x t (cid:107)
=
O ( h 20 η 2 ) , (cid:107) x t (cid:107)
=
O
( h 0 η ) , φ t = φ
+ O ( α 3max η 2 ) =
φ 0 +
O ( α max η ) , | h t − ( − 1) t − T 1 | ≤ O ( α 2max h 0 η ) .
When η is sufﬁciently small, it also holds that | h t − ( − 1) t − T 1 h T 1 | ≤ | h 0 | / 4 .
So | h T 1 − ( − 1) t h 0 | ≤ | h 0 | / 2 .
Putting all these together proves the theorem.
L.3 Drifting Phase
We deﬁne ψ φ ( x ) as in Lemma J.2.
We abuse the notation to write ψ t ( x ) = ψ φ t ( x ) , that is,
[Equation]
We also deﬁne the following notations for this subsection.
[Equation]
Lemma L.3.
If the state ( θ t , ˜ v t ) at step t is η − o (1) -bounded
and O ( | h t | η ) -misaligned, then x t = h t η v H1 ( φ t )
+ O ( h 2 t η 2 ) .
Proof.
Note that we have the decomposition x t = h t
η v H1 ( φ t ) + P H0 ( φ t ) x t + P H NZT ( φ t ) x t .
By deﬁnition of O ( | h t | η ) -misaligned state, P H NZT ( φ t ) x t =
O ( h 2 t η 2 ) .
By Lemma H.6, P H0 ( φ t ) x t =
O ( (cid:107) x t (cid:107) 22 ) .
[Equation]
Solving this equation gives (cid:107) x t (cid:107)
=
O ( | h t | η ) .
Then P H0 ( φ t ) x t =
O ( h 2 t η 2 ) , and therefore we have x t = h t η v H1 ( φ t )
+ O ( h 2 t η 2 ) .
Lemma L.4.
In the setting of Lemma H.11,
[Equation]
Proof.
Deﬁne ¯ g t =
(cid:107)∇L ( θ t ) /η
(cid:107) 2 and ¯ g t +1 =
(cid:107)∇L ( θ t +1 ) /η
(cid:107) 2 .
Since the state at step t is O ( | h t | η ) - misaligned, (cid:107) x t (cid:107) 2 ≤ O ( | h t | η ) .
By Taylor expansion of ∇L around φ t ,
[Equation]
So ¯ g t = (cid:107) H t x t /η (cid:107) 2 + O ( h 2 t η ) =
O ( h t ) .
Then Lemma K.1 implies (39) and the following approximation for ˆ u t +1 :
[Equation]
As (39) veriﬁes ˜ η t = µ t + O ( αη ) , we can use Lemma J.1 to derive the zeroth-order and ﬁrst-order approximations for θ t +1 : θ t +1 − φ t =
O ( | h t | η ) and θ t +1 − φ t =
U t x t + O
( α | h t | η 2 ) .
Then by Taylor expansion of ∇L around φ t again,
[Equation]
So ¯ g
[Equation]
+1 =
[Equation]
U
[Equation]
x
[Equation]
/η
[Equation]
2 + O ( α | h t | η ) =
O ( | h t | ) .
Then Lemma K.2 implies (40).
We can further apply Lemma J.2 to obtain the following:
[Equation]
Note that Lemma L.3 implies that x t = h t
η v H1 ( φ t )
+ O ( h 2 t η 2 ) .
Then by Lemma J.
[Equation]
Combining this with (42) gives (38).
Finally, we derive the approximation for ˆ u t +2 .
By Lemma K.2,
[Equation]
Since x t = h t η v H1 ( φ t ) + O ( h 2 t η 2 ) , for ¯ g t we have
[Equation]
Similarly, for ¯ g t +1 we have
[Equation]
So both ¯ g 2 t and ¯ g 2 t +1 can be approximated by λ H1 ( φ t ) 2 h 2 t + O ( | h t | 3 η ) .
Combining this with (43) gives
[Equation]
which implies (41)
[Equation]
Lemma L.5.
In the setting of Lemma H.11,
[Equation]
Proof.
In the following
[Equation]
we derive the approximations from (38),
[Equation]
Approximation for ˆ h t +2 .
Note that x t = ηh t v H1 ( φ t )
+ O ( h 2 t η 2 ) since the state at step t is
O ( | h t | η ) -misaligned.
For ˆ h t +2 , we have
[Equation]
where we use the fact that (cid:104) ψ t ( v H1 ( φ t )) , v H1 ( φ t ) (cid:105)
=
by Lemma J.3.
By (39), (40), ˜ η τ = µ t · (1 − 12 ηu t ) + O ( α 2 η 2 ) for τ ∈ { t, t + 1 } .
Note that µ t · λ H1 ( φ t ) = 2 .
Then for τ ∈ { t, t + 1 } ,
[Equation]
Then (1 − ˜ η t +1 λ H1 ( φ t ))(1 − ˜ η t λ H1 ( φ t )) can be approximated by
[Equation]
˜ η t +1 λ
[Equation]
( φ t ))(1 − ˜ η t λ H1 ( φ t ))
=
( −
[Equation]
ηu
[Equation]
O ( α 2 η 2 ))
[Equation]
=
− 2 ηu t
+ O ( α 2 η 2 ) .
Therefore, we have ˆ h t
+
=
(cid:0) 1 − 2 ηu t
+ O ( α 2 η 2 ) (cid:1) h t
+ O ( αh 2 t η 2 ) , which implies (44).
Approximation for P H0 ( φ t )ˆ x t +2 .
For P H0 ( φ t )ˆ x t +2 , we have
[Equation]
where we use the fact that P H0 ( φ t ) ψ t ( v H1 ( φ t ))
= 2 ∇ Γ log λ H1 ( φ t ) by Lemma J.3.
To obtain (45), we only need to note that P H0 ( φ t )( I −
˜
η t +1 H t )
( I − ˜ η t H t ) x t =
P H0 ( φ t ) x t .
Approximation for (cid:107) P H NZT ( φ t )ˆ x t +2 (cid:107) 2 .
To approximate (cid:107) P H NZT ( φ t )
ˆ x t +2 (cid:107) 2 , we note that if η is sufﬁciently small, then ˜ η t , ˜ η t +1 are sufﬁciently close to µ t .
In this case, by Lemma L.1 we have (cid:107) P
[Equation]
t )
[Equation]
I
[Equation]
˜
[Equation]
+
[Equation] [Equation]
I
[Equation] [Equation]
[Equation]
P
[Equation]
) x t
[Equation]
Combining this with (38) proves (46).
Lemma L.6.
In the setting of Lemma H.11, the approximation (25) holds for φ t +2 and
[Equation]
Proof.
By Taylor expansion,
[Equation]
By deﬁnition, Φ( φ t ) =
φ t .
By (45) and Lemma I.1,
[Equation]
Note that x t = ηh t v H1 ( φ t )
+ O ( h 2 t η 2 ) since the state at step t is O ( | h t | η )
-misaligned.
Also note that the identity ∂ 2 Φ φ t [ v H1 ( φ t ) , v H1 ( φ t )]
= 0 holds by Lemma I.2.
Then
[Equation]
Similarly, we have ∂ 2 Φ φ t [ˆ x t +2 , ˆ x t +2 ]
=
O ( | h t | 3 η 3 ) since (45) implies (cid:107) P H0 ( φ t )
ˆ x t +2 (cid:107) 2 =
O ( h 2 t η 2 ) and (46) implies (cid:107) P H NZT ( φ t )
ˆ x t +2 (cid:107)
= O ( h 2 t η 2 ) .
Now we can prove (25) by subtracting (49) with (48): φ t
[Equation]
Finally
[Equation]
for log λ H1 ( φ t +2 )
we have
[Equation]
which proves (47)
[Equation]
Proof for Lemma H.11.
We have veriﬁed (25) in Lemma L.6.
By (25) and deﬁnitions of ˆ x t +2 and x t +2 , we have
[Equation]
And we can write x t +2 − ˆ x t +2 =
φ t − φ t +
=
O ( h 2 t η 2 ) as a loose approximation.
Approximation for h t +2 .
For h t +2 , we have
[Equation]
where the fourth equality is due to (cid:104)∇ Γ log λ H1 ( φ t ) , v H1 ( φ t ) (cid:105) =
and (cid:107) φ t +2 − φ t
[Equation]
2 =
O ( h 2 t η 2 ) .
Combining this with (44) proves the claimed approximation (22).
Approximation for (cid:107) P H NZT ( φ t +2 ) x
t +2 (cid:107) 2 .
For P H NZT ( φ t +2 ) x t +2 , we have
[Equation]
where the third equality is due to x t +2 − ˆ x t +
=
φ t − φ t +
=
O ( h 2 t η 2 ) .
Combining this with (46) proves the claimed approximation (23).
Approximation for u t +2 .
Now we prove the formula for u t +2 .
Note that φ t +2 − φ t
= O ( h 2 t η 2 ) implies that µ t +2 − µ t =
O ( h 2 t η 2 ) .
Then we have
[Equation]
Note that log µ t +2 − log µ t = log (cid:16)
[Equation]
+
[Equation]
= 1 µ
[Equation]
µ
[Equation]
2 − µ
[Equation] [Equation]
O
[Equation]
( h 2 t η 2 ) 2 ) .
By (47),
[Equation]
Combining these together gives the following approximation for 1 µ t ( µ t +2 − µ t ) :
[Equation]
So η ( u t +2 − ˆ u t +2 )
= 4 η 2 h 2 t (cid:107)∇ Γ log λ H1 ( φ t ) (cid:107)
+ O ( αh 2 t η 3 ) .
Then by (41), we have
[Equation]
which proves the claimed approximation (24).
M Reduction to RMS-drift Process: The Case of Spherical Optimization
In this section, we let { ( θ t , v t ) } t ≥ 0 be a trajectory of projected gradient descent with quasi-RMSprop scheduler, and let η, β be the base learning rate and decay rate.
We follow the notations and terminologies in Appendix H.
As the analysis in the spherical case is nearly the same as the full space case, we only discuss the difference here.
M.1 Good Initialization
Proof for Lemma H.17.
Using a similar argument as in Lemma H.8, we know that the lemma holds if there is no projection in the random initialization.
But the projection only leads to an error of order
O ( r 2 ) , so the lemma holds.
M.2 Alignment Phase
Proofs for Lemma H.18 and Theorem H.19.
The proof is essentially the same as Lemma H. and Theorem H.10.
To see this, we only need to note that we only have used a linear approximation of the update rule with error O
( r 2 t ) , and the linear approximation remains unchanged if we add a projection operator (Lemma J.4).
M.3 Drifting Phase
We deﬁne ψ φ ( x ) as in Lemma J.5.
We abuse the notation to write ψ t ( x ) = ψ φ t ( x ) , that is,
[Equation]
Lemma M.1.
In the setting of Lemma H.20, the same statement as Lemmas L.4 and L.5 holds, where ψ t is interpreted as (50) .
Proof.
We can follow the argument in the proof for Lemmas L.4 and L.5, but now we are using
Lemmas J.5 and J.6 to establish the proof.
Lemma M.2.
In the setting of Lemma H.20, the same statement as Lemma L.6 holds.
Proof.
The argument is the same as Lemma L.6, but we apply Lemmas I.4 and I.5 in doing Taylor expansion for Φ( θ ) .
Proof for Lemma H.20.
Same as Lemma H.
but we invoke the spherical version of lemmas.
N Analysis of RMS-drift Process
In this section, we provide proofs for theorems in Appendix H.5.
For convenience, we deﬁne
R ( θ ) = ∇ log λ H1 ( θ ) , K t :=
(cid:112) 2 C b
+ (cid:107)∇ Γ R ( φ t ) (cid:107) 22 .
Then a C 0 -RMS-drift transition S t → S t +2 can be written as
[Equation]
N.1 Conservation of Energy
To establish the conservation of energy, we ﬁrst compute the change in energy after one transition.
Lemma N.1.
Given two drift states S
=
( h 0 , u 0 , φ 0 ) and S
= ( h 2 , u 2 , φ 2 ) in the working zone, for learning rate η and hyperparameter C b > 0 , if S 0 is α -bounded for some 1 ≤ α ≤ η − o (1) , and
S
→ S 2 is a C 0 -RMSdrift transition, then
[Equation]
Proof for Lemma N.1.
h
and h 2 have the same sign when η is small enough.
We can decompose
E ( S 2 ) − E ( S 0 ) as follows:
[Equation]
Now we bound each error term.
For δ 1 and δ 2 , we use the formula a 2 − b 2 = 2 b ( a − b )
+ ( a − b ) 2 :
[Equation]
For δ 3 , we use the Lipschitzness of ∇ Γ R ( θ ) :
[Equation]
For δ 4 , note that log(1 + z ) ≤ z + O ( z 2 )
when z = o (1) .
Then
[Equation]
Adding δ 1 , δ 2 , δ 3 , δ 4 together gives
[Equation]
So E ( S
[Equation]
− E
[Equation]
To sum up the energy change over time, we need the following lemma.
Lemma N.2.
For any M = o (( αη ) − 2 ) , if S 0 , . .
, S 2 M is an O (1) -RMS-drift process in the working zone, and S t is α -bounded for all even numbers t ≤ 2 M , then
[Equation]
Proof for Lemma N.2.
By the update rule of u t , we have
[Equation]
Since S 2 m is α -bounded for all m ≤ M , φ 2 m − φ
= O ( α 2 η 2 M ) .
By smoothness of R , (cid:107)∇ Γ R
[Equation] [Equation]
22 =
(cid:107)∇ Γ R ( φ 0 ) (cid:107) 2
+ O ( α 2 η 2 M ) , then K 2 m =
K 0 + O ( α 2 η 2 M ) .
Since S 0 and
S 2 M are α -bounded, u T − u
= O ( α ) .
Combining all these together,
[Equation]
Let Q := (cid:80) M − 1 m =
h 22 m .
Then we have
[Equation]
Rearranging the terms while noting that α 3 η 2 M = o ( α ) , we have
[Equation]
So we can estimate Q by
[Equation]
which completes the proof
[Equation]
Lemma N.3.
For an O (1) -RMSdrift process S 0 , . . .
, S 2 M in the working zone, if E ( S 0 ) ≤ α 2 for some parameter 1 ≤ α ≤ η − o (1) and M =
O (1 /η 1 . 5 ) , then E ( S t ) =
E ( S 0 ) + O ( α 2 η 0 . 5 ) for all even numbers 0 ≤ t ≤ 2 M .
Proof for Lemma N.3.
We do a bootstrap.
As E ( S 0 ) ≤ O ( α 2 ) , we can leverage Lemma N.1 to prove by induction that E ( S t ) ≤ E ( S 0 )
+ O ( α 4 η 2 M ) for all 0 ≤ t ≤ 2 M .
Then we apply Lemma N. again.
For all N ≤ M ,
[Equation]
where the third line uses Lemma N.2.
Proof for Theorem H.22.
We group the M transitions into O (1 /η 0 .
5 ) segments of length O (1 /η 1 . 5 ) .
We can do an induction with Lemma N.3 applied on each segment to show that for all t in the k -th segment, E ( S t ) ≤ (1 + O ( η 0 . 5 ))
k E ( S 0 ) .
Noting that k
= O (1 /η 0 .
and (1 + O ( η 0 . 5 ))
O (1 /η 0 .
5 ) =
O (1) ﬁnishes the proof.
N.2 Flow Approximation
Lemma N.4.
In the setting of Theorem H.23 but with M = Θ(1 /η 1 . 5 ) , if S t is α -bounded for all even numbers 0 ≤ t ≤ 2 M , then
[Equation]
Proof for Lemma N.4.
When S t is α -bounded, (cid:107) φ t +2 − φ t (cid:107)
=
O
( α 2 η 2 ) , so (cid:107) φ t − φ 0 (cid:107)
= O ( α 2 η 2 M )
=
O ( α 2 η 0 . 5 ) .
For every N ≤ M we have
[Equation]
By Lemma N.2, we have
[Equation]
Note that C 2b 2 K 20 ∇ Γ R ( φ
[Equation]
= dd t ζ (0) +
O ( α 2 η 1 / 2 ) .
Then we have
[Equation]
where the third equality uses the smoothness of ζ ( t ) .
Proof for Theorem H.23.
We group the M transitions into O (1 /η 1 / 2 ) segments of length O (1 /η 1 . 5 ) .
Then we leverage Lemma N.4 and do an induction to show that (cid:107) φ t − ζ ( tη 2 ) (cid:107) 2 ≤ O ( α 2 η 1 / 2 ) for all even numbers 0 ≤ t ≤ 2 M .
O Proofs for Linear Regression with Batch Normalization
Lemma O.1.
Assume that the regression targets y i are generated by a linear model.
For linear regression with BN, the global minimizer manifold of L ( w ) := 1 n (cid:80) ni
=1 (Φ
[Equation]
i ; w , σ y , µ y ) − y i ) 2 on the unit sphere is
[Equation]
For any global minimizer w ∈ Γ , the Hessian matrix H ( w ) of the loss is given by
[Equation]
where ˜ w :
= σ y w (cid:107) w (cid:107) Σ x as deﬁned in (4) , and z := 1 n (cid:80) ni
=
y i − µ y σ y
( x i − µ x ) .
Proof.
The model output Φ( x ; w , µ y , σ y ) can be written as
[Equation]
Then it is easy to verify that the global minimizer manifold is Γ .
Now we compute the Hessian.
Let ˜ x i := x i − µ x and q i :=
y i − µ y σ y .
Because of the use of squared loss, on Γ the Hessian can be written as the sum of outer products of gradients:
[Equation]
For each gradient we have
[Equation]
Note that Σ x ˜ w = 1 n (cid:80) ni
=1 ˜ x i ˜ x (cid:62) i ˜ w = 1 n (cid:80) ni
=1 q i ˜ x i =:
z .
Then we can simplify the gradient by
[Equation]
Now we simplify the Hessian.
[Equation]
We complete the proof by noting that
[Equation]
Proof for Theorem 5.1.
By Lemma O.1, for w ∈ S D − 1 ,
[Equation]
By simple calculation, it can be veriﬁed that the only point that has ∇ Γ (cid:107) ˜
w (cid:107) 22 = 0 is the unique point w ∗ on Γ that is a linear combination of x i − µ x .
As the shperical sharpness is bounded from below and there is only one stationary point, the sharpness-reduction ﬂow must converge on Γ and the convergence point must be w ∗ .
Since w ∗ is a linear combination of x i − µ x , the associated ˜ w ∗ should be the least square solution (without bias) of this “shifted” dataset: { ( x i − µ x , y i − µ y ) } .
In other words, ˜ w ∗ is the optimal solution of the following constrained optimization problem:
[Equation]
When w is given in the above optimization problem, there is only a unique b such that w (cid:62) x
i + b =
y i .
So we can introduce a bias to this problem without changing the optimal solution:
[Equation]
One can also easily see that this b must match with ˜ b .
Therefore, we can conclude that the sharpnessreduction ﬂow (3) ﬁnds to the optimal solution of (M) at convergence.
P Experiments
In this section, we provide experiments on matrix completion and CIFAR-10 to validate the main claim in our theory: GD+WD on scale-invariant loss persistently reduces spherical sharpness in the EoS regime (the regime where 2 / ˜ η t roughly equals to the spherical sharpness).
In addition, we validate that the generalization performance continues to improve as the spherical sharpness decreases.
See Appendices P.1 and P.2.
Then in Appendix P.3, we validate a key proof insight: the magnitude of oscillation and effective LR evolve periodically.
We also provide a series of ablation studies.
In Appendices P.4 and P.5, we demonstrate that the two key components in our theoretical setup, normalization and WD, are crucial.
In Appendix P.6, we show that the initial effective LR does not affect the ﬁnal performance when the intrinsic LR is ﬁxed.
P.1 Validation of Sharpness Reduction on Matrix Completion
First, we conduct experiments on matrix completion: there is an unknown low-rank ground-truth matrix M ∈ R d × d , and only N entries of M are known.
The goal is to recover M given the known entries.
Matrix completion has been studied via deep learning techniques (see, e.g., [20] for a survey).
Matrix completion can be connected to supervised learning as follows.
Each entry of the matrix can
[Figure]
M 50×
, 400 observations
[Figure] [Figure]
where the ground-truth matrix M ∈ R 50 × 50 is of rank 2, and the number of observations vary from 1600 to 3200 .
The test loss starts to decrease signiﬁcantly as soon as the spherical sharpness starts to decrease.
be seen as a data point, where the N observed entries constitute the training set.
As in supervised learning, given a subset of data points, the goal of matrix completion is to build a model to predict the rest of the data points.
Here we empirically study solving matrix completion using an overparameterized scale-invariant model.
We can observe that spherical sharpness is indeed decreasing as soon as the process enters the
EoS regime, and the reduction of spherical sharpness encourages low-rank.
In our experiments, we generate the ground-truth matrix
M ∈ R d × d as follows.
First, we set (cid:102)
M ← U ∗ V (cid:62)∗ for two random matrices U ∗ ,
V ∗ ∈ R d × 2 , where every entry is sampled uniformly from [ − 1 , 1] .
Then we obtain M by normalizing (cid:102) M so that the second moment of the entries is 1 , i.e., M ← ( d · (cid:107) (cid:102) M (cid:107) − 1F )
(cid:102) M .
We uniformly sample N entries of M to serve as the observations, and use Ω ⊆
[ d ] × [ d ] to denote the index set of observed entries.
Matrix completion has been studied by a line of works [41, 76, 4, 100, 79, 108] as a test-bed for the implicit regularization of gradient descent.
More speciﬁcally, they parameterize the target matrix as
W = UV (cid:62) where U , V ∈ R d × d are two trainable matrices, and run GD to minimize the squared loss L ( U , V ) :
= 1 N (cid:80) ( i,j
) ∈ Ω ( W i,j − M i,j ) 2 .
Although there is no explicit constraint on rank, GD with small random initialization can still exhibit an implicit bias towards low-rank solutions.
Inspired by this line of works, we conduct matrix completion experiments to test if GD+WD exhibits the same low-rank bias in training overparameterized models with BN.
More speciﬁcally, we parameterize the target matrix as W = BN( UV (cid:62) ) , where U ,
V ∈ R d
× d are two trainable matrices.
Given the observed positions Ω , the output of the model for a single position ( i,
j ) ∈ Ω is γσ [ UV (cid:62) ] i,j , where σ 2 :
= 1 N (cid:80) ( i,
j ) ∈ Ω [ UV (cid:62) ]
i,j is the second moment of [ UV (cid:62) ]
i,j over all observed positions, and γ is a rescaling factor.
Multiplying the factor γσ can be seen as doing BN over observed entries because it rescales the output in the same manner as BN.
But the difference is that, for the sake of simplicity, we do not subtract the mean.
To ensure the loss to be scale-invariant, we also ﬁx γ to match the second moment of observed entries, i.e., γ 2 := 1 N (cid:80) ( i,j ) ∈
Ω M 2 i,j .
To train this model, we use the standard squared loss, and we run gradient descent with LR ˆ η = 0 .and
WD ˆ λ = 0 .
01 to optimize the loss.
It is obvious that this loss is scale-invariant due to BN.
[Equation]
For an unobserved entry ( i, j ) , the model uses the same batch statistics as for observed entries, and predicts γσ [ UV (cid:62) ] i,j , where γ, σ are the same as above.
So we measure the test loss as
[Equation]
We note that the loss function has no explicit constraint on rank.
But surprisingly, in our experiments,
GD+WD tends to prefer low-rank solution as soon as the sharpness-reduction bias starts to occur.
See Figures 1 and 7 for experiments on reconstructing a rank-2 matrix of size 50 × 50 , and Figurefor experiments on reconstructing a rank-2 matrix of size 100 × 100 .
In all these experiments, we can observe that the train loss ﬁrst decreases to near zero in a short time, but the test loss remains at a high level.
Then after effective LR increases for some more steps, the dynamics enters the EoS regime, and as predicted by our theory, the spherical sharpness starts to decrease.
Meanwhile, the test loss also starts to decrease signiﬁcantly.
We notice that the gap between the second and the third largest singular values of the recovered matrix is enlarged at the same time, suggesting that this reduction of spherical sharpness encourages the recovered matrix to be low-rank.
P.2 Validation of Sharpness Reduction on CIFAR-
Now we present experiments on CIFAR-10 with crossentropy loss to validate the sharpness-reduction bias in a more realistic setting.
We run full-batch GD with accompanying WD on three different architectures: a scale-invariant variant of VGG-11, a scale-invariant variant of pre-activation ResNet20, and the standard pre-activation ResNet-20.
We ﬁx LR ˆ η = 0 .
1 and WD ˆ λ = 5 ×
10 − 4 .
See
Appendix Q.1 for more details on the architectures and training procedures.
The experiment for the scale-invariant variants of VGG-11 and ResNet-20 are presented in Figuresand 9.
Here our ResNet-20 is non-smooth due to the use of ReLU, but our VGG-11 is smooth because we choose to use Swish activation and mean pooling in this network.
For both the smooth VGGand non-smooth ResNet-20, it can be seen from the plots that in both cases the spherical sharpness has an overall tendency to decrease over time, and the test accuracy is increasing accordingly.
We note that this sharpness-reduction bias actually goes beyond the setting that our theory can directly apply: the dynamic enters the EoS regime in the very beginning of training, but our theory only analyzes the dynamic near a local minimizer manifold.
This is because that the LR here is not small enough so that θ t can approach to a local minimizer before the dynamic enters the EoS regime.
In fact, Figure 11 shows that in the scale-invariant VGG-11 experiment, the initial spherical sharpness is much larger than 2 / ˜ η 0 , then after a few steps, the spherical sharpness decreases to a level that is close to 2 /
˜ η t and the dynamic enters the EoS regime.
Besides the scale-invariant models, we also validate the sharpness-reduction bias on the standard pre-activation ResNet-20, which is not (fully) scale-invariant but only scale-invariant to a part of its parameters.
For evaluating spherical sharpness, we compute the partial Hessian only with respect to that part of parameters.
See Figure 10 for the plot and Appendix Q.1 for more experimental details.
Although our theory can only cover scale-invariant models, we can still observe the sharpnessreduction bias in experiments, and the test accuracy increases as spherical sharpness decreases.
[Figure]
[Figure] [Figure]
P.3 Periodic Behaviors
A key insight in our proof is that the magnitude of oscillation and effective LR evolve periodically in the EoS regime.
Following Section 4.2.3, we can divide the EoS regime into two sub-regimes that occur alternatively in training: (1) the sub-regime where 2 / ˜ η t is smaller than λ H1 ( φ t ) and the magnitude of oscillation | h t | keeps increasing; and (2) the sub-regime where 2 / ˜ η t is bigger than λ H1 ( φ t ) and the magnitude of oscillation | h t | keeps decreasing.
Here φ t is a projection of θ t onto the local minimizer manifold, λ H1 ( φ t ) is the spherical sharpness computed at φ t , and h t is the inner product between θ t − φ t and the top eigenvector of H ( φ t ) .
Figure 5 provides a nice validation of this periodic behavior in linear regression with BN.
We further validate the periodic behavior in matrix completion by visualizing ˜ η t in Figure 12 around the moment that the dynamic enters the EoS. We do not plot | h t | because computing φ t is inefﬁcient, but we can observe that the training loss evolves periodically, which signals that | h t | evolves periodically as well.
In CIFAR-10 experiments with realistic LR and WD, the periodic behavior still exists, but it may deviate from the regime that our theory can capture.
In most of our CIFAR-10 experiments, only the sub-regime where 2 / ˜ η t < λ H1 ( φ t ) can be observed, and the phenomenon that 2 / ˜ η t switches back and forth between being smaller and larger than the spherical sharpness does not occur anymore.
However, the change of ˜ η t and spherical sharpness still cause to the loss and gradient norm to ﬂuctuate periodically.
See Figure 13 for details.
This difference is because that the loss function of a neural net on CIFAR-10 is much less smooth than that of linear regression and matrix completion.
To make the dynamic happen in the regime that our theory describes, the LR and WD need to be very small so that θ t is sufﬁciently close to the minimizer manifold Γ , but they cannot be that small in practice due to computational inefﬁciency.
The periodic behavior can also happen in a macroscopic scale on CIFAR-10.
E.g., in Figure 2, the training loss spikes to a high value around step 50,000, but then recovers to a near-zero value after a few hundred steps.
Such a macroscopic periodic behavior may not always lead to sharpness reduction, in contrast to the microscopic periodic behavior of our interest.
But as noted in the extensive empirical study by Lobacheva et al.
[84], this macroscopic periodic behavior sometimes leads to better generalization.
They also provide a theoretical analysis for the cause assuming the
[Figure] [Figure]
gradient norm is both lower and upper bounded.
However, their explanation is not completely satisfactory because the gradient norm does not admit a lower bound when the parameter is close to the minimizer manifold Γ .
We leave it a future work to look further into the cause of macroscopic periodic behavior and how it helps generalization.
P.4 Ablation Study: Normalization
Figure 2 and many other experiments in Appendices P.1 and P.2 have shown that GD+WD on normalized nets can continue to improve test accuracy even after reaching 100% test accuracy.
In our theoretical analysis, we connect this phenomenon with the reduction of spherical sharpness during training.
Now we empirically validate that this phenomenon is indeed linked the presence of normalization layers in neural nets.
We train a VGG-11 in the same setting as Figure 2, but now we remove all the normalization layers.
See Figure 14.
The test accuracy does not increase anymore after the training accuracy reaches 100% .
Note that the spherical sharpness is meaningful only for normalized nets.
For unnormalized nets, even if GD is implicitly reducing a similar sharpness measure, such a measure cannot be strongly related to generalization because the test accuracy is not increasing accordingly.
P.5 Ablation Study: Weight Decay
Now we conduct experiments to study the effects of weight decay (WD) on sharpness reduction.
It is crucial to have non-zero WD in our theoretical analysis.
Otherwise, it is implied by
Lemma E.1 that the effective LR is monotone decreasing, but our theory only applies to the effective LRs that can
[Figure] [Figure]
[Figure] [Figure]
be viewed as quasi-RMSprop (see Deﬁnition B.5 and Theorem B.9), in which the effective LR can either increase or decrease as gradient norm changes.
A previous work by Arora et al.
[7] provides a detailed theoretical analysis in this case showing that the dynamic always stays in the stable regime after a few warm-up steps, and the parameter eventually converges to a stationary point on the unit sphere under standard assumptions in optimization.
We can verify through experiments that the requirement of WD to be non-zero is not a technical artifact, but indeed a necessity in practice to exhibit the sharpness-reduction bias.
We train a scaleinvariant VGG-11 following exactly the same hyperparameters and initialization as Figure 2, except that we set WD ˆ λ to zero.
The result is presented in Figure 15.
The spherical sharpness is no longer decreasing with time when WD is zero, and the ﬁnal test accuracy is much lower than the VGGwith WD 5 × 10 − 4 .
We further conduct experiments with smaller (but non-zero) WD than our default value 5 × 10 − 4 .
For the sake of computational efﬁciency, our experiments are conducted on a subset of CIFAR-10 images consisting of 2K images, which we call CIFAR-10-2k (see Appendix Q.1).
The result is presented in
Figure 16, from which we can observe that GD exhibits a sharpness-reduction bias as long as WD is non-zero.
But note that smaller WD leads to a slower speed of sharpness reduction.
This is an expected phenomenon because the speed of tracking the sharpness-reduction ﬂow (3) is controlled by the intrinsic LR η in :
=
ˆ η ˆ λ .
The smaller the WD, the smaller the intrinsic LR (when LR ˆ η is ﬁxed).
P.6 Ablation Study: Initial Effective Learning Rate
As we have seen in Figure 11, in training the scale-invariant VGG-11 on CIFAR-10, if ˆ η = 0 .
1 and ˆ λ = 5 × 10 − 4 , then initially the spherical sharpness is much larger than 2 / ˜ η 0 , and after a few steps, it decreases to the same level as 2 / ˜ η t , thus bringing the dynamic into the EoS regime.
Now we conduct an ablation study to see how the dynamics change with different initial effective LR.
For computational efﬁciency, we run experiments on CIFAR-10-2k (see Appendix Q.1) to train our scale-invariant variant of VGG-11.
We ﬁrst set LR ˆ η = 0 .
1 and WD ˆ λ = 5 × 10 − 4 , then we change the LR ˆ η to 10 , 10 − 3 , 10 − 5 while rescaling the WD ˆ λ inverse proportionally.
In this way, the intrinsic
[Figure]
LR remains unchanged.
See Figure 17 for the results.
For large initial effective LR, the dynamic quickly enters the EoS regime; but when the initial effective LR is as small as 10 − 5 , the dynamic is in the stable regime initially, and 2 / ˜ η t falls so slowly that the dynamic enters the EoS regime only after reaching 100% training accuracy.
In all experiments, the spherical sharpness persistently decreases in the EoS regime, and the test accuracy increases accordingly.
In the end, all the experiment runs reach nearly the same spherical sharpness and test accuracy, regardless the initial effective LR.
Q Experiment Details
All our experiments were conducted on NVIDIA RTX A5000 GPUs.
The longest experiments are the full-batch training experiments on CIFAR-10, each of which took 7 days to run.
Q.
Additional Details of
CIFAR-10 Experiments
Our implementation of full-batch GD is based on the code of Cohen et al.
Our implementations of VGGNets [105] and ResNets [46, 47] on CIFAR-10 are based on a high-starred GitHub repository of Wei Yang (bearpaw) 6 .
We do not use any data augmentation.
As it is not feasible to do BN over the full dataset under our GPU memory constraints, we use ghost batch normalization [51] instead, where we split the dataset into 50 ghost batches of size 1000 .
Some of our experiments are conducted on only a subset of CIFAR-10 consisting of 2000 images, which we call CIFAR-10-2k.
We construct the dataset by scanning the full CIFAR-10 dataset and taking the ﬁrst 200 images for each of the 10 classes.
We use the standard batch normalization with full batch for all CIFAR-10-2k experiments.
Three different neural network architectures are tested in our experiments, and we refer to them as scale-invariant VGG-11, standard ResNet-20, and scale-invariant ResNet-20 respectively.
Scale-Invariant VGG-11.
Our scale-invariant VGG-11 architecture is similar to the conﬁguration
A of the original VGGNet [105], but we make the following changes.
We add a BN layer (without afﬁne parameters) between every convolution and activation to introduce scale-invariance, We use mean pooling instead of max pooling, Swish instead of ReLU to make the training loss smooth.
We replace the ﬁnal 3 fully-connected layers with only one fully-connected layer as in Yang’s implementation.
All convolutional and fully-connected layers have no bias terms.
We add a BN layer after the last fully-connect layer, in which the afﬁne parameters are ﬁxed to γ = 310 ln 91 ≈ 1 . 353 , β =
(see Appendix R for discussion).
A visualization of the full architecture is given in Figure 18, whose output function can be shown to be scale-invariant with respect to the trainable parameters.
[Figure]
Standard ResNet-20.
We also verify the sharpness-reduction bias on a more realistic setting with
ResNet-20 and ReLU.
In some of our experiments, we use Yang’s implementation for the preactivation variant of ResNet-20.
All convolutional layers have no bias, but the ﬁnal fully-connected layer does have bias terms.
We note that the output function of this architecture is not scale-invariant to all its trainable parameters.
However, it can be shown that the output function is indeed scaleinvariant to a large part of its parameters.
More speciﬁcally, the trainable parameters can be split into two parts, w and z , and the output function remains unchanged if we replace w with c w for all c > 0 .
In our experiments, we measure the spherical sharpness by ∇ 2 θ L ( θ , z ) while letting θ := w (cid:107) w (cid:107) 2 .
See
Figure 19a for a visualization of the architecture.
Scale-Invariant ResNet-20.
The standard ResNet-20 can be scale-invariant after making minor changes.
First, we remove the afﬁne parameters in all existing BNs.
Then, we also add an extra BN (with trainable afﬁne parameters) to each shortcut for blocks with downsampling, following Li and
Arora [77].
Finally, we add a BN layer after the last fully-connect layer, in which the afﬁne parameters are ﬁxed to γ = 310 ln 91 ≈ 1 . 353 , β =
(see Appendix R for discussion).
A visualization of the full architecture is given in Figure 19b, whose output function can be shown to be scale-invariant with respect to the trainable parameters.
[Figure]
of ResNet-20, following Yang’s
[Figure] [Figure]
[Figure]
Q.2 Additional Details of Linear Regression Experiments
Our experiments on linear regression (Figure 5) follows strictly as Section 5.
For generating data, we ﬁrst sample the ground-truth weight w GT ∈ S
d − 1 uniformly, where d = 40 , and we sample the ground-truth bias as b GT ∼ N (0 , 0 . 01) , Then we sample n = 20 points from Gaussian distribution
N ( 0 , Σ ) , where Σ =
diag(1 /d, 2 /d, 3 /d, . .
To generate Figure 5, we run gradient descent with LR ˆ η = 0 .
5 and WD ˆ λ = 2 × 10 − 4 .
This particular choice of hyperparameters is for better visualization of the periodic behavior described in
Section 4.2.3.
If we enlarge LR or WD, the period will be shortened accordingly.
Q.3 Computing Spherical Sharpness
To compute the spherical sharpness at a point w ∈ R D \ { 0 } (Deﬁnition 1.1), we utilize the formula
λ 1 ( ∇ 2 L ( w (cid:107) w (cid:107) 2 ))
= (cid:107) w (cid:107) 22 · λ 1 ( ∇ 2 L ( w )) , which can be proved for all scale-invariant loss by simple calculus (see also Lemma D.1).
For computing the eigenvalues of ∇ 2 L ( w ) , we invoke the Lanczos algorithm from SciPy library ( scipy.sparse.linalg.eigsh ), where a Hessian-vector product oracle x (cid:55)→ ∇ 2 L ( w ) x is implemented with PyTorch and passed to the Lanczos algorithm as a linear operator to ﬁnd its eigenvalues.
In full-batch training of the full CIFAR-10 dataset, it is time-consuming to compute even a single
Hessian-vector product.
Following Cohen et al.
[24], we compute Hessian based on only the ﬁrst 5000 training data points.
Note that in our theory we focus on the spherical sharpness of minimizers.
However, in practice, it is usually time-consuming to compute the minimizers exactly, so in matrix completion and CIFARexperiments, we compute the spherical sharpness directly at the current parameter w t .
But for linear regression experiments (Figure 5), the spherical sharpness is indeed computed at minimizers in each step because our computational power is sufﬁcient in this setting.
Speciﬁcally, we compute the minimizer by doing projected gradient descent on S D − 1 with ﬁxed learning rate 0 .
005 until the loss decreases to 10 − 8 .
R Discussion on the Afﬁne Parameters of the Final BN
In our experiments for scale-invariant models, we add a BN to the ﬁnal linear layer and ﬁx the afﬁne parameters to be constants.
Now we analyze its effects on regression tasks with squared loss and classiﬁcation tasks with crossentropy loss, and discuss how to set the afﬁne parameters reasonably.
We consider a scale-invariant neural net in general while assuming that a BN is put as the last layer.
For input x and parameter w , let ( F 1 ( x ; w ) , . .
, F C ( x ; w ))
∈ R C be the output before the ﬁnal BN.
After the ﬁnal BN, the output is given by the following function Φ k ( w ; w ) :
[Equation]
where µ k and σ 2 k are mean and variance of { F k ( x i ; w ) } ni =1 over the training set.
R.1 Squared Loss
For linear regression experiments, we use squared loss as the loss function, and we have C =because the output is a scalar.
Then as discussed in Section 5, the mean and variance of Φ 1 ( x ; w ) are always β 1 and γ 21 regardless of w .
Therefore, to ensure that Φ has sufﬁcient representation power to express the regression targets { y i } ni =1 , we should ﬁx the afﬁne parameters in a way that β 1 and γ 21 match with the mean and variance of { y i } ni =1 .
In the general case when C ≥ 1 , the k -th output unit has mean and variance β k and γ 2 k .
We can argue similarly that we should ﬁx the afﬁne parameters in a way that β k and γ 2 k match with the mean and variance of { y i,k } ni =1 , where y i,k stands for the k -th coordinate of the regression target for the i -th data point.
In other words,
[Equation]
R.2 Crossentropy Loss
In our CIFAR-10 experiments, we ﬁx all γ k to a constant value γ , and ﬁx β k to zero.
We note that ﬁxing these afﬁne parameters in this way for an overparameterized model is equivalent to adding label smoothing to the training loss in some sense.
The basic idea is to assume the model has sufﬁcient express any function F , then Φ k can be any function of mean zero and covariance γ 2 on the training set.
If the dataset has C classes and each class has equal numbers of samples, with some efforts one can show that the minimum loss is attained when the pre-softmax logit for an input is γ ( C − 1) 1 / 2 for the correct class and − γ ( C − 1) − 1 / 2 for every wrong class.
In this case, the output probability P y ( x i ) for a class y is
[Equation]
if y
[Equation]
otherwise.
Label smoothing is a regularization technique proposed by Szegedy et al.
[109] to improve the generalization of neural nets by replacing the one-hot hard labels with soft labels that assign a probability of (cid:15)/C to each wrong class.
To see the connection to label smoothing, we can let (cid:15) :=
[Equation]
) , then P y ( x i ) =
(cid:15)/C if y (cid:54) =
y i and P y ( x i )
= 1 − (1 − 1 /C )
(cid:15) if y = y i .
In our CIFAR-10 experiments, γ = 310 ln 91 ≈ 1 . 353 .
Then by simple calculation (cid:15) = 0 .
1 , and therefore the neural net trained with γ k and β k ﬁxed to γ and 0 is encouraged to produce the output probabilities to match with the soft labels.
The minimum loss is non-zero when the afﬁne parameters are ﬁxed.
We can see from (51) that the loss attained by any minimizer is
[Equation]
In making plots for our experiments, we always subtract the original loss with its theoretical minimum value L min .
When γ = 310 ln 91 , this minimum value is L min =
ln 10091 ≈ 0 . 0943 .
Choosing γ to be 310 ln 91 is mainly because of its connection to label smoothing with (cid:15) = 0 .
1 , which is the value chosen in [109].
But other choices of γ can also lead to the same sharpness-reduction bias in practice; see Figure 20.
[Figure]