Avoiding spurious sharpness minimization broadens applicability of SAM

We thank the reviewer for their positive assessment, insightful comments, and detailed feedback. We are glad that you were able to contextualize our contribution quite spot on.

 

---

 

Same compute budget comparisons

As it stands, at an equal number of FLOPs, well-tuned AdamW does slightly outperform functional SAM. However, there are nevertheless important reasons why we believe functional SAM is still promising:
 

• Relevant Use Cases: There are practically relevant use cases, such as in data-limited regimes or where the model size is constrained (e.g. due to inference time constraints), where the better performance at fixed step count is desirable and where the extra training overhead of functional SAM may not matter as much.
   
• Solution Quality Beyond Loss: Even in the equal-FLOP setup, functional SAM’s better geometric properties in its flatter solution might be preferable than a method yielding sharper solutions, since it has been well-documented [Liu et. al., 2023] flatness of the solution correlates more robustly with downstream performance than similar values of loss.
   
• Path to Efficiency: Lastly, functional SAM is algorithmically compatible with efficient SAM approaches. Approaches like LookSAM [Liu et. al., 2022] can decrease the overhead of SAM to 5-10% while maintaining much of the benefit of the method. We hope to test this and other approaches in future work, and we believe that the overhead can be greatly reduced.

 

Comparison to other SAM variants:

• While interesting, comparing against the plethora of (effectively vision-focused) SAM variants was beyond the scope of this work.
   
• During our initial investigation into SAM's poor performance in LMs, we did experiment with common, simpler variations, such as using the unnormalized perturbation step or trying different weighted combinations of the SAM gradient and the standard gradient. However, these modifications did not appear to fundamentally resolve the performance degradation observed in language modeling tasks.
   
• In retrospect, this is perhaps not entirely surprising. Critically, none of these existing variations are explicitly designed to prioritize functional sharpness minimization over logit sharpness. Our diagnosis revealed that this preferential treatment of functional geometry is precisely what is needed to make SAM effective in LMs, requiring a more significant departure from standard SAM formulation.
   
• It would be interesting to nevertheless explore other orthogonal SAM variants, which address issues like parameter scaling sensitivity (ASAM, Kwon et al., 2021), norm choices (Tahmasebi et al.,2024) or perturbation stability (ESAM, Li & Giannakis, 2024), in relation to functional SAM as well. But since the design space of new algorithms tends to explode, we have had to stay within our scope making SAM effective in pre-training LMs.

 

Non-spurious sharpness measures that might correlate better with generalization

This is an excellent remark. We do think that something which directly measures the extent of functional curvature (like the frobenius norm of the functional Hessian) or its relative extents compared to the logit curvature to potentially be revelatory. We will try to add some measurements of this kind in the final version, but this more likely deserves a separate study of its own.

 

Vision vs Language Intuition:

Another great question. Our current hypothesis relates to the nature of the typical output distributions $p(y | x; \theta)$ in these domains.

• In many vision tasks, the probability mass often concentrates over a relatively small number of semantically related classes or visual scenes. The output distribution might be less dispersed, and manipulating logit statistics could potentially align reasonably well with improving the underlying function's robustness.
   
• In language modeling (specifically next-token prediction), the distribution over the next token is often highly dispersed and heavy-tailed, with non-negligible probability assigned to many different words. In such a setting, minimizing sharpness simply by manipulating logit statistics (e.g., making the distribution slightly peakier) might be an "easy" path for the optimizer that doesn't translate to genuine improvements in the functional geometry.

 

Suggestions on Experimental Designs Or Analyses

• Thanks for these great suggestions, we will definitely incorporate them in the camera ready version.

 

---

We hope we have able to address your concerns. We remain at your disposal should you have more questions or comments.

We thank you for your feedback and sharing the interesting works. Besides, we are pleased to hear that you find the exploration intriguing and recognize its potentially substantial ramifications.

 

1. Significance of Empirical Gains (0.03 loss):

We understand the concern about seemingly small improvements in loss. However,

• Context is Key: In large-scale LM pre-training (100M-1B+ params), improvements of 0.03-0.06 validation loss (Tables 1, 2, 3) are practically meaningful. They are comparable to or exceed gains in well-respected works on LLM optimizers (e.g., SOAP [Vyas et al., 2024], CASPR [Duvvuri et al., 2024] which are both improvements over Shampoo [Gupta et al., 2018] and report similar gains if not lower) and newer variants of Attention [Leviathan, et. al., 2025], LLM Fusion [Mavromatis, et. al. 2024], and corpus deduplication [Fig Lee, et. al., 2022], to list a few. All of these mentioned papers consider the C4 dataset, and so the differences are comparable.
   
• Statistical Significance:
  • As discussed in Lines 263-274 of Section 5.2, for prototyping we conducted our experiments using 3 random seeds. We observed that validation loss results were typically stable to the 3rd or 4th decimal place, indicating very low variance between runs. The reported results in Table 3 are averaged over 3 seeds (where the various methods rank as, 3.86 vs 3.88 vs 3.90 in Table 3 for precond. Functional SAM vs Functional SAM vs AdamW) are therefore highly significant and not due to noise. We will clarify this.
     
  • For the later experiments at the much larger parameter scale, we did have to report single seeds due to the constraints of time and corresponding costs of these experiments. However, to address this valid concern of yours, we have carried out an experiment on the 1.2 B parameter model over 3 seeds, and the averaged results in a fixed-length (50K step) settings are precond. Functional SAM 2.70 versus AdamW 2.73. This highlights that our results continue to be statistically significant even at these scales. We will include these results in the revision.
     

Hence, we can be confident that the gains delivered through our method are genuinely meaningful. In fact, it would serve us to remember that standard SAM consistently performed worse than AdamW (Fig 1). We have been able to turn things around and shown, for the first time, positive gains from SAM-style regularization in this setting over AdamW. And, the difficulty of achieving any improvement over tuned AdamW at this model scale cannot be overstated.

 

2. Preconditioning Novelty (vs. Optax, Granziol, etc.):

We appreciate the reviewer pointing out related work and common practices.

• Our contribution regarding preconditioning (Sec 4.2) should be understood specifically as addressing the potential mismatch between SAM's default Euclidean perturbation and the preconditioned geometry used by the outer optimizer (AdamW), particularly relevant for heterogeneous Transformer landscapes. Moreover, we also provide a theoretical argument (App B.1) that preconditioning can help re-balance logit/functional paths, which enriches our perspective about preconditioning as well.
   
• While general preconditioning and using Adam's state (as perhaps done implicitly in some Optax implementations) are known, what constitutes novelty in this specific context is the explicit motivation for fixing SAM's failure in LMs by aligning geometries/rebalancing paths, and the demonstration of its effectiveness especially in combination with functional SAM (Table 3 shows precond. functional SAM outperforms plain functional SAM and precond. SAM).
   
• Works like Granziol et al., Das et al., while quite intriguing, explore preconditioning for the main step, not specifically for SAM's perturbation. We will refine Sec 4.2 and Related Work to better delineate our specific contribution versus existing preconditioning concepts.

 

---

Hopefully, this addresses your pending concerns, but please let us know if you have any more questions or comments.

We thank the reviewer for their detailed comments and address some of their concerns here.

 

---

 

1. Computational Overhead:

We agree that in its current form, Functional SAM is not as FLOPs efficient as Adam.

• However, FLOPs are not the only limiting factor in training; for example, in certain scenarios industrial practitioners tend to be data limited or model size limited. In these scenarios, extra training time is acceptable for better final quality, and functional SAM can be a better fit.
   
• In addition, functional SAM is compatible with efficiency techniques like LookSAM [1]; this literature suggests that we may be able to reduce the overhead to 5-10% while maintaining most of the benefit of (Functional) SAM. We hope to pursue this avenue in future work.
   
• Our focus here was establishing the effectiveness of the functional path approach first. Our work provides a crucial understanding of SAM's limitations and offers the first validated approach to successfully apply sharpness-aware methods to large-scale LM pre-training (as noted by Reviewer 8e58).

[1] https://arxiv.org/abs/2203.02714

 

2. Concern regarding long-term performance:

Please have a look at Table 3, where we already show the gains provided by functional SAM sustain longer training durations as well. Our 1.2 B model gets 2.61 in terms of validation loss, which is very close to the 2.56 validation loss from [4] which you have alluded to. These 0.05 differences between the two works can easily be because their empirical setups, such as hyperparameters, exact architectural implementation, might not be identical.

Thus, we can be confident that functional SAM does yield long-term performance benefits as well.

 

3. Demonstrating Enhanced Functional Path:

This is a great suggestion, and we are working on measurements to show this. The experiments did not finish in time for the rebuttal but will be included in the revision.

 

4. Theoretical Support and Preconditioned SAM Motivation:

• Functional path formulation: The functional SAM update (Eq. 11) itself involves no approximation, and is an exact analogue of the original SAM update (Eq. 2), but where the contributions along the logit-sharpness path have been suppressed by design. We used the penalty SAM formulation for discussion following previous works which take advantage of the fact that penalty SAM is more amenable to theoretical analysis while simultaneously giving similar performance to original SAM. This gave us easier way to present and delineate the differences between SAM and Functional SAM.
   
• Preconditioned SAM: The motivation is twofold:
  • (1) Empirically motivated: To address the mismatch between SAM's spherical perturbation and AdamW's elliptical perturbation due to its diagonal preconditioning, which might cause issues in heterogeneous landscapes like Transformers.
     
  • (2) Theoretically motivated: Preconditioning the perturbation by approx. $H_{G}^{-1}$ (approximated by AdamW's ${M}^{-1}$) can selectively dampen the logit path ${\delta}_{logit} = H_G \epsilon^\ast$ more than the functional path $\delta_{func} = H_F \epsilon^*$, thus promoting the functional path (detailed in App B.1, where we made basic assumptions to make the argument quantitative). We will clarify this motivation in Sec 4.2.

 

5. Missing References [1-3]:

• [1] and [2] are related to SAM in that they propose more efficient variants [1] or modifications based on normalization layers [2]. But both of these are exclusively evaluated in vision, where they do not interface with the problem faced by SAM in language modeling tasks.
   
• Although ref [3] shares similar motivations as SAM and is quite interesting, they in their own words say that the “specific approaches differ significantly”.
   
• These works are firmly orthogonal to our core contribution: diagnosing the specific failure mode of SAM in LMs (logit path dominance) via a novel decomposition and proposing targeted fixes (Functional SAM, precond. SAM) that make SAM effective in this domain for the first time.

Regardless these will be good works to discuss in our paper, and we thank you for suggesting them.

---

 

Let us know if we can clarify any further points. If we have answered your concerns, please consider giving your score a second thought.

We thank you for your thorough review. We address the primary concerns below:

 

---

 

1. Significance of Performance Improvements:

• In LLM pre-training (100M-1B+ params), achieving consistent validation loss improvements of 0.03-0.06 (as seen in Tables 1, 2, 3) is highly significant. This magnitude is comparable to or exceeds gains reported in recent, well-regarded optimizers for LLMs, such as SOAP [Vyas et al., 2024, 0.04-0.06 improvement] and CASPR [Duvvuri et al., 2024, ~ 0.01-0.03]. This is also similar to the size of gains from architectural and data processing changes [Leviathan, et. al., 2025, Mavromatis, et. al. 2024, Fig Lee, et. al., 2022].

• Crucially, prior to our work, SAM consistently degraded performance compared to AdamW in this setting (Fig 1). Our methods are the first to successfully leverage SAM-style regularization for improved LM pre-training results across scales, representing a notable advance in the field.

• As clarified in our response to Reviewer Eee3, these gains are statistically significant and indeed represent genuinely meaningful improvements.

 

2. Computational efficiency concerns

• Costs relative to SAM and AdamW: functional SAM has virtually identical computational cost (FLOPs and memory) as standard SAM. Both require one forward pass, one backward pass for the initial gradient, and one backward pass (VJP) for the SAM/F-SAM gradient, resulting in ~2x the cost of AdamW per step. We will include explicitly measured time per step in the revision.

• Fixed Data Budget & Model Size scenarios: Our primary comparison point in the paper is equal steps since this is practically relevant in scenarios limited by data availability or required model size (e.g., inference constraints), where extra training time is acceptable for better final quality.

• Future Efficiency: As discussed in Section 6, Functional SAM is compatible with efficient SAM methods like LookSAM [1], offering a clear path to reducing the overhead to ~5-10% in future work. Our focus here was the fundamental advance of making any SAM variant work effectively with language models, as also identified by Reviewer 8e58.

 

3. Comparison to Other Methods:

• Improved SAM Variants: Methods like ASAM (Kwon et al., 2021), ESAM, etc., primarily target vision tasks and do not address the fundamental logit-path dominance issue we identified in language modeling. Our decomposition and Functional SAM are thus orthogonal contributions aimed at fixing SAM's failure in a new domain.

• Preconditioning Methods (Shampoo/K-FAC): Our preconditioned-SAM technique is indeed compatible with any preconditioning method, and it would be interesting to try our technique with base optimizers which use non-diagonal preconditioning. We leave this to future work.

 

4. Theoretical Claims:

• Proofs: We followed a common paradigm in deep learning research: identify an empirical issue, propose a diagnostic (logit/functional path), develop a principled fix (functional SAM), and validate empirically. Rigorous proofs for SAM-like methods are challenging and sometimes unhelpful (see response to X2he, point 2). Our consistent empirical gains across scales strongly support our hypothesis.

• Angle-SAM: We appreciate the reviewer's interest in Angle-SAM, which arises naturally from our decomposition. We de-emphasized it in the current work because it did not add additional benefits for our primary goal of making SAM effective for LLMs.

 

5. Experimental Details (Hessian Variance):

Hessian metrics (Table 4/App Table 6) were computed using standard techniques (e.g., Lanczos for $\lambda_{max}$, Hutchinson for trace) averaged over 50 batches from the validation set, where each batch is 256 sequences of length 512, and thus these metrics are aggregated over ~6.5 million tokens. We noticed that even as few as 5-10 batches already gave stable results, but in our Tables we report it with 50 batches for additional precision.

 

6. Specific Questions:

• Fine-tuning/Zero-shot: Future work, but flatter minima (which we achieve, Table 4) often correlate with better transfer and robustness [Liu et al., 2023]. Pruning results (Fig 5) also suggest improved robustness.
• PEFT Compatibility: Likely compatible, but the interaction needs study.
• Explicit Regularization (L1/L2): Yes, func. SAM can be seen as a regularizer & is compatible with like L1/L2 regularization (we already use Weight Decay of 0.1)

 

We believe our paper offers a novel diagnosis and the first effective solution for applying SAM to large-scale LM pre-training, a significant and previously unsolved problem. The empirical gains are meaningful in this context and demonstrate the success of our approach.

---

Let us know if you have additional questions; if we have answered your concerns, we hope you will consider revisiting your review score.

We thank the reviewer for their detailed comments and their positive view of our paper. We address their concerns below

---

 

Theoretical justification of the decomposition

• The decomposition into logit vs. functional sharpness is valid in any setup involving the composition of a loss function with the outputs of a parameterized function.
   
• This describes virtually all scenarios in deep learning, and there is fundamentally nothing which limits its validity to a specific architecture. We will clarify this aspect in the text to avoid any confusion.

 

Lack of theoretical guarantees

• Convergence Analysis: We appreciate the comment. However, rigorous convergence analysis for SAM-type methods is notoriously complex:
  • SAM itself has lacked a general convergence proof in the non-convex case, with initial results only recently published [8].
  • Older convergence analyses [1-5] required significant modification of the algorithm itself or strong assumptions to make progress.
  • The utility of convergence proofs for the design of SAM-like algorithms is also debatable. There is empirical and theoretical evidence that SAM is regularizing sharpness throughout training, rather than just selecting final flat minima to converge to in either the early or late time dynamics only [6].
     
• Generalization Bounds: we note that none of the existing generalization bounds can account for the ineffectiveness of SAM in language modeling. This observation highlights the potential pitfalls of pursuing generalization bounds without adequately accounting for the impact of optimization dynamics.

 

Concerns about efficiency and deployment of methods

Computational Overhead: We agree that the computational overhead needs to be improved; we believe that SAM efficiency methods like LookSAM [7] should be compatible with Functional SAM and can reduce the overhead to a more modest 5-10% — which we hope to demonstrate in future work.

• Rho $\rho$ Tuning: for large models, $\rho$ can be added to scaling studies already used for other hyperparameters (learning rate, weight decay) using detailed experiments at small scales to predict good hyperparameter values at large scales.

 

Reliance on C4

• C4 is widely used for benchmarking LLMs and using it facilitates easier comparison to those works.

• Additionally, as mentioned in the paper, C4 is a clean dataset and thus a hard testground for regularization techniques, like SAM, that aim to improve generalization. We expect the gains to be even higher when the dataset is noisy.

• Besides, C4 is a gigantic dataset significant coverage of most textual corpora on the internet. Thus, the evaluation here being dataset-specific is much less of a risk.

 

Decoder-only Transformers

We focus on this setting to be closer to the industrial use-case, as decoder-only Transformers are really the workhorse of generative LLMs. Downstream tasks can all be modeled on top of these decoder-only models, say, via in-context learning.

 

Practical NLP benefits underdeveloped

We understand your concern, and reiterate that before this work, there was not even a clear path to making SAM effective on language tasks. We believe our work has demonstrated a viable path, and hope to develop a truly practical version of the method in future work. We also recommend you to check Reviewer 8e58’s remarks, where they attest to this precise point.

 

Relation to "Penalizing Gradient Norm for Efficiently Improving Generalization in Deep Learning"

The mentioned paper only uses gradient-norm penalty as a regularizer, and does not analyze it, let alone use the decomposition of the sharpness.

 

Complexity of combining Functional-SAM with preconditioning

The implementation is a trivial and just a one-liner change in code (with negligible overhead). We don’t need to maintain any additional preconditioning statistics, but we can rely directly on that given by the base optimizer (Adam).

 

Code

Yes, we aim to release code by the camera-ready stage.

 

---

 

Let us know if you have any further questions; if we have addressed your concerns we hope you will consider revisiting your review score.

 

 

[1] M. Andriushchenko and N. Flammarion. Towards understanding sharpness-aware minimization. ICML 2022.

[2] P. D. Khanh et al. Fundamental Convergence Analysis of Sharpness-Aware Minimization. NeurIPS 2024

[3] P. L. Bartlett et al. The dynamics of sharpness-aware minimization JMLR 2023.

[4] Y. Dai et al. The crucial role of normalization in sharpness-aware minimization. NeurIPS 2023.

[5] K. Ahn et al. How to escape sharp minima with random perturbations. ICML 2024.

[6] https://proceedings.mlr.press/v202/agarwala23a

[7] https://arxiv.org/abs/2203.02714

[8] https://arxiv.org/abs/2503.02225