Unveiling m-Sharpness Through the Structure of Stochastic Gradient Noise

We thank the reviewer for their detailed evaluation and constructive suggestions. Below, we address each of the points raised.

1. Can you derive an SDE for Reweighted-SAM or justify its regularization behavior?

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

2. Why are recent lightweight SAM variants missing from experiments or discussion?

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

CIFAR-10

CIFAR-100

3. Can you show more evidence of implicit regularization?

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

4. Can the authors clarify when SAM variants are worth the extra compute?

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

We thank the reviewer for their thoughtful evaluation. Below we address each of the concerns raised.

1. Theoretical analysis is dense and may lack accessibility for a broader reader group.

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

2. No ablation comparing Reweighted-SAM directly with other SGN-regularization techniques or alternative reweighting schemes.

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

3. I think Reweighted-SAM addresses parallelization using Monte Carlo sampling. Can you explain further in detail?

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

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?

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

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?

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

We thank the reviewer for the insightful comments and address the concerns point by point below.

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.

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

2. Some sections, especially those involving the extended SDE derivations, are dense and may be difficult for readers without a strong mathematical background.

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

3. Could Reweighted-SAM be adapted to other sharpness-aware variants such as ASAM or FSAM? What adjustments would be necessary?

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

CIFAR-10

CIFAR-100

4. How sensitive is Reweighted-SAM to violations of the SGN orthogonality assumption? Have the authors empirically studied its robustness under such conditions?

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

5. How critical is the SGN reweighting parameter λ in practice? Does it require tuning across datasets, or is normalization sufficient to ensure stable performance?

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

6. Have the authors explored ways to reduce the overhead of SGN magnitude estimation, such as computing it less frequently or using alternative proxies?

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

7. To what extent does Reweighted-SAM generalize beyond supervised image classification? Have the authors considered applications in NLP or reinforcement learning?

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

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

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

We thank the reviewer for the careful reading and positive feedback. Below we address the specific points raised:

1. Could u authors comment on exact technical challenges involved in extension to two-parameter weak approximation.

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

2. Missing related work.

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

3. Tone and language adjustments.

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

4. Minor corrections.

We apologize for these errors and will correct them in the revised manuscript.

5. The empirical section could have been expanded to include more varied domains, model architectures and settings.

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

6. RWSAM's extension to other variants of sam.

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