How Sharpness-Aware Minimization Minimizes Sharpness?
Authors: Kaiyue Wen
Tengyu Ma
Zhiyuan Li
Venue: ICLR 2023 poster (2023)
Keywords: implicit bias
implicit regularization
sharpness
sharpness aware minimization

Abstract: 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.

Published as a conference paper at ICLR 2023
How DOES SHARPNESS-AWAE MINIMIZATION MINIMIZE SHARPNESS?
Kaiyue Wen
Tengyu Ma, Zhiyuan Li Computer Science Department Stanford University
Institute for Interdisciplinary Information Sciences Tsinghua University
wenky200maiils t tsinguua.... edu. cn
(tengyuma, zhiyuanlijestanfordd.
ABSTRACT
Sharpness-Aware Minimization (SAM) is highly effective regularization technique for improving the generalization of deep neural networks for various settings. However, the underlying working of SAM remains clusive because of various intriguing approxima- tions 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.
1 INTRODUCTION
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 able to find generalizable models. Recent empirical and theoretical studies suggest that generalization is well-correlatd with the sharpness of the loss landscape at the learned parameter (Keskar et al., 2016; Dinh et al., 2017: Dziugaite et al., 2017: Neyshabur et al., 2017; Jiang et al., 2019). Partly motivated by these studies, Foret et al. (2021); Wu et al. (2020); Zheng et al. (2021); Norton et al. (2021) propose penalize the sharpness of the landscape improve the generalization. We refer this method to Sharpness-Awar Aware Minimization (SAM) and focus on the version of Foret et al. (2021).
Despite its empirical success, the underlying working of SAM remains elusive because of the various in- triguing approximations made in its derivation and analysis. There are three different notions of sharpness involved SAM intends optimize the first notion, the sharpness along the worst direction, but actually implements computationally efficient notion, the sharpness along the direction of the gradient. But in the analysis of generalization, a third notion of sharpness is actually used to prove generalization guarantees, which admits the first notion as upper bound. The subtle difference between the three notions can lead to very different biases (see Figure for demonstration).
More concretely, let L be the training loss, x be the parameter and p be the perturbation radius, a hyperpa- rameter requiring tuning. The first notion corresponds to the following optimization problem (1), where we call (x)) = I.2 [.Max( L(xx) the worst-direction sharpness at x. SAM intends to minimize the original training loss plus worst-direction sharpness at Z.
min L (x), where L. (x) max L(x + pv).
However, even evaluating Maxx (2x) is computationally expensive, not to mention optimization. Thus Foret et al. (2021); Zheng et al. (2021) have introduced a second notion of sharpness, which approximates the worst- case direction in (1) by the direction of gradient, as defined below in (2). We call RA<<(x) (x) = L. As (x) - L(x) the ascent-direction sharpness at x.
min L Ase (x), where LA = L L(x + (VL())/
(2)

---

Published as a conference paper at ICLR 2023
Type of Sharpness-Aware Loss
Notation
Definition
Biases (among minimizers)
Worst-direction
Max
max/s1l12 L(x + po) x P V(,)) L (IVL(=)II
ming (v(V)((z)) (Thm G.3)
p L
Ascent-direction
ming Amin (V2 (a)) (Thm G.4)
Average-direction
I.p
E Fg [xx + PTITT
ming Tr(V2((x)) (Thm G.5)
Definitions and biases of different notions of sharpness-aware loss. The corresponding sharpness is defined as the difference between sharpness-awan loss and the original loss. Here A denotes the largest eigenvalue and Amin denotes the smallest non-zero eigenvalue.
Table 1:
For further acceleration, Foret et al. (2021); 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., VLA (z) ~ VL (I p [=)x)]l2) and derive the update rule of SAM, where 17 is the learning rate.
Sharpness- Aware Minimization (SAM): 1) = =()) V (ar +V /VL==>/IIV[[>>))) [[V((=))2)
(3)
Intriguingly, the generalization bound of SAM upperbounds the generalization error by the third notion of sharpness, called
sharpness, (z) and defined formally below.
(x)) = (x) - L(x), where T. (xx) = Rg [x -gg//g//22
(4)
The worst-case sharpness is an upper bound of the average case sharpness and thus it is a looser bound for generalization error. In other words, according to the generalization theory in Foret et al. (2021); Wu et al. (2020) in fact motivates us to directly minimize the average case sharpness (as opposed to the worst-case sharpness that SAM intends to optimize).
In this paper, we analyze the biases introduced by penalizing these various notions of sharpness as well as the bias of SAM (Equation 3). Our analysis for SAM is performed for small perturbation radius p and learning rate 1 under the setting where the minimizers of loss form a manifold following the setup of Fehrman et al. (2020); Li et al. (2021). In particular, we make the following theoretical contributions.
1. We prove that full-batch SAM indeed minimizes worst-direction sharpness. (Theorem 4.5)
2. Surprisingly, when batch size is 1. SAM minimizes average-direction sharpness. (Theorem 5.4)
3. We provide characterization (Theorems 4.2 and 5.3) of what a few sharpness regularizers bias towards among the minimizers (including all the three notions of the sharpness in Table 1), when the perturbation radius P goes to zero. Surprisingly, both heuristic approximations made for SAM lead to inaccurate conclusions: (1) Minimizing worst-direction sharpness and ascent-direction sharpness induce different biases among minimizers, and (2) SAM doesn't 't minimize ascent-direction sharpness.
The key mechanism behind this bias of SAM is the alignment between gradient and the top eigenspace of Hessian of the original loss in the latter phase of training-the angle between them decreases gradually to the level of (p). It turns out that the worst-direction sharpness starts to decrease once such alignment is established (see Section 4.3). Interestingly, such alignment is not implied by the minimization problem (2), but rather, it is an implicit property of the specific update rule of SAM. Interestingly, such an aligmment property holds for SAM with full batch and SAM with batch size one, but does not necessarily hold for the mini-batch case.
2 RELATED WORKS
Sharpness and Generalization. The study on the connection between sharpness and generalization can be traced back to Hochreiter et al. (1997). Keskar et al. (2016) observe a positive correlation between the batch size, the generalization error, and the sharpness of the loss landscape when changing the batch size. Jastrzeb- ski et al. (2017) extend this by finding a correlation between the sharpness and the ratio between learning rate batch size. Dinh al. (2017) show that one can easily construct networks with good generalization but with arbitrary large sharpness by reparametrization. Dziugaite et al. (2017); Neyshabur et al. (2017); Wei al. (2019a;b) give thcoretical guarantees on the generalization error using sharpness-relae related measures. Jiang et al. (2019) perform large-scale empirical study on various generalization measures and show that sharpness- based measures have the highest correlation with generalization.
Background on Sharpness-Aware Minimization. Foret et al. (2021); Zheng et al. (2021) concurrently pro- pose to minimize the loss at the perturbed from current parameter towards the worst direction to improve generalization. Wu et al. (2020) propose an almost identical method for a different purpose, robust gener- alization of adversarial training. Kwon et al. (2021) propose a different metric for SAM to fix the rescaling problem pointed out by Dinh et al. (2017). Liu et al. (2022) propose more computationally efficient version

---

Published as a conference paper at ICLR 2023
1-5AM
X3
(Minimize TrVP9((X))
Full-batch SAM
(Minimize (,(9)L(x))
GD on ascent-direction Loss (Minimize An(VFL(x))
01 0.2 0.4
0.6
0.8 1
X1
Figure 1: Visualization of the different biases different sharpness notions on a 4D-toy example. Let Fr, F2 R² + R+ be two positive functions satisfying that Fi > F22 on [0, 1J² For x E IR,, consider loss [(x) x2)) + Fa(m1, 22 ).. The loss L has zero loss manifold (#g x4 0} of codimension M 2 and the two non-zero eigenvalues of V L of any point x the manifold are A1 L()))( and Lxx)) = F2(21122) We test three optimization algorithms on this 4D-toy model with small learning rates. They all quickly converge to zero loss, i.e, za((),, (1) = O, and after that xı (t), xz((t) still change slowly, i.e. moving along the zero loss manifold. We visualize the loss restricted to (z3, IT) as the 3D shape at various (z1, x2)'s where 21 -z) (f), x22 (t) follows the trajectories of the three algorithms. In other words, each of the 3D surface visualize the function g
Fi ( 22)
(t), 222)) æ3, xr.) As our theory predicts, (1) Full-batch SAM (Equation 3) finds the minimizer with the smallest top eigenvalue, Fi (111, 22); (2) GD on ascent-direction loss LP (Equation 2) finds the minimizer with the smallest bottom eigenvalue, F3 (2(11 22): (3) 1-SAM (Equation 13) (with Lo (x(x) = Fi (21, )a3 and (zx = (2 x2)=zt finds the minimizer with the smallest trace of Hessian, (1 x2) + F2(11 See more details in Appendix B.
of SAM. Zhuang et al. (2022) proposes variant of SAM, which improves generalization by simultane- ously optimizing the surrogate gap and the sharpness-aware loss. Zhao et al. (2022) propose to improve generalization by penalizing gradient norm. Their proposed algorithm can be viewed as generalization of SAM. Andriushchenko et al. (2022) study a variant of SAM where the step size of ascent step is p instead They show that for simple model this variant of SAM has stronger regularization effect when batch size compared to the full-batch case and argue that this might be the explanation that SAM generalizes better with small batch sizes. More related works are discussed in Appendix A.
of [
3 NOTATIONS AND ASSUMPTIONS
For any natural number k we say a function is C* if it is k-times continuously differentiable and is 2t if its kth order derivatives are locally lipschitz. We say a subset of RR' D is compact if each of its open covers has finite subcover. It is well known that subset of R is compact if and only if it is closed and bounded. For any positive definite symmetric matrix A E DxD define (A:(A), 0z (x)))<<(() as all its eigenvalues and eigenvectors satisfying (A) Z (A)... 2 (A) and [A)l22) = 1. For any mapping F, we define 0F(x) the Jacobian where [xF(z)) (x). Thus the directional derivative of F along the vector at x can be written as OF (F()). u. We further define the second order directional derivative of F along the vectors and at x, 8² (u v), a(aF (I),,, that is, the directional derivative of OF u along the vector at x.
Given a C¹ submanifold (Definition C.1) I of IR. and a point æ E T, define Par as the projection operator onto the manifold of the normal space of I at x and P.lr = ID - -P,r. We fix our initialization as Finit and our loss function as I, : RRD - IR. Given the loss function, its gradient flow is denoted by mapping é: : DR.D x [0, oc) - RDD Here, Φ(x, r) denotes the iterate at time r of a gradient flow starting at and is defined as the unique solution of (xx,T) = x 1o VL(0(x, d)) Vr R. We further define the limiting map & as ((x) = lim Φ(x, T), that is, ((x) denotes the convergent point of the gradient flow starting from I. When L(z) is small, (xx) and x are near. Hence in our analysis, we regularly use ((x(t)) as surrogate to analyze the dynamics of ztt). Lemma 3.1 an important property of & from Li al. (2021) (Lemma C.2), which is repeatedly used in our analysis. The proof is shown in Appendix F.
Lemma 3.1. For any x at which & is defined and differentiable, we have that 04 LL)) = 0.
Recent empirical studies have shown that there 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) 8).. Motivated by this empirical discovery, we make the assumption below following Fehrman al. (2020); Li et al. (2021); Arora et al. (2022), which is theoretically justified by Cooper (2018) under a generic setting.
Assumption 3.2. Assume loss L RD + R is C2, and there exists C² submanifold T of RD that is (D -idiension for some integer 1 S M S D, where for all x E T, x is a local minimizer of and rank L(x)) M.

---

Published as a conference paper at ICLR 2023
Though our analysis for the full-batch setting is performed under the general and abstract setting, Assump- tion 3.2, our analysis for the stochastic setting uses a more concrete one, Setting 5.1, where we can prove that Assumption 3.2 holds. (see Theorem 5.2)
Definition 3.3 (Attraction Set). Let U be the attraction set TfT under gradient flow, that is, a neighborhood of I containing all points starting from which gradient flow w.r.f. loss I. converges some point in T, or mathematicallyy U {a: E (x)) exists and ((x)) l).
It can be shown that for minimum loss manifold, the rank of Hessian plus the dimension of the manifold is at most the environmental dimension D, and thus our assumption about Hessian rank essentially says the rank is maximal. Assumption 3.2 implies that U is open and & is T² on U (Arora et al., 2022, Lemma B.15).
4 EXPLICIT AND IMPLICIT BIAS IN THE FULL-BATCH SETTING
In this section, we present our main results in the full-batch setting. Section 4.1 provides characterization of explicit bias of worst-direction, ascent-dircetion, and average- dieeee-drrcction sharpness. In particular, we show that ascent- dinttdiieetionn sharpness and worst vorst-direction sharpness have different explicit biases. However, it turns out the explicit bias of ascent-direction sharpness is not the effective bias of SAM (that approximately optimizes the ascent-direction sharpness), because the particular implementation of SAM imposes additional, different biases, which is the main focus of Section 4.2. We provide our main theorem in the full-batch setting, that SAM implicitly minimizes the worst-direction sharpness, via characterizing its limiting dynamics as learning rate p and 1) goes to 0 with a Riemmanian gradient flow with respect to the top eigenvalue of the Hessian of the loss on the manifold of local minimizers. In Section 4.3, we sketch the proof of the implicit bias of SAM and identify a key property behind the implicit bias, which we call the implicit alignment between the gradient and the top eigenvector of the Hessian.
4.1 WORST- AND ASCENT-DIRECTION ASCENT-DIRECTION SHARPNESS HAVE DIFFERENT EXPLICIT BIASES
In this subsection, we show that the explicit biases of three notions of sharpness are all different under Assumption 3.2. We first recap the heuristic derivation of ascent-direction sharpness R
The intuition of approximating ItMax by Ro comes from the following Taylor expansions (Foret et al., 2021; Wu et al., 2020). Consider any compact set, for sufficiently small p. the following holds uniformly for all æ in the compact set:
L(x + pv) L(x)= = sup (po) VL(I) + p² T V'L(=)o- 2
Mux (r) sup
(5)
11
11
(L(x)
V [|+))))
Ro (x)=b(x> L(x + \ ) L(x) (1)=>>>>>>//00
VV(())1 +0
(6)
2
Here, the preference among the local or global minima is what we are mainly concerned with. Since suPpelllas T (I(x)) [|V](=)]l2 when [=)(2)) > O, the leading terms in Equations 5 and 6 are both the first order term, p [[V[(==)]22- and are the same. However, it is erroneous to think that the first order term decides the explicit bias, as the first order term [V((2)]1 vanishes at the local 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 D(p")-approximal sharpness-award loss needs to be of order p² so that we can guarantee the second- order terms in Equation 5 and/or Equation 6 be non-trivially small. Our main result in this subsection (Theorem 4.2) gives explicit characterization for this phenomenon. The corresponding explicit biases for each type of sharpness is given below in Definition 4. As we will see later, they can be derived from general notion of limiting regularizer (Definition 4.3).
minimizer of the sharpness- aware loss because VL(z)
0. Therefore, the
Definition 4.1. For x E RDD, we define (xx = \=(V=T(=))/2 (A =AM (x)))(2))) and ( = TT())(()))((I)))
Theorem 4.2. Under Assumption 3.2, let U' be any bounded open set such that its closure U C U and U Prrr C U'NI For any type E [Max, Asc, Avg) and any optimality gap A > O, there is a function R: + R with limp-o c(p) = 0, such that for all sufficiently small > and all u E U" satisfying that [(u) + Ryype (u) inf (I(x)) Rlype (x:) - Ap?
aeti
it holds inf Lxx) < (A (())))))p and that Stape (u) inf (TYDe(x)) E [[-e(p), Ae(()))
We note that 17 (x) is undefined when [V](2]]2 = 0. In such cases, set (x co.

---

Published as a conference paper at ICLR 2023
Theorem 4.2 suggests sharp phase transition of the property of the solution of min, I(x) + Rp(x)) when the optimization error drops from w(p²) to O(p²). When the optimization error is larger than w(p"), regularization effect happens and any minimizer satisfies the requirement. When the error becomes O(p)), there is non-trivial restriction on the coefficients in the second-order term.
Next we give a heuristic derivation for the above defined Stype. sharpness, the calculations are fairly straightforward and well-known in literature (Keskar et al., 2016; Kaur et al., 2022; Zhuang et al., 2022; Orvieto et al., 2022), and we sketch them here. In the limit of perturbation radius p O, we know that the minimizer of the sharpness-aware loss will also converges to r, the manifold of minimizers of the original loss L. Thus to decide to which æ E I the minimizers will converge to as - 0, it suffices to take Taylor expansion of L or Lp at each x E I and compare the second-order p² Tr (r(V² + (pp)) and RM 2 A1 +
First. for worst- and average-direction direction
coefficients, e.g., we have that R (xx 2D by Equation 5.
e²
However, the analysis for ascent-direction sharpness is more tricky because R = Do for any x E I and thus is not continuous around such z. Thus we have to aggregate information from neighborhood to capture the explicit bias of Rp around manifold I' This motivates the following definition of limiting regularizer which allows us to compare the regularization strength of R around each point on manifold I as p +
Definition 4.3 (Limiting Regularizer). We define the limiting regularizer of {Bp} as the function?
- S(x) = lim lim p-+ r-0
- Rp(z')/p²
inf
yet strictly being non- zero. By Equation 6, we have RA e² -LL''' v L(()) Here the crucial step of the proof is that because of Assumption 3.2, VL(1// [ must almost lie in the column span of V which implies that in (T(x'"T<<<(a)) ((a"//)VVV""""/VVVVVV)VV)"""""""
To minimize Epe around n, we can pick x' + x satisfying that [VL(('))l2
3 (M(VV where rank (ank((²I(zz)) = M by Assumption 3.2. The above aligmment property between the gradient and the column space of Hessian can be checked directly for any non non-negative quadratic function. The maximal Hessian rank assumption in Assumption 3.2 ensures that this property extends to general losses.
We defer the proof of Theorem 4.2 into Appendix G.I, where we develop a sufficient condition where the notion of limiting regularizer characterizes the explicit bias of R, as p 0.
4.2 SAM PROVABLY DECREASES WORST-DIRECTION SHARPNESS
Though ascent-direction sharpness has different explicit bias from worst-direction sharpness, in this subsec- tion we will show that surprisingly. SAM (Equation 3), heuristic method designed to minimize ascent- direction sharpness, provably decreases worst-direction sharpness. The main result here is an exact char- acterization of the trajectory of SAM (Equation 3) via the following ordinary differential equation (ODE) (Fquation 7), when learning rate v) and perturbation radius p are small and the initialization a0) = Tinit is in U, the attraction set of manifold r.
1 [ PA X(0) = (xInit). 3
X(7) x(0)
We assume ODE (Equation 7) has solution till time Ts, that is, Equation 7 holds for all t S Ts. We call the solution of Equation 7 the limiting flow of SAM, which is exactly the Riemannian Gradient Flow on In other words, the ODE (Equation 7) is essentially a projected gradient descent algorithm with loss > (V))()) on the constraint set I and an infinitesimal learning rate. Note >1 [x))) may not be differentiable at x if X (v²)(x)) (x))) thus to ensure Equation 7 is well-defined, we assume there is a positive eigengap for L on 1:. Assumption 4.4. For all x E T, there exists a positive eigengap, ie., (v)))(())) > de (v))
the manifold I' with respect to the loss Y₁ (V2
Theorem 4.5 is the main result of this section, which is a direct combination of Theorems I.1 and L.3. The proof is deferred to Appendix L.3.
2 Herr we implicitly assume the zeroth and first order term varnishes, which holds for three sharpness notions. 'In fact we only need to assume the positive eigengap along the solution of the ODE, If I' doesn't 't satisty Assump- tion 4.4, we can simply perform the same analysis on its submanifold {æ E 1' eigengap is positive x).
5

---

Published as a conference paper at ICLR 2023
Theorem 4.5 (Main). Let {x((t) be the iterates of full-batch SAM (Equation 3) with x(O) Timit U. Under Assumptions 3.2 and 4.4, for all 7) P such that y (n( and p/ are sufficiently small, the dynamics of SAM can be characterized in the following two phases:
Phase I: (Theorem 1.1) Full-batch SAM (Equation follows Gradient Flow with respect to L until entering an ((p)) neighborhood of the manifold I' in (1n((1 steps;
Phase I: (Theorem 1.3) Under mild don-degeyyyycy assumption (Assumption 1.2) on the initial point of phase II, full-batch SAM (Equation 3) tracks the solution X of Equation 7, the Riemannian Gradient Flow with respect to the loss A1 (()) in an O(np) neighborhood of manifold T. Quantitatively the approximation error between the iterates x and the corresponding limiting flow is (0))o(1/p)) that is,
|| (Tx(( (Tpa/()p)))))))))) ((a))l2 = O(n)n(110)).
Moreover, the angle between V (1(( 7p
and the top eigenspace of V Ts 1)) is O(p). TpP
Theorem 4.5 shows that SAM decreases the largest eigenvalue of Hessian of loss locally around the manifold of local minimizers. Phase I uses standard approximation analysis as in Hairer et al. (2008). In Phase II, as T3 is arbitrary, the approximation and alignment properties hold simultaneously for all ( along the trajectory. provided that 1) ln(1/p) and p/n are sufficiently small. The subtlety here is that the threshold of being "sufficiently small" on 1) ln(1/p) and p/1) actually depends on T3, which decreases when T3 + or - 20. We defer the proof of Theorem 4.5 to Appendix I.
As corollary of Theorem 4.5, we can also show that the largest eigenvalue of the limiting flow closely tracks the worst-direction sharpness.
Corollary 4.6. In the setting of Theorem 4.5, the difference between the worst-direction sharpness of the iterates and the corresponding scaled largest eigenvalues along the limiting flow is at most O(np² (n(1/p)). That is,
Max R
- " O(pp' In(p/)))
(8)
Since In(1/p) is assumed to be sufficiently small, the error (( pp²) is only o(p²), meaning that penalizing the top eigenvalue on the manifold does lead to non-trivial reduction of worst-direction sharpness, in the sense of Section 4.1.
Hence we can show that full-batch SAM (Equation 3) provably minimizes worst-direction sharpness around the manifold if we additionally assume the limiting flow converges to minimizer of the top eigenvalue of Hessian in the following Corollary 4.7.
Corollary 4.7. Under Assumptions 3.2 and 4.4, define U' as in Theorem 4.2 and suppose X(00)= = lim X(t)
exists and is a minimizer of X1 (z)) in U' n.r Then for all e >> 0, there exists T. > 0, such that for all p.n) such that In(1/p) and p/n are sufficiently small, we have that
D (x( (Tp(())))))) I cp² + inf I,Max xEU' " (x))
We defer the proof of Corollaries 4.6 and 4.7 Appendix I.4.
4.3 ANALYSIS OVERVIEW FOR SHARPNESS REDUCTION IN PHASE II OF THEOREM 4.5
Now we give an overview of the analysis for the trajectory of full-batch SAM (Equation 3) in Phase (in Theorem 4.5). The framework of the analysis is similar to Arora et al. (2022); Lyu et al. (2022); Damian et al. (2021), where the high-level idea is to use ((x(t)) as a proxy for xxt) and study the dynamics of ((x(t)) via Taylor expansion. Following the analysis in Arora et al. (2022) we can show Equation using Taylor expansion, starting from which we will discuss the key innovation in this paper regarding implicit Hessian- gradient alignment. We defer its intuitive derivation into Appendix I.5.
np² o (=(())
VL(x(1) [V(((()) (VL()(t))
VL(x(1)
(<((1 + (xx())
1+0 I (VL
2
Now, to understand how ())(t)) moves over time, we need to understand what the direction of the RHS of Equation 9 corresponds (o-we will prove that corresponds to the Riemannian gradient of the loss function ( at x = (x)).. To achieve this, the key is to understand the direction [[
vi(a(t))
vl(a(t)) is close to the top eigenvector of the Hessian up to sign flip, that is It turns out that we will prove [[]|²

---

Setting 5.1. Let the total number of data be M. Let fk(x) be the model output on the k-th data where fa Published as a conference paper at ICLR 2023
2 -. S (p) for some E (-1,1). We call this phenomenon Hessian-gradient
Vi(a(t)
alignment and will discuss it in more detail at the end of this subsection.
Using this property, we can proceed with the derivation (detailed in Appendix 1.5):
np² 0D (((T0/())))VV((((((((0000)))))1) OGP O0P + np²),
(xr(1 + 1)) - (((()) = - 2
(10)
Alignment. It remains to explain why the gradient implicitly aligns to the top eigenvector of the 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.2. Below we first give the formal statement of the implicit alignment on quadratic loss, Theorem 4.8 and defer the result for general case (Lemma 1.19) to appendix. Note this alignment property is an implicit property of the SAM algorithm as it is not explicitly enforced by the objective that SAM is intended to minimize, I.gue. Indeed optimizing L would rather explicitly align gradient to the smallest non zero eigenvector (See proofs of Theorem G.5)!
Implicit Hessian-gradee
Theorem 4.8. Suppose A is a positive definite symmetric matrix with unique top eigenvalue. Consider running full-batch SAM (Equation 3) on loss L(x)= == 1=²² Ax as in Equation 11 below.
x(t + 1) x(l) - (=((()) + AA())aa
(11)
Then, for almost every z(0), we have xxt) converges in direction to V1(A) up to a sign flip and (d)r(4) with n) < 1. lime- ||x(D)]22 2-1A.(A)
The proof of Theorem 4.8 relies two-phase analysis of the behavior of Equation 11, where we first show at (t) enters an invariant set from any initialization and in the second phase, we construct a potential functioo show alignment. The proof is deferred to Appendix H.
Below we briefly discuss why the case with general loss is closely related to the quadratic loss case. We claim that, in the general loss function case, the analog of Equation 11 is the update rule for the gradient:
(i(x(t))
VL(I(t+ +
+ (xx))
(V ) + Onp²²).
(12)
We first note that indeed in the quadratic case where VL(x) Ax and V2L(z) A, Equation 12 is equivalent to Equation 11 because they only differ by a multiplicative factor A on both sides. We derive its intuitive derivation into Appendix I.5.
EXPLICIT AND IMPLICIT BIASES IN THE STOCHASTIC SETTING
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 is natural to ask if we can justify the above intuition by showing the Hessian-gradicn aligmment in the stochastic setting. Unfortunately, such alignment not possible in the most general setting. Yet when the batch size is 1, we can prove rigorously in Section 5.2 that stochastic SAM minimizes stochastic worst-direction sharpness, which is the expectation of the worst-direction sharpness of loss over each data (defined in Section 5.1), which is the main result in this section. We stress that the stochastic worst-direction sharpness has a different explicit bias to the worst-direction sharpness, which full-batch SAM implicitly penalizes. When perturbation radius p + O, the former corresponds to Tr (r L(-)), the same as average- direction sharpness, and the latter corresponds to > (V²L(.)).
Below start by introducing our setting for SAM with batch size 1, or 1-SAM. We still need Assumption 3.2 in this section. We first analyze 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 p + O. In Section 5.2, we show that I-SAM penalizes the trace of Hessian. Below we formally state our setting for stochastic loss of batch size one (Setting 5.1).
is -mmoohh function and yk be the k-th label, for k 1,...,M M. We define the loss on the k-th data as [x(z) = E(fa(x)) yık) and the total loss L EMd Lv/M, where function E(y/ yy C+ -smooth in y'. We also assume for any y R. it holds that arg min EB. Ely', .y) = y and that 82 s)) 10y")" by'=y" > O. Finally, we denote the set of global minimizers of with full-ull-ra full-rank Jacobian by I and assume that it is non-empty, that is,
r 2 {I RP fx(z)) yk. Vk E [M] and (V [x(x)) k=1 are linearly independent] # 0.

---

Published as a conference paper at ICLR 2023
We remark that given training data (i.., (fa) I defined above is just equal the set of global mini- mizers, {ar RD fx( yı, Ak E [M]) except for zero measure set of labels (3k))= when fx are C smooth, by Sard's Theorem. Thus Cooper (2018) argued that the global minimizers form a differentiable manifold generically if we allow perturbation on the labels. In this work we do not make such an assumption for labels. Instead, we consider the subset of the global minimizers with full-rank Jacobian, T. A standard application of implicit function theorem implies that I defined in Setting 5.1 is indeed a manifold. (See Theorem 5.2, whose proof is deferred into Appendix E.1)
Theorem 5.2. Loss L, set I and integer M defined in Setring 5.1 satisfy Assumption 3.2.
1-SAM: We use 1-SAM as shorthand for SAM on stochastic loss with batch size 1 as below Equation 13, where & is sampled i.idd from uniform distribution on [M].
1-SAM:
xtt I) =() 7/ Lu, (ar + LL (x)/ IV, (1)]l2)
(13)
5.1 STOCHASTIC WORST-, ASCENT- AND AVERAGE- DIRECTION SHARPNESS HAVE THE SAME EXPLICIT BIASES AS AVERAGE DIRECTION SHARPNESS
Similar to the full-batch case, we use Max k. k..p to denote the corresponding sharpness-aware loss for
La and R. Kkpp p to denote corresponding sharpness for La respectively (defined as Equations 1, 2 and 4 with I. replaced by Lx). We further use stochastic worst-, ascent- and average-direction sharpness to denote E₁ [RM 1, E& [RA and E& [R Unlike the full-batch setting, these three sharpness notions have the same explicit biases, or more precisely, they have the same limiting regularizers (up to some scaling factor).
Theorem 5.3. The limiting regularizers of three notions of stochastic sharpness. denoted by ŞiMux _AvE I satisfy that SMa (xx) = (x) = - -AAEE())) = (t())((x))/2 Furthermore, define U" in the same way in Theorem 4.2 For any type E {Max, Asc, Avg). it holds that if for some u E U', L(u)= + En ppp
(u)] S inf (L(x) + EE [R
+ = ? then we have that L(u) - -infwey L(x) S ep² + +((?
ZEE"
and that (u) infreurr StPP(()) S <c+o(1).
We defer the proof of Theorem 5.3 to Appendix G.4. 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 Lg. and that the average of limiting regularizers equal to the limiting regularizer of the average regularizer by definition.
5.2 STOCHASTIC SAM MINIMIZES IVERAGE-DIRECTION SHARPNESS
This subsection aims to show that the implicit bias of 1-SAM (Equation 13) is minimizing the average- direction sharpness for small perturbation radius p and learning rate 1), which has the same implicit bias as all three notions of stochastic sharpness do (Theorem 5.3). As an analog of the analysis in Section 4.3, which shows full-batch SAM minimizes worst-direction sharpness, analysis in this section conceptually shows that 1-SAM minimizes the stochastic worst- corst-direction sharpness.
Mathematically, we prove that the trajectory of 1-SAM tracks the following Riemannian gradient flow (Equa- tion 14) with respect to their limiting regularize Tr T()))L(-)) on the manifold for sufficiently small 77 and and thus penalizes stochastic worst-direction sharpness (of batch size 1). We assume the ODE (Equation 14) has solution till time Tz.
X(r) X(0) f Pi VI(V-L((()))d X(0) = +D(iii)-
(14)
Theorem 5.4. Let (x((t) be the iterates of I-SAM (Equation 13) and xx0) = Finis E U, then under Set- ting for almost every Tmit for all and such that (1 +)+p) p) In(1 (1/mp)) is sufficiently small, with probability at least 1 O(p) over the randomness of the algorithm, the dynamics of I-SAM (Equation 13) can be split into two phases:
Phhase I (Theorem J.1): 1-SAM follows Gradient Flow with respect to until entering an Ö(pp) neighbor- hood of the manifold T in O(1(1 (1/pn))/) steps;
WWe note that RA (x) is undefined when [VL()]]2 = 0. In such cases, set Ho C0.

---

Published as a conference paper at ICLR 2023
Phase II (Theorem J.2): 1-SAM tracks the solution of Equation 14, X, the Riemannian gradient flow with respect to L(-) in ((p) neighborhood of manifold T. Quantitatively, the approximation error between the iterates x and the corresponding limiting flow X is (1) //2 p). that is,
(mp")) - (Ta)ll2 = (1) + P).
The high-level intuition for the Phase II result of Theorem 5.4 is that Hessian- -rraiien alignment holds true for every stochastic loss L along the trajectory of 1-SAM and therefore by Taylor expansion (the same argu- ment in Section 4.3), each step (xx(t) moves towards the negative (Riemannian) gradient of Å (V² Lk.) where ky is the index of randomly sampled data, or the limiting regularizer of the worst-direction sharp- ness of LR,- Averaging over long time, the moving direction becomes the negative (Riemmanian) gradient of Eg, [A1 Lv,L,, which is the limiting regularizer of stochastic worst-direction sharpness and equals to Tr(V"L) L) by Theorem 5.3.
The reason that Hessian-gradient alignment holds under Setting 5.1 is that the Hessian of each stochastic loss La at minimizers p r, L (p)
(an') 1 (===f(()) V (V<())(V[((p))) ((Lemma J.15), is exactly rank-1, which enforces the gradient V (L)x)) L (f(x))(> {(x)) to (almost) lie in the top (which is also the unique) eigenspace (x Lemma 5.5 formally states this property.
8 2o mp)
Lemma 5.5. Under Setting 5.1, for any p and k E [M]. it holds that f(a + 0 and that there is an open set V containing p. satisfying that
VL (x) Yx E (x) # = 3s Es (-1,1)
Vd((p) VL, (x) = (& + O - pll2).
Corollaries 5.6 and 5.7 below are stochastic counterparts of Corollaries 4.6 and 4.7, saying that the trace of Hessian close to the stochastic worst-direction sharpness along the limiting flow (14), and therefore when the limiting flow converges to a local minimizer trace Hessian, 1-SAM (Equation 13) minimizes the everage-direction sharpness. We defer the proofs of Corollaries 5.6 and 5.7 Appendix J.4.
Corollary 5.6. Under the condition of Theorem 5.4, we have that with probability 1 - O vp), the difference between the stochastic worst vorst-direction sharpness of the iterates and the corresponding scaled trace of Hessian along the limiting flow is at most O ((7)/4 + )p²), that is,
Eg R (x:(TT// (np (1M)
= =0())(++)))))))
Corollary 5.7. Define U¹ as in Theorem 4.2, suppose (co) = lim X(t) exists and is minimizer of t-000
L(x)) in U" nr. Then for all > O. there exists constani T, > 0, such that for all p,d such that (n+ p) (1//pp)) are sufficiently small. we have that with probability 1 - ( + vp).
Max ((TTT((10)))))) S ep² inf [ Max (z)] Ee[L
=EU'
6 CONCLUSION
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, show the explicit biases of worst,, ascent- and average-direction sharpness around the manifold of minimizers 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.
The most interesting future work to generalize the current analysis for stochastic SAM to arbitrary batch size. This is challenging because, without the alignment property which holds automatically with batch size 1, such an analysis essentially requires understanding the stationary distribution of the gradient direction along the SAM trajectory. normalization layers, momentum, and weight decay into the current analysis.
also interesting to incorporate other features of modern deep learning like
Another interesting open question is to further bridge the difference between generalization bounds and the implicit bias of the optimizers. Currently, the generalization bounds in Wu et al. (2020); Foret et al. (2020) only work for the randomly perturbed model. Moreover, the bound depends on the average sharpness with finite whereas the analysis of this paper only works for infinitesimal whether the generalization error of the model (without perturbation) can bounded from above by some function of the training loss, norm of the parameters, and the trace of the Hessian.
It's an interesting open question

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
We thank Jingzhao Zhang for helpful discussions. The authors would like to thank the support from NSF IIS 2045685.
REFERENCES
Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-awe aware minimization. In International Conference on Machine Learning, pp. 639-668. PMLR, 2022. 3
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent edge of stability in deep learning. arXiv preprint arXiv:2205.09944: 2022. 3, 4, 6, 14, 15, 16, 20, 21, 22, 58
Peter L. Bartlett, Philip M Long, and Olivier Bousquet. The dynamics of sharpness-aware minimization: Bouncing across ravines and drifting towards wide minima. arXiv preprint arXiv:2210.01513, 2022. 14,
15
Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by omstein-uhlenbeed like process. arXiv preprint UrXiv:1904.090500 2019. 14, 16
Vivek S Borkar. Stochastic approximation: a dynamical systems viewpoint, volume 48. Springer, 2009. 70
Vivek S Borkar, Jervis Pinto, and Tarun Prabhu. A new learning algorithm for optimal stopping. Discrete
Event Dynamic Systems, 19(1):91-113, 2009. 15
Jeremy M. Cohen, Simran Kaur, Yuanzhi Li, J. Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs the edge of stability, 2021. 14
Jeremy M Cohen, Behrooz Ghorbani, Shankar Krishnan, Naman Agarwal, Sourabh Medapati, Michal Badura, Daniel Suo, David Cardoze, Zachary Nado, George E Dahl, al. Adaptive gradient methods at the edge of stability. arXiv preprint 22X 14484, 2022. 14
Yaim Cooper. The loss landscape of overparameterized neural networks. arXiv preprint arXiv:1804.10200 2018. 8
Alex Damian, Tengyu Ma, and Jason Lee. Label noise sgd provably prefers flat global minimizers, 2021. 6, 14, 16
Alex Damian, Eshaan Nichani, and Jason D Lee. soff-stabilization:: The implicit bias of gradient descent at the edge of stability. arXiv preprint arXiv:2209. 15594, 2022. 14
Chandler Davis and William Morton Kahan. The rotation of eigenvectors by perturbation. ii. SIAM Journal on Numerical Analysis. 7(1):1-46, 1970. 69
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In Proceedings of the 34th International Conference Machine Crarning-Volume 70, pp. 1019- 1028. JMLR. org, 2017. 1,2
Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred Hamprecht. Essentially no barriers in neural network energy landscape. In International conference on machine learning, pp. 1309-1318. PMLR, 2018.
John Duchi and Feng Ruan. Stochastic methods for composite and weakly convex optimization problems. SIAM Journal Optimization, 28(4):3229-33) 3259, 2018. 15
Gintare Karolina Dziugaite and Daniel M Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. arXiv preprint arXiv:1703.11008, 2017. 1.2
Benjamin Fehrman, Benjamin Gess, and Arnulf Jentzen. Convergence rates for the stochastic gradient de- scent method for non-convex objective functions. Journal of Machine Learning Research, 21:136, 2020.
2, 3
10

---

Published as a conference paper at ICLR 2023
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-awarr minimization for efficiently improving generalization. arXiv preprint arXiv:2010.01412, 2020. 9, 16
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In International Conference Learning Representations, 2021. 1, 2, 4,7
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P Vetrov, and Andrew G Wilson. Loss surfaces, mode connectivity, and fast ensembling of dnns. Advances in neural information processing systems, 31, 2018.
E. Hairer, S.P. Nørsett, and G. Wanner. Solving Ordinary Differential Equations I: Nonstiff Problems. Springer Series in Computational Mathematics. Springer Berlin Heidelberg, 2008. ISBN 9783540566700. URL https: /bbook.......o/.ooo...o...//// google.
Thomas P Hayes. A large-deviation inequality for vector-valued martingales. Combinatorics, Probability and Computing, 2003. 69
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Computation, 9(1):1-42, 1997. 2
Roger A. Horn and Charles R. Johnson. Matrix analysis. Cambridge university press, 2012. 69
Stanistaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in sgd. arXiv preprint arXiv:1711.04622 2017. 2
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi. Dilip Krishnan, and Samy Bengio. Fantastic generaliza- tion measures and where to find them. arXiv preprint arXiv:1912.02175 2019. 1,2
Simran Kaur, Jeremy Cohen, and Zachary C Lipton. On the maximum hessian eigenvalue and generalization. arXiv preprint arXiv:2206. 10654, 2022.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint 1609 04836, 2016. 1, 2.5
Harold Kushner and G George Yin. Stochastic approximation and recursive algorithms and applications, volume 35. Springer Science & Business Media, 2003. 15
Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and Kwon Choi. Asam: Adaptive sharpness-awar min- imization for saleeinvariian learning of deep neural networks. In International Conference on Machine Learning, pp. 5905-5914. PMLR, 2021. 2
Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and adaptive stochastic gradient algorithms. In International Conference on Machine Learning, pp- 2101-2110. PMLR, 2017. 15
Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and dynamics of stochastic gradient
algorithms i: Mathematical foundations. The Journal of Machine Learning Research, 2019. 15
Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after sgd reaches zero loss? mathematical framework. In International Conference on Learning Representations, 2021. 2, 3, 14, 16, 21, 22
Zhouzi Li. Zixuan Wang, and Jian Li. Analyzing sharpness along gd trajectory: Progressive sharpening and edge of stability, arXiv preprini arXiv:2207. 12678, 2022. 14
Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, and Yang You. Towards efficient and scalable sharpness-aware minimization. In Proceedings of the IEEF/CVF Conference on Computer Vision and Puttern Recognition, pp. 12360-12370, 2022. 2
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normalization layers: Sharpness reduction. arXiv preprint arXiv:2206.070085, 2022. 6, 14

---

Published as a conference paper at ICLR 2023
Chao Ma and Lexing Ying. On linear stability of sgd and imput-smoothness of neural networks. Advances in Neural Information Processing Systems, 34: 14:16805-66817, 2021. 14
Chao Ma, Lei Wu, and Lexing Ying. The multiscale structure of neural network loss functions: The effect optimization and origin. arXiv preprint arXiv:2204. 1326, 2022. 14
Jan R Magnus. On differentiating eigenvalues and eigenvectors. Econometric theory, 122:1799-111 1985. 70
Stephan Mandt, Matthew D Hoffman, and David M Blei. Stochastic gradient descent approximate bayesian inference. Journal of Machine Learning Research, 18:1-35, 2017. 15
Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring generalization in deep learning. In Advances in Neural Information Processing Systems, pp. 5947-5956, 2017. 1,2
Matthew D Norton and Johannes O Royset. Diametrical risk minimization: Theory and computations. Ma- chine Learning, pp. 1-19, 2021.
Antonio Orvieto, Anant Raj, Hans Kensting, and Francis Bach. Explicit regularization in overparametrized models via noise injection. arXiv preprint arXiv:2206.04613 2022. 5
Weijie Su, Stephen Boyd, and Emmanuel Candes. A differential equation for modeling nesterov's 's accelerated gradient method: Theory and insights. In Advances in Neural Information Processing Systems, pp. 2510 2518, 2014. 15
Colin Wei and Tengyu Ma. Data-dependent sample complexity of deep neural networks via lipschitz aug-
mentation. In Advances in Neural Information Processing Systems, pp. 9722-9733, 2019a. 2
Colin Wei and Tengyu Ma. Improved sample complexities for deep networks and robust classification an all-layer margin. arXiv preprint wrXiv:1910.422800 2019b. 2
Dongxian Wu, Shu-Tao Xia, and Yisen Wang. Adversarial weight perturbation helps robust generalization.
Advances in Neural Information Processing Systems, 33:2958-29988 2020. 1. 4, 16
Lei Wu, Chao Ma, and Weinan How sgd selects the global minima over-parameterized learning: A dynamical stability perspective. Advances in Neural Information Processing Systems, 31, 2018. 14
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learn- ing requires rethinking generalization. arXiv preprint arXiv: 1oXiv:6611.005500 2016.
Yang Zhao, Hao Zhang, and Xiuyuan Hu. Penalizing gradient norm for efficiently improving generalization in deep learning, arXiv preprint xxXiv:202.003999 2022. 3
Yaowei Zheng, Richong Zhang, and Yongyi Mao. Regularizing neural networks via adversarial model per- turbation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8156-8165, 2021. 1,2
Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha Dvornek, Sekhar Tatikonda, James Dunean, and Ting Liu. Surrogate gap minimization improves sharpness-aware training. arXiv preprint arXiv:2203.00 08065, 2022. 3,5
12

---

Published as a conference paper at ICLR 2023
CONTENTS
Introduction
1
2 Related Works
2
3 Notations and Assumptions
3
4 Explicit and Implicit Bias in the Full-Batch Setting
4
4.1 Worst- and Ascent-direction Sharpness Have Different Explicit Biases
4.2 SAM Provably Decreases Worst-direction Sharpness
5
4.3
Analysis Overview For Sharpness Reduction in Phase II of Theorem 4.5
Explicit and Implicit Biases in the Stochastic Setting
5.1
Stochastic Worst-, Ascent- and Average- direction Sharpness Have the Same Explicit Biases as Average Direction Sharpness
5.2 Stochastic SAM Minimizes Average-direction Sharpness
8
Conclusion
A Additional Related Works
14
B Experimental Details for Figure 1
15
C Additional Preliminary
15
Well-definedness of SAM
16
Proof Setups
E.1 Proofs of Theorems 5.2 and E.2
20
F Properties of Limiting Map of Gradient Flow, +b
21
G Analysis for Explicit Bias
23
G.1 A General Theorem for Explicit Bias the Limit Case
G.2 Bad Limiting Regularizers May Not Capture Explicit Bias
25
G.3 Proof of Theorem G.6
26
G,4 Proofs of Corollary G.7
28
G.5 Limiting Regularizers For Different Notions of Sharpness
G.6 Proof of Theorems 4.2 and 5.3
31
H Analysis Full-batch SAM on Quadratic Loss (Proof of Theorem 4.8)
31
H.1 Entering Invariant Set
H.2 Alignment to Top Eigenvector
33
13

---

Published as a conference paper at ICLR 2023
I
Analysis for Full-hatch SAM General Loss (Proof of Theorem 4.5)
38
L11
Phase (Proof Theorem 1.1)
39
1.1.1
Tracking Gradient Flow
39
L1.2
Decreasing Loss
40
L1.3
Entering Invariant Set
................................
41
1.2
Phase II (Proof of Theorem 1.3)
48
I.2.1
Alignment to Top Eigenvector
48
1.2.2
Tracking Riemannian Gradient Flow
54
1.3
Proof of Theorem 4.5
55
1.4
Proofs of Corollaries 4.6 and 4.7
55
1.5
derivations for Section 4.3
56
Analysis for 1-SAM (Proof of Theorem 5.4)
57
J.1
Phase I (Proof Theorem J.1)
J.1.1
Tracking Gradient Flow
58
................................
J.1.2
Decreasing Loss
59
J.2
Phase II (Proof of Theorem J.2)
66
J.2.1
Convergence Near Manifold
J.2.2
Tracking Riemannian Gradient Flow
67
J.3
Proof of Theorem 5.4
J.4
Proofs of Corollaries 5.6 and 5.7
68
J.5
Other Omitted Proofs for 1-SAM
68
K
Technical Lemmas
Omitted Proofs on Continuous Approximation
A ADDITIONAL RELATED WORKS
Implicit Bias of Sharpness Minimization. Recent theoretical works (Blane et al., 2019; Damian et al., 2021; Li et al., 2021) show that SGD with label noise implicitly biased toward local minimizers with a smaller trace of Hessian under the assumption that the minimizers locally connect that normalized GD implicitly penalizes the largest eigenvalue of the Hessian. Ma et al. (2022) argues that such sharpness reduction phenomena can also be caused by multi-scale loss landscape. Lyu al. (2022) show that GD with weight decay on a scale invariant loss function implicitly decreases penalize the spherical sharpness, i.e., the largest eigenvalue of the Hessian evaluated at the normalized parameter.
manifold. Arora et al. (2022) show
Another line of works study the sharpness minimization effect of large learning rate assuming the (stochastic) gradient descent converges the end of training, where the analysis is mainly based on linear stability (Wu et al. 2018; Cohen et al., 2021; Ma et al., 2021; Cohen et al., 2022). Recent theoretical analysis (Damian et al., 2022; Li et al., 2022) show that the sharpness minimization effect of large learning rate in gradient descent do not necessarily rely on the convergence assumption and linear stability via four-phase characterization of the dynamics at the so-called Edge of Stability regime (Cohen et al., 2021).
Comparison with concurrent work Bartlett et al. (2022). Bartlett et al. (2022) prove that quadratic loss, the iterate SAM (Fquation 11) and its gradient converges the top eigenvector Hessian, which is almost the same our Theorem 4.8. Assuming such alignment for a general loss, the work of Bartlett et al.
14

---

Published as a conference paper at ICLR 2023
(2022) shows that the largest eigenvalue of Hessian decreases in the next step. This paper also proves such a Hessian-gradies alignment for general loss functions (Lemma 1.19) 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.5), which are not shown in Bartlett et al. (2022). Moreover, this paper also characterize implicit bias of stochastic SAM with batch size 1, which minimizing the average direction sharpness, while Bartlett et al. (2022) only considers the deterministic case.
Comparison with Arora et al. (2022). Our proof uses a similar framework as Arora et al. (2022). However, our analysis has its own difficulty for the following reasons. First, Arora et al. (2022) only deal with the 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, First, the alignment analysis is more complicated because we have two eeperparametees_lemmiii rate 1j and perturba- tion radius p. while Arora et al. (2022) only needs to deal with one hyperparameter, learning rate 17. Second, the mechanism of penalizing worst-direction sharpness is different, which can be seen from the dependency of the sharpness-reduction rate over learning rate 1j. In Arora et al. (2022), normalized GD reduces the sharp- ness via a second-order effect of GD and thus the sharpness is reduced by O(r²) per step. In our analysis, for fixed small perturbation radius p. the sharpness is reduced by ((p²n) per step, which is linear in 7/-
Analyzing Discrete-time Dynamics via Continuous-tine Approaches. There is a long line of research that shows the trajectory of stochastic discrete iterations with decaying step size eventually tracks the solution of some ODE (see Kushner et al. (2003); Borkar et al. (2009); Duchi et al. (2018) and the reference therein). However, those results mainly focus on the convergence property of the stochastic iterates (e.g., convergence to stationary points), while we are interested in characterizing the trajectory especially when the process is running for a long time even after the iterate reaches the neighborhood of the manifold of stationary points.
Recently there has been an effort of modeling the discrete-time trajectory of (stochastic) gradient methods by continuous-time approximation (Su et al., 2014; Mandt et al., 2017: Li et al., 2017: 2019). Notably, Li et al. (2019) presents a general and rigorous mathematical framework to prove such continuous-time approximation. More specifically, Li et al. (2019) proves for various stochastic gradient-based methods, the discrete-time weakly converges to the continuous-time one when LR 7 - 0 in ((1/n) steps. The main difference between our results with these results (e.g.. Theorem 9 in Li et al. (2019)) is that we focus on much longer training regime, i.e., T = O1) p ?) steps where the previous continuous timme approximation results no longer holds throughout the entire training. As a result, their continuous approximation is only equivalent to the Phase dynamics in our Theorems 4.5 and 5.4 and cannot capture the dynamics of SAM in Phase II, when the sharpness-reducttee rrduccion implicit bias happens. The latter requires a more -inengrieeed analysis to capture the effects of higher-order terms in 17 and p SAM Equation 3.
B EXPERIMENTAL DETAILS FOR FIGURE 1
In Figure we choose Fi (x) of 6x 8 and (x = 4(1 (1- The loss L has zero
loss manifold {a 0} and the eigenvalues of its Hessian on the manifold are Fi (x)) and (2(x) with Fi (x) >
6 > and the two non-zero eigenvalues of L of any point x on the manifold are X (V2L(z)) = Fi (11, x2) and
F2(x)) on [0, 1j2. The loss L has a zero loss manifold {2g xig = 0} of codimension M 2
X2( 2 (xx) 22).
As our theory predicts,
1. Full-batch SAM (Equation 3) finds the minimizer with the smallest top eigenvalue (x), which is 21
0, 22- = 03 = 0, 24 = 0;
2. GD on ascent-direction loss Ase (2) finds the minimizer with the smallest bottom eigenvalue, F2(x),
which is 1 = 1, .z2 1,3 = 0, .4 = 0:
3. Stochastic SAM (Fquation 13) (with To (2, y) = Fi (2:)// La (x, v) = Fz(x))²) finds the minimizer with smallest trace of Hessian, which is 21 4/5, 1/7, x3 0, = 0.
C ADDITIONAL PRELIMINARY
In this section, we introduce some additional notations and clarification before the proof.
We will first give the detailed definition of differentiable submanifold.
15

---

Published as a conference paper at ICLR 2023
Definition C.I (Differentiable Submanifold of RD.). We call a subset T C RDD C* submanifold of TR. and only for every x E r, there exists a open neighborhood U of and an invertible C* map vp : U TR such that (L' UU) = (IR" {0)) N (U))
Necessity of Manifold Assumption. 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 in the end-to-end analysis. So far almost all analysis of implicit bias for general model parameterizations relies on Taylor expansion, e.g. Blanc et al. (2019); Damian et al. (2021); 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 p. On the contrary, if the set of global minimizers are a set of discrete points, then with small perturbation radius p. implicit bias of optimizers not sufficient to drive the iterate from global minimum to the other one.
Implicit versus Explicit Bias. If an algorithm or optimizer has a bias towards certain type of global/local minima of the loss over other minima of the loss, and this bias is not encoded in the loss function, then we call such bias an implicit bias. On the other hand, a bias emerges as solely a consequence of successfully minimizing certain regularized loss regardless of the optimizers (as long the optimzers minimize the loss), we say such bias is an explicit bias of the regularized loss (or the regularizer).
As a concrete example, will prove that full-batch SAM (Equation 3) prefers local minima with certain sharpness property. The bias stems from the particular update rule of full-batch SAM (Equation 3), and not all optimizers for the intended target loss function A (Equation 2) has this bias. Therefore, it's considered as an implicit bias. As an example for explicit bias, all optimizers minimizing a loss combined with & regularization will prefer model with smaller parameter norm and this is considered as an explicit bias of & regularization.
Usage of (-) Notation: Our analysis assumes small 17 and p while treating all other problem-dependent pa- rameters as constants, such as the dimension of parameter space and the maximum possible value of deriva- w.), tives (of different orders) of loss function and the limit map Tb. In (()) w(). (()), we hide all the dependency related to the problem, e.g., the (unique) initialization Tinit. the manifold I, compact set U in Theorem 4.2, and the continuous time T3 in Theorems 4.5 and 5.4, and only keep the dependency on p and 17. For example, (p)) is placeholder for some function g(p) such that there exists problem-dependent constant C > 0, Vp > 0, [9() S C))( In informal equations such as Equation 31 in the proof sketch section, we are a bit more sloppy and hide dependency on xtt) in (+) notation as well. But these will be formally dealt with in the proofs.
Ill-dfinedness of SAM with Zero Gradient. The update rule of SAM (Equations 3 and 13) is ill-defined when the gradient is zero. However, our analysis in Appendix D shows that when the stationary point of loss L, {x| VL(x) = 0}. is a zero-measure set, for any perturbation radius p. except for countably many learning rates, full-batch SAM is well-defined for almost all initialization and all steps (Theorem D.1). A similar result is shown for stochastic SAM if the stationary points of each stochastic loss form a zero-measure set (Theo- rem D.2). Thus SAM is generically well-defined. For the sake of rigorousness, when SAM encountering zero gradients, we modify the algorithm via replacing the ill-defined normalized gradient by an arbitrary vector with unit norm and our analysis for implicit bias of SAM still holds.
D WELL-DEFINEDNESS OF SAM
In this section, we discuss the well-definedness of SAM. When Vl(x)= = 0, SAM (Equation 3) is not well- defined, because the normalized gradient L(w) is not well-defined. The main result of this section are Theorems D. and D.2, which say that (stochastic) SAM starting from random initialization only has zero probability to reach points that SAM is undefined (i.e., points with zero gradient), for all except countably many learning rates. These results follow from Theorem D.3, which is more general theorem also applicable to other discrete update rules as well, like SGD. Note results in this section does not rely on the manifold assumption, i.e., Assumption 3.2. We end this section with a concrete example where SAM is undefined with constant probability, suggesting that the exclusion of countably many learning rates are necessary in Theorems 1. and D.2.
Theorem D.1. Consider any C2 loss L with zero-measure stationary set fxr [((x) = 0). For every p > 0. except countably many learning rates, for almost all initialization and all i, the iterate of full-baich SAM (Equation 3) z(t) has non-zero gradient and is thus well well-defined.
16

---

Published as a conference paper at ICLR 2023
Theorem D.2. Consider any C² losses (T&)& with zero-measure stationary set fa: VLg(2)) 0} for each E [M. For every > O, except countably many learning rates v, for almost all initialization and all with probability one of the randomness of the algorithm, the iterate of stochastic SAM (Equation 13) x(t) has non-zero gradient and is thus well-defined.
Before present our main theorem (Theorem D.3), we need to introduce some notations first. For a map F mapping from IRD \\ Z + IRD we define that F3 IDD Z - RD as (x)( - a +)((x) for any 1) E R+
Given a sequence of functions {F"] [F we define F (x: - x f (x), , for any x E TRP. We further
define that (x) 2 " (z))
(x)) for any n 2 1 and that F", (2)2) I.
Theorem D.3. Lei Z be a closed subset of with zero Lebesgue measure and p. be any probability measure
on RD that is absolutely continuous to the Lesbegue measure. For any sequence of C¹ functions F RP R n E N+ the following claim holds for all except countably many 17 E R+
P (IEERD) 3n E F, E Z and 4O sis n 1 1, $ 2}) = 0.
Sisn-in
In other words, for almost all n (except countably many positive numbers), iteration x(t + 1) x(1) (x(()))) = F (x(0) will not enter Z almost surely, provided that z(0) is sampled from p.
Theorem D.I and Theorem D.2 follows immediately from Theorem D.3.
Proof of Theorem D.I. Let F(x) = (L + P PIVLEN and Z = {x E R (L(x)= = =0. We can easily
VLiz
check F is C on RD, Z and by assumption Z is a zero-measure set. Applying Theorem D.3 with F F for all n E N+, we get the desired results.
Proof of Theorem D.2. Let (x) = VL(x + P (La(2)
(x)il ) and Z = u E RD VLx(2) = 0}. We
can easily check F& is C² on D \ Z and by assumption is a zero-measure set. Applying Theorem D.3
with Fr = GK. for all n E N where kin is the nith data/batch sampled by the algorithm, we get the desired results.
Now we will turn to the proof of Theorem D.3, which is based on the following two lemmas.
Lemma D.A. Let Z be a closed subset of IR with zero Lebesgue measure and F RD \ Z + R
P be continuously differentiable function. Then except countably many 1i E RR+ {x E Rd Z det( = 0} is zero-measure set under Lebesgue measure.
Lemma D.5. Let Z be a closed subset of RP with zero Lebesgue measure and H R Z + R be continuously differentiable function. If {x E Hd \Z % det(0H = 0} is a zero-measure set, then for any zero-measure set Z', 11 (1') is a zero-measure set.
Proof of Theorem D.3. It suffices to prove that for every N E N+ at most for countably many
p R F (x) (x)
Z and VO < i i N 1, F (z) & z)) 0.
(15)
The desired results is immediately implied by the above claim because the countable union of countable set is still countable, and countable union of zero-measure set is still zero measure.
To prove Equation 15, we first introduce some notations. For any 7) > 0, 0 We extend the definition to set in natural way, namely F,," Th (S) A UxESF (xx for any S C RPP Under this notation, we have that
N 1, and E RD,,
define F., (1+11
(x:) (FNN-n)-1(()) (I:),, where (a((a)
F N (Z) = ( (
(x) E Z and vO < N 1,F 2))
We will prove by induction. We claim that for each 0 < < N except for countably many 17 E R+ FR (Z) has zero Lebesgue measure. The base case n = 0 is by trivial as Z is assumed to be zero mero-measure Suppose
'Though we call Equation 13 I-MAM, but our result here applies to any batch size where La can be regarded as the loss for k-thh possible batch and M is the number of the total number of batches.
17

---

Published as a conference paper at ICLR 2023
" (Z) det (z)) 0} is a zero-measure set. Next by Lemma D.5 for some E RR+ f Rd F (Z) (at, (x))
this holds for By Lemma D.4, except countably many E RT, for E B F, "
0} is zero-measure set, then n --- (Z) (F)) (F, (2)) is a 7ero-measure set. Then by induc- tion, know that except countably many 7) E R for all integer 0 S n N, F (Z) is zero-measure. Since μ is absolutely continuous to Lebesgue measure, (((F, (Z)) = 0.
We end this section with the proofs of Lemmas D.4 and D.5.
Proof of Lemma D.4. We use Ax(x)) to denote that the real part of the ith eigenvalue of the matrix OF(z) in
the descending order. Since (F(x) is continuous in x. A²(x)) is continuous in æ as well, for any i [D] and thus {æ E RD' > Z| n, let I, be the set of 17 where μu (2 RP \ Z (x(x) 1/m)) 1/n, then [La S n, because
(z) = 1/n} is measurable set. Note that for fixed i E [D], for each positive integer
[nn|
S Σ
u((f RR Z (x)) = 1/nb) - ((( \ Z 1/A(z)) E 1,) 1
n
MET
Therefore, there are at most countably many 1 E Rt, such that pffz E R Z (x = 1/n > 0.
Further note that det(aF, (x)) = 0 ž E [D] (z) 1/7, we know that there most countably many 1) E Rt, such that pl E RD
Z det(OF, = 0}) 0. This completes the proof.
Proof of Lemma D.S. Denote {x E RD! Z det(aH(x)) = = by Z". since (x)( is continuous in x as F is Ct_ z" is relatively closed in RDD Z. Since 7' is a closed set, TR. \ (Z' U 2") is open. Thus for all x RDD 3 (Z' U Z") with det( (x)) o 0, there exists open neighborhood of x; U, where for all a' E U, det( (z')) & 0, since thus det (3H(()) is continuous. This further implies H is invertible on U and its inverse (u -1 is differentiable on F(U). Therefore, ( maps any zero-measure set zero- measure set. In particular, (H y) 'Z' n (U)) is zero measure, so is
nu C (H]c) (Z' n H(U)). Now for every x E R " '2 Z we take an open neighborhood U2 S R. (Z' U (Z).. Since R is separable metric space, the open cover of R.,, (U,
(U has countable subcover, {U,] where I is countable set of TR D (Z' U ) Therefore we have that H-1 (Z')
(Z' U ) HZ') n (R.) (Z' U Z))) = Uzer H (Z') ) i U2 is a mero-measure set. Thus H (Z') is also zero-measure since Z', z" are both zero-measure. This completes the proof.
We end this section with an example where SAM is undefined with constant probability.
Theorem D.6. For any nh P 0, there is C2 loss function L R - R satisfying that (1) L has unique stationary point and (2) the set of initialization that makes SAM with learning rate v and perturbation radius P to reach the unique stationary point has positive Lebesgue measure.
Proof of Theorem D.6. We first consider the case with
1 with for x (-00, -2); for x E [-2, 2]: for x E (2, 20).
x2/2+2+ 1/2,
L(x) =
(16)
x2/2-xx++
x 1/2,
We first check L is indeed C²: I.(2) = 1122 (2))
=
-I(
and I,"(2) = (-" = 1.
Now we claim that for all (x) > 2. x(1) = 0. which is stationary point. Note that L is even and monotone increasing on [0, ).). we have VL(=)/|V((z) = sign(x). Thus for [x(()] > 1, it holds that
(a=(t) + sign(x(t)) > 2 and therefore
z(t + 1) =x(t) nL((((t) + +V((=)/\V(()))
=z(t) - L' L'(x(t) + sign(x(t)))
=x(1) - (x)))) + sign(((()) sign( (x(() + sign(x(1)))
=z(t) (t) 0.
(17)
Now we turn to the case with arbitrary positive p, It suffices to consider (z) 2 (,)()). We can use the calculation for p = 17 to verify for any |z| 2p.
1 L(a/p p- + sign(x)) x n
I' (x + psign(Lin, (x)) L' +
p psign(x)) = 1
18

---

Published as a conference paper at ICLR 2023
namely 1 psign ( 0. This completes the proof.
A common (but wrong) intuition here is that, for continuously differentiable update nule, as long as the points where the update rule is ill-defined (here it means the points with zero gradient) has zero measure, then almost surely for all initialization, gradient-based optimization algorithms like SAM will not reach exactly at any stationary point. However the above example negate this intuition. The issue here is that though differentiable map (like SAM
vVL( + P (IVL(() .) always maps the zero-measure set to zero-measure set, the preimage zero-measure set is not necessarily zero-measure, as the map I + x V PVV]+]]]V][]]]]) [T]]]]] is not necessarily invertible. The update rule of SAM is not invertible at 0 is exactly
VL(w)
vizz)
the reason of why preimage of has positive measure.
E PROOF SETUPS
In this section provide details of proof setups, including notations and assumptions/settings.
We first introduce some additional notations that will be used in the proofs. For any subset $ E IR. define (2st S) & infyes || - yll2- For any d > 0 and any subset S E IR,,, we define S & {x E TR dist(z; S) S d}. Our convention is to use K to denote compact and U to denote an open set.
D we
Below we restate our main assumption in the full-batch case and related notations in Section 3. Throughout the analysis, we fix our initialization Tinit loss function as L RDD IR.
Assumption 3.2. Assume loss L : RP + R is CA. and there exists C² submanifold I of RD that is (D M) dimeensona for some integer 1 S M S D, where for all x E r, x is a local minimizer of L. and rank V [x)) M.
Notations for Full-Batch Setting: Given any point E T, define Par the projection operator onto the manifold of the normal space of I at and P.r In Pa... Given the loss function L, its gra- RDD 00 oc) + RDD Here, ((x,T) denotes the iterate at time T
dient flow is denoted by mapping of a gradient flow starting at x and is defined as the unique solution of ((2, x - So" V (L f)) Var E RRP. We further define the limiting map of (xx.)) as (xx) = lim, 6(x; r), that is, ((x) de- notes the convergent point of the gradient flow starting from For convenience, we define (x()), (x(2)) as (((())) whenever the latter is well defined. When ztt) and I' is clear from context, Po(u())),
> ((V)L((z)) Us we also use (t) = (t)) (t)))= = vs (x(()), P.t P. (d(z(()),)
Definition 3.3 (Attraction Set). Let be the attraction set of under gradient flow, that is, neighborhood of I containing all points starting from which gradient flow w.r.. loss L converges some point in I, or mathematically, U {ar E (zx) exists and (xx) T).
Below we restate the setting for stochastic loss of batch size one in Section 5.
Setting 5.1. Let the total number of data be M. Let fi(xx) be the model output on the k-th data where J is C -mmooth function and yv, be the k-th label, for k -1....,
........ M.. We define the loss on the k-th data L,(2)) = ((f((x)), yke and the total loss L EMd Lx/M, where function Ely/ y) CA -mmoooth y'- We also assume for any y E. it holds that arg siggii/ee Ex E(y', .v) = y and that o2 w) (oy"'" by'=y O. Finally, we
denote the set of global minimizers of with full-rnkk Jacobian by T and assume that it is non-empty, that is,
2 {- E RD fx(x)) yk, vk E [M] and { are linearly independent) * 0.
Theorem 5.2. Loss set I and integer M defined in Setting 5.1 satisfy Assumption 3.2.
In analysis, we prove our main theorems in the stochastic setting under more general condition than Setting 5.1, which is Condition E.1 (on top of Assumption 3.2). The only usage of Setting 5. in the proof is Theorems 5.2 and E.2.
Condition E.1. Total loss L= E Lt- For each k E [M]. La is C², and there exists a (D 1)- dimensional C2 cu-uumannifold RR Lkk where for all x E LR x is global minimizer of Lu. L (x) = and rank( 2 La (x)) = 1. Moreover, I = n I'u for I defined in Assumption 3.2.
Theorem E.2. Setting 5.1 implies Condition E.I.
19

---

Published as a conference paper at ICLR 2023
Notations for Stochastic Setting: Since L. is rank-1 on FE for each A: [M].
write it [₂(x))
A (z) (x) for any z E J, where wi is a continuous function on with pointwise unit norm. Given
the loss function Lg, its gradient flow is denoted by mapping Φx RD x [0, 0o) + RD. Here, Φx T) denotes the iterate time T of gradient flow starting at x and is defined as the unique solution of ou 7)
Jo L (z. t)) Vz R.. We further define the limiting map 4y. (x(x) lim @& (2, T), that is, Du (x) denotes the convergent point of the gradient flow starting from Similar to Definition 3.3, we define Uy = {I E RD (x)) exists and 2u (a) ]'k] be the attraction set of T., We have that each UR is open and 4 is C UR by Lemma B.15 in Arora et al. (2022).
Definition E.3. A function L is p-PL in a set U vf E U, |VL()) 2n(((x)) - infsev L(x)). Definition E.4. The spectral 2-norm of a k-order lensor X Xi Rdd.x...dds is defined as
X
max E :|1z:/l==1
X [[s1,..,,. Tt).
Lemma E,5 (Arora et al. (2022) Lemma B.2). Given any compact set K S T, there exist r(K), ((),, E + such that
1. Kr(k) nT is compact.
KT(K) C U n (NkE(M Uk).
3. L is ((()) on KT(K)
4. inTre (A1 (AL())()))
{x(V²(())) A(K) > O.
5. inf AM (V2((x)
2 p K
> O.
infrEk)()) A1 2 Lx
2 p(K) 0.
Given compact set K C T. we further define
<(K) =EK-(K)
sup
[())ll v(K)=
sup *EKT(K)
[ T(K)
sup
|V*L(x)
=EK-(K)
/ (xx - V
{(K) = sup zEKr(K)
|V2 X(K)= = sup
Similarly, we use notations like (K). pn (K), quantities defined for stochastic loss L2 and its limiting map 4, for h E [M].
(K),, xa (K) denote the counterpart of the above
Lemma E.6 (Arora et al. (2022), Lemma B.5 and B.7). Given any compact subset K C T, let r(K) be defined in Lemma E.S, there exist 0 < h(K) < <rK) such that
sup L(x) inf L(x) S (KK) =gKh(K)
*EKM()
8
2. Vr E Kh(K), ,b()) E Kr(K)/2
3. Vx E Kh(K), - (x))l22 S
Su(K)²
c(K)u(K)
4. The whole segment (x)) lies in Kr(K), so does (bb())) for any k E [D].
The proof of the lemmas above can be found in Arora et al. (2022). Readers should note that although Arora et al. (2022) only prove these lemmas when K is special compact set (the trajectory of an ODE), all the proof does not use any property of K other than it is compact subset of I, and thus our Lemmas E.5 and E.6 hold for general compact subsets of r.
In the rest part of the appendix, for convenience we will drop the dependency on K various constants when there is no ambiguity.
E. PROOFS OF THEOREMS 5.2 AND E.2
Proof of Theorem 5.2. Define : RD + RM as [F(x)] = fa(2), E [M]. Let Tμ
span( (vV } 1 and T be the orthogonal complement of Tr in RDD. Now we apply implicit func- tion theorem on F at each x E Г. Without loss of generality (e.g. by rotating the coordinate system), we can D-M assume that 0, T2= RD- MM x (0). and that T {0} M Implicit function theorem ensures that
there are two open sets 0 E U C I D-D. AS and 0 E V C IRM and an invertible C4 map g U + V such that
F (r)
V) ((ug(w))
E U},
M is linearly independent for every x' E U x V. Thus by definition of P. it holds that I' (U V) (
where Y bi...., [9M] [M] E RRM Moreover,
n nU V) {u,g(u)) u E U}. Now for
20

---

Published as a conference paper at ICLR 2023
(1 E U x V. define 1 : U V - RD by o(u, - (u, - g(u)). We check that s is Ca and ( (U x V) = [uu v - g(u)) w= g(u). E ( [(u, 0) u u = = U x (0) = (IR DRD-M x (O)))(() This proves that I' is Ca submanifold of I of dimension D- M. (c.f. Definition C.I) Since arg / eR ((y',)) ,)
for any E IR, it is clear that Vo E T, x is a global minimizer of L. Finally we check the rank of Hessian of loss L. Note that for any I E T, 2 (z) Helg' aw) (āg')² lo'=v, Vfu(t)(VV
8 Elyl and that
(8y') y o'' UE > O, rank((" = ronk = M. This completes the proof.
Proof of Theorem E.2.
1. L = 1 Ex
M km L by definition.
2. 4k E [M]. Lx(x)= {(fx(z), we) is C as l and fa are both Ct.
3. For any I E r, by Lemma 5.5, we have fi (x) o 0. Then there exists an open neighborhood Vi2 such
that I' V22 and (x) * 0 for any Z E VR, k [M]. Then applying implicit function theorem as in the
proof of Theorem 5.2, for any & M there exists a (D - -)-dimensional C -manfolld 1% C Vt. such that for any x' E V, fa(('))
yk if and only if x' E T. As for any T C Vk, fa((²) yk. we can infer
that I C T. Then I C L T
4. For any x E T, we have f((x) = ye. which implies (x(x) = 0. Also as x E V.V.f(z)) + O.
By Lemma J.15, we have (ank(V2L(x))
F PROPERTIES OF LIMITING MAP OF GRADIENT FLOW, +
In our analysis, the property of & will be heavily used. In this section, we will recap some related lemmas from Arora et al. (2022), and then introduce some new lemmas for the stochastic setting with batch size one. Lemma F1 (Arora et al. (2022) Lemma B.6). Given any compact set K C for any x E Khh
do(x, t) di
2(L(x) L(((=)))
VL(x)
|| - (x)ll2 <
Lemma F2. Given any compact set K f.r any x E Khh
2(L(x) I(((x)))
[VL(x)]l S </I (10)llass
Proof of Lemma F22 The first inequality is by Lemma E.5 and Taylor Expansion. The second inequality is by Lemma F.I.
Lemma E3 (Arora et al. (2022) Lemmas B.16 and B.22).
d x L (z) 0, U; 04 (x) v L VL (z) 0& (x) 0² (V))( vi] P.. (A x E..
8 (z) VL))),
Lemma E44 (Arora al. (2022) Lemmas B.8 and B.9). Given any compact set K CT. any x E Kh.
5V
P P (x (xx))122 <
||
(x)
4p
VL ( (x (x- --x(()))
<
x -8-416 -(x))1²
2
IVL (x) l12
-1
-1
2v
|| ((x)]l2
[2I(((x)) (x))
"
(x))
(OL() +(x)) +1(1=~((0)11)) (x=)la).
(x) [VL(((()) (x- (x))l /
Lemma E.5 (Li et al. (2021), Lemma 4.3). For x T, 00(x) = -Ptt the orthogonal projection matrix onto the tangent space of I at Since d ( L(x) = 0.
21

---

Published as a conference paper at ICLR 2023
The proof of above lemmas can be found in Arora et al. (2022); Li et al. (2021). In the following, will first show the proof of Lemma 3.1
Proof of Lemma 3.1. Since is defined the limit map of gradient flow, it holds that for any 2 0, (1,(2, t)) = (x)). Differentiating both sides at = O, we have (#(6(2 0)) 00(x.)) completed by noting that 0z(z.1) (L) by definition of ø.
āt = 0. The proof is
ot
Lemma F... Given any compact set K TT any x E K",
|00 x V Lx (2) S (us + GxE) 4
L (x)
12 S + Gt) (1)]l2
V (x)
Proof of Lemmna F.6. By Lemma E.6 and Taylor Expansion,
[ V -
1=::: G&Ell= +x))11
(ux + - - (
this proves the first claim.
Again by Lemma E.5 and Taylor Expansion,
L (z)
L (x)
|0 1: (A
12 & [00x))V1+<((01²)
L1 (x) 1l2 + ux1l - -<(l==
L (x)
V (a)
S [100((())
112 (vx + (xE)ll> (xx)
VL (z) ||
= (vx + (E))2z 4(x)ll2
this proves the second claim.
Lemma E.7. Suppose Kh and y x V (x + (
11y 1 x < vl) ||2 + "
[1((x) - (o(x)-((y)1++ Enpl|VL ||2 unp ( 722
S <Empll= - - + - + unp² + 122
Proof of Lemma E.7. For sufficient small p. v(()) E K". By Taylor Expansion, |VL[aح|
(x)
ly 2|12 V + (IV((()) (x) la (x)) nCp
This further implies that for sufficiently small 7) and p, xy E K"
Again by Taylor Expansion,
(x)
(10(()) 2)1l2 (+0)))
[+( 21 n/2u//2
By Lemma F.3, (1)xx)(LL)))))) (z) 0 and 0 xz) V
VL (z)
0 -0) (z), VL (z)). Hence,
VL
VI((0)
[04())(y - 2)112 < ( [ [v]
lla n/2/22
{pl]VV(()) ||2 v/2
As ry K",by by Taylor Expansion,
(} (2)))12 < [ 2) + Elly x/1//2
22

---

Published as a conference paper at ICLR 2023
Putting together we have
(a:)
s = 12 1
Finally, by Lemma F.2, we have
(x) - (1))ll -
E
- -
+ + uup + EC2
This completes the proof.
Lemma E8. Suppose x & Kh and y = x V LA (x ||VVLl vL((s)
|yv - 2||2 - 7/ VL (x) 11 + n
| - ( & (V + nol (L(2)l22 + n
Proof of Lemma E8. For sufficient small p. x ([VL,())]
(L)(m)
K By Taylor Expansion,
L (z)
ly - x||2 = V x (
lla S n1 VL (z) 2l + n<p.
(z)
This further implies that for sufficiently small y) and p. ty E K" Again by Taylor Expansion,
VL (z) LL (x)
[0 - 2) S n1100(2)VLL) +
|12 + 1/22
We further have by Lemma F.1, [ L (x:)
(L) (x) (2) 1211= +(x)
VZ
+ w|| 2()) + CEl1= - b(x)112
[L(x)) (z))
Similarly,
L (x)
(p00(2) (z)
L (x) l12
L1 (z)
<11p00
l|2 + PCE [x 4(z)112
LL (x:)
V I. (x)
<1)00(((()))>>
112 + PCElLx - ((x)]12 + pullz ((x)12
LL (z)
So 10
This completes the proof.
G ANALYSIS FOR EXPLICIT BIAS
Throughout this section, we assume that Assumption 3.2 holds.
G.1 A GENERAL THEOREM FOR EXPLICIT BIAS IN THE LIMIT CASE
In this subsection we provide the proof details for section 4.1, which shows that the explicit biases three notions of sharpness are all different, using our new mathematical tool, Theorem G.6.
Notation for Regularizers. Let R TRD U (oo) a family of regularizers parameterized by p. If R2 is not well-defined at some I, then let Rz(x) oc. This convention will be useful when analyzing
23

---

Published as a conference paper at ICLR 2023
ascent-direction sharpness of T. L which is not defined when VT(x)) 0. This convention will not change the minimizers but however, when far away from manifold, there are regularizers Rp(z) of our interest that can actually e Tr()2T(a)). Therefore we make the following assumption to allow the be negative, e.g., *x) x 2D regularizer to be mildly negative.
the regularized loss. Intuitively, a regularizer should always be non-negative,
Condition G.1. Suppose for any bounded closed set B C. U, there exists C' > such that for sufficiently small p. Vx B, Rp(x)) CCp².
Definition 4.3 (Limiting Regularizer). We define the limiting regularizer of {Rp} as the function
- R, S(x) = lim lim inf
p0 ||z - Rp(z')/0².
The high-level intuition is that we want to use the notion of limiting regularizer to capture the explicit bias of Rş among the manifold of minimizers I' as p - O, which is decided by the second order term in the Taylor expansion, egg., Equation 5 and Equation 6. In other words, the hope is that whenever the regularized loss is optimized, the final solution should be in neighborhood of minimizer with smallest value of limiting regularizer S(x).. However, such hope cannot be true without further assumptions, which motivates the following definition of good limiting regularizer.
Definition G.2 (Good Limiting Regularizer). We say the limiting regularizer S of [Rp} J is good around some x' E r. if S is non-negative and continuous at x² and that there is an open set Vg- containing x", such that p2 converges uniformly to S(x) in for all eenn Va- as p - O.
for any C >>
In other words, a good limiting regularizer satisfy that for any C, E > O, there is some Pa > 0.
inf Yr EInV and p S P |5(x) |a |
R,(z))/o2<<
We say the limiting regularizer is good on I, if S is good around every point r. In such case we also say Rg admits S as a good limiting regularizer on r.
The intuition of the concept of good limiting regularizer is that, the value of the regularizer should not drop too fast when moving away from a minimizer x in its O(p) neighborhood. If so, the minimizer of the regularized loss may be (2(p) 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. (See Appendix G.2 for a counter example) We emphasize that the conditions of good limiting regularizer is natural and covers a large family of regularizers, including worst-, ascent- and average-directinn sharpness. See Theorems G.3 to G.5 below.
Theorem G.3. Worst- dirsctiin sharpness RMax admits x1 >1 LVV)())) as good limiting regularizer r and satisfies Condition G.I.
Theorem G.4. Ascent-direction sharpness Rose admits AM(V)L())/2 as a good limiting regularizer on r and satisfies Condition G.I.
Theorem G.5, Average-direction sharpness Ro admits (r(V)L())/(2)) as a good limiting regularizer on and satisfies Condition G.I.
Next we present the main mathematical tool analyze the explicit bias of regularizers admitting good limit- ing regularizers, Theorem G.6.
Theorem G.6. Let U" be any bounded open set such that its closure U' C U and U' any family of parametrized regularizers { {Rp} admitting a good limiting regularizer S(x) on and satisfying Condition G. for sufficiently small p. it holds that
r = nr. Then for
+ [[z))+R((z)) ZEU'
inf
inf L(x) - p² inf ZEU'
S(x) s so(p2)
sEU'rr
Moreover, for sufficiently small p. it holds uniformly for all E U that L(u) + RT(u) < inf (L(x) + R₂(x)) + +Opp) = R2(u)/p² "n
inf -o
xet"
'ere we implicitly assume the zeroth and first order term varnishes, which holds for all three sharpness notions.
24

---

Published as a conference paper at ICLR 2023
Theorem G.6 says that minimizing the regularized loss I(u)) + R,(u) is not very different from minimizing the original loss L(u) and the regularizer (p(u) respectively. To see this, we define the following optimality gaps
A(u) _ L(u) + Ru( inf (L(x) Rp(x)) 20 ZEt
B(u)= _ I(u) inf I(x) 20 ZEU"
C(u) _ R2 p= =Ei'Nr
inf S(x);
and Theorem G.6 implies that AA(u) B(u) p²C((u) o(p²). Moreover, A(u), B(u) are non-negative by definition, and C(u) N (1) are almost non non-negative, whenever A(u) is ((p"))aaproximatte) opti- mized.
For the applications we are interested in in this paper, the good limiting regularizer S can be continuously extended to the entire space R.". In such case, the third optimality gap has an approximate alternative form which doesn't involve Rp. namely S(u) - inf nr S(x). Corollary G.7 shows minimizing regularized loss L(u) + Rp(u) is equivalent to minimizing the limiting regularizer, Su)) around the manifold of local minimizer, r.
Corollary G.7. Under the setting of Theorem G.6, let S be continuous extension of S R For any optimality gap A > O, there is a function e : IR+ - RR+ with limp-so Ep) = 0, such that for all sufficiently small p > 0 and all U satisfying that
L(u) + Rp(u)= inf (L(x) Rp(w)) s Ap?;
=Et"
it holds that L(u)- inf L(x) S (A + <(p))p² and that
5(u) inf ((x) E [--())), + c())]
=Ei'rr
G.2 BAD LIMITING REGULARIZERS MAY NOT CAPTURE EXPLICIT BIAS
In this subsection, we provide an example where a bad limiting regularizer cannot capture the explicit bias of regularizer when P + 0, to justify the necessity of Definition G.2. Here a bad limiting regularizer is limiting regularizer which is not good.
Consider choosing (1) = L(x pe) L(x) with ||e]| = 1 as fixed unit vector. We will show minimiz- ing the regularized loss L(x) + R2(x) does not imply minimizing the limiting regularizer of R2(x) on the manifold.
By Definition 4.3 and the continuity of Rp, the limiting regularizer of Rp is Er, S(x) = lim lim inf Rx(x')/0² lim Rx(x)/b = (()(ee,, 0. 8-0
p->0 ||a =21l25r
However, for any I, we can choose x' L(x') + R2(x') L(x' + pe) = L(x) = 0.
pc, then
Therefore, no matter how small is, minimizing L(x) + Rx(x) can return solution which is p-close to any point point of I. In other words, the explicit bias of minimizing L(x) + (x) is trivial and thus is not equivalent to minimizing the limiting regularizer S on the manifold l'.
The reason behind the inefficacy of the limiting regularizer in explaining the explicit bias of R, is that S(x) is not good limiting regularizer for any E I satisfying S(x) > 0. To be more concrete, choose C = 1 and e = S(z) 22i in Definition G.2. For any x E I and sufficienlly small p > O, considering = x - pel, by Taylor Expansion,
R₂(x') L(x' + pc) L(x')
+
= - o(0²) = d
=
This implies inf |r'= Ra(x') S R2(x1) =0p²). Hence,
S(x)
R₂(x1)/0² 2 S(x) - o(1) > S(x)/2 =c. || SPp
inf
25

---

Published as a conference paper at ICLR 2023
G.3 PROOF OF THEOREM G.6
This subsection aims to prove Theorem G.6. We start with few lemmas that will be used later.
Lemma G.8. r = U n..
Proof of Lemma G.8. For any point x E U nI, there exists {= I such that limg- 2h = x. Since E U and + is continuous in U, it holds that & is continuous at H, thus lim; ((2x) = +(z) E r.
However (xxx) Tkk because Jy E Vk. Thus we know +(x) E r. Hence UnT C P. The other side is clear because I C U and CT.
Lemma G.9. Let U' be any bounded open set such that its closure U U. U n' por=U0 UT
then
Proof of Lemma G.9. By Lemma G.8. it holds that U nr U' U'. U",U U' nT CT, we have that U'n
Note that U'nr
nr CurrunT=-O nT, which completes the proof.
Lemma G.10. Let U' be any bounded open set such that its closure U S U and U' n'' CU'N. Then for all h22 > 0,3 > 0 if x E U", dist(æ, I') S + dist(r, U DT) S h2-
Proof of Lemma G.10. We will prove by contradiction. Suppose there exists h2 > 0 and {ag] U", such that limg-+20 dist (Tx, r)= 0 but Ak > 0, dist(rx,, U' 2 h2- Since U" is bounded, U is compact and thus {= has at least one accumulate point x" in U S U. Since U is the attraction set of I under gradient flow, we know that ((x") E T. Now we claim x² E r. This is because limg--oc dist D) and thus there exists a sequence of points on r2 {yk J== where limg-soa || y=ll = 0. Thus we have that x" lim ykk lime-oc ((yx) "(x"), where the last step we used that xx E U and & is continuous on U. By the definition of U, x² E - 4(x*) r, thus x² r. Then we would have x E U' nr, which is contradictory dist (IR, U' 2 dist(aR, U') 2 h2 Vk > 0. This completes the proof. O
Lemma G.11. Let U be any bounded open set such that its closure U S and U nr S U NI. Then for all ha > 0,3p1 > 0 if x E U", L(x) I infaetr L(x) + P = dist(a, (T) < h2.
Proof of Lemma G.I.. We will prove by contradiction. there exists a list of Pl, Pkay.., such that Pr: + 0 and there exists xy E U', such that L(xx) S infaeu' L(x) + Pk and dist U U')l') 2 h2. Since U" is bounded, U is compact and thus { has at least one accumulate point xf in U C U. Since L is 100
continuous in U, I(x*)- = lime- [(xx) = inf L(x). Thus x" is local minimizer of L and thus has zero gradient, which further implies that x² = ((x"). Thus x² E U nT, which is contradictory to dist(rx, U' (T) 2 dist(rR, U'N MTD 2 h22 vk > 0. This completes the proof.
Lemma G.12. Let U' be any bounded open set such that its closure T C U and U nr = NT. Suppose regularizers { admits a limiting regularizer on I, then
inf zEU
(L(x)+Rp(zT)
p² inf S(x) + inf L(x) + o(p)) zEt'NT
aeU'
Proof of Lemma G.12. First choose sufficiently small p, such that p h(U" Choose approximate minimizer S(z), Zo E U" n., such that S(zo) S infreurr S(z) + p². Then by the definition of limiting regularizers (Definition 4.3) and the assumption that U" is open. there exists I E U" satisfying that | xoll2 2/2222222 I rp < p² and (x1)))/ S(x)) S p². Thus, R2( S (()))) + -==
|| - zoll2 S x² < and vo E UT.. This further leads to Topl E U U'nT.. By Taylor expansion on L at To, would have L(x1) I L(xo) + 0[1 x l13)) infreu'nr L(x) + O(p)). Thus it holds that inf (L(x) + (x) S L(xi) + (x1x)) S p² inf S(x) + inf L(x) + 0(p*).
ZEU
zeu'nr
=EU
This completes the proof.
Lemma G.13. Let U' be any bounded open set such that its closure U C U and U nr = n'... Suppose regularizers {Rp} admits a good limiting regularizer S on I, then for all E U²,
||u Φ())ll2 OOp)) - Rp(w)) 2 p² inf S(x) o(p²).
26

---

Published as a conference paper at ICLR 2023
Proof of Lemma 1.13. Define r= rr)K) 1((K) as the constant in Lemma E.5 with K = U'nT. Note K is compact and by Lemma G.9, K cr. By Lemma E.5, have K nT I is compact set, is K n P. Since S is a good limiting regularizer for {Rah. by Definition G.2, for any x" E K n., there exists open neighborhood of x², Va- such that for any C,E1 O, there is Pa such that
inf Vx E Vg- and S Pr", S(x) 1 Ro(x')/0² <1
Note that K² i I is compact, there exists finite subset of Kth nr, {oatK, such that K nr c U.V.v. Hence for any C, E1 > 0, there is some PK mink Pze > holds that,
inf || =11-= Var E Kh n I and S PK, S(x)
Rp(a')/0
<61
(18)
We can rewrite Equation 18 as for any C > O,
sup
S(x)
inf Rp(x')/ =[)), as p 0. 0.
(19)
zEK
||' 11'-
u E U S U, we have that (u) E T. If | [[flll((ullb
((p)) then dist(u, r) < O(p). By Lemma G.10,
we have that dist(u, K) -K)=oa =0=[]]] o(1.. This further implies dist (ui), K) S dist(u, K) + dist( u) o().
Hence we have that ((u) E K" nI' for sufficiently small p. Thus we can pick r (u) in Equation 19 and C sufficiently large, which yields that
{ S
inf [/ (w--0())l25000)) Rp(n') + o(p²) S Ro(u) + o(p²),
(20)
where the last step is because I|u- 1-)((u)llll |l2 = O(p). On the other hand, we have that S(D(u)) 2 inf S(z) ((1)
(21)
=et'rr
as S is continuous on I' and dist (U ((u)) = o(1). Combining Equations 20 and 21, we have Rp(u) 2
p² inf ((x) - o(p²).
Proof of Theorem G.6. We will first lower bound L(x) + R2(x)) for x E U". Suppose Cv is the constant in
Condition G.1. Define C1 = Cortiul orr ((x)+I
We discuss by cases. For sufficiently small p.
1. I &f Kh, then by Lemma G.11, L(x) is lower bounded by a positive constant.
2. I E K" and || - (2)lla 2 Cip, then by Lemma F.1,
-|| ((x)11
L(x)= N
2 (Cur+ inf S(x) + 1)p² zEU'ni
2
This implies L(z) + Rx(x) 2 (infzeuror S(x) + 1)p² + +infreu' L(z).
3. If (2)]l2 S Cip. by Lemma G.13, R,(x)= 2 p² inf S(x) o(p²), hence
L(x) + (z)x) + o(p2) 2 inf S(x)p² + inf L(z).
zet'nn
zEG
Concluding the three cases, we have inf (I(x) + Rx(x)) >
inf M(x) + inf S(x) - o(p²). sEU'rr zeu'nr
ZEU
By Lemma G.12, we have that
(Lxx) + Rx(x)) S p² inf =EU"
inf
S(x)
inf
[(x) + o(p²).
=eU'nr
Combining the above two inequalities, we prove the main statement of Theorem G.6.
Furthermore, L(u)= + (u) & inf R₂)))) + Op?)) then by the main statement and Condi- tion G.I, we have that
L(u) inf L(x)= < inf ZEU
+ +z)) (x)) Ru(u) inf L(x) O(p²)
xEU"
zEU"
<p² inf S(x) + Cp² + + 0(p3).
se
Then by Lemma G. we have u E nT' for sufficiently small p. By Lemma F.1, we have ||u- ( = O(p). By Lemma G.13, we have (u) 2 p² inf S(z) o(p²).
27

---

Published as a conference paper at ICLR 2023
G.4 PROOFS OF COROLLARY G.7
Proof of Corollary G.7. Since L(u) + (R Rp(u) inf
zEt + + S =Ap²==(p3)) by Theorem G.6, we
have that
L(u) inf
L(x) = (A + o(1)pp
zEU"
and
R₂(x) - inf S(x) E [1o A [[]].
zEt'nI'
Thus it suffices to show Rx(x) - 3(x) = o(p²). Since L(u) - inf L(x) & (A + c()))p² = o(1). by Lemma G.II, we know (2, U' nT) = o11. Thus by Lemma F.1. || (<)] - o(1), which implies that (((()))))-(f)
v (u). Since 5 is an continuous extension,
5 - (-(((<))) =
O(I|=~ - (2))ll) = o11. Thus we conclude that 5(x) S (()()) S infeet>nr S(x) + + + o1). On the other hand, 5(x) 2 ((o(x)) o(1) 2 infaeuerr S(x) o(1), where the last step we use the fact that dist(z, U (T) = o(1). This completes the proof.
G.S LIMITING REGULARIZERS FOR DIFFERENT NOTIONS OF SHARPNESS
Proof of Theorem G.3.
1. We will first verify Condition G.I. For fixed compact set B C U, as V Lxx)l22 is continuous, there exists constant v, such that Vz E B', [|V* 22 S v. Then by Taylor Expansion,
R (x)
L(x + pv) L(x) 31
max
2 max (p(VL(x), u) + (
- up²
2 -vp²
2yp²/0 /6.
2. Now we verify SMa = \1 LV2)())) is the limiting regularizer of RyMux Let a be any point in I', by continuity of TeMax
RMMax (x')
ReMax (2x
lim lim p --00 ||
inf
lim P-0
\V
Sa
p²
p²
3. Finally verify definition of good limiting regularizer, by Assumption 3.2, S = X1 (x:/2 is non-
negative and continuous on F. For any x" E Т, choose a sufficiently small open convex set containing x² such that Vz V L[]]
For any VnI, for any x' satisfying that || 2||2 - Cp. by
Theorem K.3,
RMax (x') = max L(x' + pu) - L(x') 2 max
Po 2 >a>i (V2 - ///p//6222 p² -
- up
11
|
This implies |
inf
MMaax (x') 2 (
On the other hand, for any I E vnr,
RMax (x) max L(x + pv) L(x)
max (P(VL(z),)) + ( vp²
10,
||0|2251
max ~ u (
1101
This implies inf H|r =allallsppppppp
IeMxx (M)) S (
Thus, we conclude that
inf $1o
12/ =o E V r. indicating
M is a good limiting regularizer of RMax on r.
This completes the proof.
Proof of Theorem G.4.
28

---

Published as a conference paper at ICLR 2023
1. We will first prove Condition G.1 holds. For any fixed compact set B c U, A1 and L continuous, there exists constant C, such that Vo E B², (V2 > - and | < v. Then by Taylor Expansion.
(L)
ne - L(x + VI (z) - L(x)
2 M|V (x) 1|2 + 1²1 (x)
(L(x)
[x) ))
[V) - -vp² /6
2 - 2/))p²
2. Now we verify (x) = Tr))) is the limiting regularizer of R Let x be any point in r. Let K = {x) and choose h = h(K) as in Lemma E.5. For any x' E Kh nu"
VL())
(x') M(x + VL (x)) - L(x)
2 P| (a') 112 + 8² |[[((T) (' VL()') 12 (V)) 12-
(I()')
/2 up²/6. 1 Vz (P VLL
VL(x)
(2))
By Lemma F.,, we have v1(=)) -'((0(11))
(" " + +0 [(xx))l2) Hence
[
(2)/2) cp 0110 4 (1)))ll) up² (²
This implies lim lim inf jas' =
2 Am (V
We now show the above inequality is in fact equality. If we choose x. = x + ruM, then by Taylor Expansion,
VL(a") = VL(2) + VL
- x) 00 - =x1?)
= ruM + (
VL(2")
This implies lim,-0 O. Putting together,
vI(2)] Uag. We also have lim V2L(2,) v2L(x) and lim, V'(=")-
(x())
lim Re r00
(x")) lim L(=
L(x²))
V (x))
VL (x(")
lim (pll [ (x=') 222 + 1²(
/2 + O(up²)
1
VL (x"")
VL (xx²)
>'>> 0(p').
This implies oi lim,-o inf | jn 1 (x" limp- lim,-
RAm
Hence the limiting regularizer S is exactly Aa ((2/(-)
Finally we verify definition of good limiting regularizer, by Assumption 3.2, SMMu = Am (x)/2 is non-
negative and continuous on r. For any x² E I, choose a sufficiently small open convex set V containing
x= such that Vr ..', 1 (x) 122 S V. For any x E V nI, for any x' satisfying that ||=' x l|2 S Cp.
RA00(x') = L(=' + VL(a')
[V((')) - L(x')
2 PL (x)) 112 + / VL()')
('L (z
[V (x') /2 up²/6
(+))
7L VL ( VL (x)) /2 - up² /6.
(x')
( VL (x')
vI(=')
By Lemma F.4, we have TVLIZ' -=00""11
+ 0 2(x')12)) This implies
inf RASe(x)) 2 (V
SPp
29

---

- Published as a conference paper at ICLR 2023
On the other hand, simillar the proof in the second part, we have A (V)
inf 1p'
(2)((a)) 5
Thus, we conclude that inf Re g is good limiting regularizer of Repux on r.
=(p)) Vo e vnr, indicating
This completes the proof.
Proof of Theorem G.5.
1. We will first verify Condition G. For fixed compact set B C U, as [+)]ll is continuous, there exists constant u, such that Var Bd, [[)
u. Then by Taylor Expansion,
R2 (x)
9
L(x)
1gll
2 Myo. (2(VL(1),
9
p²(
9
T
L(x)
9
|gll
2|
2 vp" /6.
2. Now verify MMx continuity of Ro
(x) Tr(V2L(())/2 is the limiting regularizer of R Let be any point in I, by
R2 AvE (x')
It (x)
lim lim inf
lim p00
p000~ ||=d
Finally we verify definition of good limiting regularizer, by Assumption 3.2, (x) = Tr 222 is non- negative and continuous on I. For any x" E H, choose sufficiently small open convex set V containing xx such that Vo E V ,1VV & v. For any E V nT, for any x' satisfying that || - x1|2 S Cp. by Theorem K.3, R RAVE(x') = Eg Eg + +P 9
- L(x')
1gl
2 EgnN(o.)) (P(VL(=)), |19| ) + 0²
T
V L(x)
up2
12gl 2 >V - up" /6 (TTVVL(())/2DD
This implies
inf || -
R AvE (x') 2 22
On the other hand, for any I E vnr,
g L(x)
+ |gg|
< ENNNOD (D(VI(.), 9
p²
9 T
(2((2)
9
|gll
[9
2
N 9 VL((a)
9
+ uup
Ielll
2||g||
This implies
R (x') S (T) ppp
inf
Thus, we conclude that
inf SCp
(2)
= O(p), Vo E Vr, indicating
goveg is a good limiting regularizer of nA-² l..
Theorem G.14. Stochastic worst-direction sharpness E, RMM admits T((V2 as good limiting regularizer on and satisfies Condition G.I.
Proof of Theorem G.14. By Theorem E.2, Condition E.1 holds.
Easily deducted from Theorem G.3 Ag(x)) good limiting regularizer for Rm on Fx. Then as I Dr (x) is good limiting regularizer for k.p on Г. Hence S(x) = E. 22M = Tx(VL((x))22 a good limiting regularizer of EE, [RMa]]] on
30

---

Published as a conference paper at ICLR 2023
Theorem G.15, Stochastic ascent-direction sharpness R& [Re admits Tr(V2T(-)/2 good limiting regularizer on I' and satisfies Condition G.I.
Proof of Theorem G.15. By Theorem E.2, Condition E.1 holds.
Easily deducted from Theorem G.4 Ax(2r) is good limiting regularizer for RRp Ts the codi- mension of T& 1. Then as r Tay As(x) a good limiting regularizer for R&p on I H.Hence S(z) = Ex Ax(3)/M Lx))22 is a good limiting regularizer of E3 ( on
Theorem G.16. Stochastic average-direction sharpness Ex admits L()(2L())((2 as a good lim- iting regularizer on and satisfies Condition G.I.
Proof of Theorem G.16. By definition, we know that E5 k.p R The rest follows from Theorem G.5.
G.6 PROOF OF THEOREMS 4.2 AND 5.3
To end this section, we prove the two theorems presented in the main text. The readers will find the proof straight forward after we established the framework of good limiting regularizers.
Proof of Theorem 4.2. Apply Corollary G.7 Riype The mapping from R to good limiting regularizers Stype are characterized by Theorems G.3 to G.5.
Proof of Theorem 5.3. Apply Corollary G.7 on Rtyre, The mapping from to good limiting regularizers Stype are characterized by Theorems G.14 G.16.
H ANALYSIS FULL-BATCH SAM ON QUADRATIC Loss (PROOF OF THEOREM 4.8)
The goal of this section is to prove Theorem 4.8. In this section, we use A < B to indicate B - semi-definite.
A is positive
Theorem 4.8. Suppose A is a positive definite symmetric matrix with unique top eigenvalue. Consider running full-batch SAM (Equation 3) on loss L(xx) = 1xd Ax in Equation 11 below.
x(1 + = x(1) - ()(((() +
(11)
Then, for almost every z0)) we have xzt) converges in direction to v1(A) up to sign flip and (P)r(A) with qu < 1. lime-soo ||x(1]]2- 2-nAi(A)
Proof of Theorem 4.8. We first rewrite the iterate as
APr(t)
z(t + n(a))) -
11P
||ax()]l2
Define i(1) V[((()) Az(e) and we have
At)(t)
a(t =(t+)
(22)
112
We suppose A E RDD and use Ai, vi denote \;(A), (AA).
Further, we define that
D p(:D) A
Σ (A)):
1, (a [( S n23},
a,(t) ())
S a (t [|](t]]]2
t> Ti}
2-n
31

---

Published as a conference paper at ICLR 2023
By Lemma H.1, T is an invariant set for update rule Equation 22.
Our proof consists of two steps.
(1) Entering Invariant Set. Lemma H.2 implies that there exists constant Ti > 0, such that A
T1, ( S n713
(2) Alignment to Top Eigenvector. Lemmas H. and H. show that [|ā()]]22 and [āi(t) converge to which implies our final results.
2-ク
H.1 ENTERING INVARIANT SET
In this subsection, will prove the following three lemmas.
1. Lemma H.1 shows I is an invariant set for update rule (Equation 22).
2. Lemma H.2 shows that under the update rule (Equation 22), all iterates not in Iy will shrink exponentially in lz norm.
3. Lemma H.3 combines Lemmas H.1 and H.2 to show that for sufficiently large t, x(t) Nzy Lemma H.1. For t 2 0, if (A) < 1 and ((t) E J, then ä(t + 1) E I..
Proof of Lemma H.I. By (Equation 22), we have that
pliD) A? p():D)
pli:D) + 1) (I pl: " [
Hence have that
[ a(1 + 1)ll2 11(7 - pp(:p) p(j.D) A²
p(i:P) 5(1)l2
[ l
P(:D) A
$ I plu:D n 17 ||
1| 11)))
Because z(t) I,, ||x(t)l22 S
This implies,
1-
X3
23
- pli:D) n p(i.D)
7(1 - n/1 - 1) |1:(0)) (21)
< [(1 - nd - |
T.
[a(t))2
Hence, ((:-)(;:)) - مو
12 S max (1, n11 + "1P
x²
1) holds that
||
((::D)((( 1) max ((:
n>3 (1
nyy'l p(: ((f) s nn}.
where the last equality is because - nu 2 0. This above inequality exactly the definition of ((t+ 1) e and thus is proof is completed.
Lemma H.2. For & 2 0, if (A) < 1 and ((t) @ Ij then
'1 X
|| 7(1 + 1) S max 1 1- 11 1a
( n (p(:)))(())b
(23)
< max (1 a>'D, n)),,1AA)) p i(t)ll2-
Proof of Lemma H.2. Note that
p(j:D) A2
[ z(t + 1)1l2 [1(1 pli:)) 7A 1 |1(()]12
)((:D))(t)l1
Pli.D) A?
< || pli:)) nd 17 || l|2
As T(1) & IT. We have ||((1)]22 2 (p))) > nA3. hence 1) p(..) A2
pl..D) A²
+ p(j:))
nx;
32

---

Published as a conference paper at ICLR 2023
This implies that
-)yp((-)>< < < plj.)) pAi-D)nn 17 p(DDD) A?
|[(())ll
and
plimA A2
23
p(j.)) pp(. 4AA [3(()]l2
(1::D)(1-1Ap) -1D1)) " ((((t)ll2
(():D)(1 -
Hence have that
( 1)ll2 I max (1- 1- -7 11100)11 n2d,) ||P()
13
< max (1 - Mayi (71ai) (():D)(()lll
This completes the proof.
Lemma H.3. Choosing Ti = maxy (- (-1o8mx((-nADmm))) max( | 1,ā(1) I
1)), then 4I 2 T1, D
nA
Proof of Lemma H.3. We will prove by contradiction. Suppose Ej E [D] and T > Tı, such that i(T) $ ng.
By Lemma H.1, it holds that Vt < T. (t) & Ily. Then by Lemma H.2,
( S max (1 nAD,nd,) ( S n17.
which leads to contradiction.
H.2 ALIGNMENT TO TOP EIGENVECTOR
In this subsection, we prove the following lemmas towards showing that i(t) converges in direction to vj (A() up to proper sign flip.
1. Corollary H.4 show that for almost every learning rate 7 and initialization Tinit T1 (t) of 0, for every t 0. This condition is important because if ø (t) = 0 at some step t, then for any t' 2 t, 21 (t') will also be and thus alignment is impossible.
2. Lemma H.5 shows that under update rule (Equation 22), t & S = ++1 E S for sufficiently large i, where nA? the definition of S is {t]l) < 2- tt > Ti}.
3. Lemma H.9, a combination of Lemmas H.6 and H.7, shows that following update rule (Equation 22),
(t) increases for t E S.
nAa 4. Lemma H.10 shows that [[t(t)] converges to 2-mAy under Equation 22.
5. Lemma H.ll shows that | (4)]2 converges
Z-nA under Equation 22.
We will first prove that Vt, (t) * 0 happens for almost every learning rate and initialization Zmi (Corol- lary H.4), using much more general result (Theorem D.3).
Corollary H.4. Except for countably many 1) E R+, for almost all initialization Fmit x(0), holds that for all natural number t, (t) #
Proof of Corollary H.4. Let F.(x)) F(x) A(x + (TAATT); Vn E N+,x E RD and 2 = (ar E RD (z, v1) = 0}. We can easily check F is Cl on R Z and Z is a zero-measure set. Applying Theorem D.3, have the following corollary.
Lemma H.5. For & 2 0, F(((1)122 >
a(() ENT, then
2-
[ + 1)1l2 & max( - nd1
a (1- (111)1=(001)))
2x²
33

---

Published as a conference paper at ICLR 2023
Proof of Lemma H.5. Note that
A2
nA ")
}z(t)
|1((1)]12
D
Σ (1- - m)3) 2,(0)
1
==1
Consider the following two cases.
1 If for any such that [1 - -000))))((002---- 2: (11 7
then we have
||a(t + 1)]l+ S (1-
mpa (1 m))))1=(()ll
2 If there exists i, such that [11 - n7)1) [x)1) < (1 +)))]1=(0011 n7 |. then suppose WLOG,
i is the smallest among such index. As
n (1 - (1)a) < mAp - (1 ,))))| = -1- (1) ma91 We have- -n + (1-1) (1-m)i)] ma - 11-17/ (1-nAr)|(((t)l) Equivalentyy
n + tnd
| 2- 1/11 n
(24)
Combining with att E 11 = |=(())]2 - 7 we have 1) < Áد
A1
Now consider the following vertors,
(1) (mA) (1 ))11) [1((1)))))(
(t) _ ((2 - n711 - 77 ( n72 - m1 p())))
(t)+) ((nditj-1- m)++)))) - +
D
Then we have
D
|a(t + 1)112 =11 | Σ (11 , ()))),(())
(mA)
(1- - (
|0(t)l2
Σ
(1- m733) (000l2
|8(0)12
1 D+1-1
||00
|a(0)l2
By assumption, we have a(t) E Mily hence we have 11 (1)l2 = (n12 - (1 - (1 (1=0))))))
(1)122 S ((22- 1/11 m):))) ?7> m)>))
((t)))))())) S ((((n)14-11- - 71++)) - - + 1
1 1 1jjs D-i.
Using AM-GM inequality, we have
Xi+j-1+ + +
1 5
3
X + Nitj
X 1,,5 S
2
Hence
(1(2))0 S (((1)1111)) t7 (1l++))(((()111 - -112+
16 - -
$ <n11)(t\|12
+ 11
2
D-i
- x
A X1
10
(2t)))))) s 11)0(1112
2
11
3
2
34

---

Published as a conference paper at ICLR 2023
Putting together,
D+1-i
| + 1)112 - ||0(1) 2 5=1
| 112
SnAf + 129(2 - 11/1- - -m21) + 172 13-
-28
(1-m>i))]
3
x + 1
1
+ +
1
- X/
||((1]12 12
2
[T(()]2
Smii + m2a(2- - 11>1- - ()t
2
- (1- - 7/ ()))) 1/ A/
- x² + x²
1
n12
X
|a(t)l12
2113(1)
2 ni) - -(1 -11- (1-)Ai))=(t)ll=> >²> 1 x/ + xA)
1
X
EmA + m27(2 - 1711 3
|ā(t)]12
17
22X
We further discuss three cases
1. V 1 I-
+2f+n22
we have |
2- n/z > .//iv
1- ,then
3ー
||a(t + 1)llz
SmA1= + (2 - 1/111
2
- (11)))]]((0)1a ,7 12 + x1 )
" 1
3
|a(1]l
2X
3 m1) - (1 - 17211 -1-1)\1) 1// + m/
2
SMAR + 127(2- 17
17d11 my>
1/ 7/89 1
2/ + )3
2 1/11 - -mA; X
2
n7 + 1)
11
227
7 17
X1
$2 2
2x7
The second line is because (1 - -n-)n())))(0lll + >222 + x7) |a00)la
1
monotonously increase w.r.t
[(((t)]2 when |a(())l2 > ndy v 1
The last line is due to Lemma K.9.
1-
u2p+n2?
2. If 1/ 2 ny
2 2- - then
T-mAi
|| + 1)ll2
2 111 -0- (1 -1l--m))ltt)ll>>>>--83 { |a(())l2
Xo
SmAi + 1/ 1711
222
SMAR + 117(12 - 1 -
31 - (AP + 18++21)(11 (1)))))))) - -2
2
2/
3
7 - X1
2 2-mp> 1/11
2X7
The second line is because of AM-GM inequality. The last line is due to Lemma 1.
3. If nA < ndi 4NP+2
we have || nat < ndy / +x7+
then T-mA1
T-nA1
|| + 1)ll2
2
1
21
SmA + 12 - -121-
m1 - (1 (1- - x)
3
|a(t)l12
11
21²
Sm17 + 1//7(22 - n7d1
) - -11-ddimdi- - -
2
x1
3
2X
?
X4
2 2
2x²
The second line is because (1 - m1)0l1000 + 7
+ x
monotonously decrease W.r.t
[1(()) when |a(()]l2 < ndv 1
The last line is due to Lemma K.10.
35

---

Published as a conference paper at ICLR 2023
Lemma H.6. if)l S 2- it holds that [ā1(t + 1)1 2 | (t)]-
Proof of Lemma H.6. Nota that [81 (t + 1> | ny1 "7 ( 2
x?
881 (8) and that 7 Dpsttill2
A3
> 2 1729.
It
follows that - 7711 " [((ell
1. Hence we have that رāق (1 + 1)1 1 (1))
nAi
Lemma H.7. For any 2 0, if [(tt))l E 2- a(t) E Alg it holds that
| + 1)1l2 m (1 ())))(
Proof of Lemma H.7. Note that
A2
X
X
|| 7 7
[[((t)]l2
12 S max 1|1 9/, 17
ISiSD
= 27
n25).
The proof is completed by noting that (|ā(t + 1)ll - / - 1/ - 77
A²
lla (t)(t)ll
Lemma H.8. For any 2 0. if((((1)22 S クdス T holds that
|1 + 1)112
5(m - (1 + (1)
(1)
)011 )
(1 n))))l - n231
(OP
max SE(2.]]
|
fau (1 (10)2
Proof of Lemma H.8. We will discuss the movement along 21 and orthogonal to v1. First,
| + 1)1l2 II(I - p(2:D) nA - p(2.D) A2
|
| - p(2:D) n 1 p(2:D) A?
l1 ((2:
[x(t)]12
< max {|1 - nt jE[:.M]
(())2))(())l
|
Second, [t + 1)1 = ( nx?
1 + m7 | (t)). Hence we have that
|
|| + 1)
< <n1 (1 + ())))
1 (t)l2
[[L [LL
- n131 1 ))²(2
| (+)3
|
max (E[2:M]
1ya - (1 - ()=))))))
|a(t)
Lemma H.9. For t, E 8,0 f, then [ar < |. (t'))
Proof of Lemma H.9. Fort S, by Lemma H.5, t + 1 S or + & S,t+ S. We will discuss by case.
1. If + 1 E S, we can use Lemma H.6 to show [t) < [t + 1)).
2. If 1 + g S, S + 2 S, then
(nu - 1
1)1l2)
T (t + 2)
a (t)
(1
36

---

Published as a conference paper at ICLR 2023
As
(MA - (1 -1-1A1i -1--AA))-00//b//0--0000-000000--
2 |+(0)22||*( + 1)1l
T ---mm - 1)1l2)
2 (2n/1 - [ā(t + 1)1l2
= n² x --p(11111)0001000 (7))==))))]((())l
2 ((2m)) | + m)a((1 n)11)) || + 1)lla,
(1)|2
combining with Lemma H.7, we only need to prove,
2x - 1 7)>1)
2 ((2m)) (|= + nAf(1-1A1)) 7]\1)) (1727 (1 (1-n)i)(((((ll>)))
Through some calculation, this is equivalent to
(22 ( -A mAA) 2 20.
which holds for [|x())l22 S .\ Combining the two cases and using induction, we can get the desired result.
2-
Lemma H.10. ||ā(t)] converges to 2-
nd when + co.
Proof of Lemma H. By Lemma H.9, (t) increases monotonously for E S. By Lemma H.5, S is infinite. By Lemma H.2, for sufficiently large [+)(t) is bounded. Combining the three facts, we know (1) for E S converges.
Formally Ve > 0, there exists T, > 0 such that Vt, E S,t t> Te,
(t) [|2(())]2
< <1+c.
Then by Lemma H.5, Vt E S. 1 S ort 2 S, we will discuss by case. Fort 2 T..
1. Itt+le 1 S, then
[t + 1)11 [t)l(
DAA -11 (1- n))))
1 1
[[
2. If + ++16 1g S and + E S, then
|| (t + 2)1l2 || (t)ll
1
(ma) (1 - m)t) - (1- mA) + 1)112)
|ia(t) l2 (| +1)ll2
(n)
(1 m )| [()(() (n
(1-1/A1) (n2 - -1--111100))
|E(t))l2 (TA) - -(1 (-1-1\i m711) (t))l2)
?> (1 - ))) (n>?
(1 - (111)(()112)
Here in the last inequality, we apply Lemma H.7
Concluding, [|x(t)]l2 2 min 3- +e' (2- ( 1t Te, S. As VI & S,t > Te, we have
||ā(t)]22 2 nA? 2- Hence we have Vt Te, [1(())le min (2-1Af+e' (2-A1n)A1nte))
Further by Lemma H.7, > T. + |a(()]12 S n7 (1 (7t1a min
nd? ( (2-Xin)
1
Combining both bound, we have lim
||x(t]]12 2-
nd?
2x² Lemma H.11. [t)(() converges to 2-
when
37

---

Published as a conference paper at ICLR 2023
Proof of Lemma H.II. Notice that
| I(t 1) max (1-nd2-
(
17
1.11 17DD 77
(())2))(())22
When |i(t)]l2 2-1122-5-
X
X3
-1+0 1/22
1 < - <1-1 - -11 (1) S -1 -1-10 1/D
[| ((t + 1)ll2 I max(1 - .nD. - 100-1-)]1(P2-DD) i()I12
Hence for sufficiently large t. | shrinks exponentially, showing that lim ||
---
I ANALYSIS FOR FULL-BATCH SAM ON GENERAL LOss (PROOF OF THEOREM 4.5)
The goal of this section is to prove the following theorem.
Theorem 4.5 (Main). Let {x(t)} } be the iterates of full-batch SAM (Equation 3) with z(0) Tmit U. Under Assumptions 3.2 and 4.4, for all n) P such that r ln(1/p) and p/ are sufficiently small, the dynamics of SAM can be characterized in the following two phases:
SAM (Equation follows Gradient Flow with respect to L until entering an ((p) neighborhood of the manifold T in (1n(11/)) steps;
Phase I: (Theorem 1.1)
Phase H: (Theorem 1.3) Under mild non-degeneracy assumption (Assumption 1.2) on the initial point of phase II, full-batch SAM (Equation 3) tracks the solution X of Equation 7, the Riemannian Gradient Flow with respect to the loss A L()) in an ((p)) neighborhood of manifold P. Quantitatively, the approximation error between the iterates x and the correspording limiting flow is Og In(1/p)), that is,
[Tx([Ta// (ToTa/((o)))) - (Ta)ll2 = OO()n((1p))
Moreover, the angle between V (I((( 23 1) and the top eigenspace of L2L T3 1)) is O(p).
Readers may refer to Appendix E for notation.
To prove the theorem, we will separate the dynamic of SAM on general loss to two phases.
Define
(x(
0(x)2 mp))(s))) [M U,
which is the length projection of (xx) on button- A: non-zero eigenspace of (²L(((x)). We will provide a fine- mne-grained convergence bound on (r).
Theorem L.1 (Phase 1). Let {x((()} be the iterates defined by SAM ( Equation 3) and xx() = Tinis U, then under Assumption 3.2 there exists a positive number Ti independent of q and such that for any T > Tí,
it holds for all n, p such that (n) + p) (n(1/n0) is sufficiently small, we have
max max (2(x(()))) O(np²)
max Ti STTT (n(1/n)) JE(M]
[|*(x(()) (x S 0(17 p)
max Ti (n(1/np) In(1/pp)
Theorem I.1 implies SAM will converge an O(1i) neighbor of r. Notice in the time frame defined by Theorem 1.1, xt) effectively operates at local regime around & In(1/np) (n(1/np)/nl) this allows us to approximate L with the quadratic Taylor expansion of L ([[T, In(1/pp)/n) and prove the following theorem Theorem I.3.
Towards proving Theorem I.3, we need make one assumption about the trajectory of SAM, Assumption I.2.
Assumption 1.2. There exists step I, satisfying that T, In(1 /npp/// < 1. < 0(1(11 /np/n)), [(((() )(x(1)), vn (a(1))) 2 21pp)) and that [|æ((1) (a((1))l S A1(()pp 22(p2, where Ta is the constant defined in Theorem I.I.
38

---

Published as a conference paper at ICLR 2023
We remark that the above assumption is very mild only need the above conditions in Assumption 1.2 to hold for some step in (1/n)) steps after Phase I ends, and since then analysis for Phase II shows that these two conditions will hold until Phase II ends.
Theorem 1.3 (Phase II), Let {a(()) be the iterates defined by SAM (Equation 3) under Assump- tions 3.2 and 4.4, for all ?)P such that 7 pl)(1/p) and p.n is sufficiently small, further assuming that O((1 + p) In(1/np)), (3) (=(0) (z)())), in (z(0))) 2 O(p²)) and (4) ||x(0) -=(2(t))]l S (1 Q(p²), the iterates x(t) tracks the solution X of Equation 7. Quantitatively for
(1) ma Mi (x()))
O(no²), (2) [[{(x(0)) (T
1s/np"1. we have that [1*()(1)) ()l O(n lmlppp)
Moreover, the angle berween VL(x(t)) and the top eigenspace of (LL( is most O(p). Quantita- lively.
(x((t) ((x(()), = Q(1p)-
max JC(2.M]
(xx(t) ((<(t)).)(z(())) O(pp").
In this section we will define K as (t) |0 s t T3) where is the solution of Equation 7. To simplify our proof, we assume WLOG L(x) 0 for x E..
I.1 PHASE (PROOF OF THEOREM I.1)
Proof of Theorem I.I. The proof consists of three major parts.
1. Tracking Gradient Flow. Lemma 1.4 shows the existence of step tGF O(1/n) such that 2(tGF) is in subset of Kh and & (x((r)) is (17 + p) close to o (Immn)
2. Decreasing Loss. Lemma 1.6 shows the existence of step EDEC O(In(1p)/mm) such that æ(tDEc) is in ((p) neighbor of l and & (((((E)) is O( + p) In(1/p)) close to & (Xini))
3. Entering Invariant Set. Lemmas 1.11 and 1.13 shows the existence of step (INV O(1((1 (00(l//p)//) /n) such that for any t satisfying tiNv t S tINv + (Im((/n))//)) we have that rrt) E Kee( ll and (x((t) is
O( + p) (n(1/pp) /np)) close to +b (Fma).
1.1.1 TRACKING GRADIENT FLOW
Lemma 1.4 shows that the iterates sz(() tracks gradient flow to an O(1) neighbor of Lemma L.4 Under condition of Theorem I.I, there exists tcF O1/n), such that the iterate x(tGF) is O(1) close to the manifold I and (<x((eF) is O(n + p) is close to (( Quantitatively,
nh2
L(x(tar)) =
32
| - ((x(tc))) - h/4, [ (x = O(1 + p).
Proof of Lemma 1.4. Choose C = 1 /
8- Since (dimi))) limy-+00 (b( T), there exists T 0, such that [i T) (Fintt)l Ch/2. Note that
(x(()) x(t) - (LL))))((((())))) O(())
att + 1) = -\(()==n\((((() 7)V)(((t) +
(x(1)) 1
VVL
By Corollary I.3, let b(x) = -V(x), p 17 and O(p). we have that the iterates z(t) tracks gradient flow @Drin, T) in (1/n)) steps. Quantitatively for tGr = = we have that y
||œ(|GF) d(Tinis [2 O(e + p) 0(1+p)
This implies x((oF) E K*, hence by Taylor Expansion on Φ, [ + (x(mi)) 1 (( F (cb T))
< O1110((a0 - (x (12)
< On +)- p).
39

---

Published as a conference paper at ICLR 2023
This implies
[[((tr)) - ((((ttr)) ))ll
$ ( - (f To)ll2 + [[ To) (Xrini)lla + ((ima)) (xx(G))))1
+Ch/2+ p) - Ch 1 h/4.
By Taylor Expansion, we conclude that [(x(ter)) S (l(a(teF) ((x(tc=))18/2
AA²
32
1.1.2 DECREASING LOSS
Lemma 166 shows that the iterates s=(() converges to O(p) neighbor of T in (1n((1pp/m)) steps. Lemma 1.5, Under condition of Theorem 1.1, if =(() E K" and [VL((((() 2 4çp. then we have that
Mx(lI+ 1)) decreases with respect to La(()), quaniitatively we have that
(t + 1)) & L(x(t)((1- - np/8).
Moreover the movement of the projection of the iterates on the manifold is bounded, quantitatively, have that
{ 1)) +<((0))1 S 0(1))
Proof of Lemma 1.5. As sx(1) E Kh and is p-PL in K*h we have [(x(()) 2 0. As xt)) Kh, by Lemma F.7 and Taylor Expansion, we have [z(()x(( + 1]]l O1)). hence for sufficiently small (t))((((( + c KT. Using similar argument, the segment from rz() to xx(t) + vL(=(t))
is in K".
Then by Taylor Expansion on L.
L(x(t + 1)) = L(x(t) TVVL x(t)) + VVL (z(t)) (L())))
(L((())
< M(x()) ( 74 (z(()) VL (<(1) + (VL (10+)))))))))
(=(())
VL (z(()) + P V(((()) |
(25)
By Taylor Expansion on VI, we have that
VL (z()t) (L
IV xxt) + PILL (x(1)
(I()
After plugging in Equation 25, we have that L(x(1 + 1)) [(x(()) v If < (z0) G( +22000
|13
(26)
As [V(x(t)) 2 4Cp, we have that the following term is bounded.
(x()) < / (x(
( 7"' VL
cnpps <2 <
VL (x
16
After plugging in Equation 26, by Lemma F.2,
L(x(1 + 1)) < L(x(1)) 16 ([VL 1??
$ I(x(t))(1 (1/188)))))))))
As xt)) Kh,, by Taylor Expansion, we have
| (x(()) || $ Ch.
Hence by Lemma F7 and Taylor Expansion, [| + 1)) - (>(())) < Enpl (x) ||2 + udp² (z) ( = s 067))
which completes the proof.
40

---

Published as a conference paper at ICLR 2023
Lemma L.6. Under condition of Theorem I.1, assuming there exists lev such that [(((ter) wk²
32 and
a(tGF) E Kh/4, then there exists ÍDEC = +GF +OQIu(1p))mm) such that x(DDEC) is in O(p) neighbor of quantitatively, we have that
[V([(teec))] I 4Cp.
Moreover the movement of the projection of (xx())) the manifold is bounded,
((n))(1/))). |(((x((rr) (=(((n))))
Proof of Lemma 1.6. Choose ÉDEC as the minimal 2 tGF such that [VL((((De)))]l S 4Gp. Define C [in nu (64p² /ha)] O(m((1/)//))
We will first perform an induction on t & min(tDEC, lor + Cj = ler + 0(m((1/p/n)) to show that [(x(1)) S (1 11/8)/ L(x(IcF))
(+((t)) ((x(tar)) = O(7) tcr))
For tGp, the result holds trivially. Suppose the induction hypothesis holds for t. Then by F.1 and Taylor Expansion,
2[(x(trF))
(x(x(t)) xx+)|l
I h/44.
Then we have that
dist (, x()) x(( [z(((rr) -[s(l\/))))(((()))ll + +1(((<((aF)) - ((x(()() + 1|((x(()) x()]
<3h/4+ tcF) 3 O(n In(1/))) S h.
That is v11) Kh. Then as IDc. [|V[(=())]] 4Cp. Then by Lemma I.5, have that
L(x(t+ $ (1-nn/8)[(z(t) - (1= L(x(tar))
[+ - ((x(ter)) < [[(x((++1)) - + (+(((t)) (x((Gr))
3 =0(7
which completes the induction.
Now if ÉDEC 2 tGF + C = tGp + ((1((11)/n), As the result of the induction, we have that
7/ L(x((oF)) < L(a((e=) < 8p²p.
64p²
L(r(|aR + C)) s (1
8
122
By Lemma F.2, we have that [|V[[[(ter+ <( (L(x(tar+()) 4Cp, which leads contradiction.
/
Hence we have that IDEC S lar + + lGr + 00n By induction, we have that [=((<te))) ((I(tar)) = ( (IDEC tar)) = (nlpp())
This completes the proof.
I.1.3 ENTERING INVARIANT SET
We first introduce some notations that is required for the proof in this and following subsection.
Define
$ = 4(x),
A(x) VTI
/ = A(x)²,
ã (ā, Uz (y(2))
(b: Σ
vs (2(x) (x)
41

---

Published as a conference paper at ICLR 2023
Note 3 Ra VT(x)) for near the manifold T. We also use #((), A(1) and =(1) to denote x(1), A(x(()) and x(t).
Recall the original definition of Rx(x) is
M
(x(x)=
(o))) - - -
Based on the above notions, we can rephrase the notion as
(x) - n2))(x). Rz(z) = ||
We additionally define the approximate invariant set I1 as
J (1I (z)x²l I (x) + O(np²)}
Lemma 1.7, Assuming t satisfy that rrt) E K then we have that
|x((1) (x(())) | -1)z(())((ard +x)((())
Proof of Lemma 1.7. First by Lemma R4, ((x(t)) E K", hence
||ā((||| [ (((())) <(|((()) (x((())
Also
| = [V (x(())) 2 11| (()0))))))))( - -(((()))
By Lemma F.4 and Lemma E.6, we have
[|<((() (x((()) < || - (+(()) + (( (cx(0))
Ç" (x((t) la(t)
-+x(()))1 + |=(0|)
412
$ 2 [l((() +(x(1))] - 11(00))
Hence ||=(()) (xx(1)) < |+((|))
Lemma L88. Assuming & satisfy that a=()) Kh and ||x(()]12 Op)) then have that [0 + 1)) (z(())) O(p)²).
Proof of Lemma L88. By Lemma 1.7, we have ||=(t) - ((x(t)) = O(p). By Lemma F.7, we have that + <2 +x)))12 + unp² + E22n2
[ + 1)) (=((t)) < < O(1p²).
Lemma 1.9. Assuming t satisfy xzt) E k2//2 and [x((t) +(x(1))1l2 O(p), define x' as x' (t) = xx() and for T 2t, t
z'(r = ((==)==0()1---244000000001000000
-
-n0 -
Then
|| + 1) - x(1 1)l12 O(pp²)
and further if ||a(t 1) ((x(t+ + 1))ll2 -((p)) then
|| + 2) ztt + 2] O(pp²).
42

---

Published as a conference paper at ICLR 2023
Proof of Lemma 1.9. By |=(()) - -0(())))()))0000000000 O()), sx1) Kh/2, and Lemma F.,, we have that || +
xtt)|| = ((p) and hence x(t+ E K3h/44 This also implies [|=(t+ [ O()). Similarly have xxtt + 2) E k3//4
For & E (1, 2). by Taylor Expansion,
( ( ( VL
xx( + k 1) =x(1 + - (NN VVL (x(t + +)) + (
+ - k))) + O(np²)
-m0VI(((0((+44)) (L(x(++))))
[4 1 0
[V
(1(--))))M7((000((+))))(+++))))) (x((t k)))
(+0(0((((((++4)) (x(t+++
k))] +0(1²²)
Now by Lemmas 1.7 and 188 = -(x((()))))))((())
xlt + 1: 1) =x +
< (x(t + + Onpp²)
(27)
(z k]]
Now first prove the first claim, we have for k: 0, - = 0. by Lemma F.4 and eq. 27,
(=T(6()))([(() +()))
axt + 1) (t - (V2(0)))))0(((())()000))00)-)))0 ((())
Tp [==(((())((( ((( -(()))l2
O(np²)
+1++1)+0mpp)
The second claim is slightly more complex. By the first claim and Lemma F.4, we have that (x((tt(((++)))))))) 1)))
(z(t+ 1))
(1) 4 (x 1)))]la
VVL (x(++ (+=L(((((+++)))(<( 1))ll2)
(28)
We first show
(= is of order - 1))112 = 2(0?)
to show that the normalized gradient term is stable with respect to small perturbation,
[ + 1)) 0(x(t 1)))
> 2| (PPaaccHnnATTL((((00+++++++11111
((x(t + 1)))ll
(( 1)))[l
<( | (( + +1) ((x(( + 1)))112
2x(1(x((t+]) + -
2n(1((((++ + 0(x(( + +))))l2
'S
lo( |a((
(
4/p
> # (=((t(+1) ((x(t + 1))) 9(1p).
Based on Lemma F.7, we have
(x((t+ + - +
We further have by the first claim and Lemma 1.8,
-V00))/)eeDD----0---11----000000 + +0lla(=))--(((--)))44444-----))1
Vi 0(pp²)
-(((())))))+)+)))))))))))))))))) O(np²)
=0(1p²)
43

---

Published as a conference paper at ICLR 2023
This implies
-1700)((+)))((+)))a44444 +(x((+++))))))+)++++) 1)))l [V +(x)) + 0
-100)) ((x))
+ (+=[(
Combining with Equation 28, we have
VI (ar(t + 1))
(I(()))((((+1))- ((x))
IVL (z(t 1)) - +0())
By the above approximation and Equation 27,
axt + 2) x'(( + O(n)²).
Lemma L.10, Assuming satisfy that xzt) K3h//4 and ||x(t]]22 =O(p), then we have that
1)((+))-2(() (1)))(()) (1) a(t)
| ||2 O
Proof of Lemma 1.10. By Lemma 1.9, we know
*(t)
|| + 1) =(t) + nz( + noA(t)
|| - O(np²).
[((+)]|
This implies
Z(t)
- +)
[z(t)
*(t)
(l)())((t++1))-- - =() + (t( + mpA(t)
|(((t|] )I
x(t)
<cIl0(t + 1) - x(t) + m²()() + npA(t)
|| 0
(29)
|æ())]
We also have
#(t + 1) -)()(()) +1 ((x(t)))
=(A(L+1) A(0)(=(L+))) - (((((+((+)))))(( 1))) A()(((2(0)) $(((1 1)))
=0(p)²-
Plugging in Equation 29, have that
i(f) [a(t)]l ||2 O(1p²)
| + 1) =(1) + +
+ npA((1)
Lemma L.1. Under condition of Theorem I.I, assuming there exists tDEC such that a(tDEC) E Kh/2 and || [x((DeDE)) s 4Çp, then there exists tDEC2 = tDEC + O(In(11))//)))))) such that x(tDEC2) is in I n K3h/44
Furthermore, for any satisfying tpec2 <t S tDEc2 + ((m(L/n)/n)) we have that x(t) E II K3/4 and [+(x(t)) ((*((DE))) = (n(11(111)))))) (n()/)))
Proof of Lemma I.H.. For simplicity, denote C = [Ini-op
ny² 42a
((tn///)))00(((/////m/) Here the
op
quantity 0(In(1/ /p)/()))//))))))))))))))))))(()) is the same quantity in the statement of the lemma.
We will prove the induction hypothesis for fDEc S t S IDEC + 2C,
| 1) 2 np))(t), tDEC + ||ā(t]]| s (1 - (z- 1)1l,
| 1)] ~ np>i( 1), IDEC = |a((t]] S np>?(1) + O(pp²), [|)()(t)) (= S O(rp((t tDEc)),
xxt) E 33/4
44

---

Published as a conference paper at ICLR 2023
The induction hypothesis holds trivially for /. IDEC-
Assume the induction hypothesis holds for Σ t. By Lemmas R.1 and 1.7, [[ & (I]=((DEC)
((s(IDE))) < و [[VTT(((DD))) 45² p. Combining with the induction hypothesis, we have | S 10
By at)) K3M/4 and Lemma 188, we have that
| + 1)) (x(())) & O(np²).
Hence we have that
[0(x(1+ 1)) (<<((DE))) S [| + - +xx ([(x(() (x((DE)))
(t (t+ + tDEC)).
(30)
This proves the third statement of the induction hypothesis. By [[(()]] = (p) and Lemma I.10, we have that
a(t)
| + 1) =(t) nA(t))(t) npA
lla O(np²).
|](()]]
Analogous the proof of Lemmas H.1 and H.2, we have
1. If [a()]| > (t))(t), we would have
a(t)
[=((1) A(())(()))))) n0)
[[=(t]]
<| m7((1) - (1) |=((\
1
<|| max 1 - DA0 daau
|+((t)]|
<max (ay1i (=(1]| noun 771 (*(())
< <max [1)]|(((t||| ?odn ,((1(()
Hence we have
1 + 1)ll & max((1 - (111-1)))((( - npy², , + 0(1p²) s m/4 m/) [t)])
2. If |5((]]22 S np))((), then by Lemma H. have that
[|ā(t) - -n((ta(()) - (1) #(t)
2l2 S nt))(t).
|=(|)|
Hence by Lemma K.1
|| + 1)ll S (t) + Ompp) S np +1)+++ O(pp²).
Concluding the two cases, we have shown the first and second claim of the induction hypothesis holds, Hence
2 we can show that ||ā(t + 1) < A p. Then by Lemma 1.7, we have that ||x(t + 1) - 2xx(t + 1))
S sc² +
As t S DDEC + 2C fDEC (lI((1/))/)) by Equation 30,
[0(((t + 1)) (>((rec)) s 00-p21lnn))
This implies
dist(x(t + K) <dist(x(tDEc), K) + ([((tec) - 4<x(tDec))
+ +I(((((DE))) 4(=(( 1)) + ||a(t + +1) 2(x(( 1))]l
=h/2 - In (n/n))))))(n/)))))())))) _ 3h/4.
This proves the fourth claim of the inductive hypothesis. The induction is complete.
45

---

Published as a conference paper at ICLR 2023
Now define iDEC2 the minimal / 2 IDC, such that ||ā((]] np){(().
If (DEC2 > IDEC + C, then by the induction, Lemmas F.1 and 1.7,
[ + C)I $ (1 1
[[
s m [
4ç2
nn
<l|=(tDEC) (<x(tDE))
4ç2
nu²
[VL((DEc)]
40
1dnp
X (tDEC + C)np..
This contradiction. Hence we have tDEC2 I LDEC + C. By the induction hypothesis r(fDEC2) e I n K3//4
Furthermore by induction, for any t satisfying (DE22 St S IDC + 2C', we have that || S (t) O(np²).
By the induction hypothesis =(1) E I K3 and [zx()))) (<<((DD))) O(p In(1/)))
Lemma L.1. Under condition of Theorem 1.1, assuming satisfy that xtt) 1 33//4, then we have that R& (x())) 2 0 > Rx(x(t+ 1)) X² 1)mp S (Ra (z(t)) + (t)mp))
(x(()) 5O - (x(1 + 1)) S (
Proof of Lemma 1.12. As (1) E Ir, [001)22 < <np + O(1p²).
As [a(t)]22 (p), we have z(t)z(t+ + 1) C K* and ((x(t) 1)) C K".
We will begin with quantization technique separating [M] into disjoint continuous subset Si, Sp such that Vi
min KES,
(tx(t) Ax(t))
By Lemmas 1.8 and K.I, we have that for any n E [M].
(1)(1) (a + 1)1 = (((())((( (2I,(0(x(1+ I)))II
-((1)(x(()) - ((((1 1))1l)
O(npp).
This implies
min kES,JES,
[x+(++ A1(t + 1)1 2 O(1p²) 2 0.99p.
Define
p(t) StD
Σ
Un (1)
keS,
By Theorem K.3, for any k;
(²L()(z(t)) - VLL(0(I(t+ +)DD11 O(p)).
| St || 0
(++1)
By Lemma I.10, have that
z(t)
( 1) Ø(t) + nA(t)a(t) + 7t)
||2 = O(np²).
|
We will write z'(t + 1) as shorthand of z(t) nA(()a(t) - (t) a(t)
|
Now discuss by cases,
46

---

Published as a conference paper at ICLR 2023
1. If Vin 11 x(1) > maxxes, 2g(() > npmp. by Lemma H.3,
1000 1)11 S
|1pS00 50) (1+ + O(pp²)
(t)] - mpdp(t + 1)² Σi= PSO *(t) |
S (1ax((1= - mob(t + +1)]
Σ
g(0
max As + 1)1l kES;
Σ
p())
[t)]|} + (mp²)
S(o
p
a(t)l np 4r nSll Σ
I
<max((1-
Σ
p Ste
sto
a(t)]l)) + O(pp")
i=j
This further implies
2 x(t 1)11
Σ 1
stn a(t +
1)II
p
mp 13
I <max (1 muuill || Σ
p(t)
st)
z())l
2 7751
Σ
pl0
so a(1)I)
+
p 5(1 - m//i1 Σ
p(t) 5(0
(t)]l
2. If V Ep | 5(t)| S maxies, {z(() then by Lemma H.I, we have that
s0) of' (l 1)ll2
up max x²(() KES;
Hence we have that
p
a(t + 1) <
P(U stn
r'(t + 1)112 + O(1p²)
< max KES,
(Dmp + 01p²²)
< max (t + 1)mp + O(np²). KES;
This further implies
p pp() 8l 1)112 O(pl + 1)1)
Σ 1
5(1
1)112
< max KES
{t(() O(np²)
< max KES
( + 1)pp + O(1p²).
Finally taking into quantization error, as all the eigenvalue in the same group at most differ Dp, for any E Sj. we have that ( + 1) + maxkes, A2 - 1) < 2DCp + D²p²
Hence the previous discussion concludes as 1. If Rs (x())) 20 0
Ra(z(t= + 1)) + Xi(( + 1)mp $ (1 (Hx( + (tt)
47

---

Published as a conference paper at ICLR 2023
2. If (x(())) <
(x(x((+1)) S O(1p²))
Lemma I.1.. Under condition of Theorem 1.I, assuming there exists DDEC2 such that for any t satisfying (DEC2 I S tDEC2 + (1m(1/n)/m) we have that x(t) E II n K3h/4 Then there exists tnvv = (DEC2 + O(I(1/n)/n)) such that for any satisfying IINV S 1 INV (ln(1/))/m) have that
xx1) (Tketar] KTh/8_
(+(()) - (<(( || = 0 (n(])))(1)))))))))))))))))))))))
Proof of Lemma I.13. The proof is almost identical with Lemma 1.1] replacing the first two iterative hypoth- esis to Lemma 1.12 and is omitted here.
1.2 PHASE II (PROOF OF THEOREM 1.3)
Proof of Theorem 1.3. Let falign O(In(1/p)/n) be the quantity defined in Lemma 1.19.
We will inductively prove the following induction hypothesis P(t) holds for tALIGN S s Ts/np² +
x(t) E Kh/2, tALIGN STSI Stt
((((r)
((p)) (n()) FALIGN StSt
Kx(r) - (((()))))) = Onpp², fALIGN STSt
max jE[2.]]
[((x(T)) - X ((m ( (n(L/())) LALIGN STSt
P(tALIGN) holds due Lemma I.19. Now suppose P(t) holds, then æ(t + 1) E k* By Lemma I.19 again, (x(t+11--0 (xxe)((((((1)))))) 1)), on (x(t+ 1)) = ((p)) and max (x(++1) 2 (z( 1)), (x( O(no²) holds.
Now by Lemma 1.20, [0(x(r + 1)) - ((x(T)) + np² )((())),r (A)(0)/2 O +1²²²² -0
tAlIGN STST.
By Corollary L.3, let b(x) -A
np² and c c= 0 p). it holds that
[[ X(np²r))l
-((()((((AA)GG)) - 2 (2m)a) || TTmpp (0+ 17)
001 fALIGN St
This implies [|x(((1))) ( (n 1)) $ ((x++1))((((++1))) [|([++))_((((+)])]]]+++++++++++1 ([(( (1 = Ö(n ln(1/))) < h/2. Hence xxt + 1) E K4/2, Combining with P(t) holds, we have that P(t + 1) holds. The induction complete.
Now P(TTs/np²)) is equivalent to our theorem.
I.2.1
ALIGNMENT TO TOP EIGENVECTOR
We will continue to use the notations introduced in Appendix I.1.3.
We further define
S [t]|)(()||
P + O(np²)}
2 m/11
1
11?
T7il
T = [d]]=(()|
2
2 1111
+
2 2- -m1aaa
pt.
17
U (1|((()) < (t)l -
p}
2
Here the constant in depends on the constant I, and will be made clear in Lemma I.16.
For S. define next(s) as the smallest integer greater than in S.
48

---

Published as a conference paper at ICLR 2023
Lemma I.14, Under the condition of Theorem I.3, there exist constants C1 C2 < 1 independent of 1 and p. if (1)22())l2 S 2 ( 2d:108
(2-ndat0) + nda(1)
2-1 p and z(t) E (Nze(Mn(()) n K7h/8, then
n
[0(1))22 2 C1 - n111" = || + 1) I C2 2 - -m>1" P 2
Proof of Lemma 1.14. By Lemma 1.10, if we write x' (t+1) as shorthand of x(t) (())) ||((|]]* then [ + 1) - (t + 1)ll = Onn)))
#(0
Define I as (x) R2 (z) ~ 0}. Then we can find surrogate Tsur (t) such that Tsur (1) E (NHE[M)) ITqundd nnk
and | (1) =[)]|22 (np²)) We will write zuur (1 + 1) as shorthand of F (1) 1)A(() (00 (1)
npA2((t) Fa (f]ll
(t)
Let
h(t) - (2- 1)
1
1
(5 - 1)2
+1) + (1
1
(( ((S-A) A)2 +1) max G²
<2 -2-1² ((-
2t
5²
2
?² ç²
As (1) < 1, we can choose C1 < 1, such that h(1)) 1.
We can further choose C2 = max + 1)/2,
1- - 322)
< 1.
We will discuss by cases
I. If
| 2 17(1 - -))1) + (X} - P
Then
|5((0]|2 12
A2(2 1112-1)11
923 2 2-1A1 P
X7(1 - -dd1) + (Ap 21) - 11>11)
27(2 17/11)
XF(2 17,1 - -dpn)= - 2²2 (1 1721)
2 1 Xo 1 11\1
122 3C2
2
2 1
N 1-
X7 2 - 1121
2-nAi
In such case we have
5(()
Taur (L)
O(p).
I (t)
Then we have |1 (1 + 1) -1)-2((0 x1 + 1>]1 = Om1p²) ). By Lemma H.5, we have that
|| - l) < ||ā(t + 1) - I'(t + 1)]12 + || + 1) x'mm (t + 1)ll + || (t + 1)ll
2 17/
11
< max(
-m>d p 1p
2x7
np> - (1 - n> ( + (np²²)
< max(1 12 - m711)
/²2 /
1727
,(2 n111)) - (1- -):
2X1
3 2 n111" P
1/2 T
m71
< (1
3 2 - 11111
p 5 C2
1/111 p.
2. If
1711
7777
|i(())la < X²(1 - )dnn + (2} 23)(1 -111) PS S
1 m P.
Then we have
1- no + +11 - 7> [12())
< A7- X2
npbr - (1 -)))))
X
Ina13 - -(1 -11- - m7 (12))))(()
X3
<
naa -(1 (=)(112
X²
49

---

Published as a conference paper at ICLR 2023
By Lemma H.8, || (t + 1)ll2
(1)22)2
A7 Xo 12
<(np\} -(1 -1-m)i)l2(())aa)
+11
max
|
|
X9
{
|
(+)ll
(t)(1
52 - 12
(5 A)2
S(np\) - (1 (1-nAi)]1(0011))
+
max
|
|+(t)
52
5²
As
n
1/2
(1)l1 -
+
P.
2
2 - 17
For |a(()]l22 2 2- pC1,
7A?
[t)(1)
22 2 n1Tr)
+1) /C1 S
1
(
1
1
((S-
| S
12 11222
X
101
201
C2
+1)
After plugging in, we have that
nrp
mii1
| 1)12 S || 1)ll2 O(pp²) S h(()),
p + O(np²)
C2
1121 P.
m/111
This concludes the proof.
Lemma 1.1.. Under the condition of Theorem for any t 2 satisfying that (1) xx1) E ()) (2) & S, it holds that t + 1 E S.
Moreover, (+) 2 S2(p²) and || 12 & npde D(p²), then it holds that [| + 1)ll 2 2((2))
Proof of Lemma 1.15. As t S, it holds that
m77
|
2 2 - nd1 + O(np²).
By Lemma 1.10, if we write x'(1 + 1) as shorthand of a(1) A) 1)8(() mpp2(()
a(t) |=(T||| then || +
1) x' (t + 1)ll = O(pp²).
Define 1quad as (x) < 0}. Then we can find a surrogate Tsue (t) such that 2sur E (] guad n.
Khh and || (t) {1)]l22 = O(pp²). We will write d' (t 1) as shorthand of Tuur (1) - 17 (())((()rr (1) (t)
(1) Fauz [+]]l
As [a(t)] M(pp), we have
Faur (t)
i(t)
|=((1]] 1| = O(p). O(p).
Faur (1)l
Hence we have that || + 1) a (t +++1 1)ll = [| - 1)ll + [|='(t+1) - 2 (t+ (++1)|l = O(pp²)
Notice we have ||u (t) 2 2 n P for properly chosen function in the definition S, hence, by Lemma H.5
77,
| (t + 1)l12 -
n P.
This further implies + 1 S.
We also have
(Xaur (t), vi) || (t)l
Ma (t + 01)1 = Kasue (t), U2) - 17/2 (Faur (t), U2) - ?aan
mpa?
= KTaur vz) + | (t) 1)
We will discuss by cases. Let C' satisfies that C== = / (1S + At).
50

---

Published as a conference paper at ICLR 2023
1. If (Fsur(t)) - Cyp,, then we have
vEc² v (G
Ka (t + 1), 01)1 2 (Tour (t), w1 A7 -1 - s2p2)
C
2. If I (1)] 2 Cnp, then as xr(t) E Iz, we have that (s)ur(1), (1), 2 ((pp). Then as || (4) I [|ā())l2 + Onp²²) I Asnp- O(p²), we have that
Xppp
( (t 1), (( (t), v1)l( X7mp 91pp)
1) 2 s(pp)).
By previous approximation results, we have that [| + 1)ll (2(2²).
Lemma 1.16. Under the condition of Theorem 1.3, for any t 2 0 satisfying that (1) x(1) E (Nzerm)))) n
K1Gh/16, (2) t E S, it holds that next(f) is well defined and next(t) +
Proof of Lemma 1.16. Following similar argument in Lemma H.I, we have that x(++1) E (N
Ift $ S, then we can apply Lemma 1.15 to show that + + S.
Lemma 1.17. Under the condition of Theorem I.3, there exists constant C > 0 independent of n and p. assuming that (1) zrt) E (Nje(M]Iy) KTh/8, (2) E S, (3) 21p²) S [[)]( then (1ex 2 | (D) O(n²²).
Proof of Lemmn 1.17. This is by standard approximation as in previous proof and Lemma H.9. Lemma 1.18, Under the condition of Theorem 1.3, there exists constant C >
independent of 1] and P.
assuming that (1) x(t)= (Nie(mnIz) O K7/88, (2) t E S (3) D(p²)
[|1((t]| S 2 2
2-
nd?
+ n23
2- p.
then
1713
| (next(())) 2 min + En)
+
p}.
2 2
1/11
2 11/2
11
or | (xxx((exxx()) 2 min + Cn)ll (0]l,
+
p}.
2
17111
2 11/2
Proof of Lemmn 1.18. In this proof, we will sometime drop the t in (t) A(t). Applying Lemma 1.16, have (e) and (nex(te))()) are well-defined. We can suppose (nex(((() S else the result holds already.
n22
2-
By assumption, we have (1)l 2 M(p)).
Using Lemma I.10,
| + 1) - i(t) + nAū(t) + npa² i(t)
S Onpp?).
||((t]]
Denote
Z(t)
z'(t + = 1 + +nä(t)) + npd
||=())||
as the one step update of SAM on the quadratic approximation of the general loss.
Now using Lemma 1.14 and the induction hypothesis, we have for some C1 and C2 smaller than 1, |=((]] 2
C1 TAi - + || (L + 1)ll < C22-mAn P.
ax3
We will discuss by cases,
nx? 1 If [0(+)]] < C1 2 2-
If next(() = f 1, 1, then
[ + 1)ll
nptm - (1-1) (1-mA1)le (())ll
(2 - C1) - -mA1 + +cn
1
5 (1)l
[[(((]]
C₁
C
51

---

Published as a conference paper at ICLR 2023
As we have (1) S((p2), we have x1((+11 = (2(p2) then [|x( 1) 1)1 ((np²), this implies
1
2 / -
L (t + 1)1l | (t + 1)11 O(10²)
| (1 + 1)11 0
+ 1)]l (1)l
Ci
If next(t) = =t 2, define x' (t+ = z'(t+ 1) At 1 (++1)
( + 1)ll = S2(np),
|
by Lemma 1.9, we have ||a (t + 2) x(t + 2) = (np²))
| (t 2) (mpt) (1 (l [tt (1 ((1-m))
1 (f)ll
|a(t) 1)1l
(nen) (1 77,1) [()]]) (yauy (1 n711) (yauy (1- (1-111))(((
|i((1]] (aaly - ())))))(((()))))
?oas - (1 1711) (yavy (1 - -t))l
||((t|]
21 1/11)2+ 1711
(2 m)11 2 1 4Cn
C1
Combining with (1) 2 22²)), we have that
| (n)xt(()) 2 (1 +1+CC C
2 Case 2 ||+((]] > C1 nA
.> 2 P. then [ S C2 next(() = 1+1
By Lemma 1.1,, || + 1) 2 (1 - ()))
As [|t(next(t)]] < C2 .di 2-mAT similar to the first case,
| (next (t)))l (1 (next(())) 2 (1 + (t)l
In conclusion, if [|ā < 1 | 2 (1 + II (t)ll or | (Dex)(nexx(t))) 2 (1 (t)
nd?
-
n p. we would have there exists C
2-
Lemma I.19, Under the condition of Theorem 1.3, there exists constant T₃ > 0 independent of n and p. we would have that when fALIGN [Tg (n(pp)/]
Kx(t)= (z O(pp),
max ME(2:M]
(xx(t) ((x((t) O(pp²).
Further if x0t') Kh holds for t'
0, ..., fLOCLL then for satisfying AALGGN St S b,OCLL
[xx(() (xx((),, O(p)),
(xx(t) ((=(()).,(==((() ((p²²).
max JE(2:M)
Proof of Lemma 1.19. Let C' be the constant defined in Lemma I.18. By Lemma I.15, we can suppose WLOG || (0) 2 p² and 0 e S. Define
C1 e [lgi+oo|
- m711 /p)l
02
C₂ a C1 [In
0(log(1/p)/m)
max(1 - 1- 02
(2
n23
We will choose LALIGNMID as the minimal / E S, such that (ā) (1) 2
+
2-
P.
52

---

Published as a conference paper at ICLR 2023
Then by induction and Lemmas 1.17 and 1.18, easily have that for / min min(C2 + 1, [ALIGNNM)) and & E S, we have that
xxt) E K K7/8 0 (7,4))
1
| min + Cr)) (Z (0)II:
+
p}
2
2 - 1111
11/2
1
1/
or [[xx(ne min C1) |2 (0)1l,
+
p}
2
2 1/11
1122
The detailed induction is analogous to previous inductive argument and is omitted. If IALIGNMID 2 C₁ then we have for the minimal 2 C1 and f E S
1711
[t)(t)] 2 2 1711 P.
This is a contradiction and we have that AALIGNIID I C₁.
By Lemma 1.17, (next (t))ll 2 [()((t)] Onpp² > for |(: 2 2 (2- then by Lemma I.18,
nA?
+
n13 2- p and and
?a?
+ 1712
[+(t)]| (a) 2 2 7711
2 2- 1-/22/ p.
for C2 2 t (ALIGNMID-
We will then show that for / 2 LALGGNII + C iteration,
| + 1)1l < O(11²))
For Ci 2 12 (ALIGNMID,
X
XB
X
112
1 1-722- 19--0 -0 |z((t]]
S nAO 17p
S
S
||((t||
2
5²
Notice that,
X2
AB
- 1722 - 17p [|ā(+)]|
2 1 - 1122 - 17
(t)
4>3
2(2 - X2)
21 1722
(2 n722) 2 +1+
2 -1+
2 + 3\3
Ai + 3X3
25²
Hence,
(t)x)) (t+ 1)1l2 S max(1 /12
2
1
(1)
52
202
Now by Lemma K.1 and Theorem K.3,
|(22 (())))(2))))( p(2.D) (t+ 1)]l O(1p²)
[1))(t) (t s O(1p²)
[Ar(t) - +((+1)]] (mp²²
By Lemma 1.10, we have that ||~'(1 + 1) 5(L + 1) = O(np))
Combining the above, it holds that
[2 (t+ + 1)ā(t= + 1)1l < max 11²
2
()) +O(pp)
2
4C2
Hence when t FALIGN fALIGNMID C2. | 2 [0)( 2 ((P)) [|p(2 (t) S O(np²).
By xt)) E I1, we easily have [t]l = O(p)). Hence we conclude that | (t)ll Omp))
|| (())(1()))( O(pp)).
The second claim is just another induction similar to previous steps and is omitted well.
53

---

Published as a conference paper at ICLR 2023
1.2.2 TRACKING RIEMANNIAN GRADIENT FLOW
We are now ready to show that ((x(()) will track solution of Equation 7. The main principal of this proof has been introduced in Section 4.3.
Lemma 1.20. Under the condition of Theorem 1.,,for any satisfying that
xxt) E K",
[[] ((pp).
|P(2DD (1))())))))) = -(
it holds that
[0(xtt+ + 1)) - ((x(t)) + ng P ((x())) VA1 (A(t)//1 & O(np" + 1²))
Proof of Lemma 1.20. To begin with, we can approximate ( 1)) - (xx(t)) by its first order Taylor
Expansion, by Lemma F.7,
[0 + 1)) ((x(())) --0---0:))000((----------------- xx()))x())))))))))) - 0
Then by plugging the update rule and another Taylor Expansion,
VL (x)
[0 artt) x(t))-n)00(((()VV=L
VL (x))
VL (z)
(L
VV (x) 1/2112 = O(1p²)
VVL (x)
Using Lemma F.3, we have
VL (1]l0'
VL VL
|Imp@ (x0))
[v) [V) (z)
VL(()
Putting together, we have that
VI. VI
[10(xtt+1)) - (000(+=)))0o00000-n//aa200000000
[1) [v) 1/21
50(P'0² + n + (npl VL (x) I))
As we have ||ā(t)]| (mp). hence by Lemmas F.2 and 1.7, + (0)((+=))))0000000mmmm2000011110
(z(((t) VL(x(t))
VLL VL (z((())) 1/211
50(Tp" p²²
Finally, we have that
- (10) (z((()) (L((0)
VL(x(1)
(L. (V)()))) 1/2
< 0(np²)
as the angle between vL(x) and 01 it) O(p).
[[L(]]]
By Lemma E3, it holds that
Putting together we have that,
[x((( + 1)) Pm +2p
It completes the proof.
54

---

Published as a conference paper at ICLR 2023
1.3 PROOF OF THEOREM 4.5
Proof of Theorem 4.5. By Theorem 1.1, there exists constant 7i independent of 77, such that for any Ti > Ti independent of 7)s 04 it holds that
max R)(x(t)) O(np²).
max (n(L/n))SnngTT (n(1/pp)) JE[M]
max (n(1/n))Sm (m(p/pp)
[[(()(t)) - ((rmi))) - O in(1np))
By Assumption 1.2, there exists step Tj In(1/np) S MPHAAS & (n/ such that
max R (2(tpHASE) ( JE[M]
((((((P)))E) ( | O((1 p) p)((/p)))
(a(trnASE) ((Z((PHASE)) 11 (ttrnssss)) 2(p²). ||<(|HASB)\ll S (1(
Hence by Theorem I.3, if we consider a translated process with zt) x(t + (PHASE), we would have for any T3 such that the solution X of Equation 7 is well defined, we have that for t = -
(z) ((n2t))ll (g)
This implies for satisfying (tnp2(t (rHASE)) is well-defined,
= [0((x))) - O(n)n((1p))
Finally,
(t (PHASE) X(np2t)ll2 = (O(nptrHASE) = =(ph((1/m)))=0mmmm = O(1
We have that
[0(x(t)) =(nn)//\))))))
The alignment result is a direct consequence of Theorem 1.3.
I.4 PROOFS OF COROLLARIES 4.6 AND 4.7
Proof of Corollary 4.6. We will do Taylor expansion I Max By Theorem I.1 and L.3, have | (T>/p=²]) - X[TT))l = Ö(n + p) and [<<([T//pp])) - (<<(TT/nn))))) = O(p)). For conve-
nience, denote (T 21) by
R (x)
max pv p VL(x) p² Lz))0/2
1
Since max/ 511 || LLx))l O(1 (xx) = (mp). it holds that
R-Max (x) = 0² max ,7 (F p²)
1
= -2 P2 +V
= ==XX())
which completes the proof.
Proof of Corollary 4.7. We choose such that X (T.)) is sufficiently close to X(2o), such that X (X(T.) <
A (oo) e/2. By Corollary 4.6 (let T3 TR), we have that for all p. 17 such that ry ll(pp)) and p/n are sufficiently small, || (x ( - (XT)) & o(1). This further implies (n00()
< +01). We also have (TTT[///)))) (1).. Then can leverage Theorem G.6 and Theorem G.3 get the desired bound.
L(x)
55

---

Published as a conference paper at ICLR 2023
1.5 DERIVATIONS FOR SECTION 4.3
We will first show our derivation of Equation 9.
In Phase II, x(t) is (llp)-close the manifold I' and therefore it can be shown that | (<<(t)) ((p) holds for every step in Phase II. This also implies that || + 1) =(t)ll2 = ((p)) (See Lemma F.7).
Using Taylor expansion around z(t), we have that
(xx(L 1)) (x((() 0 (x(1))
(0) (11]=((1+ + xx(11))
(((((()
[=(((t)) O(T
(31)
For any I E R. applying Taylor expansion on VL P P]]V(())1>
around z, we have that
VL(x)
V + (VL(2)]
(((x) [ 2
+ e (V)()))
VL(x)
Vi(r)
-V((x) + pVP L(x)
+ O
(32)
[V [V
Using Equation 32 with x = x(t), plugging in Equation 31 and then rearranging, we have that
((x(1 + ((x(()) + np²
VL(x(() VL(x(t)
(0((())
(L(((()) [M<(0)]
vL(x(t)) + 00''p²+n'')) =-126((())VVL((() (100(=())V)(((0) ( [VL((())]]
By Lemma 3.1, we have that (z(2()))L((((( = 0. Furthermore, by Lemma F,5, we have that 0 (3 0. This implies that
0 ((((())V = (0((((())))))((((() + (()=(()) - (x((1)))) O(1p)-
Thus conclude that
np²
VL(x(() Vl(x(1)
+ (z(t + 1)) ((x(t))
2 (
(VL(()))) |
+0 +
We will then show our derivation Equation 10
(x 1)) (z(())
n1² 00(G()))FF
Vl(a(t)) VL(x(t) V 0(1²p² np²)
n1²
2 (4440000887411000044444444044414 -0, 900)
m/p²
2 (00
np²
0 (4
no²),
(10)
2
where the second last step we use the property the derivative of eigenvalue (Lemma K.7) and the last step is due to Taylor expansion of aT(-) V (V2TL) L(-)) (x(()) and the fact that [[*()(()) - .x)lll = ((P))
We will finally show our derivation of Equation 12.
The update of the gradient (Equation 12) can be viewed as an O(np²) p-perturbed version of the update of the iterate in the quadratic case. Note O(pp²) is a higher order term comparing to the other two terms, which are on the order of O(r/²) p) and ((п)) respectively. By controlling the emor terms, the mechanism and analysis of the implicit alignment between Hessian and gradient still apply to the general case. We can also show that once this alignment happens, it will be kept until the end of our analysis, which is O' steps.
56

---

Published as a conference paper at ICLR 2023
Finally, derive Equation 12 by Taylor expansion. We first apply Taylor expansion (Equation 32) the update rule of the iterate of SAM (Fquation 3):
VL((()) (L((()) + O(10²).
= -
(33)
Since phase happens in an -nigh--nighbooroooo of manifold r, we have || +1 - xx11]12 = O(pp).
Then by Equation 33 and Taylor expansion on VL(x(1 + 1)) x()), we have that
Vi(x(t+ -))=
( P²)
VL(x(t)) Onp²)) (34)
VI(x(())]
ANALYSIS FOR 1-AMM (PROOF OF THEOREM 5.4)
The goal of this section is to prove the following theorem.
Theorem 5.4. Let {x(t)} be the iterates of I-SAM (Equation 13) and x²0) = Timit E U. then under Set- ting 5.1, for almost every Tiin for all and such that (1 + p) In(1/np) is sufficiently small, with probability at least 1 (p)) over the randomness of the algorithm, the dynamics of 1-SAM (Equation 13) can be split into two phases:
Phase I (Theorem J. 1-SAM follows Gradient Flow with respect to T. until entering an Ö(10) neighbor- hood of the manifold 1 in O(I((1//n))))) (n//(//////))))))) steps;
Phase II (Theorem J.2): 1-SAM tracks the solution of Equation 14, X, the Riemannian gradient flow with respect to L()V) in Ö(1p) neighborhood of manifold T. Quantitatively, the approximation error between the iterates x and the corresponding limiting flow X is (11 /2 + p). that is,
[a<((Ts/(np"))) -
X (Ta)lla 01 + p).
As mentioned in our proof setups in Appendix E, we will prove Theorem 5.4 under more general (and weaker) condition, namely Condition E. and Assumption 3.2. The only usage of Setting 5.1 in the proof is Theorems 5.2 and E.2, which are restated below.
Theorem 5.2. Loss L. set I' and integer M defined in Setting 5.1 satisfy Assumption 3.2.
Theorem E.2. Setting 5.1 implies Condition E.
M EM Ek=1 Lg. For each k E [M]. L& is C+, and there exists a (D 1)- dimensional C² (e-summanifold of TR. D. Tk. where for all x E TR, x is global minimizer of L,k, L (x) = and rank( La (x)) 1. Moreover, I =M I
Condition E.1. Total loss
I' for I defined in Assumption 3.2.
Analogous to the full-batch setting, we will split the trajectory into two phases.
Theorem J.1 (Phase I). Let {x(t)) be the iterates defined by SAM (Equation 13) and x=0) Finit U, then under Assumption 3.2 and E.I.for almost every Tinit there exists a constant T2, it holds for sufficiently small (n + p) In 1 /p,, we have with probability - (p), there exists i & Ti In(1 In1//p)//n, such that ||x((() (z((t))l (mp) and || (2 (=(())) 0 + pp)
Theorem J.1 shows that SAM will converges to an Ö(np) neighborhood of the manifold without getting far away from (x))))). where we can perform a local analysis on the trajectory of (=x())).
Under Assumptions 3.2 and E. we have Tr(V²La(<) = Lx is differentiable for x E r,, Hence Tr(V² [r))) EM Tr (x))) is also differentiable and we have (14) is well defined for some finite
time T2.
Theorem J.2 (Phase II). Let {a(1)} be the iterates defined by SAM (Equation 13) under Assumptions 3.2 = 0(n + p). for almost every 2(0). for any T2 > till which solution of (14) X exists, for sufficiently small (n + p) In 1/ (pp)),
and E.I, assuming (1) [0(0) (x()))) = O(1p) and (2) [[ - (xx(
we have with probability 1 ((p), for all np²t T2. [((((() X(np21)l12 = 02 + p) and [|x(())
((((t))) = O(p)).
57

---

Published as a conference paper at ICLR 2023
Combining Theorems E.2, J.1 and J.2, the proof of Theorem 5.4 is clear and deferred it Appendix J.3.
Now we recall our notations for stochastic setting with batch size one. Notations for Stochastic Setting: Since La is rank-1 on & for each &o E [M], we write it as (xx (z)x)ux(z))) (x) for any Tn. where Wi is a continuous function on I'ss with pointwise unit norm. Given the loss function La, its gradient flow denoted by mapping ΦR RD [0, 00) + RD Here, (x( denotes the iterate at time T of
gradient flow starting at and is defined the unique solution of (x, 7) Va E IRP. We further define the limiting map 4 as (x(x) l Φ 7), that is, (x) denotes
Ju VL (ov (2t, ())
the convergent point of the gradient flow starting from z. Similar to Definition 3.3, we define U2 = {or e
R (x) exists and 2x (2)
] be the attraction set of T, We have that each U3 is open and 2 is e on UR by Lemma B.15 in Arora et al. (2022).
In this section we will define as (X | E [0, T3]} where is the solution of (14). We will denote h(K) in Lemma E.6 by h. Using Theorem D.3, will assume the update is always well defined.
J.1 PHASE I (PROOF OF THEOREM J.1)
Proof of Theorem J.1. The proof consists of two steps.
1. Tracking Gradient Flow. By Lemma J.3, with probability 1 p², there exists step lGF O(1/n) such that
[|<(trr) (x((or))ll s h/4.
[[((x(toF) ( lll Ö(1 + p).
2. Decreasing Loss. By Lemma J.7, with probability O(p), there exists step IDEC tGF + O(I((1p))/m) OlI((1/p)/n) such that
[VL(c(Decc)) = 0 [*(((te)) - (2mi) - [[((x(tDEc) - (x((tF))) +
[x((F()F) ||
4 (Xiii)l2
= /2 + p).
Then by Lemma J.12, with probability 1 O(p), there exists step (DEC2 = fDEC + Oln(L/np)/n) (ln(1/np)//) holds that [{<(()))2) - 4 (3((mmc22)) = O(p)). ([((EE))))) 2 [|@(<((DECa)) 2 (x((EE))) [[(((()))) 4 (Xinit)
=0 + p).
Concluding, let Ti be the constant satisfying (DEC2 S Ti In(1/np) then we have for = IDEC2 < Ti In(//pp)/n such that
||x(t) ((x(t))ll2 O(pp).
(+((()) ((rinte) ll2 = 022 + p)-
J.1.1 TRACKING GRADIENT FL.OW
Lemma J.3 shows that the iterates ztf) tracks gradient flow to O(1) neighbor of
Lemma J.3. Under condition of Theorem J.1, with probability 1 O there exists lce O(1/n), such that the iterate æ((GF) is O(1) close to the manifold I and (<((tFF) is O22 + p) is close to (Z Quantitatively,
L(x(tor)) < uh²
32
|x((or) - ((r(ter)) h/4,
[ (( 0(
Proof of Lemma J.3, Choose C C= C=
58

---

Published as a conference paper at ICLR 2023
There exists T O, such that
[ 7) 4 (Tinit) ) S Ch/2.
Consider
VL (I(1t)
z(t + 1) = x(t) - (z)
VL (z(())
= =x(t) 7y V LA (
By Theorem L.I, let b(x)) VV(x)) 7) and Op)), for sufficiently small 1) and the iterates x(t) tracks gradient flow o(Tinit: I") in O(1/n) steps in expectation, Quantitatively, with probability p², for
[o]] we have that lcF = 7
|I<(Ce)) - (z To)lla o) + P))
This implies x(tcF) Khh hence by Taylor Expansion on Φ
|[((<((FF) & (Zima) |{ - (dd(Zinit T))l
- O(1xx(toF) (DFinits l2
s 0n 1/2 +
This implies
[I<(tr)) - ((x(te))) S [|x((tr)) - aimi To)ll2 + [1 To) +
+ (1(((in))) (x(((F)))
<Ch/2 0 + p) s Ch h/4.
By Taylor Expansion,
[x((trr) S ([|x(ter) ([o(tr)))11/2 nh²
32
J.1.2 DECREASING LOsS
Lemma J.4. Under condition of Theorem J.1, assuming x(lo) E Kh/s and for any satisfying to S 555
max to + O(Im(1/np)/m) (o$rsttt(lu(1/o)/11 L(x(T)) I wh² 16
holds that
x(() E Khh Vla ST.
Moreover, have that
[+(((t)) (x((t))) 0((n + p) In(1/np)).
Proof of Lemma J.4. We will prove by induction. For to. the result holds trivially. Suppose the result holds for 1, then for any satisfying lu S t S L -- 1, by Lemmas F.I and F.8,
[0(a( 1)) o2 (x(())) _ Expll VL (z((T) ||2 unp² + (B))
0 + mp).
Also by Lemma F.1, [|x(()) (x(())) s 1/22/,, this implies,
dist(K, K, x()) <dist( (aastt(Kx :<(o)) + |x((t)) +<x((o))]l2
+ (z)(((())) - (((()))]++ - a(!)]l
<0.99h + O - tGr)) = 0.99h (n) h.
59

---

Published as a conference paper at ICLR 2023
Lemma J.5. Under condition of Theorem J.1, J. ifx(r) E Kh, then we have that
([[ & L(x(r)) T/h L(x(T)).
Moreover it holds that,
Eln (( ( I(x(r) 7/
Proof of Lemma J.5. By Lemma F.8 and Taylor Expansion,
(1 + (7)]
VL (x(T))
= L (<(T)
VT (<(7)) (((7))1] (4(T
((((T) 77 VL VL x1()) + VL (2((T)
<B
+ 2
VL (x(())
IVLA (x(7) VL (<()))
2
(L(x(T)) 7/VL
7 VVL (z(7)) lle + 1²20²p²
<L(((T)) 7 (x((T)
<L(x(T)) TZ4 L(x(T))
T/
Lemma J.6. Under condition of Theorem J.1, assuming x=lo) Ka/4 and L(x(lo)) < pbb
32. then with
probability - (p), for any t satisfying to S t S to + ((n(L1/p))/ /T), it holds that xt) Kh. Moreover, we have that
[(((((t) (=(())))) O(1 Im(1/pp)).
Proof of Lemma J.6. By Uniform Bound and Lemma J.4,
P(3to <t to + 0(m((1/m))/))) nuh
16
to (0t((((((n)m/m)
μh2
and L(x(r)) < nh²
P(L(x(t)) 2
vto
S
16
16
i=
to (0tO0((((/m))/m)
P(L(x(t)) 2 μh²
and I(T) E K", vto ST
16
t=
Consider each term, and applying uniform bound again,
P(L(x(t)) 2 μh?
and x(T) E K²h Vto STSt-1
16
P(L(x(0) > wh²
and L(x(7)) < uh?
16
32
1I 12 2 + 1, nhi
wh2
and
> [(x(7')>>
16
32
and - - 7" 2 ",,(("" E k*).
Then if we consider each term, we have that it is bounded by P(L(x(1)) 2 ph²
1h2
16 and VI - MV-12/7+1,
32
nh2
and
Vt 1 2 " T, E k L(x(T)) <
32
60

---

Published as a conference paper at ICLR 2023
Define coupled process [LT + = 1)= I 1)) and
In (xxr))) L(T' -nnp2
if if - in - 1)) 2 In(
T(7')
32
if otherwise.
Then clearly
P(L(x(1)) 2 ph²
ph²
and 2 (22+22+11((((7))
16
32
whh²
and 2 2 T, x((")) E k ((((t)) -
32
< P(L(() 2 In( nuh
16 ))
Consider a fixed 7' satisfying T + < +' S L. By Lemma J.5, we have that
[(((r' + 1)) [xx(7') S -1p/2.
Hence É(I) + mptl/2 a super martingale.
Further it holds that if L(x(r' - 1)) 2 ( nk?
32 then
M(x(T' - 1)) - [(x(T')) = 0(1 - x(-')l)=
Using the smoothness at log(x) at nh² 32
which is positive constant,
| + i(>')l| S O(1) S Cn.
Here C is constant independent of This implies [(x(T + 1)) S
16v2
Now by Azuma-Hoeffüle bound (Lemma K.4), we have that
a²
P(i(() - i(T + +)+ (1 - ( 1)mu/2 > a) & 2exp(
8(t 5 1)(C 12
pa² With In( 1n ) + (t - 1)1//22 (In + (t (t--- - 1) have that
P(ī(t)= In( the
(In + (t T - 1)n
)) exp(
16
32(C + 22)
In 2(t - T -1)pg
< 2eexp(
8(C 11)2.11
Hence have
P3t I Sts O((((((/p))/)) 2 nuh²
16
2n - - |n2 (1/np/ /p/p//////))) p. 8(C + 11)
20(2exp(
Hence with probability 1 1 p. [(x(t)) have completed our proof.
S
ah² Vta t to + ((m((1ppp//)) combining with Lemma J.4, 16
Lemma J.7. Under condition of Theorem J.1, assuming there exists foF such that L(x(ter)) <
wh² 32 and
(tGr) E Kh//4 then with probability 1- (p)). there exists tDec = tgr + (Oln(1/))))) such that x(tDec) is in ((p) neighbor o. quantitatively, have that
[V((((ne)))] S 4Cp-
Moreover the movement of the projection of (x(-)) on the manifold is bounded,
[[((z((oF) - (x((DE))) O(17 + ln(
61

---

Published as a conference paper at ICLR 2023
Proof of Lemma J.7. For simplicity of writing, define T₁ - 258
(1n((1/p)/m)
By Lemma J.6, we may assume axt) K for top I S T tcp.
Define indicator function as
A(t) - (VL 2 VA tGF]
By Lemma J.5, we have that,
E[L(x(t + 1)) + 1)] (xL(t + ()) $ (1 1/ ())(()(t))A(t)).
We can then conclude that with Σ3 Ti + tGF. using Lemma E.2,
/ [(x((or)) & Sup².
Spp²RA((2++1) + <
+ +)] <
2
We have
EA(T2 + 1) S
This implics A(22 - 1) = 0 with probability 1 O()), which indicates the existence of DEC- The second claim is direct application of Lemma J.6.
Lemma J.8 (A general version of Lemma 5.5). Under Assumption 3.2 and Condition E.I. for x E Kh and PEC, L. = As (p)un(p))u there exists E (1,
V La (x)
swx(p) + O(1 - pll2)-
V LE (x)
Further if |f (x p)l | 2/13/2 then s s= sign (a p)). This implies
VL (I)
V [I) (x 2 swm (x p) O(I -pl12)
2 I1 (x p)ll2 O(1 lr - 2/13/2)
Proof of Lemma J.8. We will calculate the direction of vL((2) compare them to get our result.
[Vxx(]] using two different approximations and
1. According to Lemma F.4
(=) + +(II= (x)2))) L (x) (= +x
L (x)
Suppose V (x)) = A+ ( (x))wx ( (Q then
VL (x)
L (z) wy o (z)) O (x(x)12)
As (p) = Ax(p) (p) using Davis-Kahan Theorem K.3, would have Es {-1,1},
such that |ww (xx())) - (w)(w)) 2 - (1) pll2.
LL (z)
swn (p) 0(11x(=) - pllz + || pll2).
LL (x)
According to Lemma F.1, we have || $x(x)]12 < [Vln(x))l < Cll=-plla, This implies,
"
L (x:)
L (x) sw + +[1lz pllz)-
(35)
Equation 35 is our first statement.
62

---

Published as a conference paper at ICLR 2023
2. By Taylor expansion at p.
LE (x) = (())))) px + - -pl1))
That being said, when |o (a p)] 2 || - p12 /// we have
L (x) Ammww (x p)ll2 S O - -pll))
VI (x) || 2 W (a p)ll2 - O(1 pll2)= 2 0211 p12
Concluding,
L1 (a:
A (x p) Am (I - p)l |/2 <0(1x -1ll//2
L (x)
Hence we have
VL (z)
L (I) = sign("? (x p-)) + O - -pll2
(36)
Comparing (35) and (36), we have sign(w>(p) (x p)) when m (rr - p)] 2 || pll2 3/2
Lemma J.9. Under condition of Theorem J.1, for
constant Ci > C2 > 0 independent of n, P. if) E Kh constant C > 0 independent of ny P, there exists
and Cimp v [1((1) - (xx(1))) S Cp,, then we have
that
Ex (xx(t + 1)) x=t)] ||=(() + (x(1))]l C2np.
Proof of Lemma J.9. By Lemma F.2, |z(() (x((t))| =O))). Hence we have that by Taylor Expansion,
L1 (z(())
+l = x(1) L (z(() +
L P1
(x((1)
LA (a((1)
=x(() V (x((1) V2 L (x(())
O11p²)
L (z<(t)
V L. (x(())
= xt) L (x(1)) - 7p Here Ag, u% indicates A₁ (o(r()))) wx (+x()))
+ O(ne²).
VL (=(())
Notice that given ||a(t)) ((x(t))] = O(p). by Lemma F.8, we have that [| + 1)) - (1)())] O(1p²) ||0(| + - x((]]] ((1p).
This implies 1((+1)E
Further by Taylor Expansion, (Ix(x(()) = Am (x((1- - (((=(()))+((())
By Lemma J.8, we have for some (t) 1,,
LL (x(t))
wk
LV (x()) (k(t)u 0(1x(t)) (=(())12
We also have
(t) # sign(wz (=(1) -(x(())))))) 1oll (x(() (=(1)))] < |a(()) (,(()))
(37)
Concluding,
+ -xx(1)+(((++1)
(xx(1) - (c((())) 1 (a(1) (((())) Ap m0w wa O(p²).
After we take square and expectation,
El||((t 1) (x + 1)) xx))]
2n²²
2n²p²
<||z(t) - 0(x(t)]1
Aglny (=(() ((((t))) +
A
M
M
2 " M
M
A& wi (I((t= - 1 (2(1))) -2 nP
A (t)w (x(t) ((x()))
M k=1
M
+ O0mp7 - ((x(t))) + ))P)))
63

---

Published as a conference paper at ICLR 2023
We will then carefully examine each positive term,
2n7
Ai (x((1) ((x(1))))= 2M1 (0000000000000000000000000000000
- -(00)((0000000000=====000000 +(x(1))
M
2n²² p² M
A - =
M
k=1
This implies,
Bl|la( + 1) (x((t + 1))1 x(t)
<|>(() - (((())] - nP M
M Σ
( (x(t) (()(t)))
+ +0(n02llz(t) (xx())) + p²).
We will now lower bound Σxe1 (t) (x(t) +(x(t))) By Equation 37,
M
M
(Esu(t)ur (x(1) +(x(t))) N
A (z((1) ((x(1))) - 2 2 Ax|l=(t) (x((t)))
k=1
M Au (I(t) - - (=((t)))
2
For
E Aulluy (zr(1) (x( by Lemma Lemma F.4.
M
Aelluy (a((t) (<(())))22 2
A -
(x(()) (x(()) (1)00=/())))) (x((t))
= [ -
2 (110)( (<x())))
- 2 -
- (x(()))
Concluding, we have that
M
Assu (t) (x(t) -=((t)))) 2 ml|x(() (x(1))l222
So
E[]|((1+ (x((++ 1)>1 ax(
- <(1a(() +xx())11 1m9 ||-(() - (<((t))
M
(0(m)))((()))))))))))))))))))( - (<((t)) + +
[[]x((1) - (xx()))l - 2
The inequality holds if ||a(()) (=(())) > Cinp.
Finally by Jenson's 's Inequality,
E[||a(1 + 1) (x(( + 1))laa()) < [|x(1) (x(1))l C2np.
Lemma J.10. Under condition of Theorem J.1, for any constant C > 0 independent of n) P. there exists constant C3 > independent of 1), p. a() E Kh and [x((t) #(x(t)]] < Cp, then we have that
||| + 1) ((x(t 1)) (z(() ((x(t)))l S Camp.
64

---

Published as a conference paper at ICLR 2023
Proof of Lemma J.10. This is direct application of Lemma F.8.
Lemma J.11. Under condition of Theorem J.1, assuming s<(a) E Kh/2 and [|<((o) ((+(()))|| - f(n p)
(nnpplnnn////// (L/np)) n O(p), then with probability O(p). for any for some fixed function f and f(n. p) E Q(p (nmplln(g
satisfying to S + (lm((L(p)m/)) it holds that [| - (xx(t)) S 2f( p). Moreover, we have that [[(()(()) (x)(()))) = O(1 + p) In(1/np)).
Proof of Lemma J.1I. By Lemma J.6, we have that x(I) E Kh for any satisfying that t0 S lo + ((I((1/pp)//) and with probability - (p) will suppose this hold for the following deduction.
By Uniform Bound,
P(3 S1S to + O(In(1/no)/))) [=((1) (x((1)) 2 25(1 p))
16+0 (
P([x(()) - ([(=(())(<(())]] 2 p)) and ||((T) ((x(T)) s 2f(()
St
Consider each term and apply Uniform bound again, ((()a(()) - क (x(())) 2 p)) and [|((() (n(r)-<(zzf)) +(x(T))[ | (% st
v4 -((=(t)-<(<(t)))) 2 p)) and - $ f(n.))
==to and ((,p) s |m((²) ((((())) S (2
Then if we consider each term, is bounded by (([)=[()) - (xx(()) 2 2/(1))) p)) and p) < |x((-)) - I - (=(r)) - p)n.
st
(38)
Now let C' be the positive constant satisfying 2f p) S Cp, suppose C₁, C2 are the constants corresponds to C' in Lemma J.9 and C3 is the constant correspond to C in Lemma J.10. By definition C3 > Cz.
Define a coupled process 5( + 1) = u(r+ + 1) and
||0((²) - (x(T'))) 8(7') = (5(T' - 1) C2mp, if otherwiiee.
if - 1) = || - 1) - D(x(1' - 1))122 > f(n.p)
Now clearly Equation 38 is bounded by P(t(t) 2 p).
As [[)( O(T'= 1) C2np by Lemma J.9 and (|ū(r') - 1)ll I Canp by Lemma J.10. This
implies [[( - C2npr' is a super martingale. By Azuma-Hoefiding bound(Lemma K.4), we have
ha
P(T(() 2 7(t C2npl1 - 1)+ + h) s 2exp( 4(t 4 - 22p
Choosing h C2np(t 1|=1 + - 4<<(7 1))1l + 2f((pp) P(i(t + 2 (f)
(C2nol1 - -T) - ||-((T + 1) <<((T 1))1 +))|+++200 2f( p))
<2exp(
8(t - (C3 + 2n p²
(Czmolt - 7) + (11,p)/23
<2expl
4(1-T) (C + 122
C2f(1,)) p)
<2exp(
S 1/ 20
2(C's + C2)
We then have
P(3 5t5 to + (ll((Lnp))/)) ||=(1) - ((x(1)] 2 2f(1.))) p)) S
65

---

Published as a conference paper at ICLR 2023
Lemma J.12. Under condition of Theorem J. assuming there exists IDEC such that x((DEC) E Khh/2 and [V((((DE))) & 4Cp, then with probability - O(p), there exists /DE22 = IDEC + (0(n(1/np))))))) such that (x<()))2) (tDEC2)) S ((p)).
Furthermore, for any satisfying IDEC2 S t S IDEC2 + @(In(I/np)/n) we have that [[(((()) - ((x((DR))) Opp² In(1/np)).
Proof of Lemma J.12. We have that xxt) E K² (Lemma J.6) and ||x(t) (x(t))) S Op for some constant C (Lemma J.11) for any satisfying that IDEC S t S IDEC + (In(1/np)/n) with probability O(p)
and we will suppose this holds for the following deduction. The second statement then follows directly from Lemma F.8.
Let C1, C2 be the constant in Lemma J.9 corresponding to C, For simplicity of writing, define Ti -
C'lnl C = ((In(1/np)/n) /n). Define indicator function as A()) = [[]=(()) - (=(())) > Ciip, VI T 2 LGF)
Can
By Lemma J.9, we have that,
Bi||z(t + 1) (x((t + (t)]IAA 1)] s [||| +
= [[] ((z(t))] C2mPE[A(t) s [[](((1) +(=(()])4(([(11 Can)
We can then conclude that with 7: = Ti + IDEC, using Lemma F.2, CinpEA(T22 + +) s [T] + 1) 4 (x((2 - 1))1l2A(33 +
C27 nTh T (I=(E)))) (<<(tDE))) S Cinp'.
s (1
This implies A(T2 1) with probability O(p), which indicates the existence of 1DEC2.
J.2 PHASE II (PROOF OF THEOREM J.2)
Proof of Theorem J.2. We will inductively prove the following induction hypothesis P(() holds with proba- bility 1 Og1 p) for 1 < Tz/np² + 1,
- - xxr) Kh//2 stt
[x((7) - (x(r))]2 S
O(m))
[}(x(T)) 000 p)
P(O) holds trivially, Now suppose P(I) holds, then x(1 11) E K*. By Lemma J.13, we have that with probability - 0(1 1| 0( + 1))]| - 21|=(()) (() 0(1p).
Now we have
2| x(x 0(pp) 1+
< K"h , 55111
By Lemma J.14, it holds that
[ + 1)) - ((x(r)) + n P (((())) VA1 LLL (o(x(T))) /21 s Ö(pp' + pz)
P1
As
PA ( VA1 V L (9(z(t))) P P VI(V-L((((())))
By Theorem L.1, let b(x) = - (T)
(x(x) = (00)) (z)), ) p = upp and
Onn + p). it holds that, with probabilit Onn² )²))
(+(x(r)) - x(1p²7))
=0(1 -
( (ximit)) + Tamp² + / log22 (p n)Ta)
This implies = + 0 + p) < h/2. Hence x(( + 1) Kh/2 Combining with P()) holds with probability 1 - 0(12>²t))
(t+1)( -((t ( Xx =
we have that P(L + 1) holds with probability 1 O(r' + 1)). The induction is complete.
Now P( P([Ts/np²1) is equivalent to theorem.
66

---

Published as a conference paper at ICLR 2023
J.2.1 CONVERGENCE NEAR MANIFOLD
Lemma J.13. Under condition of Theorem J.2, assuming x(t) v St to + ((1/pp²) and [I<(t)) - (xx(()))] S f(n. p) for some fixed function f and [(n, E (1 (p), then with probability O p'), for any t satisfying to I t S to + O(1/no"), it holds that |x(() +(x(t))] S pf.
Proof of Lemma J.13. The proof is almost identical to Lemma J.11 and omitted.
J.2.2 TRACKING RIEMANNIAN GRADIENT FLOW
Lemma J.14. Under the condition of Theorem J.2, for any t satisfying that x(t) E Khh and
|x((t) ((x(t)) (nnpln((L/pp))
It holds that
11 + 1)) - (x((t) + np² P ((2(0)), VA (f(x(t))) s Ö(np² 1²)
Proof of Lemma J.14. We will abbreviate k, by l in this proof.
By Taylor Expansion,
xlt + =++1==x(0 - L/LL (a(1) + VL (x(t))
LL (z(())
LL (x(t))
= =x(t) - L (z(t)) - 7 (x(t))
LL (z(t)
n VL (a(t))
VL (I(t))
LL (I(t)) | |VL, (z(t)) 1/2 + O(pp²).
Now as ||x(1)= (((()))ll2 Ö(1p). by Lemma F.8, it implies | + 1) 2t)]l = O(pp).
Then we have
[0 + 1)) ((x(t)) (((())((((++
- -(()))))(((()))((()))))) s 11) O31²p²)
Using Lemma F.6, we have
(I) (z(1))
((((1)) -(())((())))) ( 017i0 +n²²)
(())==(1)
VL (x(())
(=n00 (x(1)
VL, (x(1) [|2 O(pll=((1) +(x(1))l22 ÖGr20² + n
Hence
[1((t - -)))) - (2)
VL, (I(t)) VL (x(t))
Ö + LL. (x(1)) (z(1) 1/212 1/211
VL (I
Notice finally that by Lemma J.8,
VL (x(1)) (Le
( L (x(t) V (x( 10 -
= To (x)()l)
Hence have
((x(t+ + (x((t)) = P VAi (PP 06? + n0²)
This completes the proof.
67

---

Published as a conference paper at ICLR 2023
J.3 PROOF OF THEOREM 5.4
Proof of Theorem 5.4. By Theorem J.1, there exists constant Ti independent of Ps such that there exists IPHASE S Ti In(1/np)/n, with probability O(p), holds that
[<((hAsS)) - 2 (=(((HHs)) = ((p)).
[((((tPHASE) ((iin)) 0(1)
Hence by Theorem J.2, if we consider a translated process with 2t) x(t + (PHASE), we would have for any T3 such that the solution X of Equation 14 is well defined, we have that for 1i
np
|+ (t)) - ()n (n)
This implies for satisfying X (ngp(t (PHASE]) is well-defined,
[|{(x(t)) - X(n0² (pHAs)))] = 02) + p).
Finally, as
X (1- (ommssn)))x(mmp))))))
X(nx2t))22 O(nn = 100
Ö(p).
We have that We also have
[0(x(t)) X(np²t)ll2 0 p).
[x(()) - (ox((((x)))))))) (np))
by Theorem J.2.
J.4 PROOFS OF COROLLARIES 5.6 AND 5.7
By Theorem J.1 and J.2, we have Proof of Corollary 5.6. We will do Taylor expansion on <(Ta///p = Ö11 22 + p). For
[|x[[Ta// (TTTa//>])> (Ta)lll = (122/2 + p) and [|<<<((Z//np]1)
convenience, we denote 2 [Ts/np²1) by
Max Ex 1e k,
(x) = max [ VL +
||
Since olls1 10 V (z)x = O 4 (z))l)) = 00) holds that,
Eg [
Max (3x) = p2 max
(=L(E))/2)+ 114))
||||
= =p'Ex max
V= + -1//))
= =((T((((T)))2+++(( +44)))p7)))
Proof of Corollary 5.7. We choose T. such that XTT) is sufficiently close to X(oo), such that Tr(X < Tr(X (co)) e/2. By corollary 5.6 (let T3 T.), have for all p. such that (n p) In(1 (1 (1/p)) is sufficiently small, [|E MMax (o(([/((o))) - (TT) < o(1). This further implies
(TTr(X(0)))/21) s cp²/2 /22/2210 o(1). We also have L(x([Te/(Mpp)]] infzeu T(x))= = (1). Then we leverage Theorems G.6 and G.14 get the desired bound.
J.5 OTHER OMITTED PROOFS FOR 1-SAM
We will use ('y,, 9k) and (" yk) to denote
dE(y' dy y= and
L 3p2
dzy"
1z'=
Lemma J.15. Under Setting 5.1, fix k E for any p satisfying ((f( yn) have that
2
68

---

Published as a conference paper at ICLR 2023
Proof of Lemma J.15. ((fx(p), wb) = implies (f=(()) 0. Then by Taylor Expansion,
V2 (D) VzE((A(D), me)
- 3.
+ 3))
=f'((e(()) uK V
This concludes the proof.
Proof of Lemma 5.5. By Lemma J.15, as L(p) = 1/4 E T [p),,,, we have
1
v= M
(ōy')² (x) V
By definition of L in Setting 5.1, we have for any p E r, (v/x()) are linearly independent, which implies that fi (p) + O for any p E T.
For any p E as V(s(p) is well defined and continuous at p, there exists open ball V containing p such [v
that for any x E V, V (x) 12 2 C > 0 and | V lll S C2 for some constants C₁ and C2.
Suppose V 7 (z) * 0, then as by Taylor Expansion,
(Lx(2) = (f) yk Vz V/ii) M
We have (a)(a [ [v
Va(a)
C2llx -el1-~11 pll. which completes the proof.
We note that the alignment result in Lemma 5.5 is not directly used in our proof. Instead, we use its gener- alized version Lemma J.8 which holds under holds under a more general condition than Setting 5.1, namely Condition E.I.
K TECHNICAL LEMMAS
Lemma K.1 (Corollary 4.3.15 in Horn et al. (2012)). Ler L, Dx be symmetric and non-negative with eigenvalues d 2 AD and d 2 ... 2 Åα. then for any i,
1² A S E E112
Definition K.2 (Unitary invariant norms). A matrix norm || l| the space of matrices in Rexd is unitary invariant if for any matrix K E Rpxd ||KKW||. = || for any unitary matrices U E RPxP, W E Rddd
Theorem K.3. [Davis-Kahan sin(O) theorem (Davis et al., 1970)] Let Σ.E E RPxP be symmetric, with eigenvalues A1 2 ... 2 App and À₁ 2 ... 2 Sn respectively. Fix 1 S S.5s5 let d 2 s-r+1 1 and let V (or, .., v2) E Rxd and V (0r, By+1..... (2) Rpxd have orthonormal columns satisfying Σo,- Agyz and Edg 2903 for j == r, 1, ,,...,,. Define A min fmax X, S (0, Ar where dn 2 oo and ip+1 2 0o, we have for any unitary invariant norm || I|+.
|| sin @(V, V) S || El|.-
Here @(V, Rd with @(V, V)s, arccos o, for any j [d] and O(V, V)., 0 all i # E [d]. T1 2 02 2 2 da denotes the singular values of TV [sin θ]i defined as sin(O.;). Lemma K.4 AAzuma-Hoeffding Bound). Suppose [Z, JNEN is super-martingale, suppose -a & Zi+1
Z, < then for all n > 0, a > O, we have
P(Z, Zo 2 a) < 2 exp((a²) ((2((a(B)))
Lemma K.5 AAzuma-Hoefliee
Bound, Vector Form, Hayes (2003)). Suppose (ZalneN is a RD valued martingale, suppose [ Z I a, then for all n > 0, a > O, we have
7, Zoll2 2 (1 + a)) < 2 exp(11 a² 22/2m).
In other words, for any 0 < 5 < 1, with probability at least 8, we have that
2e
2e 2n log
|| Zn - Zol2 5a
1
2n log
< 2
69

---

Published as a conference paper at ICLR 2023
Lemma K.6 (Discrete Gronwall Inequality, Borkar (2009)). Let {a(())eN be sequence of nonnegative real numbers, (#n)net be sequence of positive real numbers and C, I. > O scalars such that for all n,
x(t) S C
Σ an.(n).
n=
Then for T; E'no On it holds that x(t + 1) S CeLT.
Lemma K.7 (Magnus (1985)). Let A RDD + R be any C¹ symmetric matrix function and x² E RR satisfying X1 (A(x*)) > A2 (A(x))) and 01 be the top eigenvector of A(x"). It holds that V (A((A()))l (vv A(z))1))2ze)
We then present some of the technical lemmas we required to prove Lemma H.5.
Lemma K.8. If0 < <e< b-a
+222 av c(1-bb) 2 2 cb then a > 2 cb S !
+b²
Proof of Lemma K.8. Notice that
a² + 62
a² + 262
2 cb² + ca²
cb² + ca² 2 cb
00
1 cb
2(1 cb)
2 cb
So
VI ob
v1 cb
As b-a 6.
we have
cb
So
+
1 1+
b
a
The above inequality implies 2 jb. As c <
b 4
cb S
1
Lemma K.9. When 0 < b,0 <
b.00c <c< b-a
have
c(a² + b²)
- ca² 1
2 ca cb
cb2
ab² + 2 cb - ca) - (1 - cb)
(a
+ (²)
2 - ca - cb
2
(a² + 122)
2 cb
Proof of Lemma K.9. Equivalently, we are going to prove
2 ab
ca cb
(1-cb)b² 2 ca cb 2 cb
+ 2 1a
+ +b)
2 d2(2 cb (0)
2 - - cb
(a² + b2)
Further simplifying, we only need prove
(1 - cb)eab?
a - cb
1
a
> ca +
(2 ca- cb)
22 cb) ca eb)
2 ca cb
2(a² + P2
We have the following auxiliary inequalities.
cb)b > a
ab
2²
>
2 ca
cb
a+b b-a
a+
b-a
a + 62
a + 62
[1 (1-ob)b
cb
00
cb
2
a
+ 1/2
70

---

Published as a conference paper at ICLR 2023
Using the above auxiliary inequalities we have
(1 - cb)cab
+ a² 1 cb
1
at
(2 ca cb)
(2 (2 ca -bb
- e0
cb
a² 2(a² + 112)
ca²b
+ 1 1 (2 ca
db)
1 ca
cb)
(2 - (2 - ca - cb)
- ca eb
ca²b
ca' + 6) cb)
2 00'
(2 - (2 ea - cb)
2(2 00 - cb)
b
ca cb)
1
+
2 ca
(2 -2-cb) (2 - ca cb)
2(2 - ca cb)
3
1
1 1 cb
(2 - cb)² +
2 d)
=3 (2 db) + 6 2 2(2 ob)²
- (cb)² ob + +20
Lemma K.10. When 0 < a < <<<
h-a 5
uv 2(1 r[b- 2 2- =2-ca-ch
a24262
cb'
we have
cb co cb ca) (1- - ca a
2
bb²
1
ct²
2
62
cb
Proof of Lemma K.10. Equivalently, are going prove,
cb² + 2 (2 cb 00) - 122
a²
+
cb
b²
2 ca) < cb
b²
a
cb + a² (1 (1
cb
212
We have the auxiliary inequality
1 > 1/ + cb 4
c-c
Hence
2
6²
cb + a1 cb-
c0) <
3
2 - cb
25
2
6²
a
cb3
tob3 2 (1
ob
(00) ca) <
2
21
313 ba?
12
4
a
a
+
2
21
1 Case 1, If 36t
4 baz
3 < 0, then
313
2
8
1
ha 3
0))
5
+
d²
3
2
2b2
2 Case 2, If 36a
baz
3
2 3
0, then
36
2
62
4
ba?
2
3
5
4
3
2
212
b
0
353
bd
3
=
2
3
<
(6² 2)22
12
4
3
212
313
2
2
ba
< (b a)(b + a)²
13
4
206 ba (b a) + a)² <
2
3
13
40
+(b - ()(((( + +b) - (a + b)2) <
+
2
3
6² 4
(b- (bb 2 3
Using Lemma K.8,a > ( (b(b-))(b+o) (62 - (b a)
< 6b
a)
71

---

Published as a conference paper at ICLR 2023
h-a Lemma K.11. When 0 S a S 6,00 1 c
0F 6F,
12 2
aV
2[1-cb)
a24.222
2-c0-o8' we huve
2 ca) 2ca (02 (62+ 2
cb?
cb + ca
a1) - cb) v 2 - cb
3
Proof of Lemma K.11. Define
2
1 nn² cbb)
F(a) _ (2 - cb =(a)
-20/ (B2 +
b- b² 2 01
a² 262
a² + b²
S(c,b) & {a|0- s a b.0 <c
6.0
a
2
6²
2(1 cb)
ca
Aminn (c,b) - inf S(c. b)
amaz (c,b) - sup S(c, b) - b - cb²
Consider
dF(a) da
2
2
- cb
2a(2 cb ca)
oa²
-2 [62
=a²)(1 cb) - a²
3
2
112 n
d F(a) da²
2
4
4
1- cb
1- cb
a
2(2 cb
c0)
3
C0
0
V
8 + a
2a
18 + 0
2(6 + u
v1 cb
cb
2 4 2eb 4ea 30
62 5
Define u - cb. v 2 b. then u v 1.11 1.
(()
1
2
2u 4uc 3/1 "
da
V2
2 4 2u 4u(1 - 11) - -311 u
2
(1 u)
2 4u 6u - 3/1 u
2
As
(1-w)² 2
+
1211
[1-1)2+1
u.we have
2
(a)
2
4u²
6u + 4 3(1 - u) = 4u² + 33 >
da²
The above inequality shows that F(a) is convex w.r.t to for Omin (c. b) S S amen (c, b). Hence F(a) < max (F (c, b)), F F(amar (c. b))). Below we use amins amar as shorthands for amin (c. b) 0mman (c, b).
For F(amin), ) we have amin
min 26 2(1-bb)
amin +b²
This implies
2 m-n cb
1
(a + b²)
1 [ camin cb + a ( ammn + +0²) (a + b2)
2 2 Camin - cb
20min V (b2 2 mmmn (1 cb) = =(1 cb),
Hence using Lemma K.9,
2
c(a +b²
1 - min
+ b2) 2 - C0min cb
F(amin)) 2min (2 cb
camin) )= (1 cb) Camin cb
co2 co
3
2a2 + 62)
1 cb²
< C 2 cb
c
72

---

Published as a conference paper at ICLR 2023
For F(amaz). know that amaz must satisfy at least of the following three equalities and we discuss three cases one by one.
1. amar
donn +26
oony +ba
comax -cb in this case we simply redo the calculation in Part 1.
2(1-cb)
2. b² amaz
o² +20
This implies
c11-b))
1 (02 + 2d am cb) = =(1- -cb]b² amaal (amax + b²)
1 b2
20
Hence using Lemma K.10,
2 (l0ma) - (1 - eb)ct? - ca @ 1 + b²)
1
F(uman) umax (2 cb
3
82
1 cb²
db²)
cb
3. cb² = b- - amax. Define e - b cb = - v. Note that 1 - cb = Emer and b² 2 amaa
blo
moer +2b2
2aman
These imply amaz + 2amaz
2b < 0 + amaz < 0.9b. This implies v S 0.9. By Lemma K.8,
0.5 S
As U E (0.5, 0.9). it holds that
82 w(1 11+00 u) + 1 S 2 (1+
2 )o.
This implies
22 - (1 u)
? v)u) - 201 (1+ )u S
y2
-
3
2
1+0
2 - cb
Finally,
2
F(amax) = aman (2 cb
c0 20max // (B2
{maa)) (1 cbb
3
2
02 =b² 022 - (1 - v)
2 - (1 - -))b)- 201 /
(1+ +
02
2 )0
$ b²6
1
1).
3
2 - cb
In conclusion, holds that,
F(a) S max (F (amin (c, b)), F (aman b))
cb
< - 2 - cb
cb²).
L OMITTED PROOFS ON CONTINUOUS APPROXIMATION
In this section we give general approximation result (Theorem L..) between a continuous-time flow (Equa- tion 39) and a discrete-time (stochastic) iterates (Equation 40) in some compact subset of R denoted by K. This result is used multiple times in our analysis for full-batch SAM and 1-SAM. Let b K - R is Ci-lipschitz function, that is, Var, x' E K, it holds that [lb(xr) - b(z')ll I C1 ||x x' ll.. Let by be mappings
from K R D for [M] satisfying that b(x) 1 ΣM, (x() for all x E K.
We consider the continuous-time flow X [0, T + K, which is the unique solution of dx (T) ((
(39)
and the discrete-time iterate {a(1))ee which approximately satisfy
x(t - 1) ~ xtt) + pbu, (z()),,
(40)
'Though we believe this approximation result is folklore, we cannot find a reference under the exact setting ours. For completeness, we provide quick proof in this section.
73

---

Published as a conference paper at ICLR 2023
where K2 is independently sampled from uniform distribution over [M] for each N and sx() is deter- ministic function of .a.....,k--. We use Fr to denote the o-algebra generated by ko...., k- and F, to denote the filtration (Fe)seN. Thus x(t) is adapted filtration F,,. Note b undefined outside K, thus in the analysis we only consider the process stopped immediately leaving K, that is, x (t)) & x(min(t, tx ),, where 1x & [t' E N x<(() & K). <xt)) in for all f. 2 O, then l oc. It easy to verify that ix is stopping time with respect to the filtration F... For convenience, we denote X (x(T) = X (min(T, ptppg) as the stopped continuous counterpart of xx
Theorem L.1. Suppose there exist constants C2,6,6 satisfying that
1.
lbbx()) < C2.for for any E K and k E [M]; [x) b(x))11 I Cs. for any x E K and le E [M];
2.
3. be, (z(t)) =(t+1)-=r(t m(t) I c. for all
P
Then for any integer
T/p and
8 1, with probability at least 1- holds that
K |I*(()) (t) xx (pt)l Hp.
max OStST/p
2cT where Hp [|x(0) (()) + C,CTTp+ 2C3 v pf log
eT.
bp
Proof of Theorem I.I. Summing up Equation 39 and Equation 40, for any 1 l.. we have that
( (0)) = / pt
((((T))
(41)
and that
z(t) - x00) = Σ x(t' + x(t')
(42)
P==0
Denote |1=((1) - (m(()l2 by Er, we have that for S lk.
Fre - Fo
pl
6(X(T))dr
E(x(t'
(z(t' + - -x(tt)
P==
1-1
1-1
b(X(T)) - p Σ b(x(pt'))
Σ 6(X - Σ b(x(tt)
p
(4)
(B)
+
((x(t))
byp (z(t'))
+
Σ
by
(x(t))
Σ
(xlt' + +1) x=(z)))
(43)
p'=
20
(D)
Below we will proceed by bounding the four terms (A), (B), (C) and (D) in Equation 43.
1. Note that for any 0 S < T. we have that
( = -10" ((X(s)dds
f (b(((s))l -)
Thus, by Cr-lipschitzness of 6,
(A) = | pr ((x(F) - ((((L=/l))) $ pr
f -
SC s C2pTT.
2. By definition of E, and C₁ Ci-lipschiteness of b, we have that (B) < Cip Σ Ey.
74

---

Published as a conference paper at ICLR 2023
3. We claim that for any 0 < < 1, have that for probability at least 6, it holds that
2eT
(C) - -=220 2pI log
(44)
óp
Below we prove our claim. We denote p 20= ) b(x((()) -u-q >min(t,tx) bby (x((t)) by St, which is martingale with respect filtration F., since ik a stopping time. Note
|| S - kEEK)
max
[lb(x) bb(x))l22 - C3.
by Azuma-Hoefdding inequality (vector form, Lemma K.5), it holds that for any 0 S t S T/p and 1, with probability least 8
2e ... 1/pp 1, we conclude that with
SA.l2 - 2Cary 21 log 8
Applying union bound on the above inequality over = .....,
probability at least 1 8, (C) < 2C 2T /p log 2nT 203 v 2Tp log 2-T
4. We have that
z(t' 1) x(t)
(D) S p Σ bK
(x(t') (x,(t
S pte ET'
Combining the above upper bounds for (A), (B), (C) and (D), we conclude that for any 0 < t S min(T/p.tx))
E, S Hp.6 + CiP Σ Ev.
(45)
Applying the discrete gronwall inequality (Lemma K.6) on Equation 45, we have that
Ex Hp.s Cipt S Cp.
which completes the proof.
Corollary L.2. If minog dist (is) K) Hp. then with probability at least 1 5, tk [T/p) and therefore
|x((1) - ((x))l S H Hi 0SAST/P
max
Proof of Corollary L.2. By Theorem L.1, we know with probability at least 6, have that
| (1) - x (pt)]l s C.T.
max OS/ST/P
Therefore (x (1), R.D K) 2 d (pl), R D K)- dist (X (pl), x (1)) > 0 for any ( S T/p. which implies x ex(t) $ R K, or equivalently, x² (t) K. Thus we conclude that tx 2 [T/p]. 0
Corollary L.3, Suppose M- = and there exist constants Cz, c > satisfying that
I. [b(z)]] for any x E K;
2. b(x) =(t++))=(() P
S c, for all K.
Then for any le E N such that kp S T, it holds that
II.KK (1) (1)-xx(xu)) (pt)) < H.T
max
0StST/P
where H, [|x(0) - X(O)Ila + C.CTT+
Therefore, similar to Corollary L.2, if minosrss dist dist)))) K) > H1eCIT, then it holds that tg > [Tpp) and that
||x(t) - X(pt)l < CIT
max
0S1ST/P
75

---

Full-batch SAM, Phase (Lemma 1.4) Published as a conference paper at ICLR 2023
Proof of Corollary L.3. For any (0, 1}, choosing C3 0 and by Theorem I..1, have that
I
max 0S
1 x (t)
(pt)) - H
2 1 1-6.
Since å can be any number in (0, 1]. the above probability is exactly 1.
We end this section with a summary of applications of Theorem L.I and corollary L.3 in our proofs (Table 2).
Setting
by
P
Full-batch SAM, Phase II (Theorem I.3)
πp² -0. (
+ 1]
1-SAM, Phase (Lemma J.3)
VL+(+)) -
1-SAM, Phase (Theorem J.2)
np?
p + "j
Table 2: Summary of applications of Theorem L.1 and corollary L.3 in our analysis
76


Reviews:
[{"review_id": "duqBzDgliry", "reviewer": "Reviewer_ReRk", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Despite the success of SAM in many applications, the existing theory for SAM is still insufficient to explain its impressive empirical performance. To bridge the gap, this paper rigorously proposes three types of sharpness: worst-direction, ascent-direction, and average-direction, which correspond to the notions of sharpness used in the objective, the actual SAM implementation, and the term used in the generalization analysis of SAM, respectively. They demonstrate that these three types of sharpness lead to different explicit biases in the deterministic case. Moreover, they show that the full-batch SAM minimizes the worst-direction sharpness though it adopts the inexact computationally efficient variants in realization, owing to the implicit alignment between the gradient and the maximal eigenvalue of Hessian. Besides, stochastic SAM with batch size 1 minimizes the average-direction sharpness, which directly benefits the generalization performance. ", "summary_of_the_review": "This paper is novel and solid, which solves an open question in the theoretical area and provides a new aspect in sharpness and generalization analysis.", "strength_and_weaknesses": "**Strengths:**\n\n* The paper is very well-written. I enjoy reading this paper. Instead of stating the theoretical stuff alone, they elaborate various notions and theorems with the high-level idea throughout the paper. \n\n* Three notions of sharpness, their explicit biases, and correlations to SAM discovered in this paper are interesting and could stimulate future research in this area.  \n\n* The toolboxes used in the analysis are interesting, which borrow ideas from the recent literature on edge of stability in SGD. \n\n**Weakness:**\n\n* It is not very clear why assuming the minimizers of loss form a manifold following (Fehrman et al. 2020; Li et al. 2021; Arora et al. 2022) is important. What is the limitation and advantage of this assumption? Please clarify. \n* The proof techniques largely follow (Arora et al. 2022). So it would be helpful to clarify the new challenges and how they are being solved in this paper.  \n* In Section 5.1, it is not clear why the Hessian of each individual loss is rank-1 although I am aware of the Hessian of the total loss is rank-M. Does the individual loss need to be independent?\n* This paper is mainly theoretical sound, but it can benefit from more empirical validations beyond the toy example. For example, adding a real simulation to verify different convergence properties of deterministic and stochastic SAM will be useful. If not, adding references that can support the theoretical results is also helpful. \n\nMinor issues:\n* Sharpness in stochastic setting seems not well defined. \n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "WknvuyTZSVK", "reviewer": "Reviewer_G4Gz", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper studies the properties of sharpness aware minimization through a second order Taylor approximation, in the setting where the perturbation radius tends to zero. Under many strong assumptions, some sort of equivalence (in the limit) between different notions of sharpness-aware-minimization and the spectra of the Hessian of the loss function is derived. The stochastic setting is also studied where two common notions of SAM are shown to be equivalent. The main ideas of the proofs follow the arguments in the particular case of quadratic functions, which are then extended to the general case by means of a second order Taylor approximation. Some of the results are in fact presented only for the quadratic case.\n\n**after rebuttal**\nAuthors have made an effort to address some of my concerns. Although I do not feel all of them have been addressed in full, I am not opposed to it being accepted in current form. I can increase my score to \"marginally above acceptance threshold\". The main drawback of this paper is that its hard to read and the appendix is so long and full of equations that **it's unrealistic that anyone will review it in full for a conference cycle**. Hence, my score should be seen as reflecting the fact that the main 9 pages present a coherent story with (potentially) significant results, but **trusting their validity is left to the authors and potential readers**. As a result, I will also decrease my confidence.", "summary_of_the_review": "The paper requires a good effort to improve readability and be more accessible to a wider audience. The appendix was not carefully written which makes it hard to review. Some of the results might be already understood, but there is no mention of the works that already establish some relation between sharpness and the spectra of the Hessian.", "strength_and_weaknesses": "The main weakness is that the notation and results are presented with really complicated math, hurting the readability. The results are basically a study on sharpness aware minimization on quadratic functions (in this setting the results are more clear and readable), and the added complexity comes from extrapolating such results to the general case by using a second-order Taylor expansion. I will list some further issues:\n\n1. Definition of $U$ in page 3 should be made clear as it appears in other statements. This is problematic as theorem 4.3 which is supposed to be important, is no longer self-contained which hurts readability and clarity\n2. The definition of the \"Limiting regularizer\", as it contains a limit when $\\rho \\to 0$, simply has the goal of making the second order Taylor approximation to have a vanishing error. This facilitates the analysis and makes the function behave like a quadratic function $L(x)=x^\\top A x$. Indeed all the arguments become simple and clear if one just studies this family of functions. Most of the results presented like Theorems 4.5, 4.6 and 4.7 follow easily in this case. The arguments dealing with the approximation error become highly technical and confusing. This hurts readability. I would say the authors could expose the results more clearly if they just explained the arguments for the case of quadratic functions, and leave all the technicalities of the approximation to the appendix.\n3. Another trick the authors use for simplifiyng the analysis is assuming that the solution of the regularized objective is also a solution of the unregularized objective. However this is never clearly stated, and is probably not true except in the limit when $\\rho \\to 0$ (in which case there is no sharpness-aware minimization). This leads to confusing claims like the ones in the paragraph after equation (5): \"The first order term vanishes when we try to minimize the *regularized objective* because every minimizer of the *original loss* has zero gradient\". Again, this is potentially false as there is no reason to believe that a solution of the regularized objective is also a solution of the unregularized objective.\n4. Also in the paragraph following (5) the authors claim \"both the first order terms are the same since $\\sup_{\\|v\\| \\leq 1} v^\\top \\nabla L(x) = \\| \\nabla L(x) \\|_2$\", however this is not true as the supremum in equation (4) contains a second order term of the form $v^\\top \\nabla L(x) v$, hence the supremum is not necessarily the supremum of the first term only. This claim is misleading.\n5. Issues with notation: In page 4 a regularizer is defined as a positive function as implied by the notation $R_\\rho: \\mathbb{R}^D \\to \\mathbb{R}^+ \\cup \\{\\infty\\}$, however immediately after this they discuss regularizers that might be negative? the notation should be fixed.\n6. In definition 4.2 it should be noted that $S$ can take the value $\\infty$ in some cases. Furthermore there is a typo $\\epsilon > 0, \\epsilon > 0$.\n7. The appendix is rather careless with a lot of notation that is not introduced with care. For example already at the first page of the appendix (page 13) we have in Lemma B.5 some letter $C$ that then does not appear at all in the body of the lemma. After a lot of thinking I assume the authors meant $K$ instead of $C$. But the problem does not stop there. Item 3 in lemma B.5 states that $L$ is $\\mu$-PL in some set $K^r$ where $K$ is some compact set. This is clearly false in general. I had to do my best effort to track the issue and indeed $K$ is not an arbitrary compact set $K \\subseteq \\Gamma$ but has a particular structure that is defined in the reference Arora et al. 2022, but which is not reproduced here. Indeed I had to go to the appendix in that reference to decipher that $K$ here corresponds to $Y$ in the reference, which is defined as \"the trajectory of limiting flow Equation (4)\", none of which is clarified in the appendix of this work. I believe the authors simply copied and pasted without care. Hence, the beginning of the appendix is already unreadable and hard to review.\n8. Lemma B.6 has the same issue where there is some letter $C$ that actually refers to $K$ (I believe).\n9. To me, it is not clear how the second order Taylor approximation can be used for $L^\\text{asc}$ as in equation (5): Note that this is only possible in a neighborhood of points where $\\nabla L (x) \\neq 0$, as this term appears as a denominator. However the authors then start arguing that they study points where $\\nabla L(x)$ vanishes? I don't think this is really rigorous as equiation (5) assumes that the taylor expansion has an error that is uniformly bounded by the term $\\mathcal{O}(\\rho^3)$ but in fact this term can depend on $x$.\n10. Assumption 4.1 holds trivially as soon as the regularizer is continous, as any continuous function attains a minimum over a compact set. $R^\\max$ is continuous. Hence it is trivial that assumption 4.1 holds and it is not necessary to go over the complicated argument in Step 1 in the proof of theorem 4.5 (appendix page 17). The same goes for theorem 4.7. \n11. There are some missing citations and some of the results like theorem 4.5 and 4.7 might be well-known, see the clarity/quality/novelty/reproducibility section for details.\n\nAnother weakness is that there is no clear impact from the results presented. Even though they are interesting, what practical conclusions can we draw from this? Having a better understanding of SAM could be achieved by studying the case of quadratic functions, but the heavy math to extrapolate the results to the general case through Taylor approximation, at least in my case, does not improve my understanding of what SAM does further. In any case the results hold in the regime where $\\rho \\to 0$ which means when the regularization is almost negligible. I think this will be of interest only in a small niche in the community. I would say the authors could achieve more impact and more interest if they would rewrite the results with simpler notation and just exposing the core ideas in the context of quadratic functions, while relegating the notation-heavy stuff to the appendix in case someone wants to delve further.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "0pIhc5oa7c", "reviewer": "Reviewer_pjDp", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Sharpness-aware training has been widely adopted for its generalization performance.\nThis paper provides the theoretical analysis for sharpness-aware minimization (SAM).\nThis paper categorizes the SAM in three ways, worst-direction, ascent-direction, and average direction.\nBesides, they provide explicit bias for each way.\nThe paper is clear and theoretically sound.", "summary_of_the_review": "I carefully read the paper and the main statement of the theorem.", "strength_and_weaknesses": "Strength\n1. This paper categorizes the SAM in three ways, and they explain each method in detail.\n2. The paper claims the novel findings and their explicit bias.\n3. There are no experimental results for the real-world dataset, but they provide an intuitive toy experiment to support the claim.\n\nWeakness and Questions\n1. There is no experiment for the real-world datset.\n2. This paper raises a new perspective to analyze the SAM. It would be more interesting if the paper provided a new approach from a new perspective.\n3. This paper claims that batch size=1 SAM corresponds to the average-direction sharpness. In the SAM paper, they introduce the m-sharpness, and m=1 shows a strong performance improvement compared to the full-batch SAM.\nI wonder about the author's opinion about the average-direction sharpness that performs better than worst-direction sharpness in real-world experiments. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster