Title: Unveiling m-Sharpness Through the Structure of Stochastic Gradient Noise
Venue: NeurIPS 2025
Accepted: yes
Keywords: Optimization

Sharpness-aware Minimization

Stochastic Differential Equations
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n### 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).\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-02"}

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

STRENGTHS AND WEAKNESSES:
**Strengths**
- Solid theoretical framework using advanced SDE analysis to explain m-sharpness.
- Strong empirical results across datasets and models, including robustness to label noise.

**Weaknesses**
- Theoretical analysis is dense and may lack accessibility for a broader reader group.
- No ablation comparing Reweighted-SAM directly with other SGN-regularization techniques or alternative reweighting schemes.

QUESTIONS:
- I think Reweighted-SAM addresses parallelization using Monte Carlo sampling. Can you explain further in detail?
- 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?
- 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
TITLE: Official Review
CLARITY: 2
IMPACT: 3
ORIGINALITY: 2
DATE: 2025-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n*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.\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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”\n\nQUESTIONS:\nPlease 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.\n\nLIMITATIONS:\nthe 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.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}