A NEW CHARACTERIZATION OF THE EDGE OF STABILITY BASED ON A SHARPNESS MEASURE AWARE OF BATCH GRADIENT DISTRIBUTION
A BSTRACT
For full-batch gradient descent (GD), it has been empirically shown that the sharpness, the top eigenvalue of the Hessian, increases and then hovers above 2 / (learning rate), and this is called “the edge of stability” phenomenon.
However, it is unclear why the sharpness is somewhat larger than 2 / (learning rate) and how this can be extended to general mini-batch stochastic gradient descent (SGD).
We propose a new sharpness measure (interaction-aware-sharpness) aware of the interaction between the batch gradient distribution and the loss landscape geometry.
This leads to a more refined and general characterization of the edge of stability for SGD.
Moreover, based on the analysis of a concentration measure of the batch gradient, we propose a more accurate scaling rule, Linear and Saturation Scaling
Rule (LSSR), between batch size and learning rate.
I NTRODUCTION
For full-batch GD it has been empirically observed that the sharpness, the top eigenvalue of the
Hessian, increases and then hovers above 2 / (learning rate) (Cohen et al., 2021) as the training proceeds.
This observation can provide a link between two empirical results regarding generalization, (i) using larger learning rates for GD can generalize better (Bjorck et al., 2018;
Li et al., 2019b;
Lewkowycz et al., 2020; Smith et al., 2020) and (ii) minima with low sharpness tend to generalize better (Hochreiter & Schmidhuber, 1997; Keskar et al., 2017).
This observation has a significant implication in existing neural network optimization convergence analyses since it is contrary to the frequent assumption that ‘the learning rate is less than 2 /β (here β is an upper bound of the Hessian top eigenvalue)’, which ensures the decrease in the training loss (Nesterov, 2003; Schmidt, 2014;
Martens, 2014; Bottou et al., 2018).
Even though the training loss evolves non-monotonically over short timescales due to the violation of the assumption, interestingly, the loss is observed to decrease over long timescales consistently.
This regime in which GD typically occurs has been referred to as ‘the edge of stability (EoS)’
(Cohen et al., 2021).
There remain many aspects that are not clearly explained about the EoS regime.
For example, it is not clear why and to what extent the sharpness hovers above 2 / (learning rate).
Moreover, the inherent mechanism is not yet elucidated for the unstable optimization to occur at the EoS consistently while prevented from entirely diverging.
How this phenomenon can be generalized beyond GD, especially to mini-batch SGD, is still an open question.
In this paper we provide a new characterization of the EoS for SGD, which can serve as an answer to the above questions.
As a tool to analyze the optimization process of SGD, we first propose a sharpness measure of neural network loss landscape aware of SGD batch gradient distribution (hence capturing the interaction between SGD and the loss landscape), which we refer to as the interaction-aware-sharpness (IAS) (Section 2).
Based on this measure, we define the stable and unstable regions in the neural network parameter space.
We then scrutinize both theoretically and empirically the transition process of the iterate from the stable to the unstable region (Section 4.1) and the mechanism to escape from the unstable region, i.e., how the optimization can occur at the
EoS. We interpret the latter mechanism based on the non-quadraticity of the loss and the presence of asymmetric valleys in the loss landscape (He et al., 2019)
(Section 4.2).
Based on these analyses, we propose the notion of implicit interaction regularization (IIR), i.e., the
IAS is implicitly bounded during SGD, as an implicit bias of SGD (Section 4.3).
The value that IAS is bounded by is the ratio of a concentration measure of the batch gradient distribution of SGD to the learning rate.
This is a more refined characterization of the EoS, as it shows that IAS does not hover above a certain value, but rather hovers around.
More importantly, it can be naturally applied to SGD since we do not make any impractical assumptions on the batch size or learning rate.
Our new characterization of the EoS leads to a novel scaling rule between batch size and learning rate, from the idea of preserving a similar level of IIR (Section 5).
This scaling rule, referred to as the Linear and Saturation Scaling Rule (LSSR), recovers the well-known linear scaling rule (LSR) (Jastrz˛ebski et al., 2017; Masters & Luschi, 2018;
Zhang et al., 2019; Shallue et al., 2018; Smith et al., 2020; 2021) for small batch sizes and reduces to no scaling (due to saturation) for large batch sizes.
G RADIENT D ISTRIBUTION AND L OSS L ANDSCAPE
In this section, we review some concepts required for further discussion.
See Appendix A for a quick reference for the notations.
To simplify the notations, we often omit the dependence on some variables and the subscript of the expectation operation when clear from the context.
For a learning task, we use a parameterized model with model parameter θ ∈ Θ ⊂ R m .
Then we train the model using training data D
= { x i } ni =1 and a loss function ℓ ( x ; θ ) .
We denote the (total) training loss by L ( θ ) ≡ 1 n (cid:80) ni =1 ℓ ( x i ; θ ) for training data D .
At time step t , we update the parameter θ t using GD: θ t +1 =
θ t − η ∇ θ L ( θ t ) with a learning rate η > 0 , or using SGD: θ t +1 =
θ t − ηg ξ ( θ t ) with a mini-batch gradient g ξ ( θ t ) ≡ 1 b (cid:80) x ∈B tξ ∇ θ ℓ ( x ; θ t ) ∈
R m for a mini-batch B tξ ⊂ D of size b ( 1 ≤ b ≤ n ).
Here, we use the subscript ξ to denote the random batch sampling procedure.
Now, we are ready to introduce some important matrices, C b , S b , and H .
First, we define the covariance C b ( θ ) ≡ Var ξ [ g ξ ( θ )]
=
E ξ (cid:104) ( g ξ ( θ ) − E ξ [ g ξ ( θ )]) ( g ξ ( θ ) − E ξ [ g ξ ( θ )]) ⊤ (cid:105) ∈
R m × m and the second moment S b ( θ ) ≡ E ξ [ g ξ ( θ ) g ξ ( θ ) ⊤ ] ∈ R m × m of the mini-batch gradient g ξ ( θ ) over batch sampling for a batch size 1 ≤ b ≤ n .
The covariance C b and the second moment S b satisfy not only
C b = S b − S n but also the following equation (Hoffer et al., 2017; Li et al., 2017; Wu et al., 2020):
[Equation]
sampling
[Equation]
replacement
[Equation]
replacement
[Equation]
We provide a self-contained proof of (1) in Appendix B.1.
We note that, for sampling without replacement, many previous works approximate γ n,b ≈ 1 assuming b ≪ n (Jastrz˛ebski et al., 2017; Hoffer et al., 2017; Smith et al., 2021), but we consider the whole range of 1 ≤ b ≤ n ( 0 ≤ γ n,b ≤ 1 with γ n, 1 = 1 and γ n,n = 0 ).
Second, we define the Hessian H ( θ ) =
∇ 2 θ L ( θ ) =
n (cid:80) n i =1 ∇ 2 θ ℓ ( x i ; θ ) ∈
R m × m and denote the i -th largest eigenvalue and its corresponding normalized eigenvector by λ
i ( H ) ∈ R and q i ( H ) ∈
R m , respectively, for i
= 1 , · · · , m .
The operator norm ∥
H ∥ ≡ sup ∥ u ∥ =1 ∥
Hu ∥ of H is equivalent to the top eigenvalue λ 1 .
We emphasize that C b and S b represent the stochasticity of the batch gradients, and H represents the loss landscape geometry.
Therefore, we can write one of our goals as follows: we aim to understand how the loss landscape geometry ( H ) and the gradient distribution ( S b ) interact with each other during SGD training.
We investigate this “interaction” in terms of matrix multiplication HS b .
To be specific, we consider the trace tr( HS b ) and its normalized value tr( HS b ) tr( S b ) , and we call the latter interaction-aware sharpness :
Definition 1 (Interaction-Aware Sharpness (IAS)) .
[Equation]
Here, tr( HS b ) ≤ ∥ H ∥ tr( S b ) , i.e., ∥ H ∥ S b ≤ ∥ H ∥ , and the equality holds only when every g ξ is aligned in the direction of the top eigenvector q 1 of H .
R ELATED W ORK
Some studies investigate the interaction between the gradient distribution and the loss landscape geometry represented by tr( HS b ) in the context of escaping efficiency (Zhu et al., 2019, Section 3.1), stationarity (Yaida, 2019, Section 2.2), and convergence (Thomas et al., 2020, Section 3.1.1).
However, they require some additional assumptions like stochastic differential equation (SDE) approximation of SGD (Zhu et al., 2019), the existence of a stationary-state distribution of the model parameter (Yaida, 2019, Section 2.3.4), and strong convexity of the training loss function (Thomas et al., 2020), respectively.
In this paper, we provide a new insight into the interaction tr( HS b ) without these assumptions.
Convergence of full-batch GD ( b = n ) has been instead analyzed with an upper bound on the interaction tr( HS n ) with further assumptions for the stable optimization, such as β -smoothness of the objective and 0 < η < 2 β (e.g., η = 1 β ) (Nesterov, 2003; Schmidt, 2014; Martens, 2014;
Bottou et al., 2018).
However, it may lose useful information of the interaction between H and
S n .
Moreover, when we train a standard neural network with GD in practice, ∥ H ∥ ( ≤ β ) increases in the early phase of training and the iterate enters the EoS where ∥ H ∥ ⪆ 2 η , i.e., η ⪆ 2 ∥ H ∥ ≥ 2 β .
This contradicts with the assumption for stable optimization and the iterate exhibits unstable behavior with a non-monotonically decreasing loss (Xing et al., 2018; Wu et al., 2018; Cohen et al., 2021).
We further extend this discussion of unstable dynamics for GD at the EoS to the case of SGD.
From the generalization perspective, many studies focus on the implicit bias of SGD toward a better generalization (Neyshabur, 2017; Zhang et al., 2021;
Gunasekar et al., 2017; Soudry et al., 2018;
Jastrz˛ebski et al., 2020; 2021; Barrett & Dherin, 2021; Smith et al., 2021).
There are mainly two factors known to correlate with the generalization performance: the batch gradient distribution during training (Hoffer et al., 2017;
Jastrz˛ebski et al., 2017; Smith & Le, 2018; Zhu et al., 2019) and the sharpness of the loss landscape at the minimum (Hochreiter & Schmidhuber, 1997; Keskar et al., 2017; Dinh et al., 2017; Jiang et al., 2020; Foret et al., 2021; Kwon et al., 2021).
We provide a link between the batch gradient distribution and the sharpness that the model is implicitly regularized to have a low sharpness when the second moment of the batch gradient is large (see Section 4.3).
We provide further discussion to reconcile our arguments with some previous studies in Appendix D
A NEW CHARACTERIZATION OF THE E DGE OF S TABILITY 4. U NSTABLE O PTIMIZATION
Using the second-order Taylor expansion of the total training loss L ( θ ) at θ t , the change in the loss
L t =
L ( θ t ) as the SGD iterate moves from θ t to θ t +1 at time step t can be expressed as follows: [Equation]
[Equation]
where δ t = θ t +1 − θ t = − ηg ξ .
Given θ t , the expected loss difference over batch sampling ξ is
[Equation]
[Equation]
( Taking E ξ for the both sides of (3)) (5)
[Equation]
[Equation]
where ϵ =
O ( E ξ [ ∥ δ t ∥ 3 ]) and E ξ [ X ] is the conditional expectation of X given θ t .
We consider a bounded region with a finite maximum loss near and including θ t in which the loss is approximately quadratic.
Then, when the following instability condition is met within the region, the loss (close to the expected loss) tends to increase and the iterate tends to escape from the region:
[Equation]
[Figure]
we mark the end point with ‘x’ when the iterate enters the unstable region.
For mean squared error (MSE) loss (bottom right), we plot the graph for a few more steps after the iterate enters the EoS. We train 6CNN on
CIFAR-10-8k with η = 0 .
02 (see Remark at the end of Section 4.2).
Theorem 1.
For SGD on a quadratic L , the expected loss increases, i.e., E ξ [ L t +1 ] − L t > 0 , if and only if θ t satisfies the instability condition tr( HS b ) tr
( S n ) > 2 η .
Furthermore, if the batch gradient g ξ is normally distributed, then the following inequalities hold for any positive x > 0 :
[Equation]
where β
= η 2 γ n,b b
[Equation]
The proof is deferred to Appendix B.2.
From the above theorem, we define unstable and stable region :
[Equation]
respectively.
It has been empirically shown that, for full-batch GD, ∥ H ∥ increases and then hovers above 2 /η
and Cohen et al.
(2021) mark the EoS with { θ ∈ Θ : ∥ H ( θ ) ∥
= 2 η } , but we mark with
[Equation]
We emphasize the superiority of our ∂ S over the previous EoS that (i) ∂ S provides a clearer “edge” since it considers the interaction between the gradient direction and the Hessian and that (ii) it is more general since it applies to SGD with any batch size 1 ≤ b ≤ n .
Figure 1 (top row) empirically validates (7), showing the normalized loss difference E ξ [ L t +1 ] − L t tr( S n ) against tr( HS b ) tr( S n ) in the early phase of training before entering the unstable region.
This result implies that the training loss L ( θ ) is approximately locally quadratic, i.e., ϵ ≈ 0 , in the early phase.
Especially, for full-batch GD ( b = n ), the instability condition can be rewritten as ∥ H ∥ S n > 2 η
and we have the following relationship between the loss difference L t +1 − L t and ∥ H ∥ S n from (7):
[Equation]
[Figure] [Figure] [Figure]
Figure 1 (bottom row) shows ∥ H soars from 0 in a few steps after ∥ H ∥ exceeds 2 η , satisfying (12) with ϵ ≈ 0 , before the iterate enters the unstable region.
This result is consistent with the following theorem for a quadratic loss L and generalized momentum GD with ( β 1 , β 2 ) :
[Equation]
where β 1 , β 2 ∈ [0 , 1) .
Here, we have vanilla GD when β 1 = β 2 = 0 , Polyak momentum when β 1 ∈ (0 , 1) and β 2 =
(Polyak, 1963), and Nesterov momentum when β 1 = β 2 ∈ (0 , 1) (Nesterov, 1983).
We will focus on the vanilla GD as it can be easily extended to the generalized momentum variants (see Appendix C.6 for details).
The proof is deferred to Appendix B.3.
Theorem 2.
For generalized momentum GD with ( β 1 , β 2 ) on a quadratic L , if 0 < λ i <
η γ ( β 1 , β 2 ) <
λ 1 for all i ̸
= 1 where γ ( β 1 , β 2 )
= 1+ β 1 1+
β 2 , then q ⊤ 1 δ t oscillates and diverges with the exponential growth of | q ⊤ 1 δ t | =
Θ( e ct ) for some c > 0 .
Moreover, | cos( q 1 , δ t ) | and ∥ H ∥
S n increase to 1 and λ 1 , as t → ∞ , respectively, with 1 − | cos( q 1 , δ t ) | , λ 1 − ∥
H ∥ S n =
O ( e − 2 ct ) .
Note that γ ( β 1 , β 2 ) = 1 for the vanilla GD.
To summarize, if ∥ H ∥ exceeds 2 η , then ∥ H ∥
S n increases towards ∥ H ∥ with the exponential convergence rate and also exceeds 2 η in a few steps, i.e., the iterate enters the unstable region U .
Together with Theorem 1, if we consider a bounded subregion V ⊂ U with a finite maximum loss near and including θ t in which the loss is approximately quadratic, then the iterate tends to escape from the region V (Nar & Sastry, 2018;
Wu et al., 2018; Cohen et al., 2021).
4. N ON - QUADRATICITY , A SYMMETRIC V ALLEYS AND THE E DGE OF S TABILITY
In the previous section, we have shown that the training loss is approximately locally quadratic before the iterate enters the unstable region.
However, after the iterate enters the unstable region, i.e., tr( HS b ) tr( S n ) reaches and exceeds 2 η , the step size is relatively large for the sharp loss landscape so that the iterate jumps across the valley (Jastrz˛ebski et al., 2019), and the higher-order terms ϵ in (7) and (12) become non-negligible and cause a different behavior of the iterate than in the stable region.
Figure 2 shows empirical evidences for the non-quadraticity .
After the SGD/GD iterate enters the unstable region, when the instability condition tr( HS b ) tr( S n ) > 2 η is met, the normalized increase in the loss (cid:12)(cid:12)(cid:12) E ξ [ L t +1 ] − L t tr( S n ) (cid:12)(cid:12)(cid:12) is often smaller than η 2 2 (cid:12)(cid:12)(cid:12) tr( HS b ) tr( S n ) − 2 η (cid:12)(cid:12)(cid:12) from (7) and (12) (blue line), which is the case for a locally quadratic function.
This results in a gentle slope less than η 2 2 .
[Figure]
||
[Figure]
||
[Figure]
n || H ||
[Figure]
|| H || S n <
[Figure]
||
H || S n >
[Figure]
-min
L
L || H || S n <|| H || S
n >
We hypothesize that due to this non-quadraticity of the training loss, the iterate is discouraged from staying within the unstable region.
Note that, for a globally quadratic loss, when the iterate is in the unstable region it diverges within the unstable region.
Figure 3 demonstrates the asymmetric valley (He et al., 2019) that one side is sharp and the other is flat.
In Figure 3 (left), we evaluate the directional sharpness ∥
H α ∥ S n along the gradient descent direction − η ∇ L ( θ )
where H α ≡ H ( θ − αη ∇ L ( θ )) for α ∈ 14 × [1 , 2 , 3 , 4 , 5] , and compare ∥ H α ∥ S n ( θ ) with ∥ H ∥ S n ( θ ) .
At the sharp side, it has a high ∥ H ∥ S n > 2 η (blue) with the gradient ∇ L and the top eigenvector q 1 ( H ) of the Hessian being highly aligned (cf.
Theorem 2).
However, when the loss landscape gets far from being quadratic, the Hessian and its top eigenvector can change abruptly, q 1 ( H α ) would not always be aligned with q 1 ( H ) and ∇ L ( θ ) , and ∥ H α ∥ S n tends to decrease.
This would be a possible explanation for the tendency of decreasing and then oscillating ∥ H ∥ S n .
See Appendix C.3 for detailed empirical evidences of the above arguments.
Figure 3 (right) similarly shows that when the iterate is at a sharp side of the valley, it tends to jump to the other side of a flatter area, and vice versa.
To summarize, we make the following observations for GD in order: (i) ∥
H ∥ increases in the beginning (the progressive sharpening (Jastrz˛ebski et al., 2019; 2020; Cohen et al., 2021)), (ii)
∥
H ∥ exceeds 2 η , (iii) the gradient ∇ L becomes more aligned with the top eigenvector q 1 ( H ) in a few steps, (iv)
∥
H ∥
S n reaches the threshold 2 η and the iterate jumps across the valley, (v)
∥
H ∥
S n tends to decrease due to the non-quadraticity, and it repeats this process, while ∥ H ∥ S n oscillating around 2 η .
We observe a similar behavior with oscillating tr( HS b ) tr( S n ) around 2 η for SGD.
It requires further investigation into the exact underlying mechanisms (e.g., the progressive sharpening) (Arora et al., 2022; Li et al., 2022; Damian et al., 2022; Zhu et al., 2022) and we leave it as a future work.
Remark (Experiments in Section 4.1 and 4.2) .
We report the experimental results using vanilla
SGD/GD without momentum and weight decay, constant learning rate, and no data augmentation.
We train a simple 6-layer CNN (6CNN, m = 0 .
51 M) on CIFAR-10-8k where D ATASET - n denotes a subset of D ATASET with |D| = n and k = 2 10 = 1024 .
See Appendix C.1-C.3 for the results from other datasets, learning rates and networks (ResNet-9 with m = 2 .
3 M (He et al., 2016) and
WRN-28-2 with m = 36 M (Zagoruyko & Komodakis, 2016)
where m = dim( θ ) ).
4. I MPLICIT I NTERACTION R EGULARIZATION (IIR)
In the previous sections, we have shown the SGD iterate is implicitly discouraged from staying in the unstable region.
Now, we are ready to investigate this property from the regularization perspective.
First, to understand the effect of batch size b on the batch gradient distribution, we define the following concentration measure ρ b :
[Figure]
We also observe a sharp increase in ∥ H ∥ S n right after ∥ H ∥ exceeds 2 η .
Curves are plotted for every step.
We train a model on CIFAR-10-8k ( n = 2 13 ) using cross-entropy loss with η = 0 .
02 , respectively.
c
: We plot curves and when trained with various b .
After a few steps, the curves in (c) reach the threshold 2 ρ b η (see (d) together) which increases as b becomes larger when b ≪ n = 2 13 , and saturates to
ρ b η ≈ 2 η when b is large.
Curves are smoothed for visual clarity.
Definition 2 (a concentration measure of the batch gradient) .
We define ρ b as the ratio of the squared norm of the total gradient ∥∇ L ∥ 2 to the expected squared norm of the batch gradients
E ξ [ ∥ g ξ ∥ 2 ] , i.e.,
[Equation]
Here, we can write ∥∇ L ∥ 2 = ∥
E ξ [ g ξ ] ∥ 2 and thus the ratio ρ b = ∥
E ξ [ g ξ ] ∥ 2 E ξ [ ∥ g ξ ∥ 2 ] ≤ 1 is similar to the square of the mean resultant length ¯ R 2 b ≡ ∥ E ξ [ g ξ ∥
g ξ ∥ ] ∥ 2 ≤ 1 of the batch gradient g ξ (Mardia et al., 2000), especially when std ξ [ ∥ g ξ ∥ ] ≪ E ξ [ ∥ g ξ ∥ ] (see Appendix C.5 for empirical evidences).
Both ρ b and ¯
R 2 b are concentration measures and have lower values when the batch gradients g ξ are more scattered.
Therefore, it is natural to expect that the ratio ρ b is small for a small batch size b , and we will revisit this in more detail in the following section (cf.
We also note that ρ n = ¯ R 2 n = 1 .
Yin et al.
) call 1 / ( nρ 1 ) the gradient diversity
Now, we can rewrite the instability condition tr( HS b ) tr( S n ) > 2 η (multiplying both sides by ρ b ) as follows: [Equation]
[Equation]
From Theorem 1 and 2, the instability condition (16) implies that IAS
∥ H ∥ S b is implicitly regularized and bounded to be less than 2 ρ b η .
We name this Implicit Interaction Regularization (IIR) .
We argue that the upper constraint 2 ρ b η in IIR is crucial in determining the SGD dynamics.
With a low constraint,
SGD strongly bounds IAS
∥ H ∥ S b .
We also note that IIR affects not only the magnitude ∥ H ∥ but also the directional interaction.
In other words, IIR discourages the batch gradients from aligning with the top eigensubspace of the Hessian that is spanned by a few largest eigenvectors of the Hessian (cf.
Gur-Ari et al.
Figures 4(a)-4(b) show that, for GD ( ρ n = 1 ), IAS ∥
H ∥ S n (red) oscillates around 2 η and exhibits IIR.
This result is consistent with Cohen et al.
(2021) that ∥ H ∥ hovers above 2 η for GD.
This is because, as mentioned earlier, 2 η ≈ ∥
H ∥ S n ≤ ∥ H ∥ and the equality holds only when the gradient ∇ L and the top eigenvector q 1 of H are aligned, which is in general not the case.
For this reason, IIR provides a tighter relation and more clearly identifies the EoS than Cohen et al.
These results are also consistent with Theorem 2 that ∥ H ∥ S n suddenly increases from 0 to 2 η in a few steps after ∥ H ∥ exceeds 2 η (see Appendix C.3-C.4 for more).
Moreover, IIR also applies to a general SGD training with 1 ≤ b ≤ n .
Figure 4(c)-4(d) show IIR for SGD with different batch sizes.
The upper bound 2 ρ b η of ∥ H ∥ S b according to IIR is higher when using a larger batch size, but limited to less than 2 η ( ρ b ≤ 1 ).
We will further discuss this behavior with an investigation of ρ b in the following section.
L INEAR AND S ATURATION S CALING R ULE (LSSR)
In this section, we first introduce two previous scaling rules on how to tune the learning rate for varying batch sizes.
Then we explain why they fail and propose a new scaling rule based on IIR.
The ratio b/η of batch size b to learning rate η has long been believed as an important factor influencing the generalization performance, and the test accuracy has observed to be similar when trained with the same ratio b/η = b ′ /η
′ , i.e., b ′ =
kb and η ′ =
kη for k > 0 .
This is called the linear scaling rule (LSR) (Krizhevsky, 2014; Goyal et al., 2017; Jastrz˛ebski et al., 2017; Smith & Le, 2018; Zhang et al., 2019).
They argued that LSR holds because θ t + k − θ t =
− ηb (cid:80) k
− 1 i =0 (cid:80) x ∈B t
+ i ξ
[Equation]
[Equation]
t
[Equation]
i ) ≈ −
[Equation]
(cid:80)
[Equation]
(cid:80)
[Equation]
ℓ ( x ; θ t ) assuming ∇ θ ℓ ( θ t + i ) ≈ ∇ θ ℓ ( θ t ) for 0 ≤ i < k , where B t : t + k ξ ≡ ∪ k − 1 i =
B t + i ξ and |B t : t + k ξ | =
kb = b ′ .
However, the assumption is not accurate since the gradient oscillates mostly with a negative cosine value cos( g ξ ( θ t ) , g ξ ( θ t +1 )) <between two consecutive gradients after entering the EoS (see Figure 24-25 in Appendix C.3).
Moreover, LSR is known to fail when the batch size is large (Jastrz˛ebski et al., 2017; Masters & Luschi, 2018; Zhang et al., 2019; Smith et al., 2020; 2021).
On the other hand, Krizhevsky (2014); Hoffer et al.
(2017) proposed the square root scaling rule (SRSR) with another ratio √ b/η to keep the covariance of the parameter update constant for b ≪ n based on Var ξ [ ηg ξ ]
= η 2 C b = γ n,b η 2 b
[Equation]
However, Shallue et al.
(2018) showed that both LSR and SRSR do not hold in general.
Based on the analysis of IIR with a new ratio 2 ρ b /η in the previous section, we explore why LSR fails in the large-batch regime and provide a more accurate rule to achieve similar generalization performance of the models trained with various choices of batch size and learning rate pairs ( b, η ) .
To this end, we investigate the concentration measure ρ b = tr( S n ) / tr( S b ) .
By combining two equations, C b = S b − S n (by definition) and C b = γ n,b b ( S 1 − S n ) in (1), we can obtain S b =
C b + S n = γ n,b b S
+ (1 − γ n,b b ) S n .
Therefore, we have tr( S b ) = γ n,b b tr( S 1 )
+ (1 − γ n,b b ) tr( S n ) , which leads to the following equation:
[Equation]
[Equation]
(17) from (15) where ρ = ρ 1 = tr( S n ) / tr( S 1 ) .
Note that ρ is (much) smaller than 1 because ∇ θ ℓ ( x i ) has different direction for each x i and tr( S n ) =
∥∇ L ∥ 2 = ∥ 1 n (cid:80) i ∇ θ ℓ ( x i )
∥ 2 ≤ 1 n (cid:80) i ∥∇ ℓ θ ( x i ) ∥ 2 = tr( S 1 ) .
In other words, 1 /ρ is (much) larger than 1 (see Appendix C.5).
Figure 5 (left) demonstrates a new scaling rule with the ratio ρ b /η , called the Linear and Saturation
Scaling Rule (LSSR), with the two regimes that (i) ρ b is almost linear when b ≪ n (linear regime) and (ii) ρ b saturates when b is large (saturation regime).
It depends on which part of the denominator
[Figure]
accuracy for models trained with a large number of pairs of b η ) on 8k 8k , STL 4k, and Tiny
32k from left right, bottom).
It does not follow either LSR or
, but tends follow
We also plot
Note that they all and thus a
1 indicates relationship.
( ∗ ) in (17) dominates the other.
First, when b ≪ n , then γ n,b /b is not very small and the first term γ n,b b 1 ρ dominates the second term 1 − γ n,b b since 1 ρ ≫ 1 .
Second, as b becomes large, γ n,b /b ≈ 0 and the second term ( ≈ 1 ) dominates the first term.
Thus, ρ b saturates to 1 and is not linearly related to b , and LSR is no longer valid.
The above arguments also hold for the batches sampled with replacement where the only modification is γ n,b = 1 , ∀ b in (17).
Figure 5 (right) empirically supports LSSR with the test accuracies when trained with various combinations of pairs ( b, η ) .
To be specific, the optimal learning rate is almost linear to b when b is small, but it saturates when b is large.
We also plot ρ b /η =
C (the yellow dashed curve) for some ρ and C which shows a theoretical prediction of pairs ( η, b ) that yield the optimal performance.
Note that Figure 8 of Shallue et al.
(2018, Section 4.7) shows similar “linear and saturation” behaviors supportive of LSSR on other datasets (see also
Figure 7 of Zhang et al.
(2019, Section 4.3)).
Remark (Experiments in Section 4.3 and 5) .
We train models using vanilla SGD/GD without momentum and weight decay, constant learning rate, and no data augmentation.
For Figure 5, we use subsets of the datasets CIFAR-10 (Krizhevsky & Hinton, 2009), CIFAR-100 (Krizhevsky &
Hinton, 2009), STL-10 (Coates et al., 2011), and Tiny-ImageNet (a subset of ImageNet (Deng et al., 2009) with 3 × 64 × 64 images and 200 object classes).
We use a large number of epochs (800) and batch normalization (Ioffe & Szegedy, 2015) to achieve a zero training error even with a large b and a small η .
In the lower right corner (red area) of each heatmap in Figure 5 (right), when b is too large or η is too small so that ∥ θ t +1 − θ t ∥ =
η ∥ g ξ ∥ is too small, it requires an exponentially large number of steps for the iterate to enter the EoS. Thus, in this case, the assumption in Goyal et al.
(2017), ∇ θ ℓ ( θ t ) ≈ ∇ θ ℓ ( θ t + i ) for 0 ≤ i < k , approximately holds and the reasoning on LSR is valid.
However, this only holds for a non-practical ( b, η ) which shows a suboptimal performance.
See Appendix C.4-C.5 for the results from other networks and hyperparameters.
C ONCLUSION
From an analysis of unstable dynamics of SGD and the instability condition, we clearly mark the edge of stability with the interaction-aware sharpness ∥ H ∥ S b and show the presence of the implicit regularization effect on the interaction between the gradient distribution and the loss landscape geometry (IIR).
Moreover, introducing the concentration measure ρ b of the batch gradient, we link the second moment of the gradient distribution and the sharpness of the loss landscape, and propose a new scaling rule called Linear and Saturation Scaling Rule (LSSR).
Due to the simplicity of the analysis, we hope that our insights will motivate the future work toward understanding various learning tasks. A CKNOWLEDGEMENTS
SL is also the lead author and the corresponding author.
This work was done while SL was at Korea
Institute for Advanced Study.
SL was supported by a KIAS Individual Grant (AP083601) via the Center for AI and Natural Sciences at Korea Institute for Advanced Study. C. Jang was supported in part by IITP Artificial Intelligence Graduate School Program for Hanyang University funded by MSIT (Grant No. 2020-0-01373) and NRF/MSIT (Grant No. 2021M3E5D2A01019545).
This work is supported by the Center for Advanced Computation at Korea Institute for Advanced Study
R EPRODUCIBILITY S TATEMENT
We refer the reader to the following pointers for reproducibility:
• Codes to reproduce some Figures: Supplementary Material.
• Proofs of the claims: Appendix B. • Experimental settings : Remarks at the end of Sections 4.2 and 5, and Appendix C. • Additional Figures for other hyperparameter settings: Appendix C.1-C.6.
A N OTATIONS
We summarize the notations for a quick reference.
t ∈ N time step θ ∈ Θ [Equation] ( or indexed θ [Equation]
; dim [Equation] = m model parameter x [Equation]
X ( or indexed x i ) training sample D = { x i } ni [Equation] |D
[Equation] = n training data ℓ ( x
[Equation]
θ ) loss function
L [Equation]
θ [Equation] ≡ [Equation]
n n (cid:88) i =
[Equation]
(cid:88) x ∈D ℓ ( x ; θ ); L t =
L ( θ t ) (total) training loss η >learning rate
B ⊂ D ( or indexed B [Equation] , B tξ ) batch b = |
[Equation]
|
[Equation]
≤ n batch size g
[Equation] [Equation]
x
[Equation]
ξ
[Equation]
x
[Equation]
ξ ∇ ℓ ( x ; θ ) batch gradient θ t
[Equation]
GD θ
[Equation]
SGD
[Equation]
displacement/velocity vector
[Equation]
the covariance of the batch gradient
S
[Equation]
the second moment of the batch gradient γ n,
[Equation]
= (cid:40)
[Equation]
for sampling without replacement for sampling with replacement sampling coefficient
[Equation]
Hessian
[Equation]
i
[Equation]
the Euclidean ℓ 2 -norm of a vector u
[Equation]
∥
Au
[Equation]
∥
[Table]
the spectral (operator) norm of a matrix
A
[Table]
[Equation]
[Table]
the i -th largest eigenvalue of the Hessian
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
the corresponding i -th eigenvector of the Hessian
[Table]
[Equation]
[Table]
the trace of a square matrix
A
[Table]
concentration measure of the batch gradient
[Table]
concentration
[Table]
of the per
[Table]
example gradient
[Equation]
unstable regime
[Equation]
stable regime
[Equation]
the edge of stability
[Table]
hyperparameters for generalized momentum GD
B P ROOFS AND R EMARKS B. P ROOF OF (1)
We provide a proof of (1) to make the paper self-contained.
Similar proofs are given in Li et al.
Hoffer et al.
; Wu et al.
Proof.
We start with
[Equation]
where ∇L = [ ∇ ℓ 1 , · · · , ∇ ℓ n ] ∈ R m × n , ℓ i = ℓ ( x i ) , the random vector w = [ w 1 , · · · , w n ] ⊤ , each element of which represents 1 b × “how many times the index i is sampled in B ξ ”, and v = bw .
In case of sampling with replacement, we have v = v (1) + · · · + v ( b ) where v ( i ) represents sampling of a single sample.
Thus, Var ξ [ v ] =
b Var[ v (1) ] .
We have E [ v (1) ]
= 1 n 1 and
[Equation]
23) where |B (1) | = 1 .
Thus,
[Equation]
In case of sampling without replacement, we have E ξ [ v ] = bn 1 and
E
[Equation]
else , (25) where C ( n 1 , r 1 ) is the number of r 1 -combinations from a set of n 1 elements.
This leads to
[Equation] [Equation]
and
Var
[Equation] [Equation]
Putting the two cases together, from (22), (24) and (29), we have
[Equation]
[Equation]
where
[Equation]
B. P ROOF OF T HEOREM
First, we expect the variance of the loss at the next iteration ( t + 1) to be small for a sufficiently large b .
Lemma 3.
If the batch gradient g ξ is normally distributed and the loss L is quadratic, then
[Equation]
where v = ( I − ηH ) ∇ L .
Proof.
Since L is quadratic
[Equation]
we have the following:
[Equation]
Then
[Equation]
the variance of L t +1 is Var ξ
[ L t +1 | θ t ] =
E ξ [ Q 2 ξ ] − E ξ [ Q ξ ] 2 where Q ξ = − ηv ⊤ ε + η 2 2 ε ⊤ Hε .
First, we can obtain the square of the expected value of Q ξ as follows:
[Equation]
where the last equation holds since Σ is symmetric and (cid:80) i,j H i,j Σ i,j = (cid:80) i,j H i,j Σ j,
i = (cid:80) i
[ H Σ] i,i = tr( H Σ) .
Second, we have the expected value of the square of Q ξ as follows:
[Equation] [Equation]
and, since E ξ [ ε i ε j ε k ] = 0 and E [ ε j ε k ε j ′ ε k ′ ]
=
E [ ε j ε k ] E [ ε j ′ ε k ′ ]
+ E [ ε j ε j ′ ]
E [ ε k ε k ′ ]
+ E [ ε j ε k ′ ]
E [ ε j ′ ε k ] by Isserlis’ theorem (Isserlis, 1918) for zero-mean normal random vector ε ,
E
[Equation]
[Equation]
Finally, we have the variance
[Equation] [Equation]
From Lemma 3, we can easily derive the following theorem from Chebyshev’s inequality:
Theorem 4.
For SGD on a quadratic loss L , the expected loss at the next step ( t + 1) is
[Equation]
Further, if the batch gradient g ξ is normally distributed, then given δ > 0 and α ≥ (cid:112) β/δ , we have, with probability of at least 1 − δ ,
[Equation]
where β =
[Equation]
and
[Equation]
=
[Equation]
However, we can replace this inequality with a better exponential inequality.
To do so, we need a generalized Lemma 6 of the following Lemma 5 (Laurent & Massart, 2000):
Lemma 5 (Laurent & Massart (2000)) .
For i.i.d.
Gaussian variables ( Y 1 , · · · , Y D ) with mean 0 and variance 1, the following inequality holds for any positive x :
[Equation]
where Z = (cid:80) i a i Y 2 i .
Lemma 6.
For i.i.d.
Gaussian variables ( Y 1 , · · · , Y D ) with mean 0 and variance 1, the following inequality holds for any positive x :
[Equation]
where Z = (cid:80) i a i Y
i + c i Y i .
Proof.
For a standard normal random variable Y i ∼ N (0 , 1)
[Equation]
let ψ be the Cramér transform of
Y 2 i + c
i Y i − 1 :
[Equation]
where the last inequality holds for 0 < λ < 12 since − λ − 12 log(1 − 2 λ )
= − λ +
[Equation]
(cid:80) [Equation] [Equation] k ≥
[Equation]
k
[Equation]
Therefore,
[Equation]
We can obtain (69) from the following (Birgé & Massart, 1998): if log E [exp( λZ )] ≤ b 1 λ 2 1 − b 2 λ , then, for any positive x > 0 , P ( Z ≥ √ b 1 x + b 2 x ) ≤ exp( − x ) .
Also, for − 12 < λ < 0 , ψ ( λ ) ≤ λ 2 (1+ c 2 i 2 ) and thus log E [exp( λZ )]
= ( ∥ a ∥ 2 + ∥ c
∥ 2 2 ) λ 2 which leads to (70).
Theorem 1.
For SGD on a quadratic L , the expected loss increases, i.e., E ξ [ L t +1 ] − L t > 0 , if and only if θ t satisfies the instability condition tr( HS b ) tr
( S n ) > 2 η .
Furthermore, if the batch gradient g ξ is normally distributed, then the following inequalities hold for any positive x > 0 :
[Equation]
where β
= η 2 γ n,b b (cid:16) v ⊤
[Equation] [Equation]
+
[Equation]
tr( HC 1 HC 1 ) (cid:17) and v = ( I − ηH ) ∇ L .
Proof.
The first part of the statement is trivial from (7) with ϵ = 0 for a quadratic
L.
Now, we focus on the concentration of L t +1 from its expected value: L t +1 − E
[ L t +1 ] =
Z − E [ Z ] for
Z = Q ξ = − ηv ⊤ ε + η 2
ε ⊤ Hε where v =
( I − ηH )
∇ L and ε ∼ N (0 , C b ) from (38).
For
C b = γ n,b b BB ⊤ , we can represent ε = (cid:113) γ n,b b B ˜ ε with ˜ ε ∼ N (0 , I ) and
[Equation]
[Equation]
[Equation] [Equation]
(68) where
A =
η 2 γ n,b 2 b B ⊤ HB =
Q diag( a ) Q ⊤ , c = − η (cid:113) γ n,b b Q ⊤ B ⊤ v and Y =
[ Y 1 , · · · , Y m ] = Q ⊤ ˜ ε .
Therefore, from the rotation invariance, { Y i } are i.i.d.
Gaussian variables with mean 0 and variancelike ˜ ε .
From Lemma 6, we have the following inequalities:
[Equation]
Further, we have
[Equation]
and
[Equation]
These lead to ∥ A ∥ 2 F + ∥ c
∥
[Equation]
[Equation]
v = β 2 and then (9) follows from (70).
Moreover, we have ∥ A ∥ ≤ η 2 γ n,b 2 b ∥
H ∥∥ C 1 ∥ which leads to (8) from (69).
B. P ROOF OF T HEOREM
We consider the following linear homogeneous second-order equation for the sequence { y t } :
[Equation]
and its the characteristic equation
[Equation]
Lemma 7 (Theorem 2.35 in Elaydi (2005)) .
All solutions of (74) (i) oscillate (about zero) if and only if (75) has no positive real roots, and (ii) converges to
if and only if the solutions of (75) x 1 and xsatisfy | x 1 | , | x 2 | < 1 .
Lemma 8 (Theorem 2.37 in Elaydi (2005)) .
The conditions
[Equation]
C1)
[Equation]
C
[Equation]
C1) are necessary and sufficient for the equilibrium point (solution) of (74) to be asymptotically stable (i.g., all solutions converges to 0).
Theorem 2.
For generalized momentum GD with ( β 1 , β 2 ) on a quadratic L , if 0 < λ i <
η γ ( β 1 , β 2 ) <
λ 1 for all i ̸
= 1 where γ ( β 1 , β 2 )
= 1+ β 1 1+
β 2 , then q ⊤ 1 δ t oscillates and diverges with the exponential growth of | q ⊤ 1 δ t | =
Θ( e ct ) for some c > 0 .
Moreover, | cos( q 1 , δ t ) | and ∥ H ∥
S n increase to 1 and λ 1 , as t → ∞ , respectively, with 1 − | cos( q 1 , δ t ) | , λ 1 − ∥
H ∥ S n =
O ( e − 2 ct ) .
We provide a proof of Theorem 2 which is modified from the proofs in Appendix A in Cohen et al.
Proof.
Put a quadratic training loss L ( θ ) = 12 θ ⊤ Hθ + b ⊤ θ + c .
The update rules for generalized momentum GD with ( β 1 , β 2 ) on this quadratic function L ( θ ) are:
[Equation]
77) and thus, for a pair of eigenvalue/eigenvector ( q, λ ) the update rules for the quantities ˜
θ t = q ⊤ θ t +λ q ⊤ b and ˜ δ t = q ⊤ δ t are:
[Equation]
where a = ηλ > 0 and this leads to:
[Equation]
Note that ˜ θ t and ˜ δ t have the same characteristic equation.
Thus, by Lemma 8, ˜ δ t is asymptotically stable when the following three conditions (C1-3) hold:
[Equation]
C3 and C1 always hold because 1 − β 1 > 0 and
[Equation]
Lastly, since
[Equation]
The asymptotic convergence condition (C2) is equivalent to
[Equation]
Therefore, along the direction of q 1 , the sequence { q ⊤ 1 δ t } diverges, and along the direction of q i ( i > 1 ), the sequence { q ⊤ i δ t } converges to 0.
The discriminant D ( a 0 ) for the characteristic equation ϕ ( x ; a 0 ) is positive as shown below where a 0 = 2 γ ( β 1 , β 2 ) , i.e., ϕ ( x ; a 0 ) has two distinct real solutions.
[Equation]
(94) implies that D ( a ) is convex quadratic with respect to a .
Thus, to show that D ( a ) > 0 ( ϕ ( x ; a ) has two distinct real solutions) for all a > a 0 , it is sufficient to show that D ′ ( a ) ≥ 0 for a > a 0 .
The following inequality holds
[Equation]
if a ≥ 1+ β 1 − β 2 + β 1 β 2 (1+ β 2 ).
And, D ′ ( a ) ≥ 0 for a > a 0 since a > a 0 > 1+ β 1 − β 2 + β 1 β 2 (1+ β 2 )from
[Equation]
Now, for a > a 0 , we want to show that x 1 for the dominant solution ( | x 1 | > | x 2 | ) is negative
so that the general solution ˜ δ t =
c 1 x t 1 + c 2 x t 2 = x t(cid:20)
c
1 + c 2 (cid:16) x 2 x 1 (cid:17) t (cid:21) oscillates where c i , x
i ∈ R and i ∈ { 1 , 2 } .
The sum of the two solutions of ϕ ( x ; a ) is x 1 + x 2 = − p 1 and this is negative for a > a from the following:
[Equation]
Thus, we have that x 1 for the dominant solution is negative which leads to an oscillatory behavior of q ⊤ 1 δ t .
And it has the exponential growth of | q ⊤ 1 δ t | =
Θ( | x 1 | t ) =
Θ( e ct ) for c = ln | x 1 | > 0 since [Equation] for t
[Equation]
t [Equation] for some t
[Equation]
( λ 1 violates (92) with λ 1 > 2 η γ ( β 1 , β 2 ) which implies | x 1 | > 1 ).
This makes δ t to be asymptotically aligned with the top eigenvector
q 1 of H , i.e., lim t
[Equation]
˜ δ
[Equation]
= q ⊤ i δ
[Equation]
Moreover, we can obtain the exponential convergence ( 1 − | cos( q 1 , δ t ) |
=
O ( e − 2 ct ) ) as follows:
[Equation]
[Equation]
for t ≥ t 1 for some t 1 .
[Equation]
The last inequality holds because a t = | ˜ δ (1) t /s t | diverges to ∞ and lim x →∞ x 2 − x √ x 2 + 1+1 = 12 ,
Moreover, we have
[Equation] [Equation]
[Equation]
θ
[Equation] [Equation]
118) ∇ θ L ( [Equation]
t [Equation] ∇ θ
[Equation]
(cid:88)
[Equation]
i
[Equation]
where ˜ θ ( i ) t = q ⊤ i θ t + q ⊤
i
λ i b .
Therefore, we have λ 1 − ∥ H ∥ S n = O ( e − 2 ct ) from the following:
[Equation]
for t ≥ t 2 for some t
[Equation]
since lim t →∞ ˜ θ ( i ) t = 0 for i > 1 .
Remark.
Due to the exponential convergence of ∥ H ∥ S n to λ 1 , it only takes a few steps (5-20) for ∥
H ∥ S n to exceed 2 η γ ( β 1 , β 2 ) (see Appendix C.3).
Remark.
In practice, ∥ H ∥
= λ 1 keeps increasing after it exceeds 2 η γ ( β 1 , β 2 ) .
Therefore, we may relax the assumption as the eigenvalues λ 1 ( θ t ) > 2 η γ ( β 1 , β 2 )+ ϵ 1 and ϵ 2 <
λ
i ( θ t ) <
2 η γ ( β 1 , β 2 ) − ϵfor i ̸
= 1 may change within the bounds over t for some ϵ 1 , ϵ 2 , ϵ 3 > 0 ( q i ’s are fixed).
We can draw the same conclusion except that the limit ∥ H ∥
S n may not exist because of varying ∥ H ∥ according to t , but we can ensure that ∥ H ∥ S n eventually exceeds 2 η γ ( β 1 , β 2 ) .
C E XPERIMENTAL S ETTINGS AND A DDITIONAL F IGURES
We report the experimental results using vanilla SGD/GD without momentum and weight decay, constant learning rate, and no data augmentation.
We use a simple 6-layer CNN (6CNN, m = 0 .
51 M),
ResNet-9 (He et al., 2016) 3 ( m = 2 .
3 M ), WRN-28-2 (Zagoruyko & Komodakis, 2016) ( m = 36 M ).
We use subsets of the datasets CIFAR-10/100 (Krizhevsky & Hinton, 2009) 4 , STL-10 (Coates et al., 2011) 5 , and Tiny-ImageNet 6 (a subset of ImageNet (Deng et al., 2009) with 3 × 64 × 64 images and 200 object classes) where D ATASET - n denotes a subset of D ATASET with |D| = n and k= 2 10 = 1024 .
Moreover, we use PyHessian (Yao et al., 2020) 7 to compute the Hessian-vector product (e.g., H ∇ L ), the top eigenvalue λ 1 and its corresponding eigenvector q 1 of the Hessian.
For these computations, we use the power iterations with a batch size of 2k, the tolerance of 0.001, and the maximum iteration of 100.
C. F IGURE
In Figure 1 and Figure 2, we plot E ξ [ L t +1 ] − L t tr( S n ) against tr( HS b ) tr( S n ) , which is equivalent to L t +1 − L t tr( S n ) against ∥ H ∥ S n for GD.
Therefore, we expect the following linear relationship with the slope η 2 2 and the x-intercept 2 η when the training loss L is locally quadratic, i.e., ϵ = 0 :
[Equation]
Figure 1 shows the behavior in the early phase, until the iterate enters the edge of stability.
For GD, they are plotted after ∥ H ∥ exceeds 2 η after which
∥
H ∥ S n starts to increase from 0 to 2 η in a few steps.
For cross-entropy loss, we mark the end point with ‘x’ when the iterate enters the unstable regime.
[Figure]
[ L t + 1 ]
[Figure]
t ) /
t
[Figure]
S
[Figure]
[Figure]
( L t +
L t ) / t r (
S n )
[Figure]
GD
[Figure] [Figure]
[Figure]
[Figure]
s t e p
[Figure] [Figure]
L t
[Figure]
L t
[Figure]
t
[Figure]
t )
[Figure]
See caption of Figure 1 for more details.
[Figure]
t +] L t ) / t r (
[Figure]
n )
[Figure] [Figure]
Figure 2 shows the non-quadraticity (represented by deviation from the blue line in Figure 2 and the following figures) of the training loss function L after the iterate enters the edge of stability
[Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
F IGURE
AND T HEOREM
Figures 20-21 provide additional information of Figure 3.
After these figures, we focus on providing some empirical evidences of Theorem 2 that ∥ H ∥ S n , | cos( q 1 , ∇ L ) | and | q ⊤ 1 ∇ L | increase in a few steps after ∥ H ∥ exceeds 2 η .
[Figure]
[Figure] [Figure]
[Figure] [Figure]
Step 1.
train loss 1.
|| H t
|| S n /||
H t || |cos( q 1 , L )| cos( L ( t ), L ( t +1 )) |cos( q 1 ( H t +1 ), q 1 ( H t ))
|
||
H t || S n
Step 1.
train loss 1.
|| H t
|| S n /||
H t || |cos( q 1 , L )| cos( L ( t ), L ( t +1 )) |cos( q 1 ( H t +1 ), q 1 ( H t ))
|
||
H t || S n
Step 1.
train loss 1.
|| H t
|| S n /||
H t || |cos( q 1 , L )| cos( L ( t ), L ( t +1 )) |cos( q 1 ( H t +1 ), q 1 ( H t ))
|    || H t || S n
Step 1.
train loss 1.
|| H t
|| S n /||
H t || |cos( q 1 , L )| cos( L ( t ), L ( t +1 )) |cos( q 1 ( H t +1 ), q 1 ( H t ))
|    || H t || S n
C.4.
[Figure] [Figure]
details
[Figure] [Figure]
| H || S n || H
||
C.4. SGD
We demonstrate two types of IIR, (i)
∥
H ∥ S b ≤ 2 ρ b η and (ii) tr( HS b ) tr( S n ) ≤ 2 η in Figure 30 and 31, respectively.
They are equivalent
[Figure]
but the latter shows the effects of IIR more clearly as it has the fixed threshold 2 η regardless of b
[Figure] [Figure]
With the upper bound 2 η (red dashed line), this shows the effects of IIR more clearly than Figure 30.
ResNet-9, CE, η = 0 .
08 , and b ∈ { 2 12 , 2 11 , · · · , 2 7 } .
[Figure]
[Figure]
Figure 34-35 provide some additional information of Figure 5.
Figure 36 shows that 1 /ρ ≈is much larger than 1 and
that std[ ∥ g ξ ∥ ] is 2 - 3 × smaller than E
ξ [ ∥ g ξ ∥ ] even in the case of b = 1 .
Therefore, we use the approximation ∥ g ξ ∥
≈ E ξ [ ∥ g ξ ∥ ]
[Figure]
and thus the square of the mean resultant length is similar to the concentration measure ρ b as shown in the following approximation:
[Equation]
[Figure]
[Equation]
[Figure] [Figure]
[Figure]
C. G ENERALIZED MOMENTUM VARIANTS
For generalized momentum variants of GD with ( β 1 , β 2 ) , we have
[Equation]
[Equation]
When ∥ H ∥ is very small in the beginning, we may approximate that L is linear where ∇ θ L ( θ ) =
g is constant with respect to θ in some region.
Then δ t converges to δ satisfying the following equations:
[Equation] [Equation]
and thus
[Equation]
Therefore, together with Theorem 2, we generalize (12) with the following approximation:
[Equation]
where γ ( β 1 , β 2 ) = 1+ β 1 1+2 β 2 .
In Figures C.6 and C.6, we empirically validate (131) similar to Figuresand 4, respectively.
[Equation]
[Equation]
1e      ( L t +
L t ) / t r ( S n ) 1eslope= {+
 [Equation] [Equation]+} 2 ( 1 , 2 ) =
=0.01,( 1 , [Equation] =(0.1,
[Equation]
s t
e
p
4 || H || S n 1e
( L t +
L t ) / t r ( S n ) 1eslope= 2 2 ( 1 , 2 )(1 1 ) 2 ( 1 , 2 ) =
=0.01,( 1 , 2 )=(0.1,0.1)
s t e p
4 || H || S n 1e
( L t +
L t ) / t r ( S n ) 1eslope= 2 2 ( 1 , 2 )(1 1 ) 2 ( 1 , 2 ) =
=0.01,( 1 , 2 )=(0.2,0)
s t e p
4 || H || S n 1e
( L t +
L t ) / t r ( S n ) 1eslope= 2 2 ( 1 , 2 )(1 1 ) 2 ( 1 , 2 ) =
=0.01,( 1 , 2 )=(0.2,0.2)
s t e p
4 || H || S n 1e
( L t +
L t ) / t r ( S n ) 1eslope= 2 2 ( 1 , 2 )(1 1 ) 2 ( 1 , 2 ) =
=0.01,( 1 , 2 )=(0.4,0)
s t e p
4 || H || S n 1e
( L t +
L t ) / t r ( S n ) 1eslope= 2 2 ( 1 , 2 )(1 1 ) 2 ( 1 , 2 ) =
=0.01,( 1 , 2 )=(0.4,0.4)
s t e p
[Figure]
D D ISCUSSION
We provide a new insight on the link between the batch gradient distribution and the sharpness of the loss landscape.
In this section, we reconcile our arguments with some previous studies.
Jastrz˛ebski et al.
(2017) explain the optimization behavior of SGD with the SDE approximation dθ t =
−∇ L ( θ t ) dt + (cid:112) ηb C 1 / 2 1 dW ( t ) of the SGD where W is an m -dimensional Brownian motion.
Therefore, the same ratio ηb = η ′ b ′ leads to the same SDE, which implies LSR.
Moreover, a large ηb implies a large diffusion in SDE, which has been linked with the escaping efficiency from a sharp local minimum in Zhu et al.
Our arguments are free from any other problems raised for the
SDE-based analyses of SGD which assume vanishing learning rates (Mandt et al., 2016; 2017; Hu et al., 2019; Li et al., 2017; 2019a; Jastrz˛ebski et al., 2017; Smith & Le, 2018; Chaudhari & Soatto, 2018), e.g., the mismatch to practical finite learning rate regime or the inherent theoretical issues in the SDE approximations (Yaida, 2019; Li et al., 2021).
We instead argue that a large second moment tr( S b ) (compared to tr( S n ) ) and a large η lead to a low constraint 2 ρ b /η on the interaction-aware sharpness.
Wu et al.
(2020) empirically show that what is important for the generalization performance of a neural network is not the class to which the gradient distribution belongs, but the second moment of the distribution.
This is consistent with our arguments with the interaction tr( HS b ) and the concentration measure ρ b = tr( S n ) / tr( S b ) , because they depend on the second moment S b , not on the class of the gradient distribution.
Recently, Li et al.
(2021) suggest a necessary condition that the “noise-to-signal ratio” needs to be large for LSR (and the SDE assumption) to hold.
This is consistent with our result on the linear regime (where b and ρ b are small) because the noise-to-signal ratio is approximately the inverse of the “signal-to-noise” ratio ρ b = tr( S n ) / tr( S b ) , but defined for an equilibrium distribution.
We provide not only the necessary condition but also the sufficient condition for LSR with a novel scaling rule
LSSR applicable to every batch size including where LSR fails (the saturation regime). [Figure]
=
[Figure] = [Figure]
[Figure]
b =
[Figure]
Step 0.
[Figure]