Explicit Eigenvalue Regularization Improves Sharpness-Aware Minimization

Authors: haocheng  luo
tuan  truong
tung  pham
mehrtash  harandi
dinh  phung
trung  le
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 14:49:55

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.

Main paper: 
 introduction: 
 Understanding the generalization of deep learning algorithms is one of the core challenges in modern machine learning. Overparameterization makes the loss landscape of neural networks highly nonconvex, often featuring numerous global optima, while simple gradient-based algorithms surprisingly tend to find solutions that generalize well. A body of empirical and theoretical work suggests that the "flatness" or "sharpness" of the minima is a promising explanation for generalization (Hochreiter and Schmidhuber, 1997; Keskar et al., 2016; Dinh et al., 2017; Jiang et al., 2019; Xie et al., 2020; Liu et al., 2023b) , 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 et al., 2022; Ahn et al., 2023; Tahmasebi et al., 2024) .Inspired by research on flatness and generalization, recent work by Foret et al. (2021) proposed Sharpness-Aware Minimization (SAM), a dual optimization method that perturbs parameters before performing gradient descent to enhance generalization performance by minimizing sharpness. Although SAM has demonstrated empirical success across various fields (Foret et al., 2021; Kaddour et al., 2022) , theoretical analysis of the principles underlying its success remains limited. The work by Compagnoni et al. (2023) explains SAM's generalization advantage as an 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 are either somewhat simplified or rely on overly idealized assumptions, leading to a noticeable gap between theory and practice. This gap limits their ability to 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., 2022; Damian et al., 2022) . We extend the previous PAC-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 et al. (2023) and revealing that the complex third-order terms play a 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 a new algorithm, Eigen-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 of 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 (Alquier et al., 2016) .• We propose a third-order SDE (Theorem 4.1) to model the dynamics of SAM. This approach achieves a lower approximation error compared to the previous second-order SDE by Compagnoni et al. (2023) and additionally infers a close relationship between perturbationeigenvector alignment and sharpness reduction (Corollary 4.1.1).• We introduce a novel algorithm, Eigen-SAM, based on our theoretical insights and experimental observations (Section 5). This method aims to enhance alignment between the perturbation vector and the top eigenvector, resulting in a more effective reduction of sharpness.• We validate our theory and the effectiveness of the proposed algorithm through comprehensive experiments (Section 6). 
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., 2021) is limited. The authors only presented a PAC-Bayes generalization bound, which is effective only for 0-1 loss. Subsequently, Andriushchenko 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 et al. (2023) 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 a connected manifold. Given sufficient training time and infinitesimally small η and ρ, they rigorously proved that SAM's dynamics would track the trajectory of a Riemannian flow with respect to sharpness, achieving the same sharpness-reducing 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 ρ.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 by Li et al. (2017) , who proposed a 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 et al., 2021) . In this paper, we follow this mathematical framework, reusing some of its notations and definitions. 
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 a given mini-batch γ, we define the mini-batch loss as f γ (x) with parameters x ∈ R d . The generalization loss is defined asf D (x) = E γ∼D [f γ (x)],while the empirical loss is defined asf S (x) = E γ∼S [f γ (x)]. Since we primarily analyze and discuss the empirical loss in this paper, we drop the dependency on S for simplicity, denoting the empirical loss as f (x) when no ambiguity arises.We use ∥ • ∥ to denote the Euclidean norm. The k-th order derivative of the loss f at x is denoted by∇ k f (x), which is a symmetric k-tensor in (R d ) ⊗k when x ∈ R d . We denote by λ 1 (∇ 2 f (x)) the largest eigenvalue of the Hessian matrix ∇ 2 f (x) and v 1 (∇ 2 f (x)) its corresponding unit eigenvector, with ∥v 1 (∇ 2 f (x))∥ = 1. Additionally, we use ∇ 3 f (x)(u, v) ∈ R dto represent the application of the symmetric 3-tensor ∇ 3 f (x) along directions u and v.SAM (Foret et al., 2021) seeks flat minima by minimizing the perturbed loss:min x max ∥ϵ∥≤1 f (x + ρϵ),where ρ is a predefined hyperparameter controlling the radius of the perturbation. Solving the inner maximization problem leads to ϵ SAM (x) = ∇f (x) ∥∇f (x)∥ . Differentiating the perturbed loss with respect to x, we get:∇f (x + ρϵ SAM (x)) = d(x + ρϵ SAM (x)) dx ∇f (x)| x+ρϵ SAM (x) ≈ ∇f (x)| x+ρϵ SAM (x) .In the last approximation, Foret et al. (2021) ignore the dependency of ϵ SAM (x) on x, leading to faster computational efficiency and higher generalization performance. Applying SAM in the stochastic case, the SAM iteration for mini-batch γ k is summarized as:EQUATIONWe use ∇f γ k (•) and ∇f γ k 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:EQUATION)3.3 SDE approximation for SGD and SAM Li et al. (2017) developed the following SDE to approximate discrete SGD:dX t = -∇f (X t )dt + √ η(Σ 1,1 (X t ))1 2 dW t , where W t is standard Brownian motion. Compagnoni et al. (2023) applied this framework to analyze the dynamics of SAM, deriving the following second-order SDE for SAM:EQUATIONwhere Σ a,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 Eq.3 as the second-order SDE since it includes up to second-order partial derivatives in both the drift and diffusion coefficients.In this paper, we use the largest eigenvalue λ 1 (∇ 2 f (x)) as a measure of sharpness, similar to prior works (Lyu et al., 2022; Arora et al., 2022; Wen et al., 2022; Damian et al., 2022) . Geometrically, the top eigenvalue of the Hessian matrix at a 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. (2021) at the minima. Note that at the minima, ∇f (x) = 0; thus,max ∥ϵ∥≤1 f (x + ρϵ) -f (x) ≈ max ∥ϵ∥≤1 ρϵ ⊤ ∇f (x) + ρ 2 2 ϵ ⊤ ∇ 2 f (x)ϵ = ρ 2 λ 1 (∇ 2 f (x)) 2 .Another possible choice is the trace of 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 Foret et al. (2021) . For any δ ∈ (0, 1) and σ > 0, with a probability over 1 -δ over the choice of S ∼ D n , we have(x) ≤ E ϵ∼N (0,σ 2 I d ) f D (x + ϵ), as inf D (x) ≤ f S (x) + dσ 2 2 λ 1 ∇ 2 f S (x) + Cd 3 σ 3 6 + L 2 √ n d log 1 + ∥x∥ 2 dσ 2 + O(1) + 2 log 1 δ + 4 log (n + d).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 SAMIn this chapter, 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 of our framework. Finally, in Section 4.3, we propose a corollary linking perturbation-eigenvector alignment with eigenvalue regularization, furthering our understanding of the implicit regularization effects inherent in SAM. 
heuristic derivations for the third-order sde: 
 We begin by examining the drift coefficient in Compagnoni et al. (2023) (Eq. 3):-∇f(X t ) -ρE ∇ 2 f γ (X t )∇f γ ∥∇f γ ∥ = -∇f (X t ) -ρ∇E∥∇f γ (X t )∥,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 a third-order Taylor expansion. The missing cubic term in the Taylor expansion is:ρ 2 2 E ∇ 3 f γ (X t )(∇f γ , ∇f γ ) ∥∇f γ ∥ 2 = ρ 2 2 ∇E ∇f ⊤ γ ∇ 2 f γ (X t )∇f γ ∥∇f γ ∥ 2 .The equality holds because ∇f γ is treated as a perturbation vector independent of X t , as SAM's implementation does not involve differentiating with respect to it (See Section 3.2 for a 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-gradient alignment during training. Recent findings indicate that during training, the gradient implicitly aligns with the top eigenvector of the Hessian matrix under certain conditions: (1) training with normalized full-batch gradient descent (Arora et al., 2022) ; (2) a 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 is crucial for interpreting the third-order term in our drift coefficient. Specifically, when the gradient is highly aligned with the top eigenvector, we have:ρ 2 2 E ∇ 3 f γ (X t )(∇f γ , ∇f γ ) ∥∇f γ ∥ 2 ≈ ρ 2 2 E∇ 3 f γ (X t )(v 1 (∇ 2 f γ (X t )), v 1 (∇ 2 f γ (X t ))) = ρ 2 2 E∇λ 1 (∇ 2 f γ (X t )),where the final equality follows from the properties of differentiating eigenvalues (see Magnus (1985) 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.In this subsection, we present the general formulation of the SDE for SAM. We refer to our SDE (Eq. 4) as the third-order SDE, to distinguish it from the second-order SDE (Eq. 3). For the complete statements and proofs, we refer the reader to Appendix B. Theorem 4.1. (Third-order SDE for SAM, Informal Statement of Theorem B.4) Let 0 < η < 1, T > 0, N = ⌊T /η⌋, and {x k : k ≥ 0} denote the sequence of discrete SAM iterations defined by Eq. 1. Define {X t : t ∈ [0, T ]} as the stochastic process satisfying the SDEdX t = -∇ f SAM (X t )dt + √ η(Σ SAM (X t )) 1 2 dW t , X 0 = x 0 (4) with f SAM (X t ) := f (X t ) + ρE∥∇f γ (X t )∥ + ρ 2 2 E ∇f ⊤ γ ∇ 2 fγ (Xt)∇fγ ∥∇fγ ∥ 2 , Σ SAM (X t ) := Σ 1,1 (X t ) + ρ(Σ 1,2 (X t ) + Σ 1,2 (X t ) ⊤ ) + ρ 2 Σ 2,2 (X t ) + 1 2 (Σ 1,3 (X t ) + Σ 1,3 (X t ) ⊤ ) ,where Σ a,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).Under sufficient regularity conditions, let ρ = O(η 1 3 ). Then, {X t : t ∈ [0, T ]} is an order-1 weak approximation of {x k : k ≥ 0}, i.e., for any test function g of at most polynomial growth, there exists a constant C independent of η such thatmax k=0,1,...,N |Eg(x k ) -Eg(X kη )| ≤ Cη.Our proof relies on the third-order Taylor expansion of f γ k X t + ρ ∇fγ k ∥∇fγ k ∥ , carefully matching the first-and second-order conditional moments and quantifying the errors for higher-order terms. Our third-order SDE reveals that SAM's implicit bias includes a complex combination of second-order and third-order terms, with scales of ρ and ρ 2 2 , respectively. Compared to Compagnoni et al. (2023) , our theorem offers two main advantages: first, we allow ρ to take larger values (η 2023)), which is more consistent with real-world settings; equivalently, our SDE achieves a lower approximation error for a fixed ρ. Second, our SDE explicitly captures SAM's implicit bias on the Hessian matrix, manifesting as the gradient of the Hessian in the gradient's direction. Additionally, the diffusion coefficient in Eq. 4 implies 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.1 3 compared to η 1 2 in Compagnoni et al. (The implicit bias introduced by the third term in the drift coefficient of the SDE (Eq. 4), i.e.,ρ 2 2 E ∇ 3 fγ (Xt)(∇fγ ,∇fγ ) ∥∇fγ ∥ 2, remains difficult to understand. As discussed in Section 4.1, if we assume 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 sharpness. Next, we quantify and formalize this heuristic approach. Define Align(ϵ, v 1 ) := 1 -min s∈{±1} ∥ ϵ ∥ϵ∥ -s • v 1 ∥ as a measure of alignment between the perturbation vector ϵ and the top eigenvector v 1 . It is worth noting that +v 1 and -v 1 are equivalent eigenvectors, which is why we define alignment as the maximum over both +v 1 and -v 1 . Corollary 4.1.1. Recall that ϵ SAM γ is defined in Eq. 2. Let s * denote the direction scalar, i.e.,s * = arg min s∈{±1} ∥ ϵ ∥ϵ∥ -s • v 1 ∥.Under the same conditions as in Theorem 4.1, and assuming a positive eigenvalue gap (see Assumption B.2 for a definition), we have the following:EQUATIONwhereEQUATIONwhere∇ f SAM ρ 2 (X t ) := ∇f (X t ) + ρE s * • λ 1 (∇ 2 f γ (X t ))v 1 (∇ 2 f γ (X t )) + ρ 2 2 E∇λ 1 (∇ 2 f γ (X t )).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 1 -O(ρ), 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, ρ, and ρ 2 2 . 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-or 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 A. Furthermore, if the alignment is at least 1 -O(ρ 2 ), then the gradient of the gradient norm oscillates in the direction of the top eigenvector. Notably, Bartlett et al. (2023) derived a 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 specific initialization. When these conditions are met, they require an alignment of the gradient with the top eigenvector of at least 1 -O(ηρ) = 1 -O(ρ 4 ). In comparison, our SDE framework is more general.Comparison with Wen et al. (2022) . Wen et al. (2022) 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 a much longer training time Θ(η -1 ρ -2 ) compared to our Θ(η -1 ). Thus, our SDE corresponds to the initial phase of their analysis regarding time scale, during which they do not conclude any implicit bias. In contrast, our SDE (Eq. 5), which indicates that the implicit bias comprises three components with different scales, provides richer insights in this phase. Second, they require ηln(1/ρ) to be sufficiently small, causing ρ to be exponentially smaller than η, whereas our theory accommodates a more practical range, ρ = O(η 1 3 ).5 Eigen-SAM: an explicit regularization method for the top eigenvalue of the Hessian 
failure of perturbation-eigenvector alignment in practice: 
 Within the theoretical framework of Section 4, a natural question arises: Is the perturbationeigenvector alignment sufficient in practice for SAM to effectively minimize sharpness? Unfortunately, we have empirically verified that such alignment may be poor in practice, even for relatively simple models. As a 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. (2021) ; Deng et al. (2024) , on CIFAR-10 ( Krizhevsky et al., 2009) . Figure 1 illustrates two 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.To address the issue of poor alignment, we propose a novel algorithm called Eigen-SAM, which aims to 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 in the drift coefficient can effectively minimize the largest eigenvalue. To achieve this alignment, we estimate the top eigenvector of the Hessian matrix once every p mini-batch steps (with p = 100 in our implementation) using the power method for q iterations (with q = 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 an estimate v of the top eigenvector, we decompose it into components parallel and perpendicular to the gradient direction, then add the perpendicular component to the perturbation vector to enhance alignment explicitly:EQUATION)where α is a hyperparameter that controls the strength of the explicit alignment. Since +v and -v are equivalent eigenvectors, we include sign(⟨∇f γ (x), v⟩) to determine the direction of v. Here, we always choose v to have a smaller angle with the gradient. The full algorithm is presented in Algorithms 1 and 2. In Appendix D, we provide an in-depth discussion of the theoretical properties of Eigen-SAM, including sufficient conditions for improving alignment (Proposition D.1) and its convergence rate (Theorem D.2), which is comparable to that of SAM.Algorithm 1 Power iteration to estimate the top eigenvector 1: Initialize v as a random unit vector 2: for t 2 = 1, 2, . . . , q do 3:Compute Hessian-vector productv = ∇ 2 f (x) • v 4: Normalize v: v ← v/∥v∥ 5: end for 6: return v Algorithm 2 Eigen-SAM 1: for t 1 = 1, 2, . . . , T do 2:Compute mini-batch loss f γ (x k )3:if t 1 mod p = 1 then 4:Use Algorithm 1 to estimate v 5:end if 6:Compute the perturbation ϵ t per Eq. 7 and 87:Perturb the model: xt = x k + ρϵ t 8:Update parameters:x k+1 = x k - η∇f γ (x k ) 9: end for 10: return x kThe 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 is approximately 2 + q p to 2 + 2q p 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) . 
numerical simulation of the third-order sde: 
 In this subsection, we validate the approximation error between our proposed SDE (Eq. 4) and the discrete SAM algorithm (Eq. 1). We trained a fully-connected network with one hidden layer, consisting of 784 hidden units and using GeLU activation, on the MNIST dataset (Deng, 2012) . The training was conducted with η = 0.01 and ρ = 0.2 (where ρ ≈ η 1 3 ). 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 2 .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 a more precise representation of the discrete SAM dynamics, thus enhancing the theoretical understanding of SAM.To evaluate the effectiveness of Eigen-SAM, we applied it to several image classification tasks on benchmark datasets, including CIFAR-10 ( Krizhevsky et al., 2009) , CIFAR-100 (Krizhevsky et al., 2009) , Fashion-MNIST (Xiao et al., 2017) , and SVHN (Netzer et al., 2011) . For these tasks, we used ResNet-18 (He et al., 2016) , ResNet-50 (He et al., 2016) , and WideResNet-28-10 (Zagoruyko and Komodakis, 2016) models.We selected SGD as 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 an initial learning rate of 0.1 for CIFAR-10, Fashion-MNIST, and CIFAR-100, and 0.01 for SVHN, adjusting the learning rate over time with a cosine schedule. The weight decay was set to 5 × 10 -5 , and the momentum was 0.9. Detailed hyperparameter settings are provided in Appendix E.The test set performance, reported in Table 1 along with the 95% confidence interval, shows that Eigen-SAM consistently achieves state-of-the-art performance across various datasets and models, validating its effectiveness and robustness. 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 PyTorch's official repository1 . For SAM and Eigen-SAM, we used an initial learning rate of 0.01 and trained for 4k steps, while for SGD, we trained for 8k steps. Table 2 shows the test accuracy, where Eigen-SAM consistently outperforms the baselines. We investigated the impact of varying the hyperparameter α on the test accuracy of Eigen-SAM.We conducted experiments on ResNet-18 with CIFAR-100, testing a range of α values, as shown in Figure 3 . We observed that the test accuracy peaks at α = 0.2, yielding a 0.9% improvement in test accuracy compared to α = 0, which corresponds to the standard SAM. These results suggest that α is a robust hyperparameter, as its variations do not cause significant performance fluctuations, while consistently enhancing performance.In Table 3 and Table 4 in Appendix F, we demonstrate how larger values of p affect generalization performance and observe that setting p 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 1, which typically requires only a few steps to converge. Figure 3 shows the Hessian spectrum at the 200th epoch for ResNet-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.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: We specify them in 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 to a 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, or as supplemental material. 
conclusion: 
 In this work, we analyzed the training dynamics of SAM using a 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 Eigen-SAM, an 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 achieves state-of-the-art performance across various tasks, validating both its accuracy and effectiveness.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?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, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run 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).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It 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 a 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 of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?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 of compute required for each of the individual experimental runs as well as 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 or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: Our research conduct with the NeurIPS 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 circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
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 is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or 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.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are 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 to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: We only use public datasets.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 to 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 a best faith effort.12. Licenses 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 use 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 a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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. 
 
 Appendix: 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. (2022) suggests that SAM's implicit regularization primarily occurs near the minima manifold, where ∇f (x) ≈ 0 and the quadratic term dominates the perturbed term. To test this, we designed a counterexample algorithm called Reverse-SAM, which applies SAM's ascent step using a negative normalized gradient, i.e., ϵ = -∇fγ (x) ∥∇fγ (x)∥ . If this theory holds in practice, then Reverse-SAM 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), a regularization method with a long history, tracing back to the works of Barrett and Dherin (2020) and Drucker and Le Cun (1991) . The second-order SDE proposed by Compagnoni et 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 is accurate in practice, EGR with a regularization coefficient equal to ρ should exhibit performance comparable to SAM.We trained ResNet-18 on CIFAR-10 to compare the performance of these two counterexample algorithms with SAM. Figure 4 shows 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.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. 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 α = (α 1 , α 2 ..., α n )• |α| := α 1 + α 2 + ... + α n • α! := α 1 !α 2 !...α n ! • For x = (x 1 , x 2 ..., x n ) ∈ R n , x α := x α1 1 x α2 2 ...x αn n • For a multi-index β, ∂ |β| β f (x) := ∂ |β| ∂x β 1 1 ∂x β 2 2 ...∂x βn n f (x)We denote the partial derivative with respect to x i by ∂ ei . for all x ∈ R d . Moreover, for each integer α ≥ 1 we denote by G α the set of α-times continuously differentiable functions R d → R which, together with its partial derivatives up to and including order α, belong to G. This definition originates from the field of numerical analysis of SDEs (Mil'shtein, 1986) . In the case of g(x) = ||x|| j , the bound restricts the difference between the j-th moments of the discrete process and those of the continuous process. We write O(ρ α1 η α2 ) to denote that there exists a function K ∈ G independent with ρ, η, such that the error terms are bounded by Kρ α1 η α2 . Assumption B.1. Assume that the following conditions on f , f i and their gradients are satisfied:• ∇f , ∇f i satisfy a Lipschitz condition: there exists L > 0 such that∥∇f (x) -∇f (y)∥ + n i=1 ∥∇f i (x) -∇f i (y)∥ ≤ L∥x -y∥;• f , f i and its partial derivatives up to order 7 belong to G;• ∇f , ∇f i satisfy a growth condition: there exists M > 0 such that∥∇f (x)∥ + n i=1 ∥∇f i (x)∥ ≤ M (1 + ∥x∥);• g and its partial derivatives up to order 6 belong to G; See Li et al. (2017) for a detailed discussion on Assumption B.1.To prove the Corollary 4.1.1, 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.For all k > 0, λ 1 (∇ 2 f (x k )) > λ 2 (∇ 2 f (x k )), and λ 1 (∇ 2 f (X kη )) > λ 2 (∇ 2 f (X kη )).In the proof of our main theorem, we will utilize the following two auxiliary lemmas.Lemma B.1. (Lemma 1 (Li et al., 2017 )) Let 0 < η < 1. Consider a stochastic process {X t : t > 0} satisfying the SDEdX t = b(X t )dt + η 1 2 σ(X t )dW t with X 0 = x ∈ R dand b, σ together with their derivatives belong to G. Define the one-step difference ∆ = X η -x, then we haveEQUATIONEQUATION)all the functions are evaluated at x.Lemma B.2. (Theorem 2 and Lemma 5 (Mil'shtein, 1986 ) )Under Assumption B.1, we define the onestep difference for the stochastic process ∆ = X η -x. if in addition there existsEQUATIONEQUATIONThen, for each g ∈ G, there exists a constant C such thatmax k=0,1,...N |Eg(x k ) -Eg(X kη )| ≤ CηNext, we will prove a lemma that controls the moments of the discrete process for SAM: Lemma B.3. Under Assumption B.1, let 0 < η < 1. We define:∂ ei f SAM (x) := ∂ ei f (x)+ρE j ∂ 2 ei+ej f γ (x)∂ ej f γ (x) ∥∇f γ (x)∥ + ρ 2 2 E jk ∂ 3 ei+ej +e k f γ (x)∂ ej f γ (x)∂ e k f γ (x) ∥∇f γ (x)∥ 2 .In addition, we define the one-step difference for the discrete-time algorithm as ∆ = x 1 -x. Then we have:EQUATION)All functions are evaluated at x.To evaluateE ∆i = -E ∂ ei f γ x + ρ ∥∇fγ (x)∥ ∇f γ (x), we start by analyzing∂ ei f γ x + ρ ∥∇fγ (x)∥ ∇f γ (x), which is the partial derivative in the direction e i . Taking the Taylor expansion, we have:EQUATION)where the residual term is defined in Folland (2005) . For some constant c ∈ (0, 1), it holds thatEQUATION)Now, we observe thatK i (x) := ∂ 4 ei+α f γ x + cρ ∇fγ (x) ∥∇fγ (x)∥ ρ |α| ∇fγ (x) ∥∇fγ (x)∥ α α!is a finite sum of products of functions that, by assumption, are in G. We can rewrite Equation ( 19) asEQUATIONwhich implies thatEQUATIONwhere Ki (x) = EK i (x).Therefore, we have ∀i = 1, 2, . . . , d,E ∆i = -∂ ei f SAM (x)η + O ηρ 3 .To prove the second statement, by definition, we haveEQUATIONEQUATION)Finally, it is clear thatEQUATIONEQUATIONNow we are ready to state the theorem and prove it.Theorem B.4. (Stochastic modified equations) Under Assumption B.1, let 0 < η < 1, T > 0, N = ⌊T /η⌋. Let x k ∈ R d , 0 ≤ k ≤ N denotethe sequence of SAM iterations defined by Equation 1. Define {X t } as the stochastic process satisfying the SDEEQUATION3 ) Then, {X t : t ∈ [0, T ]} is an order-1 weak approximation of {x k : k ≥ 0}, i.e. for each g ∈ G, there exists a constant C independent of η such thatmax k=0,1,...N |Eg(x k ) -Eg(X kη )| ≤ CηProof. We will check the conditions in Lemma B.2. As we apply Lemma B.1, we make the following choices:b(x) = -∇ f SAM (x) σ(x) = Σ SAM (x) 1 2First, for ∀i = 1, ..., d, we haveEQUATIONTherefore, we have that for someEQUATIONSecond, for ∀i, j = 1, ..., d, it holds thatEQUATIONE ∆i ∆jEQUATIONConsequently, we have for someEQUATIONThird, for ∀i j = 1, ..., d, we haveEQUATIONTherefore, for some K 3 (x) ∈ G, we haveEQUATIONAdditionally, for some K 4 (x) ∈ G, ∀i j = 1, ..., d,E 3 j=1 ∆ij LemmaB.3 ≤ K 4 (x)η 2 (39)By combining the four equations, Eq. 32, 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:ρ 2 2 E ∇f ⊤ γ ∇ 2 f γ (X t )∇f γ ∥∇f γ ∥ 2 - ρ 2 2 E v 1 (∇ 2 f γ (X t )) ⊤ ∇ 2 f γ (X t )v 1 (∇ 2 f γ (X t )) = O(ρ 3 ) (40)2. Under the supposition, the second statement is immediately followed by using the fact:EQUATION, andEQUATIONEQUATIONC Proof of Theorem 3.1Theorem 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 f D (x) ≤ E ϵ∼N (0,σ 2 I d )f D (x + ϵ), similar to Foret et al. (2021) . For any δ ∈ (0, 1) and σ > 0, with a probability over 1 -δ over the choice of S ∼ D n , we havef D (x) ≤ f S (x) + dσ 2 2 λ 1 ∇ 2 f S (x) + Cd 3 σ 3 6 + L 2 √ n d log 1 + ∥x∥ 2 dσ 2 + O(1) + 2 log 1 δ + 4 log (n + d).Proof. We use the PAC-Bayes theory in this proof. In PAC-Bayes theory, x follows a distribution, denoted by P , and we express the expected loss over x as follows:f D (P ) = E x∼P f D (x) f S (P ) = E x∼P f S (x)For any distribution P = N (0,σ 2 P I d ) and Q = N (x, σ 2 I d ) over x ∈ R d, where P is the prior distribution and Q is the posterior distribution, use the general PAC-Bayes theorem in Alquier et al. (2016) , for all β > 0, with a probability at least 1 -δ, we haveEQUATIONwhere Ψ is defined asΨ(β, n) = log E x∼P E S exp β f D (x) -f S (x) .When the loss function is bounded by L, then by using the Hoeffding's inequality we have:Ψ(β, n) ≤ β 2 L 2 8n .The task is to minimize the second term of RHS of (44), we thus choose β = 8n(KL(Q∥P )+log 1 δ )L . Then the second term of RHS of ( 44) is equal toKL(Q∥P ) + log 1 δ 2n × L.The KL divergence between Q and P , when they are Gaussian, is given by formulaKL(Q∥P ) = 1 2 dσ 2 + ∥x∥ 2 σ 2 P -d + d log σ 2 P σ 2 .For given posterior distribution Q with fixed σ 2 , to minimize the KL term, the σ 2 P should be equal to σ 2 + ∥x∥ 2 /d. In this case, the KL term is no less than d log 1 + ∥x∥ 2 dσ 2 . Thus, the second term of RHS isKL(Q∥P ) + log 1 δ 2n × L ≥ d log 1 + ∥x∥ 2 dσ 2 4n × L ≥ L when ∥x∥ 2 2 > σ 2 exp(4n/d) -1 .Hence, for any ∥x∥ 2 > σ 2 exp(4n/d) -1 , we have the RHS is greater than the LHS, the inequality is trivial. In the remainder of the proof, we only consider the case:EQUATION)Distribution P is Gaussian centered around 0 with variance σ 2 P = σ 2 + ∥x∥ 2 /d, which is unknown at the time we set up the inequality, since x is unknown. Meanwhile we have to specify P in 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., 2021) . We set c = σ 2 1 + exp(4n/d)P j = N 0, c exp 1 -j d I d P := P j : j = 1, 2, . . .Then the following inequality holds for a particular distribution P j with probability 1 -δ j withδ j = 6δ π 2 j 2 f D Q ≤ f S Q + 1 β KL(Q∥P j ) + log 1 δ j + Ψ(β, n) .Use the well-known equation:∞ j=1 1 j 2 = π 2 6, then with probability 1 -δ, the above inequality holds with every j. We pickj * := 1 -d log σ 2 + ∥x∥ 2 /d c = 1 -d log σ 2 + ∥x∥ 2 /d σ 2 (1 + exp{4n/d}) .Therefore,1 -j * = d log σ 2 + ∥x∥ 2 /d c ⇒ log σ 2 + ∥x∥ 2 /d c ≤ 1 -j * d ≤ log σ 2 + ∥x∥ 2 /d c + 1 d ⇒ σ 2 + ∥x∥ 2 /d ≤ c exp 1 -j * d ≤ exp(1/d) σ 2 + ∥x∥ 2 /d ⇒ σ 2 + ∥x∥ 2 /d ≤ σ 2 P j * ≤ exp(1/d) σ 2 + ∥x∥ 2 /d .Thus the KL term could be bounded as followKL(Q∥P j * ) = 1 2 dσ 2 + ∥x∥ 2 σ 2 P j * -d + d log σ 2 P j * σ 2 ≤ 1 2 d(σ 2 + ∥x∥ 2 /d) σ 2 + ∥x∥ 2 /d -d + d log exp(1/d) σ 2 + ∥x∥ 2 /d σ 2 = 1 2 d log exp(1/d) σ 2 + ∥x∥ 2 /d σ 2 = 1 2 1 + d log 1 + ∥x∥ 2 dσ 2For the term log 1 δ j * , with recall that c = σ 2 1 + exp(4n/d) andj * = 1 -d log σ 2 +∥x∥ 2 /d σ 2 (1+exp{4n/d}) , we have log 1 δ j * = log (j * ) 2 π 2 6δ = log 1 δ + log π 2 6 + 2 log(j * ) ≤ log 1 δ + log π 2 6 + 2 log 1 + d log σ 2 1 + exp(4n/d) σ 2 + ∥x∥ 2 /d ≤ log 1 δ + log π 2 6 + 2 log 1 + d log 1 + exp(4n/d) ≤ log 1 δ + log π 2 6 + 2 log 1 + d 1 + 4n d ≤ log 1 δ + log π 2 6 + log(1 + d + 4n).Hence, the inequalityf D (Q) ≤ f S (Q) + KL(Q∥P j * ) + log 1 δ j * 2n × L ≤ f S Q + L 2 √ n 1 + d log 1 + ∥x∥ 2 dσ + 2 log π 2 6δ + 4 log(n + d) ≤ f S Q + L 2 √ n d log 1 + ∥x∥ 2 dσ 2 + O(1) + 2 log 1 δ + 4 log(n + d). E ϵ∼N (0,σ 2 I d ) f D x + ϵ ≤ E ϵ∼N (0,σ 2 I d ) f S x + ϵ + L 2 √ n d log 1 + ∥x∥ 2 dσ 2 + O(1) + 2 log 1 δ + 4 log(n + d).Using the assumption that fD (x) ≤ E ϵ∼N (0,σ 2 I d ) f D x + ϵ , we reach f D (x) ≤ E ϵ∼N (0,σ 2 I d ) f S x + ϵ + L 2 √ n d log 1 + ∥x∥ 2 dσ 2 + O(1) + 2 log 1 δ + 4 log(n + d).Using the second-order Taylor expansion for f S x + ϵ , we obtainf S x + ϵ = f S x + ϵ T ∇ x f S x + 1 2 ϵ T ∇ 2 f S (x) ϵ + 1 6 i1,i2,i3 ∂ 3 f S (x + tϵ) ∂x i1 ∂x i2 ∂x i3 ϵ i1 ϵ i2 ϵ i3 ≤f S x + ϵ T ∇ x f S x + 1 2 λ 1 ∇ 2 f S (x) ∥ϵ∥ 2 2 + 1 6 i1,i2,i3 ∂ 3 f S (x + tϵ) ∂x i1 ∂x i2 ∂x i3 ϵ i1 ϵ i2 ϵ i3 ,where t ∈ [0, 1]. Thanks to E ϵ∼N (0,σ 2 I d ) ∥ϵ 2 ≤ E ϵ∼N (0,I d ) ∥ϵ 2 = dσ 2 , we haveE ϵ∼N (0,σ 2 I d ) f S x + ϵ ≤ f S x + dσ 2 2 λ 1 ∇ 2 f S (x) + Cd 3 6 E ϵ1∼N (0,σ 2 ) [|ϵ 1 |] E ϵ2∼N (0,σ 2 ) [|ϵ 2 |] E ϵ3∼N (0,σ 2 ) [|ϵ 3 |] ≤ f S x + dσ 2 2 λ 1 ∇ 2 f S (x) + Cd 3 6 E ϵ1∼N (0,σ 2 ) ϵ 2 1 1/2 3 = f S x + dσ 2 2 λ 1 ∇ 2 f S (x) + Cd 3 σ 3 6 . By the assumption f D (x) ≤ E ϵ∼N (0,σ 2 I d )f D (x + ϵ), we reach f D (x) ≤ f S x + dσ 2 2 λ 1 ∇ 2 f S (x) + Cd 3 σ 3 6 + L 2 √ n d log 1 + ∥x∥ 2 dσ 2 + O(1) + 2 log 1 δ + 4 log(n + d).In this subsection, first we show that Eq. 8 can indeed improve the alignment for moderate α. Let ω := cos( ∇fγ (x) ∥∇fγ (x)∥ , v), without loss of generalization, we suppose ω > 0. We only consider the case ω > √ 2 2 , since in the case 0 < ω ≤ √ 2 2 , any α > 0 will enhance the alignment. Using fundamental mathematics to solve the inequality, we obtainProposition D.1. Let ω > √ 2 2 , for any α ∈ 0, 2ω √ 1-ω 2 2ω 2 -1 , we have cos( ∇fγ (x) ∥∇fγ (x)∥ + αv ⊥ , v) > ω.Proposition D.1 shows our update can indeed improve the alignment for a wide range of α. For example, if cos(∇f γ (x), v) = 0.8, then α can be any value in (0, 3.43); if cos(∇f γ (x), v) = 0.9, then α 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 β-smooth function f satisfying f * = inf x f (x) > -∞, let ∆ := f (x 0 ) -f * ,and assume the mini-batch variance is bounded by σ 2 . Under Eigen-SAM, starting at x 0 with any perturbation size ρ > 0 and step size η = min{ 1 2β ,√ ∆ √ βσ 2 T } to minimize f ,we have1 T T -1 t=0 E ∥∇f γ (x t )∥ 2 ≤ O β∆ T + βσ 2 ∆ √ T + β 2 (ρ 2 + α 2 )Proof. By the definition of β-smoothness, we haveEf (x t+1 ) ≤ Ef (x t ) -ηE⟨∇f (x t ), ∇f (x t + ρϵ)⟩ + βη 2 2 E∥∇f (x t + ρϵ)∥ 2 ≤ Ef (x t ) -ηE⟨∇f (x t ), ∇f (x t + ρϵ)⟩ + βη 2 E∥∇f (x t + ρϵ)∥ 2 + E∥∇f (x t + ρϵ) -∇f (x t )∥ 2 ≤ Ef (x t ) -ηE⟨∇f (x t ), ∇f (x t + ρϵ)⟩ + βη 2 E∥∇f (x t + ρϵ)∥ 2 + σ 2 = Ef (x t ) - η 2 E∥∇f (x t )∥ 2 - η 2 E∥∇f (x t + ρϵ)∥ 2 + η 2 E∥∇f (x t ) -∇f (x t + ρϵ)∥ 2 + βη 2 E∥∇f (x t + ρϵ)∥ 2 + σ 2 ≤ Ef (x t ) - η 2 E∥∇f (x t )∥ 2 + β 2 η 2 E∥x t -(x t + ρϵ)∥ 2 + βσ 2 η 2 = Ef (x t ) - η 2 E∥∇f (x t )∥ 2 + β 2 η(ρ 2 + α 2 ) 2 + βσ 2 η 2 .Rearranging this inequality and deviding both sides by ηT 2 , we have1 T T -1 t=0 E∥∇f (x t )∥ 2 ≤ 2 ηT (Ef (x 0 ) -Ef (x T )) + β 2 (ρ 2 + α 2 ) + 2βσ 2 η ≤ 2∆ ηT + β 2 (ρ 2 + α 2 ) + 2βσ 2 η. By using η = min{ 1 2β , √ ∆ √ βσ 2 T}, we complete the proof.For hyperparameter ρ, we follow the guidelines by Foret et al. (2021) , setting ρ to 0.05 for CIFAR-10 and Fashion-MNIST, 0.01 for SVHN, and 0.1 for CIFAR-100. We ensure that SAM and Eigen-SAM use the same ρ for a 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 α = 0.2 works the best for almost all cases. 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 on NVIDIA RTX 4090 24GB GPUs. A limitation of this work is the additional memory and time required to estimate the top eigenvalue of the Hessian matrix. Improving the efficiency of Eigen-SAM is a direction for future research.Answer: [Yes] Justification: Yes, our code is open source.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.https://pytorch.org/vision/main/models/generated/torchvision.models.vit_b_16.html 


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

Reviews: review: paper_summary: paper_summary: ### Summary:



The 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. 

### Main Results: 
- Theorem 4.1: relating max e.v. of Hessian to generalization via PAC-Bayesian bounds.
- Theorem 4.2: a 3rd order SDE for SAM, that is used later to study top e.v. alignment.
strength_weakness: strength_weakness: ### Cons:

- the need for 3rd order SDEs for SAM are not well justified
- the current version can be improved in terms of clarity/typos
questions: questions: ### Questions/Comments:


This 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. 


My comments are as follows:

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

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




#### Other comments:




- line 79 -- I suggest using $\|$ instead of $||$ to denote norms.

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


- line 84 -- It might be useful to emphasize that you use the 2-norm of vectors to define SAM.

- line 85 -- typo "aslo"

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

- line 156 -- "SDE 5" and "SDE 3" are using confusing notation.

- line 160 -- typo $x_\in$

- 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$?

- line 163 -- why $\Sigma^{2,2}$ does not depend on $x$? Is this a typo?

- line 167 -- I suggest replacing $g \in G$ with just a few words explaining it. This makes that part clearler. 

- line 195 -- it looks like to be typo, what does the middle term mean in the formula?

- Experiments in Figure 1 look noisy -- how many times did you iterate it? 

- line 290 -- typo

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

- 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:
        - A Universal Class of Sharpness-Aware Minimization Algorithms, ICML 2024
        - How to escape sharp minima with random perturbations, ICML 2024
limitations: limitations: N/A
paper_summary: paper_summary: The authors consider sharpness aware minimization problem and derive third-order SDE governing SGD dynamics.

As 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.
strength_weakness: strength_weakness: 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.

Bound (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.
questions: questions: How to practically verify assumptions on the loss? Do the assumptions of the main theorems hold in conducted empirical study?
limitations: limitations: -
paper_summary: paper_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.
strength_weakness: strength_weakness: 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. 

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.

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.

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?

**Minor:**

- Typo in line 249: should be $v$
- Typo in line 290: should be $\rho$

[1] Wen, Kaiyue, Tengyu Ma, and Zhiyuan Li. "How does sharpness-aware minimization minimize sharpness?." arXiv preprint arXiv:2211.05729 (2022).

[2] Kaur, Simran, Jeremy Cohen, and Zachary Chase Lipton. "On the maximum hessian eigenvalue and generalization." Proceedings on. PMLR, 2023.
questions: 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: limitations: The authors have mentioned in Section A.5 that Eigen-SAM requires additional computational time.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

C. Additionally, a comprehensive analysis of computational costs is required.
questions: questions: A. Could you offer theoretical evidence elucidating the poor alignment in Section 5.1?

B. What is the computational cost in comparison to other algorithms? Can Eigen-SAM be feasibly implemented in large-scale models?
limitations: limitations: Please refer to the Weaknesses.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
- The lack of publicly available implementation may pose a barrier to community adoption.
- 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.
- Showing effectiveness in more practical settings like ImageNet would impact my score.
- Comparisons with other optimization methods using eigenvalues would strengthen this paper.
questions: 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?
- Is Figure 4(a) showing sensitivity analysis for alpha rather than rho? This is a bit confusing to me.
- 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?
- 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?
- 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?
- For Table 1, can you also show the results of how much the maximum eigenvalue improved?
limitations: 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.
- A significant limitation is that experiments are conducted only on ResNet models, and the effectiveness on more modern models like ViT is unknown.

Comments: title: Official Comment by Reviewer B3hY
Official Comment by Reviewer B3hY
Official Comment by Reviewer B3hY
Thanks for the rebuttal
comment: Dear Reviewer B3hY,

Thank you for responding to our message and giving us the opportunity to share our thoughts. While we appreciate the contributions of CR-SAM and Fisher SAM and are willing to discuss these works in the related work section, we believe that the theoretical contributions of these papers and ours are different.

CR-SAM aims to minimize the trace of the Hessian matrix, whereas Fisher SAM seeks to find adversarial models within a ball around $\theta$ with respect to the Fisher distance. Moreover, the theories in these papers primarily involve the extension of PAC-Bayes in SAM for their specific cases. As mentioned before, one of the main contributions of our paper is to advance the theories of [1, 2], which are two significant theoretical works in sharpness-aware minimization, i.e., ours offers a better way to study the dynamics of SAM and the convergence for more realistic $\rho$.

We further build on these theories to propose a practical method that demonstrates our claims: (i) using third-order expansion, we can approximate the dynamics of SAM more accurately, and (ii) our approach can obtain smaller top eigenvalues, leading to improved performance comparing to SAM.

We believe that it is not appropriate to use the experimental standards of CR-SAM/Fisher SAM for our case, especially since experiments on ImageNet are not affordable for many due to resource limitations. Moreover, in addition to [1, 2], there are many published works in the SAM domain that did not conduct experiments on ImageNet.  

[1] Wen, K., Ma, T., and Li, Z. (2022). How does sharpness-aware minimization minimize sharpness? CoRR, abs/2211.05729.

[2] 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 25209–25253. PMLR.
Thank you for your response and I aknowledge for the additional experiments you provided. 

However, I must reiterate my earlier point regarding the scope of your experiments. The two related works[1,2] I mentioned that improve SAM have demonstrated their methods on ImageNet scale datasets. In contrast, the studies you referenced focus more on theoretical understanding or exploring phenomena from different perspectives, rather than proposing new methods. While you have proposed Eigen-SAM and provided evidence of its effectiveness, the limited experimental result, specifically,  CIFAR scale experiments with ResNet-based and ViT-based architecture do not, in my view, sufficiently establish its efficacy.

I encourage you to consider this aspect in further revisions or future work.

[1] CR-SAM: Curvature Regularized Sharpness-Aware Minimization

[2] Fisher SAM: Information Geometry and Sharpness Aware Minimisation
Dear Reviewer B3hY,

We appreciate your feedback on our work and thank you for acknowledging our developed theories.

While we agree with your point that further validation of this work on more practical datasets and across multiple model architectures could strengthen our findings, we believe that the number of experiments across multiple model architectures (ResNet, WideResNet-28-10, VIT-B-16: please refer to **the pdf of the global response**) we have conducted in our paper, along with the additional ones in the rebuttal, is appropriate for a theory-based paper. In your comparison with [1, 2], it’s worth noting that [1] did not include any experiments, while [2] performed some experiments to demonstrate certain behaviors. We kindly suggest that the reviewer take this aspect into consideration.

[1] Wen, K., Ma, T., and Li, Z. (2022). How does sharpness-aware minimization minimize sharpness? CoRR, abs/2211.05729.

[2] 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 25209–25253. PMLR.

Best regards,
I would like to express my sincere appreciation for detailed response and the considerable effort authors have invested in this research. I acknowledge authors' theoretical contributions are commendable.

However, my primary concern remains with the empirical evaluation. I maintain my current score as I believe that the empirical results presented, while promising, do not meet the soundness required for acceptance at a conference of NeurIPS's caliber. Specifically, the limited experimental setup and the observed marginal improvements, though noteworthy, may not sufficiently demonstrate the broader applicability and effectiveness of the proposed approach.

I recommend further validation of this work on more practical datasets and across multiple model architectures. Such additional evaluations would strengthen authors' findings and potentially make authors work more suitable for submission to future venues.
Dear Reviewer B3hY,

We believe that we have addressed most of your questions and concerns, except for conducting additional experiments on ImageNet due to our limited computational resources. We hope you will give greater consideration to the theoretical contributions of our paper, which establish a connection between generalization loss and the top eigenvalue of the Hessian matrix. Additionally, we have developed a third-order SDE of SAM to study its dynamics and identify the factors important for minimizing the top eigenvalue.

We would like to note that our theories remarkably advance two important theoretical works in SAM [1] and [2]. Specifically, comparing to [1], firstly, their theory requires a much longer training time $Θ(η^{−1}ρ^{−2})$, compared to our $Θ(η^{−1})$. Therefore, our SDE corresponds to one phase one of their main analysis in terms of the time scale, during which they do not draw any conclusions about implicit bias. In contrast, our SDE (6), which indicates the implicit bias consists of three components with different scales, is much richer in this phase. Secondly, they require $ηln(1/ρ)$ be sufficiently small, leading ρ to be exponentially smaller than η, whereas our theory works for a more practical range, $ρ = O(η^{1/3}$). Comparing to [2], our theorem has the following two advantages: firstly, we allows ρ to take larger values ($η^{1/3}$ compared to $η^{1/2}$ in [2]), which is more consistent with real-world settings, or equivalently, our SDE has a lower approximation error for fixed ρ; secondly, our SDE reveals SAM’s implicit bias on the Hessian matrix, which occurs as the gradient of the Hessian along the direction of the gradient. 

 

[1] Wen, K., Ma, T., and Li, Z. (2022). How does sharpness-aware minimization minimize sharpness? CoRR, abs/2211.05729.

[2] 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 25209–25253. PMLR.

Please let us know if there is anything else we can do to increase your recognition of our work.
Thank you for your thoughtful feedback and for bringing up the comparison with CR-SAM and Fisher SAM. However, while CR-SAM aims to regularize the trace of the Hessian, in high-dimensional settings, the trace often differs significantly from the top eigenvalue, and Fisher SAM does not work on eigenvalues because it aims to find the perturbation w.r.t. the geometry specified by the Fisher distance. Therefore, it is difficult to compare these methods beyond final performance. Nevertheless, we are willing to discuss CR-SAM and Fisher SAM in related work.

Due to limited computational resources, we acknowledge that we were unable to expand the scale of our experiments as much as possible. However, the theoretical developments/derivations are the main contributions of our work. We hope that you can focus more on the theoretical aspects, as better understanding the working principles of SAM can often guide the design of more principled algorithms, which is also beneficial to the community.
Thank you for your positive feedback and for increasing the score. We will include the experimental results in the manuscript as you suggested.
Thank you for your positive feedback and for increasing the score. We will incorporate the explanation of the power method and update Figure 1 with averaged results from repeated experiments to reduce noise.
Thank you for your response and for addressing my comments. I strongly recommend the authors to add the explanation on the power method to their paper. Moreover, please replace Figure 1 with new experiments showing less noisy observations (e.g., by repeating the experiments and taking averages).

Given that most of my comments have been addressed, I have decided to increase the score of this paper.
Thanks for the response. Regarding the suggestion to compare with other optimization methods using eigenvalues, there are SAM variants that constrain the maximum eigenvalue, such as those proposed in [1] (Hessian Trace, which regularize eigenvalue indirectly) and [2] (Eigenvalue of FIM). Although these methods are derived differently, evaluating the performance of similar approaches that aim to improve SAM is important for the community. Additionally, since the current experimental results are limited to a narrow set of models, there are concerns about the broader effectiveness of the proposed method. I will maintain my scores.

[1] CR-SAM: Curvature Regularized Sharpness-Aware Minimization

[2] Fisher SAM: Information Geometry and Sharpness Aware Minimisation
Thank you for conducting the additional experiments and providing the discussions in response to the review. The authors' rebuttal has addressed most of the previously raised concerns, and I am inclined to give a positive review score. I suggest that the experimental results be included in the manuscript.
Thanks for the detailed rebuttal and explanation. I have no further questions.