Explicit Eigenvalue Regularization Improves Sharpness-Aware Minimization
Authors: Haocheng Luo
Tuan Truong
Tung Pham
Mehrtash Harandi
Dinh Phung
Trung Le
Venue: NeurIPS 2024 poster (2024)
Keywords: Optimization
Sharpness-Aware Minimization
stochastic differential equation
Primary area: optimization_for_deep_networks

Abstract: Sharpness-Aware Minimization (SAM) has attracted significant attention for its effectiveness in improving generalization across various tasks. However, its underlying principles remain poorly understood. In this work, we analyze SAM’s training dynamics using the maximum eigenvalue of the Hessian as a measure of sharpness and propose a third-order stochastic differential equation (SDE), which reveals that the dynamics are driven by a complex mixture of second- and third-order terms. We show that alignment between the perturbation vector and the top eigenvector is crucial for SAM’s effectiveness in regularizing sharpness, but find that this alignment is often inadequate in practice, which limits SAM's efficiency. Building on these insights, we introduce Eigen-SAM, an algorithm that explicitly aims to regularize the top Hessian eigenvalue by aligning the perturbation vector with the leading eigenvector. We validate the effectiveness of our theory and the practical advantages of our proposed approach through comprehensive experiments. Code is available at https://github.com/RitianLuo/EigenSAM.

Explicit Eigenvalue Regularization Improves Sharpness-Aware Minimization
Haocheng Luo', Tuan Truong", Tung Pham", Mehrtash Harandi', Dinh Phung', Trung Le
Monash University, Australia VinAl Research, Vietnam
haocheng. luo@monash. edu, ttantm22@vvnaai...... io, v. tunggh4@vinai. 1o, mehrtash. harandi @monash edu, dinh. phungonhnomon.s edu, trunglm@monash. edu
Abstract
Sharpness-Aware Minimization (SAM) has attracted significant attention for its effectiveness in improving generalization across various tasks. However, its un- derlying principles remain poorly understood. In this work, we analyze SAM's training dynamics using the maximum eigenvalue of the Hessian as a measure of sharpness and propose third-order stochastic differential equation (SDE), which reveals that the dynamics are driven by a complex mixture of second- and third- order terms. We show that alignment between the perturbation vector and the top eigenvector is crucial for SAM's effectiveness in regularizing sharpness, but find that this alignment is often inadequate in practice, which limits SAM's efficiency. Building on these insights, we introduce Figen-SAM, an algorithm that explicitly aims to regularize the top Hessian eigenvalue by aligning the perturbation vector with the leading eigenvector. We validate the effectiveness of our theory and the practical advantages of our proposed approach through comprehensive experiments. Code available at https://githuub.cm. wwtpi/uu..ooomtttiaiaiii hat.://bb..oomtttaaaoooooggaaaaa
Introduction
Understanding the generalization of deep learning algorithms is one of the core challenges in modern machine learning. Overparameterizztion makes the loss landscape of neural networks highly non- convex, often featuring numerous global optima, while simple gradient-based algorithms surprisingly tend to find solutions that generalize well. body of empirical and theoretical work suggests that the "flatness" or "sharpness" of the minima is promising explanation for generalization (Hochreiter and Schmidhuberi 1997] Keskar et al.) 2016 Dinh et al.) (011 Jiang et al. 2019 Xie et al. 2020 Liu [ and the implicit bias of optimization algorithms drives them toward flatter solutions,
thereby ensuring good generalization (Blanc et al. (2020) Wen et al. 2022 Arora et al.
2022 Damian
al. 2022 Ahn et 2023; Tahmasebi al. (024)..
Inspired by research on flatness and generalization, recent work by Foret et al. (2021 proposed Sharpness-Aware Minimization (SAM), dual optimization method that perturbs parameters before performing gradient descent enhance generalization performance by minimizing sharpness. Al- though SAM has demonstrated empirical success across various fields (Foret et al.) 2021 Kaddour et al. 2022), theoretical analysis the principles underlying its success remains limited. The work by Compagnoni al. (2023) explains SAM's generalization advantage
implicit minimization of the gradient norm, while Wen et al. (2022) suggests that SAM regularizes the Hessian spectrum near the minima manifold. However, existing theories either somewhat simplified or rely on overly
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

idealized assumptions, leading to noticeable gap between theory and practice. This gap limits their ability fully explain the advantages of SAM (see Appendix [A] or contemporaneous work Song et al.) (2024) for empirical evidence).
In this paper, we consider a widely used measure of sharpness: the largest eigenvalue of the Hessian matrix (Lyu et al. 2022 Arora et al. 2022; Wen et al. [022 Damian et al. (022). We extend the previous Bayes theory to demonstrate the importance of this measure for generalization. Our main contribution is an in-depth analysis of the training dynamics of SAM, expanding on the second-order Stochastic Differential Equation (SDE) proposed by Compagnoni al. (2023) and revealing that the complex third-order terms play crucial role in shaping SAM's ' implicit bias. Under ideal conditions, where the perturbation vector aligns well with the top eigenvector of the Hessian matrix, these terms effectively reduce the sharpness of the loss function. However, our experiments show that this alignment does not hold in real-world settings, limiting SAM's ability to effectively regularize sharpness. Based on our theoretical findings and experimental observations, we propose new algorithm, Bigen-SAM, which intermittently estimates the top eigenvector of the Hessian matrix and incorporates its component orthogonal to the gradient into the perturbation, enabling explicit regularization the top Hessian eigenvalue.
We summarize our contributions as follows:
We prove a new theorem (Theorem 3.1 to establish the relationship between the top eigenvalue and generalization error, building on the general PAC-Bayes theorem AAlquerr ct al. 2016).
We propose a third-order SDE (Theorem44 4.1 to model the dynamics of SAM. This approach achieves lower approximation error compared to the previous second-order SDE by Compagnoni et al. (2023) and additionally infers a close relationship between perturbation- eigenvector alignment and sharpness reduction (Corollary 4.1.I).
theoretical insights and ex- perimental observations (Section 5). This method aims to enhance alignment between the perturbation vector and the top eigenvector, resulting in more effective reduction of sharpness.
We introduce novel algorithm, Eigen-SAM, based
We validate our theory and the effectiveness of the proposed algorithm through comprehen- sive experiments (Section 6).
2 Related work
Theoretical understanding of SAM. SAM has garnered widespread attention for its significant improvements in generalization performance; however, the theoretical analysis provided in the original paper (Foret et al.] (022) limited. The authors only presented PAC-Bayes generalization bound, which is effective only for 0-1 loss. Subsequently, Andrrusshhnuuuncneecee and Flammarion (2022) attempted to further understand the success of SAM by restricting the network structure to diagonal linear networks and providing an implicit bias for SAM. They also established the first convergence result for SAM. Bartlett al. (202) conducted a detailed study of SAM's dynamics for quadratic loss, suggesting that it oscillates between the two sides of the minimum in the direction of greatest curvature and drifts toward flatter minima. However, the assumption of quadratic or locally quadratic loss settings is not realistic for practical deep learning models. More recently, Wen et al. (2022) extended the analysis of SAM's dynamics to general loss functions, assuming that all global minima form connected manifold. Given sufficient training time and infinitesimally small vy and p. they rigorously proved that SAM's dynamics would track the trajectory of a Ricmannian flow with respect to sharpness, achieving the same sharpness-rdducing effect. On a different front, Compagnoni et al. (2023) applied the continuous- time approximation framework from Li et al. (2017) to analyze SAM dynamics, concluding that SAM implicitly minimizes the norm of the gradient scaled by p.
Continuous-time approximations for discrete algorithms. Substantial research demonstrates that the trajectory of stochastic discrete iterations with decaying step sizes will ultimately follow the solution of certain Ordinary Differential Equations (ODEs) (Harold et al., 1997 Borkar et al. 2009) Duchi and Ruan 2018). Further developments in understanding deep learning algorithms were made Ly al. (2017). 1, who proposed general and rigorous mathematical framework for continuous-time approximations, deriving SDEs for SGD and its various variants. They also provided experimental

---

evidence supporting the reasonableness of continuous- time approximations in real-world models (Li al. 202T). In this paper, follow this mathematical framework, reusing some of its notations and definitions.
3 Preliminaries
3.1 Notations
We start by introducing the notation used throughout our paper. Let S denote the training set, sampled from the true data distribution D. For given mini-batch 7, define the mini-bate loss as (x) with parameters empirical loss is defined as fs (x) = B [f (z)]. Since we primarily analyze and discuss the empirical loss in this paper, drop the dependency S for simplicity, denoting the empirical loss as f() when no ambiguity arises.
RR The generalization loss is defined as fD (z) E,~D [f, (z)). while the
We use || to denote the Euclidean norm. The k-th order derivative of the loss f at x is denoted by v f(z), which is symmetric k-tensor in (Rd) ( when E RR We denote by ( (v))(() the largest eigenvalue of the Hessian matrix f2 and v1 f(x)) its corresponding unit eigenvector, v f())(, v) IRd to represent the application of the symmetric 3-tensor 8 f(x) along directions and v.
with || (V2 (x=))l 1. Additionally,
3.2 Sharpness-Aware Minimization
SAM (Foret et al. 2021) seeks flat minima by minimizing the perturbed loss:
min max f(x + pc),
# ||el|
where p predefined hyperparameter controlling the radius of the perturbation. Solving the inner maximization problem leads to SAM (x) [V to x, we get:
vf(x) Differentiating the perturbed loss with respect
V Vf( + pe SAM (x)) d(x + pe SAM (x)) V V
da
(z) (z (AM
In the last approximation, Foret al. 2202 ignore the dependency of SAM (I) on z, leading to faster computational efficiency and higher generalization performance. Applying SAM in the stochastic case, the SAM iteration for mini-batch 7 is summarized as:
Vnn
Tx+1 = F3 V/ f-e xx +
S
We use Vz and V to distinguish between those needing and not needing backpropagation of gradients, matching the algorithm in practice. Thus, the perturbation vector of mini-batch SAM can be written as:
Vfr
SAM
||
3.3
SDE approximation for SGD and SAM
Li
2017) developed the following SDE to approximate discrete SGD:
dX, = (XX + <n("' (Xx)) dW,,
where W, is standard Brownian motion. Compagnoni al. (2023) applied this framework analyze the dynamics of SAM, deriving the following second-order SDE for SAM:
dx, = ()(x)) - P V2 (x,,
)) )at + vT
( ( (X) (x))
dw,
|V
(3)
where Ea.b denotes the covariance matrix of the a-th and b-th terms in the Taylor expansion of the perturbed loss (see Appendix B for the full expression). We refer to Fq the second-order SDE since it includes up to second-order partial derivatives in both the drift and diffusion coefficients.

---

3.4 Choice of sharpness measure
In this paper, we use the largest eigenvalue A1 fV²f(z) as a measure of sharpness, similar to prior works (Lyu et al. 2022 Arora et al., 2022] Wen et al. 2022 Damian et al. (022). Geometrically, the top eigenvalue the Hessian matrix given point represents the maximal curvature of the loss function along any direction. Moreover, it is closely related to the concept of sharpness (defined as the maximum perturbed loss difference) used by Foret et al. (202T) at the minima. Note that the minima, (x())=
0; thus,
0² VS(x) (V)(
max f(x + pE) f(x) H ISI
max De | 11
م (V
Another possible choice the trace the Hessian matrix. However, the trace of the Hessian scales with the dimensionality of the parameters, complicating cross-model comparisons and limiting this measure's applicability.
We further establish a PAC-Bayes theorem that bounds the generalization error through the top eigenvalue of the Hessian matrix. This theorem is based on the general PAC-Bayes theorem Alquier et al. (2016) and applies to bounded losses, not limited to 0-1 loss as in the work of Foret et al. (2021)
Zhuang et al. (2022).
Theorem 3.1. (Generalization Bound) Assume that the loss function is bounded by L, and the third-order partial derivative of the loss function is bounded by C. Additionally, we assume fo(z) S (, fD(x + c), as in Foret ef al. (2021). For any 5 E 10,1) and a >> 0, with probability
over over the choice of "'",, , we have
JD (x) S Js (x) + do²
A1 (v (x)) Cd
2
T.
||-||2
1 + 4 4log (n + d).
dlog 1+
+ +(L) + 2log 8
n
do2
where n is the number of samples.
We defer the proof to Appendix C Theorem 3.1 suggests that minimizing the top eigenvalue of the Hessian matrix is crucial for improving generalization ability.
4 Third-order SDE reveals implicit regularization in SAM
In this section, we delve into the discussion and derivation of the third-order SDE continuous-time approximation for SAM. In Section 4.1 we present heuristic derivations that provide intuitive insights into our approach. Following this, Section 4.2 offers a formal third-order SDE approximation for
SAM, establishing the mathematical rigor our framework. Finally, in Section 4.3 we propose cearbaationeigenveette
corollary linking perturbation eigenvector alignment with eigenvalue regularization, furthering our understanding of the implicit regularization effects inherent in SAM.
4.1 Heuristic derivations for the third-order SDE
We begin by examining the drift coefficient in Compagnoni et al. (2023) (Eq. 3):
(x,) =-V PVE VAA())))
V2
(()))
pE
|V
where the second term indicates that SAM penalizes trajectories with large loss gradients. However, this formulation does not reveal an implicit regularization effect concerning sharpness, specifically regarding the top eigenvalue of the Hessian matrix. Thus, understanding the implicit bias on the Hessian matrix requires third-order Taylor expansion. The missing cubic term in the Taylor expansion is:
V
V
2
(X.)
E
VE 2
||
||

---

The equality holds because V is treated as perturbation vector independent of X,, as SAM's implementation does not involve differentiating with respect to it (See Section 3.2 for detailed description of this part). This third-order term suggests that SAM employs an additional gradient measurement to compute a specific third derivative: the gradient of the second derivative along the direction of the gradient.
Hessian-gradien alignment during training. Recent findings indicate that during training, the gradient implicitly aligns with the top eigenvector of the Hessian matrix under certain conditions: (I) training with normalized full-batch gradient descent (Arora al., 2022]; (2) locally quadratic loss landscape (Bartlett et al. (2023); (3) training with SAM when very close to the minimizer manifold
(Wen et al. 2022). This alignment phenomenon crucial for interpreting the third-order term in our drift coefficient. Specifically, when the gradient is highly aligned with the top eigenvector, we have:
p² E
V
(V) V.+.)
0² EV' fy (V
~
2
|
2
(V
2
where the final equality follows from the properties of differentiating eigenvalues (see Magnus 19855 for a detailed discussion).
If this alignment phenomenon holds, we can conclude that the implicit bias of the drift coefficient aligns with the gradient of the top eigenvalue of the Hessian, thereby implicitly minimizing sharpness.
4.2 Formal third-order SDE approximation for SAM
In this subsection, we present the general formulation of the SDE for SAM. We refer to our SDE (Eq. 4 the third-order SDE, distinguish it from the second-order SDE (Eq. statements and proofs, we refer the reader to Appendix B
3) For the complete
Theorem 4.1. (Third-order SDE for SAM, Informal Statement of Theorem B.4 Let 0 < 1,
0, N= [T/n). and {ax: : k 2 0} denote the sequence of discrete SAM iterations defined by Eq. Define {X, : [), as the stochastic process satisfying the SDE
dX, = - --VjAAM JSAM (Xx)dl + VT(E (x.))adW,
Xo xo
(4)
V + E VIA with JSAM (X,) = f(X.) ME
1V61
+ (X)( + (E + 1 (x)* (x)())))))
ESAM (X,M(X,))= :- = + ((EL)
where Eabb denotes the covariance matrix of the a-th and b-th terms in the Taylor expansion of the perturbed loss (see Appendix for the full expression).
Under sufficient regularity conditions, les (ng)). Then, (X, :q:te E [0, is an order-I weak approximation of (xx k 0}, i.e., for any test function g of at most polynomial growth, there exists a constant C independent of n such that
max k=0,1,....
[Eg(zz) - Ex(Xxn) S Cn..
Vhy Our proof relies on the third-order Taylor expansion of fyx (X, + P PIVJ-y carefully matching the first- and second-ordes -oder conditional moments and quantifying the errors for higher-order terms. Our third-order SDE reveals that SAM's implicit bias includes complex combination of second-order and third-order terms, with scales of and e- respectively. Compared to Compagnoni et al. (0223
our theorem offers two main advantages: first, we allow to take larger values (17 compared 17 in Compagnoni et al. (2023)) ), which is more consistent with real-world settings; equivalently, our SDE achieves lower approximation error for fixed p. Second, our SDE explicitly captures SAM's implicit bias on the Hessian matrix, manifesting as the gradient of the Hessian in the gradient's s direction. Additionally, the diffusion coefficient Eq. Hiimplies that SAM injects additional noise in the form of the covariance of the higher-order terms in the Taylor expansion of the perturbed loss. This curvature-dependent noise aligns with recent studies (Gatmiry et al. 2024a(b). (). which demonstrate that label noise in SGD exhibits similar behavior to SAM.
5

---

4.3 Perturbationn-eieevveooo alignment and eigenvalue regularization
The implicit bias introduced by the third term in the drift coefficient of the SDE (Eq. +. i.e., LE F' ((X.VVV)))))) remains difficult to understand. As discussed in Section 4.1 if we assume IV
that the perturbation vector is well-aligned with the top eigenvector, we can interpret the cubic term as the gradient of the top eigenvalue, leading to an implicit bias that decreases sharpmess. Next, we quantify and formalize this heuristic approach. Define Align(e, 21) == 1 min [ a vi as a measure of alignment between the perturbation vector and the top eigenvector vn It is worth noting that + and - are equivalent eigenvectors, which is why we define alignment as the maximum over both +1 and -11-
Corollary 4.1.1. Recall that E, E.AM is defined in Eq. 2 Let s* denote the direction scalar, i.e..
8' arg (= T positive eigenvalue gap (see Assumption [.. for definition), we have the following:
vill- Under the same conditions as in Theorem 4.1 and assuming
1. If Align(c. AM 21 (V f,( (X))) 2 1 O(p), then the SDE (Eq. 4) becomes
dx, - J (xx)dt + (SAA > >.W,
(5)
where VJS FSAM (X,) == VX(()) + VVEI (X,
(V2
If Align( SAM 11 (V2 (V2)())))) 2 1 O(p2), then the SDE (Eq. 4 becomes
dX, -VAAAM (X,)dt (TESAM
AMY jdd
(6)
where VfSAM = V + pE[s (x))))) (V2 (X,)) E' VRX
The proof is deferred to Appendix B In Corollary 4.1.1 we rigorously formalize our intuition from Section 4.1 If the alignment is at least Opp). we conclude that the SAM trajectory comprises three components: the gradient of the loss, the gradient of the gradient norm, and the gradient of the top eigenvalue, with respective scales 1, p. and e- Under this well-aligned condition, the SAM trajectory minimizes the loss while implicitly regularizing both the gradient norm and sharpness. This demonstrates SAM's complex implicit bias. which is not solely influenced by second- third-order terms, as summarized in previous work (Wen et al. 2022. Compagnoni et al. 2023). For empirical evidence supporting this observation, we refer readers to Appendix
Furthermore, if the alignment is at least 1 - Opp)) ), then the gradient of the gradient norm oscillates in the direction of the top eigenvector. Notably, BBrrlettt et al. (2023) derived similar discrete SAM dynamic under specific conditions (Theorem 20), where the parameter trajectory oscillates in the direction of the top eigenvector while regularizing the leading eigenvalue. However, their conditions are stricter than ours, assuming that the parameters are already close to the minimum and requiring a specifie initialization. When these conditions are met, they require alignment of the gradient with the top cigenvector of at least - O(pp) = O(pt). In comparison, our SDE framework is more general.
Comparison with Wen et al. (2022). Wen et al. (002) derived an implicit bias similar to our cubic term in the third-order SDE (Eq. 5 for SAM through the analysis of the Riemannian flow near the minimizer manifold. However, our work fundamentally differs from theirs. First, their theory requires much longer training time (17 "p" compared to our O1) Thus, our SDE corresponds the initial phase of their analysis regarding time scale, during which they do not conclude any implicit SDE (Eq. 5). which indicates that the implicit bias comprises three components
bias. In contrast, with different scales, provides richer insights in this phase. Second, they require r)ln(1p)) to be sufficiently small, causing p to be exponentially smaller than 7,, whereas our theory accommodates more practical range, p Onz)).
5 Eigen-SAM: an explicit regularization method for the top eigenvalue of the Hessian
5.1 Failure of perturbattoon-iigenvecoo alignment in practice
Within the theoretical framework of Section 4
a natural question arises: the perturbation-
eigenvector alignment sufficient in practice for SAM to effectively minimize sharpness? Unfortunately.

---

Lxpan
SAM
SCD
Alignment
70-
60-
0.5
5000 10000
15000
pa56o
2520s
n0o0p
31000
aonnn
9000
19000
15300
25000
25008
Step
(a) Perturbation-eigeneeete alignment
(b) Top eigenvalue of Hessian
Figure 1: Alignment and top eigenvalue for a 6-layer CNN model trained on CIFAR-10 10. The left panel shows the trend of alignment during SAM training; the shaded area represents the 95% confidence interval. The right panel displays the trend the top eigenvalue over the course training.
we have empirically verified that such alignment may be poor in practice, even for relatively simple models. As result, the regularization effect on the largest eigenvalue, as discussed in Corollary 4.1.1 may not be clearly observable in practical scenarios.
To investigate this phenomenon, we trained a 6-layer SimpleCNN model, as used in Jastrzebski et al. (202T): Deng al. (2024), on CIFAR-10 (Krizhevsky al., 2009). Figure T illustrates key findings from our experiments. The left panel shows that the alignment between the perturbation vector and the top eigenvector is indeed poor during training. Consequently, the right panel reveals that SAM is unable to efficiently minimize the top eigenvalue when alignment is weak. These results highlight the limitations of SAM in real-world scenarios where ideal alignment cannot be assumed.
5.2 Proposed method: Eigen-SAM
To address the issue of poor alignment, we propose novel algorithm called Eigen-SAM, which aims explicitly align the perturbation vector with the top eigenvector. This approach makes SAM's update closer to the SDE approximation in Eq. 5 where the third-order term the drift coefficient can effectively minimize the largest eigenvalue. To achieve this alignment, estimate the top eigenvector of the Hessian matrix once every p mini-batch steps (with 100 in our implementation) using the power method for g iterations (with g 5 in our implementation). This strategy of intermittently estimating the Hessian matrix has been shown to be effective in practice (Liu et al. 2023a).
After obtaining estimate è of the top eigenvector, decompose it into components parallel and perpendicular to the gradient direction, then add the perpendicular component to the perturbation vector to enhance alignment explicitly:
é + TT
VS
Eigen-SAM SAM
+ +a- sign( 1)
(8)
|
where a is a hyperparameter that controls the strength of the explicit alignment. Since + and -v equivalent eigenvectors,
include sign((V f- (x), 6)) to determine the direction of 0. Here, we always choose ê to have a smaller angle with the gradient. The full algorithm is presented in Algorithms [ In Appendix D we provide in-depth discussion of the theoretical properties of Eigen-SA SAM, including sufficient conditions for improving alignment (Proposition) D.I and its convergence rate (Theorem D.2), which comparable that of SAM.

---

Algorithm Power iteration to estimate the top eigenvector
Algorithm 2 Eigen-SAM
1:
for /1 = 1........ do
Initialize i as a random unit vector
2: Compute mini-batch loss (1x)
2: for t2 1,2,..., do
3:
f1 mod then
3: Compute Hessian-vector product
Use Algorithm Ioo estimate
f²r(()
5:
end if
4: Normalize t- 0/llö|1
Compute the perturbation c per Eq.
5: end for
and
6: return
Perturb the model: 3; 3 + pc:
8: Update parameters: 26+1
V (Tx)
9: end for
10: return 2k
5.3 Analysis of additional computational overhead
The additional overhead in Eigen-SAM arises from running the Hessian-vector product q times every p steps to estimate the top eigenvector. The Hessian-vector product requires roughly 1 2 times the time needed to compute the gradient, so the overhead of our algorithm approximately p to 2 + times that of SGD, compared to 2 times for standard SAM.. For larger models, the computation time for the Hessian-vector product remains nearly constant. For a detailed analysis of the computation cost of Hessian-vector products, we refer readers to Dagréou et al. (2024).
2+ 9
Experiments
6.1 Numerical simulation of the third-order SDE
In this subsection, we validate the approximation error between our proposed SDE (Eq. and the discrete SAM algorithm (Eq. D. We trained a fully-connected network with one hidden layer, consisting of 784 hidden units and using GeLU activation, on the MNIST dataset (Deng. 2012p. The training was conducted with 17 = 0.01 and = 0.2 (where R 17 During training, we carefully tracked several key metrics, including training loss, test loss, test accuracy, parameter norm, gradient norm, and the top eigenvalue of the Hessian, as shown in Figure
Our results demonstrate that the approximation error of our third-order SDE is significantly lower than that of the previous second-order SDE. Specifically, the curves of our SDE closely match those of the discrete SAM across all tracked metrics, underscoring the accuracy and reliability of our approximation. This close alignment suggests that our proposed continuous-time approximation provides more precise representation of the discrete SAM dynamics, thus enhancing the theoretical understanding of SAM.
6.2 Image classification from scratch
To evaluate the effectiveness of Eigen-SAM, we applied it to several image classification tasks on benchmark datasets, including CIFAR-10 10 (Krizhevsky et al.] 2209), CIFAR-100 (Krizhevsky et al.) 2009), Fashion-MNIST (Xiao et al. 2017), and SVHN (Netzer et al. 2011). For these tasks, we used R8 He et al. 2016), ResNer-50 (He et al. 2016), and 2edt (Zagoruyko and Komodakis 2016) models.
We selected SGD the base optimizer and applied basic data augmentation techniques, including horizontal flips, padding by four pixels, and random cropping. The batch size was set to 256, with training conducted for 200 epochs. We used initial learning rate of 0.1 for CIFAR- 10, - and CIFAR-100 100, and 0.01 for SVHN, adjusting the learning rate over time with cosine schedule. The weight decay was set to 5 x 10-5, and the momentum 0.9. Detailed hyperparameter settings are provided in Appendix
The test set performance, reported in Table 1 along with the 95% confidence interval, shows that Figen-SAM consistently achieves state-of-the-art performance across various datasets and models, validating its effectiveness and robustness.

---

018-
dicrete SAM
discrate SAM
SDE
order
thirc-onder SDliours
14-
chird order (00)))
Voanngg 100
12
discrete SAM
8.00
Scco no order SDE shind onder Steuurr)
823
200
420
050 Grep
023
1000
400
1000
230
660
1000
Spep
Step
(a) Training Loss
(b) Test Loss
(c) Test Accuracy
14-
15:00
18:35
24
1625
6ão
Step
Steo
Step
(d) Parameter Norm
(c) Gradient Norm
(f) Top Eigenvalue
Figure Training dynamics of discrete SAM, second-order SDE, and third-order SDE during training. Metrics include training loss, test loss, test accuracy, parameter norm, gradient norm, and the top Hessian eigenvalue. Each plot illustrates how each approach affects loss dynamics and key stability metrics.
Table 1: Test accuracy on CIFAR-10, CIFAR-100, Fashion-MNIST. SVHN. Architecture Method
SVHN.
CIFAR-10 CIFAR-100 Fashion-MNIST SVHN
ResNet18
SGD 94.840.2
74.6
94.940.2
96.1_o.1
SAM
95.540.1
77.440.2
95.440.1
96.3_o.1
Eigen-SAM
95.940.2
78.3
95.611.2
96.5
ResNet50
SGD
95.0+0.1
76.640.2
94.8+0.1
96.1_0.1
SAM
95.640.2
79.0+0.2
95.440.1
96.4_o.1
Eigen-SAM
96.24.0.1
79.7
95.7
96.620.1
11deResNee-28-100
SGD
95.710.1
79.8
95.1 10.1
96.2 0.1
SAM
96.51..1
82.0
95.610.1
96.4_ox
Higen-SAM
96.8
82.8 10.1
95.9 15.910.1
10.1
6.3 Finetuning
We evaluated performance by fine-tuning a ViT-B-16 model Dosovitskiy et al. (2020) pre-trained on ImageNet for CIFAR-10 and CIFAR-100. We used the checkpoint provided by Torch'''s h official repository!]
[1] For SAM and Eig, we used an initial learning rate of 0.01 and trained for 4k steps, while for SGD, trained for 8k steps. Table 2 shows the test accuracy, where Eigen-SAM consistently outperforms the baselines.
Table 2: Test accuracy for fine-tuning V6 pretrained ImageNet- IK on CIFAR-10 and CIFAR-100.
Architecture Method
CIFAR-10 CIFAR-100
ViT-B-16
SGD
98.0=0.1
88.6-0.1
SAM
98.41.<0.1
89.5
Higen-SAM
98.5 10.1
89.8 89.8
16.html

---

6.4 Sensitivity analysis
We investigated the impact varying the hyperparameter a on the test accuracy of Eigen- Sigen-SAM. We conducted experiments on ResNet-18 18 with CIFAR-100, 100, testing range of values, shown in 0.2, yielding 0.9% improvement in test
Figure We observed that the test accuracy peaks accuracy compared to 0, which corresponds to the standard SAM. These results suggest that a robust hyperparameter, as its variations do not cause significant performance fluctuations, while consistently enhancing performance.
In Table and Table Appendix F we demonstrate how larger values of p affect generalization performance and observe that setting to 1000 (resulting in less than 1% additional overhead) retains most of the performance gains. In Figure 5 in Appendix F we demonstrate the convergence speed of Algorithm which typically requires only a few steps converge.
Eigen-SAM
SAM
152
0.
accuracy
Demsity
0.3
I
T
27.8
27.6
272
95AM0
0.05
0.7
-2
2
Alpha
Eigenvalue
(a) Sensitivity Analysis of a
(b) Spectrum of Hessian
Figure Left: Sensitivity analysis of at the blue lines indicate the confidence interval. Right:
Spectrum of the Hessian at the end of training.
Hessian spectrum
Figure 3 shows the Hessian spectrum at the 200th epoch for ResNet-18 18 trained on CIFAR-100 using SAM and Eigen-SAM. We observe that the model trained with Eigen-SAM has both a smaller top eigenvalue and trace, with more eigenvalues concentrated near zero. This observation aligns with our motivation for proposing Eigen-SAM and explains why Eigen-SAM generalizes better than SAM.
Conclusion
In this work, we analyzed the training dynamics of SAM using third-order SDE, identifying the alignment between the perturbation vector and the top eigenvector as a crucial factor for effective sharpness regularization. However, our empirical analysis showed that this alignment is often poor in practice. Building on our theoretical framework and experimental insights, we proposed Figen- SAM, algorithm that intermittently estimates the top eigenvector of the Hessian matrix and incorporates its component orthogonal to the gradient into the perturbation, explicitly regularizing the top Hessian eigenvalue. Extensive experiments demonstrated that our third-order SDE yields a smaller approximation error than previous models and that Eigen-SAM SAM achieves state-of-the-art performance across various tasks, validating both its accuracy and effectiveness.
Acknowledgements
We would like to thank Zehang Deng for helpful discussions. We also appreciate the constructive feedback provided by the anonymous reviewers. This work was supported by ARC DP23 grant DP230101 176 and by the Air Force Office of Scientific Research under award number FA2386-23-1- 4044.
10

---

References
Ahn, K., Jadbabaic, A., and Sra, S. (2023). How to escape sharp minima with random perturbations. arXiv preprint arXiv:2305.15659
Alquier, P., Ridgway, J., and Chopin, N. (2016). On the properties of variational approximations of gibbs posteriors. Journal of Machine Learning Research, 17(236).
Andriushchenke, M. and Flammarion, N. (2022). Towards understanding sharpness-aware minimiza- tion. In International Conference on Machine Learning, pages 639-668. PMLR.
Arora, S., Li, Z., and Panigrahi, A. (2022). Understanding gradient descent on the edge of stability in deep learning. In International Conference on Machine Learning, pages 948-1024, PMLR.
Barrett, D. G. and Dherin, B. (2020). Implicit gradient regularization. 116.
arXiv preprint
Bartlett, L., Long, P. M., and Bousquet, O. (2023). The dynamics of sharpness -wwre aware minimization: Bouncing across ravines and drifting towards wide minima. Journal of Machine Learning Research, 24(316):1336.
Blanc, G., Gupta, N., Valiant, G., and Valiant, P. (2020). Implicit regularization for deep neural networks driven by an ornstein-uhlenbeck like process. In Conference on learning theory, pages 483-513. PMLR.
Borkar, V. S., Pinto, J., and Prabhu, T. (2009). A new learning algorithm for optimal stopping. Discrete Event Dynamic Systems, 19:91-113.
Compagnoni, E. M., Biggio, L.,, Orvieto, A., Proske, F. N., Kersting, H., and Lucchi, A. (2023). An sde for modeling sam: Theory and insights. In International Conference on Machine Learning. pages
PMLR.
Dagréou, M., Ablin, P., Vaiter, S., and Moreau, T. (2024). How compute hessian-vector products? In ICLR Blogposts 2024. https://ccrr-looposts..s http:///elllopssgggiiioo/////e/d
Damian, A., Nichani, E., and Lee, J. D. (2022). Self-stabilization: The implicit bias of gradient descent at the edge of stability. arXiv preprint 15555:2209.15594.
Deng, L. (2012). The mnist database handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 29(6):141-142,
Deng, Z.,, Sun, R., Xue, M., Wen, S., Camtepe, S., Nepal, S., and Xiang, Y. (2024). Leakage-resilient and carbon-neutral aggregation featuring the federated ai-enabled critical infrastructure. arXiv preprint arXiv:2405.152.
Dinh, L., Pascanu, R., Bengio, S.. and Bengio, Y. (2017). Sharp minima can generalize for deep nets. In International Conference on Machine Learning, pages 1019-1028. PMLR.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., al. (2020). An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprini arXiv:2010. /0000///9999 1929.
Drucker, H. and Le Cun, Y. (1991). Double backpropagation increasing generalization performance. In IJCNN-91-Seatle International Joint Conference on Neural Networks, volume pages 145-150. IEEE.
Duchi, J. C. and Ruan, F. (2018). Stochastic methods for composite and weakly convex optimization problems. SIAM Journal on Optimization, 28(4):3229-322 3259..
Dziugaite, K. and Roy, D. M. (2017). Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. arXiv preprint arXiv: 1008.
Folland, (2005). Higher-order derivatives and taylor's formula several variables. Preprint, pages 1-4.

---

Foret, P., Kleiner, A., Mobahi, H., and Neyshabur, B. (2021). Sharpness aware minimization for efficiently improving generalization. In 9th International Conference on Learning Representations. ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net.
Gatmiry, K., Li, Z., Ma, T., Reddi, S., Jegelka, S., and Chuang, C.-Y. (2024a). What is the inductive bias of flatness regularization? a study of deep matrix factorization models. Advances in Neural Information Processing Systems, 36.
Gatmiry, K., Li, Z., Ruiz, L., Reddi, J., and Jegelka, (2024b). Simplicity bias of SGD via sharpness minimization.
Harold, Kushner, G., and Yin, G. (1997). Stochastic approximation and recursive algorithm and applications. Application of Mathematics, 35(10).
He, K., Zhang, X., Ren, S., and Sun, J. (2016). Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778.
Hochreiter, S. and Schmidhuber, J. (1997). Flat minima. Neural computation, 9(1):1-42
Jastrzebski, S., Arpit, D., Astrand, O., Kerg, G. B., Wang, H., Xiong, C., Socher, R., Cho, K., and Geras, K. J. (2021). Catastrophic fisher explosion: Early phase fisher matrix impacts generalization. In International Conference on Machine Learning, pages 4782-4784. PMLR.
Jiang. Y., Neyshabur, B., Mobahi, H., Krishnan, D., and Bengio. S. (2019). Fantastic generalization measures and where to find them. arXiv preprint arXiv: 191200278.
Kaddour, J., Liu, L.,, Silva, R., and Kusner, M. J. (2022). When do flat minima optimizers work? Advances in Neural Information Processing Systems, 35: 35:16577-16595.
Keskar, N. S., Mudigere, D., Nocedal, J., Smelyanskiy, M., and Tang, P. T. P. (2016). On large-batch batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv: (0di::6600.0000 0409644436
Krizhevsky, A., Hinton, G., et al. (2009). Learning multiple layers of features from tiny images.
Li, Q. Tai, C., and Weinan, E. (2017). Stochastic modified equations and adaptive stochastic gradient algorithms. In International Conference on Machine Learning, pages 2101-2110. PMLR.
Li, 7.,, Malladi, S., and Arora, S. (2021). On the validity of modeling sgd with stochastic differential equations (sdes). Advances in Neural Information Processing Systems, 34:12712-1272
Liu, H., Li, Z.,, Hall, D., Liang, P., and Ma, T. (2023a). Sophia: scalable stochastic second-order optimizer for language model pre-training. arXiv preprint arXiv:2305.14342.
Liu, H., Xie, S. M., Li, Z., and Ma, T. (2023b). Same pre-training loss, better downstream: Implicit bias matters for language models. In International Conference on Machine Learning, pages PMLR.
Lyu, K., Li, Z., and Arora, S. (2022). Understanding the generalization benefit of normalization layers: Sharpness reduction. Advances in Neural Information Processing Systems, 35:34689-34700
Magnus, R. (1985). On differentiating eigenvalues and eigenvectors. Econometric theory, 1(2):179- 191.
Mil'shtein, G. (1986). Weak approximation of solutions of systems of stochastic differential equations. Theory of Probability & Its Applications, 30(4):750-77) 766.
Netzer, Y., Wang, T., Coates, A., Bissacco, A., Wu, B., Ng. A. Y., et al. (2011). Reading digits in natural images with unsupervised feature learning. NIPS workshop on deep learning and unsupervised feature learning, volume 2011, page 7. Granada, Spain.
Si, D. and Yun, C. (2024). Practical sharpness-aware minimization cannot converge all the way to optima. Advances in Neural Information Processing Systems, 36.
Song, M., Ahn, K., and Yun, C. (2024). Does sgd really happen in tiny subspaces?
12

---

Tahmasebi, B., Soleymani. A., Bahri, D., Jegelka, S., and Jaillet, P. (2024). A universal class of sharpness-aware minimization algorithms. arXiv preprint arXiv:2406.0366 03682..
Wen, K., Ma, T., and Li, Z. (2022). How does sharpness-aware minimization minimize sharpness? CoRR, abs/221 1.05729.
Xiao, H., Rasul, K., and Vollgraf, R. (2017). Fashion-mnist: novel image dataset for benchmarking machine learning algorithms. anXiv preprint arXiv:1708.07747
Xie, Z., Sato, I., and Sugiyama, M. (2020). diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. arXiv preprint arXiv:2002.03495.
Zagoruyko, S. and Komodakis, N. (2016). Wide residual networks. arXiv preprint arXiv:1655.0711
Zhuang, J., Gong, B., Yuan, L., Cui, Y., Adam, H., Dvoreek, N., Tatikonda, S., Duncan, J., and Liu, T. (2022). Surrogate gap minimization improves sharpness-aware training. arXiv preprint arXiv:2203.080%
13

---

Appendix
A Additional empirical evidence
In this section, we present empirical evidence highlighting discrepancies between existing theories and practical observations. We demonstrate this by designing several "counterexample" algorithms and comparing their performance with SAM.
First, the theory developed by Wen et al. (022) suggests that SAM's implicit regularization primarily occurs near the minima manifold, where (z) R and the quadratic term dominates the perturbed term. To test this, we designed counterexample algorithm called Reverse-SAM, which applies SAM's ascent step using negative normalized gradient, i.c., E (f)(z) (x) If this theory holds in practice, then Reverse-SAl -SMM should perform similarly to SAM, as the sign of the perturbation vector would not affect the implicit bias induced by the quadratic term.
The second counterexample algorithm we consider is Explicit Gradient Regularization (EGR), regularization method with a long history, tracing back to the works of Barrett and Dherin (2020) and Drucker and Cun (199T). The second-order SDE proposed by Compagnoni al. (2023) suggests that SAM's implicit bias is equivalent to optimizing the objective f(x) + f(x) Therefore, if this second-order SDE model accurate in practice, EGR with a regularization coefficient equal to should exhibit performance comparable to SAM.
We trained ResNet-18 18 on CIFAR-10 10 to compare the performance of these two counterexample algorithms with SAM. Figure[4ssove the training and test losses. Additionally, the test accuracies for SAM, EGR, and Reverse-SAM were 95.5% = 0.1%, 95.0% = 0.2%, and 94. 94.4% = 0.2%, respectively, Both counterexample algorithms show noticeable performance gaps from SAM across all metrics, highlighting the limitations of existing theories in explaining SAM's practical outcomes.
Reverse-SAM
SAM
1.79
SAM
Reverse-SMM
EGR
EGG
1.58
1.20
0.5
0.58
0.2
(a) Training Loss
(b) Test Loss
Figure 4: Comparison of training loss and test loss metrics across algorithms.
B Proofs for SDEs (Section 4)
We follow the notation of multi-index, which is commonly used in SDE literature:
A multi-index is an n-tuple of non-negative integers a = (ar, Q2...., On))
||| := 01 02 +... On
a! a lag!..a..
For (01, F2..., Fnn RR _a 20 .1 x.22 ..... an
For multi-index 3,
a f(x) := dz 0z Du f(x)
a18l
We denote the partial derivative with respect x; by dag.
Following the SDE framework by Mil' shl'is 1986); et al. (017)) Compagnoni al. (2023 we
have the following definitions and assumptions:
14

---

Definition B.1. Let G denote the set of continuous functions IR" - R of at most polynomial growth, i.e. g E G if there exists positive integers K1 Kg > 0 such that
|g(x)| S K1 + ||0||
for all x E Rd. Moreover, for each integer a 2 we denote by Go the set of a-limes continuously differentiable funcfions TRd + R which, together with its partial derivatives up and including order o, belong to G.
This definition originates from the field of numerical analysis of SDEs shtein 1986). In the case of g(x) = ||>||", the bound restricts the difference between the j-th moments of the discrete process and those of the continuous process. We write 0) to denote that there exists a function K G independent with p. 7, such that the enror terms are bounded by K 1/
Assumption B.1. Assume that the following conditions f. fs and their gradients are satisfied:
f. fi satisfy Lipschitz condition: there exists L > 0 such that
IVf(=)-Vffi)) + Σ [V ff(ii)l S L|| a|1i
f. fs and its partial derivatives up to order 7 belong to G;
f. V fi satisfy a growth condition: there exists M > 0 Ouch that
[V(() +
[f((a)] & M(1 + ||>]]);;
g and its partial derivatives up order belong to G;
See Li et al. (2017) for detailed discussion Assumption B.1
To prove the Corollary 4.1.I we need an additional assumption to ensure that the top eigenvalue is
differentiable. Note that this assumption is common in recent works (Damian et al. 2022 Wen et al. 2022) and easy to satisfy.
Assumption B.2. (Eigenvalues gap) For all k > O, (V(((((a)) X1 (V2 f(Xr)()
Ae ( and
In the proof of our main theorem, we will utilize the following two auxiliary lemmas.
Lemma B.1. (Lemma ei 2017) Let 0 < 11 < 1. Consider stochastic process {X, >>>0>00 0} satisfying the SDE
dx, = b(X,) + + (X.)dW.
with X, = x E IRd and b, a together with their derivatives belong to G. Define the one-step difference
A = X,, then we have
EA; = bin + 2
Σ b byde,by by
n² +000
1......,,
(9)
EA,A, = [b.b, + tooT
Vi,j = 1.....d. .d: n² +0
(10)
( = 0 E
V 3, = 1....dd
(11)
1=1
all the functions are evaluated Lemma B.2. (Theorem 2 and Lemma 5 (Mil'shtein 1986) Wnder Assumption B.1 we define the one-
step difference for the stochastic process A X,, in addition there exists K K2, K3, KĄ E G
15

---

so that
EA; EA, & Ki vi=1,...,
(12)
E E K2(x)n Vi.j = 1...,,....,d;
(13)
E
A -E E II A <Ks()) Vs Vi, {l....,d};
(14)
E
II
A
$ Ka(x)) Wig (I.....d).
(15)
j=1
Then, for each E G, there exists constant C' such that
max
[Bg(xx) (X(()))) S On
0=
Next, we will prove lemma that controls the moments of the discrete process for SAM:
Lemma B.3. Under Assumption B.I ler <
1 1. We define:
Σ & te, fx(
jk O
de, jSAM
= (sA)
(
E
[
2
(v)
In addition, we define the one-step difference for the discrete-time algorithm as A Then we have:
1. EÅ; = - -c, FSAM (x)m + ( Vi vi=
(16)
2.
EA,A; dc. FSAM (z)de, JSaM (x)n²² + (6) ,2 ()'p') = d: (17)
de,
SAM
3. E A1
V2 (L......, d].
(18)
j=1
All functions are evaluated at x.
Proof. To evaluate EÅ;
we start by analyzing fr (a+ [للل Vf((z)). which is the partial derivative in the direction e,, Taking the Taylor
-E (en (x+ VI
Vi(())]
expansion, we have:
0xf [+)[
de.f x [ff(()] vh(())
=x)
Σ
82 (z)
|a|=1
1 Σ
0ala(z)) a (1VF2(()))
|a|=2
((()) + R
(
V[((())
(19)
where the residual term defined in Folland (2005), For some constant c e (0, 1), it holds that
V1)(w) اlaم [.)())) 04 aetat (s (p]l
(x) [V]]
R
P
Σ
(20)
Now, we observe that
04 Httaf
1 cp
Vi(la) (P]V)*)))]]
p/ام V/(aa)
(x)
a!
16

---

is a finite sum of products of functions that, by assumption, are in G. We can rewrite Equation 19 as
(a)(())
De.fa x [ff(())] V(((a)) = (x)
+
Σ
8h (z))
(x)
Oal((a)
a
+ Σ
I3 seltasa(a)p
2 1a|=2
[[
+ p (2)x,,
(21)
which implies that
E (a+ [x)l V4((a) = -0ef(()) + +B
Σg 82 te,e, te,te, (1()
(x)
1²
Ei
(1)))
+
2
[x)
+ p² R((x))
(22)
where K((x)= Ex,(z)). Therefore, we have Vi 1,2, .d.
EÅ, = -o, -.. JSAM 1
To prove the second statement, by definition, we have
A,) =f ( (()))))))))))))))))))) [13
+ p +  (2]](0)) + +
(23)
= = (2) (z)+0(1ppp)
(24)
, = =EĀEA, + Cov(A, As)
(25)
= G. JSaM (z)0.,, jSAM (x))? + 1 (x)
(26)
Finally, it is clear that
E
A, = 0 Vs i (l....,d) d}
(27)
j=1
= (1²).
V*23, e
(28)
Now we are ready to state the theorem and prove
Theorem B.4. (Stochastic modified equations) Under Assumption B.I let 0 < <,<,,>> N,N [T/n]. Let Iy E IRd, 0 < k < N denote the sequence of SAM iterations defined by Equation Define {Xe) as the stochastic process satisfying the SDE
SAM dx, = - JSAM (X,)dt + STM (X.))
(29)
where (X;A((x)) == f(x,) + PE
VI. 14NDV E
IV
ESAM [3,00_))))---,),--00"""",)),,,,3 (2 (E2: 1 (1)")
17

---

(x,( = E [V)(x)) - (X(() - -VXX) V(((I))
(X,) =E (V -AA--AA 24400) (B
(I
(X
|4
|
V2 [XX)V _
x2, :=
IVAI
( VAJ
E (X,( = (V((x.)-V((x))) E
|V
|
Additionally, let us take
0m)
Then, {X, :t [], is order-1 weak approximation of {æ k 2 0}, i.e. for each g E G, there exists a constant C independent of such that
[B((x)) Ex(xx,)) S On
O.1......NN
Proof. We will check the conditions in Lemma B.2 As we apply Lemma B.1 choices:
we make the following
b(x) V jSAM (x) oxx) = E (x)
First, for Vi = -I,... d, we have
EA;
Lemurd 8.
de. SSAM (x)n + O (np²)
(30)
EA
Lemond
F7
0e. JSAM (x)n (np²)
(31)
Therefore, we have that for some K Ki(x) EG
E. RA {1()))
(32)
Second, for Vi,j = 1, ..., d, it holds that
E
Lemme B
De, ESAM (z) ESAM (x) ESAM 12 7 +
(33)
E
De. FSAM (x)8., FSAM (SAM()))) + SAAM 2 00
(34)
Consequently, we have for some K2(4) G
EA,A, EA,A < (2(x))²
(35)
Third, for Vig- = 1,.., d. we have
E
Lemma 1 (173)
(36)
E
A
Lemmo 1
O (1²2
(37)
1=1
Therefore, for some Kz(x) E G. we have
E
A
I1
A,
S Ka(z))²
V323 3
(38)
1=1
18

---

Additionally, for some (xxx) E G, Vij 1....... d,
Levamd
E
A
(
(39)
By combining the four equations, Eq.
Eq.
35
Eq.
38
Eq.
39
we complete the proof.
Proof of Corollary 4.1.1 1. Under the supposition, the first statement is immediately followed by
using the fact:
(
e² ". - O -0(0)) (40)
|V
2. Under the supposition, the second statement is immediately followed by using the fact:
(X,
L (7 - (41)
LE
2
f
2
and
(1) V
PE
P [1
O(p²).
(42)
|
V
(x)
PE
-p
A ( 1 [v))()))] O(p²)
(43)
C Proof of Theorem 3.1
Theorem C.1. (Generalization Bound) Assume that the loss function is bounded by L, and the
third-order derivative of the loss function is bounded by C. Additionally, we assume fo(x) S o,, (x(( x + e), similar to Foret et al. 2022). For any 8 E (0,1) and > 0, with probability over 8 over the choice of D we have
d
(V (z)) Cd.
to Ss
2
L
(1+
||=1|2 + 0((1) + 2log 8 + log (n + d).
dlog
2 n/
do2
Proof. We use the PAC-Bayes theory in this proof. In PAC-Bayes theory, follows distribution, denoted by P. and we express the expected loss over x as follows:
fo(P) E [z)
fs(P) Ex~ [x)(
NO00 TZT))) and Q = NVe,, Ta) over x E IRd, where is the prior distribution and Q is the posterior distribution, the general PAC-Bayes theorem in Alquier et al. (2016), for all B > 0, with a probability at least 1 - 5, we have
For any distribution
1 fo( < fs( [L(Q]|P) + log g + (B.n))
1
(44)
where is defined as
(3, = log En. exp {0[ -
When the loss function is bounded by L, then by using the Hoeffding's 's inequality we have:
B²
(83 <
8n
19

---

The task is to minimize the second term of RHS of (44), we thus choose 3 = (an((LiQQ))))toox)) Then the second term of RHS of (44) is equal to
L
KL(QI + log á L.
2n
The KL divergence between Q and P. when they are Gaussian, is given by formula
1
do² + ||||² d dlog of 02 o
KL(QII
2
For given posterior distribution Q with fixed a², to minimize the KL term, the oZ. should be equal to a² + ||x]|//d In this case, the KL term is no less than
|0
d dog 11+
do
Thus, the second term of RHS is
KL(QIIP) + log
x L
d dlgg
+
||=||
do
x LL
2n
4n
when ||0:| > o²{ exp(4n/d) - 1}. Hence, for any | > o&{ (epp(4n/d) }]. we have the RHS greater than the LHS, the inequality trivial. In the remainder of the proof, we only consider the case:
| < o² (exx(44n/d) - 1).
(45)
Distribution P is Gaussian centered around O with variance of o² + ||=+|//d, which is unknown the time we set up the inequality, since æ unknown. Meanwhile we have to specify advance, since P is the prior distribution. To deal with this problem, we apply the union bound technique (Dziugaite and Roy 2017 Foret et al., (22T)) We set
o² exx((4n/d))
Pj = ((0.ccxp ( 11 La)
(P) ; 1.
Then the following inequality holds for particular distribution Pj with probability 1 6, with
6, 7 7
68
[T( S (s( + KL(Q1I))) + log 5, + (33
Use the well-known equation: holds with every j. We pick
Σラ fs
a² 6 then with probability 6, the above incquality
² + |,x1///d
o + |x/1//d
1 = d log
1 dlog 0² + exp(4n/dd)
Therefore.
0² + 1//dd
1- jr
dlog
0² + ||0||///1 /d
--
0² + || 1
log
$ $log
d
d
(1-5)
0² || dd cexp
s exxp(1/d) + 11=11//d|
0² + 111x1///55 or S Sexp(1/0) (02 + ||=11//4.
20

---

Thus the KL term could be bounded as follow
1
d + |0:1² d + dloog
KL(QIIP3-)
2
T.--
x²
1
d(o² + |//11//d)
(xp +
+|/0|///0
$ 2
0² ||x12/d d dlog
o²
1
dlog exp(1/d)(o² [1ol1(//d]]
2
a²
1 + dlogg (1 (1+ |1:12
)
2
do²
For the term log 8,- with recall that c (1 + exp(4n/d)) and
1 - dlog 2² (1xppp 4n/d))
0 +|1=||
we have
log (
n
log
log
+ log
+ 2 log(j")
85
66
x²
0² exp(4n/d))
log
log
+ 2lo
1 dlog
6
0² + 1|01///
π²
- log
log
2 log
d (1 + exp(4n/d))
2
[ + [++d(1= d(1 + )
S
log
log
+ 2 log
6
π²
S log
log
log(1 + d + 4n).
6
Hence, the inequality
KL(Q)IP,.)
log 1 85*
Jo (Q) S 1s (Q)
x L
2n
$ fss
||:11
n²2
+ 2/m V
+ dlog (1+
+ 2 log
+ 1lgg(n + dd)
do
65
$ fs (0)
||0||
+ 2/72 V dllog (1 +
+ 0(1) + 2oo + 4loogn + d).
do2
[fo x +c) s (EnN(o,
fs x + +c]
+ L
|0/1/
dlog (1 do? ) + 0(11 2log + +loog( + d).
2/n
Using the assumption that fo(x) s E,-N(0,021)) [f0 x + +e)]. we reach
fo (x) S Ban [s (x +c)+
L
|
1 5 4log(n + d).
+ 2/ni V dlog (1
+ 0(1) + 2 log
do2
21

---

Using the second-orde order Taylor expansion for fs + c), we obtain
+c) fs(x))=x=)>)) - e Vafs(()) ( (z) Σ
8 (or + 1e)
Ess Cis
dz,, Uzi, dris
<fs(z)+ + c Vo.s(()) + E (V' (z)) | +
Σ
8¹ fs (x te)
Ess 6ig Eigs
Dzi Ou Owiss
$1.$2.33
where E [0, 1]. Thanks to o0 [||2]] S Be-N(o. [|]]]] = do we have
+ do²
(e-N(0,02t,)) [fs x (=((=))) s fs fz(
A1 (V fs (x))
2
Cd
E
B(,,(0,022) ||c 1] E 122 [|< BegrNN(a. [|ss]]
+ da²
\1 (V ss [s))) + (EeINN(o.09) [[]] 1/2
Cd
Js
x
2
6
do
Cd
= Is
1
+
\1 ( fs (x) + 6
2
By the assumption fo(x) s (1.2)10001 + <), we reach
do²
X1 (V (z) + Cd
(x)x) S fs 2) +
2
6
|
+ 2 2/7 V dlog (1+
+ 0(1) + 2 log 3 + 4 4log(n + d).
do2
D Theoretical Properties of Eigen-SAM
In this subsection, first we show that Eq. 8 can indeed improve the alignment for moderate Let = cos( (m) [+]] v), without loss of generalization, we suppose w 0. We only consider the case
vI > 2 since in the case < w S vI 2 any a > 0 will enhance the alignment. Using fundamental
mathematics to solve the inequality, we obtain
Proposition D.I. Let L2 for any (0,
2wv1 20
we have cos(
Vf(2) (x] avi,v) v) w.
Proposition D.I shows our update can indeed improve the alignment for a wide range of a. For example, if cos( V (o,(V) (x), v) 0.8, then can be any value in (0, 3.43; if cos( (os(V)(())) v) 0.9, then a can be any value in (0, 1.27).
Next, inspired by Si and Yun 2024). we have the following convergence rate for stochastic Eigen- SAM on non-convex function:
Theorem D.2. (Convergence rate) Consider a B-smooth f-smmothh f satisfying f = =inff f(x) > -00,
ler A = f(o)) - f*, and assume the mini-batch variance is bounded by o². Under Eigen-SAM,
starting at ao with any perturbation size p > 0 and step size n = min 23
VI. minimize
VDo2T
we have
Σ EIVV (a.)) so
AA T
+
V
B VT
2 + +
22

---

Proof. By the definition of B-smoothness, we have
_nz 2
Ef((+1) (/ +
B V pE)]
- < - (/
+ + V
I [B((z)) V/ Bn (ziV + pe)l1 pe)1 + 0²)
Ef(x)) V/ p=)l
pe)l1 + "EIV Vf(2 + Bm [ pe)ll + o²)
S <Bf(x)) 7 (xx;)) 8
EZ, (x+ pe)l1 + BB²11
g2m(p² + a²)
= BS(x)) - "E]VJ(<>))'*
2 + B0"11
Rearranging this inequality and deviding both sides by 2 we have
T T
(Bf(zo) - -f((T)) + (2 + a²) + 2300
7/
2A
s nT + B²2 + a2)) 2
By using 17 min
V V
}, complete the proof.
E Additional Experimental Details
For hyperparameter p. follow the guidelines by Forct et al. 2021), setting p 0.05 for CIFAR-10 and
0,01 for SVHN, and 0.1 for CIFAR-I00, We ensure that SAM and Eigen-SAM use the same p for fair comparison. Additionally, we tune the hyperparameter for Eigen-SAM over (0.05, 0.1, 0.2) using 10% of the training set as a validation set. We find that a = 0.2 works the best for almost all cases. Therefore, to demonstrate that Eigen-SAM does not require extensive hyperparameter tuning. We run three independent repeat experiments with different weight initializations and data shuffling. Because SAM and Eigen-SAM require twice the runtime, we allow SGD to train for twice the number of epochs. All our experiments were conducted NVIDIA RTX 4090 24GB GPUs.
report the performance with a 0.2 for all experiments
F Additional experiment results
Table 3: Test accuracy on CIFAR-10 for different values of (interval steps for estimating eigenvec- tors) using Eigen-SAM. Architechture 1=100
p=200
p=500
p=1000
SAM
ResNet18
95.910.2
95.910.1
95.710.2
95.81o.1
95.5
WideResNet-28-100 10
96.8
96.7.0.1
96.8
96.710.1
96.5
23

---

Table 4: Test accuracy on CIFAR-100 for different values p (interval steps for estimating eigenvec-
tors) using Eigen-SAN SAM.
Architechture
p=100
p=200
p=500
p=
SAM
ResNet18
78.3
78.2
78.220.2
78.3
27.4
WideResNet-28-11 10 82.840.1 82.7=0.2
82.7
82.640.1
82.0
Allhmmntt
Algnmmnt
6.4
(a) ResNet18
(b) WideResNet-28-10
Figure 5: The effect of the number of Hessian-vector product steps in Algorithm (power iteration) on the alignment of the estimated vector with the top eigenvalue. The dataset is CIFAR-100, and the models ResNet18 and WideResNet-28-10 at mid-training stage (100th epoch).
G Limitaions
limitation of this work the additional memory and time required estimate the top eigenvalue of the Hessian matrix. Improving the efficiency of Eigen- SAM is direction for future research.
24

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification:
We accurately assess and describe our contributions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss our limitation at the end of the appendix.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
25

---

Justification:
We provide all assumptions and our proof is complete.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We provide all details for experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
26

---

Answer: [Yes]
Justification:
Yes, our code is open source.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https guwdes/dde/eeee Cooeeeoodeesssseeeeeee 7/nips.com htte://allllliiiiee////eeeuaassl ssiniiiiccyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: We specify them the main paper. Additional details are provided in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification:
We give a confidence level of 95%.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
27

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide the information of the computational resource.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
Our research conduct with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our paper focuses on fundamental research in deep learning theory and algorithms.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
28

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We only use public datasets.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We correctly cite authors for models and data.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
29

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We communicate the details of the dataset/code/mooel..
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
We do not contain human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We do not contain human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
30


Reviews:
[{"review_id": "U57dhK8jKr", "reviewer": "Reviewer_eUoc", "summary": "### Summary:\n\n\n\nThe authors first provide a theoretical connection between max e.v. of Hessian and generalization via PAC-Bayesian bounds. Then, they propose a 3rd-order SDE for SAM, which has a lower approximation error compared to existing 2nd-order SDEs in the literature. They argue that the perturbation-eigenvector alignment is important, so they propose the Eigen-SAM algorithm which estimates max e.v. and improves alignment using power method. The paper concludes with experiments showing improvement over SAM and SGD. \n\n### Main Results: \n- Theorem 4.1: relating max e.v. of Hessian to generalization via PAC-Bayesian bounds.\n- Theorem 4.2: a 3rd order SDE for SAM, that is used later to study top e.v. alignment.", "questions": "### Questions/Comments:\n\n\nThis is an interesting paper with improvements over the previous works on understanding SAM dynamics and implicit bias. I'm still not well convinced about the need for 3rd order data for SAM analysis. The reuslts are interesting but the paper is lacking a clear introduction/motivation for this study. Also, the provided Eigen-SAM algorithm looks extensive since it relies on Hessian-vector products. The limitations should be discussed clearer in the early pages of the paper. \n\n\nMy comments are as follows:\n\n- I recommend rewritting Theorem 4.1 with exactly parsing the role of each term and also redefining some key quantities there. For example, what is $N$? I believe it is the number of samples (based on previous PAC-Bayesian bounds) but nothing explained there. \n\n- It's worth mentioning that Theorem 4.1 does not work for over parameterized models (as expected). Becasue to make the upper bound smaller one needs to have a very large number of samples. \n\n\n\n\n#### Other comments:\n\n\n\n\n- line 79 -- I suggest using $\\|$ instead of $||$ to denote norms.\n\n- line 81 -- is $\\nabla^3 f(x)(u,v)$ a vector? Is it a multilinear function? How does it relate to the symmetric tensor $\\nabla^3 f(x)$? This part of notations is a bit unclear.\n\n\n- line 84 -- It might be useful to emphasize that you use the 2-norm of vectors to define SAM.\n\n- line 85 -- typo \"aslo\"\n\n- line 92 -- the notation here is a bit weird, though I can totally understand what you mean that. I don't have any suggestion right now but it can be confusing for unfamiliar readers. \n\n- line 156 -- \"SDE 5\" and \"SDE 3\" are using confusing notation.\n\n- line 160 -- typo $x_\\in$\n\n- line 163 -- Is the expectation over the randomness of $X_t$? Then, it just adds some constants to $f(X_t)$ (i.e., independent of $X_t$) and then will vanish in Equation (5)? I think $X_t$ must be replaced with $x$ and then in Equation (5) $x$ is replaced with $X_t$?\n\n- line 163 -- why $\\Sigma^{2,2}$ does not depend on $x$? Is this a typo?\n\n- line 167 -- I suggest replacing $g \\in G$ with just a few words explaining it. This makes that part clearler. \n\n- line 195 -- it looks like to be typo, what does the middle term mean in the formula?\n\n- Experiments in Figure 1 look noisy -- how many times did you iterate it? \n\n- line 290 -- typo\n\n- How did Algorithm 1 estimate the top eigenvector? Can you explain this with a simple example? I mean, the power method in words/quick examples. \n\n- Missing references: I suggest discussing the following two related papers in the future version of the paper. One of them is related to random perturbations, decomposition, and trace regularization, and the other one is related to SAM algorithms with general sharpness measures:\n        - A Universal Class of Sharpness-Aware Minimization Algorithms, ICML 2024\n        - How to escape sharp minima with random perturbations, ICML 2024", "limitations": "N/A", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "### Pros:\n\n- the paper is super easy to read\n- notations are well defined, the problem is clearly formulated, and the results are explicitly stated \n\n\n- provides a deeper understanding of SAM convergence and implicit bias", "weaknesses": "### Cons:\n\n- the need for 3rd order SDEs for SAM are not well justified\n- the current version can be improved in terms of clarity/typos", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough and insightful feedback, which has been invaluable to the improvement of our work. Below we address your comments:\n\n* **line 81 -- is $\\nabla^3f(x)(u,v)$ a vector? Is it a multilinear function? How does it relate to the symmetric tensor $\\nabla^3f(x)$?** Yes $\\nabla^3f(x)(u,v)$ is a vector with the same size with $x$, it can be written as $(\\nabla^3f(x)(u,v))_k=u^T(\\nabla^3f(x))_kv$, it is bilinear with respect to the variables $u,v$.\n\n* **Suggestions regarding Theorem 4.1** We appreciate your suggestions regarding the presentation of Theorem 4.1. We will rephrase it in future versions to make it more rigorous and readable. You mentioned that Theorem 4.1 does not work for over parameterized models. It is correct. Because we need to bound the KL divergence between the prior and the posterior, dependence on the number of parameters in the bound is unavoidable if we do not impose additional constraints on the prior and posterior. The reason we did not choose to add more constraints is that we wanted to remain consistent with other PAC-Bayes type bounds related to SAM (e.g. [1], [2], [3]).\n\n* **line 163 -- Is the expectation over the randomness of $X_t$? Then, it just adds some constants to $f(X_t)$ and then will vanish in Equation(5)? I think $X_t$ must be replaced with $x$ and then in Equation (5) is replaced with $X_t$?** No, we mentioned in Section 3.1 when introducing the notation that unless otherwise specified, our expectations are with respect to minibatch $\\gamma$. Therefore, $f(X_t)$ is still a function of $X_t$. $X_t$ and $x_k$ are different stochastic processes, the former is the solution of the SDE, while the latter is obtained by running the discrete algorithm.\n\n* **line 163 -- why $\\Sigma^{2,2}$ does not depend on $x$? Is this a typo?** Yes, it is a typo. In fact, it should depend on $X_t$.\n\n* **line 195 -- it looks like to be typo, what does the middle term mean in the formula?** It is not a typo. $s^*$ and $\\lambda_1$ are both scalars. We will rewrite this expression to make it more readable.\n\n* **Experiments in Figure 1 look noisy -- how many times did you iterate it?** We did not run repeated experiments because we only aimed to obtain a qualitative result, and the overall trend is clear.\n\n* **How did Algorithm 1 estimate the top eigenvector? Can you explain this with a simple example? I mean, the power method in words/quick examples.** The power method is an iterative technique used to find the dominant eigenvalue and corresponding eigenvector of a square matrix. It starts with a random unit vector $v$, at every iteration, the vector $v$ is multiplied by the matrix and normalized. The convergence is geometric, with ratio $\\lambda_1/\\lambda_2$. We refer to Figure 1 in our rebuttal PDF, where we demonstrate the convergence speed of this algorithm in practice.\n\n* **Missing references** We will cite the relevant and interesting concurrent work you mentioned in future versions.\n\n* **All typos and misuses of notation** We are very grateful for your careful help in finding them. We deeply apologize for the errors caused by the tight submission deadline, and we will correct all typos and misuses of notation in future versions.\n\nReferences\n\n[1] P. Foret, A. Kleiner, H. Mobahi, and B. Neyshabur, \"Sharpness-Aware Minimization for Efficiently Improving Generalization,\" in *Proc. Int. Conf. Learning Representations (ICLR)*, 2021, doi: 10.48550/arXiv.2010.01412.\n\n[2] Zhuang, J., Gong, B., Yuan, L., Cui, Y., Adam, H., Dvornek, N., ... & Liu, T. (2022). Surrogate gap minimization improves sharpness-aware training. arXiv preprint arXiv:2203.08065.\n\n[3] Li, B., & Giannakis, G. (2024). Enhancing sharpness-aware optimization through variance suppression. Advances in Neural Information Processing Systems, 36.", "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": "", "extra_scores": {}}, {"review_id": "2zGH28VhxY", "reviewer": "Reviewer_4ndr", "summary": "The authors consider sharpness aware minimization problem and derive third-order SDE governing SGD dynamics.\n\nAs a result, they obtained a bound on the generalization error of SAM, showing that SAM trajectories favor flat minimal with the smaller largest eigenvalue of the Hessian, implying that flatter minima can improve generalization.", "questions": "How to practically verify assumptions on the loss? Do the assumptions of the main theorems hold in conducted empirical study?", "limitations": "-", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Third-order approximation by SDE of underlying SGD dynamics and explicit bounds. It is well-written and rigorous.", "weaknesses": "The full list of assumptions on the loss is not listed. Following the proofs, I can see that it should satisfy uniform boundness and the existence of a non-zero spectral gap w.r.t. the largest eigenvalue of the Hessian. I would appreciate it if the authors specified a list of assumptions and examples of losses that satisfy them.\n\nBound (4) does not necessarily imply improved generalization. Reading the text 'crucial' etc., creates an impression as if it is. I would appreciate it if it would be mentioned that the paper proposes a possible explanatory mechanism.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "I appreciate the time and effort you've taken to review our manuscript. Below are our responses to your concerns:\n\n* **The full list of assumptions on the loss is not listed.** The assumptions for Theorem 4.1 can be found in Section A.2 of the appendix, including the continuity of the third derivatives, assumptions on the parameter domain, and the boundedness of the loss function $f$ on the parameter domain. The assumptions for Theorem 4.2 and Corollary 4.2.1 can be found in Assumption 1 in Appendix A.1, which include the continuous differentiability of $f$, at most polynomial growth rate, and Lipschitz conditions. We will make these lists of assumptions more clearly visible in future versions.\n\n* **Bound (4) does not necessarily imply improved generalization.** We need to clarify that bound (4) is an upper bound on the generalization error, which is used to link the largest eigenvalue of the Hessian matrix with the generalization error. Even though there is a correlation, we acknowledge that this does not mean a smaller eigenvalue necessarily implies a smaller generalization error. In practice, we rely on our proposed algorithm, Eigen-SAM, to regularize the largest eigenvalue rather than directly optimizing the bound on the generalization error.\n\n* **How to practically verify assumptions on the loss? Do the assumptions of the main theorems hold in conducted empirical study?** Our assumptions are quite mild and do not introduce additional assumptions compared to previous theoretical work. We refer you to some prior empirical studies (e.g. [1], [2]) that demonstrate how these assumptions can be validated in practice, but this is beyond the scope of this paper.\n\nReferences\n\n[1] Fazlyab, M., Robey, A., Hassani, H., Morari, M., & Pappas, G. (2019). Efficient and accurate estimation of lipschitz constants for deep neural networks. Advances in neural information processing systems, 32.\n\n[2] Khromov, G., & Singh, S. P. (2024). Some Fundamental Aspects about Lipschitz Continuity of Neural Networks. In The Twelfth International Conference on Learning Representations.", "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": "", "extra_scores": {}}, {"review_id": "cVCuCXPCDF", "reviewer": "Reviewer_WYKu", "summary": "The paper focuses on improving Sharpness-Aware Minimization (SAM) by introducing a novel approach called Eigen-SAM. It establishes a theoretical connection between the top eigenvalue of the Hessian matrix and generalization error using an extended PAC-Bayes theorem. The authors derive a third-order stochastic differential equation (SDE) to model the dynamics of SAM, revealing lower approximation error compared to second-order SDEs. The paper emphasizes the importance of perturbation-eigenvector alignment in reducing sharpness. Eigen-SAM intermittently estimates the top eigenvector to enhance this alignment, leading to better sharpness minimization. Extensive experiments demonstrate the effectiveness of Eigen-SAM in improving test accuracy and robustness over standard SAM and SGD across multiple datasets and model architectures.", "questions": "1. Does third-order SDE shown in Figure 2 & 3 estimate the top eigenvector on each mini-batch or once every $k$ mini-batch?", "limitations": "The authors have mentioned in Section A.5 that Eigen-SAM requires additional computational time.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper provides a strong theoretical foundation by connecting the top eigenvalue of the Hessian to generalization error and developing a third-order SDE for SAM.\n\n2. The paper introduces Eigen-SAM, which improves upon standard SAM by focusing on perturbation-eigenvector alignment.\n\n3. Experiments across datasets and models validate the theoretical insights and demonstrate the practical benefits of Eigen-SAM.", "weaknesses": "1. The paper claims that it's the first work establishing the relationship between the top eigenvalue and generalization error. However, some works [1,2] have discussed related topics. \n\n2. The generalization bound given in Theorem 4.1, based on the result of the original SAM paper, considers any random perturbation rather than worst-case perturbation as SAM does in practice. Hence, the theoretical conclusion might be less instructive.\n\n3. Empirical results are inadequate. Test on large-scale datasets like ImageNet-1K is common in SAM-related works. That would be better if there were more experiments on other tasks like noisy label task.\n\n4. Eigen-SAM requires extra computational cost for estimating the top eigenvector. While the authors claim that it only needs 5%-10% overhead, is the overhead related to the size of model architectures? Whether the estimation need more frequent if tested on large-scale datasets?\n\n**Minor:**\n\n- Typo in line 249: should be $v$\n- Typo in line 290: should be $\\rho$\n\n[1] Wen, Kaiyue, Tengyu Ma, and Zhiyuan Li. \"How does sharpness-aware minimization minimize sharpness?.\" arXiv preprint arXiv:2211.05729 (2022).\n\n[2] Kaur, Simran, Jeremy Cohen, and Zachary Chase Lipton. \"On the maximum hessian eigenvalue and generalization.\" Proceedings on. PMLR, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your helpful comments and insightful suggestions. We address your questions one by one below:\n\nWeaknesses:\n\n**1. The paper claims that it's the first work establishing the relationship between the top eigenvalue and generalization error. However, some works [1,2] have discussed related topics.** We would like to clarify that, to the best of our knowledge, we are the first to establish theoretical generalization bounds based on the top eigenvalue. The paper [1] you mentioned analyzes the dynamics and implicit bias of the SAM algorithm but does not establish a connection with generalization. The paper [2] you mentioned is purely experimental work, which is significantly different from our theoretical bound.\n\n**2. The generalization bound given in Theorem 4.1, based on the result of the original SAM paper, considers any random perturbation rather than worst-case perturbation as SAM does in practice. Hence, the theoretical conclusion might be less instructive.** Our Theorem 4.1 aims to establish a relationship between the top eigenvalue of the Hessian matrix and generalization. We acknowledge that it may not be practically instructive. In practice, we rely on our proposed algorithm, Eigen-SAM, to regularize the largest eigenvalue rather than directly optimizing the bound on the generalization error.\n\n**3. Empirical results are inadequate. Test on large-scale datasets like ImageNet-1K is common in SAM-related works. That would be better if there were more experiments on other tasks like noisy label task.** Our main goal is to expand the existing theories on understanding sharpness-aware minimization (SAM) and to bridge the substantial gap between current theory and practice, rather than to propose a highly practical, and very scalable optimizer to replace SAM. Therefore, we did not conduct experiments on large datasets, such as ImageNet. Nevertheless, in the rebuttal PDF, we show the results of fine-tuning on CIFAR-10/CIFAR-100 using ViT-B-16 pretrained on ImageNet, showing that Eigen-SAM consistently outperforms original SAM. Additionally, we did not conduct experiments related to label noise because we have not developed a theory concerning label noise.\n\n**4. Eigen-SAM requires extra computational cost for estimating the top eigenvector. While the authors claim that it only needs 5%-10% overhead, is the overhead related to the size of model architectures? Whether the estimation need more frequent if tested on large-scale datasets?** Our additional overhead comes from running the Hessian-vector product p times every k steps to estimate the top eigenvector. The Hessian-vector product roughly takes 1-2 times the time required to compute the gradient, so the overhead of our algorithm is approximately $(2+p/k)$ to $(2+2p/k)$ times that of SGD, compared to 2 times for standard SAM. For larger models, the computation time for the Hessian-vector product hardly increases. We refer to [1] for a more detailed analysis of the computation cost of Hessian-vector products. In Table 2 and Table 3 of our rebuttal PDF, we show how larger values of k affect generalization performance and observe that even setting k to 1000 (resulting in less than 1% additional overhead) can retain most of the performance gains. Regarding the question of whether larger models require more frequent estimation of the top eigenvector, we show in Figure 2 of our rebuttal PDF the rate of decay in the quality of the estimated top eigenvector. For larger models, this decay rate only slightly increases and generally decays relatively slowly.\n\n**5. typos** We are very grateful for your careful help in finding them. We deeply apologize for the errors caused by the tight submission deadline, and we will correct all typos and misuses of notation in future versions.\n\n**Questions: Does third-order SDE shown in Figure 2 & 3 estimate the top eigenvector on each mini-batch or once every $k$\n mini-batch?** In the numerical simulation experiments of the SDE (Figure 2,3), we estimate the eigenvector in each mini-batch.\n\nReferences\n\n[1] Dagréou, et al., \"How to compute Hessian-vector products?\", ICLR Blogposts, 2024.", "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": "", "extra_scores": {}}, {"review_id": "qTkddDMKmX", "reviewer": "Reviewer_zC7c", "summary": "This study establishes a theoretical connection between the top eigenvalue of the Hessian matrix and generalization error through an extended PAC-Bayes theorem. Highlighting the importance of perturbation-eigenvector alignment in mitigating sharpness, this study introduces Eigen-SAM. This method intermittently estimates the top eigenvector to enhance alignment, leading to enhanced sharpness minimization. Experimental findings demonstrate improvements in both test accuracy and robustness when compared with conventional SAM and SGD methods.", "questions": "A. Could you offer theoretical evidence elucidating the poor alignment in Section 5.1?\n\nB. What is the computational cost in comparison to other algorithms? Can Eigen-SAM be feasibly implemented in large-scale models?", "limitations": "Please refer to the Weaknesses.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "This research introduces theoretical advancements aimed at enhancing Sharpness-Aware Minimization (SAM) and proposes an algorithm called Eigen-SAM based on these theoretical foundations. Additionally, supportive experimental evidence is presented in this study.", "weaknesses": "The experimental findings exhibit limited strength, indicating the necessity for additional analysis and further experimentation:\n\nA. The experimental results do not sufficiently demonstrate the superiority of Eigen-SAM over SAM.\n\nB. The experiments were conducted using smaller scale models and datasets, where concerns regarding generalization are relatively minimal.\n\nC. Additionally, a comprehensive analysis of computational costs is required.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed and insightful feedback. We address your concerns below:\n\n* **The experimental findings exhibit limited strength, indicating the necessity for additional analysis and further experimentation: A. The experimental results do not sufficiently demonstrate the superiority of Eigen-SAM over SAM. B. The experiments were conducted using smaller scale models and datasets, where concerns regarding generalization are relatively minimal.** Our main goal is to expand the existing theories on understanding sharpness-aware minimization (SAM) and to bridge the substantial gap between current theory and practice, rather than to propose a highly practical, and very scalable optimizer to replace SAM. Therefore, we did not conduct experiments on large datasets, such as ImageNet. Nevertheless, in Table 1 of the rebuttal PDF, we show the results of fine-tuning on CIFAR-10/CIFAR-100 using ViT-B-16 pretrained on ImageNet, showing that Eigen-SAM consistently outperforms original SAM. \n\n* **C. Additionally, a comprehensive analysis of computational costs is required. Questions B. What is the computational cost in comparison to other algorithms? Can Eigen-SAM be feasibly implemented in large-scale models?** Our additional overhead comes from running the Hessian-vector product p times every k steps to estimate the top eigenvector. The Hessian-vector product roughly takes 1-2 times the time required to compute the gradient, so the overhead of our algorithm is approximately $(2+p/k)$ to $(2+2p/k)$ times that of SGD, compared to 2 times for standard SAM. For larger models, the computation time for the Hessian-vector product hardly increases. We refer to [1] for a more detailed analysis of the computation cost of Hessian-vector products. In Table 2 and Table 3 of our rebuttal PDF, we show how larger values of k affect generalization performance and observe that even setting k to 1000 (resulting in less than 1% additional overhead) can retain most of the performance gains. Regarding the question of whether larger models require more frequent estimation of the top eigenvector, we show in Figure 2 of our rebuttal PDF the rate of decay in the quality of the estimated top eigenvector. For larger models, this decay rate only slightly increases and generally decays relatively slowly.\n\n* **Questions: A. Could you offer theoretical evidence elucidating the poor alignment in Section 5.1?** We are sorry that we cannot provide a theoretical proof for this phenomenon. This precisely highlights the gap between the theory and practice of SAM.\n\nReferences\n\n[1] Dagréou, et al., \"How to compute Hessian-vector products?\", ICLR Blogposts, 2024.", "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": "", "extra_scores": {}}, {"review_id": "zStfD3zBNa", "reviewer": "Reviewer_B3hY", "summary": "This paper proposes a novel method, Eigen-SAM, as an improvement to Sharpness-Aware Minimization (SAM). The paper theoretically elucidates the relationship between the top eigenvalue of the Hessian matrix and generalization error and models the dynamics of SAM using a third-order stochastic differential equation (SDE). Eigen-SAM aims to more effectively minimize sharpness by intermittently estimating the top eigenvector of the Hessian matrix. The effectiveness of Eigen-SAM is validated through experiments across multiple small datasets.", "questions": "- More specific guidelines on the frequency of top eigenvector estimation and parameter selection in Eigen-SAM implementation would be helpful. How are these parameters selected?\n- Is Figure 4(a) showing sensitivity analysis for alpha rather than rho? This is a bit confusing to me.\n- Does the fact that v updates only once every 100 steps in equation 8 imply stability in the inner product of the gradient with the direction of the maximum eigenvalue? Could you provide an intuitive explanation?\n- Are parameters other than alpha being tuned? If 90% of the training set is used as training data and the remaining 10% as validation data for Eigen-SAM, and no validation data is used for other optimizers, is my understanding correct?\n- Wen et al. (2022) interpret SAM learning with full batch training as regularizing the maximum eigenvalue. Could increasing the batch size improve alignment for SAM? Additionally, how does batch size affect alignment, maximum eigenvalue, and test accuracy in Eigen-SAM?\n- For Table 1, can you also show the results of how much the maximum eigenvalue improved?", "limitations": "- The authors acknowledge the increased computational cost of implementing Eigen-SAM, which may limit its feasibility in some applications. Future research should explore methods to reduce this computational cost and develop more efficient implementations. It would also be beneficial to consider cases where the applicability of the proposed method may be limited.\n- A significant limitation is that experiments are conducted only on ResNet models, and the effectiveness on more modern models like ViT is unknown.", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Soundness\n- The technical claims, experimental methods, and research methodology are robust, and the central claims are well-supported by evidence.\n- The necessity of the third-order SDE is well-explained in section 4.\n\nPresentation\n- The presentation is generally clear, providing context by comparing with existing research.\n\nContribution\n- This study contributes to the field by deepening the theoretical understanding of Sharpness-Aware Minimization and proposing a new algorithm that enhances practical performance.\n\nOriginality\n- Novelty by proposing Eigen-SAM as an improvement to SAM, theoretically demonstrating the relationship between the top eigenvalue of the Hessian matrix and generalization error through a third-order SDE.", "weaknesses": "- Experiments are limited to ResNet models, raising questions about the method's effectiveness with other models. Comparative experiments with ViT and other architectures are strongly recommended.\n- The lack of publicly available implementation may pose a barrier to community adoption.\n- Demonstrating ablation studies on parameters such as k (interval for power method) and p (iterations for power method) would make the method more useful for the community.\n- Showing effectiveness in more practical settings like ImageNet would impact my score.\n- Comparisons with other optimization methods using eigenvalues would strengthen this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful comments and constructive suggestions. Below we address your comments:\n\nWeaknesses:\n\n* **Experiments are limited to ResNet models, raising questions about the method's effectiveness with other models. Comparative experiments with ViT and other architectures are strongly recommended.** In Table 1 of the rebuttal PDF, we additionally show the results of fine-tuning on CIFAR-10/CIFAR-100 using ViT-B-16 pretrained on ImageNet, showing that Eigen-SAM consistently outperforms original SAM.\n\n* **The lack of publicly available implementation may pose a barrier to community adoption.** We are committed to supporting the research community and plan to open-source our code upon acceptance of the paper.\n\n* **Demonstrating ablation studies on parameters such as k (interval for power method) and p (iterations for power method) would make the method more useful for the community.** In Table 2 and Table 3 of our rebuttal PDF, we show how larger values of k affect generalization performance and observe that even setting k to 1000 (resulting in less than 1% additional overhead) can retain most of the performance gains. In Figure 1 of the rebuttal PDF, we demonstrate the efficient convergence speed of Algorithm 1, which requires minimal adjustment to p.\n\n* **Showing effectiveness in more practical settings like ImageNet would impact my score.** Due to computational costs and time constraints, we were unable to complete experiments on ImageNet during the rebuttal period, but we are very willing to include it in our future plans.\n\n* **Comparisons with other optimization methods using eigenvalues would strengthen this paper.** To the best of our knowledge, we are the first to use eigenvalue regularization in a SAM variant.\n\nQuestions:\n\n* **More specific guidelines on the frequency of top eigenvector estimation and parameter selection in Eigen-SAM implementation would be helpful. How are these parameters selected?** Among our hyperparameters, alpha is tuned on the validation set. There is no consistent selection principle for k and p. A larger k improves the quality of the estimated top eigenvector during training but incurs higher computational costs, requiring a trade-off between the two. For more experiments on hyperparameter selection, please refer to Tables 2 and 3 as well as Figure 1 in our rebuttal PDF.\n\n* **Is Figure 4(a) showing sensitivity analysis for alpha rather than rho? This is a bit confusing to me.** This is a typo, and we apologize for the confusion it caused. We will correct it in future versions.\n\n* **Does the fact that v updates only once every 100 steps in equation 8 imply stability in the inner product of the gradient with the direction of the maximum eigenvalue? Could you provide an intuitive explanation?** Yes, it means that the decay rate of the alignment between the estimated top eigenvector and the real top eigenvector is relatively slow. We refer to Figure 2 in our rebuttal PDF for empirical evidence of this phenomenon.\n\n* **Are parameters other than alpha being tuned? If 90% of the training set is used as training data and the remaining 10% as validation data for Eigen-SAM, and no validation data is used for other optimizers, is my understanding correct?** For other hyper-parameters, such as learning rate and perturbation radius, we followed previous work and kept them consistent across all optimizers.\n\n* **Wen et al. (2022) interpret SAM learning with full batch training as regularizing the maximum eigenvalue. Could increasing the batch size improve alignment for SAM? Additionally, how does batch size affect alignment, maximum eigenvalue, and test accuracy in Eigen-SAM?** Previous work [1] has empirically shown that increasing the batch size can actually increase the sharpness of the model, thereby harming generalization performance, highlighting the gap between the theory and practice of SAM. Our theory holds for any batch size, unlike [2], which is only applicable to specific batch sizes. Therefore, we do not have theoretical results related to batch size and have not conducted batch size related experiments.\n\n* **For Table 1, can you also show the results of how much the maximum eigenvalue improved?** Yes, here is the table of the top eigenvalue:\n|Model| Method | CIFAR-10 | CIFAR-100 |\n|----------|----------|----------|----------|\n|ResNet-18|SAM | 5.0 | 5.7 |\n||Eigen-SAM | 4.2 | 4.9 |\n|ResNet-50|SAM | 4.2 | 4.3 |\n||Eigen-SAM | 3.8 | 3.7 |\n|WideResNet-28-10|SAM | 1.2 | 1.7 |\n||Eigen-SAM | 1.0 | 1.4 |\n\nReferences\n\n[1] Andriushchenko, M., & Flammarion, N. (2022, June). Towards understanding sharpness-aware minimization. In International Conference on Machine Learning (pp. 639-668). PMLR.\n\n[2] Wen, K., Ma, T., & Li, Z. (2023). How Sharpness-Aware Minimization Minimizes Sharpness?. In The Eleventh International Conference on Learning Representations.", "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": "", "extra_scores": {}}]

Author rebuttal:
We would like to thank all reviewers for their careful and thoughtful feedback. Some concerns from the reviewers are centered on our experimental section, therefore, we would like to provide a general response here. Firstly, we would like to clarify that our paper is more theoretical in nature. Our main goal is to expand the existing theories on understanding sharpness-aware minimization (SAM) and to bridge the substantial gap between current theory and practice, rather than to propose a highly practical, and very scalable optimizer to replace SAM. Nevertheless, we still hope to provide some supporting experiments in our rebuttal PDF to further support our algorithm, including: (1) Fine-tuning on CIFAR-10/CIFAR-100 using ViT-B-16 pretrained on ImageNet, where we have presented the test accuracies in Table 1, showing that Eigen-SAM consistently outperforms original SAM; (2) Conducting a sensitivity analysis on the hyperparameter k (interval steps for estimating eigenvectors), with the experimental results reported in Tables 2 and 3. We found that setting k=1000, while the additional computational cost is less than 1%, retains most of the performance gain; (3) Showing the convergence rate of Algorithm 1 used for estimating the top eigenvalue in Figure 1; (4) Reporting the rate of decay in the quality of the estimated top eigenvectors as training progresses in Figure 2. We hope these experiments will adequately support the effectiveness of our algorithm and address some of the reviewers' concerns. We will respond individually to each reviewer's specific issues in the rebuttals below. Thank you for your patience.

Meta-review: The authors consider sharpness aware minimization problem and derive third-order SDE governing SGD dynamics. All authors agreed on the theoretical contribution in the paper, while the experiments could be improved. Reviewers and authors conducted constructive discussions and I believe that this paper is a useful contribution to deep learning optimization.
Decision: Accept (poster)