H OW D OES S HARPNESS -A WARE M INIMIZATION M INIMIZE S HARPNESS ?
Anonymous authors
Paper under double-blind review
A BSTRACT
Sharpness-Aware Minimization (SAM) is a highly effective regularization technique for improving the generalization of deep neural networks for various settings.
However, the underlying working of SAM remains elusive because of various intriguing approximations in the theoretical characterizations.
SAM intends to penalize a notion of sharpness of the model but implements a computationally efficient variant; moreover, a third notion of sharpness was used for proving generalization guarantees.
The subtle differences in these notions of sharpness can indeed lead to significantly different empirical results.
This paper rigorously nails down the exact sharpness notion that SAM regularizes and clarifies the underlying mechanism.
We also show that the two steps of approximations in the original motivation of SAM individually lead to inaccurate local conclusions, but their combination accidentally reveals the correct effect, when full-batch gradients are applied.
Furthermore, we also prove that the stochastic version of SAM in fact regularizes the third notion of sharpness mentioned above, which is most likely to be the preferred notion for practical performance.
The key mechanism behind this intriguing phenomenon is the alignment between the gradient and the top eigenvector of Hessian when SAM is applied.
I NTRODUCTION
Modern deep nets are often overparametrized and have the capacity to fit even randomly labeled data (Zhang et al., 2016).
Thus, a small training loss does not necessarily imply good generalization.
Yet, standard gradient-based training algorithms such as SGD are able to find generalizable models.
Recent empirical and theoretical studies suggest that generalization is well-correlated with the sharpness of the loss landscape at the learned parameter (Keskar et al., 2016;
Dinh et al., 2017;
Dziugaite & Roy, 2017; Neyshabur et al., 2017; Jiang et al., 2019).
Partly motivated by these studies, Foret et al.
; Wu et al.
; Zheng et al.
(2021) propose to penalize the sharpness of the landscape to improve the generalization.
We refer this method to Sharpness-Aware Minimization (SAM) and focus on the version of Foret et al.
(2021) in this paper.
Despite its empirical success, the underlying working of SAM remains elusive because of the various intriguing approximations made in its derivation and analysis.
There are three different notions of sharpness involved – SAM intends to optimize the first notion, the sharpness along the worst direction but actually implements a computationally efficient notion, the sharpness along the direction of the gradient.
But in the analysis, a third notion of sharpness is actually used to prove generalization guarantees, which admits the first notion as an upper bound.
The subtle difference between the three notions can lead to very different explicit biases.
(see Figure 1 for demonstraion)
More concretely, let L be the training loss, x be the parameter and ρ be the perturbation radius , a hyperparameter requiring tuning.
The first notion corresponds to the following optimization problem (1), where we call R Max ρ ( x ) =
L Max ρ ( x ) − L ( x ) the worst-direction sharpness at x and thus
SAM is intended to minimize the original training loss plus the worst-direction sharpness at x .
[Equation]
However, even evaluation of L Max ρ ( x ) is computationally expensive, not to mention optimization.
Thus Foret et al.
; Zheng et al.
(2021) propose to approximate the worst perturbation direction by the direction of the gradient and implement the second notion of sharpness, which corresponds
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
to (2).
We call R Asc ρ ( x ) =
L Asc ρ ( x ) − L ( x ) the ascent-direction sharpness at x . min
[Equation]
Asc
[Equation]
( x ) , where L Asc ρ ( x ) =
L (cid:18) x
+ ρ ∇ L ( x ) ∥∇ L ( x ) ∥(cid:19) .
Intriguingly, in the generalization analysis of SAM by Foret et al.
; Wu et al.
(2020), the first notion of sharpness, i.e. , the worst-direction sharpness, is only used for upper bounding the third notion of sharpness via the PAC Bayesian theory (McAllester, 2003).
We call the third notion R Avg ρ ( x ) =
L Avg ρ ( x ) − L ( x ) the average-direction sharpness at x , where L Avg ρ ( x ) =
E g ∼ N (0 ,I ) L ( x + ρg/ ∥ g ∥ ) .
For further acceleration, Foret et al.
; Zheng et al.
(2021) omit the gradient through other occurrence of x and approximate the gradient of ascent-direction sharpness by gradient taken after one-step ascent, i.e. , ∇ L Asc ρ ( x ) ≈ ∇ L (cid:16) x
+ ρ ∇ L ( x ) ∥∇ L ( x ) ∥ 2 (cid:17) and derive the update rule of SAM, where η is the learning rate.
[Equation]
In this paper, we analyze the explicit bias of various notions of sharpness and the optimization trajectory of SAM.
Our analysis is performed for small perturbation radius ρ and learning rate η under the setting where the minimizers of loss form a manifold following Fehrman et al.
; Li et al.
In particular, we make the following theoretical contributions.
We prove that full-batch SAM does minimize worst-direction sharpness.
(Theorem 4.9) 2.
Surprisingly, when batch size is 1, SAM minimizes average-direction sharpness.
(Theorem 4.9) 3.
We characterize the explicit biases of three notions of sharpness among minimizers when perturbation radius ρ goes to zero.
(Theorems 4.5 to 4.7, also see Table 1)
Surprisingly, both heuristic approximations made for the update rule of SAM lead to inaccurate solutions, that is, (1) minimizing worst-direction sharpness and ascent-direction sharpness induce different biases among minimizers, and (2) SAM doesn’t minimize ascent-direction sharpness.
The key mechanism behind this implicit bias of SAM is the implicit alignment between the gradient and the top eigenvector of Hessian when running SAM.
R ELATED W ORKS
Sharpness and Generalization.
The study on the connection between sharpness and generalization can be traced back to Hochreiter & Schmidhuber (1997).
Keskar et al.
(2016) observed a positive correlation between the batch size, the generalization error, and the sharpness of the loss landscape when changing the batch size.
Jastrzebski et al.
(2017) extends this by finding a correlation of the sharpness and the ratio between learning rate to batch size.
Dinh et al.
(2017) show that one can easily construct networks with good generalization but with arbitrary large sharpness by reparametrization.
Dziugaite & Roy (2017); Neyshabur et al.
Wei & Ma (2019) give theoretical guarantees on the generalization error using sharpness-related measures.
Jiang et al.
(2019) perform a largescale empirical study on various generalization measures and showed that sharpness-based measures have the highest correlation with generalization.
Background on Sharpness Aware-Minimization.
Foret et al.
; Zheng et al.
(2021) concurrently proposed to minimize the loss at the perturbed from current parameter towards the worst direction to improve generalization.
Wu et al.
(2020) propose the almost identical method for a different purpose, robust generalization of adversarial training.
Kwon et al.
(2021) proposed a different metric for SAM to fix the rescaling problem pointed out by Dinh et al.
Liu et al.
(2022) propose an more computationally efficient version of SAM.
Zhao et al.
(2022) proposed to improve generalization by penalizing gradient norm.
Their proposed algorithm can be viewed as a
[Figure]
generalization of SAM.
Andriushchenko & Flammarion (2022) studied a variant of SAM where the step size of ascent step is ρ instead of ρ ∥∇ L ( x ) ∥ 2 .
They showed that for a simple model this variant of
SAM has a stronger regularization effect when batch size is 1 compared to the full-batch case and argued that this might be the explanation that SAM generalizes better with small batch sizes.
N OTATIONS AND A SSUMPTIONS
For any integer k , we define C k as the set of k times continuously differentiable functions.
For any mapping F , we define ∂F ( x )[ u ] and ∂ 2 F ( x )[ u, v ] as the first and second order directional derivative of F ( x ) along the direction u (and v ).
Given a differential submanifold Γ of R D and a point x ∈ Γ , define P x, Γ as the projection operator onto the manifold of the normal space of Γ at x and P ⊥ x, Γ = I D − P x, Γ .
We fix our initialization as x init and our loss function as L : R D
→
R .
Given the loss function, the gradient flow can be defined as mapping ϕ : R D × [0 , ∞ ) →
R D satisfying ϕ
( x, τ ) =
x − (cid:82) τ 0 ∇ L ( ϕ ( x, t ))
dt .
We further define the limiting map Φ : R D →
R D as Φ( x ) =
lim
τ →∞ ϕ
( x, τ ) .
For any positive definite symmetry matrix A ∈ R D × D , define { λ i ( A ) , v i ( A ) } i ∈
[ D ] as all its eigenvalues and eigenvectors satisfying λ 1 ( A ) ≥ λ 2 ( A ) ... ≥ λ 3 ( A ) and ∥ v i ( A ) ∥ 2 = 1 .
Our analysis assumes sufficiently small η and ρ and uses O ( · ) to hide constant.
Following Fehrman et al.
; Li et al.
(2021); Arora et al.
, we make the below assumption.
Assumption 3.1.
Assume loss L : R D
→ R belongs to C 4 , and there exists a manifold Γ that is
D − M dimensional C 2 − submanifold of R D for some integer 1 ≤ M ≤ D , where for all x ∈ Γ , x is a global minimizer of L , L ( x ) = 0 and rank ( ∇ 2 L ( x ))
=
M .
The smoothness assumption is meet with networks with smooth activation functions and the existence of the manifold is due to the vast overparameterization of the modern neural network.
The full rank assumption is necessary for the analysis to guarantee the differentiability of Φ .
Definition 3.2 (Attraction Set) .
Let U be the attraction set of Γ under gradient flow, that is, a neighborhood of Γ containing all points starting from which gradient flow w.r.t.
loss L converges to some point in Γ , or mathematically, U ≜ { x ∈ R D | Φ( x ) exists and Φ( x ) ∈ Γ } .
Assumption 3.1 implies that U is open and Φ is in C 3 on U (from Lemma B.15 Arora et al.
E XPLICIT AND I MPLICIT B IAS IN THE F ULL - BATCH S ETTING
Section 4.1 provides a general theorem to properly analyze the explicit bias of various notions of sharpness among different minimizers.
We then apply our machinery on ascent-direction sharpness and worst-direction sharpness and show that they have different explicit biases.
In Section 4.
we provide our main theorem in the full-batch setting, that SAM implicitly minimizes the worstdirection sharpness, via characterizing its limiting dynamics as learning rate ρ and η goes to 0 with a Riemmanian gradient flow with respect to the top eigenvalue of the Hessian of the loss on the manifold of global minimizers.
In Section 4.3 we sketch the proof of the implicit bias of SAM and identified a key property behind the implicit bias, which is the implicit alignment between the gradient and the top eigenvector of the Hessian throughout the training.
4. W ORST - AND A SCENT - DIRECTION S HARPNESS HAVE D IFFERENT E XPLICIT B IAS
The intuition of approximating R Max ρ by R Asc ρ comes from the following Taylor expansions (Foret et al., 2021; Wu et al., 2020), which works for small ρ :
[Equation]
Here, the preference among the global minima is what we are mainly concerned with.
Since sup ∥ v ∥ 2 ≤ 1 v ⊤ ∇ L ( x ) =
∥∇ L ( x ) ∥ 2 when ∥∇ L ( x ) ∥ 2 > 0 , the leading terms in Equations (4) and (5) are both the first order term, ρ ∥∇ L ( x ) ∥ 2 , and are the same.
However, it is erroneous to think that the first order term decides the explicit bias, as the first order term ∥∇ L ( x ) ∥ 2 vanishes at the global minimizers of the loss L and thus the second order term becomes the leading term.
Any global minimizer x of the original loss L is an O ( ρ 2 ) -approximate minimizer of the sharpness-aware loss because ∇ L ( x ) = 0 .
Therefore, the sharpness-aware loss needs to be of order
ρ 2 so that we can guarantee the second-order terms in Equation (4) and/or Equation (5) to be non-trivially small.
In this section, we aim to understand under Assumption 3.1 what the explicit biases of various notions of sharpness among different minimizers are.
Theorem 4.3 will be our main theoretic tool to analyze the explicit bias for small perturbation radius ρ .
Below we first introduce some notations and assumption.
Notation for Regularizers.
Let R ρ : R D → R ∪{∞} be a family of regularizers parameterized by ρ .
If R ρ is not well-defined at some x , then we let R ρ ( x ) = ∞ .
This convention will be useful when analyzing ascent-direction sharpness R Asc ρ =
L Asc ρ − L which is not defined when ∇ L ( x ) = 0 .
This convention will not change the minimizers of the regularized loss.
Intuitively, a regularizer should always be non-negative, but however, when far away from manifold, regularizer R ρ ( x ) can actually be negative, e.g. , R Avg ρ ( x ) ≈
ρ 2 2 D Tr[ ∇ 2 L ( x )] .
Therefore we make the following assumption to allow the regularizer to be mildly negative.
Assumption 4.1.
Suppose for any bounded closed set B ⊂ U , there exists C > 0 , such that for sufficiently small ρ , ∀ x ∈ B, R ρ ( x ) ≥ − Cρ 2 .
The following concept of limiting regularizer is of crucial role in our analysis.
Definition 4.2 (Limiting Regularizer) .
We define the limiting regularizer of { R ρ } as the function
[Equation] [Equation]
We say the limiting regularizer S of { R ρ } is a good around some x ∗ ∈ Γ , if there is an open set V containing x ∗ , such that S is a non-negative continuous function in V and for any ϵ > 0 , there is some ρ x ∗ > 0 , it holds that ∀ x ∈ Γ ∩ V, ρ ≤ ρ x ∗ , (cid:12)(cid:12) S ( x ) − inf ∥ x ′ − x ∥ 2 ≤ ϵ · ρ R ρ ( x ′ ) /ρ 2 (cid:12)(cid:12) <
ϵ.
We say the limiting regularizer S is good on Γ , if S is good around every point x ∈ Γ .
The high-level intuition behind the definition of limiting regularizer is that we want to use it to capture the second order term in the Taylor expansion of regularizer R ρ when ρ → 0 .
When the second order term is continuous in x , the definition of S ( x ) can also be simplified as R ρ ( x ) /ρ 2 .
The intuition of the concept of a good limiting regularizer is that, the regularizer should not change very fast, especially in an O ( ρ ) neighborhood of the minimizer.
If so, the minimizer of the regularized loss may be Ω( ρ ) away from any minimizer to reduce the regularizer at the cost of increasing the original loss, which makes the limiting regularizer unable to capture the explicit bias of the regularizer.
Theorem 4.3.
Let U ′ be any bounded open set such that its closure U ′ is contained in U and that
U ′ ∩ Γ = U ′ ∩ Γ .
Then for any family of parametrized regularizers { R ρ } admitting a good limiting regularizer on Γ and satisfying Assumption 4.1 and any ϵ ≥ 0 , there is a ρ 0 > 0 , such that for all u ∈ U ′ and ρ < ρ 0 , it holds that
[Equation]
For the applications we are interested in in this paper, the good limiting regularizer S can be continuously extended to the entire space R D .
In such a case, the implication of “ = ⇒ ” of Theorem 4.3 also admits the following alternative form which doesn’t involve R ρ .
Corollary 4.4 implies minimizing regularized loss L ( x ) + R ρ ( x ) is equivalent to minimizing the limiting regularizer of { R ρ } ρ , S ( x ) on the global minimizer manifold Γ .
Corollary 4.4.
Under the setting of Theorem 4.3, let S be an continuous extension of S to R d , if
L ( u ) + R ρ ( u ) ≤ inf x ∈ U ′
( L ( x ) + R ρ ( x ))
+ ϵρ
+ o ( ρ 2 ) , then we have that L ( u ) − inf x ∈ U ′
L ( x ) =
O ( ρ 2 ) and that (cid:12)(cid:12) S ( u ) − inf x ∈ U ′ ∩ Γ S ( x ) (cid:12)(cid:12) =
ϵ
+ o (1) .
Corollary 4.4 suggests a sharp phase transition of the property of the solution of min x L ( x )+ R ρ ( x ) when the optimization error drops from ω ( ρ 2 ) to O ( ρ 2 ) .
when the optimization error is larger than ω ( ρ 2 ) , no regularization effect happens and any minimizer satisfies the requirement.
When the error becomes O ( ρ 2 ) , there is a non-trivial restriction on the (extended) limiting regularizer.
Below we show all of worst-, ascent- and average-direction sharpness admit good limiting regularizers.
Thus we can apply Corollary 4.4 to characterize their explicit biases, which are all different.
Theorem 4.5.
Worst-direction sharpness R Max ρ satisfies Assumption 4.1 and admits λ 1 ( ∇ 2 L ( · )) /as
a good limiting regularizer on Γ .
Theorem 4.6.
Ascent-direction sharpness R Asc ρ satisfies Assumption 4.1 and admits λ M ( ∇ 2 L ( · )) /as
a good limiting regularizer on Γ .
Theorem 4.7.
Average-direction sharpness R Avg ρ satisfies Assumption 4.1 and admits
Tr( ∇ 2 L ( · )) / (2 D ) as a good limiting regularizer on Γ .
When R ρ is continuous at some x ∈ Γ , the definition of S ( x ) can be simplified as lim ρ → 0 R ρ ( x ) /ρ 2 .
Worst- and average- direction sharpness fall into this type and the limiting regularizer can be solved straightforwardly by Taylor expansion, which is well-known in literature (Keskar et al., 2016; Kaur et al., 2022; Zhuang et al., 2022; Orvieto et al., 2022).
Theorems 4.5 and 4.7 can be viewed as a restatement of the existing results in our manifold setting.
Our major contribution in this subsection is the analysis for ascent-direction sharpness, which is more tricky as R Asc ρ ( x ) = ∞ and thus is not continuous for any x ∈ Γ .
To minimize R Asc ρ around x , we can pick x ′ → x to make ∥∇ L ( x ) ∥ 2 → 0 but not equal to 0.
By (5), we have
R Asc ρ ( x ′ )
≈ ρ 2 / 2 · ∇ L ( x ′ )
⊤ ∇ 2 L ( x ) ∇ L ( x ′ )
/ ∥∇ L ( x ′ )
∥ 2 .
Here the crucial step of the proof is that because of Assumption 3.1, ∇ L ( x ) / ∥∇ L ( x ) ∥ 2 must almost lie in the column span of ∇ 2 L ( x ) , and thus implies inf x ′ ∇ L ( x ′ )
⊤ ∇ 2 L ( x ) ∇ L ( x ′ )
/ ∥∇ L ( x ′ )
∥ 2 ρ → 0 → λ M ( ∇ 2 L ( x )) .
The above alignment property between the gradient and the column space of Hessian can be checked directly for any nonnegative quadratic function and the maximal Hessian rank assumption in Assumption 3.1 ensures this property extends to general losses.
4. SAM P ROVABLY D ECREASES W ORST - DIRECTION S HARPNESS L OCALLY
Though ascent-direction sharpness has different explicit bias from worst-direction sharpness, in this subsection we will show that surprisingly, SAM, an heuristic method designed to minimize ascentdirection sharpness, provably decreases worst-direction sharpness.
The main result here is an exact characterization of the trajectory of SAM (3) via the following ordinary differential equation (ODE) (6), when learning rate η and perturbation radius ρ are small and the initialization x (0)
=
x init is in U .
We call the solution of (6) the limiting flow of SAM, which is exactly the Riemannian Gradient Flow on the manifold Γ with respect to λ 1 ( ∇ 2 L ( · )) .
In other words, the ODE (6) is essentially
[Equation]
Note λ 1 ( ∇ 2 L ( x )) may not be differentiable at x if λ 1 ( ∇ 2 L ( x ))
=
λ 2 ( ∇ 2 L ( x )) , thus to ensure the (6) is well-defined, we assume there is a positive eigengap for L on Γ .
Assuming ODE (6) has a solution till time T 3 , we have Theorem 4.9, which is the main theorem of this section.
Assumption 4.8.
For x ∈ Γ , there exists a positive eigengap,i.e., λ 1 ( ∇ 2 L ( x )) > λ 2 ( ∇ 2 L ( x )) .
Theorem 4.9 (Main, Theorems C.31 and C.32 stated informally) .
Let { x ( t ) } be the iterates of
SAM (3) with x (0) =
x init ∈ U , then under Assumptions 3.1 and 4.8, for all η, ρ such that η ln(1 /ρ )
and ρ/η are sufficiently small, the dynamics of SAM can be split into two phases:
• Phase I: SAM follows Gradient Flow with respect to L until entering an O ( ηρ ) neighborhood of the manifold Γ in ˜ O
( 1 η ) steps; • Phase II: SAM tracks the solution X of (6) , the Riemannian Gradient Flow with respect to λ 1 ( ∇ 2 L ( · )) in the O ( ηρ ) neighborhood in the sense
that max 0 ≤ T ≤ T 3 ∥
Φ( x ( ⌈ T/ ( ηρ 2 ) ⌉ ))
− X ( T ) ∥
=
O (( η + ρ ) log(1 /ηρ )) .
Moreover, the angle between ∇ L ( x ( t )) and v 1 ( ∇ 2 L ( x ( t )))
is at most O ( ρ ) .
Theorem 4.9 shows that SAM decreases the largest eigenvalue of Hessian of loss locally around the manifold of global minimizers.
As a corollary of Theorem 4.9, we can also show that the largest eigenvalue of the limiting flow closely tracks the regularized training loss.
Corollary 4.10.
For all T ′ 3 > 0 , for all ρ, η such that η ln(1 /ρ ) and ρ/η are sufficiently small we have ∀ T ′ 3 < ηρ 2 t ≤ T 3 , ∥ R Max ρ ( x ( t ))
− ρ 2 λ 1 ( ∇ 2 L ( X ( ηρ 2 t )) / 2 ∥ =
˜ O ( ηρ 2 )
Recall Theorem 4.6 shows that the largest eigenvalue of Hessian is the limiting regularizer of the worst-direction sharpness , leveraging the equivalence relationship in Theorem 4.3, below we show that full-batch SAM provably minimizes worst-direction sharpness if we additionally assume the limiting flow converges to a local minimizer of the top eigenvalue of Hessian.
Corollary 4.11.
Define U ′ as in Theorem 4.3, suppose X ( ∞ ) =
lim t →∞ X ( t ) exists and is a minimizer of λ 1 ( ∇ 2 L ( x )) in U ′ ∩ Γ , for all ϵ > 0 , there exists a constant T > 0 , then for all ρ, η such that η ln(1 /ρ ) and ρ/η are sufficiently small we have L Max ρ ( x ( ⌈ T/ ( ηρ 2 ) ⌉ ))
≤ ϵρ 2 + inf x ∈ U ′
L Max ρ ( x )
4. P ROOF S KETCH FOR P HASE II IN T HEOREM 4.
Now we sketch the proof for the ODE-based characterization of the trajectory of SAM in Phase II.
The framework of the analysis is similar to Arora et al.
; Lyu et al. (2022), where the high-level idea is to use Φ( x ( t )) as a proxy for x ( t ) and study the dynamics of Φ( x ( t )) via Taylor expansion, which turns out to be dependent on the alignment between the gradient and the eigenvectors of the
Hessian.
In particular, like Arora et al.
; Lyu et al. (2022), we show that the gradient aligns to the top eigenvector of Hessian, and thus encourages the SAM dynamics to reduce the top eigenvalue of the Hessian.
Taylor Expansion on Φ .
In Phase II, it can be shown that ∥ x ( t ) − Φ( x ( t )) ∥
= O ( ηρ ) holds for every step, this implies ∥ x ( t + 1) − x ( t ) ∥ 2 =
O ( ρη ) .
(See Lemma C.12)
Therefore we have that
[Equation]
Now we apply Taylor expansion on ∇ L (cid:16) x
+ ρ ∇ L ( x ) ∥∇ L ( x ) ∥ 2 (cid:17) around x and get
[Equation]
Lemma 4.12 (Lemma B.16 Arora et al.
For x
[Equation]
x
[Equation]
− ∂ 2 Φ( x )[ ∇ L ( x ) , ∇ L ( x )] .
Now we plug (8) into (7) and simplify the expression using Lemma 4.12.
We have that
[Equation]
where in the last step we use the property that ∥ x ( t ) − Φ( x ( t ))
∥
=
O ( ηρ ) , which further implies ∥∇ L ( x ( t ))
∥ 2 =
O ( ηρ ) in Phase II (See Lemma C.7).
Lemma 4.13.
For x ∈ Γ , ∂ Φ( x ) =
P ⊥ x, Γ , ∂ Φ( x ) ∇ 2 L ( x ) = 0 , where P ⊥ x, Γ is the orthogonal projection matrix of the tangent space of Γ at x .
Next we use Lemma 4.13 to further simplify (9).
Since the entire phase II happens in an
O ( ηρ ) -neighborhood of manifold Γ , we have that ∥ x ( t ) − Φ( x ( t ))
∥ 2 = O ( ηρ ) , thus both (cid:13)(cid:13) ∂
Φ( x ( t )) ∇ 2 L ( x ( t ))
(cid:13)(cid:13) and ∥∇ L ( x ( t ))
∥ 2 are O ( ηρ ) .
So far, the proof is almost completed, with the implicit alignment between gradient and top eigenvector of Hessian being the last missing piece.
Suppose we have (cid:13) (cid:13)(cid:13) ∇ L ( x ) ∥∇ L ( x ) ∥ 2 − v 1 ( ∇ 2 L ( x ))
(cid:13)(cid:13)(cid:13) =
O ( ρ ) , then it holds that
[Equation]
x (
[Equation]
+ 1)
[Equation]
where the second to last step we use the following property about the derivative of eigenvalue (Lemma C.69) and the last step is due to Taylor expansion.
Implicit Hessian-gradient Alignment.
It remains to explain why the gradient implicitly aligns to the top eigenvector of Hessian, which is the key component of the analysis in Phase II.
The proof strategy here is to first show alignment for a quadratic loss function, and then generalize its proof to general loss functions satisfying Assumption 3.1.
Below we first give the formal statement of the implicit alignment on quadratic loss Theorem 4.14.
Note this alignment property is an implicit property of the SAM algorithm – it is not due to the fact that SAM is an approximation of GD on
L Asc ρ , because optimizing L Asc ρ would rather make the gradient align to the smallest eigenvector!
Theorem 4.14.
Suppose A is a positive definite symmetric matrix with unique top eigenvalue.
Consider running SAM on loss L ( x ) :
= 12 x T Ax as (11) , then for almost every x (0) , we have x ( t ) converges in direction to v 1 ( A ) and lim t →∞ ∥ x ( t ) ∥
= ηρλ 1 ( A ) 2 − ηλ 1 ( A ) with ηλ 1 ( A ) < 1 .
[Equation]
Equivalently, we can reformulate the update rule through the lens of the gradient ∇ L ( x ) =
Ax :
[Equation] [Equation]
and the result of Theorem 4.14 becomes the alignment between gradient and the top eigenvector of
Hessian.
This property can be generalized to the analysis of general loss as the dynamic of SAM near a fixed point on Γ is similar to the quadratic case with small perturbation.
To see this, we apply
Taylor expansion on the update rule of SAM (3):
[Equation]
Since phase II happens in an O ( ηρ ) -neighborhood of manifold Γ , we have ∥ x ( t + 1) − x ( t ) ∥ 2 =
O ( ηρ ) .
Then by (13) and Taylor expansion on ∇ L ( x ( t + 1)) at x ( t ) , we have that
[Equation]
Equation (14) is a O ( ηρ 2 ) -perturbed version of the update rule in the quadratic case.
Note this is a higher order term comparing to the other two terms, which have orders Θ( η 2 ρ ) and Θ( ηρ ) respectively, the implicit alignment between Hessian and gradient happens for the same reason as in the quadratic case.
We further show once this alignment happens, it will be kept until the end of our analysis, which is Θ( η − 1 ρ − 2 ) steps.
E XPLICIT AND I MPLICIT BIAS IN THE S TOCHASTIC S ETTING
In practice, people usually use SAM in the stochastic mini-batch setting, and the test accuracy improves as the batch size decreases (Foret et al., 2021).
Towards explaining this phenomenon, Foret et al.
(2021) argue intuitively that stochastic SAM minimizes stochastic worst-direction sharpness.
Given our results in Section 4, it’s natural to ask if we can justify the above intuition by showing the implicit Hessian-gradient alignment in the stochastic setting.
Unfortunately, such alignment is not possible in the most general setting.
For example, take a simple quadratic loss L ( x ) =
L 1 ( x )+ L 2 ( x ), where L k ( x ) = 0 .
5 x ⊤ A k x
and A k is a positive definite matrix for k = 1 , 2 .
If A 1 and A 2 have different top eigendirection, then no x can simultaneously satisfy that ∇ L k ( x ) =
A k x aligns to the top eigenvector of ∇ 2 L k ( x )
=
A k .
Yet when the batch size is 1, we can prove rigorously that stochastic SAM minimizes stochastic worst-direction sharpness (Section 5.2), which is the main result in this section.
Below we start by introducing our setting for SGD with batch size 1.
We still need Assumption 3.
in this section.
We first by analyzing the explicit bias of the stochastic ascent- and worst-direction sharpness in Section 5.1 via the tools developed in Section 4.1.
It turns out they are all proportional to the trace of hessian as ρ → 0 .
In Section 5.2, we show stochastic SAM locally decreases trace of
Hessian.
Setting Let f k ( x ) be the model output on k th data where f k is a C 4 -smooth function and y k be the i th label for l = 1 , . . .
, M .
We define the loss on k th data as L k ( x ) =
ℓ ( f k ( x ) , y k ) and the total loss L = (cid:80) Mk =1 L k /M , where ℓ ( y ′ , y ) is a C 4 -smooth function satisfying the following properties 1 : (1).
arg min y ′ ∈ R ℓ ( y ′ , y ) =
y , for any y ∈ R ; (2).
d 2 ℓ ( y ′ ,y )
d 2 y ′ | y ′
=
y > 0 , for any y ∈ R ; (3).
There exists x 0 , such that for any k ∈ [ M ] , it holds that f k ( x 0 ) =
y k .
5. S TOCHASTIC W ORST - AND A SCENT - DIRECTION S HARPNESS HAVE S AME E XPLICIT B IAS
Below we specify stochastic worst-direction sharpness and stochastic ascent-direction sharpness as E k
[ R Max k,ρ ] =
E k
[ L Max k,ρ ] − L and E k
[ R Asc k,ρ ] =
E k
[ L Asc k,ρ ] − L .
Here L max k,ρ and L asc k,ρ is defined as in Equations (1) and (2) with L replaced by L k .
Unlike the full-batch setting, these two sharpness have same explicit bias, or more precisely, they have the same limiting regularizers.
We omit the result on the stochastic average-direction sharpness as it is the same as its counterpart in the fullbatch case.
Theorem 5.1.
Stochastic worst-direction sharpness E k
[ R Max k,ρ ] satisfies Assumption 4.1 and admits
Tr( ∇ 2 L ( · )) / 2 as a good limiting regularizer on Γ .
Theorem 5.2.
Stochastic ascent-direction sharpness E k
[ R Asc k,ρ ] satisfies Assumption 4.1 and admits
Tr( ∇ 2 L ( · )) / 2 as a good limiting regularizer on Γ .
Unlike in the full-batch setting where the implicit regularizer of ascent-direction sharpness and worst-direction sharpness have different explicit bias, here they are the same because there is no difference between the maximum and minimum of its non-zero eigenvalue for rank-1 Hessian of each individual loss, and that the average of limiting regularizers is equal to the limiting regualrizer of the average regularizers by definition.
5. S TOCHASTIC SAM M INIMIZES S TOCHASTIC W ORST - DIRECTION S HARPNESS
Stochastic SAM: Recall L k is the loss on k th data, we use stochastic SAM to denote the following update rule, where k t is sampled i.i.d from uniform distribution on [ M ] .
[Equation]
The goal of this section is to show stochastic SAM tracks the following Riemannian gradient flow with respect to Tr( ∇ 2 L ( · )) on the manifold for sufficiently small η and ρ ,
[Equation]
The fundamental difference between stochastic SAM (16) with batch size 1 and that in the general setting is that the rank of the Hessian of each stochastic loss at minimizers is only rank-1, which enforces the gradient ∇ L k ( x ) ≈ ∇ 2 L k ( x )( x − Φ( x )) to (almost) lie in the direction of the top eigenvalue of the Hessian.
Lemma 5.3 formally states this property.
For convenience, we denote d 2 ℓ ( y [Equation] ,y k ) d 2 y ′ | y
[Equation]
= f k ( x [Equation] ∥∇ f [Equation]
∇ f k
[Equation]
x
[Equation]
/ ∥∇ f k ( x
[Equation]
∥ as Λ k ( x ) , w k ( x ) .
Lemma 5.3.
Under Assumption 3.1, for any x ∈
U and p ∈ Γ , ∇ 2 L k ( p ) =
Λ k ( p ) w k ( p ) w k ( p ) T and ∃ s ∈ { 1 , − 1 } ∇ L k ( x ) ∥∇ L k ( x ) ∥
=
sw k ( p ) + O ( ∥ x − p ∥ )
With Lemma 5.3, we can show Φ( x ( t + 1))
− Φ( x ( t ))
=
− ηρ 2 P ⊥ t, Γ ∇ (Λ k ( x )) / 2 + ˜ O
η 2 ρ 2 ) when ∥ x ( t + 1) − x ( t ) ∥
=
˜ O ( ηρ ) , via a similar but slightly more complicated analysis on Φ( x ( t + 1)) − Φ( x ( t )) as in Section 4.2.
As we have E [Λ k ( x )]
= ( (cid:80) k
Λ k ( x )) /M
=
Tr( ∇ 2 L ( x )) , x ∈ Γ , we have that locally SAM is essentially performing a stochastic projected gradient descent with respect to Tr( ∇ 2 L ( · )) .
It still remains to prove stochastic SAM gets ˜ O ( ηρ ) close of the manifold, which is addressed by Theorem C.50.
Put everything together, we have the following main theorem for this section.
Theorem 5.4.
Let { x ( t ) } be the iterates defined by SAM (16) and x (0) =
x init ∈ U , then under
Assumption 3.1, for all η and ρ such that ( η + ρ ) log(1 /ηρ ) is sufficiently small, with probability at least 1 − O ( √ ρ + √ η ) , the dynamics of SAM can be split into two phases:
• Phase I: Stochastic SAM follows Gradient Flow with respect to L until entering an O ( ηρ ) neighborhood of the manifold Γ in ˜ O
( 1 η ) steps; • Phase II: Stochastic SAM tracks the solution X of (17) , the Riemannian Gradient Flow with respect to Tr( ∇ 2 L ( · )) , with time scaling as ∥ Φ( x ( ⌈ T 3 / ( ηρ 2 ) ⌉ ))
− X ( T 3 ) ∥ =
O ( η 1 / 4 + ρ 1 / 4 ) .
Corollaries 5.5 and 5.6 below are stochastic counterparts of Corollaries 4.10 and 4.11, which says that the trace of Hessian of the limiting flow of well-tracks the stochastic worst-direction sharpness, and therefore when the limiting flow converges to a local minimizer of trace of Hessian, stochastic
SAM minimizes the stochastic worst-direction sharpness.
Corollary 5.5.
Under the condition of Theorem 5.4, we have that with probability 1 − O ( √ η + √ ρ ) , the difference between the stochastic worst-direction sharpness of the iterates and the corresponding scaled trace of Hessian along the limiting flow is at most O (cid:0) ( η 1 / 4 + ρ 1 / 4 ) ρ 2 (cid:1) , that is,
[Equation]
Corollary 5.6.
Define U ′ as in Theorem 4.3, suppose X ( ∞ ) =
lim t →∞ X ( t ) exists and is a minimizer of Tr( ∇ 2 L ( x ))) in U ′ ∩ Γ . Then for all ϵ > 0 , there exists a constant T ϵ > 0 , such that for all ρ, η such that ( η + ρ ) ln(1 /ηρ ) are sufficiently small, we have that with probability 1 − O ( √ η + √ ρ ) ,
[Equation]
C ONCLUSION
In this work, we have performed a rigorous mathematical analysis of the explicit bias of various notions of sharpness when used as regularizers and the implicit bias of the SAM algorithm.
In particular, we show the explicit biases of worst-, ascent- and average-direction sharpness around the manifold of minimizers are minimizing the largest eigenvalue, the smallest nonzero eigenvalue, and the trace of Hessian of the loss function.
We show that in the full-batch setting, SAM provably decreases the largest eigenvalue of Hessian, while in the stochastic setting when batch size is 1,
SAM provably decreases the trace of Hessian.
C ONTENTS 1 Introduction
2 Related Works 3 Notations and Assumptions 4 Explicit and Implicit Bias in the Full-batch Setting 4.
Worst- and Ascent-direction Sharpness have Different Explicit Bias . . . . . . . .
SAM Provably Decreases Worst-direction Sharpness Locally . .
Proof Sketch for Phase II in Theorem 4.9 . . . . . . . . . . . . . . . . . . . . . .
5 Explicit and Implicit bias in the Stochastic Setting 5.
Stochastic Worst- and Ascent-direction Sharpness have Same Explicit Bias . . . .
Stochastic SAM Minimizes Stochastic Worst-direction Sharpness . .
6 Conclusion
A Additional Related Works
B More Discussions of Manifold Assumption
C Proof
C.1 Proof Setups . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.2 Properties of Limiting Map of Gradient Flow, Φ . . . . . . . . . . . . . . . . . . .
C.3 Explicit Bias . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.4 Full-batch SAM on Quadratic Loss (Proof of Theorem 4.14) . . . . . . . . . . . .
C.4.
Entering Invariant Set . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.4.
Alignment to Top Eigenvector . . . . . . . . . . . . . . . . . . . . . . . .
C.5 Full-batch SAM on General Loss (Proof of Theorem 4.9) . . . . . . . . . . . . . .
C.5.
Phase I
Proof of Theorem C.31) . . . . . . . . . . . . . . . . . . . . . . .
C.5.
Phase II (Proof of Theorem C.32) . . . . . . . . . . . . . . . . . . . . . .
C.5.
Proofs of Corollaries 4.10 and 4.11 . . . . . . . . . . . . . . . . . . . . .
C.
[Table]
Analysis for 1-SAM (Proof of Theorem 5.4) . . . . . . . . . . . . . . . . . . . . .
C.6.
Phase I
Proof of Theorem C.50) . . . . . . . . . . . . . . . . . . . . . . .
C [Table]
Phase II (Proof of Theorem C.51) . . . . . . . . . . . . . . . . . . . . . .
C.6.
Proofs of Corollaries 5.5 and 5.6 . . . . . . . . . . . . . . . . . . . . . . .
C.6.
Other Omitted Proofs for 1-SAM . . . . . . . . . . . . . . . . . . . . . .
C.7 Technical Lemmas . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D Omitted Proofs on Continuous Approximation
A A DDITIONAL R ELATED W ORKS
Sharpness Minimization as Implicit Bias.
Recent theoretical works Blanc et al.
; Damian et al.
; Li et al.
(2021) showed that SGD with label noise is implicitly biased to local minimizers with a smaller trace of Hessian.
Arora et al.
(2022) showed that normalized GD implicitly decreases the largest eigenvalue of the Hessian.
Lyu et al.
(2022) showed that GD with weight decay on a scale invariant loss function implicitly decreases the spherical sharpness, i.e. , the largest eigenvalue of the Hessian evaluated at the normalized parameter.
In a concurrent work, Bartlett et al.
(2022) prove that on quadratic loss, the iterate of SAM (Equation (11)) and its gradient converges to the top eigenvector of Hessian, which is almost the same as our Theorem 4.14.
Assuming such alignment for a general loss, the work of Bartlett et al.
(2022) shows that the largest eigenvalue of Hessian decreases in the next step.
This paper also proves such a Hessian-gradient alignment for general loss functions and an end-to-end theorem showing that the largest eigenvalue of Hessian and worst-direction sharpness decrease along the trajectory of SAM (Theorem 4.9), which are not shown in Bartlett et al.
Moreover, this paper also characterizes implicit bias of stochastic SAM with batch size 1 , which is minimizing the average-direction sharpness, while Bartlett et al.
(2022) only considers the deterministic case.
Comparison with Arora et al. (2022) .
Our proof uses similar framework as Arora et al.
However our analysis has its own difficulty for the following reasons.
First, Arora et al.
(2022) only deal with deterministic case, while our analysis extends to stochastic SAM as well (Section 5).
Second, our analysis for the deterministic case is different from that of Arora et al.
(2022) in the following two aspects:
The alignment analysis is more complicated because we have two hyperparameters,learning rate η and perturbation radius ρ , while Arora et al.
(2022) only needs to deal with one hyperparameter, learning rate η . 2.
The mechanism of penalizing worst-direction sharpness is different, which can be seen from the dependency of the sharpness-reduction rate over learning rate η .
In Arora et al.
(2022), normalized GD reduces the sharpness via a second-order effect of GD and thus the sharpness is reduced by O ( η 2 ) per step.
In our analysis, for fixed small perturbation radius ρ , the sharpness is reduced by O ( ρ 2 η ) per step, which is linear in η .
B M ORE D ISCUSSIONS OF M ANIFOLD A SSUMPTION
Validity.
Recent empirical studies have shown that there are essentially no barriers in loss landscape between different minimizers, that is, the set of minimizers are path-connected (Draxler et al., 2018; Garipov et al., 2018).
Motivated by this empirical discovery, we make the assumption below following Fehrman et al.
; Damian et al.
; Li et al.
(2021), which is theoretically justified by Cooper (2018) under a generic setting.
Importance.
The connectivity of the set of local minimizers implied by the manifold assumption above allows us to take limits of perturbation radius ρ → 0 while still yield interesting and insightful implicit bias results.
So far almost all analysis of implicit bias for general model parameterizations relies on Taylor expansion, e.g. Blanc et al.
; Damian et al.
; Li et al.
(2021); Arora et al.
(2022), so does the derivation of the SAM algorithm (Foret et al., 2020; Wu et al., 2020).
Thus it’s crucial to consider small perturbation size ρ .
On the contrary, if the set of global minimizers are a set of discrete points, then with small perturbation radius ρ , implicit bias of optimizers is not sufficient to drive the iterate from global minimum to the other one.
Limitation.
As proved by Cooper (2018), the manifold setting holds generically for regression problems with smooth loss.
The manifold assumption doesn’t hold for networks with ReLU (which is non-smooth) and classification problems with softmax loss (where an overparametrized network may achieve arbitrarily small but non-zero loss).
In the latter case, there may not even exist a global minimizer, not to mention a manifold of minimizers.
C P ROOF C. P ROOF S ETUPS
In this section we provide details of our proof setups, including notations and assumptions/settings.
We first introduce some additional notations that will be used in the proofs.
For any subset S ∈ R D , we define Dist( x, S ) ≜ inf y ∈ S
∥ x − y ∥ 2 .
For any d > 0 and any subset S ∈ R D , we define
S d ≜ { x ∈ R D | Dist( x, S ) ≤ d } .
Our convention is to use K to denote a compact set and U to denote an open set.
Below we restate our main assumption in the full-batch case and related notations in Section 3.
Throughout the analysis, we fix our initialization as x init , our loss function as L : R D
→ R .
Assumption 3.1.
Assume loss L : R D
→ R belongs to C 4 , and there exists a manifold Γ that is
D − M dimensional C 2 − submanifold of R D for some integer 1 ≤ M ≤ D , where for all x ∈ Γ , x is a global minimizer of L , L ( x ) = 0 and rank ( ∇ 2 L ( x ))
=
M .
Notations for Full-Batch Setting: Given any point x ∈ Γ , define P x, Γ as the projection operator onto the manifold of the normal space of Γ at x and P ⊥ x, Γ = I D − P x, Γ .
Given the loss function L , its gradient flow is denoted by mapping ϕ : R D × [0 , ∞ ) → R D .
Here, ϕ ( x, τ ) denotes the iterate at time τ of a gradient flow starting at x and is defined as the unique solution of ϕ ( x, τ )
= x − (cid:82) τ 0 ∇ L ( ϕ ( x, t )) dt , ∀ x ∈ R D .
We further define the limiting map of ϕ ( x, · ) as Φ( x ) =
lim
τ →∞ ϕ
( x, τ ) , that is, Φ( x ) denotes the convergent point of the gradient flow starting from x .
For convenience, we define λ
i ( x ) , v i ( x ) as λ i ( ∇ 2 L (Φ( x ))) , v i ( ∇ 2 L (Φ( x )))
whenever the latter is well defined.
When x ( t ) and Γ is clear from context, we also use λ
i ( t ) :
= λ i ( x ( t )) , v i ( t ) :
= v i ( x ( t )) , P ⊥ t, Γ := P ⊥ Φ( x ( t )) , Γ , P t, Γ := P Φ( x ( t )) , Γ .
Definition 3.2 (Attraction Set) .
Let U be the attraction set of Γ under gradient flow, that is, a neighborhood of Γ containing all points starting from which gradient flow w.r.t.
loss L converges to some point in Γ , or mathematically, U ≜ { x ∈ R D | Φ( x ) exists and Φ( x ) ∈ Γ } .
Below we restate the setting for stochastic loss of batch size one in Section 5.
Setting C.1.
Let f k ( x ) be the model output on k th data where f k is a C 4 -smooth function and y k be the i th label for l = 1 , . . .
, M .
We define the loss on k th data as L k ( x ) =
ℓ ( f k ( x ) , y k ) and the total loss L = (cid:80) Mk =1 L k /M , where ℓ ( y ′ , y ) is a C 4 -smooth function satisfying the following properties 2 : (1).
arg min y ′ ∈ R ℓ ( y ′ , y ) =
y , for any y ∈ R ; (2).
d 2 ℓ ( y ′ ,y )
d 2 y ′ | y ′
=
y > 0 , for any y ∈ R ; (3).
There exists x 0 , such that for any k ∈ [ M ] , it holds that f k ( x 0 ) =
y k .
In our analysis, we prove our main theorems in the stochastic setting under a more general condition than Setting C.1, which is Condition C.2 (on top of Assumption 3.1).
The only usage of Setting C. in the proof is Lemma 5.3 and theorem C.49.
Condition C.2.
Total loss L = 1 M (cid:80) Mk =1 L k .
For each k ∈ [ M ] , L k is C 4 , and there exists a ( D − 1) -dimensional
C 2 -submanifold of R D , Γ k , where for all x ∈ Γ k , x is a global minimizer of
L k , L k ( x ) = 0 and rank ( ∇ 2 L k ( x ))
= 1 .
Moreover, Γ ⊂ ∩ Mk =1 Γ k for Γ defined in Assumption 3.1.
Notations for Stochastic Setting: Since L k is rank- 1 on Γ k for each k ∈ [ M ] , we can write it as
L k ( x ) =
Λ k ( x ) w k ( x ) w ⊤ k ( x ) for any x ∈ Γ , where w k is a continuous function on Γ with pointwise unit norm.
Given the loss function L k , its gradient flow is denoted by mapping ϕ k : R D × [0 , ∞ ) → R D .
Here, ϕ k ( x, τ ) denotes the iterate at time τ of a gradient flow starting at x and is defined as the unique solution of ϕ k
( x, τ ) =
x − (cid:82) τ 0 ∇ L k ( ϕ k ( x, t )) dt , ∀ x ∈ R D .
We further define the limiting map Φ k as Φ k ( x ) =
lim
τ →∞ ϕ k ( x, τ ) , that is, Φ k ( x ) denotes the convergent point of the gradient flow starting from x .
Similar to Definition 3.2, we define U k
= { x ∈ R D | Φ( x ) exists and Φ k ( x ) ∈
Γ k } be the attraction set of Γ i .
We have that each U k is open and Φ k is C 2 on U k by Lemma B. in Arora et al.
Definition C.3.
A function L is µ -PL in a set U iff ∀ x ∈ U , ∥∇ L ( x ) ∥ 22 ≥ 2 µ ( L ( x ) − inf x ∈ U L ( x )) .
Definition C.4.
The spectral 2-norm of a k -order tensor X i 1 ,...,i k ∈ R d 1 × ...
× d k is defined as
[Equation]
Lemma C.5 (Arora et al.
Lemma B.2) .
Given any compact set K ⊆ Γ , there exist r ( K ) , µ ( K ) , ∆( K ) ∈ R + such that
K r ( K ) ∩ Γ is compact.
K r ( K ) ⊂
U ∩ ( ∩ k U k ) .
L is µ ( K ) -PL on K r ( K ) .
inf x ∈ K r ( K ) ( λ 1 ( ∇ 2 L ( x ))
− λ 2 ( ∇ 2 L ( x )))
≥ ∆( K ) > 0 . 5.
inf x ∈ K r ( K ) λ M ( ∇ 2 L ( x )) ≥ µ ( K ) > 0 . 6.
inf x ∈ K r ( K ) λ 1 ( ∇ 2 L k ( x )) ≥ µ ( K ) > 0 .
Given compact set K ⊂ Γ , we further define
[Equation]
Similarly, we use notations like ζ k ( K ) , ν k ( K ) , Υ k ( K ) , ξ k ( K ) , χ k ( K ) to denote the counterpart of the above quantities defined for stochastic loss L k and its limiting map Φ k for k ∈
[ M ] .
Lemma C.6 (Arora et al.
(2022), Lemma B.5 and B.7) .
Given any compact subset K ⊂ Γ , let r ( K ) be defined in Lemma C.5, there exist 0 < h
( K ) <
r ( K ) such that
sup x ∈ K h ( K ) L ( x ) − inf x ∈ K h ( K ) L ( x ) ≤ µ ( K ) ρ 2 ( K ).
∀ x ∈ K h ( K ) , Φ( x ) ∈ K r ( K ) / 2 . 3.
∀ x ∈ K h ( K ) , ∥ x − Φ( x ) ∥ 2 ≤ 8 µ ( K ) 2 ζ
( K ) ν ( K ) .
The whole segment x Φ( x ) lies in K r ( K ) , so does x Φ k ( x ) , for any k ∈
[ D ] .
The proof of the lemmas above can be found in Arora et al.
Readers should note that although Arora et al.
(2022) only prove these lemmas when K is a special compact set (the trajectory of an ODE), all the proof does not use any property of K other than it is a compact subset of Γ , and thus our Lemmas C.5 and C.6 hold for general compact subsets of Γ .
In the rest part of the appendix, for convenience we will drop the dependency on K in various constants when there is no ambiguity.
C. P ROPERTIES OF L IMITING M AP OF G RADIENT F LOW , Φ
In our analysis, the property of Φ will be heavily used.
In this section, we will recap some related lemmas from Arora et al.
and then introduce some new lemmas for the stochastic setting with batch size one.
Lemma C.7 (Arora et al.
Lemma B.6) .
Given any compact set K ⊂ Γ , for any x ∈ K h ,
[Equation]
Lemma C.8.
Given any compact set K ⊂ Γ , for any x ∈ K h ,
[Equation]
Proof.
The first inequality is by Lemma C.5 and Taylor Expansion.
The second inequality is by Lemma C.7.
Lemma C.9 (Arora et al.
Lemmas B.16 and B.22) .
Given any compact set K ⊂ Γ , for any x ∈ K h ,
[Equation]
Lemma C.10 (Arora et al.
Lemmas B.8 and B.9) .
Given any compact set K ⊂ Γ , for any x ∈ K h ,
[Equation]
The proof of above lemmas can be found in Arora et al.
Lemma C.11.
Given any compact set K ⊂ Γ , for any x ∈ K h ,
[Equation]
Proof.
By Lemma C.6 and Taylor Expansion,
[Equation]
this proves the first claim.
Again by Lemma C.5 and Taylor Expansion,
[Equation]
this proves the second claim
[Equation]
Lemma C.12.
Given any compact set K ⊂ Γ , for any x ∈ K h and y = x − η ∇ L (cid:16) x
+ ρ ∇ L ( x ) ∥∇ L ( x ) ∥ (cid:17) ,
[Equation]
Proof.
For sufficient small ρ
[Equation]
x + ρ ∇ L ( x ) ∥∇ L ( x ) ∥ ∈ K r .
By Taylor Expansion,
[Equation]
This further implies that for sufficiently small η and ρ , xy ∈ K r .
Again by Taylor Expansion,
[Equation]
By Lemma C.9, ∂ Φ( x ) ∇ L ( x ) = 0 and ∂ Φ ( x ) ∇ 2 L ( x ) ∇ L ( x ) = − ∂ 2 Φ ( x ) [ ∇ L ( x ) , ∇ L ( x )] .
Hence,
[Equation]
As xy ∈ K r , by Taylor Expansion,
[Equation]
Putting together we have
[Equation]
Finally, by Lemma C.8, we have
[Equation]
Lemma C.13.
Given any compact set K ⊂ Γ , for any x ∈
K h and y = x − η ∇
[Equation]
Proof.
For sufficient small ρ
[Equation]
x + ρ ∇ L k ( x ) ∥∇ L k ( x )
∥ ∈ K r .
By Taylor Expansion,
[Equation]
This further implies that for sufficiently small η and ρ , xy ∈ K r .
Again by Taylor Expansion,
[Equation]
We further have by Lemma C.7,
[Equation]
Similarly
[Equation]
C. E XPLICIT B IAS
We will first prove the generic Theorem 4.3 and then apply this theorems to characterize a variety of regularizers.
Lemma C.14.
Γ = U ∩ Γ .
Proof of Lemma C.14.
For the first claim, for any point x ∈ U ′ ∩ Γ , there exists { x k } ∞ k =1 ∈ Γ such that lim k →∞ x k = x .
Since x ∈ U ′ ⊆
U and Φ is continuous in U , it holds that Φ is continuous at x , thus lim k →∞ Φ( x k ) =
Φ( x ) ∈ Γ .
However Φ( x k ) =
x k because x k ∈ Γ , ∀ k .
Thus we know x
= Φ( x ) ∈ Γ .
Lemma C.15.
Let U ′ be any bounded open set such that its closure U ′ ⊆
U .
If U ′ ∩ Γ ⊆ U ′ ∩ Γ , then U ′ ∩ Γ = U ′ ∩ Γ .
Proof of Lemma C.14.
By Lemma C.14, it holds that U ′ ∩ Γ =
U ′ ∩ U ∩ Γ = U ′ ∩ Γ .
Note that
U ′ ∩ Γ ⊆
U ′ , U ′
∩ Γ ⊆ Γ , we have that U ′ ∩ Γ ⊆
U ′
∩ Γ = U ′ ∩ Γ , which completes the proof.
Lemma C.16.
Let U ′ be any bounded open set such that its closure U ′ ⊆
U .
Further assume
U ′ ∩ Γ = U ′ ∩ Γ , then for all h 2 > 0 , ∃ ρ 0 > 0 if x ∈ U ′ , dist( x, Γ) ≤ ρ 0 ⇒ dist( x, U ′ ∩ Γ) ≤ h
Proof.
We will prove by contradiction.
Suppose there exists h 2 > 0 and { x k } ∞ k =1 ∈ U ′ , such that lim k →∞ dist( x k , Γ) =
but ∀ k > 0 , dist( x k , U ′ ∩ Γ) ≥ h 2 .
Since U ′ is bounded, U ′ is compact and thus { x k } ∞ k
=1 has at least one accumulate point x ∗ in U ′ ⊆
U .
Since U is the attraction set of Γ under gradient flow, we know that Φ( x ∗ ) ∈ Γ .
Now we claim x ∗ ∈ Γ .
This is because lim k →∞ dist( x k , Γ) =
and thus there exists a sequence of points on Γ , { y k } ∞ k =1 , where lim k →∞ ∥ x k − y k
∥
= 0 .
Thus we have that x ∗ =
lim k →∞ y k
= lim
k →∞ Φ( y k ) =
Φ( x ∗ ) , where the last step we used that x ∗ ∈ U and Φ is continuous on U .
By the definition of U , x ∗ ∈ U ⇐⇒ Φ( x ∗ ) ∈ Γ , thus x ∗ ∈ Γ .
Then we would have x ∗ ∈ U ′ ∩ Γ , which is contradictory to dist( x k , U ′ ∩ Γ) ≥ dist( x k , U ′ ∩ Γ) ≥ h 2 , ∀ k > 0 .
This completes the proof.
Lemma C.17.
Let U ′ be any bounded open set such that its closure U ′ ⊆
U .
then for all h 2 > 0 , ∃ ρ 1 > 0 if x ∈ U ′ , L ( x ) ≤ ρ 1 ⇒ dist( x, U ′ ∩ Γ) ≤ h
Proof.
We will prove by contradiction.
If there exists a list of ρ 1 , ..., ρ k , ... , such that ρ k → 0 and there exists x k ∈ U ′ , such that L ( x k ) ≤ inf x ∈ U ′
L ( x ) + ρ k and dist( x k , U ′ ∩ Γ) ≥ h 2 .
Since U ′ is bounded, U ′ is compact and thus { x k } ∞ k
=1 has at least one accumulate point x ∗ in U ′ ⊆
U .
Since
L is continuous in U , L ( x ∗ ) =
lim k →∞ L ( x k ) =
inf x ∈ U ′
L ( x ) .
Thus x ∗ is a local minimizer of
L and thus has zero gradient, which further implies that x ∗ =
Φ( x ∗ ) .
Thus x ∗ ∈ U ′ ∩ Γ , which is contradictory to dist( x k , U ′ ∩ Γ) ≥ dist( x k , U ′ ∩ Γ) ≥ h 2 , ∀ k > 0 .
This completes the proof.
Lemma C.18.
Let U ′ be any bounded open set such that its closure U ′ ⊆
U and U ′ ∩ Γ = U ′ ∩ Γ .
Suppose regularizers { R ρ } admits a limiting regularizer S on Γ , then
[Equation]
Proof of Lemma C.18.
First choose sufficiently small ρ , such that ρ < h ( U ′ ∩ Γ) . Choose an approximate minimizer of S ( x ) , x 0 ∈ U ′ ∩ Γ , such that S ( x 0 ) [-/ρ 
-]≤ inf x ∈ U ′
∩ Γ S ( x )
+ ρ 2 .
Then by the definition of limiting regularizers (Definition 4.2) and the assumption that U ′ is open, there exists x 1 ∈ U ′ satisfying that ∥ x 1 − x 0 ∥ 2 ≤ r ρ < ρ 2 and R ρ ( x 1 ) /ρ 2 − S ( x 0 ) ≤ ρ 2 .
Thus, R ρ ( x 1 ) ≤ ρ 2 S ( x 0 ) + ρ 4 . As ∥ x 1 − x 0 ∥ 2 ≤ ρ 2 < h and x 0 ∈ U ′ ∩ Γ . This further leads to x 0 x 1 ∈ U ′ ∩ Γ h .
By Taylor expansion on L at x 0 , we would have L ( x 1 ) ≤ L ( x 0 )+ O ( ∥ x 0 − x 1 ∥ 22 ) = inf x ∈ U ′ ∩ Γ L ( x )
+ O (
ρ 4 ) .
Thus it holds that [Equation] This completes the proof. Lemma C.19.
Let U ′ be any bounded open set such that its closure U ′ ⊆ U and U ′ ∩ Γ = U ′ ∩ Γ . Suppose regularizers { R ρ } admits a good limiting regularizer S on Γ , then for all u ∈ U ′ [-∩ Γ S ( x ) + ϵρ 
-], [Equation]
Proof of Lemma C.19. Define K = dist( u, U ′ ∩ Γ) . Define r = r ( K ) , h = h ( K ) as the constant in Lemma C.5 with K = U ′ ∩ Γ . Note K is compact and by Lemma C.15, K = U ′ ∩ Γ ⊂ Γ .
By Lemma C.5, we have K r ∩ Γ is a compact set, so is K h ∩ Γ . Since S is a good limiting regularizer for { R ρ } , by Definition 4.2, for any x ∗ ∈ K h ∩ Γ , there exists open neighborhood of x ∗ , V x ∗ such that for any C, ϵ 1 > 0 , there is a ρ x ∗ such that
[Equation] Note that K h ∩ Γ is compact, there exists a finite subset of K h ∩ Γ , { x k } k , such that K h ∩ Γ ⊂ ∪ k V x k . Hence for any C, ϵ 1 > 0 , there is some ρ K = min k ρ x k > 0 , it holds that,
[Equation] We can rewrite Equation (18) as for any C > 0 , [Equation] As u ∈ U ′ ⊆
U , we have that Φ( u ) ∈ Γ . If ∥ u − Φ( u ) ∥ 2 = O ( ρ ) , then dist( u, Γ) ≤ O ( ρ ) . By Lemma C.16, we have that dist( u, K ) = o (1) .
This further implies dist(Φ( u ) , K ) ≤ dist( u, K ) + dist(Φ( u ) , u ) =
o (1) .
Hence we have that Φ( u ) ∈ K h
∩ Γ for sufficiently small ρ .
Thus we can pick x = Φ( u ) in Equation (19) and C sufficiently large, which yields that
[Equation] where the last step is because ∥ u − Φ( u ) ∥ 2 = O ( ρ ) . On the other hand, we have that [Equation]
as S is continuous on Γ and dist( U ′ ∩ Γ , Φ( u ))
=
o (1) . Combining Equations (20) and (21), we have R ρ ( u ) ≥ ρ 2 inf x ∈ U ′ ∩ Γ S ( x ) − o ( ρ 2 ) . Proof of Theorem 4.3 .
We will first prove the ⇒ direction. By Lemma C.18, it holds that [Equation] [Equation]
which implies R ρ ( u ) ≤ ρ 2 inf x ∈ U ′ ∩ Γ S ( x ) + ρ 2 ϵ + o ( ρ 2 ) . Moreover, combining Equation (22) with Assumption 4.1, we have that
[Equation] [Equation] By Equation (23) and lemma C.7, ∥ u − Φ( u ) ∥ 2 = O ( ρ ) . Then by Lemma C.19, we have R ρ ( u ) ≥ ρ 2 inf x ∈ U ′
∩ Γ S ( x ) − o ( ρ 2 ) .
Putting this back into Equation (22), we have L ( u ) ≤ inf x ∈ U ′ ∩ Γ L ( x ) + ϵρ 2 + o ( ρ 2 ) . This implies, [Equation]
Finally notice that Equation (22) is equivalent to
[Equation]
Putting together the above claims we have our results.
For the ⇐ side, we will first lower bound L ( x ) + R ρ ( x ) for x ∈ U ′ .
Suppose C U
′ is the constant in Assumption 4.1 Define [Equation]
We discuss by cases. For sufficiently small ρ , If x ̸∈ K h , then by Lemma C.17, L ( x ) is lower bounded by a positive constant. If x ∈ K h and ∥ x − Φ( x ) ∥ 2 ≥ C 1 ρ , then by Lemma C.7,
[Equation] This implies L ( x )
[Equation] R
[Equation] x [Equation] x ∈
[Equation]
If ∥ x − Φ( x ∥ ≤ C 1 ρ
, by
Lemma C. , R ρ ( x ) ≥ ρ 2 inf x ∈ U ′ ∩ Γ S ( x ) − o ( ρ hence
[Equation] The result then follows directly from this lower bound.
Theorem 4.3 implies Corollary C.20 saying that the minimum of the regularized loss is approximately the sum of the minimum of loss and the minimum of limiting regularizer.
Corollary C.20.
Under the setting of Theorem 4.3,
[Equation]
Proof of Corollary C.20.
In Theorem 4.3, choose ϵ = 0 , and choose u such that L ( u ) + R ρ ( u ) ≤ inf x ∈ U ′
( L ( x ) + R ρ ( x ))
+ o (
ρ 2 ) , the result is then clear.
We will then prove the Corollary 4.4.
Proof of Corollary 4.4.
By Theorem 4.3, we have ∥ L ( u ) − inf x ∈ U ′
L ( x ) ∥
=
O ( ρ 2 ) .
This implies ∥ u − Φ( u ) ∥
=
O ( ρ ) and also dist( U ′ ∩ H, Φ( u ))
=
o (1) .
We also have ρ 2 S (Φ( x ))
− o ( ρ 2 ) ≤ R ρ ( u ) ≤ ρ 2 inf x ∈ U ′
∩ H S ( x ) .
We have ∥ S (Φ( x ))
− inf x ∈ U ′
∩ H S ( x ) ∥
=
o (1) .
This further implies ∥ S [Equation] Φ [Equation]
x )) − inf [Equation] U ′ ∩
[Equation]
Proof of Theorem 4.5.
We will first verify Assumption 4.1.
For fixed compact set B ⊂ U , as ∥∇ 3 L ( x ) ∥ 2 is continuous, there exists constant ν , such that ∀ x ∈ B 1 , ∥∇ 3 L ( x ) ∥ 2 ≤ ν .
Then by Taylor Expansion,
[Equation]
Now we verify S Max ( x ) =
λ 1 ( ∇ 2 L ( · )) / 2 is the limiting regularizer of R Max ρ .
Let x be any point in Γ , by continuity of R Max ρ ,
[Equation]
Finally we verify definition of good limiting regularizer, by Assumption 3.1, S Max ( x ) =
λ 1 ( x ) /is non-negative and continuous on Γ .
For any x ∗ ∈ Γ , choose a sufficiently small open convex set
V containing x ∗ such that ∀ x ∈ V 1 , ∥∇ 3 L ( x ) ∥ 2 ≤ ν .
For any x ∈ V ∩ Γ , for any x ′ satisfying that ∥ x ′ − x ∥ 2 ≤ Cρ , by Theorem C.67,
[Equation]
This implies inf
[Equation]
On the other hand, for any x ∈ V ∩ Γ ,
[Equation]
This implies
[Equation]
Thus, we conclude that (cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
cating S Max is a good limiting regularizer of R Max ρ on Γ .
Proof of Theorem 4.6.
We will first prove Assumption 4.1 holds.
For any fixed compact set B ⊂ U , as λ 1 ( ∇ 2 L ) and ∥∇ 3 L ∥ is continuous, there exists constant C , such that ∀ x ∈ B 2 , λ 1 ( ∇ 2 L ) >
− ζ and ∥∇ 3 L ( x ) ∥ <
ν .
Then by Taylor Expansion,
[Equation]
Now verify S ( x ) =
Tr( ∇ 2 L ( · )) / 2 is the limiting regularizer of R Asc ρ .
Let x be any point in Γ .
Let K = { x } and choose h = h ( K ) as in Lemma C [Equation]
For any x ′ ∈ K h ∩ U ′ ,
[Equation]
By Lemma C.10, we have ∇ L ( x ′ ) ∥∇ L ( x
[Equation]
∥ =
[Equation]
x
[Equation]
x
[Equation]
− Φ
[Equation]
x [Equation] )) ∥∇ [Equation]
Φ
[Equation]
x ′ ))
[Equation]
x ′ − Φ
[Equation] x [Equation] ∥ 2 + O ( νµ ∥ x ′
[Equation]
Φ
[Equation]
x ′ )
∥ 2 ) .
Hence
R Asc
[Equation]
This implies lim ρ
[Equation]
We now show the above inequality is in fact equality.
If we choose x ′′ r = x + rv M , then by Taylor
Expansion
[Equation]
This implies
[Equation]
We also have lim r → 0 ∇ 2 L ( x ′′ r ) = ∇ 2 L ( x ) and
lim r → 0 ∇ L ( x ′′ r )
= 0 .
Putting together,
[Equation]
L [Equation] x [Equation]
[Equation]
+ O ( ρ
[Equation]
This implies
[Equation]
Hence the limiting regularizer S is exactly λ M ( ∇ 2 L ( · )) / 2 .
Finally we verify definition of good limiting regularizer, by Assumption 3.1, S Max ( x ) =
λ M ( x ) /is non-negative and continuous on Γ .
For any x ∗ ∈ Γ , choose a sufficiently small open convex set
V containing x ∗ such that ∀ x ∈ V 1 , ∥∇ 3 L ( x ) ∥ 2 ≤ ν .
For any x ∈ V ∩ Γ , for any x ′ satisfying that ∥ x ′ − x ∥ 2 ≤ Cρ ,
[Equation]
By Lemma C.10, we have ∇ L ( x ′ )
∥∇ L ( x ′ )
∥
= ∇ 2 L ( Φ( x ′ ) )
( x ′ − Φ( x ′ )) ∥∇ 2 L (Φ( x ′ ))
( x ′ − Φ( x ′ ))
∥ 2 + O ( νµ ∥ x ′ − Φ( x ′ ) ∥ 2 ) .
This implies inf
[Equation]
On the other hand, simillar to the proof in the second part, we have inf
[Equation]
ρ 2 λ M ( ∇ 2 L ( x )) / 2 + O ( ρ 3 ) .
Thus, we conclude that (cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
- cating S Max is a good limiting regularizer of R Max ρ on Γ .
Proof of Theorem 4.7.
We will first verify Assumption 4.1.
For fixed compact set B ⊂ U , as ∥∇ 3 L ( x ) ∥ 2 is continuous, there exists constant ν , such that ∀ x ∈ B 1 , ∥∇ 3 L ( x ) ∥ 2 ≤ ν .
Then by Taylor Expansion,
[Equation]
Now we verify S Max ( x ) =
Tr( ∇ 2 L ( · )) / 2 D is the limiting regularizer of R Avg ρ .
Let x be any point in Γ , by continuity of R Avg ρ ,
[Equation]
Finally we verify definition of good limiting regularizer, by Assumption 3.1, S Avg ( x ) =
Tr( x ) / 2 D is non-negative and continuous on Γ .
For any x ∗ ∈ Γ , choose a sufficiently small open convex set V containing x ∗ such that ∀ x ∈ V 1 , ∥∇ 3 L ( x ) ∥ 2 ≤ ν .
For any x ∈ V ∩ Γ , for any x ′ satisfying that ∥ x ′ − x ∥ 2 ≤ Cρ , by Theorem C.67,
[Equation]
This implies
[Equation]
On the other hand, for any x ∈ V ∩ Γ ,
[Equation]
This implies
[Equation]
Thus, we conclude that (cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
indicating S Avg is a good limiting regularizer of R Avg ρ on Γ .
Proof of Theorem 5.1.
By Theorem C.49, Condition C.2 holds.
Easily deducted from Theorem 4.5 Λ k ( x ) is a good limiting regularizer for R max k,ρ on Γ k .
Then as Γ ⊂ Γ k , Λ k
( x ) is a good limiting regularizer for R max k,ρ on Γ .
Hence S ( x ) = (cid:80) k Λ k
( x ) / 2 M =
Tr( ∇ 2 L(x)) / 2 is a good limiting regularizer of E k [ R Max k,ρ ]( x ) on Γ .
Proof of Theorem 5.2.
By Theorem C.49, Condition C.2 holds.
Easily deducted from Theorem 4.6 Λ k ( x ) is a good limiting regularizer for R asc k,ρ on Γ k as the codimension of Γ k is 1.
Then as Γ ⊂ Γ k , Λ k
( x ) is a good limiting regularizer for R max k,ρ on Γ .Hence
S ( x ) =
(cid:80)
k Λ k
( x ) / 2 M =
Tr( ∇ 2 L ( x )) / 2 is a good limiting regularizer of E k
[ R Asc k,ρ ]( x ) on Γ .
C. F ULL - BATCH SAM ON Q UADRATIC L OSS (P ROOF OF T HEOREM 4.14)
The goal of this section is to prove Theorem 4.14.
In this section, we use A ≺ B to indicate B − A is positive semi-definite.
Theorem 4.14.
Suppose A is a positive definite symmetric matrix with unique top eigenvalue.
Consider running SAM on loss L ( x ) :
= 12 x T Ax as (11) , then for almost every x (0) , we have x ( t ) converges in direction to v 1 ( A ) and lim t →∞ ∥ x ( t ) ∥
= ηρλ 1 ( A ) 2 − ηλ 1 ( A ) with ηλ 1 ( A ) < 1 .
[Equation]
Proof of Theorem 4.14.
We first rewrite the iterate as
[Equation]
Define ˜ x ( t ) ≜ ∇ L ( x ( t ))
ρ =
Ax ( t ) ρ , and we have
[Equation]
We suppose A ∈ R D × D and use λ i , v i to denote
λ i ( A ) , v i ( A ) .
Further, we define that
[Equation]
By Lemma C.21, I j is an invariant set for update rule Equation (24).
Our proof consists of two steps.
(1) Entering Invariant Set.
Lemma C.22 implies that there exists constant T 1 > 0 , such that ∀ t > T 1 ,
∥ P ( j : D ) ˜ x
( t ) ∥ 2 ≤ ηλ 2 j (2) Alignment to Top Eigenvector.
Lemmas C.29 and C.30 show that ∥ ˜ x ( t ) ∥ 2 and | ˜ x 1 ( t ) | converge to ηλ 21 2 − ηλ 1 , which implies our final results.
C.4. E NTERING I NVARIANT S ET
In this subsection, we will prove the following three lemmas.
Lemma C.21 shows I j is an invariant set for update rule (Equation (24)).
Lemma C.22 shows that under the update rule (Equation (24)), all iterates not in I j will shrink exponentially in ℓ 2 norm.
Lemma C.23 combines Lemmas C.21 and C.22 to show that for sufficiently large t , x ( t ) ∈
∩ j I j .
Lemma C.21.
For t ≥ 0 , if ηλ 1 ( A ) < 1 and ˜ x ( t ) ∈
I j , then ˜ x
( t + 1) ∈
I j .
Proof of Lemma C.21.
By (Equation (24)), we have that
[Equation]
Hence we have that
[Equation]
t [Equation] Because ˜ x ( t ) ∈ I j , ∥ ˜ x ( t ) ∥ 2 ≤ ηλ 2 j 1 − ηλ j .
This implies,
I
[Equation]
Furthermore
[Equation]
Hence it holds that
[Equation]
where the last equality is because 1 − ηλ j ≥ 0 .
This above inequality is exactly the definition of ˜ x
( t + 1) ∈
I j and thus is proof is completed.
Lemma C.22.
For t ≥ 0 , if ηλ 1 ( A ) < 1 and ˜ x
( t )
̸∈ I j , then
[Equation]
Proof of Lemma C.22.
Note that
[Equation]
As ˜ x
( t )
̸∈ I j , We have ∥ ˜ x
( t ) ∥ 2 ≥ ∥ P ( j : D )
˜ x ( t ) ∥ 2 > ηλ 2 j , hence η P ( j : D )
A 2 ∥
˜ x
( t ) ∥ 2 ≺ η P
[Equation]
j
[Equation]
D
[Equation]
j ≺ P
[Equation] j : D
[Equation]
This implies that
[Equation]
and
P (
[Equation]
Hence we have that
[Equation]
This completes the proof.
Lemma C.23.
Choosing T 1 = max j (cid:16) − log max(1 − ηλ D ,ηλ j ) max( ∥ ˜ x (0) ∥ 2 ηλ 2 j , 1) (cid:17) , then ∀ t ≥ T 1 , D > j ≥ 1 , ˜ x ( t ) ∈ I j
Proof of Lemma C.23.
We will prove by contradiction.
Suppose ∃ j ∈
[ D ] and T > T 1 , such that ˜ x ( T ) ̸∈ I j .
By Lemma C.21, it holds that ∀ t < T, ˜ x
( t )
̸∈ I j .
Then by Lemma C.22, [Equation]
which leads to a contradiction.
C.4. A LIGNMENT TO T OP E IGENVECTOR
In this subsection, we prove the following lemmas towards showing that ˜ x ( t ) converges in direction to v 1 ( A ) up to a proper sign flip.
Lemma C.24 shows that under update rule (Equation (24)), t ̸∈ S ⇒ t + 1 ∈ S for sufficiently large t , where the definition of S is { t |∥ ˜ x ( t ) ∥ 2 ≤ ηλ 21 2 − ηλ 1 , t > T 1 } .
Lemma C.28, a combination of Lemmas C.25 and C.26, shows that following update rule (Equation (24)), ˜ x 1 ( t ) increases for t ∈ S . 3.
Lemma C.29 shows that ∥ ˜ x ( t ) ∥ converges to ηλ 21 2 − ηλ 1 under Equation (24).
Lemma C.30 shows that ∥ ˜ x 1 ( t ) ∥ 2 converges to ηλ 21 2 − ηλ 1 under Equation (24).
Lemma C.24.
For t ≥ 0 , if ∥ ˜ x
( t ) ∥ 2 > ηλ 21 2 − ηλ 1 , ˜ x
( t ) ∈
∩ I j , then
[Equation]
Proof of Lemma C.24.
Note that
[Equation]
Consider the following two cases.
1 If for any i , such that (cid:12)(cid:12) (1 − ηλ 1 ) ∥
˜ x ( t ) ∥ 2 − ηλ 21 (cid:12)(cid:12) ≥ (cid:12)(cid:12) (1 − ηλ i ) ∥
˜ x ( t ) ∥ 2 − ηλ 2 i (cid:12)(cid:12) , then we have ∥ ˜ x ( t + 1) ∥ 2 ≤ (cid:12)(cid:12) (1 − ηλ 1 ) ∥
˜ x ( t ) ∥ 2 − ηλ
[Equation]
= ηλ 21 −
[Equation]
[Equation]
1 ) ∥
˜ x ( t ) ∥ 2 .
If there exists i , such that (cid:12)(cid:12) (1 − ηλ 1 ) ∥
˜ x ( t ) ∥
[Equation]
− ηλ 21 (cid:12)(cid:12) < (cid:12)(cid:12) (1 − ηλ i ) ∥
˜ x ( t ) ∥ 2 − ηλ 2 i (cid:12)(cid:12) , then suppose
WLOG, i is the smallest among such index.
As
[Equation]
We have − ηλ 2 i + ( [Equation] i
[Equation] 21 − [Equation] [Equation]
∥ ˜ x ( t ) ∥ 2 .
Equivalently,
[Equation]
Combining with ˜ x ( t ) ∈ I 1 ⇒ ∥
˜ x ( t ) ∥ 2 ≤ ηλ 21 , we have η < λ 1 − λ i λ 21 .
Now consider the following vertors,
[Equation]
Then we have
[Equation]
(cid:0) (1 − ηλ j ) ∥ ˜ x ( t ) ∥ 2 − ηλ 2 j (cid:1) ˜ x j ( t ) v j ∥≤∥ ∥ ˜ x ( t ) ∥i − 1 (cid:88)
j = (cid:0) ηλ 21 − (1 − ηλ 1 ) ∥ ˜ x ( t ) ∥ 2 (cid:1) ˜ x j ( t ) v j ∥ + ∥ ∥ ˜ x ( t ) ∥
D (cid:88)
j = i (cid:0) (1 − ηλ j ) ∥ ˜ x ( t ) ∥ 2 − ηλ 2 j (cid:1) ˜ x j ( t ) v j ∥≤ ∥ ˜ x ( t ) ∥
D +1 − i (cid:88)
j = ∥ v ( j ) ∥ By assumption, we have ˜ x ( t ) ∈ ∩ I j , hence we have ∥ v (1) ( t ) ∥ 2 = ( ηλ 21 − (1 − ηλ 1 ) ∥ ˜ x ( t ) ∥ 2 ) ∥
˜ x ( t ) ∥ 2 , ∥ v (2) ( t ) ∥ 2 ≤ η ((2 − ηλ 1 − ηλ i ) ∥
˜ x ( t ) ∥ 2 − ηλ 2 i − ηλ 21 ) λ
2 i , ∥ v (2+ j ) ( t ) ∥ 2 ≤ η (( ηλ i + j − 1 − ηλ i + j ) ∥ ˜ x ( t ) ∥ 2 − ηλ 2 i + j + ηλ 2 i + j − 1 ) λ 2 i + j , 1 ≤ j ≤ D − i .
Using AM-GM inequality
[Equation]
we have
[Equation]
Hence
[Equation]
Putting together,
[Equation]
We further discuss three cases
If ηλ i (cid:113) 12 λ 2 i + λ 21 1 − ηλ 1 < ηλ 21 + ηλ [-
-]2 i 2 − ηλ 1 − ηλ i , we have ∥ ˜ x
( t ) ∥ 2 > ηλ 21 + ηλ 2 i 2 − ηλ 1 − ηλ i > ηλ i (cid:113)
12 λ 2 i + λ 21 1 − ηλ 1 ,then ∥ ˜
[Equation]
The second line is because (1 − ηλ 1 ) ∥
˜ x ( t ) ∥
+ η 2 λ 2 i ( 12 λ 2 i + λ 21 ) 1 ∥
˜ x ( t ) ∥ 2 monotonously in
- crease w.r.t ∥
˜ x ( t ) ∥ when
∥ ˜ x ( t ) ∥ 2 > ηλ i (cid:113) [-
-]12 λ 2 i + λ 21 1 − ηλ 1 .
The last line is due to technical lemma
Lemma C.
If ηλ 21 ≥ ηλ i (cid:113)
λ 2 i + λ
[Equation]
1 ≥
[Equation]
21 +
[Equation]
[Equation]
[Equation]
[Equation]
i , then
[Equation]
The second line is because of AM-GM inequality.
The last line is due to technical lemma Lemma C.73.
If ηλ 21 < ηλ i (cid:113)
λ 2 i + λ 21 1 − ηλ 1 , we have ∥ ˜ x
( t ) ∥ 2 < ηλ 21 <
ηλ i (cid:113)
λ 2 i + λ 21 1 − ηλ 1 , then
[Equation]
The second line is because (1 − ηλ 1 ) ∥
˜ x ( t[-) ∥ 
-]
) ∥ 2 + η 2 λ 2 i ( 12 λ 2 i + λ 21 ) 1 ∥ ˜ x
( t ) ∥ 2 monotonously decrease [-w.r.t ∥ ˜ x ( t ) ∥ 2 when ∥ ˜ x ( t ) ∥ 2 < ηλ i (cid:113) 
-]w λ
Lemma C.25.
if ∥ ˜ x
( t ) ∥ 2 ≤ ηλ 21 2 − ηλ 1 , it holds that | ˜ x 1 ( t + 1) | ≥ | ˜ x 1 ( t ) | .
Proof.
We have | ˜ x 1 ( t + 1) | = | 1 − ηλ 1 − η λ 21 ∥
˜ x ( t ) ∥ 2 || ˜ x 1 ( t ) | .
We also have η λ 21 ∥
˜ x ( t ) ∥ 2 > 2 − ηλ 21 , hence 1 − ηλ 1 − η λ 21 ∥
˜ x ( t ) ∥ 2 < − 1 .
Hence we have that | ˜ x 1 ( t + 1) | > | ˜ x 1 ( t ) | .
Lemma C.26.
For any t ≥ 0 , if ∥ ˜ x
( t ) ∥ 2 ≤ ηλ 21 2 − ηλ 1 , ˜ x ( t ) ∈
∩ I j , it holds that
[Equation]
Proof.
The result holds as
[Equation]
Lemma C.27.
For any t ≥ 0 , if ∥ ˜ x
( t ) ∥ 2 ≤ ηλ 21 1 − ηλ 1 , it holds that ∥ ˜ [Equation] t [Equation]
[Equation]
Proof.
We will discuss the movement along v 1 and orthogonal to v 1 .
First,
[Equation]
Second
[Equation]
Hence we have that ∥ ˜ x ( t + 1) ∥≤ ( λ 21 − (1 + ηλ 1 ) ∥ ˜ x
[Equation]
Lemma C.28.
For t, t ′ ∈ S, 0 ≤ t ≤ t ′ , then | ˜ x 1 ( t ) | ≤ | ˜ x 1 ( t ′ ) | .
Proof.
For t ∈ S , by Lemma C.24, t + 1 ∈ S or t + 1 ̸∈ S, t + 2 ∈ S .
We will discuss by case.
If t + 1 ∈ S , we can use Lemma C.25 to show | ˜ x 1 ( t ) | ≤ | ˜ x 1 ( t + 1) | . If t + 1 ̸∈ S, t + 2 ∈ S , then [Equation]
As ( ηλ [Equation] [Equation] ) ∥ ˜
[Equation]
( t [Equation] [Equation] [Equation]
[Equation]
( t
[Equation]
1) ∥ 2 ) [Equation] [Equation] ) ∥ 2 ∥ ˜
[Equation]
( t + 1) ∥ ⇐⇒ η
[Equation]
λ 41 − ηλ 21 (1 − ηλ 1 )( ∥ ˜ x ( t ) ∥ 2 + ∥ ˜ x ( t + 1) ∥ 2 ) ≥ (2 ηλ 1 − η 2 λ 21 ) ∥ ˜ x ( t ) ∥
[Equation]
combining with Lemma C.26, we only need to prove,
[Equation]
Through some calculation, this is equivalent to
[Equation]
which holds for ∥ ˜ x ( t ) ∥ 2 ≤ ηλ 21 2 − ηλ 1 .
Combining the two cases and using induction, we can get the desired result.
Lemma C.29.
∥ ˜ x ( t ) ∥ converges to ηλ 21 2 − ηλ 1 when t → ∞
Proof.
By Lemma C.28, | ˜ x 1 ( t ) | increases monotonously for t ∈ S .
By Lemma C.24, S is infinite.
By Lemma C.22, for sufficiently large t , | ˜ x 1 ( t ) | is bounded.
Combining the three facts, we know ˜ x 1 ( t ) for t ∈ S converges.
Formally ∀ ϵ > 0 , there exists T ϵ > 0 such that ∀ t, t ′ ∈ S, t ′ > t [Equation]
T [Equation]
< [Equation] + [Equation]
Then by Lemma C.24, ∀ t ∈ S, t + 1 ∈ S or t + 2 ∈ S , we will discuss by case.
For t ≥ T ϵ ,
If t + ∈ S , then
[Equation]
If t + 1 ̸∈ S and t + 2 ∈ S , then
[Equation]
ϵ
[Equation]
∥ ˜ x 1 ( t + 2) ∥
[Equation]
˜
[Equation]
( ηλ 21 − (1 − ηλ 1 ) ∥
˜ x ( t ) ∥ 2 )( ηλ 21 − (1 − ηλ 1 ) ∥ ˜ x ( t + 1) ∥ 2 ) ∥ ˜ x ( t [Equation]
2 ∥ ˜
[Equation]
[Equation]
+ 1) ∥≥ ( ηλ 21 − (1 − ηλ 1 ) ∥
˜ x ( t ) ∥ 2 ) (cid:0) ηλ 21 − (1 − ηλ 1 ) (cid:0) ηλ 21 − (1 − ηλ 1 ) ∥
˜ x ( t ) ∥ 2 (cid:1)(cid:1) ∥
˜ x ( t ) ∥ 2 ( ηλ 21 − (1 − ηλ 1 ) ∥ ˜ x ( t ) ∥ 2 ) = ηλ 21 − (
[Equation] 1 ) (cid:0) ηλ 21 − (
[Equation]
∥ ˜ x ( t ) ∥. Here in the last inequality, we apply Lemma C.26. Concluding, ∥ ˜ x ( t ) ∥ 2 ≥ min (cid:16) ηλ 21 2 − ηλ 21 + ϵ , η 2 λ 31 (2 − λ 1 η ) λ 1 η + ϵ (cid:17) , ∀ t > T ϵ , t ∈ S .
As ∀ t ̸∈ S, t > T ϵ , we have ∥ ˜ x
( t ) ∥ 2 ≥ ηλ 21 2 − ηλ 21 .
Hence we have ∀ t > T ϵ , ∥
˜ x ( t ) ∥ 2 ≥ min (cid:16) ηλ 21 2 − ηλ 21 + ϵ , η 2 λ 31 (2 − λ 1 η ) λ 1 η + ϵ (cid:17) .
Further by Lemma C.26, ∀ t > T ϵ + 1 , ∥ ˜ x ( t ) ∥ 2 ≤ ηλ 21 − (1 − ηλ 1 ) min (cid:16) ηλ 21 2 − ηλ 21 + ϵ , η 2 λ 31 (2 − λ 1 η ) λ 1 η + ϵ (cid:17) .
Combining both bound, we have lim t →∞ ∥
˜ x ( t ) ∥ 2 = ηλ 21 2 − ηλ 1 .
Lemma C.30.
∥
˜ x 1 ( t ) ∥ 2 converges to ηλ 21 2 − ηλ 1 , when t → ∞ .
Proof.
Notice that ∥
[Equation]
When ∥ ˜ x ( t ) ∥ 2 >
[Equation]
, −
[Equation] +
[Equation]
Hence for sufficiently large t , ∥ P (2: D ) ˜ x
( t ) ∥ 2 shrinks exponentially, showing
that lim t →∞ ∥ ˜ x 1 ( t ) ∥ 2 = ηλ 21 2 − ηλ 1 .
C. F ULL - BATCH SAM ON G ENERAL L OSS (P ROOF OF T HEOREM 4.9)
The goal of this section is to prove the following theorem.
Theorem 4.9 (Main, Theorems C.31 and C.32 stated informally) .
Let { x ( t ) } be the iterates of
SAM (3) with x (0) =
x init ∈ U , then under Assumptions 3.1 and 4.8, for all η, ρ such that η ln(1 /ρ )
and ρ/η are sufficiently small, the dynamics of SAM can be split into two phases:
• Phase I: SAM follows Gradient Flow with respect to L until entering an O ( ηρ ) neighborhood of the manifold Γ in ˜ O
( 1 η ) steps; • Phase II: SAM tracks the solution X of (6) , the Riemannian Gradient Flow with respect to λ 1 ( ∇ 2 L ( · )) in the O ( ηρ ) neighborhood in the sense
that max 0 ≤ T ≤ T 3 ∥
Φ( x ( ⌈ T/ ( ηρ 2 ) ⌉ ))
− X ( T ) ∥
=
O (( η + ρ ) log(1 /ηρ )) .
Moreover, the angle between ∇ L ( x ( t )) and v 1 ( ∇ 2 L ( x ( t ))) is at most O ( ρ ) .
To prove the theorem, we will separate the dynamic of SAM on general loss L to two phases.
Define
[Equation]
which is the length projection of x − Φ( x ) on button − k non-zero eigenspace of ∇ 2 L (Φ( x ))
We will provide a fine-grained convergence bound on R j ( x ) .
Theorem C.31 (Phase I) .
Let { x ( t ) } be the iterates defined by SAM ( Equation (3) ) and x ( t ) =
x init ∈ U , then under Assumption 3.1 there exists a constant T 1 , such that for any T ′ 1 > T 1 , it holds for all η, ρ such that ( η + ρ ) ln(1 /ηρ ) is sufficiently small, we have
[Equation]
Theorem C.31 implies SAM will converge to an O ( ηρ ) neighbor of Γ .
Notice in the time frame defined by Theorem C.31, x ( t ) effectively operates at a local regime around Φ( ⌈ T 1 ln(1 /ηρ )
/η ⌉ ) , this allows us to approximate L with the quadratic Taylor expansion of L at Φ( ⌈ T 1 ln(1 /ηρ ) /η ⌉ ) and give us the following theorem.
Theorem C.32 (Phase II) .
Let { x ( t ) } be the iterates defined by SAM (Equation (3) ) under Assumptions 3.1 and 4.8, for all η, ρ such that η ln(1 /ρ )
and ρ/η is sufficiently small, further assuming that (1) max j R j ( x (0))
=
O ( ηρ 2 ) , (2) ∥
Φ( x (0)) − Φ( x init ) ∥
=
O (( η + ρ ) ln(1 /ηρ )) and (3) |⟨ x (0) − Φ( x (0)) , v 1 ( x (0))
⟩| ≥ Ω( ρ 2 ) , the iterates x ( t ) tracks the solution X of Equation (6) .
Quantitatively for t =
⌈ T 3 /ηρ 2 ⌉ , we have that
[Equation]
Moreover the angle between ∇ L ( x ( t )) and the top eigenspace of ∇ 2 L (Φ( x ( t ))) is at most O ( ρ ) .
Quantitatively,
[Equation]
In this section we will define K as { X ( t ) | 0 ≤ t ≤ T 3 } where X is the solution of Equation (6).
To simplify our proof, we assume WLOG L ( x )
= 0 for x ∈ Γ .
C.5. P HASE I (P ROOF OF T HEOREM C.31)
The goal of this section is to prove Theorem C.31.
Proof.
The proof consists of three major steps.
Tracking Gradient Flow Lemma C.33 shows the existence of step t GF =
O (1 /η ) such that x ( t GF ) is in a subset of K h and Φ( x ( t GF ))
is O
( η + ρ ) close to Φ( x init ) .
Decreasing Loss Lemma C.35 shows the existence of step t DEC
= O (ln(1 /ρ )
/η )
such that x ( t DEC ) is in O ( ρ ) neighbor of Γ and Φ( x ( t DEC )) is O
(( η + ρ ) ln(1 /ρ )) close to Φ( x init ) .
Entering Invariant Set Lemmas C.40 and C.42 shows the existence of step t INV =
O (ln(1 /ρη ) /η )
such that for any t satisfying t INV ≤ t ≤ t INV + Θ(ln(1 /η ) /η ) , we have that x ( t ) ∈
∩ k ∈
[ M ] I k and Φ( x ( t )) is O
(( η + ρ ) ln(1 /ηρ ))
close to Φ( x init ) .
Tracking Gradient Flow Lemma C.33 shows that the iterates x ( t ) tracks gradient flow to an O (1) neighbor of Γ .
Lemma C.33.
Under condition of Theorem C.31, there exists t GF =
O (1 /η ) , such that the iterate x ( t GF ) is O (1) close to the manifold Γ and Φ( x ( t GF )) is O
η + ρ ) is close to Φ( x init ) .
Quantitatively,
[Equation]
Proof.
Choose C = 14 (cid:113) µζ .
There exists T > 0 , such that
[Equation]
Consider
[Equation]
By Theorem D.5, by choosing p = η and q = η + ρ , the iterates x ( t ) tracks gradient flow ϕ ( x init , T ) in O (1 /η ) steps.
Quantitatively for t GF =
Tη ⌉ , we have that
[Equation]
This implies x ( t GF
) ∈ K h , hence by Taylor Expansion on Φ ,
[Equation]
This implies
[Equation]
By Taylor Expansion,
[Equation]
Decreasing Loss Lemma C.35 shows that the iterates x ( t ) converges to an O ( ρ ) neighbor of Γ in
O (ln(1 /ρ ) /η ) steps.
Lemma C.34.
Under condition of Theorem C.31, if x ( t ) ∈ K h and ∥∇ L ( x ( t ))
∥ ≥ 4 ζρ , then we have that L ( x ( t + 1)) decreases with respect to L ( x ( t )) , quantitatively, we have that
[Equation]
Moreover the movement of the projection of the iterates on the manifold is bounded, quantitatively, we have that
[Equation]
Proof.
As x ( t ) ∈ K h and L is µ -PL in K h , we have L ( x ( t )) ≥ 0 .
As x ( t ) ∈ K h , by Lemma C.12 and Taylor Expansion, we have ∥ x ( t ) x ( t + 1) ∥ =
O ( η ) .
hence for sufficiently small η , x ( t ) x ( t + 1) ⊂ K r .
Using similar argument, the segment from x ( t ) to x ( t ) + ρ ∇ L ( x ( t ))
∥∇ L ( x ( t ))
∥ is in K r .
Then by Taylor Expansion on L ,
[Equation]
By Taylor Expansion on ∇ L , we have that
[Equation]
After plugging in Equation (27), we have that
[Equation]
As ∥∇ L ( x ( t ))
∥ ≥ 4 ζρ , we have that the following term is bounded.
[Equation]
After plugging in Equation (28), by Lemma C
[Equation]
[Equation]
As x ( t ) ∈ K h , by Taylor Expansion, we have
[Equation]
Hence by Lemma C.12 and Taylor Expansion,
[Equation]
Lemma C.35.
Under condition of Theorem C.31, assuming there exists t GF such that L ( x ( t GF )) ≤ µh 2 32 and x ( t GF ) ∈
K h/ 4 , then there exists t DEC = t GF + O
(ln(1 /ρ )
/η ) , such that x ( t DEC ) is in
O ( ρ ) neighbor of Γ , quantitatively, we have that
[Equation]
Moreover the movement of the projection of Φ( x ( · )) on the manifold is bounded,
[Equation]
Proof.
Choose t DEC as the minimal t ≥ t GF such that ∥∇ L ( x ( t DEC ))
∥ 2 ≤ 4 ζρ .
Define C = ⌈ ln 1 − ηµ 8 (64 ρ 2 /h 2 )
⌉ =
O (ln(1 /ρ ) /η ) .
We will first perform an induction on t ≤ min { t DEC , t GF + C } =
t GF + O (ln(1 /ρ )
/η )
to show that
[Equation]
For t = t GF , the result holds trivially.
Suppose the induction hypothesis holds for t .
Then by C. and Taylor Expansion,
[Equation]
Then we have that
[Equation]
) ∥ ≤ 3 h/ 4 + O ( η 2 ( t − t GF ))
= 3 h/ 4 + O ( η ln(1 /ρ ))
≤ h .
That is x ( t ) ∈ K h .
Then as t ≤ t DEC , ∥∇ L ( x ( t ))
∥ 2 ≥ 4 ζρ .
Then by Lemma C.34, we have that
L ( x ( t + 1)) ≤ (1 − ηµ/ 8) L ( x ( t )) ≤ (1 − ηµ/ 8) t +1 − t GF L ( x ( t GF )) , ∥ Φ( x ( t + 1)) − Φ( x ( t GF ))
∥ ≤ ∥ Φ( x ( t + 1)) − Φ( x ( t ))
∥ + ∥ Φ( x ( t )) − Φ( x ( t GF ))
∥ ≤ O ( η 2 ( t − t GF )) .
The induction is complete.
Now if t DEC ≥ t GF + C = t GF + O
(ln(1 /ρ )
/η ) , As the result of the induction, we have that
[Equation]
By Lemma C.8, we have that
[Equation]
This is a contradiction.
Hence we have that t DEC ≤ t GF + C =
t GF + O
(ln(1 /ρ ) /η ) .
By induction, we have that ∥ Φ( x ( t DEC )) − Φ( x ( t GF )) ∥ = O ( η 2 ( t DEC − t GF ))
= O ( η ln(1 /ρ )) .
Entering Invariant Set We will introduce some notations that is required for the proof in this and following subsection.
Define
[Equation]
Note ˜ x
≈ ∇ L ( x ) for x near the manifold Γ .
We also use ˜ x ( t ) , A ( t ) and ˆ x ( t ) to denote ˜ x ( t ) ,
A ( x ( t )) and ˆ x ( t ) .
Recall the original definition of R j ( x ) is
[Equation]
Based on the above notions, we can rephrase the notion R as
[Equation]
We additionally define the approximate invariant set I j as
[Equation]
Lemma C.36.
Under condition of Theorem C.31, assuming t satisfy that x ( t ) ∈ K h , then we have that
[Equation]
Proof.
First by Lemma C.10, Φ( x ( t ))
∈ K r , hence
[Equation]
Also
[Equation]
By Lemma C.10 and Lemma C.6, we have
[Equation]
Hence ∥ x ( t ) − Φ( x ( t ))
∥ ≤ 2 µ ∥ ˜ x
( t ) ∥ .
Lemma C.37.
Under condition of Theorem C.31, assuming t satisfy that x ( t ) ∈ K h and ∥ ˜ x ( t ) ∥ 2 =
O ( ρ ) , then we have that
[Equation]
Proof.
By Lemma C.36, we have ∥ x ( t ) − Φ( x ( t )) ∥
=
O ( ρ ) .
By Lemma C.12, we have that
[Equation]
Lemma C.38.
Under condition of Theorem C.31, assuming t satisfy x ( t ) ∈ K h/ 2 and ∥ x ( t ) − Φ( x ( t ))
∥ 2 =
O ( ρ ) , define x ′ as x ′ ( t ) =
x ( t ) and
[Equation]
Then
[Equation]
and further if ∥ x ( t + 1) − Φ( x ( t + 1)) ∥ 2 = Ω( ηρ ) , then
[Equation]
Proof.
For simplicity, we will sometime omit the t in x ( t ) .
By ∥ x ( t ) − Φ( x ( t ))
∥
=
O ( ρ ) , x ( t ) ∈
K h/ 2 , and Lemma C.12, we have that ∥ x ( t + 1) − x ( t ) ∥ =
O ( ηρ ) and hence x ( t +1)
∈ K 3 h/ 4 .
This also implies ∥ x ( t +1) − Φ( x ( t +1))
∥
= O ( ρ ) .
Similarly we have x ( t + 2) ∈ K 3 h/ 4 .
For k ∈ { 1 , 2 } , by Taylor Expansion,
[Equation]
k[-)) − Φ( x ( t )) ∥ 
 = O ( ηρ 2-] )
[Equation]
Now by Lemmas C.36 and C.37, ∥ Φ( x ( t + k )) − Φ( x ( t
[Equation]
∥ 2 = O ( ηρ 2 ) ,
[Equation]
Now we first prove the first claim, we have for k = 0 , ∥ x ( t + k ) − x ′ ( t + k ) ∥ 2 = 0 , by Lemma C. and eq.
[Equation]
The second claim is slightly more complex.
By the first claim and Lemma C.10, we have that
[Equation]
[Equation]
We first show ∥∇ 2 L (Φ( x ( t +1)))
( x ( t +1) − Φ( x ( t +1)))
∥ 2 is of order ∥
x ( t +1) − Φ( x ( t +1))
∥ 2 = Ω( ρ 2 ) to show that the normalized gradient term is stable with respect to small perturbation,
[Equation]
[Equation]
[Equation]
Based on Lemma C.12, we have
[Equation]
We further have by the first claim and Lemma C.37,
[Equation]
This implies
[Equation]
Combining with Equation (30) [Equation] we have
[Equation]
By the above approximation and Equation (29)
[Equation]
Lemma C.39.
Under condition of Theorem C.31, assuming assuming t satisfy that x ( t ) ∈ K 3 h/and ∥ ˜ x ( t ) ∥ 2 = O ( ρ ) , then we have that
[Equation]
Proof.
By Lemma C.38, we know
[Equation]
This implies
[Equation]
We also have
[Equation]
Plugging in Equation (31), we have that
[Equation]
Lemma C.40.
Under condition of Theorem C.31, assuming there exists t DEC such that x ( t DEC ) ∈
K h/ 2 and ∥∇ L ( x ( t DEC ))
∥ ≤ 4 ζρ , then there exists t DEC2 = t DEC + O (ln(1 /η ) /η ) , such that x ( t DEC2 ) is in I 1 ∩ K 3 h/ 4 .
Furthermore, for any t satisfying t DEC2 ≤ t ≤ t DEC2 + Θ(ln(1 /η ) /η ) , we have that x ( t ) ∈
I 1 ∩ K 3 h/ 4 and ∥ Φ( x ( t ))
− Φ( x ( t DEC ))
∥
=
O ( ρ 2 ln(1 /η )) .
Proof.
For simplicity, denote C = ⌈ ln 1 − ηµ ηµ 3 4 ζ 2 ⌉ + Θ(ln(1 /ρ ) /η )
=
O (ln(1 /η ) /η ) .
Here the quantity Θ(ln(1 /ρ ) /η ) is the same quantity in the statement of the lemma.
We will prove the induction hypothesis for t DEC ≤ t ≤ t DEC + 2 C , 
[Equation]
The induction hypothesis holds trivially for t = t DEC .
Assume the induction hypothesis holds for t ′ ≤ t .
By Lemmas C.7 and C.36, ∥ ˜ x ( t DEC ) ∥ 2 ≤ ζ ∥ x ( t DEC ) − Φ( x ( t DEC ))
∥ ≤ ζµ ∥∇ L ( x ( t DEC ))
∥ ≤ 4 ζ 2 µ ρ .
Combining with the induction hypothesis, we have ∥ ˜ x
( t ) ∥ ≤ 4 ζ 2 µ ρ .
By x ( t ) ∈ K 3 h/ 4 and Lemma C.37, we have that
[Equation]
Hence we have that
[Equation]
This proves the third statement of the induction hypothesis.
By ∥ ˜ x ( t ) ∥
=
O ( ρ ) and Lemma C.39, we have that
[Equation]
Analogous to the proof of Lemmas C.21 and C.22, we have
If ∥ ˜ x
( t ) ∥ > ηρλ 21 ( t ) , we would have
[Equation]
Hence we have
[Equation]
If ∥ ˜ x
( t ) ∥ 2 ≤ ηρλ 21 ( t ) , then by Lemma C.21, we have that
[Equation]
Hence by Lemma C.
[Equation]
Concluding the two cases, we have shown the first and second claim of the induction hypothesis holds.
Hence we can show that ∥ ˜ x
( t +1) ∥ ≤ 4 ζ 2 µ ρ .
Then by Lemma C.36, we have that ∥ x ( t +1) − Φ( x ( t + 1))
∥ ≤ 8 ζ 2 µ 2 ρ .
As
[Equation]
η ) /η ) , by Equation (32),
[Equation]
This implies dist
[Equation]
This proves the fourth claim of the inductive hypothesis.
The induction is complete.
Now define t DEC2 the minimal t ≥ t DEC , such that ∥ ˜ x
( t ) ∥ ≤ ηρλ 21 ( t ) .
If t DEC2 > t DEC + C , then by the induction, Lemmas C.7 and C.36,
[Equation]
This is a contradiction.
Hence we have t DEC2 ≤ t DEC + C .
By the induction hypothesis x ( t DEC2 ) ∈
I 1 ∩ K 3 h/ 4 .
Furthermore by induction, for any t satisfying t DEC2 ≤ t ≤ t DEC + 2 C , we have that
[Equation]
By the induction hypothesis x ( t ) ∈ I 1 ∩ K 3 h/ 4 and ∥ Φ( x ( t ))
− Φ( x ( t DEC ))
∥
=
O ( ρ 2 ln(1 /ρ )) .
Lemma C.41.
Under condition of Theorem C.31, assuming assuming t satisfy that x ( t ) ∈
I 1 ∩
K 3 h/ 4 , then we have that
[Equation]
Proof.
As x ( t ) ∈ I 1 , ∥
˜ x ( t ) ∥ 2 ≤ ζηρ + O ( ηρ 2 ) .
As ∥ ˜ x ( t ) ∥ 2 = O ( ρ ) , we have x ( t ) x
( t + 1) ⊂ K h and Φ( x ( t ))
Φ( x ( t + 1))
⊂ K r .
We will begin with a quantization technique separating [ M ] into disjoint continuous subset S 1 , ..., S p such that ∀ i ̸
= j ,
[Equation]
By Lemmas C.37 and C.65, we have that for any n ∈ [ M ] ,
[Equation]
This implies
[Equation]
∈ S
j | λ k ( t + 1) − λ l ( t + 1) | ≥ ρ − O ( ηρ 2 ) ≥ 0 .
99 ρ .
Define
P ( t ) S ( i ) ≜ (cid:88) k ∈ S i v n ( t ) v n ( t ) T .
By Theorem C.67, for any k , ∥ P ( t ) S k − P ( t +1)
[Equation]
k ∥
[Equation]
O
[Equation]
∥∇ 2 L (Φ( x ( t ))
[Equation]
2 L (Φ( x ( t + 1)
[Equation]
ρ )
=
O ( ηρ ) .
By Lemma C.39, we have that
[Equation]
We will write x ′ ( t + 1) as shorthand of ˜ x
( t ) − ηA ( t )˜
x ( t ) − ηρA 2 ( t ) ˜ x ( t ) ∥
˜ x ( t ) ∥ .
Now we discuss by cases,
If (cid:113)(cid:80) pi =
j ∥ P ( t ) S ( i ) ˜ x ( t ) ∥ 2 > max k ∈ S
j λ 2 k ( t ) ηρ > µ 2 ηρ , by Lemma C. (cid:118)(cid:117)(cid:117)(cid:116) p (cid:88) i = j ∥ P ( t ) S ( i )
˜ x
( t + 1) ∥ 2 ≤ (cid:118)(cid:117)(cid:117)(cid:116) p (cid:88) i =
j ∥ P ( t ) S ( i ) x ′ ( t + 1) ∥ 2 + O ( ηρ 2 ) ≤ max { (cid:0) 1 − ηλ D ( t + 1) (cid:1)
∥ p (cid:88) i = j P ( t ) S ( i ) ˜ x
( t ) ∥ − ηρλ D ( t + 1) 2 ∥ (cid:80) pi =
j
[Equation]
) S ( i ) ˜ x ( t ) ∥
∥
˜ x ( t ) ∥ , η max k ∈ S j λ k
( t + 1) ∥ p (cid:88) i =
j
P ( t ) S ( i ) ˜ x ( t ) ∥} + O ( ηρ 2 ) ≤ max { (cid:0) 1 − ηµ (cid:1) ∥ p (cid:88) i =
j
P ( t ) S ( i ) ˜ x ( t ) ∥ − ηρµ 3 2 ζ , ηζ ∥ p (cid:88) i =
j
P ( t ) S ( i ) ˜ x ( t ) ∥} + O ( ηρ 2 ) .
This further implies
[Equation]
If
[Table]
pi =
j ∥ P
[Equation]
t )
[Equation]
i ) ˜
[Equation]
t
[Equation]
[Table]
[Equation]
≤ [Table] k ∈ S
j λ 2 k
[Table]
t ) ηρ , then by Lemma C.21, we have that ∥
[Equation]
i =
j
[Equation]
Hence we have that
[Equation]
k
[Equation]
S
[Equation]
This further implies
[Equation]
Finally taking into quantization error, as all the eigenvalue in the same group at most differ Dρ , for any i ∈ S j , we have that − λ 2 i ( t + 1) + max k ∈ S
j λ 2 k ( t + 1) ≤ 2 Dζρ + D 2 ρ 2 .
Hence the previous discussion concludes as
If R k ( x ( t )) ≥R
k ( x ( t + 1)) + λ 2 k ( t + 1) ηρ ≤ (1 − ηµ )( R k ( x ( t )) + λ 2 k ( t ) .ηρ ) 2.
If R k ( x ( t )) <
R k
( x ( t + 1))
≤ O ( ηρ 2 ) .
Lemma C.42.
Under condition of Theorem C.31, assuming there exists t DEC2 such that for any t satisfying t DEC2 ≤ t ≤ t DEC2 + Θ(ln(1 /η ) /η ) , we have that x ( t ) ∈
I 1 ∩ K 3 h/ 4 .
Then there exists t INV =
t DEC2 + O (ln(1 /η ) /η ))
such that for any t satisfying t INV ≤ t ≤ t INV + Θ(ln(1 /η ) /η ) , we have that
[Equation] [Equation]
Proof.
The proof is almost identical with Lemma C.40 replacing the first two iterative hypothesis to Lemma C.41 and is omitted here.
C.5. P HASE II (P ROOF OF T HEOREM C.32)
The goal of this section is to prove Theorem C.32.
Proof.
First by Lemma C.47, there exists t ALIGN such that ∇ L ( x ( t ALIGN )) aligns with v 1 ( t ALIGN ) with O ( ρ ) angle.
Then we can inductively use Lemma C.48 and theorem D.5 with p = ηρ 2 and q = η + ρ to show that x ( t ) tracks the Riemmannian gradient flow with approximation error O ( η ln(1 /ρ )) and hence maintains in K h and Lemma C.47 to show that x ( t ) keeps alignment with v 1 ( t ) for T 3 /ηρ 2 step.
Alignment to Top Eigenvector We will continue to use the notations introduced in Appendix C.5.1.
We further define
[Equation]
Here the constant in O depends on the constant in I j and will be made clear in Lemma C.44.
For s ∈ S , define next( t ) as the smallest integer greater than s in S .
Lemma C.43.
Under the condition of Theorem C.32, there exist constants C 1 , C 2 < 1 independent of η and ρ , if ∥ ˜ x 1 ( t ) ∥ 2 ≤ 12 (cid:16) ηλ 1 ( t ) 2 2 − ηλ 1 ( t ) + ηλ 2 ( t ) 2 2 − ηλ 2 ( t ) (cid:17) ρ and x ( t ) ∈
∩
j I j ∩ K 7 h/ 8 , then
[Equation]
Proof.
By Lemma C.39, if we write x ′ ( t + 1) as shorthand of ˜ x
( t ) − ηA ( t )˜
x ( t ) − ηρA 2 ( t ) ˜ x ( t ) ∥
˜ x ( t ) ∥ , then ∥
˜ x ( t + 1) − x ′ ( t + 1) ∥
=
O ( ηρ 2 ) .
Define I quad j as { x | R j ( x ) ≤ 0 } .
Then we can find a surrogate x sur ( t ) such that x sur ( t ) ∈
∩
j I quad j ∩ K h and ∥ x sur ( t ) − ˜ x
( t ) ∥ 2 =
O ( ηρ 2 ) .
We will write x ′ sur ( t + 1) as shorthand of x sur ( t ) − ηA ( t ) x sur ( t ) − ηρA 2 ( t ) x sur ( t ) ∥ x sur ( t ) ∥ .
Let
[Equation]
As h (1) < 1 , we can choose C 1 < 1 , such that h ( C 1 ) < 1 .
We can further choose C 2 = max { ( h ( C 1 ) + 1) / 2 , 1 − µ 2 3 ζ 2 } < 1 .
We will discuss by cases 1.
If
[Equation]
Then
[Equation]
In such case we have
[Equation]
Then we have ∥ x ′ sur ( t + 1) − x ′ ( t + 1) ∥
=
O ( ηρ 2 ) .
By Lemma C.24, we have that ∥
[Equation]
If
[Equation]
Then we have | − ηρλ 2 D + (1 − ηλ D ) ∥ ˜ x ( t ) ∥ 2 | ηρλ 21 − (1 − ηλ 1 ) ∥
˜ x ( t ) ∥≤ λ 21 − λ 2 D λ. | ηρλ 22 − (1 − ηλ 2 ) ∥ ˜ x ( t ) ∥ 2 | ηρλ 21 − (1 − ηλ 1 ) ∥ ˜ x ( t ) ∥ 2 ≤ λλ 21 .
By Lemma C.27, ∥ x ′ ( t + 1) ∥≤ ( ηρλ 21 − (1 − ηλ 1 ) ∥
˜ x ( t ) ∥ 2 ) (cid:115) ∥
˜ x 21 ( t ) ∥∥
˜ x ( t ) ∥
22 + (1 − ∥ ˜ x 21 ( t ) ∥∥
˜ x ( t ) ∥ 22 ) max { λ 21 − λ 2 D λ, λλ 21 } ≤ ( ηρλ 21 − (1 − ηλ 1 ) ∥
˜ x ( t ) ∥ 2 ) (cid:115) ∥
˜ x 21 ( t ) ∥∥
˜ x ( t ) ∥
22 + (1 − ∥ ˜ x 21 ( t ) ∥∥
˜ x ( t ) ∥ 22 ) max { ζ 2 − µζ, ( ζ − ∆)ζ} .
As ∥ ˜ x 1 ( t ) ∥ 2 ≤ 1(cid:18) ηλ2 − ηλ 1 + ηλ2 − ηλ(cid:19) ρ
For ∥ ˜ x ( t ) ∥ 2 ≥ ηλ 21 2 − ηλ 1 ρC 1 , ∥ ˜ x 1 ( t ) ∥∥
˜ x ( t ) ∥ 2 ≤  (cid:18) λ 2 2 (2 − ηλ 1 ) λ 21 (2 − ηλ 2 ) +(cid:19) /C
1 ≤ 1(cid:18) λ 2λ 21 +(cid:19) /C 1 ≤ 1 2 C(cid:18) ( ζ − ∆)ζ+(cid:19) .
After plugging in, we have that
[Equation]
This concludes the proof.
Lemma C.44.
Under the condition of Theorem C.32, assuming that (1) x ( t ) ∈ ∩ j I j , (2) t ∈
S ∩ K 15 h/ 16 then next( t ) is always well defined and next( t ) ≤ t + 2 .
Proof.
If t + 1 ̸∈ S , then
[Equation]
Following similar argument in Lemma C.21, we have that x ( t + 1) ∈ ∩
j I j .
By Lemma C.39, if we write x ′ ( t + 2) as shorthand of ˜ x ( t + 1) − ηA ( t + 1)˜ x
( t + 1) − ηρA 2 ( t + 1) ˜
x ( t +1) ∥
˜ x ( t +1) ∥ , then ∥
˜ x ( t + 2) − x ′ ( t + 2) ∥
=
O ( ηρ 2 ) .
Define I quad j as { x | R j ( x ) ≤ 0 } .
Then we can find a surrogate x sur ( t + 1) such that x sur ( t + 1) ∈
∩
j I quad j ∩ K h , and ∥ x sur ( t + 1) − ˜ x
( t + 1) ∥ 2 =
O ( ηρ 2 ) .
We will write x ′ sur ( t + 2) as shorthand of x sur ( t + 1) − ηA ( t + 1) x sur ( t + 1) − ηρA 2 ( t + 1) x sur ( t +1)
∥ x sur ( t +1) ∥ .
As ∥ ˜ x ( t + 1) ∥
=
Ω( ηρ ) , we have
[Equation]
Hence we have that ∥ ˜ x
( t +2)
− x ′ sur ( t +2) ∥ =
∥
˜ x ( t +2) − x ′ ( t +2)
∥
+ ∥ x ′ ( t +2) − x ′ sur ( t +2) ∥ =
O ( ηρ 2 )
Notice we have ∥ x sur ( t ) ∥ 2 ≥ ηλ 21 2 − ηλ 1 ρ for properly chosen function in the definition S , hence, by Lemma C.
∥ x ′ sur ( t + 1) ∥ 2 ≤ ηλ2 − ηλ 1 ρ.
This further implies t + 2 ∈ S .
Lemma C.45.
Under the condition of Theorem C.32, there exists constant C > 0 independent of η and ρ , assuming that (1) x ( t ) ∈
∩ j I j , (2) t ∈ S ∩ K 7 h/ 8 (3) Ω( ρ 2 ) ≤ ∥ ˜ x 1 ( t ) ∥ , then ∥ ˜ x 1 (next( t ))
∥ ≥ ∥
˜ x 1 ( t ) ∥ − O ( ηρ 2 ) .
Proof.
This is by standard approximation as in previous proof and Lemma C.28.
Lemma C.46.
Under the condition of Theorem C.32, there exists constant C > 0 independent of η and ρ , assuming that (1) x ( t ) ∈
∩ j I j , (2) t ∈ S ∩ K 7 h/ 8 (3) Ω( ρ 2 ) ≤ ∥ ˜ x 1 ( t ) ∥ ≤ 12 (cid:16) ηλ 21 2 − ηλ 1 + ηλ 22 2 − ηλ 2 (cid:17) ρ , then ∥
˜ x 1 (next( t ))
∥ ≥ min { (1 + Cη ) ∥
˜ x 1 ( t ) ∥ , 1(cid:18) ηλ2 − ηλ 1 + ηλ2 − ηλ(cid:19) ρ } , or ∥ ˜ x 1 (next(next( t )))
∥ ≥ min { (1 + Cη ) ∥
˜ x 1 ( t ) ∥ , 1(cid:18) ηλ2 − ηλ 1 + ηλ2 − ηλ(cid:19) ρ } .
Proof.
In this proof, we will sometime drop the t in λ k ( t ) or A ( t ) .
Applying Lemma C.44, we have next( t ) and next(next( t )) are well-defined.
We can suppose ∥ ˜ x 1 (next( t ))
∥ 2 ≤ 1 2 (cid:16) ηλ 21 2 − ηλ 1 + ηλ 22 2 − ηλ 2 (cid:17) , else the result holds already.
By assumption, we have ∥ ˜ x 1 ( t )
∥ ≥ Ω( ρ 2 ) .
Using Lemma C.39,
[Equation]
Denote
[Equation]
as the one step update of SAM on the quadratic approximation of the general loss.
Now using Lemma C.43 and the induction hypothesis, we have for some C 1 and C 2 smaller than 1 , ∥
[Equation]
We will discuss by cases, 1 If
[Equation]
ρ
If next( t ) =
t + 1 , then ∥
[Equation]
As we have ˜ x 1 ( t )
=
Ω( ρ 2 ) , we have x ′ 1 ( t +1) =
Ω( ρ 2 ) , then as ∥ ˜ x 1 ( t +1) − x ′ 1 ( t +1) ∥
[Equation]
O ( ηρ 2 ) , this implies ∥
[Equation]
[Equation]
˜ x 1 ( t ) ∥ .
If next( t ) =
t +2 , define x ′ ( t +2)
= x ′ ( t +1) − ηAx ′ ( t +1)
− ηρA 2 x ′ ( t +1)
∥ x ′ ( t +1) ∥ , as ∥ ˜ x 1 ( t +1) ∥
=
Ω( ηρ ) , by Lemma C.38, we have ∥ x ′ ( t + 2) − ˜
x
( t + 2) ∥
=
O ( ηρ 2 ) .
∥ x ′ ( t +
[Equation]
Combining with | ˜ x 1 ( t ) | ≥ Ω( ρ 2 ) , we have that ∥ ˜
[Equation]
1 (next( t ))
∥ ≥ (1 + Cη ) ∥
˜ x 1 ( t ) ∥ 2 Case 2 ∥
˜ x ( t ) ∥ > C 1 ηλ 21 2 − ηλ 1 ρ , then ∥
˜ x ( t + 1) ∥
≤ C 2 ηλ 21 2 − ηλ 1 ρ , next( t ) =
t +
By Lemma C.45, ∥ ˜ x 1 ( t + 1) ∥ ≥ (1 − Cη ) ∥ ˜ x 1 ( t ) ∥ .
As ∥ ˜ x (next( t ))
∥ ≤ C 2 ηλ 21 2 − ηλ 1 , similar to the first case, ∥
˜ x 1 (next(next( t )))
∥ ≥ (1 + 4 Cη ) ∥
˜ x 1 (next( t ))
∥ ≥ (1 + Cη ) ∥
˜ x 1 ( t ) ∥ .
In conclusion, if ∥ ˜ x 1 ( t ) ∥ ≤ 12 (cid:16) ηλ 21 2 − ηλ 1 + ηλ 22 2 − ηλ 2 (cid:17) ρ , we would have there exists C >∥ ˜ x 1 (next( t ))
∥ ≥ (1 + Cη ) ∥
˜ x 1 ( t ) ∥ or ∥ ˜ x 1 (next(next( t )))
∥ ≥ (1 + Cη ) ∥
˜ x 1 ( t ) ∥ .
Lemma C.47.
Under the condition of Theorem C.32, there exists constant T 2 > 0 independent of η and ρ , we would have that when t = t ALIGN =
⌈ T 2 ln(1 /ρ )
/η ⌉ ,
[Equation]
Further if x ( t ′ ) ∈ K h holds for t ′ = 0 , 1 , ..., t LOCAL , then for t satisfying t ALIGN ≤ t ≤ t LOCAL |⟨ x ( t
[Equation]
− Φ( x ( t
[Equation]
[Equation]
[Equation]
x
[Equation]
[Equation]
=
Θ( ηρ
[Equation]
[Equation]
Proof.
Let C be the constant defined in Lemma C.46.
Suppose WLOG ∥
˜ x 1 (0)
∥ ≥ ρ 2 .
Define
[Equation]
We will choose t ALIGNMID as the minimal t ∈ S , such that ∥ ˜ x 1 ( t ) ∥ [Equation] [Equation]
[Equation]
ρ .
Then by induction and Lemmas C.45 and C.46, we easily have that for t ≤ min { C 2 +1 , t ALIGNMID } and t ∈ S , we have that
[Equation]
or ∥ ˜ x 1 (next( t ))
∥ ≥ min { (1 + Cη [Equation] ∥
[Equation]
The detailed induction is analogous to previous inductive argument and is omitted.
If t ALIGNMID ≥
C 1 , then we have for the minimal t ≥ C 1 and t ∈ S
[Equation]
This is a contradiction and we have that t ALIGNMID ≤ C 1 .
As we have ∥ ˜ x 1 (next( t ))
∥ ≥ ∥
˜ x 1 ( t ) ∥ − O ( ηρ 2 ) for ∥ ˜ x 1 ( t )
∥ ≥ 12 (cid:16)
[Equation]
for C 2 ≥ t ≥ t ALIGNMID .
We will then show that for t ≥ t ALIGNMID + C 1 iteration, ∥ P (2: D )
¯ x ( t + 1) ∥ ≤ O ( ηρ 2 ) .
For C 1 ≥ t ≥ t ALIGNMID ,
[Equation]
Notice that,
[Equation]
Hence,
[Equation]
Now by Lemma C.65 and Theorem C.67,
[Equation]
By Lemma C.39,
[Equation]
Combining,
[Equation]
Hence when t = t ALIGN =
t ALIGNMID + C 2 ,
[Equation]
By x ( t ) ∈ I 1 , we easily have ∥ ˜ x 1 ( t ) ∥
=
O ( ηρ ) .
Hence we conclude that
[Equation]
The second claim is just another induction similar to previous steps and is omitted as well.
Tracking Riemannian Gradient Flow We are now ready to show that Φ( x ( t )) will track the solution of Equation (6).
The main principal of this proof has been introduced in Section 4.3.
Lemma C.48.
Under the condition of Theorem C.32, for any t satisfying that
[Equation]
it holds that
[Equation]
Proof.
To begin with, we can approximate Φ( x ( t +1))
− Φ( x ( t )) by its first order Taylor Expansion, by Lemma C.12,
[Equation]
Then by plugging in the update rule and another Taylor Expansion,
[Equation]
Using Lemma C. [Equation] we have
[Equation]
Putting together [Equation] we have that
[Equation]
As we have ∥ ˜ x ( t ) ∥
=
Θ( ηρ ) , hence by Lemmas C.8 and C.36,
[Equation]
Finally [Equation] we have that
[Equation]
as the angle between ∇ L ( x ) ∥∇ L ( x ) ∥ and v 1 ( t ) is O ( ρ ) .
By Lemma C. [Equation] it holds that
[Equation]
Putting together we have that,
[Equation]
It completes the proof.
C.5. P ROOFS OF C OROLLARIES 4.10 AND 4.
Proof of Corollary 4.10.
We will do a Taylor expansion on L Max ρ .
By Theorem C.31 and C.32, we have ∥ x
[Equation]
and
[Equation] x ( ⌈ T 3 /ηρ 2 ⌉ )))
∥ 2 =
O ( ηρ ) .
For convenience, we denote x ( ⌈ T 3 /ηρ 2 ⌉ ) by x .
[Equation]
Since max ∥ v ∥ 2 ≤ 1 ∥ v T ∇ L ( x ) ∥ 2 = O ( ∥ x − Φ( x ) ∥ 2 ) =
O ( ηρ ) , it holds that
[Equation]
which completes the proof.
Proof of Corollary 4.11.
We choose T such that X ( T ϵ ) is sufficiently close to X ( ∞ ) , such that λ 1 ( X ( T ϵ )) ≤ λ 1 ( X ( ∞ ))
+ ϵ/ 2 .
By Corollary 4.10 (let T 3 = T ϵ ), we have that for all ρ, η such that η ln(1 /ρ ) and ρ/η are sufficiently small, ∥ R Max ρ ( x ( ⌈ T ϵ / ( ηρ 2 ) ⌉ ))
− ρ 2 λ 1 ( X ( T ϵ )) / 2 ∥ ≤ ˜ o
This further implies ∥ R Max ρ ( x ( ⌈ T ϵ / ( ηρ 2 ) ⌉ ))
− ρ 2 λ 1 ( X ( ∞ )) / 2 ∥ ≤ ϵρ 2 + o (1) .
We also have
L ( x ( ⌈ T ϵ / ( ηρ 2 ) ⌉ )) − inf x ∈ U ′
L ( x ) =
o (1) .
Then we can leverage Theorem 4.3 and Theorem 4.5 to get the desired bound.
C. A NALYSIS FOR 1-SAM (P ROOF OF T HEOREM 5.4)
The goal of this section is to prove the following theorem.
Theorem 5.4.
Let { x ( t ) } be the iterates defined by SAM (16) and x (0) =
x init ∈ U , then under
Assumption 3.1, for all η and ρ such that ( η + ρ ) log(1 /ηρ ) is sufficiently small, with probability at least 1 − O ( √ ρ + √ η ) , the dynamics of SAM can be split into two phases:
• Phase I: Stochastic SAM follows Gradient Flow with respect to L until entering an O ( ηρ ) neighborhood of the manifold Γ in ˜ O
( 1 η ) steps; • Phase II: Stochastic SAM tracks the solution X of (17) , the Riemannian Gradient Flow with respect to Tr( ∇ 2 L ( · )) , with time scaling as ∥ Φ( x ( ⌈ T 3 / ( ηρ 2 ) ⌉ ))
− X ( T 3 ) ∥ =
O ( η 1 / 4 + ρ 1 / 4 ) .
As mentioned in our proof setups in Appendix C, we will prove Theorem 5.4 under a more general (and weaker) condition, namely Condition C.2 and Assumption 3.1.
The only usage of Setting C. in the proof is Theorem C.49 and lemma 5.3.
Theorem C.49.
Condition C.2 holds under Setting C.1 and assumption 3.1. Proof of Theorem C.49. L = 1 M (cid:80) Mk =1 L k by definition. ∀ k ∈
M
L k
( x ) = ℓ ( f k ( x ) , y k )
is C 4 as ℓ and
f k
are both C
For any x ∈ Γ , by Lemma 5.3, we have ∇ f k ( x ) ̸ = 0 . Then there exists an open neighborhood V k such that Γ ⊂ V k and ∇ f k ( x ) ̸ = 0 for any x ∈ V k , k ∈ [ M ] . Then applying implicit function theorem, for any k ∈ M there exists a ( D − 1) -dimensional
C 4 -manifold Γ ′ k ⊂ V k , such that for any x ′ ∈ V , f k ( x ′ ) = y k if and only if x ′ ∈ Γ ′ k .
As for any x ∈ Γ ⊂ V k , f k ( x ′ ) = y k , we can infer that Γ ⊂ Γ ′ k . Then Γ ⊂ ∪ Mk =1 Γ k . For any x ∈ Γ k , we have f k ( x ) =
y k , which implies L k ( x ) = 0 . Also as x ∈ V , ∇ f k ( x ) ̸ = 0 . By Lemma C.64, we have rank( ∇ 2 L ( x )) = 1 .
Analogous to the full-batch setting, we will split the trajectory into two phase.
Theorem C.50 (Phase I) .
Let { x ( t ) } be the iterates defined by SAM (Equation (16) ) and x (0) =
x init ∈ U , then under Assumption 3.1 and C.2, for almost every x init , there exists a constant T 1 , it holds for sufficiently small ( η + ρ ) ln 1 /ηρ , we have with probability 1 − O ( √ ρ ) , min T 1 ln(1 /ηρ )
/η ≥ t ∥ x ( t ) − Φ( x ( t ))
∥ 2 = O ( ηρ ) and max T 1 ln(1 /ηρ )
/η ≥ t ∥
Φ( x init ) − Φ( x ( t ))
∥ 2 =
O (( η + ρ ) ln 1 /ηρ ) .
Theorem C.50 shows that SAM will converges to an ˜ O ( ηρ ) neighborhood of the manifold without getting far away from Φ( x (0)) , where we can perform a local analysis on the trajectory of Φ( x ( t )) .
Under Assumptions 3.1 and C.2, we have Tr( ∇ 2 L k ( x ))
=
λ 1 ( ∇ 2 L k ( x )) is differentiable for x ∈ Γ i .
Hence Tr( ∇ 2 L ( x )) =
(cid:80) Mk =1 Tr( ∇ 2 L k ( x )) is also differentiable and we have (17) is well defined for some finite time T 2 .
Theorem C.51 (Phase II) .
Let { x ( t ) } be the iterates defined by SAM (Equation (16) ) under Assumptions 3.1 and C.2, assuming (1) ∥ x (0) − Φ( x (0))
∥ 2 =
O ( ηρ ) and (2) ∥
Φ( x init ) − Φ( x (0))
∥ 2 =
O (( η + ρ ) ln 1 / ( ηρ )) , then for almost every x (0) , for any T 2 > 0 till which solution of (17) X exists, for sufficiently small ( η + ρ ) ln 1 / ( ηρ ) , we have with probability 1 − O ( ηρ ) , for all ηρ 2 t < T 2 , ∥ Φ( x ( t ))
− X ( ηρ 2 t ) ∥ 2 = O (( η + ρ ) ln 1 /ηρ ) and ∥ x ( t ) − Φ( x ( t ))
∥
= O ( ηρ ln 1 /ηρ ) .
Combining Theorems C.49 to C.51, the proof of Theorem 5.4 is clear.
Now we recall our notations for stochastic setting with batch size one.
Notations for Stochastic Setting: Since L k is rank- 1 on Γ k for each k ∈ [ M ] , we can write it as L k ( x ) =
Λ k ( x ) w k ( x ) w ⊤ k ( x ) for any x ∈ Γ k , where w k is a continuous function on Γ k with pointwise unit norm.
Given the loss function L k , its gradient flow is denoted by mapping ϕ k : R D × [0 , ∞ ) → R D .
Here, ϕ k ( x, τ ) denotes the iterate at time τ of a gradient flow starting at x and is defined as the unique solution of ϕ k
( x, τ ) =
x − (cid:82) τ 0 ∇ L k ( ϕ k ( x, t )) dt , ∀ x ∈ R D .
We further define the limiting map Φ k as Φ k ( x ) =
lim
τ →∞ ϕ k ( x, τ ) , that is, Φ k ( x ) denotes the convergent point of the gradient flow starting from x .
Similar to Definition 3.2, we define U k
= { x ∈ R D | Φ( x ) exists and Φ k ( x ) ∈
Γ k } be the attraction set of Γ i .
We have that each U k is open and Φ k is C 2 on U k by Lemma B.15 in Arora et al.
In this section we will define K as { X ( t ) | t ∈ [0 , T 3 ] } where X is the solution of (17).
We will assume the update is always well defined.
C.6. P HASE I (P ROOF OF T HEOREM C.50)
The goal of this section is to prove Theorem C.50.
Proof.
The proof consists of two major steps.
Tracking Gradient Flow Lemma C.52 shows the existence of step t GF =
O (1 /η ) such that x ( t GF ) is in a subset of K h and Φ( x ( t GF ))
is O
( η + ρ ) close to Φ( x init ) .
Decreasing Loss Lemmas C.56 and C.61 shows existence of step t DEC2 = O (ln(1 /ηρ ) /η )
such that x ( t DEC2 ) is in a O ( ηρ ) neighbor of K h and Φ( x ( t DEC2 )) is O
(( η + ρ ) ln(1 /ηρ ))
close to Φ( x init ) .
Tracking Gradient Flow
Lemma C.52.
Under condition of Theorem C.50, with probability 1 − O ( √ ρ + √ η ) , there exists t GF =
O (1 /η ) , such that the iterate x ( t GF ) is O (1) close to the manifold Γ and Φ( x ( t GF )) is
O ( η + ρ ) is close to Φ( x init ) .
Quantitatively,
[Equation]
Proof.
Choose C = 14 (cid:113) µζ .
There exists T > 0 , such that
[Equation]
Consider
[Equation]
By Theorem D.5, by choosing p = η and q = η + ρ , for sufficiently small η and ρ , the iterates x ( t ) tracks gradient flow ϕ ( x init , T ) in O (1 /η ) steps in expectation, Quantitatively, with probability 1 − √ η − √ ρ , for t GF =
⌈ T 0 η ⌉ , we have that
[Equation]
This implies x ( t GF
) ∈ K h , hence by Taylor Expansion on Φ ,
[Equation]
This implies
[Equation]
By Taylor Expansion,
[Equation]
Decreasing Loss
Lemma C.53.
Under condition of Theorem C.50, assuming x ( t 0 )
∈
K h/ 4 and for any t satisfying t 0 ≤ t ≤ t
+ O (ln(1 /ηρ ) /η ) ,
[Equation]
16 , it holds that
[Equation]
Moreover, we have that
[Equation]
Proof.
We will prove by induction.
For τ = t 0 , the result holds trivially.
Suppose the result holds for t − 1 , then for any τ satisfying t 0 ≤ τ ≤ t − 1 , by Lemmas C.7 and C.13, ∥ [Equation] = O
[Equation]
+
[Equation]
Also by Lemma C.7, ∥ x ( t ) − Φ( x ( t ))
∥ 2 ≤ h/ 2 √ 2 , this implies, dist
[Equation]
Lemma C.54.
Under condition of Theorem C.50, if x ( τ ) ∈ K h , then we have that
[Equation]
Moreover it holds that,
[Equation]
Proof.
By Lemma C.13 and Taylor Expansion,
[Equation]
Lemma C.55.
Under condition of Theorem C.50, assuming x ( t 0 )
∈
K h/ 4 and L ( x ( t 0 )) ≤ µh 2 32 , then with probability 1 − O ( ρ ) [Equation] for any t satisfying t 0 ≤ t ≤ t [-
-]0 + O (ln(1 /ηρ ) /η )
[Equation] it holds that x ( t ) ∈ K h .
Moreover, we have that
[Equation]
t 0 ))
∥
=
O (( η + ρ ) ln(1 /ηρ )) .
Proof.
By Uniform Bound and Lemma C.53,
[Equation]
and L ( x ( τ [Equation] ) ≤ µh 2 16 , ∀ t 0 ≤ τ ≤ t − 1) [Equation]
[Equation]
and x ( τ ) ∈ K h , ∀ t 0 ≤ τ ≤ t − 1)
Consider each term, and applying uniform bound again,
[Equation]
and x ( τ ) ∈ K h , ∀ t 0 ≤ τ ≤ t
[Equation]
[Equation]
and L ( x ( τ )) ≤ µh 2and ∀ t − 1 ≥ τ ′ ≥ τ + 1 , µh 2 16 > L ( x ( τ ′ ))
> µh 2and ∀ t − 1 ≥ τ ′′ ≥ τ, x ( τ ′′ ) ∈ K h ) .
Then if we consider each term, we have that it is bounded by
[Equation]
and ∀ t − 1 ≥ τ ′ ≥ τ
[Equation]
1 , L ( x ( τ ′ )) > µh 2and ∀ t − 1 ≥ τ ′′ ≥ τ, x ( τ ′′ ) ∈ K h | L ( x ( τ )) ≤ µh 2 32 ) .
Define a coupled process ˜ L ( τ + 1) =
ln L ( x ( τ + 1)) and ˜ L ( τ ′ )
=
(cid:40)
ln L ( x ( τ ′ )) , if ˜ L ( τ ′ − 1) =
ln L ( x ( τ ′ − 1)) ≥ ln( µh 2 32 ) , ˜ L ( τ ′ − 1) − ηµ/ 2 , if otherwise.
Then clearly
P ( L ( x ( t )) ≥ µh 2and ∀ t ≥ τ ′ ≥ τ + 1 , L ( x ( τ ′ ))
> µh 2and ∀ t ≥ τ ′′ ≥ τ, x ( τ ′′ ) ∈ K h | L ( x ( τ )) ≤ µh 2 32 ) ≤ P (˜ L ( t ) ≥ ln( µh 2 16 )) .
Consider a fixed τ ′ satisfying τ + 1 ≤ τ ′ ≤ t .
By Lemma C.54, we have that
[Equation]
Hence ˜ L ( t ) + ηµt/ 2 is a super martingale.
Further it holds that if L ( x ( τ ′ − 1)) ≥ ( µh 2 32 ) , then [Equation] =
[Equation]
Using the smoothness at log( x ) at µh 2 32 which is a positive constant,
[Equation]
Here C is a constant independent of η .
This implies ˜ L ( x ( τ + 1))
≤ µh 2 16 √
Now by Azuma-Hoeffding bound (Lemma C.68), we have that
[Equation]
With a = ln( µh 2 16˜ L ( τ +1) )
+ ( t −
[Equation] [Equation]
2 ≥ (ln
[Equation]
+
[Equation]
t
[Equation]
[Equation]
) / 2 , we have that
[Equation]
Hence we have
[Equation]
Hence with probability 1 − ρ , L ( x ( t )) ≤ µh 2 16 , ∀ t 0 ≤ t ≤ t 0 + O (ln(1 /ηρ ) /η ) , combining with Lemma C.53, we have completed our proof.
Lemma C.56.
Under condition of Theorem C.50, assuming there exists t GF such that L ( x ( t GF )) ≤ µh 2 32 and x ( t GF ) ∈
K h/ 4 , then with probability 1 − O ( ρ ) , there exists t DEC = t GF + O
(ln(1 /ρ )
/η ) , such that x ( t DEC ) is in O ( ρ ) neighbor of Γ , quantitatively, we have that
[Equation]
Moreover the movement of the projection of Φ( x ( · )) on the manifold is bounded,
[Equation]
Proof.
For simplicity of writing, define T 1 ≜ ⌈ 2ln h 2 256 ρ 3 µ ηµ [Equation] O
[Equation]
By Lemma C.55, we may assume x ( t ) ∈ K h for t GF ≤ t ≤ T 1 + t GF .
Define indicator function as
[Equation]
[Equation]
By Lemma C.54, we have that,
[Equation]
We can then conclude that with T 2 = T 1 + t GF , using Lemma C.8,
[Equation]
We have
[Equation]
This implies A ( T 2 + 1) = 0 with probability 1 − O ( ρ ) , which indicates the existence of t DEC .
The second claim is a direct application of Lemma C.55.
Lemma C.57 (A general version of Lemma 5.3) .
Under Assumption 3.1 and Condition C.2, for x ∈ K h and p ∈ C, ∇ 2 L k ( p ) =
Λ k ( p ) w k ( p ) w k ( p ) ⊤ , there exists s ∈ { 1 , − 1 } ,
[Equation]
k
[Equation]
Further if | w ⊤ k ( x − p ) | ≥ ∥ x − p ∥ 3 / 2 2 , then s = sign( w ⊤ k ( x − p )) .
This implies
[Equation]
Proof.
We will calculate the direction of ∇ L k ( x ) ∥∇ L k ( x ) ∥ using two different approximations and compare them to get our result.
According to Lemma C.10,
[Equation]
Suppose ∇ 2 L k (Φ k ( x ))
=
Λ k
(Φ k ( x ))
w k (Φ k ( x ))
w k (Φ k ( x )) ⊤ , then ∇ L
[Equation]
( x
[Equation]
As ∇ 2 L k ( p ) = Λ k ( p ) w k ( p ) w k ( p ) ⊤ , using Davis-Kahan Theorem C.67, we would have ∃ s ∈ {− 1 , 1 } , such that ∥ w k (Φ k ( x )) − sw k ( p ) ∥ 2 ≤ ζ ∥ Φ k
[Equation]
According to Lemma C.7, we have ∥ x − Φ k
( x ) ∥ 2 ≤ ∥∇ L k ( x ) [Equation] [Equation] ≤ ζ
[Equation]
x
[Equation]
This implies,
[Equation]
Equation (33) is our first statement.
By Taylor expansion at p ,
[Equation]
That being said, when | w ⊤ k ( x − p ) | ≥ ∥ x − p ∥ 3 / 2 2 , we have ∥∇ [Equation] k
[Equation]
x
[Equation]
− Λ
[Equation]
k w ⊤ k ( x − p )
∥ 2 ≤ O ( ∥ x − p ∥ 22 ) .
∥∇ L k ( x ) ∥ ≥ ∥ Λ k
[Equation]
k
[Equation]
⊤ k
[Equation]
p )
[Equation]
[Equation]
[Equation]
x
[Equation]
p ∥
[Equation]
[Equation]
Concluding,
[Equation]
Hence we have ∇ L k
[Equation]
x ) [Equation] Comparing (33) and (34), we have s = sign( w k ( p ) ⊤ ( x − p )) when | w ⊤ k ( x − p ) | ≥ ∥ x − p ∥ 3 / 2 2 .
Lemma C.58.
Under condition of Theorem C.50, for any constant C > 0 independent of η, ρ , there exists constant C 1 > C 2 > 0 independent of η, ρ , if x ( t ) ∈ K h and C 1 ηρ ≤
∥ x ( t ) − Φ( x ( t ))
∥ ≤
Cρ , then we have that
[Equation]
Proof.
By Lemma C.8, ∥ x ( t ) − Φ( x ( t ))
∥
=
O ( ρ ) .
Hence we have that by Taylor Expansion,
[Equation]
Here Λ k , w k indicates Λ k (Φ( x ( t ))) , w k (Φ( x ( t ))) .
Notice that given ∥ x ( t ) − Φ( x ( t ))
∥
=
O ( ρ ) , by Lemma C.13, we have that ∥ Φ
[Equation]
x ( t + 1)
[Equation]
This implies x ( t + 1) ∈ K r .
Further by Taylor Expansion, ∇ L k ( x ( t ))
=
Λ k w k w ⊤ k
( x ( t ) − Φ( x ( t )))
+ O ( ρ 2 ) .
By Lemma C.57, we have for some s k ( t ) ∈ {− 1 , 1 } .
[Equation]
We also have
[Equation]
Concluding,
[Equation]
After we take square and expectation,
[Equation]
We will then carefully examine each positive term, 2 η
[Equation]
This implies,
[Equation]
O ( ηρ 2 ∥ x ( t ) − Φ( x ( t ))
∥ + η 2 ρ 2 )
[Equation]
We will now lower bound (cid:80) Mk =1 Λ k s k ( t ) w ⊤ k
( x ( t ) − Φ( x ( t ))) .
By Equation (35),
M (cid:88)
[Equation]
[Equation]
For
[Equation]
[Equation]
Concluding, we have that
M (cid:88)
[Equation]
So
[Equation]
The inequality holds if ∥ x ( t ) − Φ( x ( t ))
∥ 2 > C 1 ηρ .
Finally by Jenson’s Inequality,
[Equation]
Lemma C.59.
Under condition of Theorem C.50, for any constant C > 0 independent of η, ρ , there exists constant C 3 > 0 independent of η, ρ , if x ( t ) ∈ K h and ∥ x ( t ) − Φ( x ( t ))
∥ ≤ Cρ , then we have that
[Equation]
Proof.
This is a direct application of Lemma C.13.
Lemma C.60.
Under condition of Theorem C.50, assuming x ( t 0 )
∈
K h/ 2 and ∥ x ( t 0 ) − Φ( x ( t 0 ))
∥
≤ f ( η, ρ ) for some fixed function f and f ( η, ρ ) ∈ Ω( ηρ ln 2 (1 /ηρ ))
∩ O ( ρ ) , then with probability 1 − O ( ρ ) , for any t satisfying t 0 ≤ t ≤ t
+ O (ln(1 /ηρ ) /η ) , it holds that ∥ x ( t ) − Φ( x ( t ) ∥ ≤ 2 f ( η, ρ ) .
Moreover, we have that
[Equation]
Proof.
By Lemma C.55, we have
that x ( t ) ∈
K h for any t satisfying that t 0 ≤ t ≤ t 0 +
O (ln(1 /ηρ ) /η ) and with probability 1 − O ( ρ ) we will suppose this hold for the following deduction.
By Uniform Bound,
[Equation]
Consider each term and apply Uniform bound again,
[Equation]
ρ ))
[Equation]
∥ x
[Equation]
τ ) − Φ( x ( τ )) ∥ ≤ f ( η, ρ )
[Equation]
and f ( η, ρ ) ≤ ∥ x ( τ ′ ) − Φ( x ( τ
[Equation]
∥
[Equation]
[Equation]
τ
[Equation]
t
[Equation]
[Equation]
Then if we consider each term, it is bounded by
[Equation]
Now let C be the positive constant satisfying 2 f ( η, ρ ) ≤ Cρ , suppose C 1 , C 2 are the constants corresponds to C in Lemma C.58 and C 3 is the constant correspond to C in Lemma C.59.
By definition C 3 > C 2 .
Define a coupled process ˜
y ( τ + 1) =
y ( τ + 1) and
[Equation]
if otherwise .
Now clearly Equation (36) is bounded by P (˜ y ( t ) ≥ 2 f ( η, ρ )) .
As E [˜ y ( τ ′ )] ≤ ˜ y ( τ ′ − 1) − C 2 ηρ by Lemma C.58 and ∥ ˜ y ( τ ′ ) − ˜ y
( τ ′ − 1)
∥ ≤ C 3 ηρ by Lemma C.59.
This implies ∥ ˜
y ( τ ′ )
∥ − C 2 ηρτ ′ is a super martingale.
By Azuma-Hoeffding bound(Lemma C.68), we have
[Equation]
Choosing h
= C 2 ηρ ( t − τ − 1) − ∥
x ( τ + 1) − Φ( x ( τ + 1)) ∥ + 2 f ( η, ρ )
[Equation]
˜
[Equation]
We then have
[Equation]
Lemma C.61.
Under condition of Theorem C.50, assuming there exists t DEC such that x ( t DEC ) ∈
K h/ 2 and ∥∇ L ( x ( t DEC ))
∥ ≤ 4 ζρ , then there exists t DEC2 = t DEC + O (ln(1 /ηρ ) /η ) , such that ∥ x ( t DEC2 ) − Φ( t DEC2 ) ∥
≤ O ( ηρ ) .
Furthermore, for any t satisfying t DEC2 ≤ t ≤ t DEC2 + Θ(ln(1 /ηρ ) /η ) , we have that ∥ Φ( x ( t ))
− Φ( x ( t DEC ))
∥
=
O ( ρ 2 ln(1 /ηρ )) .
Proof.
We have
that x ( t ) ∈
K h (Lemma C.55) and ∥ x ( t ) − Φ( x ( t ))
∥ ≤ Cρ for some constant
C (Lemma C.60) for any t satisfying that t DEC ≤ t ≤ t DEC + O (ln(1 /ηρ ) /η )
with probability 1 − O ( ρ ) and we will suppose this holds for the following deduction.
The second statement then follows directly from Lemma C.13.
Let C 1 , C 2 be the constant in Lemma C.58 corresponding to C , For simplicity of writing, define
T
[Equation]
Define indicator function as
[Equation]
By Lemma C.58, we have that,
[Equation]
We can then conclude that with T 2 = T 1 + t DEC , using Lemma C.8,
[Equation]
This implies A ( T 2 + 1) = 0 with probability 1 − O ( ρ ) , which indicates the existence of t DEC2 .
C.6. P HASE II (P ROOF OF T HEOREM C.51)
The goal of this section is to prove Theorem C.50.
Proof.
We can inductively use Lemma C.63 and theorem D.5 to show that x ( t ) tracks the Riemmannian gradient flow and maintains in K h with probability 1 − O ( √ η + √ ρ ) and Lemma C.62 to show that x ( t ) is in an ˜ O ( ηρ ) neighbor of Γ for T 3 /ηρ 2 step.
Convergence Near Manifold
Lemma C.62.
Under condition of Theorem C.51, assuming x ( t ) ∈ K h , ∀ t 0 ≤ t ≤ t
+ O (1 /ηρ 2 ) and ∥ x ( t 0 ) − Φ( x ( t 0 ))
∥
≤ f ( η, ρ ) for some fixed function f and f ( η, ρ ) ∈ Ω( ηρ ln 2 (1 /ηρ ))
∩ O ( ρ ) , then with probability 1 − O ( ηρ ) , for any t satisfying t 0 ≤ t ≤ t
+ O (1 /ηρ 2 ) , it holds that ∥ x ( t ) − Φ( x ( t ) ∥ ≤ 2 f ( η, ρ ) .
Proof.
The proof is almost identical to Lemma C.60 and is omitted.
Tracking Riemannian Gradient Flow
Lemma C.63.
Under the condition of Theorem C.51, for any t satisfying that x ( t ) ∈ K h and ∥ x ( [Equation] [Equation] ( x (
[Equation]
( ηρ
[Equation]
2 (1 /ηρ )) .
It holds that
[Equation]
Proof.
By Taylor Expansion,
[Equation]
) ∇ L k ( x ( t ))
∥∇ L k ( x ( t ))
∥ − ηρ 2 ∂ 2 ( ∇ L k )[ ∇ L k ( x ( t ))
∥∇ L k ( x ( t ))
∥ , ∇ L k ( x ( t )) ∥∇ L k ( x ( t ))
∥ ] / 2 + O ( ηρ 3 ) .
Now as ∥ x ( t ) − Φ( x ( t ))
∥ 2 = ˜ O
( ηρ ) , by Lemma C.13, it implies
[Equation]
Then we have ∥
[Equation]
Using Lemma C.11, we have
[Equation]
Hence
[Equation]
Notice finally that by Lemma C.
[Equation]
Hence we have
[Equation]
This completes the proof.
C.6. P ROOFS OF C OROLLARIES 5.5 AND 5.
Proof of Corollary 5.5.
We will do Taylor expansion on E k
[ L Max k,ρ ]( x ) .
By Theorem C.50 and C.51, we have ∥ x ( ⌈ T 3 /ηρ 2 ⌉ ) − X ( T 3 ) ∥ [-
-]2 =
O ( η 1 / 4 + ρ 1 / 4 ) and ∥ Φ( x ( ⌈ T 3 /ηρ 2 ⌉ )) − x ( ⌈ T 3 /ηρ 2 ⌉ )
∥ 2 =
O ( η 1 / 4 + ρ 1 / 4 ) .
For convenience, we denote x ( ⌈ T 3 /ηρ 2 ⌉ ) by x .
[Equation]
Since max ∥ v ∥≤ 1 | v ⊤ ∇ L k ( x ) | =
O ( ∥ x − Φ( x ) ∥ )
=
O ( η 1 / 4 + ρ 1 / 4 ) , it holds that,
[Equation]
Proof of Corollary 5.6.
We choose T ϵ such that X
( T ϵ ) is sufficiently close to X ( ∞ ) , such that
Tr( X ( T ϵ )) ≤
Tr( X ( ∞ ))
+ ϵ/ 2 .
By corollary 5.5 (let T 3 = T ϵ ), we have for all ρ, η such that ( η + ρ ) ln(1 /ηρ ) is sufficiently small, ∥ E k [ R Max k,ρ ]( x ( ⌈ T ϵ / ( ηρ 2 ) ⌉ ))
− ρ 2 Tr( X ( T )) / 2 ∥ 2 ≤ o (1) .
This further implies ∥ E k
[ R Max k,ρ ]( x ( ⌈ T ϵ / ( ηρ 2 ) ⌉ ))
− ρ 2 Tr( X ( ∞ )) / 2 ∥ 2 ≤ ϵρ 2 / 2 + o (1) .
We also have L ( x ( ⌈ T ϵ / ( ηρ 2 ) ⌉ )) − inf x ∈ U ′
L ( x ) =
o (1) .
Then we can leverage Theorems 4.3 and 5.1 to get the desired bound.
C.6. O THER O MITTED P ROOFS FOR 1-SAM
We will use ℓ ′ ( y, y k ) and ℓ ′′ ( y, y k ) to denote dℓ ( y ′ ,y k ) dy ′
| y ′
=
y and d 2 ℓ ( y ′ ,y k ) dy ′| y ′
=
y .
Lemma C.64.
Under Setting C.1, fix k ∈ [ M ] , for any p satisfying ℓ ( f k ( p ) , y k ) = 0 , we have that
[Equation]
Proof.
ℓ ( f k ( p ) , y k ) = 0 implies ℓ ′ ( f k ( p ) , y k )
= 0 .
Then by Taylor Expansion,
[Equation]
This concludes the proof.
Proof of Lemma 5.3.
By Lemma C.64, as L ( p ) = 1 M (cid:80) M k =1 L k ( p ) , we have
[Equation]
By definition of Γ in Setting C.1, we have for any p ∈ Γ , {∇ f k ( p ) } nk =1 are linearly independent, which implies that ∇ f k ( p ) ̸ = 0 for any p ∈ Γ .
For any p ∈ Γ , as ∇ f k ( p ) ∥∇ f k ( p ) ∥ is well defined and continuous at p , there exists a open ball V containing p such that for any x ∈ V , ∥∇ f k ( x ) ∥ 2 ≥ C 1 > 0 and ∥∇ [ ∇ f k ( x ) ∥∇ f k ( x ) ∥ ] ∥ 2 ≤ C 2 for some constants Cand C 2 .
Suppose ∇ L k ( x ) ̸ = 0 , then as by Taylor Expansion,
[Equation]
We have ∇ L k ( x ) ∥∇ L k ( x ) ∥ = ∇ f k ( x ) ∥∇ f k ( x ) ∥ = ∇ f k ( p ) ∥∇ f k ( p ) ∥
+ C 2 ∥ x − p ∥ , which completes the proof.
We note that the alignment result in Lemma 5.3 is not directly used in our proof.
Instead, we use its generalized version Lemma C.57 which holds under holds under a more general condition than
Setting C.1, namely Condition C.2.
C. T ECHNICAL L EMMAS
Lemma C.65 (Corollary 4.3.15 in Horn & Johnson (2012)) .
Let Σ , ˆΣ ∈ R D
× D be symmetric and non-negative with eigenvalues λ 1 ≥ ... ≥ λ D and ˆ λ 1 ≥ ... ≥ ˆ λ D , then for any i ,
| ˆ λ i − λ i | ≤ ∥ Σ − ˆΣ ∥ Definition C.66 (Unitary invariant norms) .
A matrix norm ∥ · ∥ ∗ on the space of matrices in R p × d is unitary invariant if for any matrix K ∈ R p × d , ∥ UKW ∥ ∗ = ∥ K ∥ ∗ for any unitary matrices U ∈ R p × p , W ∈ R d × d .
Theorem C.67.
[Davis-Kahan sin( θ ) theorem (Davis & Kahan, 1970)] Let Σ , ˆΣ ∈ R p
×
p be symmetric, with eigenvalues λ 1 ≥ . . .
≥ λ p and ˆ λ 1 ≥ . .
≥ ˆ λ p respectively.
Fix 1 ≤ r ≤ s ≤ p , let d ≜ s − r + 1 and let V = ( v r , v r +1 , . .
, v s
) ∈ R p ×
d
and ˆ V =
(ˆ v r , ˆ v r +1 , . . .
, ˆ v s ) ∈
R p × d have orthonormal columns satisfying Σ v j = λ
j v j and ˆΣˆ v j = ˆ λ j ˆ v j for j = r, r + 1 , . .
. , s. Define ∆ ≜ min (cid:110) max { 0 , λ s − ˆ λ s +1 } , max { 0 , ˆ λ r − 1 −
λ r } (cid:111) , where ˆ
λ 0 ≜ ∞ and ˆ λ p +1 ≜ −∞ , we have for any unitary invariant norm ∥ · ∥ ∗ , ∆ ·
∥ sin Θ( ˆ V , V ) ∥ ∗ ≤ ∥ ˆ Σ − Σ ∥ ∗ .
Here Θ( ˆ V , V ) ∈ R d × d , with Θ( ˆ V , V ) j,j = arccos σ j for any j ∈ [ d ] and Θ( ˆ V , V ) i,j = 0 for all i ̸ = j ∈ [ d ] .
σ 1 ≥ σ 2 ≥ · · · ≥ σ d denotes the singular values of ˆ V ⊤ V. [sin Θ] ij is defined as sin(Θ ij ) .
Lemma C.68 (Azuma-Hoeffding Bound) .
Suppose Z n is a super-martingale, suppose − α ≤
Z i +1 − Z i ≤ β , then for all n > 0 , a > 0 , we have
[Equation]
Lemma C.69 (Magnus (1985)) .
Let A : R D → R D
× D be any C 1 symmetric matrix function and x ∗ ∈ R D satisfying λ
( A ( x ∗ )) > λ 2 ( A ( x ∗ )) and v 1 be the top eigenvector of A ( x ∗ ) .
It holds that ∇ λ 1 ( A ( x )) | x = x ∗ = ∇ ( v ⊤ 1 A ( x ) v 1 ) | x = x ∗ .
We then present some of the technical lemmas we required to prove Lemma C.24.
Lemma C.70.
If 0 < c < b − a b 2 , a (cid:113) a 2 +2 b 2 2(1 − cb ) ≥ a 2 + b 2 2 − ca − cb , then a > 12 b, cb ≤
Proof.
Notice that
[Equation]
So
[Equation]
As c < b − a b 2 , we have 1 > 1 − cb > ab .
So
[Equation]
The above inequality implies a ≥ 12 b .
As c < b − a b 2 , cb ≤
Lemma C.71.
When 0 < a < b, 0 < c < b − a b 2 , we have
[Equation]
Proof.
Equivalently, we are going to prove
[Equation]
Further simplifying
[Equation]
we only need to prove
[Equation]
We have the following auxiliary inequalities,
[Equation]
Using the above auxiliary inequalities we have
[Equation]
Lemma C.72.
When 0 < a < b, 0 < c < b − a b 2 , a (cid:113) a 2 +2 b 2 2(1 − cb ) ≥ a 2 + b 2 2 − ca − cb , we have
[Equation]
Proof.
Equivalently, we are going to prove,
[Equation]
We have the auxiliary inequality
[Equation]
−
[Equation]
[Equation]
Hence cb
[Equation]
[-Case 
-][Table] [Table] [Equation] [Equation]
[Table] [Table]
[Equation]
[-Case 
-][Table]
, then
[Equation]
Using Lemma C.70, a > b 2 , ( b − a ) 2 ( b + a )
= ( b 2 − a 2 )( b − a ) ≤ b 2 ( b − a ) ≤ b [-
-][Equation]
Lemma C.73.
When 0 ≤ a ≤ b, 0 ≤ c ≤ b − a b 2 , b 2 ≥ a (cid:113) a 2 +2 b 2 2(1 − cb ) ≥ a 2 + b 2 2 − ca − cb , we have
[Equation]
Proof.
Define
[Equation]
Here we suppose WLOG S a ( c, b ) ̸ = ϕ .
Consider
[Equation]
( a )
[Equation]
− cb b 2 + 12 a
Define u ≜ cb, v ≜ ab , then u + v ≤ 1 .
[Equation]
As (cid:113) (1 − u [Equation] ≥ √
[Equation]
u ,we have
[Equation]
The above inequality shows that F ( a ) is convex w.r.t to a for a min ( c, b ) ≤ a ≤ a max ( c, b ) .
Hence F ( a ) ≤ max ( F ( a min ( c, b )) , F ( a max ( c, b )))
Part 1 We abuse the notation and use a min a shorthand for a min ( c, b ) .
We have a min (cid:113)
[Equation]
ca min − cb . This implies
[Equation]
Hence using Lemma C.71,
[Equation]
Part 2 We abuse the notation and use a max a shorthand for a max ( c, b ) .
It’s not easy to see which boundary condition a max satisfy, hence we will discuss by cases.
[Equation]
, in this case we simply redo the calculation in Part 1.
2 Case 2, b 2 = a max (cid:113) a 2 max +2 b 2 2(1 − cb ) .
This implies
[Equation]
Hence using Lemma C.72,
[Equation]
and
[Equation] ≥ a max (cid:113) [Equation]
This implies a 3 max + 2 a max b 2 − 2 b 3 ≤ 0 ⇒ a max < 910 b .
Define v ≜ a max b , cb = 1 − v
Then by Lemma C.70, 12 ≤ v ≤
[Equation]
We will prove the following inequality,
[Equation]
In fact we can directly show
[Equation]
for v ∈ [0 . 5 , 0 . 9] .
Hence we have
[Equation]
D O MITTED P ROOFS ON C ONTINUOUS A PPROXIMATION
Our main results Theorems 4.9 and 5.4 characterize the dynamics of full-batch SAM and 1-SAM using continuous approximations for both the first phase and second phase.
In this section we provide the omitted details for such continuous approximation.
Our analysis is mainly based on Li et al.
Below we first restate their main theorem adapted to our setting.
Assumption D.1.
The function b : R D → R D is uniform Lipschitz, that is, sup x,y ∈ R D
∥ b ( x ) − b ( y ) ∥ 2 ∥ x − y ∥.
Definition D.2.
Let G denote the set of continuous functions R D
→ R of at most polynomial growth, i.e. g ∈ G if there exists positive integers κ 1 , κ 2 > 0 such that
[Equation]
for all x ∈ R d .
Moreover, for each integer α ≥ 1 we denote by G α the set of α -times continuously differentiable functions R D
→ R which, together with its partial derivatives up to and including order α , belong to G .
Moreover, if g depends on additional parameters, we say g ∈ G α if the constants κ 1 , κ 2 are independent of these parameters, i.e., g ∈ G α uniformly.
Finally, the definition generalizes to vector-valued functions coordinate-wise.
Let ψ b ( x, t ) , t ∈ R , t ≥ 0 be the unique solution of
[Equation]
and let ψ ( x ) be a random variable depending on x and ψ ( x, k ) , k ∈ N ≜ ( ψ )
k ( x ) .
Note here we abuse the notation and assume the randomness in application of ψ is independent.
For example, ϕ ( x ) could be used to denote the one-step update of 1-SAM, ψ ( x, 1) ≡ ψ ( x ) =
x − η ∇ L k
x + ρ ∇ L k ( x ) ∥∇ L k ( x ) ∥ 2 ) where k is independently sampled, and ψ ( x, 2) denote the (random) iterate of 1-SAM achieved after two steps of updates.
In this section (including the following theorem), ψ depends on η and ρ , but we omit the explicit dependency on η and ρ for notational convenience.
Below is a modified version of Theorem 3 in Li et al.
(2019), adapted to our setting.
The main difference between the original version and our adaption is that our discrete update ψ ( x ) have two hyperparameters η and ρ , where the proper continuous time scaling depends on both hyperparameters.
In detail, in Phase I, the target continuous dynamics is the gradient flow and in Phase II, the target continuous dynamics is the Riemmannien gradient flow on manifold.
(See Table 2 for a summary.)
Therefore, we use p ( η, ρ ) to denote the amount of continuous time that one-step discrete update corresponds to, which is the η in Li et al.
We further use q ( η, ρ ) to denote the total error.
Though the original theorem statement in Li et al.
(2019) only applies to the case q ( η, ρ ) =
p α ( η, ρ ) , their proof indeed extends to the more general case, p α ( η, ρ )
=
O ( q ( η, ρ )) , as stated below.
Theorem D.3 (Adapted version of Theorem 3, Li et al.
Let T > 0 , α ≥ 1 be an integer, p, q : R + × R + → R + be two functions satisfying that p ( η, ρ )
=
O ( η + ρ ) , Ω( q ( η, ρ ))
=
p α ( η, ρ ) =
o (1) , and N =
⌊ T/p ( η, ρ ) ⌋ .
Let ∆( x ) ≜ ψ ( x ) − x and (cid:101) ∆( x ) ≜ ψ b ( x, p ( η, ρ )) − x .
Under
Assumption D.1, suppose further that the following conditions holds:
(1) There exists a K 1 ∈ G independent of η, ρ such that
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) E s (cid:89)
j = ∆ ( i j ) ( x ) − s (cid:89)
j = (cid:101) ∆ ( i j ) ( x ) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) ≤ K 1 ( x ) p ( η, ρ ) q ( η, ρ ) , for s = 1 , 2 , . .
α and
E
α
+1 (cid:89)
j = (cid:12)(cid:12) ∆ ( i j ) ( x ) (cid:12)(cid:12) ≤ K 1 ( x ) p
α +1 ( η, ρ ) , for all i j ∈ { 1 , . . .
, D } .
(2) For each m ≥ 1 , the 2 m -moment of ψ ( x, k ) is uniformly bounded with respect to k , η and ρ , i.e., there exists K 2 ∈ G , independent of η, ρ, k , such that
[Equation]
for all k = 0 , . . .
, N ≡ ⌊ T/p ( η, ρ ) ⌋ .
Then for each g ∈ G α +1 , there exists a constant C > 0 , independent of η, ρ , such that
[Equation]
[Equation]
Theorem C.
[Equation]
C.
[Equation]
Theorem C. ηρρ + η − ∂ Φ
[Equation]
∇
Tr
[Equation]
∇
[Equation]
L
[Equation]
Theorem D.4 (Lemma 29, Li et al.
{+Suppose there is a function p : R + × R + → R + and a constant c > 
+}
such that (cid:13)(cid:13) ψ ( x ) − x (cid:13)(cid:13) 2 ≤ C (1+ | x | ) p ( η, ρ ) for all x ∈ R D , then for any fixed T >and any m ≥ 1 , E | x k | m exists and is uniformly bounded in η and k = 0 , . . .
, N =
⌊ N/p ( η, ρ ) ⌋ .
Combining Theorems D.3 and D.4, we have the following theorem:
Theorem D.5.
Let T > 0 , p, q : R + × R + → R + be two functions satisfying that p ( η, ρ ) =
O ( η + ρ ) , q ( η, ρ ) =
Ω( p ( η, ρ )) =
o (1) , and N =
⌊ T/p ( η, ρ ) ⌋ .
Let ∆( x ) ≜ ψ ( x ) − x and (cid:101) ∆( x ) ≜ ψ b ( x, p ( η, ρ )) − x .
Suppose there is an open neighborhood V of { ψ b ( x, t ) }
{+Tt =
+}
that b is uniform Lipschitz in V and ∥
∆( x ) ∥ 2 p ( η,ρ ) is bounded in V .
Further suppose that the following conditions holds: There exists a constant K independent of η, ρ such that for all x ∈ V ,
[Equation]
Then there exists a constant C > 0 , independent of η, ρ , such that for each k = 0 , . . .
, ⌊ T/p ( η, ρ ) ⌋ ,
[Equation]
Moreover, when ψ ( x ) is deterministic, there exists a constant C > 0 , independent of η, ρ , such that
[Equation]
Proof of Theorem D.5.
The proof invokes Theorem D.3 with α = 1 and g t ( x ) =
∥ x − ψ b ( x, t ) ∥for each 0 ≤ t ≤ T and g i ( x )
= x i for each i ∈
[ M ] .
Since there is an open neighborhood V of { ψ b ( x, t ) } Tt =0 that b is uniform Lipschitz in V , without loss of generality, we can assume that b is uniform Lipschitz in R D , because by classical results approximation results, (cid:13)(cid:13)
ψ ( x, k ) − ψ b ( x, kp ( η, ρ ))
(cid:13)(cid:13) 2 will be larger than constants for at most exponential small probability and thus we can just modify b outside V such that it is uniform Liptshitz.
This will not affect our final high probability bounds.
By Theorem D.4, we know that the condition (2) of Theorem D.3 holds.
For the same reason,, we can assume (cid:13)(cid:13)(cid:13) E ∆( x ) − (cid:101) ∆( x ) (cid:13)(cid:13)(cid:13)
≤
Kp ( η, ρ ) q ( η, ρ ) holds for each x ∈ R D , otherwise we can modify ψ ( x ) outside V .
We can also assume ∥ ∆( x ) ∥ 2 p ( η,ρ ) is not only bounded in V , but in the entire R D .
Since we set α = 1 , it holds that
[Equation]
for all i j ∈ { 1 , . . .
, D } .
Now by Theorem D.3, choose g ( x ) =
∥ x − ψ b ( x, kp ( η, ρ ))
∥ 22 for each k , we know Then there exists a constant C > 0 , independent of η, ρ , such that
[Equation]
By Markov inequality, we know that the probability of (cid:13)(cid:13) ψ ( x, k ) − ψ b ( x, k · p ( η, ρ ))
(cid:13)(cid:13) 2 ≥
C (cid:112) q ( η, ρ ) is at most (cid:112) q
( η, ρ ) .
When ψ ( x ) is deterministic, setting g i
= x i for each i ∈
[ D ] , we know that there exists a constant
C > 0 , independent of η, ρ , such that
[Equation]