Unveiling m-Sharpness Through the Structure of Stochastic Gradient Noise
Authors: Haocheng Luo
Mehrtash Harandi
Dinh Phung
Trung Le
Venue: NeurIPS 2025 poster ()
Keywords: Optimization
Sharpness-aware Minimization
Stochastic Differential Equations
Primary area: optimization

Abstract: Sharpness-aware minimization (SAM) has emerged as a highly effective technique to improve model generalization, but its underlying principles are not fully understood. We investigate m-sharpness, where SAM performance improves monotonically as the micro-batch size for computing perturbations decreases, a phenomenon critical for distributed training yet lacking rigorous explanation. We leverage an extended Stochastic Differential Equation (SDE) framework and analyze stochastic gradient noise (SGN) to characterize the dynamics of SAM variants, including n-SAM and m-SAM. Our analysis reveals that stochastic perturbations induce an implicit variance-based sharpness regularization whose strength increases as m decreases. Motivated by this insight, we propose Reweighted SAM (RW-SAM), which employs sharpness-weighted sampling to mimic the generalization benefits of m-SAM while remaining parallelizable. Comprehensive experiments validate our theory and method.

Unveiling m-Sharpness Through the Structure of Stochastic Gradient Noise
Haocheng Luo
Mehrtash Harandi Dinh Phung Trung
Monash University, Australia
(haocheng luo, mehrtash. harandi, tinh.phung trunglm)@monee (atg:lmm@mmmaash.coo
Abstract
Sharpness-awwan minimization (SAM) has emerged highly effective technique to improve model generalization, but its underlying principles are not fully under- stood. We investigate m-sharppess,, where SAM performance improves monotoni- cally as the micro-batch size for computing perturbations decreases, a phenomenon critical for distributed training yet lacking rigorous explanation. We leverage extended Stochastic Differential Equation (SDE) framework and analyze stochastic gradient noise (SGN) characterize the dynamics of SAM variants, including n-SAM and m-SAM. Our analysis reveals that stochastic perturbations induce an implicit variance-basee based sharpness regularization whose strength increases as decreases. Motivated by this insight, we propose Reweighted SAM (RW-SAM), which employs sharpness- tharpness-weighted sampling to mimic the generalization benefits of m-SAM while remaining parallelizable. Comprehensive experiments validate our theory and method.
Introduction
In machine learning, gradient-based optimization algorithms aim to minimize the following loss function:
min f(x) = n Σ fi(x), #EB
where Rdd denotes the parameter, fz) represents the loss on the i-th sample, i ranges from 1 to n, and n is the size the training set. We primarily focus on stochastic algorithms, where at each step k, an index set 7/k of fixed cardinality In| is sampled uniformly at random. We denote the mini-batch loss by fzu (x) = Ii Σier fx().
We investigate the recently proposed Sharpness- Aware Minimization (SAM) (Foret et al.
2021
which has achieved remarkable success in various application domains (Foret et al. 2021 K Kwon
ct al. 202T Kaddour et al. [2022] Liu et al. 2022a Qu et al. 2022; Wang et al. 2024 | Nguyen et al.
2024 Hoang-Anh al.
2025 [Singh et al. 2025). It seeks flat minima by minimizing the perturbed
loss minger* f(z + pc (x)), where the perturbation c* (x) is defined as the solution to
max (VI(z),)), |
which admits the closed-form expression e (x) V<<>//I Here p > is hyperparameter controlling the perturbation radius. The algorithm, referred to as n-SAM, computes its perturbation using the full-batch gradient. Its update at iteration is
Ik+1 17 V frx 2x + f((e))
(V
(3)
where 7) > the learning rate, p the perturbation radius.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Calculating the perturbation on the entire training dataset at each step is prohibitively expensive. Therefore, Forett Forre et al. (202I) suggest estimating the perturbation using a mini-batch, resulting in SAM commonly used in practice. We refer to the practical SAM algorithm as mini-batch SAM to distinguish it from other variants. The update rule can be summarized:
Vln (IR)) (2x)
2k+1 = 23 -
fT -
I +
(4)
Interestingly, it has been observed that although mini-batch SAM was proposed computationally efficient variant, it exhibits remarkable generalization ability. In contrast, the original n-SAM (3)
offers little no improvement in generalization (Foret 2021 Andriushchenko and Flammarion 2022).
m-SAM and m-sharpness. m-SAM refers to dividing mini-batch of data into disjoint micro-batches of size m. and independently computing perturbations and gradients for each dicro-batch which are then combined to update the parameters. It has been widely observed that the practical performance of m-SAM improves monotonically as m decreases, a phenomenon known as m-sharpness (Foret et al.) 2021 Behdin al. 2022 Andriushchenko and Flammarion, 2022). It is worth noting that when m is smaller ler than the batch size, the perturbation must be computed sequentially across micro-batches,,
which cannot be parallelized and therefore introduces substantial additional computational overhead, although smaller m typically leads to improved generalization performance.
The update rule for m-SAM can be written as:
Jk+1 3g
Tm
Σ
VII,
Ek +
JI (xx)
Iz (Tx)
(5)
IGC Kil=
where Zj are disjoint subsets of 7k. each with size m.
In synchronous data-parallel (multi-GPU) training with D devices and per-device batch size b, SAM is typically implemented by computing the perturbation locally on each device using its own data and then aggregating the perturbed gradients across devices; this implementation is exactly instance of m-SAM with m = b. The portu-perturration design avoids an additional cross-device synchronization in the inner perturbation step (only the outer gradient aggregation requires all-reduce), thereby reducing per-step communication overhead. The empirical observation of m-sharpness has thus become a practical comerstone for deploying SAM scale.
To provide a theoretical explanation for m-sharpness, we extend the recent Stochastic Differential Equation (SDE) framework of Li et al. (2019;; Compagnoni et al. (2023): Luo et al. (2025) by jointly tracking both 1) and p to arbitrary expansion orders, providing a unified basis for analyzing SAM and its variants. Under this framework, we derive closed-form drift terms for three unnormalized SAM (USAM) variants, including n- A-M, mini-batch USAM, and m-USAM, revealing how stochastic gradient noise (SGN) drives implicit sharpness regularization and closely correlates with generalization performance. We further extend our analysis to the three normalized (vanilla) SAM variants and observe similar noise-induces pattern, albeit without closed-form solutions. Motivated by our theory, we propose a sample-reweighting method that uses the magnitude of the SGN as an importance measure.
Our contributions are threefold:
We develop an extended SDE framework that simultaneously tracks both y and p arbitrary orders, and use it to derive SDEs with controllable error terms for n-USAM/SAM,, mini-batch USAM/SAM, and M-USAM//AM.
We provide a theoretical explanation of the m-sharpness phenomenon, showing how SGN induces an implicit variance-regee regularization term in the drift, whose strength increases m decreases. We further present empirical evidence demonstrating that this effect is strongly correlated with generalization performance.
We introduce Reweighted SAM (RW-SAM), an adaptive reweighting mechanism that assigns larger weights to samples with higher SGN magnitudes, thereby strengthening implicit sharpness regularization; its superior generalization confirmed through comprehensive experiments.

---

2 Related Work
Sharpness-Awan Minimization. Sharpness-Aware Minimization (Foret et al. 2021) has attracted increasing attention due to its consistent improvements in generalization across a wide range of tasks. A growing body of work has been devoted to analyzing and enhancing SAM, including
investigations into its generalization principles (Andriushchenko and Flammarion, 2022, Möllenhoff
and Khan 2022 Wen et al.) 2022 2023 Agarwala and Dauphin| [ et al. 2024 Luo
et al. 2025) and convergence properties (Khanh et al. 2024 Oikonomou and Loizou| (025). exploring
its applications in various domains, and developing algorithmic variants further improve both
generalization (Kwon et al. 2
et al., [0022 Liu et al. [002bb] Kim et al. 2022 Nguyen
al.,
2023a bj
et al.] 2024c Wu al. 2024] Lahmasebi al. 2024; Truong al. 2024 Li et
2024b
2025
Phan et 2002) and computational efficiency (Du et al. [202]]] Liu et al. 2022a Du
et al.
2022 Mordido et al., 2023 Tan et al. 2024 Xie al. (024)..
m-sharpness. m-sharpness has long been mysterious phenomenon in the field of SAM-related research and was first introduced in the original work of Foret al. (2021). They observed that although SAM theoretically aims to minimize the perturbed loss over the entire training set, its computationally efficient variant, mini-batch SAM, which computes perturbations only at the mini- batch level, outperforms n-SAM, which applies perturbations at the full-batch level. More generally, the generalization performance of m- mn-AMM improves monotonically as m decreases. This phenomenon was further confirmed through extensive experiments in a single- setting by Annriusshhhhno and Flammarion (2022), and in multi-GPU setting by Behdin et al. (2022). Although Andriushchenko and Flammarion (2022) proposed several hypotheses to explain it, they were later invalidated by their own experiments, and the underlying cause of this phenomenon remains an open question. is important to note that our definition of m-sharpness follows the original work Foret et al. (2021) and the pioneering contributions of (Andriushchenko and Flammarion (2022). Some studies have adopted different definitions. For example. Wen et al. (2022) refer to the deterministic algorithm as n-SAM and SAM with a batch size of as T-SAM, whereas Behdin al. (2022) define m as the number of divided micro-batches.
Structure of stochastic gradient noise. In expectation, a widely accepted assumption that the stochastic gradient serves as an unbiased estimator of the full-batch gradient (Jastrzebski et al. 2017 Zhu et al. 2018 [HooChen et al. [02T] Ziyin al., (02T). Regarding the covariance of SGN, Simsekli et al. (2019) assumed it to be isotropic. However, this view was later challenged by Xie et al. (2020) and et al. (202I), who argued that Siisseli et al. (2019) were actually analyzing gradient noise across different iterations rather than noise arising from mini-batch sampling. They provided extensive evidence supporting the idea that the latter can be well modeled as a multivariate Gaussian variable with an
ptappi///paaamteeeeeeeeedee covariance structure. Furthermore, Xie et al. (2023) conducted statistical tests on the Gaussianity to support this perspective.
3 Theory
3.1
Notation and assumption
In this paper, we denote by || the Euclidean norm, and the expectation operator E is taken with respect to the random index set unless otherwise stated. We assume that the stochastic gradient is an unbiased estimator of the full gradient and possesses finite second moment.
Assumption 3.1. We assume that sampling an index i uniformly at random yields i.i.d, stochastic gradients
f((x) = V f(z) + (x(x,, where == - Vz((2) is the full gradient and (xx denotes the SGN. We further assume [(((()) = O. (C) = V(z),
where
V(z)
v Vf())) vf(x) Vf)=
An important consequence of this assumption is E V |/ (V((())), which we will use repeatedly.
3

---

3.2 Overview of two-parameter approximation
We extend the existing SDE framework for SAM (Cmmaagnnni et al. 2023 Luo et al. 2025p. which can only track to order 1, while ours can jointly track two parameters 17 and to arbitrary expansion orders, thereby decoupling their convergence rates and enabling precise control of the overall approximation error. Specifically, 1/ governs the higher-order terms in Dynkin's s formula, while p captures the remainder term arising from the Taylor expansion (see the detailed formulations in Appendix A). By employing a Dynkin expansion instead of a full Itó-Taylor expansion, it avoids the proliferation of terms and provides streamlined approach to controlling the remainder error in the two-parameter setting. Another major advantage of this approach is that it allows us to let y and p tend to zero at independent rates, rather than being constrained to a fixed ratio as in the al. (2023) and Luo et al. (2025). The definition of a two-parameter weak approximation of order (a, 5) is as follows.
work of @ooopagnonii
Definition 3.2 poo-paraeeter weak approximation). Let T 0, 0.0000 1,0 11.00- and set
N = [T/n]- Let
{ and (X.hee(o.r)
be a discrete-time and a continuous time stochastic process, respectively. We say that X₃ is an order- (ar, (3) weak approximation of s& if, for every g E Ga+1, there exists constant C 0, independent of v). P. such that
T (()(
B[g(Tx)] s C(n + pO-)).
max
OSKEN
3.3 SDE approximation for USAM variants
We begin by considering USAM (Andriuhchhenko and Flammarion, 2022 Compagnoni et al. 2023 Dai al. 2024 Zhou et al. 2024), which is widely studied theoretically friendly variant of SAM and can achieve comparable performance to SAM in practice. The update rules for the different variants the USAM algorithm are shown below. We will see that for USAM, all expressions the drift term are in closed form, providing an intuitive understanding. In Section 3.4 conclusions the standard SAM.
we will extend our
mini-batch USAM: Fh++1 = Th. V (I + V Sn (mx))
(6)
n-USAM:
Tk+1 = Ik- - V (Ix + vf(x)))
(7)
m-USAM:
Tk+1 Tk
hl
Σ
(Ix
(8)
I,CT
Theorem 3.3 (Mini-batch USAM SDE informal statement of Theorem B.2 adapted from Theorem
3.2 of Compagnoni et al. 2023)). Under Assumption 3./and mild regularity conditions, the solution of the following SDE (9) is an order-(1,1) weak approximation of the discrete update of mini-batch
USAM (6) with batch size by]:
dx, -v -V(((V)))) " |IV((X)1)= 2/1
(V(((x)) di (AFUSAM(X.)JWV
implicit regudarization
Based on Eq. (9) can be observed that under our Assumption 3.1 the implicit regularization term of mini-batch USAM (6) can be divided into the gradient of two parts: the squared norm of the full-batch gradient and the trace of the SGN covariance. The latter, which arises additionally from Theorem 3.2 (Compagnoni al.) (023). is due the use of random batches for perturbation in mini-batch USAM. We will see that if we use a deterministic perturbation, the regularization effect on the SGN covariance will disappear, as stated in the following theorem:
Theorem 3.4 (n-USAM SDE informal statement of Theorem B.4). Under Assumption 3.1 and mild
regularity conditions, the solution of the following SDE (10) is an order-(1,1) weak approximation of the discrete update of n- USAM 7 with batch size hy]:
dx, - -x(x-V(((X)) P (V V USAM SEN-USSMMMX.))
(10)

---

As observed in Table Appendix J USAM and SAM exhibit similar behavior under full-batch perturbations, meaning that n-USAM cannot improve generalization performance like mini-bate USAM. Therefore, we argue that the sharpness regularization effect of mini-batch USAM actually stems from the last term of its drift term, which the gradient of the trace of the SGN covariance, i.e., Vtr(V(X.))
Having understood that n-USAM lacks the sharpness regularization benefits brought by SGN in the drift term, we analyze another contrasting variant, m-USAM, which uses smaller micro-batches to compute the perturbation. Within the framework of the SDE approximation, we can derive the following theorem.
Theorem 3.5 (m-USAM SDE informal statement of Theorem B.J Under Assumption 3. and mild regularity conditions, the solution of the following SDE 11 is order- (1 weak approximation of the discrete update of m-USAM (
dxx - ((((X), + [V(())) + P (x.)((( di
mm m -SA-USAM (Xi)WW,
(11)
2m
It is worth noting that Theorem 3.3 is recovered by setting hy|- From this SDE approximation,
see that m- m-USAM's advantages arise from two sources. First, amplifies the sharpness regularization in the drift term: the coefficient
the SGN covariance changes from p/ (2(2 to p/ (with m < 17D. Second, the diffusion term in m-USAM is reduced by a factor of m//7 compared to mini-batch USAM with batch size m. Since the diffusion term captures the random fluctuations that counteract implicit regularization, shrinking it enhances the stability of the method.
3.4 SDE approximation for SAM variants
We now turn to the analysis of (normalized) SAM. With the addition of the normalization factor, the situation becomes much more complex because the expectation of the (unsquared) norm does not have clementary expression. However, the overall pattern remains similar to the case of USAM. We first present the SDE approximation theorem for SAM variants, then highlight their key differences from the unnormalized version.
Theorem 3.6 (Mini-batch SAM SDE - informal statement of Theorem adapted from Theorem 3.5 of Compagnoni et al. (023)). Under Assumption and mild regularity conditions, the solution of the Tollowing SDE (12) is an order-(1, 1) weak approximation of the discrete update of mini-batch
SAM with batch size 7
dx, (f((())))
P E b
VV(X))) dt (NESAM(X.)dVV
(12)
Theorem 3.7 (n-SAM SDE - informal statement of Theorem C.4). Under Assumption 3.1 and mild
regularity conditions, the solution of the following SDE 13 is an order-(1, 1) weak approximation of the discrete update of n-SAM ( with batch size In]:
dX. = -J((())) el P +
D n SAM (X,)dW,
(13)
Theorem 3.8 (m-SAM SDE informal statement of Theorem C.7p. Under Assumption 3.7 and mild
regularity conditions, the solution of the following SDE (14 is an order-(1, 1) weak approximation of the discrete update of m-SAM (5)
dx, - P m
E
Σ iEZ,
VS(().I))
me h
mm SAM (X.)WW,
(14)
Unlike the unnormalized algorithm, these SAM regularization terms cannot be expressed simple functions of the full gradient and SGN covariance. Nonetheless, the following proposition clarifies how the regularization terms in the SDEs depend on the mini/micro batch size, revealing an inverse relationship between the norm and the batch size.
Proposition 3.9. Under Assumption 3.1 the following inequalities hold:
/VVI//)))) tr(V(x))
V f [2)]1 < (x)) < [f BEIE,
[V
hl
5

---

where E,, denotes the SGN in Furthermore, V ( a log-concave distribution, then we have (x))] monotonically increases as Iyy decreases.
A complete proof of Proposition 3.9 given in Appendix [ Notably, log-concave distributions include many of the standard distributions of interest, such as the Gaussian and exponential distributions.
Based on the above theorems and propositions, we observe that mini-batch SAM (like USAM) regularizes the magnitude of the SGN. In contrast, n-SAM loses this noise-regularization effect, leading to degraded performance, whereas m-SAM amplifies it and thus achieves improved results. One intuitive explanation that in larger batches, these noise terms tend to cancel each other out, causing the averaged stochastic gradient to concentrate around its expectation and diminishing the contribution of SGN.
3.5 Benefits of SGN Covariance Regularization for Generalization
Building on our observation that USAM/SAM variants impose different strengths of sharpness regularization, i.e., variance-based regularization of the SGN, we investigate how directly regularizing the covariance of the SGN further enhances generalization performance.
We address this from two perspectives. First, following the work Neu al. (02 Wang and Mao (2021). bound on the generalization error can be decomposed into a sum of munual-information terms, among which the "trajectory term" is controlled by the covariance of the SGN. This implies that
by implicitly regularizing the covariance of the SGN during training, one can reduce the cumulative trajectory term over time, thereby tightening the overall generalization bound.
On the other hand, we examine the algorithm's 's dynamics as it approaches convergence in the late stages of training. As the parameters approach the minimum, the loss and the full gradient the training set diminish, causing the gradient the SGN covariance to dominate the drift term. Near a local minimum, it is well established (Jastrzebski al. 2017 Daneshmand al. 2018: Zhu et al.) 2018
Xie et al. 2020 2023) that for the negative log-likelihood loss we have
V(z) TI Σ V=( FIM(2) V+f())
(15)
where FIM(z) denotes the empirical Fisher information matrix and (2f(x) the Hessian matrix. Moreover, an empirical result by Xie al. (2020) shows that for neural networks, this relationship remains approximately valid even far away from the local minima. Therefore, in the late stages of training. regularizing the trace of the SGN covariance is approximately equivalent regularizing the
trace of the Hessian, which is widely regarded as good measure of sharpness that has been observed
to correlate strongly with generalization performance (Keskar et al. 2017 Blane al. 2020 Wen
let al. 2022 Arora et al. 2022; Damian et al.) 2022; Ahn et al. 2023 Tahmasebi et al. 2024).
To verify the dynamics around the minima, we consider the setting from the work of Liu al. (2020); ;
Damian et al. (202T), where initialization is performed at a bad minimum. We use the checkpoint provided by Damiin et al. (202I). The learning rate is set 1c-3, under which SGD has been shown to struggle to escape poor minima. We do not employ any explicit regularization techniques. For SAM, we set p 5e-3 and compare the performance of m-SAM with different values of m. In Figures [ and 2 we can clearly observe that as m decreases, the regularization effect on the SGN covariace is significantly strengthened, which in turn accelerates the escape from poor consistent with our theoretical analysis.
4 Practical Method
In this section, we will demonstrate how to translate the theoretical insights gained from our SDE approximations into practical method. While m-SAM" Jachivess substantially better generalization than mini-batch SAM by strengthening the regularization of the SGN magnitude, its inherently se- quential nature creates serious parallelization bottleneck (see Table 7 in Appendix forr performance and time cost). This is because it must sequentially compute the perturbation for each micro-batch
In the multi multi-GPU setting, we use "m-SAM" to denote the scenario where < per-device batch size, which inherently not parallelizable well.

---

0.60
12.0
m=g
11.
Test Accuracy
11.6
0.54
m-44
of
11.2
0.50
0.48
0
10
20
30
40
50
20
30
40
50
Epoch
Epoch
Figure 1: Speed of escaping poor minima, mea- sured by test accuracy.
Figure 2: Variance of SGN over iterations.
and then individually backpropagate to obtain the perturbed gradient. Building on theoretical insights, a natural question arises:
Can we design a parallelizable algorithm that preserves the generalization advantages of m-SAM? One important observation is, if we further assume that the SGN approximately orthogonal to the full gradient, or that the norm of the full gradient is negligible compared to that of the SGN (as often occurs in the late stages of training when the model approaches convergence), which commonly arises in signal-to wew-ttoooooo--uitee jet al., 2023 Allen-ZZh and Li 2023 Han et al. 2024] Huang et al. 2024,] Li et al. 2024a), then from the following equation observe that the norm of the stochastic gradient is directly related to the norm the SGN:
analyses (Cao et al., 2022 Delassi et al. 2022 Zou et al. [223]] Huang
[+)())) VIV 11&.112.
(16)
This decomposition reveals that samples with larger gradient norms carry gigher-mannitude SGN. Drawing inspiration from importance sampling, which assigns a weight to each sample proportional to its "importance",
adapt this idea emphasize more important samples when computing SAM's perturbation (2) In our framework, we quantify importance by the magnitude of the SGN: samples exhibiting larger SGN contribute weight in the perturbation. Specifically, we introduce an adaptive weighting mechanism in which each weight y; reflects the importance of sample i when computing the perturbation vector. This reweighting strategy defines probability distribution P = {p: over the sampled indices i E thereby modulating the influence of each sample. We refer to the resulting algorithm Reweighted SAM (RW-SAM). The objective of RW-SAM's perturbation is formulated as follows:
sharpness regularization and therefore deserve greater
P4 V. H(P) X
max PEA
max [el] SI
(17)
where A denotes the probability simplex, the entropy function used to prevent all weights from concentrating on single sample, and À is hyperparameter maintain balance between emphasis and diversity. Notably, as X Objective (17) degenerates to mini-batch SAM's perturbation (2)
Solving Objective
17 for yields:
Pi Vz(z))
(18)
EI Pi f((x))
Since Objective [7
does not have a closed-form solution for P, we propose solving its relaxation:
H((P)
max
Σ Pil [p()\
(19)
PEA
EET
Objective (99 corresponds to the well-known Gibbs distribution (see derivation in Section H which is given by:
exp(A [+)(z)))
p
(20)
exp (V
7

---

Broadly speaking, Eq. (20) corresponds to assigning higher weights to samples with larger stochastic gradient norms, where X controls the concentration of this distribution. practice, the optimal value of À depends on the scale of per-sample gradient norms. Therefore, we normalize the estimated gradient norms before applying the exponential function, making the algorithm's 's performance less sensitive to the choice of X. As observed in Section 5.4, normalization significantly reduces the need for tuning A.
Additionally, we propose using finite-difference method combined with Monte Carlo sampling to avoid per-sample gradient-norm estimation via backpropagation. The formula is as follows (see derivation in Appendix|Gr
1
fi(x + 82g) Sí(x)
[V) ~
(21)
6
where 2 E IRd is a Rademacher random vector, 5 is a small constant. This estimator requires only Q additional forward passes, as we conveniently obtain the loss for each sample in single forward pass, making it an efficient way to estimate the gradient norm for each sample. To minimize additional computational overhead, we set Q 1 in our experiments, consistent with common practice in deep learning (Kingma et al. 2013: Gal and Ghahramani 2216; Ho et al. 2020 Malladi et al. 2023). and found that this choice suffices to achieve a non-trivial performance improvement. However, unlike common implementations (Kingma et al. 2013 Ho et al.) 2020 Malladi et al.) (023). we propose using Rademacher instead of Gaussian perturbations to minimize the variance of the Monte Carlo estimator. Since Rademacher variables have a fixed expected squared norm, they achieve the optimal variance (according to Theorem 2.2 in Ma and Huang (2025). The pseudocode for our algorithm is presented in Algorithm T (see Appendix lb.
5
Experiments
5.1 Training from scratch
We evaluate three optimization methods: SGD, SAM and RW-SAM CIFAR-10 and CIFAR-100
(Krizhevsky et al. 2009), training three models from scratch: ResNet-18, ResNet-50 (He et al., 2016). and WideResNet-22 28- 10 (Zagoruyko and Komodakis 2016). We use batch size of 128 and cosine learning rate schedule with an initial learning rate of O.I. SAM and RW-SAM are trained for 200 epochs, while SGD is trained for 400 epochs. We apply momentum of 0.9 and a weight decay of 5e-4, along with standard data augmentation techniques, including horizontal flipping, padding by four pixels, and random cropping.
For SAM and RW-SAM, we set = 0.05 for CIFAR-10 and p = 0.1 for CIFAR-I00, In the case of RW-SAM, we determine â through finite-difference estimation on the model before training, based on the estimated error. Specifically, we consider 8 E (le-5, 1e-4, 1e-3) and use 8 = le-3 for ResNet-18 and ó le-4 for both ResNer-50 and WideResNet-28-100 For the additional hyperparameter À in RW-SAM, we performed grid search over (0.25.0.4 0.5, 1.0, 2.0} on a validation set and found that 0.5 consistently yielded strong performance across experiments.
Table 1: Test accuracy comparison on CIFAR-10 and CIFAR-100 with different optimizers.
Model
CIFAR-10
CIFAR-100
SGD
SAM
RW-SAM
SGD
SAM
RW-SAM
ResNet-18
95.62 = 0.03
95.99 + 0.07
96.24 + 0.05
78.91 + 0.18
78.90 = 0.27
79.31 + 0.28
ResNet-50
95.64 + 0.37
96.06 + 0.04
96.34 + 0.04
79.55 + 0.16
80.31 + 0.35
80.83 + 0.05
WideResNet
96.47 + 0.03
96.91 + 0.02
97.11 + 0.05
81.55 0.15
83.25
0.07
83.52 0.08
For large-scale experiments, we train ResNet-50 on ImageNet-1K (Deng al. 2009) for 90 epochs with an initial learning rate of 0.05. For both SAM and RW-SAM, we use the same p = 0.05. For the additional hyperparameter À in RW-SAM, set À 0.25. All other hyperparameters remain the same as those used on CIFAR-10/100.
In this section, consistent with common practice, use "SAM" refer the mini-batch batch SAM.
8

---

We repeated three independent experiments and reported the mean and standard deviation of the test accuracy in Table and Table 2a We observe that RW-SAM consistently outperforms the baselines across various models, as well as on both small and large datasets.
Analysis of Computational Overhead. RW-SAM requires additional forward pass estimate per-sample gradient norms, leading to approximately 1/6 more training overhead compared to vanilla SAM, forward pass typically takes about half the time of backward pass (Kaplan) 2022). For a report of the additional wall-clock time overhead, please see Table g in Appendix U However,
RW-SAM matches the performance of m-SAM m = 64 without incurring its nearly two-fold training overhead (See Table in Appendix D) This highlights the efficiency of RW-SAM in balancing computational cost and performance.
Table 2: (a) Test accuracy on ImageNet- KK with different optimizers; (b) Test accuracy fine-tuning ViT-B/16 on CIFAR-10/100 with different optimizers.
SGD
SAM
RW-SAM
Iet-INet-IK
SGD
SAM
RW-SAM
CIFAR-10
98.24
0.05
98.40
0.02
98.58
0.02
ResNet-50
76.67 0.05
77.16
0.04
77.37 0.05
CIFAR-100
88.71 0.10
89.63
0.12
0.09
(a)
(b)
5.2 Fine-tuning
We fine-tune a ViT-B/16 model (Dosovitskiy et al.
2020), pre-trained on ImageNet-IK, on CIFAR-10
and CIFAR-I00. We train for 20 epochs with iniiill learning rate of 0.01. Other hyperparameters are the same as those in training from scratch. The results are summarized in Table (2] We also fine-
tune pretrained DistilBERT model (Sanh al. (019) on the GLUE benchmark (Wang et al. 2018). We use AdamW (Loshchilov and Hutter (2199) as the base optimizer. The detailed hyperparameter
settings are provided in Table of Appendix and the results are presented Table We observe
that RW-SAM consistently outperforms SAM in both fine-tuning experiments.
Table 3: Performance comparison on GLUE tasks using different optimizers.
Optimizer
CoLA
MNLI
MRPC
QNLI
QQP
RTE
SST-2
STS-B
Average
AdamW
0.538
0.825
0.900
0.884
0.868
0.628
0.914
0.864
0.804
SAM
0.517
0.824
0.900
0.894
0.871
0.610
0.911
0.870
0.800
RW-SAM
0.560
0.826
0.904
0.896
0.872
0.625
0.915
0.870
0.809
5.3
Robustness to label noise
Foret et al. (2021) have shown that SAM exhibits robustness to label noise. Motivated by this, we evaluate RW-SAM's performance on CIFAR-10 with labels randomly flipped at specified noise ratios. We train ResNet-16 88 and report clean test accuracies in Table As the noise ratio increases, RW-SAM maintains remarkably strong performance. In particular, at an 80% noise ratio, RW-SAM achieves 16% absolute accuracy improvement over SAM.
5.4 Additional experiments
Hyperparameter sensitivity. We train ResNet- on CIFAR- 100 using RW-SAM to evaluate its sensitivity to different values of X. As summarized in Ta RW-SAM demonstrates robustness to the choice of À and consistently outperforms the baselines within a reasonable range.
Trace of stochastic gradient covariance. According to our theory, we compare the trace of the stochastic gradient covariance matrix at convergence for different algorithms, which, near the minimum, closely approximates the Hessian matrix (See Eq. (15)). The results, shown in Table 6 indicate that compared to SGD and SAM, RW-SAM indeed converges to minimum with smaller stochastic gradient covariance magnitude.

---

Table 4: Performance comparison CIFAR-
Table 5: RW-SAM X-sensitivity
10 across different noise ratios
0.25
0.5
1.0
200
Noise Ratio
SGD
SAM
RW-SAM
79.09 1 0.1
79.31
0.288
79.12 1 0.33
79.03
0.
20%
87.54 1: 0.20
90.01
0.09
90.34 + 0.20
40%
86.40
Table 6: Trace of the gradient covariance
60%
76.64 +: 0.32
78.79
0.24
81.52 + 0.31
80%
46,53 + 1.13
37.69 3.12
53.17 = 3.70
Optimizer
Trace
SGD
572.39 24.15
SAM RW-SAM
198.40 = 6.20
177.79 5.10
Conclusion
In this work, conducted comprehensive theoretical analysis of SAM and its variants through an enhanced SDE modeling framework. Our findings reveal that the structure of SGN plays a crucial role in implicit regularization, significantly influencing generalization performance. Based analysis, we proposed Reweighted SAM, an adaptive weighting mechanism for perturbation, which empirically validated through extensive experiments. Our study provides deeper understanding of the dynamics of SAM-based algorithms and offers new perspectives on improving their generalization performance.
Acknowledgement
Trung Le, Mehrtash Harandi, and Dinh Phung were supported by the ARC Discovery Project grants DP23010117 176 and DP250100262. Trung Le and Mehrtash Harandi were also supported by the Air Force Office of Scientific Research under award number TA9550-22-S-0001.
References
Agarwala, A. and Dauphin, Y. (2023). Sam operates from home: eigenvalue regularization as dynamical phenomenon. In International Conference Machine Learning, pages 152-168. PMLR.
Ahn, K., Jadbabaie, A., and Sra, S. (2023). How escape sharp minima with random perturbations. arXiv preprint urXiv:2305. 15659.
Allen-Zhu, Z. and Li, Y. (2023). Towards understanding ensemble, knowledge distillation and self-distillation in deep learning. In International Conference on Learning Representations (ICLR).
Andriushchenko, 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. International Conference on Muchine Learning, pages 948-1024. PMLR.
Behdin, K., Song, Q. Gupta, A., Durfee, D., Acharya, A., Keerthi, S., and Mazumder, R. (2022). Improved deep neural network generalization using m-sharpness-aware minimization. arXiv preprint arXiv:22/2.0434.
Blane, 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.
Cao, Y., Chen, Z.,, Belkin, M., and Gu, Q. (2022). Benign overfitting in two-layer convolutional neural networks. Advances Neural Information Processing Systems, volume 35, pages 25237-25250.
Compagnoni, E. M., Biggio, L.,, Orvieto, A., Proske, F. N., Kersting, H., and Lucchi, A. (2023). An sde for modeling sam: Theory and insights. International Conference on Machine Learning, PMLR.
pages
10

---

Dai, Y., Ahn, K., and Sra, S. (2024). The crucial role of normalization in sharpness- awaree aware minimiza- tion, Advances in Neural Information Processing Systems, 36.
Damian, A., Ma, T., and Lee, J. (2021). Label noise sgd provably prefers flat global minimizers. Advances in Neural Information Processing Systems, 34:27449-27461.
Damian, A., Nichani, E., and Lee, J. D. (2022). Self-stabilization:: The implicit bias of gradient descent at the edge of stability, arXiv preprint arXiv: 15594.
Daneshmand, H., Kohler, J., Lucchi, A., and Hofmann, T. (2018). Escaping saddles with stochastic gradients. In International Conference Machine Learning, pages 1155-1164. PMLR.
Deng, J., Dong. W., Socher, R., Li. L.-J., Li, K., and Fei-Fei, L. (2009). Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248 255. lece.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenbom, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al. (2020). An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprini arXiv:2010.11 201...1999.9
Du, J., Yan, H., Feng, J., Zhou, J. T., Zhen, L.., Goh, R. S. M., and Tan, V. Y. (2021). Eff- cient sharpness-awars minimization for improved training of neural networks. arXiv preprint 10.031..
Du, J., Zhou, D., Feng. J., Tan, V., and Zhou, J. T. (2022). Sharpness-aware training for free. Advances in Neural Information Processing Systems, 35:23439-23455 -34511.
Evans, L. C. (2012). An introduction to stochastic differential equations, volume 82. American Mathematical Soc.
Folland, G. (2005). Higher-order derivatives and taylor's formula in several variables. Preprint, pages 1-4.
Foret, P., Kleiner, A., Mobahi, H., and Neyshabur, B. (2021). Sharpness-awart minimization for efficiently improving generalization. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OvenReview.net.
Gal, Y. and Ghahramani, Z. (2016). Dropout as bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learming. pages 1050-1059. PMLR.
Han, A., Huang. W., Cao, Y., and Zou, D. (2024). On the feature learning in diffusion models. arXiv preprint urXiv:2412.0102..
HaoChen, J. Z., Wei, C., Lee, J., and Ma, T. (2021). Shape matters: Understanding the implicit bias of the noise covariance. In Conference on Learning Theory, pages 2315-2357. PMLR.
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.
Ho, J., Jain, A., and Abbeel, P. (2020). Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851.
Hoang-Anh. D., Le, P. T., Cai, J., and Do, T.-T. (2025). Sharpness- data generation for zero-shot quantization. arXiv preprint arXiv:2510.07018.
Huang, W., Han, A., Chen, Y., Cao, Y., Xu, Z., and Suzuki, T. (2024). On the comparison between multi-modal and single-modal contrastive learning. arXiv preprint MxXiv:241..02877
Huang, W. Shi, Y., Cai, Z., and Suzuki, T. (2023). Understanding convergence and generalization in federated learning through feature learning theory. In The Twelfth International Conference on Learning Representations (ICLR).
Jastrzebski, S., Kenton, 7., Arpit, D., Ballas, N., Fischer, A., Bengio, Y., and Storkey, A. (2017). Three factors influencing minima in sgd. arXiv preprint 4rXiv:1711.04623.

---

Jelassi, S., Sander, M., and Li, Y. (2022). Vision transformers provably learn spatial structure. In Advances in Neural Information Processing Systems, volume 35, pages 37822-37836.
Kaddour, J., Liu, L., Silva, R., and Kusner, M. (2022). When do flat minima optimizers work? Advances in Neural Information Processing Systems, 35: 16577-16555
Kaplan, J. (2022). Notes contemporary machine learning for physicists. Lecture Notes, Department of Physies and Astronomy, Johnss Hopkins University. Available Online https://sites. krieg er. jhu. mpora ryMLf orPhy sicis is. pdf.
Keskar, S. Mudigere, D., Nocedal, J., Smelyanskiy, M., and Tang, P. P. (2017). On large-batch training for deep learning: Generalization gap and sharp minima.
Khanh, P., Luong, H.-C., Mordukhovich. B., and Tran, D. (2024). Fundamental convergence
analysis of sharpness-awann minimization. Advances in Neural Information Processing Systems, 37:13149-3312.
Kim, M., Li, D., Hu, S. X., and Hospedales, T. (2022). Fisher sam: Information geometry and sharpness aware minimisation. In International Conference on Machine Learning, pages 11 148- 11161. PMLR.
Kingma, D. P., Welling, M., et al. (2013). Auto-encoding variational bayes.
Krizhevsky, A., Hinton, G., al. (2009). Learning multiple layers of features from tiny images.
Kwon, J., Kim, J., Park, H., and Choi, I. K. (2021). Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks. In International Conference on Machine Learning, pages 5905-5914. PMLR.
Li, B. and Giannakis, G. (2024). Enhancing sharpness- optimization through variance suppres- sion. Advances in Neural Information Processing Systems, 36.
Li, B., Huang, W., Han, A., Zhou, Z., Suzuki, T., Zhu, J., and Chen, J. (2024a). On the optimiza- tion and generalization of two-layer transformers with sign gradient descent. arXiv preprint arXiv:2410.04870
Li, B., Zhang, Y., and Giannakis, G. B. (2025). Vasso: Variance suppression for sharpness-aware minimization.
Li, Q. Tai, C., et al. (2019). Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations. Journal of Machine Learning Research, 20(40):1-47
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, T., Zhou, P., He, Z., Cheng, X., and Huang, X. (2024b). Friendly sharpness- minimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 5631-5640.
Li, T., Zhou, T., and Bilmes, J. (2024c). Reweighting local mimina with tilted sam.
Li, Z., 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-12725.
Liu, S., Papailiopoolos, D., and Achlioptas, D. (2020). Bad global minima exist and sgd can reach them. Advances in Neural Information Processing Systems, 33:8543-8552.
Liu, Y., Mai, S., Chen, X., Hsieh, C.-J,, and You, (2022a). Towards efficient and scalable sharpness- aware minimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages
Liu, Y., Mai, S., Cheng, M., Chen, X., Hsieh, C.-J., and You, Y. (2022b). Random sharpness-awee minimization. Advances in neural information processing systems, 35:24543-245555
Loshchilov, L. and Hutter, F. (2019). Decoupled weight decay regularization.
12

---

Luo, H., Truong, T., Pham, T., Harandi, M., Phung, D., and Le, T. (2025). Explicit eigenvalue regularization improves sharpness-ave aware minimization. In The Thirty- eighth Annual Conference on Neural Information Processing Systems.
Ma, S. and Huang, H. (2025). Revisiting zeroth-order optimization: Minimum-variance two-point estimators and directionally aligned perturbations. In The Thirteenth International Conference on Learning Representations.
Malladi, S., Gao, T., Nichani, E., Damian, A., Lee, J. D., Chen, D., and Arora, S. (2023). Fine-tuning language models with just forward passes, Advances in Neural Information Processing Systems, 36:53038-33 53075.
Mil'shtein, G. (1986). Weak approximation of solutions of systems of stochastic differential equations. Theory of Probability & Its Applications, 30(4):750-766.
Möllenhoff, T. and Khan, M. E. (2022). Sam as an optimal relaxation of bayes. arXiv preprint arXiv:2210.01662.
Mordido, G., Malviya, P., Baratin, A., and Chandar, S. (2023). Lookbehind-sam: steps back, step forward. arXiv preprint WrXiv:2307.16700
Müller, A. and Stoyan, D. (2002). Comparison methods for stochastic models and risks. (No Title).
Neu, G., Dziugaite. G. K., Haghifam, M., and Roy, M. (2021). Information -hoomatiio-thhoortiic generalization bounds for stochastic gradient descent. In Conference on Learning Theory, pages 3526-3545. PMLR.
Nguyen, V.-A., T., Bui, A., Do, T.-T., and Phung, D. (2023a). Optimal transport model distribu- tional robustness. Advances in Neural Information Processing Systems, 36:24074-24087.
Nguyen, V.-A,, Tran, Q., Truong, T., Do, T.-T., Phung, D., and Le, (2024). Agnostic sharpness- aware minimization. arXiv preprint arXiv:2406.07107
Nguyen, V.-A., Vuong, T.-.., Phan, H., Do, T.-T,, Phung, D., and Le, T. (2023b). Flat seeking bayesian neural networks. Advances in Neural Information Processing Systems, 36:30807-30820.
Oikonomou, D. and Loizou, (2025). Sharpness-aware minimization: General analysis and improved rates. arXiv preprint arXiv:2503.02225.
Phan, H., Tran, L., Tran, Q. Tran, N., Truong, T., Lei, Q., Ho, N., Phung, D., and Le, T. (2025). Beyond losses reweighting; Empowering multi-task learning via the generalization perspective. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 2440-2455
Qu, Z., Li, X., Duan, R., Liu, Y., Tang, B., and Lu, Z. (2022). Generalized federated learning via sharpness aware minimization. In International conference on machine learning, pages 18250- 18280. PMLR.
Sanh, V., Debut, L., Chaumond, J., and Wolf, T. (2019). Distilbert, distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108.
Si, D. and Yun, C. (2024). Practical sharpness-aware minimization cannot converge all the way to optima. Advances in Neural Information Processing Systems, 36.
Simsekli, U., Sagun, L., and Gurbuzbalaban, M. (2019). A tail-index analysis of stochastic gradient noise in deep neural networks. In International Conference on Machine Learning, pages 5827- 5837. PMLR.
Singh, P., Mobahi, H., Agarwala, A., and Dauphin, Y. (2025). Avoiding spurious sharpness minimization broadens applicability of sam. arXiv preprint arXiv:2502.02407..
Springer, J. M., Nagarajan, V., and Raghunathan, A. (2024). Sharpness-awar minimization enhances feature quality via balanced learning. arXiv preprint arXiv:2405.2044 20439.
Tahmasebi, B., Soleymani, A., Bahri, D., Jegelka, S., and Jaillet, P. (2024). universal class of sharpness-awand aware minimization algorithms. arXiv preprint 036602
13

---

Tan, C.. Zhang, J., Liu, J., and Gong, Y. (2024). Sharpness-awend aware lookahead for accelerating convergence and improving generalization. IEEE Transactions on Pattern Analysis and Machine Intelligence.
Truong, T., Tran, Q. Pham-Ngoc, Q. Ho, N., Phung. D., and Le, T. (2024). Improving generalization with flat hilbert bayesian inference. arXiv preprint arXiv:2410.04196.
Wang, A., Singh, A., Michael, J., Hill, F., Levy, O., and Bowman, S. R. (2018). Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv:1804.00744
Wang, Y., Zhou, K., Liu, N., Wang, Y., and Wang, X. (2024). Efficient sharpness-aware minimization for molecular graph transformer models. arXiv preprint arXiv:2406.13137.
Wang. and Mao, (2021). On the generalization of models trained with sgd: Information-theores bounds and implications. arXiv preprint arXiv:21
Wen, K., Li, Z., and Ma, T. (2023). Sharpness minimization algorithms do not only minimize sharpness to achieve better generalization. Advances in Neural Information Processing Systems, 36:1024-1035.
Wen, K., Ma. T., and Li, Z. (2022). How does sharpness-aware minimization minimize sharpness? CoRR, abs/221 1005229.
Wu, T., Luo, T., and Wunsch II, D. C. (2024). Cr-sam: Curvature regularized sharpness-aware minimization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages
Xie, W., Pethick, T., and Cevher, V. (2024). Sampa: Sharpness-aware minimization parallelized. Advances in Neural Information Processing Systems, 37:51333-51357.
Xie, Z., Sato, I., and Sugiyama, M. (2020). diffuusio theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. arXiv preprint arXiv:2002.03499.
Xie, Z., Tang, Q.-., Sun, M., and Li, P. (2023). On the overlooked structure of stochastic gradients. Advances in Neural Information Processing Systems, 36:66257-66776.
Zagoruyko, S. and Komodakis, N. (2016). Wide residual networks. arXiv preprint arXiv:1605.07146.
Zhou, Z.., Wang, M., Mao, Y., Li, B., and Yan, J. (2024). Sharpness-aware minimization efficiently selects flatter minima late in training. arXiv preprint arXiv:2410.10373.
Zhu, Z., Wu, J., Yu, B., Wu, L., and Ma, J. (2018). The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. arXiv preprint 1rXiv:1803.00195.
Zhuang, J., Gong, B., Yuan, L., Cui, Y., Adam, H., Dvorekk, N., Tatikonda, S., Duncan, J., and Liu, T. (2022). Surrogate gap minimization improves sharpness-aware training. arXiv preprint arXiv:2203.080665
Ziyin, L., Liu, K., Mori, T., and Ueda, M. (2021). Strength of minibatch noise in sgd. arXiv preprint 152.5.5
Zou, D., Cao, Y. Li, Y., and Gu, Q. (2023). Understanding the generalization adam in learning neu- ral networks with proper regularization. In International Conference on Learning Representations (ICLR).
14

---

A General theory for two-parameter weak approximation.
Let T > 17 E (0, min (.). and N [T/n]. We consider the general discrete iteration
Øk+1 = 24 + 17 hh (zk, 7k.
"k. 17 p). To E R.,
A: 0, I,
(22)
and its corresponding continuous-time approximation by the SDE
dx.- P) dt (ño aW,,
Xo mo. Ee [0,
(23)
We denote Xx = Xky, and the one-step changes:
(x) 21 -2
((x) == X1 r.
(24)
Following the SDE framework by Mil' shtein (9866; Li et al.
(017):
Compagnoni et al. 20233
Luo
al. (2025), we have the following definition:
Definition A.1. Let G denote the set of continuous functions Rd - R of at most polynomial growth, ie. g E G if there exists positive integers K1, > such that
|g(x)) S K1 (1 + ||=:|2=;))
for all æ E Rd Moreover, for each integer a 2 1 we denote by Go the set of a a-times continuously differentiable functions Fd > which, together with its partial derivatives up and including order a, belong to G.
This definition comes from the field of numerical analysis of SDEs (Mil'shtein, 1986). 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 (17 p²) to denote that there exists a function K E G independent of p, such that the error terms are bounded by Kippp²
Theorem A.2. (Adaptation of Theorem 3 in Li et (2019) Let T O, 1J E (0, min (1, T)). N = [T/n] Let a 2 1 be an integer. Suppose farther that the following conditions hold:
(i) There exists K1 E G. independent of 17, p. such that for each s = 1,2, ...... and any indices i....., [I.....d).
E
A (z) - E II
(o) j=1
s (1) + =
1=1
and
a++
EE I (x))
s Ki(x) (1) apt np&+)).
(ii) For each m 2 the 2m-moment of a& uniformly bounded and n,, iee. there exists K2 E G, independent of and k. such that
E E||02|1 2m S K2(x), == O.O.,...,,N
Then, for each E G*+1, there exists constant C > O, independent of Ps such that
max OSKSN
(g([z)
Eg((x)) s C)f + )+1).
By substituting Assumption (i) in the penultimate step of the proof in Theorem 3 in Li et al. (2019) with Assumption (i) of our Theorem A.2 the same argument goes through and yields the desired conclusion. Hence, we omit the proof.
Next, we state the key lemma for the two-parameter weak approximation. By employing a Dynkin expansion (see, e.g. [vans (2012)) instead of a full Itô-Taylor expansion, it avoids the proliferation of terms and provides streamlined approach to controlling the remainder error in the two-parameter setting.
15

---

Lemma A.3 TTo--aaaaetter Dynkin (semigroup) expansion). Let 2b G2 and suppose the drift admits the expansion
B b(x P) Σ
pT pm (x) + O(p
a(z)=oofe
m=0
Define
Am (x)x = 6122 (x)(()) (m = 0......... As {a))))= = [ (i5) 0g (x))
Suppose further that bm (m = 0, 1,,,... 9)..,,9,0000 E G2a, then for any nonnegative integers a, B
8
E [[x)),,]]]
Pm
E
Ae, Ar, ... (e) + O(1 + O(pp
n!
n=
M=U
IE 10
33
Here [e] = Ei ((. Ci and each multi-index & = (fi...., .n) contributes factor pla, and
indices & = 0 or E = A contribute no p-power (via Au or Aa).
Proof. Let
L ((x) (b(x,)) p) (x) + & [000]](00 07
be the infinitesimal generator of the diffusion. Since
b(x, p) = pm bm (z) + 0(p
m==
may write
L Σ pmh Am + As O(p+t),
where Am and As are as in the statement. By Dynkin's formula (or equivalently the semigroup expansion),
E [(()),)) = (6)L
(x) (x)
> Σ 1 n!
(')(s) + 0(1at1).
==
remains expand each power L", By the multinomial theorem,
L. nm Am + As + 0(p
Σ
Σ LE(..A.......0f [d
A Arz Ac, + 0(p+1)).
m=0
m=0
Hence
17th
17t (2((x)= m!
Σ
Ar, -. (c,((x) + O(1*+1) + 0l)
n!
m=0
EE(O,A.......9 BY" 11|=m
Summing over 0, 1, and collecting the remainders 0(1 and O(pp++1) yields exactly the claimed two-parameter expansion.
Since our focus in this paper is the order-(1,1) weak approximation, we now present the one-step approximation lemma for SDEs the case a = B I, as follows. For readers interested in higher-order wooparraeter weak approximations, it is sufficient to apply higher-order truncations of the Dynkin and Taylor expansions in the two lemmas below and then match the corresponding moments at each order.
Lemma A.4 (One-step moment estimates up to 1//, ,p for SDEs). Suppose the drift admits the expansion
b(x, = bu(x) + (x) (z)+0(p²))
(x) = oo(x),
and assume by, b1, 0o E G². Let ((x) be the one-step increment defined in 22 Then:
16

---

(0
- = pb. +0(1)) + O(1p²).
(i)) ((())))) (1) 1/ 1/ (z) (x) + Σ 0 (1.8)
(x) (z) (0)))())x))) (1)((b000))
(2) (x) (x)
bz)) 0(12p²) 0(1²).
(ii)) E [T = 0
Proof. For each 2, and any choice of indices i ,,,,.,, iss define the test function
(x(+) II (E)) 2(5)))
Since v. E (RRd) with at most polynomial growth, we may invoke Lemma A.2 with truncation orders a and 1. This yields.
[x)) = +
Σ RE(O,A)
Arb(z) +
(i) First moment. Here
(1(2) = 2()) - I(()
Hence
= = + (x)
E
(z) + 0(7²) + O
proving (i).
(ii) Second moment. Now
22(2)) (Z(0) I(3) x(1).
follows that
E (z) ( 11 Σ
(i,b) (i,b)]
+1 o(
p12pp(6000 110000000000000000000000 10 +000)))))))) + 0
=
proving (ii).
(iii) Third moment. Finally,
(b)(2)
II (z0) - m(,)),
1=1
and since each nonzero term in the expansion has total order n 2 3, Lemma A.2 gives
= [ = [N++(()++))] O(1²),
establishing (iii). This completes the proof.
Similar to the continuous-time setting, require the following one-step error lemma for the discrete algorithm in the case a
1:
Lemma A.S (One-step moment estimates up 1/=,p² for the discrete algorithm). Suppose the discrete update 22 admits the expansion
h(z, 7) ho(a, 7) + (2, + +(p))
and assume h h1 G2. Let ((x) be the one-step increment defined in 24 Then:
(i) [T) = = (x) + do h (x) O
(ii) ([) (11) (x) + +2 (z)) + O
(1
(4) (z) h (I) + 20.0 (x)) - n2p((00 (z) (2) (x)
h (x) 4
(z) (<)
17

---

(ii) (1 O(n").
where ha(x)) (2 7), 1o (1) Bhi (x, Eoo(x) ho), 7)), (x,
Cov(ho( 7), h1 (xx
Proof. Recall that
(x)) = (((.7))))
Hence for each coordinate i,
(i) First moment.
B =
R) (x, (z,7) (p.))())))) 0(12)
= n (z) ph) + O(T
(ii) Second moment.
E (()(((z) (z) = En (100 (I, (z, + ph
+ +11200
5 (I, (z,>)) + p (E, (z,7)
+ (M,
(x) (x) (1))
(1)) h (x) - (x) =
(iii) Third moment. Since ho, hn E G2 implies that all moments up to order three finite and On)), we have
Am = 0
This completes the proof.
B
SDE approximation for USAM variants
Recall that the update rules of USAM variants are defined by:
mini-batch USAM: Tk+1 = Ik V (Ix + f (xx))
(25)
n-USAM:
T++1
Tk
Viy (Ix + p ((()))
(26)
m-USAM:
Tk+1 = I4
q m
Σ
Vlz + P
(27)
hr
In this section, we impose the following growth assumption on the functions f and fy: Assumption B.1. The functions ff and fi belong the class G
B.1
Mini-batch USAM
For the mitc- USAM algorithm 25 we define the continuous-time approximation X, as solution to the following SDE:
dx, -V--VUUUA USAM (Xt) dt + V LUSAM dWe
(28)
where
JUSAM (X₂) = f(x.) |F(x))11 2/>l
(x,))
18

---

EUSAM (Xi) == (x) + [X,
(29)
E0.0 (V V (V V.(()))
Eo,1 B[( ( - (F
Theorem B.2 (mini-batch USAM SDE, adapted from Theorem 3.2 of @Compagnoni et al. (2023) )
Under Assumptions 3.1 and 1.1 ler 0 < 17 < 1, T O, and N = [T [T/ Denote by { 1k= the
mini-batch USAM iterates in (6. and let (0, be the solution of the SDE 28 Suppose:
(i) The functions
USAM
P
(r(V))
and EUSAM
2/7 2/1
are Lipschitz on Rd.
(ii) The mapping
1x, = = -V
satisfies, almost surely, the Lipschitz condition
(x) (vy(w)) I Ly || vll V2,yER"
where L, > 0 a.s. and [[L < xo for every 1.
Then (X, t [0, is an order-(1, weak approximation of {Ik}, namely: for each g E G2, there exists a constant C' 0, independent of 7,, p. such that
[[g(z=)] B[[((Xn)]] S (n
max OSKSN
Lemmas A.1, A.2 and A.5 in Compagnoni al. 2023) with two-parameter Resionn--Thoorem A.2 Lemma A.4 and
Proof Sketch. Theorem B. follows by replacing the
Lemma imposing the extra global Lipschitz conditions to guarantee existence and uniqueness of the strong solution, and using Assumption 3.T expand the drift term. We therefore omit the routine algebraic details.
B.2 n-USAM
For the n- U-US algorithm, define the continuous-timg upproximation X, as the solution to the following SDE:
USAM (X.) dt + / V. (Xe) dW,,
dX, -V -V mn.USAM
(30)
where
PLUSAM (X,) = f(X.) + (((())))
yMMUSAM (X2) = (X,) + + + (X.)).
(31)
En.o (x) = (V - V/(X.)) VS(X.)) ]
(x, EE vx(x)) -VP-))))) -(X "].
We begin by deriving, via the following lemma, a one-step error estimate for the n-USAM discrete algorithm, which will be used to prove the main approximation theorem.
Lemma B.3 (One-step moment estimates for n-USAM up to 1 pt)). Under Assumptions 3. and
B.l
Define
0, f.-USAM = - + p p (+)00,,
3
Let A(x) be the one-step increment defined in 24
Then:
(i) [[A()())] = -n (n.AMM))((())) (
19

---

(i) ([)
((.f()));f(1) +
y (z) (i5)
('((as(())) Σ +1100)0000+++0000000 ( + O
(iii P [+)())) - 0
1=1
Proof. Recall that the n-USAM update is
17 Sm (I + V f(x)))
Tk+1 2k
so the one-step increment
(x) = Ik+1 n,( 7,p),
where we define
7, := - (
By Taylor's theorem with integral remainder (Folland 2005 we have, for each 7
V[((+)VV(())) =
V + R(2,7,))
where
R(x, = f (--
fV())] dt.
Here D (v) denotes the third-order tensor of partial derivatives of f, at and D Dd (u) v) its bilinear action on vectors u, v.
Because fy E GO, there exists a polynomially bounded function K ((x) E G such that
|ID' < K(x), Vy with |1y - x|| f=)]l
Hence
[R( p)ll < / (L=o)|)*(((++)OV((()[]]]][(]
s K(x) p² 2 [
uniformly in r. Accordingly. a Taylor expansion in gives
(f (+)
Hence
h(z, 7,P) ho( + phi(o,7) + O(p²).
with
ho(z, == (x),
(f(())
By Assumption 3.1 and Assumption B.1 each ho, hi E G2,,nnd
E[ho(æ, 7)] -z)( [[h v) 22 V(())
We may therefore apply Lemma A.5 with these bog h1, which yields exactly the three moment expansions up to 17 p'
In Lemma B. we derived one-step moment estimates for the n-USAM discrete algorithm and, via Lemma A.4 for corresponding SDE update (26). These estimates demonstrate that the first- and second-order moments satisfy the matching conditions of Theorem A.2 Together with the uniform
moment bounds from Lemma D.2 we are now ready to establish the main weak-approximation theorem for n-USAM.
Theorem B.4 (n-USAM SDE). Under Assumptions 3.1 and B.1 let 0 < < 0, and N T//) Denote by (*x) N N the n- n-USAM iterates in and (Xt)tejo,ry be the solution of
the SDE 30
Suppose:
20

---

(i) The functions
fT-UUAM
((5+11V517) and VEE-USAM
are Lipschitz on IR
(i) The mapping
(1(2)) -V
satisfies, almost surely, the Lipschitz condition
I (a) Vh,(u)) I L, 1| vll. VE,EER,
where Ly > 0 a.s. and B[L < xo for every m 2 1.
Then (X+ I. E }) is an order-(1, weak approximation of (rk), namely: for each g E G2, there exists constant C O, independent of 7), p. such that
[[((ax)) - E[g((xn)] & c()+
max OSASN
Proof. First, verify that SDE 30 admits unique strong solution. By assumption, both the drift and diffusion coefficients are globally Lipschitz, which turn implies linear-growth condition. Therefore, Theorem D.I applies and yields the existence and uniqueness of strong solution [0.
Then, by Lemmas A.4 complete.
B.3 and
D.2 all the conditions of Theorem A.2 are satisfied, and the proof is
Remark B.5. The Lipschitz conditions to ensure that the SDE has unique strong solution with uniformly bounded moments. It is possible to appropriately relax them we allow weak solutions Miil'shtein 1986).
B.3 m-USAM
For the n-USAM algorithm, we define the continuous-time approximation X, as the solution the following SDE:
dx, (f(((X)) ? [IVF(X)))
tr(V(X.)))
(X))dW,
(32)
2m
h
where
E. EM-USAM (00AMM(X):00) = Ex,o(X,) Fo, (x,)
(33)
(X,) = (V,z(X,) V((x)) vf(X,))
Ex, = E Vx(X)) (XX EV
We begin by deriving, via the following lemma, one-step error estimate for the m- m-USAM discrete algorithm, which will used to prove the main approximation theorem.
Lemma B.6 (One-step moment estimates for m-USAM up to r7', p²)) Under Assumptions 3.1 and
B.
Define
0 m (x) = 0,f(x pE Σ (0
Let A(x) be the one-step increment defined in 124). Then:
(i) [[A(][))]] = - a (m)
(ii) ((A (()())))()(()))
Im.USAA(()))
(2))
)²p((a.(z)) Σ 01) 0,f(x) Σ (1 0(12p3).
21

---

3
(ii) I P = 0
=
Proof. Recall that the m-USAM update
17
Σ
(z + p VV (Ix))
1k+1
hl I, 1,
so the one-step increment
(Φ) Ik+1 - -gk = =1,1(,
where we define
h(z, P)
Σ
Viz, (Ix + p V fI, (2x)).
I,C%,.Z,//=m
By Taylor's theorem with integral remainder (Folland 2005 we have, for each 7
VZz + VZz
Σ C mm
(1x) R(.,1.))
I,
7y/=m
where
R(2,7 RB000000111DDDP//4400000)00101)0 VVr (x)] dt.
Here Sz (v) denotes the third-order tensor of partial derivatives of fT, at y. and Dd SI, (V)(u, its bilinear action on vectors u, v.
Because JI, G3, there exists a polynomially bounded function K(x) E G such that
D (x,(1l)) S K(x), Vy with by xll S p V fr, (z)ll-
Hence
[( & ( ('(1-t) (I' tpV |/ dt s K(x) p² 2 (V
0
uniformly in y. Accordingly, Taylor expansion in p gives
V(((+))((()) VV +
Hence
h(z,7,p) ho( + phi(o,7) + O(p²).
with
ho(z, 7) =
m
fr, (xx); := = m
hl
Ctkk
I, T,C 1
Tyl=m
By Assumption 3.1 and Assumption B.1 each ho, h1 E G2, and
E[ho(x,>)] -v) Elhr(x, = - E fz, (x) (z)
We may therefore apply Lemma A.5 with these ho, h1s which yields exactly the three moment expansions up to ppt
In Lemma B.6,
we derived one-step moment estimates for the m-USAM discrete algorithm and, via
Lemma A.4 for corresponding SDE update (27). These estimates demonstrate that the first- and
second-order moments satisfy the matching conditions of Theorem A.2 Together with the uniform
moment bounds from Lemma D.2 we are now ready to establish the main weak-approximation theorem for m- m-USAM.
22

---

Theorem B.7 (m-USAM SDE). Under Assumptions 3.1 and B.1 let 0 < <7 <π< 0, and N n Denote by {ax N k=o the m-USAM iterates and let {Xt]xe(0,T] be the solution of
the SDE 32
Suppose:
(i) The functions
fu.
v(f 2 -
1 +
tr(()) 2m
and Vm
are Lipschitz on RR..
(ii) The mapping
(x) hl I,C,,
m
Σ
satisfies, almost surely, the Lipschitz condition
|V (a) Vh,(w)l & Ly -+11~1111 vll
where L, > 0 a.s. and B[L < 3o for every m 2 1.
Then (X, 1 [0, is an order-(1, 1) weak approximation of (rk), namely: for each g E G2, there exists constant C 0, independent of 1), ps such that
[[g(xx)) E[o(Xen)] s c((++
max OSRSN
Proof. First, we verify that SDE 32 admits a unique strong solution. By assumption, both the drift and diffusion coefficients globally Lipschitz, which in turn implies linear-growit condition. Therefore, Theorem D.I applies and yields the existence and uniqueness a strong solution on [0,
Then, by Lemmas A.4 B.6 and complete.
D.2 all the conditions of Theorem A.2 are satisfied, and the proof
Remark B.8. The Lipschitz conditions are to ensure that the SDE has unique strong solution with uniformly bounded moments. It is possible appropriately relax them if we allow weak solutions (Mil'shtein) 1986).
C SDE approximation for SAM variants
Recall that the update rules of SAM variants defined by:
V/n (1rk)
mini-batch SAM: T6+1 = V VSJa (xx +
(34)
V fr (xxx)l
V((x))
n-SAM:
Fk+1 = 2k
V (x +
[
(35)
V fzz, (Ix)
S-SAM:
ZE+1 =
m
Σ
Vzz, (Ix +
(36)
V [I*)
T, CT mm
C.1
Mini-batch SAM
For the mini-batch SAM algorithm 34 we define the continuous time approximation X, as the solution to the following SDE:
dx, -V- SAM (X.) dt + V E.SM (...)
SVSSAM
(37)
where
SAM (X,) = f [X
SAM (Xi) = (x. ( + Eo1 (X:)).
(38)
23

---

(X,)(X,) E (V) V.(X)) (((()) vx(())) ' (x,1 := := -V (X
E V (XX)VXXX "'
(V
(x,)l
Remark C.I (On normalization at critical points). Note that SAM is ill-defined when the gradient is zero. To solve this, we may replace the denominator by || lle VII- 11 + e² with fixed E > 0, which also a common implementation in practice. Under the standing assumption that Vf is and C All local Taylor or
L-Lipschitz, the resulting coefficients are globally O(L/s)-Lipsche weak-approx arguments and moment bounds in Appendix[[continuu [Continue to hold with constants depending on but independent of 7/ p. The stated orders in 7) P are unaffected.
Theorem C.2 (mini-batch SAM SDE, adapted from Theorem 3.5 of Compagnoni et al. (202)
Under Assumptions 3./ and B.1 ler 0 < fj < T > O, and N = [T/1]. Denote by [== the mini-batch SAM iterates in 4). and let (Xt) be the solution of the SDE (37) Suppose:
(i) The functions
ff
V(f pE (
and EAA
Lipschitz IRd.
(ii) The mapping
hx(x) -Vf(++ V2((2)
satisfies, almost surely, the Lipschitz condition
(x) (n)(w) S Ly || vll VE,yER
where Ly > 0 a.s. and B[L < oo for every m 2 1.
Then (X, :t [0, is order-(1, weak approximation of {Ik), namely: for each g E G?, there exists constant C' > 0, independent of 1) such that
E[g(xx)] - B[g(XEn)]] S c
max OSKSN
Proof Sketch. Theorem C.2 follows by replacing the single-paramees Lemmas A.1, A.2 and A.14
in Compagnoni et al. 2023) with our two-parameter versions- versions-Theorem Theorem AA.2| Lemma [A.4 and Lemma imposing the extra global Lipschitz conditions to guarantee existence and uniqueness of the strong solution. We therefore omit the routine algebraic details.
C.2 n-SAM
For the n-SAM algorithm, we define the continuous-timg approximation X, the solution the following SDE:
dX, = -V (X1) dt + EN.SAM dW,,
(39)
where
Pm. (X,) = f(X,) (((X))
SAM Sn.SAM(X,)) = Eo.o P + (x.))
(40)
No (X,, = (V,,(x,)) - vx(x)) ( (x,) - vs(x))) !
(x. :- (X - VS(X)) -
S(x)))
vf(X,) VV(X,.I
)7..
We begin by deriving, via the following lemma, a one-step error estimate for the n-SAM discrete algorithm, which will be used to prove the main approximation theorem.
24

---

Lemma C.3 (One-step moment estimates for n-SAM up to n/, , Under Assumptions 3.1 and B.1 define
0,f(x)
0, (x):: + P Σ (a)
[Vf
j=1
Let A(x) be the one-step increment defined in 24
Then:
(i) [[A()(=]] S, pa O
(ii) (() (()(E)
n² (0.f(1)0;((z)
+
M (x)
v² (o.((s)) Σ (1)(()
[V{()) +zf((()
a.((x)
Σ f()))
[V}(x]] 001) p))
01f(x)
(iii) E
II A (x)( 0
j=1
Proof. The n-SAM update
Tk-1 Ik - VT (Ik + [[]((+]]]
vf(mm)
so with
A(x) 2k+1 Ia 17 h)(,,.
vf(z)
MΦ, 7, ())
Vf + pu(x)),
u(x) |Vf[()]|
By Taylor's theorem with integral remainder v f(=)) we have
(Folland 2005), for each y and writing u(x)
Vf +u(x)) (z+(2) + ( R(2,,)
where
R(x, 7. / (1-101111) (1- D2 +pt()))) p p(()]] dt.
Here (y) denotes the third-order tensor of partial derivatives of Jn at y, and D (x)(e,, w] its bilinear action vectors v, w.
Because fy E G3, there exists a polynomially bounded function K (x) E G such that
s K(x), Vy with |yy xI1 S p [|((x)]]]
Hence
|R(2,7,))] ( (1 (1- (())lt(()) 0
uniformly in Hence
M(x,7,p) = ho(x, p
with
V
ho(a, >): == -Vf(()), h(r,))=-vv911
[
By Assumptions
3.1 and B.l each ho, hr G2 and
[ho ()) -V((x),
-=) (>)())
| fV((z) [z)
The same application Lemma A.5 then yields the moment expansions (i)(iii) up to order 7/' pp as stated.
0
25

---

In Lemma we derived one-step moment estimates for the n-SAM discrete algorithm and, via Lemma A.4 corresponding SDE update (35). These estimates demonstrate that the first- and second-order moments satisfy the matching conditions of Theorem A.2 Together with the uniform moment bounds from Lemma D.2 theorem for n-SAM.
we are now ready to establish main apak-approx
Theorem C.4 (n-SAM SDE), Under Assumptions 3.1 and B. let < <
O, and N nl Denote by {rw]k=o the n-SAM iterates in and let {X1)ee(o.T) be the solution of the SDE (39) Suppose:
(i) The functions
V M ((S+PIVii)) and V.
are Lipschitz on R
(ii) The mapping
Vf(x) (x(x)) (Vf + [Vf
satisfies, almost surely, the Lipschitz condition
(x) Vhy(v)l & L, || - vll.
VE,EER'
where L, > 0 and E 3o for every m 1.
Then (X+ : [] an order-(1 1) weak approximation of (rk), namely: for each g e G2, there exists constant C. >> 0, independent of n, such that
[[g(z=)] - B[g(Xxn)] - ((
max OSKSN
Proof. First, verify that SDE 339 admits unique strong solution. By assumption, both the drift and diffusion coefficients are globally Lipschitz, which in turn implies a linear-grows condition. Therefore, Theorem D.I applies and yields the existence and uniqueness of strong solution [0,
Then, by Lemmas A.4 complete.
C.3 and D.2 all the conditions of Theorem A.2 are satisfied, and the proof is
Remark C.5. The Lipschitz conditions are to ensure that the SDE has unique strong solution with uniformly bounded moments. It is possible to appropriately relax them if we allow weak solutions
Mill shtein 1986).
C.3
m-SAM
For the m-SAM algorithm, we define the continuous-time approximation X, as the solution to the following SDE:
dx, (F((()) LEI Σ (((())))
+
mi
[n (E SAM- (X,)) (A0-(AM())))
aW.
(41)
IEI.
where
m 5m-SAM (X,) Eo,u( ((Ea.)))))))))))) (x)()))))
(42)
Ex,0(Xx) =E (VIz(X,)) - (((X)) -VXx)))))))(()))))))))))))))))))
vS(X.))
No.1 (X., = EE v(())) (xx,
Vz((X,) [ (Xx)
[ 1.
fX()
We begin by deriving, via the following lemma, one-step error estimate the m-sam discrete algorithm, which will be used to prove the main approximation theorem.
26

---

Lemma C.6 (One-step moment estimates for m m-SAM up to to ,p)), Under Assumptions 3.7 and B.1 define
0ffz(z)
d,f f.f (x) == 0.f(2) + p Σ fz((())
j=1
Let ((x) be the one-step increment defined in 24
Then:
(i) [[Am[z)] fm. (x) + O(n
(i) (1))()
Em-SAM (x))
(13)
d
0xfz(z)=
d
E Σ
8 (z)(())
[
+ 0zʃ(x)
E
(z))))
fz)(()))]
O
(iii) E II
A ( = - -
j=1
Proof. Recall that the m-SAM update is
Vfz, (Ix)
17 m
Σ
VIz JR P
Fk++
VV (rx)
I, [Z, =m
so the one-step increment
A(x) = Ik+1 - -x = =
where
Vzz(())
m
bl I, |m
V (x)
By Taylor's theorem with integral remainder (Folland 2005). for each subset index Zj.
Vzr (x)
fzz (z)
Vzz (2+
fzz, (z)
V fz, [z)
where
R(x,7,00 (1-1) D2 (x (a
1P [P
V
Vzz (x)
fzz((z)
V (w) 7 [+)] dt.
Here Dd fz, (v) is the third-order derivative tensor of fz, at y. and D fz, (v) [u, u] its action (u, v).
Since Jz, E G3, there is k(x) E G' polynomially bounded so that
|| JT, (v)ll S K(x)) whenever || - S
Viz (x) [V (2)
Hence
[R(e, 7, P)1) < f (1- K(x) lpe [للل 1 {(( p²= 0
Vi2((r) (z)
uniformly in y. Thus a Taylor expansion in gives
h(x, 7) = ho(o, 7) + phi (
with
m ho(x, 7) = h
Σ (z, (1(2,7) = m
Σ
2z)2(()
V (x) V (z)
hl
I, Cr.. II, =m
I CY., [
By Assumptions 3.1 and B.1 each ho, hu E G2, and
[[ho 7)] (xx), (a 7)] -B -B
Vz,
(z)
V (x)
Applying Lemma A.5 to these ho, hn yields exactly the three moment estimates up to 1/ pp as claimed.
0
27

---

In Lemma C.6 we derived ste-step moment estimates for the m-SAM discrete algorithm and, via Lemma A.4 corresponding SDE update (36). These estimates demonstrate that the first- and second-order moments satisfy the matching conditions of Theorem A.2 Together with the uniform moment bounds from Lemma D.2 theorem for m-SAM.
we are now ready to establish main aeak-appro
Theorem C.7 (m-SAM SDE), Under Assumptions and 1 lei 0 < イ 1, , O, and N = (T/nl. Denote by { k= the m-SAM Hterates in and let {Xt)LEJOTT) be the solution of the SDE 4T. Suppose:
(i) The functions
V PP = V(
Σ Vhill) and MMMMAM
IEZ,
[Z]
are Lipschitz on IRd.
(ii) The mapping
(zz
(x)
m
Σ
Vz +
hy
satisfies, almost surely, the Lipschitz condition
Vi (ax (w)ll)) S L, || vll
Vx,yeR"
where L,, > 0 a.s. and [[ Do for every m 2
Then {X+ E [0, an order-(1, 1) weak approximation of { namely: for each g E G2, there exists constant C' > 0, independent of 1) such that
[[g(xx)] E[g(Xen)] S c((++
max OCKEN
Proof. First, verify that SDE 41 admits unique strong solution. By assumption, both the drift and diffusion coefficients are globally Lipschitz, which in turn implies linear-growth condition. Therefore, Theorem D.I applies and yields the existence and uniqueness of strong solution [0.
Then, by Lemmas complete.
A.4
C.6 and D.2 all the conditions of Theorem A.2 are satisfied, and the proof is
Remark C.8. The Lipschitz conditions are to ensure that the SDE has unique strong solution with uniformly bounded moments. It is possible to appropriately relax them we allow weak solutions Mil'shtein 1986).
D Auxiliary Lemmas
Theorem D.1 (Existence and Uniqueness of Strong Solutions (Evans 2012)). Let b Rd + Rd and
Rd Fdxm be measurable functions satisfying:
(i) Global Lipschitz. There exists a constant L such that
[lb(z)- - bly)l + ||a(x) - a(y)l I L 1 vll. Vx,yeRs
(ii) Linear growth. There exists a constant K such that
[[b()]|2 + [|a(2))12 < K (1+ + |x]l?),
Vr VrER
Let Xo be va--ue random variable with Xo [1 < sc. Then the SDE dX, b(X,) dt + (X,) dW.. Xo given,
admits unique strong solution { [Xth22 satisfying
B sup
|X, [1] < <c, VT 0.
DS$ST
28

---

Lemma D.2. (Li et al. 2019 Let {ø : k 0} be the generalized iterations defined in 22 Suppose
[1(2, S Ly (1 I
where Ln > 0 almost surely and
00 for all 2 Then for any fixed T and any m 2 1, the moment E [[]=x|1"]]] exists and is uniformly bounded in both 7 and k = 0, 1, ......,,,,, N. where [//]]-
E Proof of Proposition 3.9
We will use the following lemma on convex order to prove Proposition 3.9 whose proof can be found
in classical textbooks on stochastic order, such Müller and Stoyan 2002) Lemma E.1 (Convex-order Monotonicity). Let XI..... be i.i.d. random vectors Rd with a log-concave density. For each integer < k n, define
1 k
S
X
Then for any S le m and any convex cb TRd - R.
E[o(Sm)) & ((x))
Proof of Proposition 3.9 Lower bound: Applying Jensen's inequality, [=)]] s B (x))l..
[V)())]
Upper bound: By Cauchy- Sirohy-Schwazzieee inequality,
|V/ tr(V
||| < V VEII
hn
Combining both bounds completes the proof of the first statement.
For the second statement, we apply Lemma E.1 to the convex function || |||] which concludes the proof.
F Additional Related Works
Theoretical understanding of SAM. Although SAM and its variants have achieved remarkable success in various practical applications (Foret et al. 2021 Kwon al. [021]] Kaddour al.) 2022 Li et al.) 2024b) [i and Giannakis (2024), the theoretical understanding behind them remains limited. The pioneering work ol/Anddrusschenko and Flammarion (2022) provided the first theoretical framework for understanding SAM, covering its convergence properties and implicit bias in simple network structures, while also systematically illustrating several empirical phenomena. Subsequently, Si and Yun 2024) extended the convergence analysis to various deterministic and stochastic settings. in-depth analysis of the dynamics of SAM using the SDE framework previously developed by [i et al. (2017). leading to a deeper understanding of its implicit bias. On the other hand, |Wen et al. (2022) investigated the implicit bias of SAM by analyzing its slow ordinary differential equation (ODE) behavior near the minimizer manifold, demonstrating how SAM drifts toward flatter minima. More recently, Zhou et al. (2024) studied the late-stage behavior of SAM using stability analysis, showing its advantage in escaping sharp minima.
Compagnoni al. (203)); Luo et al. (2025) conducted
G Derivation of the Finite Difference Estimator in Equation
We start with the first-order Taylor expansion around a:: fz x + 82) - f((z)
Vf(())== O()),
8
where = (+1)* is a Rademacher random vector, and Si is assumed twice differentiable so that the remainder is of order O()).
29

---

Step 1: Square both sides.
f f((x)) (V 2(V((()) 0 0(0)
Often we simply write this as
(Vi(()) 0(5))² =(()))
(V)(()))=)² 0(6) (VI(()) 0(8))
Step 2: Take expectation over 2. Because 2 has independent { components, we have
Eg (Vi((z)T 2) = |[/a(')]|
(standard Rademacher property).
Hence,
E. ( (Alettba) ) (V(A(r)))))) [))()))=)]]]] + 0(62) ||V 0
Thus the mean-squared estimate of the finite difference quotient differs from [ by ( bias term, implying that the estimator is approximately unbiased as - 0. We compare two d-dimensional random vectors 2 E Rdd
Rademacher: each component 2j is independently #1 with probability 1/2. Standard Gaussian: each component žz is i..d. W/O, 1).
Both have E[zy] = 0 and E[2]] 1, so E[( T u) = ||w||2 for any E Rd. We look at the fourth moment [2 "] relevant to the variance in many difte-iifference gradient estimators.
Rademacher case.
Since 23 I,
d
(= (=Tu)²= =
Uzzÿ)
Σ j.k
UjUh Zj
/=1
Then
(z (z]b)²=== =
By 2yzk
3.k
and using |[=11] = 1[ (when j & k) with zero cross terms odd product, one obtains
relatively small constant factor. Detailed calculation yields
[( [o)
Σ j=l
+6
u s <3100|
(for d >> 1).
jck
Gaussian case. If NNO, then B[= 3 and 16|
1[23
for o k. One can show
B[(z
[1)=J 3||w||
Hence the constant factor in front of [|||| is exactly 3 for Gaussian.
Conclusion. For both Rademacher and Gaussian vectors, [B( ||o|2 However, when analyzing higher-order moments (e.g. (2 u)= that affect the variance of many difte-diffeeence or mandom-direction estimators. the Rademacher distribution can yield smaller constant factor. This often leads to reduced variance and tighter theoretical bounds for the same sample size.
H Derivation of the Gibbs Distribution 20
Consider the following maximization problem:
H(P)
Σ pi
[+(())]
PEA
30

---

where A is the probability simplex (i.e., E, = 1 and p: 0), (P) E; P: In pi is the entropy
term, and is given constant.
Construct the Lagrangian. We introduce the constraint E.Pr = with a Lagrange multiplier o::
((P,a) Σee [Vf(a)] + ( p. Imp.) + ( -p-
Differentiate w.r.t pj and zero. Taking partial derivatives with respect to pt and setting them dc Op,
to zero,
[Vf[(z)]| (Iup: 1) 0.
Therefore,
Vf((2)] p; a = p. = =epp A [V + +1-10 +1-da).
A
Enforce the normalization. The constraint E; = 1 fixes the value of a it amounts to one overall normalization factor in the denominator. Hence the solution takes the well-known Gibbs distribution form:
exp (A (x))l)
p
Σ, exp A [z)()))
I Algorithm: Reweighted SAM
Algorithm 1 Reweighted SAM
1: while not converged do
2: Forward pass obtain Jr, (xx)
3: for qorq=1... 1. Q do
4: Estimate per-sample gradient norm using Eq. 21
5: end for
6: Normalize estimated per-sample gradient norm
7:
Compute weight p" using Eq. (20)
8:
Compute perturbation e using Eq. 18
9:
Compute perturbed gradient 9k ff (2x pex)
10:
Update model parameters: #k+1 = 3% 77k
11:
end while
J Additional Experiment Results
All experiments were run on NVIDIA RTX 4090 GPUs.
Algorithm
Test Accuracy
Time/Epoch
Mini-batch SAM
78.90 0.27%
13.56
Algorithm
Test Accuracy
Time/Epoch
RW-SAM
79.31 0.28%
15.21
Mini-bulch USAM
78.94 + 0.45%
12.98
m-SAM (m 8)
80.72 0.12%
175.45
m-USAM (m 8)
80.66 + 0.04%
173.77
m-SAM (m 16)
80.47 0.09%
92.22
m-USAM 16)
90.86
m-SAM (07 32)
80.02 0.06%
49.87
m-USAM (m 32)
80.02 + 0.09%
47.09
M-SAM (m 64)
79.35 0.11%
26.44
m-USAM 64)
79,16 0.04%
24.15
n-SAM
78.15 0.19%
n-USAM
78.63 + 0.06%
Table 7: Left: performance and time cost of SAM, RW-SAM, m-SAM (with varying m), and n-SAM CIFAR-100 100. Right: performance and time cost of USAM, m m-USAM (with varying m), and n-USAM; with ResNet-18 18 on CIFAR-100.
31

---

Table 8: wall-clock time overhead of RW-SAM
ResNet-18
ResNet-50
WedeRes
SAM
13.6
43.0
97.7
RW-SAM
15.2
50.9
112.4
We present the experimental results of applying the proposed reweighting strategy ASAM wK et al. 2021) in Table The results demonstrate consistent improvements, and further investigation into effectiveness of applying the reweighting strategy to different SAM variants remains an interesting direction for future work.
Table 9: Test accuracy (%) comparison between ASAM and RW-ASAM CIFAR-IO and CIFAR- 100.
Method
CIFAR-10
CIFAR-100
ResNet-18
ResNet-50
ResNet-18
ResNet-50
ASAM
95.86 0.14
96.12 0.23
79.17 0.14
80.27 0.33
RW-ASAM
96.02 0.08
96.43 + 0.17
79.46 0.25
80.65 + 0.16
Table 10: Hyperparameter settings for fine-tuning DistilBERT on GLUE tasks.
Task
CoLA
MNLI
MRPC
QNLI
QQP
RTE
SST-2
STS-B
WNLI
Batch size
32
64
32
64
32
64
32
Learning rate
2c-5
3e-5
2c-5
3e-5
3e-5
2c-5
3e-5
2e-5
2e-5
Epochs
3
8
3
3
8
LR scheduler
Linear
Warmup ratio
0.1
Max sequence length
256
(for SAM and RW-SAM)
0.05
32

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We summarize the paper's contribution in both abstract and introduction.
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 analyze computational overhead as a limitation.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those 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]
33

---

Justification:
We provide the full set of assumptions and rigorous proof.
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 describe the proposed algorithm in detail.
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?
34

---

Answer: [No]
Justification: All experiments use publicly available datasets, but our code is not yet available. We plan release it upon acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm ec/
public/guidese https//ooooeeuaasssssooooooooo
for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips. //nips.c//ppl wwwppbblccgguiecuceccccceooooope Detgguddsse/Ceebbbiiiinn htt.SSbbmiiiiooooiiiiiooooi for more details.
The authors should provide instructions 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. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We describe all experimental details.
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 report the mean and standard deviation over repeated experiments.
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).
35

---

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 figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures 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 to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We describe them.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU 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 't make it into the paper).
9. Code of ethics
Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Question:
Answer:
[Yes]
Justification:
All aspects of the research conform to the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in 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 fundamental research does not have societal impacts.
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 socictal 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.
36

---

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 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: The paper poses no such risks.
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
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?
Question:
Answer: [Yes]
Justification: All open-source libraries and public datasets properly cited in the paper.
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.
37

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent 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. Crowdsoureing 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: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing 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 the main paper.
According the NeurIPS Code of Ethics, 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:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending 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 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 the institution conducting the review.
16. Declaration of LLM usage
38

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs any important, original, or non-standard components.
Guidelines:
The answer NA
that the method development in this research does involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy for what should or should be described.
39


Reviews:
[{"review_id": "Uoks0hwWCV", "reviewer": "Reviewer_vX9u", "summary": "This paper provides a theoretical and empirical investigation into m-sharpness, the phenomenon that Sharpness-Aware Minimization (SAM)  improves generalization more as the micro-batch size $m$ used for perturbation decreases (i.e. using more smaller disjoint batches for the SAM inner maximization). The authors extend the stochastic differential equation (SDE) framework for analyzing stochastic gradient descent (SGD) to two parameters (the learning rate $\\eta$ and SAM perturbation radius $\\rho$) and derive continuous-time approximations for several SAM variants. Building on these approaximations, the paper’s contributions are: (1) a rigorous SDE-based analysis of SAM, revealing a variance-driven implicit sharpness regularization mechanism underlying m-sharpness; (2) a unifying comparison of SAM, m-SAM, and related variants (normalized vs. unnormalized) in terms of their dynamics; and (3) a new sharpness-weighted data sampling technique (Reweighted SAM) inspired by the theory, with empirically demonstrated efficacy.", "questions": "Please see the weakness section. \n\ncould u authors comment on exact technical challenges involved in extension to two-parameter weak approximation. Prima-facie it looks a simple enough natural extension following up on previous works such as compognoni 2023 but i will reserve my judgement.", "limitations": "the authors have mentioned computation overhead, but this paragraph could have been expanded with additional comments about other variants of sam and RWSAM's extension to those.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\nThe paper is technically rigorous and the mathematical derivations in the paper appear sound and build upon established theory for interpreting SGD as a stochastic differential equations. The authors formalize a two-parameter weak convergence expansion (in Appendix A) that allows $\\eta \\to 0$ and $\\rho \\to 0$ at independent rates. This generalization of prior one-parameter analyses (e.g. Compagnoni et al. 2023and Luo et al. 2025) is non-trivial and is executed with rigor. In particular, Theorem 3.3 (for mini-batch USAM) and Theorem 3.5 (for m-USAM) are derived by applying Dynkin’s formula and carefully controlling error terms $O(\\eta^{\\alpha}\\rho^{\\beta})$. The resulting SDEs are also  intuitively reasoanble. The key finding – that stochastic gradient noise introduces a $\\nabla \\operatorname{tr}(V)$ regularization term – is a deep insight that advances our understanding of why SAM (especially with small $m$) finds flatter minima.\n\nAdditional algorithmic contribution — Rather than just providing post-hoc analysis, the authors use the theory to inspire a new method. RW-SAM is conceptually simple (importance-sampling the batch for perturbation based on gradient norm) yet effective. It addresses the key limitation of m-SAM (lack of parallelization) by staying with a single mini-batch but simulating the effect of smaller $m$. This idea is novel, and the experiments show it works: RW-SAM consistently outperforms standard SAM by a small but non-negligible margin (e.g., +0.4–0.5% on CIFAR-100, see Table 1) with only ~10-20% training overhead, and is much faster than true m-SAM (which requires serial micro-batch updates).\n\nThe empirical results are interesting and demonstrate that RW SAM consisntely outperforms vanilla SAM, though the improvements seem minor compared to more recent variations of SAM. \n\nThe paper is clearly written and well-organized. The theoretical results with their assumptions, development of new results and their context with known results from previous papers are all well presented. \n\nWeaknesses: \n\nThe related works are largely well-cited except a couple of important misses such as Keskar et al. (2017), who originally identified the connection between batch size, sharpness, and generalization. And the Vasso paper by Li and Giannakis which is also very relevant. mSam explorations by Behdin et al seems also relevant (msam: Micro-batch-averaged sharpness-aware minimization)\n\nSome of the dense sections could be written better and made more accessible but this is a highly technical paper and warrants some background e.g. in stochastic calculus to fully understand and appreciate. i have not gone through all the proofs etc.\n\n\nSome of the language could be toned down as it reads as over claiming e.g. the gradient variance is certainly important  as found by the analysis but would not fully explain the generalization. . It might be worth explicitly acknowledging that m-sharpness holds empirically in many cases but not universally (they cite a result in Wen et al. 2022 about transformers where it fails).\n\nThe empirical section could have been expanded to include more varied domains, model architectures and settings. \n\nMinor : \n\nTheorem 3.4 — incorrect equation is references (7) instead of (10)\n\nLine 289 : “tables” -> table.\n\nThere are some grammar issues .e.g “Motivated by this, we  evaluate RW-SAM maintains”", "was_revised": true, "final_justification": "I will stick to my rating. I dont have major concerns about the paper but because of its relatively incremental nature i dont see many discussions happening around this paper at the conference which is why I am hesistant to give a higher rating.", "rebuttal": "We thank the reviewer for the careful reading and positive feedback. Below we address the specific points raised:\n\n**1. Could u authors comment on exact technical challenges involved in extension to two-parameter weak approximation.**\n\nA key technical challenge lies in our need to simultaneously track both $\\eta$ and $\\rho$ in Dynkin expansion (Lemma A.3), and to develop a matched two‑parameter moment argument for both the continuous SDE and the discrete update (Lemmas A.4–A.5). In doing so, we must carefully control all mixed remainder terms uniformly for any relative scaling $\\eta,\\rho\\to0$, which is more complex than the single‑parameter case. Compared to prior SDE analyses for SAM (e.g., Compagnoni et al., 2023), which directly invoke the weak‑approximation result for $\\eta$ from Li et al. (2019), we avoid the combinatorial explosion of cross‑terms arising from an Itô–Taylor expansion (see Lemma 28 in Li et al. 2019) by using our generator‑based approach—resulting in a more concise derivation and a precise quantification of each order’s contribution.\n\n**2. Missing related work.**  \n\nWe have already cited Behdin et al. (2023). In the revised manuscript, we will also include discussion and citations for Keskar et al. (2017) and Li & Giannakis (2023).\n\n**3. Tone and language adjustments.** \n\nWe will soften overly strong claims about generalization, explicitly acknowledge the empirical limits of m-sharpness, reduce the density of formulas in the theory section, and add more intuitive explanations to improve readability, as suggested by the reviewer.\n\n**4. Minor corrections.**\n\nWe apologize for these errors and will correct them in the revised manuscript.\n\n**5. The empirical section could have been expanded to include more varied domains, model architectures and settings.**  \n\nDuring the rebuttal period, we added experiments on three GLUE tasks using DistilBERT with AdamW as the base optimizer trained for 10 epochs, and report the median results over five independent runs below. We find that, although SAM sometimes underperforms (e.g., on RTE), RW‑SAM consistently improves upon SAM. In the revised manuscript, we will include results for the full GLUE benchmark.\n\n| Method   | SST‑2 | RTE    | STS‑B (Pearson / Spearman) |\n|----------|-------|--------|----------------------------|\n| AdamW    | 90.7  | **63.9** | 86.9 / 86.8                |\n| SAM      | 91.3  | 61.4   | 87.0 / 86.9                |\n| RW‑SAM   | **91.7** | 62.5   | **87.2 / 87.0**              |\n\n**6. RWSAM's extension to other variants of sam.**\n\nWe have already obtained preliminary experimental results by applying our proposed reweighting scheme to other SAM variants, and we will include a discussion of these combined approaches in the revised manuscript."}, {"review_id": "nrNhM5r96v", "reviewer": "Reviewer_19PB", "summary": "The paper investigates the *m-sharpness* phenomenon in Sharpness-Aware Minimization (SAM), where generalization improves as the micro-batch size used for perturbation decreases. Extending the stochastic differential equation (SDE) framework to jointly model learning rate and perturbation magnitude, the authors provide theoretical insight into how stochastic gradient noise (SGN) induces implicit regularization in SAM and its variants. Motivated by this analysis, they propose Reweighted-SAM, an algorithm that adaptively reweights samples based on estimated SGN magnitude. This method aims to reproduce the benefits of m-SAM while remaining parallelizable. Experiments on standard vision benchmarks support the method’s efficacy.", "questions": "* How sensitive is Reweighted-SAM to violations of the SGN orthogonality assumption? Have the authors empirically studied its robustness under such conditions?\n* Could Reweighted-SAM be adapted to other sharpness-aware variants such as ASAM or FSAM? What adjustments would be necessary?\n* How critical is the SGN reweighting parameter λ in practice? Does it require tuning across datasets, or is normalization sufficient to ensure stable performance?\n* Have the authors explored ways to reduce the overhead of SGN magnitude estimation, such as computing it less frequently or using alternative proxies?\n* To what extent does Reweighted-SAM generalize beyond supervised image classification? Have the authors considered applications in NLP or reinforcement learning?", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "*Strengths*\n\n* The paper offers a rigorous and well-motivated theoretical framework that explains the role of SGN in implicit sharpness regularization.\n* Reweighted-SAM is a novel algorithm derived directly from the theoretical findings and is more computationally efficient than m-SAM.\n* Empirical results are thorough and demonstrate consistent gains across various datasets, architectures, and settings.\n* The analysis provides a compelling explanation for the generalization gap between full-batch and micro-batch SAM variants.\n\n*Weaknesses*\n\n* The theoretical analysis depends on strong assumptions, including orthogonality between SGN and the full gradient, and properties such as log-concavity and Lipschitz continuity, which may not hold in practice.\n* Some sections, especially those involving the extended SDE derivations, are dense and may be difficult for readers without a strong mathematical background.\n* Although Reweighted-SAM improves over m-SAM in efficiency, it still introduces nontrivial computational overhead, and the scalability to large models remains unclear.\n* Baseline comparisons omit several recent or stronger alternatives in sharpness-aware optimization, limiting the empirical context of the results.", "was_revised": true, "final_justification": "I have raised my score (from 4 to 5) based on the authors' strong rebuttal, which resolved my initial concerns. They provided new experiments demonstrating broader applicability to other optimizers and domains and clarified that their theoretical framework does not strictly require the strong assumptions I had questioned. The paper's core contributions—providing novel theoretical insight into m-sharpness and proposing a new algorithm—are significant. The computational overhead is a minor limitation but represents a reasonable trade-off for the performance gains and improved efficiency over m-SAM. The paper's strengths now clearly outweigh its weaknesses.", "rebuttal": "We thank the reviewer for the insightful comments and address the concerns point by point below.\n\n**1. The theoretical analysis depends on strong assumptions, including orthogonality between SGN and the full gradient, and properties such as log-concavity and Lipschitz continuity, which may not hold in practice.**\n\nWe would like to clarify that these assumptions are used in different contexts rather than all at once. Specifically: Lipschitz continuity is assumed only to guarantee the existence and uniqueness of a strong solution to the SDE; if weak solutions are acceptable, this requirement can be relaxed. Log‑concavity of the gradient distribution appears only in the extra statement of Proposition 3.9 to ensure monotonicity with respect to batch size; even without log‑concavity, the main proposition still holds. The orthogonality between the full gradient and noise is merely an intuitive assumption and is not required by our core theory. In practice, we only need the norm of the stochastic gradient to approximate the norm of the gradient noise—which is a good approximation for USAM or in late‐training when the full‐batch gradient norm is small—so strict orthogonality is unnecessary.\n\n**2. Some sections, especially those involving the extended SDE derivations, are dense and may be difficult for readers without a strong mathematical background.**\n\nIn the revised version, we will reduce the density of formulas in the theory section, add more intuitive explanations to improve readability, and provide additional background on SDEs to make the content easier for readers to follow.\n\n**3. Could Reweighted-SAM be adapted to other sharpness-aware variants such as ASAM or FSAM? What adjustments would be necessary?**\n\nOur data‑reweighting approach is complementary to parameter‑level variants such as ASAM and FSAM, so it can be combined with them directly. In the rebuttal, we added experiments on ASAM with reweighting (termed RW‑ASAM); the results below clearly demonstrate that our method integrates effectively with ASAM and improves performance. In the revised manuscript, we will include a discussion of how our reweighting method can be paired with other SAM variants and present more extensive experimental results.\n\n### CIFAR-10\n| Method   | ResNet-18        | ResNet-50        |\n|---------|------------------|------------------|\n| ASAM    | 95.86 ± 0.14     | 96.12 ± 0.23     |\n| RW-ASAM | **96.02 ± 0.08** | **96.43 ± 0.17** |\n\n### CIFAR-100\n| Method   | ResNet-18        | ResNet-50        |\n|---------|------------------|------------------|\n| ASAM    | 79.17 ± 0.14     | 80.27 ± 0.33     |\n| RW-ASAM | **79.46 ± 0.25** | **80.65 ± 0.16** |\n\n**4. How sensitive is Reweighted-SAM to violations of the SGN orthogonality assumption? Have the authors empirically studied its robustness under such conditions?**\n\nAs noted in our response to Point 1, strict orthogonality between the full gradient and noise is not required—what matters is that the norm of the stochastic gradient closely approximates the norm of the noise, which empirically holds in mid‑to‑late training. For instance, on CIFAR‑100 with ResNet‑18, we observed that, for most samples, the ratio of noise norm to gradient norm exceeds 0.98. While we have not conducted a dedicated robustness study on this orthogonality assumption, our noisy‑label experiments (Table 3) implicitly stress it: injected label noise strongly violates any near‑orthogonality, yet Reweighted‑SAM still outperforms all baselines. In particular, at an 80% noise ratio, Reweighted‑SAM achieves a 16% absolute accuracy gain. These results indicate that our method does not critically depend on the SGN–gradient orthogonality assumption.\n\n**5. How critical is the SGN reweighting parameter λ in practice? Does it require tuning across datasets, or is normalization sufficient to ensure stable performance?**\n\nAs demonstrated in Section 5.4 and Table 4, our algorithm is largely insensitive to the choice of λ: the default setting is generally robust, thanks to the normalization step in our method.\n\n**6. Have the authors explored ways to reduce the overhead of SGN magnitude estimation, such as computing it less frequently or using alternative proxies?**\n\nWe did not explore this in our current implementation because we estimate SGN magnitude on a per‑sample basis, and reducing the estimation frequency risks producing overly coarse reweighting—especially on large datasets. Exploring techniques to mitigate the extra overhead of our reweighting scheme is a direction for future work.\n\n**7. To what extent does Reweighted-SAM generalize beyond supervised image classification? Have the authors considered applications in NLP or reinforcement learning?**\n\nDuring the rebuttal period, we added experiments on three GLUE tasks using DistilBERT with AdamW as the base optimizer trained for 10 epochs, and report the median results over five independent runs below. We find that, although SAM sometimes underperforms (e.g., on RTE), RW‑SAM consistently improves upon SAM. In the revised manuscript, we will include results for the full GLUE benchmark. \n\n| Method   | SST‑2 | RTE    | STS‑B (Pearson / Spearman) |\n|----------|-------|--------|----------------------------|\n| AdamW    | 90.7  | **63.9** | 86.9 / 86.8                |\n| SAM      | 91.3  | 61.4   | 87.0 / 86.9                |\n| RW‑SAM   | **91.7** | 62.5   | **87.2 / 87.0**              |\n\nIn reinforcement learning, there have been very recent efforts to adapt SAM (Lee & Yoon, 2025), but it is not yet popular in RL. Since RL setups differ fundamentally from supervised learning, developing a rigorous theoretical framework for SAM in RL will require additional technical advances and represents a promising direction for future work.\n\nLee, H. K., & Yoon, S. W. (2025, April). Flat reward in policy parameter space implies robust reinforcement learning. In The Thirteenth International Conference on Learning Representations."}, {"review_id": "P1Rg3OBIzC", "reviewer": "Reviewer_TwRD", "summary": "The authors explore the phenomenon of m-sharpness in SAM, where reducing the micro-batch size during perturbation leads to better generalization. Using a detailed SDE framework, they analyze various SAM variants and uncover that SGN inherently contributes to implicit sharpness regularization. Based on this insight, they propose Reweighted-SAM that adaptively samples data based on SGN magnitude to achieve strong generalization while maintaining efficiency. Extensive experiments on CIFAR and ImageNet show that Reweighted-SAM outperforms standard SAM without the high overhead of m-SAM.", "questions": "- I think Reweighted-SAM addresses parallelization using Monte Carlo sampling. Can you explain further in detail?\n- Line 228. The orthogonality assumption between SGN and full gradient seems strong. What happens if this is not held? And does it have any effects on the parallelization characteristic of Reweighted-SAM?\n- Too many proofs and maths, will you disclose your code during the review process so we can have a further check on the approach?", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n- Solid theoretical framework using advanced SDE analysis to explain m-sharpness.\n- Strong empirical results across datasets and models, including robustness to label noise.\n\n**Weaknesses**\n- Theoretical analysis is dense and may lack accessibility for a broader reader group.\n- No ablation comparing Reweighted-SAM directly with other SGN-regularization techniques or alternative reweighting schemes.", "was_revised": true, "final_justification": "I have no further questions after the authors' responses. I will not change my rating.", "rebuttal": "We thank the reviewer for their thoughtful evaluation. Below we address each of the concerns raised.\n\n**1. Theoretical analysis is dense and may lack accessibility for a broader reader group.**\n\nIn the revised version, we will reduce the density of formulas in the theory section, add more intuitive explanations to improve readability, and provide additional background on SDEs to make the content easier for readers to follow.\n\n**2. No ablation comparing Reweighted-SAM directly with other SGN-regularization techniques or alternative reweighting schemes.**\n\nTo our knowledge, ours is the first algorithm to incorporate a reweighting mechanism into SAM. Existing SGN-based regularization methods typically struggle to match SAM’s performance and cannot be directly integrated into the SAM framework.\n\n**3. I think Reweighted-SAM addresses parallelization using Monte Carlo sampling. Can you explain further in detail?**\n\nWe assign each sample an importance weight related to its SGN norm, and use Monte Carlo sampling to estimate these per-sample gradient norms with only a single forward pass over the minibatch. This MC-based weight estimation is the key enabler that makes our approach fully parallelizable.\n\n**4. Line 228. The orthogonality assumption between SGN and full gradient seems strong. What happens if this is not held? And does it have any effects on the parallelization characteristic of Reweighted-SAM?**\n\nThe orthogonality between the full gradient and noise is merely an intuitive assumption. In fact, we only require that the stochastic‐gradient norm closely approximates the noise norm, which empirically holds in mid-to-late training. For instance, on CIFAR‑100 with ResNet‑18, we observed that, for most samples, the ratio of noise norm to gradient norm exceeds 0.98. While we have not conducted a dedicated robustness study on this orthogonality assumption, our noisy‑label experiments (Table 3) implicitly stress it: injected label noise strongly violates any near‑orthogonality, yet Reweighted‑SAM still outperforms all baselines. In particular, at an 80% noise ratio, Reweighted‑SAM achieves a 16% absolute accuracy gain. These results indicate that our method does not critically depend on the SGN–gradient orthogonality assumption. The orthogonality does not compromise the algorithm’s parallelizability.\n\n**5. Too many proofs and maths, will you disclose your code during the review process so we can have a further check on the approach?**\n\nDue to this year’s NeurIPS policy, we are unable to upload code during the rebuttal period; however, we will be happy to release our implementation if the paper is accepted."}, {"review_id": "kwXrNEBdoL", "reviewer": "Reviewer_Gf2z", "summary": "This paper studies why smaller micro-batches in Sharpness-Aware Minimization (SAM) improve generalization and proposes an efficient optimizer that captures this effect. The main contributions are the followings:\n\n- Develops a unified SDE-based analysis of SAM/USAM variants, linking smaller micro-batch size (m) to stronger variance-based regularization.\n- Shows that the covariance of stochastic gradient noise (SGN) approximates the Hessian near minima, connecting SGN variance to sharpness.\n- Demonstrates that this SGN-induced regularization explains the m-sharpness effect observed empirically.\n- Proposes **Reweighted-SAM**, an efficient optimizer that assigns higher weight to samples with larger gradient norms.\n- Reweighted-SAM avoids the high compute cost of small-m SAM by requiring only one extra forward pass.\n- Empirically outperforms or matches SAM and m-SAM across CIFAR-10/100, ImageNet-1K, and ViT fine-tuning tasks.", "questions": "1. Can you derive an SDE for Reweighted-SAM or justify its regularization behavior? The current theory stops before the proposed optimizer. It’s unclear whether Reweighted-SAM introduces new drift/diffusion terms or still follows the same implicit regularization as SAM.  \n\n2. Why are recent lightweight SAM variants missing from experiments or discussion?\nThe paper does not compare Reweighted-SAM against other one-pass or efficient SAM variants. This omission makes it difficult to assess how competitive your method is in the broader landscape of sharpness-aware optimizers.  Given the paper’s theoretical focus, extensive experiments aren’t expected—but a brief discussion of lightweight SAM variants is important to contextualize Reweighted-SAM within existing work.\n\n3. Can you show more evidence of implicit regularization? Currently, the only supporting number is the trace of the gradient noise covariance at convergence.\n\n4. Can the authors clarify when SAM variants are worth the extra compute?\nReweighted-SAM adds only one forward pass, but this can still be costly at scale. Since SAM methods are generally more expensive than SGD or Adam, it would help to explain which real-world tasks justify this overhead.", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "### Strengths\n- The paper provides a solid SDE-based analysis for various SAM and USAM variants, including closed-form expressions for the unnormalized case.\n\n- It offers a clear explanation for the m-sharpness effect by linking it to the variance of stochastic gradient noise, which acts as an implicit regularizer.\n\n- The proposed method shows how to retain the benefits of small-m SAM with only one additional forward pass.\n\n### Weaknesses\n\n- The proposed optimizer feels detached from the SDE analysis. The paper does not derive an SDE for Reweighted-SAM, so it is unclear whether the new method truly follows the same noise-regularization dynamics or introduces different effects.\n\n- The paper only compares Reweighted-SAM against SGD, standard mini-batch SAM, and a few SAM variants. It omits newer, lightweight sharpness-aware optimizers such as **LookSAM** [4], **ESAM** [1], **GSAM** [5], **Momentum-SAM** [2], and **ASAM** [3]. Without these baselines, it is difficult to assess how Reweighted-SAM stacks up against the broader state of the art. I believe a comparison with some of these methods or at least a discussion on these prior works would be helpful.\n\n- The authors do not derive an SDE for Reweighted-SAM, nor do they analyze its drift or diffusion behavior. Apart from reporting a single scalar (the trace of the SGN covariance at convergence), there are no deeper implicit-regularization diagnostics leaving it unclear whether Reweighted-SAM truly follows the same variance-regularization dynamics as SAM.\n\n---\n\n### References  \n\n[1] Du, J., Yan, H., Feng, J., Zhou, J. T., Zhen, L., Goh, R. S. M., & Tan, V. Y. F. (2022). Efficient Sharpness-Aware Minimization for Improved Training of Neural Networks. Proceedings of the International Conference on Learning Representations (ICLR). \n\n[2] Becker, M., Altrock, F., & Risse, B. (2024). Momentum-SAM: Sharpness-Aware Minimization without Computational Overhead (arXiv:2401.12033). arXiv preprint. \n\n[3] Kwon, J., Kim, J., Park, H., & Choi, I. Y. (2021). ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks. In Proceedings of the 38th International Conference on Machine Learning (ICML) (pp. 6060–6069). \n\n[4] Liu, Y., Mai, S., Chen, X., Hsieh, C.-J., & You, Y. (2022). Towards Efficient and Scalable Sharpness-Aware Minimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). \n\n[5] Zhuang, J., Gong, B., Yuan, L., Cui, Y., Adam, H., Dvornek, N., Tatikonda, S., Duncan, J. S., & Liu, T. (2022). Surrogate Gap Minimization Improves Sharpness-Aware Training. In Proceedings of the 39th International Conference on Machine Learning (ICML) (pp. 26848–26872).", "was_revised": true, "final_justification": "I appreciate that some of my concerns were addressed. That said, the reweighted SAM still feels disconnected from the main theoretical contribution. I believe my current score (borderline accept) is a fair assessment given the current state of the paper.", "rebuttal": "We thank the reviewer for their detailed evaluation and constructive suggestions. Below, we address each of the points raised.\n\n**1. Can you derive an SDE for Reweighted-SAM or justify its regularization behavior?**\n\nReweighted-SAM is essentially a heuristic algorithm inspired by our theoretical analysis. Although we can formally derive an SDE for Reweighted-SAM, the result contains highly complex coupling terms between the per-sample weights $p_i$ and the gradients, making the equation neither intuitive nor particularly informative for understanding the implicit regularization.\n\n**2. Why are recent lightweight SAM variants missing from experiments or discussion?**\n\nOur paper is primarily focused on studying the m-sharpness phenomenon in SAM and developing a theoretical framework to understand it, which is why we did not cover the full range of empirical SAM variants.  We acknowledge this oversight.\nOur data‑reweighting approach is complementary to parameter‑level variants such as ASAM and GSAM, so it can be combined with them directly. In the rebuttal, we added experiments on ASAM with reweighting (termed RW‑ASAM); the results below clearly demonstrate that our method integrates effectively with ASAM and improves performance. In the revised manuscript, we will include a discussion of how our reweighting method can be paired with other lightweight SAM variants, including LookSAM, ESAM, GSAM, Momentum-SAM, and ASAM, and present more extensive experimental results.\n### CIFAR-10\n| Method   | ResNet-18        | ResNet-50        |\n|---------|------------------|------------------|\n| ASAM    | 95.86 ± 0.14     | 96.12 ± 0.23     |\n| RW-ASAM | **96.02 ± 0.08** | **96.43 ± 0.17** |\n\n### CIFAR-100\n| Method   | ResNet-18        | ResNet-50        |\n|---------|------------------|------------------|\n| ASAM    | 79.17 ± 0.14     | 80.27 ± 0.33     |\n| RW-ASAM | **79.46 ± 0.25** | **80.65 ± 0.16** |\n\n**3. Can you show more evidence of implicit regularization?**\n\nIn our paper, we quantify implicit regularization via the trace of the SGN covariance matrix according to our theory. Beyond the final trace values, Figures 1 and 2 in Section 3.5 plot its evolution throughout training and demonstrate a strong correlation with generalization performance. In the revised manuscript, we will include analogous learning curves for Reweighted-SAM—visually illustrating how its SGN trace evolves and induces implicit regularization over time (we are unable to upload these figures here due to NeurIPS policy this year).\n\n**4. Can the authors clarify when SAM variants are worth the extra compute?**\n\nWe anticipate that in overparameterized models—where training loss can be driven to near zero—the extra cost of stronger implicit regularization is often worthwhile. In such settings, the abundance of global minima means methods like SAM and Reweighted-SAM can steer optimization toward flatter minima that generalize better (as discussed in our Section 3.5). Conversely, for tasks where convergence speed is paramount—such as LLMs—the overhead of SAM and Reweighted-SAM may outweigh their benefits."}]

Author rebuttal:

Meta-review: The authors investigate sharpness-aware minimization (SAM) via a phenomenon known as m-sharpness where the performance of SAM improves monotonically as the micro-batch size for computing perturbations decreases. The authors use a SDE framework to analyze this phenomenon. They find that the stochastic noise introduced during SAM perturbations inherently induces a variance-based sharpness regularization effect. The paper then introduces a reweighted SAM. 

All the reviewers were positive about the paper. They thought the results were theoretically rigorous and interesting. Several reviewers increased their scores after the rebuttal from the authors. Some additional work cited needs to be included as well as a few additional experiments that appeared to be performed during the rebuttal.
Decision: Accept (poster)