Avoiding spurious sharpness minimization broadens applicability of SAM
Authors: Sidak Pal Singh
Hossein Mobahi
Atish Agarwala
Yann Dauphin
Venue: ICML 2025 poster (2025)
Keywords: Sharpness Aware Minimization
SAM
Hessian
LLMs
Generalization
Primary area: deep_learning->algorithms
TL;DR: SAM does not work for LLMs; we diagnose the issue and fix it, resulting in better generalization.

Abstract: Curvature regularization techniques like Sharpness Aware Minimization (SAM) have shown great promise in improving generalization on vision tasks. However, we find that SAM performs poorly in domains like natural language processing (NLP), often degrading performance --- even with twice the compute budget. We investigate the discrepancy across domains and find that in the NLP setting, SAM is dominated by regularization of the logit statistics --- instead of improving the geometry of the function itself. We use this observation to develop an alternative algorithm we call Functional SAM, which regularizes curvature only through modification of the statistics of the overall function implemented by the neural network, and avoids spurious minimization through logit manipulation. Furthermore, we argue that preconditioning the SAM perturbation also prevents spurious minimization, and when combined with Functional SAM, it gives further improvements. Our proposed algorithms show improved performance over AdamW and SAM baselines when trained for an equal number of steps, in both fixed-length and Chinchilla-style training settings, at various model scales (including billion-parameter scale). On the whole, our work highlights the importance of more precise characterizations of sharpness in broadening the applicability of curvature regularization to large language models (LLMs)

Avoiding spurious sharpness minimization broadens applicability of SAM
Sidak Pal Singh
Hossein Mobahi?
Atish Agarwala 2 Yann Dauphin 2
Abstract
commonly held belief is that it is important for models to converge to well-behaved and robust solutions. The 'regu- larity' of the model is often obtained using regularization techniques, which - even in the day and age of LLMs remain an indispensable part of any training algorithm.
Curvature regularization techniques like Sharp- ness Aware Minimization (SAM) have shown great promise in improving generalization on vi- sion tasks. However, we find that SAM performs poorly in domains like natural language process- ing (NLP), often degrading performance even with twice the compute budget. We investigate the discrepancy across domains and find that in the NLP setting, SAM is dominated by regularization of the logit statistics - instead of improving the geometry of the function itself. We use this observation develop an alternative algorithm which regularizes
Some of the most prominent regularization methods include weight decay (Krogh & Hertz, 1991). dropout (Sri- vastava et al., 2014), data augmentation (Ciregan al., 2012; Krizhevsky et al., 2012), Mixup (Zhang, 2017), and curvature- basedd controis (Foret et al., 2020; Wu et al., 2020). In recent years, curvature regularization techniques have gained popularity due to their effectiveness in promoting generalization. These techniques bias learning dynamics to areas of lower curvature (i.e., less sharp regions) in the loss landscape (Chaudhari et al., 2017: Keskar et al., 2017; Foret et al., 2020; Pittorino et al., 2020; Wu et al., 2020). The origins of these curvature or sharpness minimization techniques can be traced back to the classical ideas of minimum description length (Rissanen, 1978; Hinton & Van Camp. 1993; Hochreiter & Schmidhuber, 1997). Lately, their development has been inspired by their success in a large-scale correlational study (Jiang et al., 2019) and in the NeurlPS generalization competition (Jiang et al., 2020).
call FUNCTIONAL-SAN curvature only through modification of the statistics of the overall function implemented by the neural network, and avoids spurious minimization through logit manipulation. Further- more, we argue that preconditioning the SAM perturbation also prevents spurious minimization, and when combined with FUNCTIONAL-SAN it gives further improvements. Our proposed algorithms show improved performance over ADAMW and SAM baselines when trained for an equal number of steps, in both fixed-length and Chinchilla-style training settings, at various seale).
Sharpness minimization has demonstrated significant im- provement on vision tasks. In contrast, there has not been much uptake of these methods in NLP' and, especially for as cornerstone (Brown et al., 2020) a task in NLP as language modeling. Curiously enough, we observe that Sharpness Aware Minimization (SA one of the best studied sharp- regularization methods (Foret al., 2020), shows poor performance here; indeed, its validation metrics are typi- cally worse than ADAMW throughout training (Figure 1), despite using more computation per step. Hence, this raises the following questions which form the basis of our work:
model scales (including On the whole, work highlights the importance of more precise characterizations of sharpness in broadening the applicability of curvature regularization to large language models (LLMs).
1. Introduction
One of the most fundamental questions in machine learning research is: how do we train models that are useful beyond their training data? This question arises in multiple scenar- ios from generalizing to unseen samples, dealing with distribution shift, and fine-tuning on specific domains. A
What are the reasons for SAM's poor perfor- mance in language modeling? How can they be addressed to successfully yield generalization ben- efits, while being equal in cost SAM?
'Google Research Google DeepMind. Correspondence to: Sidak Pal Singh <ssidak@ @ google.com>>.
Towards this end, in Section 3, we perform a novel analysis
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
'The most notable exception of SAM in NI.P is in the fine- tuning scenario (Bahri et al., 2022), where the parameters con- strained to move smaller distances by the very nature of the setup.

---

Avoiding spurious sharpness minimization broadens applicability of SAM
Nanodo 23.9M trained on C4
seek model the input-output relation via a neural network fo(x): Rd - T with learnable parameters f RR,, such that fo R y (xxx)) E D. We take the usual route of empirical risk minimization (Vapnik, 1991) and consider 0- == argmingge (()) with L()) == 1 Σi= {(z;0) 0) and where, S = } is the training set of size n and & denotes the loss function. Hereafter we will consider the loss to be cross-entropy, which is the most popular choice; however, our analyses extend to other loss functions well.
5.00
SAM AdamW
4.75
Validation Loss
4.50
3.55
4.25
3.50
4.00
3.45
3.75
8000
10000
3.50
Under the above setup, Foret et al. (2020) formulates sharpness-aware minimization as the following min-max problem: ming max turbation of the parameters. In particular, the inner maxi- mization is approximated to first order in the perturbation,
2000
4000
6000
8000
10000
Training steps
|ell<p ( + <), where e denotes per-
Figure 1: Evaluation loss curves of ADAMW and SAM for Nanodo decoder-only Transformer model (Liu et al., 2024) on the C4 dataset (Raffel al., 2020).
max L()) + t VoL((),
(I)
|el|sp
of the path SAM takes to sharpness reduction, and show that can be split into two contributions - one which modifies the logit statisties to reduce sharpness, and the other which modifies the geometry of the function itself. We measure these two contributions and find that in vision (where SAM works well) the contributions are relatively balanced: in contrast, in language modeling settings the logit path to sharpness minimization dominates.
which subsequently yields E'(0) = P VoCO) (OLCO//TOOO) VoL(O) as the optimal perturbation. In SAM (Foret et al., 2020), the authors propose making an update along the direction,
VoL(O)
SAA VoC(0 + pE'),
(2)
(ee(())
3. The Dual Routes to Sharpness Minimization
We hypothesize that the functional path to sharpness mini- mization needs to be amplified in the language setting, and in Section 4, develop novel sharpness minimization algo- rithm called FUNCTIONAL-SAN that achieves this. Addi- tionally, we motivate another algorithm, PRECONDITIONED SAM, based on preconditioning, and give a theoretical ar- gument that shows it promotes the functional path.
In this section, we develop a diagnostic tool to understand the failures of SAM in language modeling settings. Our theoretical analysis shows that there are two possible routes to sharpness minimization via SAM- the logit path and the functional path. We derive quantities which can be used to measure the extent to which each route is active. We find that each path is relatively balanced in vision, but in language modeling settings (where SAM performs poorly) the logit path overwhelms the functional path.
In Section 5, we show that FUNCTIONAL- AM and PRE- CONDITIONED SAM provide improvements over baseline ADAMW when trained on the C4 dataset using the Nan- odo Transformer codebase (Liu et al., 2024). Moreover, FUNCTIONAI-SAM and PRECONDITIONED SAM can be combined to yield maximal gains. This resulting com- bination consistently improves validation metrics in both fixed number of steps as well as Chinchilla-like scaling set- tings (Hoffmann et al., 2022) at a variety of model sizes, spanning three orders of magnitude.
3.1. The Penalty Formalization
We begin by looking at PENALTY-SAMA an alternative ver- sion of SAM that is often used in theoretical analyses (An- driushchenko & Flammarion, 2022; Dauphin et al., 2024) and whose gradient matches SAM to first order p
min (()) + + P [VVo(()) 0
We conclude with an extensive discussion, in Section 8, of additional SAM variants that may be useful in other domains as well as important future directions implied by our work.
SP
We can think of the added gradient norm term as a sharpness penalty SP. To understand how SP influences optimization, investigate the structure of its gradient:
② Vo (P P 00 Vo((O))
0 ( a0 voC(O)
2. Setup and Background
c())).
(4)
Let us assume that we are given data points E Z drawn i..d. from some (unknown) distribution D. where the sam- plcs z are input-output tuples (x, y), with the input x E Rdd of dimension and the targets y E R of dimension K. We
The gradient itself can be decomposed by the chain rule as
VeL(O) = VoF VFL((),
(5)
2

---

Avoiding spurious sharpness minimization broadens applicability of SAM
where F(0)= = (x)( ... fo(xx) E R Km collates the output over the entire dataset. The first term (() represents the Jacobian (derivative of the model function outputs with respect parameters); while the second term VpC(0) is the derivative of the loss with respect to the outputs and which comes out to be the difference of the (softmax-ed) logits and the targets.
Tlogit: Trunc: Teross. where the three sum to 1:
||1oogit| 1|
1 |
(ogitt
V SP [| Trune
V8 SP !
(9)
Teross := 2 {Slogie: &zune)
Vo SP||2
Hence based on their values, we can realize whether sharp- ness minimization will prioritise reduction of logit sharp- ness more that of functional sharpness, as well as how correlated they are by looking at the cross term.
Using the product rule, we can rewrite Eqn. 4 as,
a (VEC(O)]
(()
(0) +
00
SP gradient from logit perturbation 6zogst
3.3. Composition of Sharpness Gradient in Practice
00 (oF(O)) VF((
EE"
(6)
We will now analyze SAM's 's behavior in language model- ing, focusing on how the sharpness gradient composition reveals key differences between its application in language and vision tasks. More concretely, we consider next-token prediction task in the case of language using the C4 dataset and image classification in the case of vision. Since the typ- ical vocabulary sizes in language is in the order of tens of thousands, so besides ImageNet-1K, we adopt other datasets like JFT (Sun et al., 2017) and ImageNet-211 K (Rid- nik et al., 2021) to make the settings further comparable in terms of number of outputs. For both settings, we em- play Transformer-based networks, Nanodo (Liu al., 2024), which is simplified version of GPT-2 (Radford et al., 2019). in language modeling and Vision Transformer (VIT, Doso- vitskiy et al., 2021) for vision tasks. Furthermore, in both cases, we train with ADAMW as the optimizer and measure the normalized sharpness gradient contributions (Eqn. 9) throughout training using exact Hessian- -eector products. We present the results for Nanodo, C4 as well as ViT with ImageNet-IK K and JFT in Figure 2, and that on ImageNet- 21K in Appendix A.
gradient from perturbing function Jacobian == ggune
The first term &iggit represents the logit-path to sharpness minimization - directly optimizing the sharpness via the effect of the logits on the loss. In contrast, the second term Orune represents the functional-pauth to sharpness minimiza- tion- that is, sharpness minimization via modification of the Jacobian statistics.
3.2. Contribution of Logit and Functional paths
We can explicitly relate the logit and functional paths to sharpness minimization (frogie and Ornn in Eqn. 6) with decomposition of the Hessian of the loss H Vgc.
Vo SP = diogit + Orunc= = Ho + Hp E" = HE
(7)
Hessian (Schraudolph, 2002), namely: where, we use the Gauss-Newton Decomposition of the
=z=Ho HG + Hr
Vofo vofo(x) (o
Observations. Comparing these figures, we find a stark con- trast between the language and vision settings. For vision, we find that the Tlogit starts close to 0 but that Tlogit and Trune quickly become comparable for most of the training process (Figure 2, leftmost and second from left). In con- trast for language, the logit-related gradient fraction Tlogit is close 1 while the sharpness gradient related to the functional part is much smaller Figure 2 (second from right and rightmost). In all cases, Tcross tends to be negative through most of training (see Figure 4 in the Appendix). This suggests that the two paths to sharpness regularization are antagonistic - taking one path moves you against the other. In language modeling, the dominance of the logit path, combined with negative Taross. means that the overall contribution from the gradient of SP is unaligned, or even anti-aligned, with the functional path to sharpness.
n
K
(8)
12
k=1
where, the Generalized Gauss Newton (GGN) (HG) and the functional Hessian (Hp) (Singh et al., 2021) are the component matrices of the loss Hessian (H.). The GGN term captures the curvature of the linearized model: in con- trast, the functional Hessian (also known as the Nonlinear Modeling Error) captures curvature due to model second derivatives (Dauphin et al., 2024).
Thus, we can interpret the gradient through the logit and functional paths bused on the Hessiun component they de- pend upon the GGN and functional Hessian respectively.
Normalized composition of Sharpness Gradient. To better gauge which mode of sharpness gradient domi- nates, we will measure the following natural quantities,
Logit Sharpness vs Functional Sharpness. The above observations suggest that in NLP settings, SAM is heav- ily biased towards minimizing sharpness by reducing the
3

---

Avoiding spurious sharpness minimization broadens applicability of SAM
VIT, ImageNet1k
VIT JFT
10
10
Sharpness composition
Sharpness composition
10
10-1
10
10-2
10
10
7540
15080
22620
30160
37700
10680
21380
32040
42720
53400
Training Steps
Training Steps
23.9 M. Nanodo, C4
1208.01 M, Nanodo, C4
100
10
Sharpness composition
Sharpness composition
10-1
10-2.
10-2
MM
M M
Tugi
Touec
10
0
1600
3200
4800
6400
41150
62300
123450
164600
Training Steps
Training Steps
Figure 2: Normalized sharpness contributions Tiogit and Trunc show dramatically different trends across modalities. For ViT trained on ImageNet- IK (leftmost) and JFT (second from left), Tlogit starts near 0 but quickly increases comparable magnitude as Trune- For Transformer models trained on C4 (second from right and rightmost), Tlogit " Trune after the first few steps of training. This suggests that the pathways sharpness regularization are more imbalanced in NLP compared to vision settings, which may contribute to the poor performance of SAM in NLP settings. Teross (plotted in Appendix A.1) is usually negative, suggesting the two methods of sharpness regularization tend to be antagonistic.
gradient of the loss with respect to the logits. There is not much sharpness reduction happening due to making the function more belledehaved. While both of these con- tribute to decreasing sharpness, the underlying mechanisms show interesting differences. Recalling the relation of logit- sharpness to the Gauss-Newton term HG, we see that simple but spurious way to decrease it is by making the network over-confident about its predictions. This is be- cause of the presence of the term V3 L, which equates to diag (p:) P.P.T. where p; = softmax(f (x;))) And thus as becomes more one-hot (irrespective of leading to the correct output or the incorrect), the logit sharpness will get reduced.
These observations lead to the hypothesis that if we could encourage reduction of functional sharpness lieu of logit sharpness, then sharpness minimization might fact work for language modeling too. In the next section, op- erate under this hypothesis and derive simple but principled modifications of the SAM update rule.
4. Algorithms to promote the functional path to sharpness minimization
From the PENALTY-S SAM formulation in Eqn. 3 and the form of the corresponding sharpness gradients in Equ. 7, we can see that the functional path to sharpness can, in principle, amplified by manipulating òoegue and drune- However, while PENALTY-SAM algorithm, is less robust than SAM due ill-behaved second derivatives (Dauphin et al., 2024). In contrast to FENALTY-SAM. the perturbation approach in SAM does not explicitly involve derivvtives in the gradient. Therefore, it important to develop strategies that promote
In contrast, the functional sharpness is connected the func- tional Hessian, and the ways of decreasing it (such via modeling the target better p; y;; or reducing the second- derivative of the function with respect to the parameters ||3 + 0 and hence encouraging a functional simplic- ity) are intuitively more desirable, even though setting a particular proportion might be unclear.
competitive with the SAM
4

---

Avoiding spurious sharpness minimization broadens applicability of SAM
the functional path using the methodology taken in SAM.
heterogeneity and improves conditioning.
In this section, we discuss two strategies one direct and the other indirect- - that promote the functional path to sharpness minimization, using algorithms which maintain the benefits of the SAM formulation.
Naively combining SAM and ADAMW creates potential mismatch- - the SAM perturbation is carried out with re- speet to the unpreconditionee geometry. We hypothesize that this mismatch offers grounds for spurious sharpness minimization. To rectify this, we consider preconditioning the perturbation with the inverse of the same second- moment statistics M from ADAMW, resulting in the update:
4.1. FUNCTIONAL-SAM
The simplest way to promote the functional path would be to use a perturbation which aligns more with otune than diogit. We can find such a perturbation by decom- posing the SAM update rule from Eqn. 2, as follows:
g PRECOND SAM = VoL (0 ( 1 (0) VoC())) , (12)
where 7 == p/ p/| Vol(O))l. and g PRECOND SAM is then passed to the rest of Adam in lieu of the gradient ()
VoL(O + pe*) = (oF VD pe)) VoL(O)
Gzogie up first order in p
Another motivation for PRECONDITIONED SAM is that the gradients often align with the principal eigenspaces of HG (Gur-Ari et al., 2018). This would amplify Óoogit over Ogunc, since ólogit = HG E" and E" is parallel to the gradient. Preconditioning E* by H&¹ reduces this effect, thereby promoting the functional path over the logit path (see Appendix B.1 for a more detailed argument). Although H would be the best preconditioner in this regard, is expensive to estimate generally and ADAMW already gives us diagonal estimator in its own preconditioner M- additional cost.
(FF(++pee) ( + Vo)( + O(p²)
Seune up first order
(10)
The first term is &logit up first order in p. and the sec-
ond term is Orune up to first order in p. This suggests the following update rule which we call FUNCTIONAL-SAM,
FUNC-SAM = - F(0 + pc") PP((0)
(11)
where, we discard the &logit contribution. Further, FUNCTIONAL-SAN, uses the same perturbation E" as the SAM formulation, but only perturbs the Jacobian - thus emphasizing the functional path to sharpness as desired. This update rule can be implemented as efficiently as regular SAM using the very same Jector-Jacobian product opera- tions that are used to compute gradient in most autodiffer- entiation frameworks. In Appendix C, we provide the JAX code snippets for SAM and FUNCTIONAL-SAN AM, demon- strating that the difference in their implementation is a mat- ter of a few lines. Further, like SAM, FUNCTIONAL-SAM remains compatible with methods like ADAMW which take gradient and then further process it.
To conclude, this algorithm (Eqn. 12) has only marginally higher computational cost than standard SAM + ADAMW, and provides an indirect way to improve the contribution of the functional path to sharpness minimization.
5. Empirical Evaluation
5.1. Setup
Training Lengths. In order to evaluate language mod- els of multiple sizes, we consider (pre-)traude )training them in two scenarios: (a) when the training length is kept fixed across scales and (b) when the training length is adjusted as per compute-ope corupute-ppimmaltyy considerations (Kaplan et al., 2020; Hoffmann et al., 2022).
Overall, the update rule in Eqn. 11 has the same cost as the SAM update rule, and keeps the benefit of the finite- differences based perturbation approach to sharpness estima- tion that keeps SAM robust meeting all our initial goals.
(a) Fixed-Length Training Scheme. In this setup, we train all the models for 10K steps, which amounts to seeing roughly total of 1.3 billion tokens. The results for this setting are presented in Table 2. (b) Chinchilla-style Training Hori- zon Scheme. Unlike traditional image-classificatoon based scenarios, training language models often involves deter- mining compute-optimal scaling laws, and whereby models are trained corpus size in proportion to their parame- ters. Specifically, we follow the 20x over-training policy suggested in the Chinchilla (Hoffmann et al., 2022) training regime, and use it together with the fixed-depth scheme, as considered in Everett et al. (2024). The corresponding results are shown in Table 3.
4.2. PRECONDITIONEN SAM
In language modeling, Transformers (Vaswani, 2017) are almost exclusively trained with ADAMW or other adaptive methods, and SGD-based training is known to be signifi- cantly worse (Liu et al., 2020). This is commonly attributed to the presence of heterogeneity in the gradients (Liu et al., 2020; Noci et al., 2022; Pan & Li, 2023) and the curva- ture (Zhang et al., 2024; Ormaniec et al., 2024; Jiang et al., 2024) across different "modules" (layer types, layers at dif- ferent depths, etc.), with the idea that ADAMW alleviates
5

---

Avoiding spurious sharpness minimization broadens applicability of SAM
Training Details and Hyperparameters. In both scenar- ios, we consider a batch size of 256 sequences, of maxi- mum length 512, and evaluate model at 5 different sizes: 2M (for prototyping). 23.9M, 42.5M, 117.9M, and 1208M in terms of non-embedding parameters (see details in Ap- pendix A.6), and trained with ADAMW (Kingma & Ba, 2017; Loshchilov & Hutter, 2019) as the underlying opti- mizer on the C4 dataset (Raffel et al., 2020). We used a decoupled weight decay parameter set to 10-4 for all ex- perimental settings. We use the Nanodo (Liu et al., 2024) framework to implement these minimal decoder-only Trans- former models, in Flax (Heek et al., 2024) together with JAX (Bradbury et al., 2018).
efits from preconditioning; we hypothesize that the func- tional route to sharpness also benefits from reduction in heterogeneity across modules and better matching between inner and outer optimization geometries, Hereafter, we will use the preconditioned variant for FUNCTIONAL-SS AM, in- stead of just the plain FUNCTIONAL-SAN -AM, to reduce the computational costs associated with testing at larger model scales. This relationship merits further study, which is out- side the scope of our work.
5.3. Results at Multiple Model Scales
Hyperparameter Setup. For thoroughness, at each model size separately, we simultaneously tune learning rate and perturbation strength p in our experiments, even though this may be difficult at even bigger model sizes or when faced with computational constraints. The upside of this much harder testing ground is to provide us the key assurance that the resulting gains cannot be obtained through other means such as strong hyperparameter tuning of baselines. These results can be found in Tables 2 and 3.2
5.2. Comparison of direct and indirect approaches
Before carrying out an extensive evaluation across differ- ent model scales, we do initial prototyping on a smaller model size. This lets us save computational resources and scale up the most promising methods. In particular, we are interested in knowing which methods out of direct and indi- rect approaches, and their combinations, are most relevant. In Table 1, we present results on a model with 2M non- embedding parameters trained in the fixed-length regime.
In the case of billion-plus parameter model, this procedure however amounted to extremely high computational costs, and so we instead tuned the learning rate for the baseline A first. Subsequently, we tuned for all SAM meth- ods based on this corresponding optimal learning rate for AdamW.
Table I: Effect of Preconditioning the SAM perturbation on FUNC-SAM with 2M parameter model when trained for 50K steps and at peak learning rate of 0.001. Lower is better.
Observations. We find that in both the training regimes, our proposed algorithms, namely PRECOND FUNCTIONAL.- SAM and PRECOND S SAM, significantly outperform SAM, at all the model scales in Tables 2 and 3. The gains typi- cally range from 10-3 3 - -10 2, are largest for 117.9M parameters. We find that SAM performs the worst of the lot, even worse than ADAMW while being 2x computa- tionally expensive. The numbers reported for SAM repre- sent the smallest p tested (0.1), since it monotonically gets worse with increasing p. Overall, we find that PRECOND FUNCTIONAL-S/ AM achieves the best results, followed by a mix of PRECOND SAM and ADAMW.
METHOD
EVAL LOss
ADAMW
3.901
SAM
3,910
PRECOND SAM
3.889
FUNC-SAM
3.880
PRECOND FUNC-SAM
3.861
We notice that the plain version of FUNCTIONAL-SAM. outperforms both ADAMW and preconditioned SAM. But, FUNCTIONAL-SAM can be further improved by using pre- conditioning alongside, yielding significant improvement in terms of evaluation loss. Also, it should be noted that the standard deviation across different seeds is typically around ~ 0.002. Moreover, mid-third decimal differences in evaluation loss are typical of the gains provided by new optimization methods for pre-training LLMs look like, e.g., in CASPR (Duvvuri et al., 2024), and even in newer variants of attention (Leviathan et al., 2024), as well as when fus- ing I.LMs (Mavromatis et al., 2024) or when deduplicating training corpus (Lee et al., 2021). Thus, the kind of gains shown in Table are indeed significant.
At the billion scale, PRECOND SAM and SAM show susceptibility to numerical instabilities (yielding NaNs) across training regimes, whereas PRECOND FUNCTIONAL-SAM is significantly more robust. In the relatively cheaper fixed-length runs where we could carry out additional investigation, we find that warming up the perturbation radius or a lower learning rate choice can sometimes mitigate these issues for PRECOND SAM and SAM, however, they are still outperformed by our best
Foor brevity, we omit the standard deviation for the basclines in these Tables. But in the Appendix, the reader can also find the full table with standard deviations for all baseline methods well. The standard deviation for the baselines results nevertheless very similar.
The experiments suggest that FUNCTIONAL--SAN also ben-
6

---

Avoiding spurious sharpness minimization broadens applicability of SAM
Table 2: Evaluation loss comparison of different methods (10K steps) training setup, where the op-
5.4. Discussion
in a timal learning rate is found for each model size separately. Results are averaged over seeds. For completeness, the optimal values of the perturbation strength for PRECOND FUNCTIONAL-SAL ALMM at the various scales are respectively 0.5, 0.5, 0.4, 0.4. Lower is better.
The above-mentioned results demonstrate that the issue with SAM in NLP can be successfully resolved through PRECOND FUNCTIONAL-SAM. The significant improve- ments in validation performance above are especially intriguing if we bear in mind that these are obtained (a) even when using a clean corpus such as C4 and (b) training in an online fashion where no batch is seen more than once due to massive size of the C4 corpus (even when training for longer duration like in the Chinchilla setup). These ideal conditions which may not always hold. This suggests that further improvements in generalization may occur when training on noisy corpora or in multipass settings; we leave exploration of these potential effects for future work.
PRECOND
PRECOND SAM
SIZE
SAM
ADAMW
FUNC-SAM
23.9 M
3.425 =0.5x10-
3,431
3,450
3.430
42.5 M
3.344 #2,2x10-3
3.350
3.366
3.349
117.9 M
3.218
10-3
3.224
3.257
3.228
1208 M
3.048 #1.x100-8
3.059
NaN
3.056
All in all, these results confirm the benefits imparted by FUNCTIONAL-SAM and PRECONDITIONED SAM over SAM, and these improvements in generalization over ADAMW restore the promise of sharpness regularization.
method, PRECOND FUNCTIONAL-SAM, which obtains a lower evaluation loss and does not need a perturbation warm-upinn in either training regime.
6. Ablatiooo Studies
Perturbation strengths. the above-mentioned results, the best values for SAM typically occur around perturbation radius م 0.1, which, in our experiments, happens to be the smallest nonzero perturbation radius considered. However, if we employ larger perturbation radii, the performance of SAM rapidly deteriorates, as shown in Figure 3 (left), which moreover suggests that the optimal value of perturbation size p is 0 - i.c., not using SAM at all. In contrast, the optimal value of perturbation for PRECOND FUNCTIONAL-SAM tends to be much larger, across all values of learning rate shown in Figure 3 (right). This explains why we needed to modify the relative magnitudes of the sharpness contributions using PRECOND FUNCTIONAI,-S AM the logit term degrades performance at even small p. overwhelming the potential gains from the functional path to sharpness minimization.
Table 3: Evaluation loss comparison of different methods in Chinchilla like training (Everell et al., 2024), where the opti- mal learning rate is found for each model size separately. All results have been averaged over 5 seeds, except for the more expensive billion parameter runs which are based on a single seed. For completeness, the optimal values of the perturba- tion strength for PRECOND FUNCTIONAL-SAN at the vari- ous scales are respectively 0.7, 0.5, 0.4, 0.5. Lower is better.
PRECOND
PRECOND
SIZE
SAM
ADAMW
FUNC-SAM
SAM
23.9 M
3.492 =2.9x10-
3.498
3.516
3.497
42.5 M
3.340 +2.0x10
3.348
3.359
3.344
117.9 M
3.070 12-2x10-3
3.074
3.094
3.079
Non-linearity Choice. While all the above experimental results utilize GeLU non-linearity, we also carry out exper- iments with ReLU as the choice of non-linearity. These results can be found in Table 8 of the Appendix, but the key observation is that similar improvements are observed in the evaluation loss when using PRECOND FUNCTIONAL-SAM for ReLU based architectures as well.
1208 M
2.557
2.562
NaN
2.561
We note that for the 1208M parameter models, was tuned more coarsely due to computational costs. Despite less tuning, the gains for PRECOND FUNCTIONAL-SAN persist at this scale the Chinchilla training horizon experiments.
1.1. Sharpness of Final Solutions
In Tables and 7 of the appendix, we also illustrate how the results fare if common learning rate of 0.0001 is used across all model scales. Interestingly, we find that PRECOND FUNCTIONAL-SAM is much less sensitive to the lack of rightly tuned learning rate than the baselines. As a matter of fact, in such settings we find the corresponding gains to be even bigger.
Here, we confirm that the generalization benefits imparted by PRECOND FUNCTIONAL- SAM are brought about convergence a solution with lower curvature, as shown in the Table 4 for the 23.9M model. Similar results also be found other model sizes, and Table 9 of the Appendix shows them for the 117.9N model.
7

---

Avoiding spurious sharpness minimization broadens applicability of SAM
0,0055
0.0055
0.0078
0.0078
Loss
Eval
352
0.0156
0.158
0.0211
0.0221
350
00
2.
Figure 3: Effect of increasing perturbation strength for SAM (left) and PRECOND FUNCTIONAI-SAN (right) across various peak learning rates, in an equal compute setup. Each cell shows the evaluation loss, averaged over 5 seeds, at the corresponding values of learning rate and perturbation radius. The best performing value of perturbation strength at a given learning rate is marked by white squares, while the overall best learning rate and perturbation strength pair is marked in green. We see that with SAM any non-zero p does worse than the baseline (p = 0), while PRECOND FUNCTIONAL.-S with the same compute costs, shows improvements all values of peak learning rates (Nanodo trained C4, 23.9M parameters, Chinchilla like setup).
Table 4: Comparison of different methods based on Hessian He and GGN Hg maximum eigenvalue and trace for the 23.9M model trained per Chinchilla like training setup, with peak learning rate of 0.001. Lower is better for all metrics. The best entry is in bold, the second best is underlined.
2024; Li & Giannakis, 2024). Some of these methods have shown small improvements vision tasks. We believe our decomposition approach is orthogonal to this line of research. Therefore, the obtained PUNCTIONAL-SAN al- gorithm is substantially different from prior work.
Studies exploring preconditioning for SAM. Other work has also suggested that the perturbation step in SAM should be taken in an alternative geometry. Our approach to preconditioning is most similar to Fisher SAM (Kim et al., 2022), and the concurrent work of Zhang et al. (2025) which describes a more general preconditioning scheme for SAM. Our key insight is that useful take the SAM pertur- bation the exact same geometry used by the optimizer, which can be accomplished for negligible cost in the case of Adam and its variants. Furthermore, our work is primarily driven by the problem of making SAM work in language modeling, which is far from the focus these other works.
METHOD
EVAL Loss
Amax (Hc)
(r(H
tr (r(H))
ADAMW
3.688
10.61
4897.52
4745.03
SAM
3,706
2.71
3324.58
3231.46
PRECOND 3.663
5.62
3182.78
3097.87
SAM
PRECOND 3.631
6.20
2687.12
2503.86
FUNC-SAM
Role of the indefinite Hessian term. At a more conceptual level, our study aligns with recent works (Singh et al., 2021; Dauphin al., 2024) which underscore paying more impor- tance to the functional Hessian, the understudied indefinite term of the Hessian in the Gauss-Newton decomposition, as opposed to focusing solely on the positive semi-definite GGN term as suggested by prior studies (Sagun et al., 2018; Papyan, 2019; Jacot et al., 2020). The decomposition the sharpness gradient into logit and functional modes, along with the demonstrated significance of FUNCTIONAL-SAA in NLP tasks (which is closely tied to the functional Hes- sian), highlights the risks of over-reliance on the GGN and the corresponding outlier spectrum of the Hessian -ecaaa--sss espe- cially in the context regularization of sharpness.
We notice that all sharpness minimization methods yield lower value of the maximum cigenvalue and the trace, as compared ADAMW. Interestingly, we also find that SAM results in the lowest Hessian maximum eigenvalue, though it performs much worse than ADAMW when evaluating the validation set. This further highlights how SAM, by default, in language modeling tasks is set up to minimize sharpness spuriously. In contrast, we see that for both of our proposed methods, the improved generalized performance comes hand-in-hand with better landscape properties of their solutions.
7. Related work
SAM NLP. Prior works building on SAM NLP have been restricted the fine-tun fine-tuning setting (Bahri et al., 2022). domain transfer (Sherborne et al., 2024) small-scale machine translation setups (Li & Giannakis, 2024). To
Improved variants of SAM in Vision. An extensive line of work has attempted to propose better definitions of sharp- ness - particularly those which are less sensitive to details of parameterization (K won al., 2021: Tahmasebi et al.,
8

---

variants Avoiding spurious sharpness minimization broadens applicability of SAM
the best of our knowledge, SAM has hitherto not been successfully applied to language modeling, particularly in any scaling setting.
simple demonstration with one-shot pruning in Figure 5 where we see that FUNCTIONAL- shows more grace- ful degradation as opposed to ADAMW with increasing sparsity. We leave a detailed study to future work.
8. Discussion and Future Work
Efficiency. A common drawback of sharpness minimization methods is that they require twice the gradient computation per step, and hence are twice as expensive as compared to other optimizers like ADAMW. In this paper, our singular focus has been to address the ineffectiveness of SAM in language modeling, which we have been able to carry out successfully through the proposed FUNCTIONAL- AAAAAL--MM SAM and its preconditioning variants. However, these algorithms still share the same 2x computational burden like SAM itself, and thus additional work is required in the future to make it suitable for deployment. Thankfully, since FUNCTIONAL- S AAM is structured similarly to SAM, most advancements in efficient implementations of SAM should be useful for FUNCTIONAL-SA AM as well. Given the plethora of recent work in this area for vision (Du et al., 2022; Liu et al., 2022; Becker et al., 2024; Xie et al., 2024), we are optimistic about the development of efficient FUNCTIONAL- for language modeling in the near future.
There are several important aspects that we would like to elaborate on, which could spark interesting future work:
Interpolating smoothly between paths to sharpness mini- mization. In this work, we focused on FUNCTIONAL-SAM due to our diagnosis of issues in language modeling. How- ever, one can imagine that there might be other domains gives spurious minimization,
where FUNCTIONAL-SAN and the alternative "LMGIT-AMM may need to be empha- sized. Our implementation of FUNCTIONAL-SAM can be extended to a more continuous " NGLE-SAM" which can smoothly interpolate between the extremes, which we discuss in detail in Appendix D. Though we did not find A ANGLE-S AM necessary in the language modeling setting, there might be other dateell-dtttsttopimnize triples where it would prove beneficial.
The Perturbation Scope. One of the key takeaways from this work is how scoping the perturbation to the level of the function Jacobian suffices to enable the benefits of SAM in NLP tasks. This scoping can be generalized to any set of paths or branches in the computational graph. For example, the output of network with residual connections can be decomposed into multiple streams like fxx) = fi (x) + fe(x) +. +m fm(() (x) and practitioners can choose to either have the perturbation go through all or some of them. This perspective opens the door to many more creative levels of perturbation scoping and new regularization techniques.
Data-bound scenarios and model size constraints. In any case, we would like to emphasize that our primary compar- ison point in the paper is that at equal number of steps which is becoming an increasingly relevant scenario. This is because industrial settings are often bound by data avail- ability or a limited model size (e.g.. inference constraints). Therefore, here extra training time is acceptable for better final quality.
9. Conclusion
Optimal Perturbation Transfer. At the largest scales, extensive uning of the perturbation radius is unfeasible; FUNCTIONAL-SAAM and its variants will benefit from improved parameterizations that optimally transfer the per- turbation radius across model scales. The recently explored layerwise perturbation scaling regime for SAM in vision tasks (Haas et al., 2024) may be promising for language too.
Our work thoroughly highlights the ineffectiveness of SAM for language modeling and uncovers the underlying reasons behind such an occurrence. We show that this arises because sharpness reduction in NLP settings is prone logit ma- nipulation, and hence spurious sharpness minimization, rather than being driven by promoting the simplicity of the network function. Based on this insight, we propose address- ing this issue via FUNCTIONAL-SAN and preconditioning. We demonstrate that both these simple but principled modi- fications to SAM restore its generalization properties across multiple model scales.
Utility for Downstream tasks and deployment. An added advantage of training with sharpness minimization methods is that the resulting flatter solutions can adapt more gracefully to the post-training pipeline, like downstream fine-tuning model compression. These benefits have been demonstrated in prior work such as (Liu et al., 2023), where a lower Hessian trace at the solution has been shown to correlate better with performance on downstream tasks, and we expect similar benefits to also hold with FUNCTIONAL- SAM and its variants, given the obtained flatter solutions (see Tables 4, 9). Likewise prior work (Na et al., 2022) has also advocated the use of such methods when subsequent model compression is intended, and we also present a very
More broadly, we believe that novel foray into functional and logit modes of sharpness reduction will reinvigorate the existing research into SAM, and pave the way for advanced curvature regularization techniques. Lastly, we are excited about the nuanced characterization of sharpness introduced here and hope that it advances our fundamental understand- ing of sharpness, its dynamics, and the broader nature of loss landscapes.
9

---

Avoiding spurious sharpness minimization broadens applicability of SAM
Impact Statement
for image recognition scale, 2021. URL https:
/ //arxii.rrg/abs//0000
This paper presents work whose goal to advance the field of Machine Learning. There are many potential societal consequences of our work, none which feel must be specifically highlighted here.
Du, J., Zhou, D., Feng, J., Tan, V. Y., and Zhou, J. T. Sharpness-aware training for free. arXiv preprint Sti::2205./4088. 2022.
Duvvuri, S., Devvrit, F,, Anil, R., Hsieh, C.-J,, and Dhillon, I. S. Combining axes preconditioners through kronecker approximation for deep learning. In The Twelfth International Conference on Learning Represen- lations, 2024.
References
Andriushchenko, M. and Flammarion, N. Towards derstanding sharpness-aware minimization, 2022. URL org/ abs/ 2200. 06232.
Bahri, D., Mobahi, H., and Tay, Y. Sharpness-aware mini- mization improves language model generalization, 2022.
Everett, K., Xiao, L., Wortsman, M., Alemi, A. A., Novak, R., Liu, P. J., Gur, I., Sohl-Dickstein, J., Kaelbling, L. P., Lee, J., et Scaling exponents across parameterizations and optimizers. arXiv preprint arXiv:2407.05872, 2024.
URL https::///// /aaaiviiiiv.////.///aaaa/aaaaava orgo http/bbb/211000 08529.
Becker, M., Altrock, F., and Risse, B. Momentum-sam: Sharpness aware minimization without computational overhead, 2024. URL https: /ttp:://aavvvvvo org/ abs/ 2401.12033.
Foret, P., Kleiner, A., Mobahi, H., and Neyshabur, B. awaree minimization for efficiently improving generalization. arXiv preprint arXiv:0010001412 2020.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D.. Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q. JAX: composable transformations of Python+NumPy programs, 2018. URL http:/ / thub. com/ cww/jax-mmaaeecccceeeeeaae jax-ma/ jaxx
Gur-Ari, G., Roberts, D. A., and Dyer, E. Gradient descent happens in tiny subspace. arXiv preprint 2018.
Haas, M., Xu, J., Cevher, V., and Vankadara, L. L.C. pP2 Ef- fective sharpness aware minimization requires layerwise perturbation scaling, 2024. URL org/
Brown, T. B., Mann. B., Ryder, N., Subbiah, M., Kaplan, J. Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert -osse A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D. M., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. Language models are few-shot learners, 2020. URL https:// arxiv. wwt//bs//////oooo.ooo...o...eeet aawaaas20000 14165.
00075.
Heek, J., Levskaya, A., Oliver, A., Ritter, M., Rondepierre, B., Steiner, A., and van Zee, M. Flax: A neural network library and ecosystem for JAX, 2024. URL http: github. com/ gwwgle/ flax.
Hinton, G. E. and Van Camp, D. Keeping the neural net- works simple by minimizing the description length of the weights. In Proceedings of the sixth annual conference on Computational learning theory, pp. 5-13, 1993.
Chaudhari, P., Choromanska, A., Soatto, S., LeCun, Y., Bal- dassi, C., Borgs, C., Chayes, J., Sagun, L., and Zecchina, R. Entropy-sgd: Biasing gradient descent into wide val- leys, 2017. URL httpt:tt///// //ap:ii///r/...//eeeeeeeae http///xxii..ggaas/s/asaa/ wwt./boo////eeeooooooooo 01838.
Hochreiter, S. and Schmidhuber, J. Flat minima. Neu- ral Computation, 9: 9:1-42, 1997. URL https: //api. semanticscholes org :orpusID: 733161.
Ciregan, D., Meier, U., and Schmidhuber, J. Multi-columm deep neural networks for image classification. In 2012 IEEE conference on computer vision and pattern recogni- tion, pp. 3642-3649. IEEE, 2012.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., de Las Casas, D.,, Hendricks, L. A., Welbl, J., Clark, A.. Hennigan, T., Noland, E., Millican, K., van den Driessche, G., Damoc, B., Guy, A., Osindero, S., Simonyan, K., Elsen, E., Rae, J. W., Vinyals, O., and Sifre, L. Training compute-optimal large language models, 2022. URL https: http///xxxvvvr..e abs / 11.20...5555
Dauphin, Y. N., Agarwala, A., and Mobahi, H. Neglected hessian component explains mysteries in sharpness regu- larization, 2024.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner. T., Dehghani, M., Minderer, M.. Heigold, G., Gelly. S., Uszkoreit, J., and Houlsby, N. An image is worth 16x 16 words: Transformers
Jacot, A., Gabriel, F., and Hongler, C. The asymptotic spectrum of the hessian of dnn throughout training, 2020. URL htt.tt://// //apxii////xvvv///// http///aaxvv.//aaae///
0. 02875.
10

---

Avoiding spurious sharpness minimization broadens applicability of SAM
Jiang, K., Malik, D., and Li, Y. How does adaptive optimiza- tion impact local neural network geometry? Advances in Neural Information Processing Systems, 36, 2024.
Li, B. and Giannakis, G. Enhancing sharpness-awart aware op- timization through variance suppression. Advances in Neural Information Processing Systems, 36, 2024.
Jiang, Y., Neyshabur, B., Mobahi, H., Krishnan, D., and Ben- gio, S. Fantastic generalization measures and where to find them, 2019. URL https: /ttp/aaxiivv/..////v//eee/.a.av. org/ abs/ 1912.0217
Liu, H., Xie, S. M., Zi, and Ma, T. Same pre-training loss, better downstream: Implicit bias matters for language models, 2023. URL, https: /oppnrvviee net// forum? Fuu cwABMu.
Jiang Y., Foret, P., Yak, S., Roy, D. M., Mobahi, H., Dz- iugaite, G. K., Bengio, S., Gunasekar, S., Guyon, I,, and Neyshabur, B. Neurips 2020 competition: Pre- dieting generalization in deep learning, 2020. URL ht https: arxii.iiivv..... org/
Liu, L., Liu, X., Gao, J., Chen, W., and Han, J. Understand- ing the difficulty of training transformers. arXiv preprint arXiv:2004. 008-2444440 2020.
Liu, J., Novak, R., Lee, J., Wortsman, M., Xiao, I.., Ev- erett, K., Alemi, A. A., Kurzeja, M., Marcenac, P., Gur, I., Korbblihh, S., Xu, K., Elsayed, G., Fischer, I., Pen- nington, J., Adlam, B., and Dickstein, J.-S. Nanodo: A minimal transformer decoder-only language model im- plementation in JAX., 2024. URL com con/ggoole--oo hatpoggleedeeemmin/naaneeeeee
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language mod- els. 2020. URL https:ttpt.... / arxiv. orgoor////oooo. hwagg//b//// 2001. 08361.
Keskar, N. S., Mudigere, D., Nocedal, J., Smelyanskiy, M., and Tang, P. T. P. On large-batch training for deep learning: Generalization gap and sharp minima, 2017.
Liu, Y., Mai, S.. Chen, X., Hsieh, C.-J., and You, Towards efficient and scalable sharpness-aware minimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 12360-12370, 2022.
URL https: http//axxivvv/// org. abs 1609. 04836.
Kim, M., Li, D., Hu, S. X., and Hospedales, T. Fisher sam: Information geometry and sharpness aware minimisation. In International Conference on Machine Learning, pp. 11148-11161. PMLR, 2022.
Loshchilov, and Hutter, Decoupled weight decay regu- larization, 2019. URL Ut htt.tttsss/aaaii..rrrrrtteeee 1711.05100
Mavromatis, C., Karypis, P., and Karypis, G. Pack of Ilms: Model fusion at test-time via perplexity optimization. arXiv preprint arXiv:2404.1155: 2024.
Kingma, P. and Ba, J. Adam: method for stochastic op- timization, 2017. URL httpt::t//// /tetts://arxiv,.. org/ / 1412. 6980.
Na, C., Mehta, S. V., and Strubell, E. Train flat, then compress: Sharpness-awarr minimization learns more compressible models. In Findings of the Asso- ciation for Computational Linguistics: EMNLP 2022, pp. 4909 4936. Association for Computational Lin- guistics, 2022. doi: 10. 10.16655ww/20222.nnnnn--mmm 361. URI http: / 2022 findings-nnnniii peaiiinnnsemmll...m 361.
Krizhevsky, A., Sutskever, I., and Hinton, G. E. Imagenet classification with deep convolutional neural networks. Advances in neural information processing systems, 25, 2012.
Krogh, A. and Hertz, J. simple weight decay can im- prove generalization. Advances in neural information processing systems, 4, 1991.
Noci, L., Anagnostidis, S., Biggio, L.,, Orvieto, A., Singh, S. P., and Lucchi, A. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. Advances in Neural Information Processing Systems, 35: 27198-27211, 2022.
Kwon, J., Kim, J., Park, H., and Choi, I. K. Asam: Adaptive sharpness-aware minimization for scale-invariant learn- ing of deep neural networks. In International Conference on Machine Learning, pp. 5905-5914, PMLR, 2021.
Lee, K., Ippolito. D., Nystrom, A., Zhang, C., Eck, D., Callison-Burchl C., and Carlini, N. Deduplicating train- ing data makes language models better. arXiv preprint MoXiv:2107.06499, 2021.
Ormaniec, W., Dangel, F., and Singh, S. What does mean be transformer? insights from theoretical hessian analysis, 2024. URL https:/: http://axxvvvvroee abs / 221224100 10986.
Leviathan, Y., Kalman, M., and Matias, Y. Selec- tive attention improves transformer. arXiv preprint arXiv:2410.02703 2024.
Pan, Y. and Li, Y. Toward understanding why adam con- verges faster than sgd for transformers. arXiv preprint arXiv:2306.0020 2023.

---

Avoiding spurious sharpness minimization broadens applicability of SAM
Papyan, V. The full spectrum of deepnet hessians at scale: Dynamies with sgd training and sample size, 2019. URL atttps://axxx /aaxivv..v////////////////xaatv org/ abs/ 11111.111. 07062.
Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, I., and Salakhutdinov, R. Dropout: simple way to prevent neural networks from overfitting. The journal of machine learning research, ((()::199-11958 2014.
Pennington, J. and Bahri, Y. Geometry of neural net- work loss surfaces via random matrix theory. In Precup, D. and Teh, Y. W. (eds.), Proceedings of the 34th International Conference on Machine Learn- ing, volume 70 of Proceedings of Machine Learning Research, pp. 3798-2806. PMLR, 06-11 Aug 2017. / /pt.p/eeoocceeeeeeeeeeee mlr paa...presssvv penningtonl7a. html.
Sun, C., Shrivastava, A., Singh, S., and Gupta, A. Revisit- ing unreasonable effectiveness of data in deep learning era, 2017. URL https: /arxiv.ivvv./////// org/ abs/1707. 02968.
Tahmasebi, B., Soleymani, A., Bahri, D., Jegelka, S., and Jaillet, P. A universal class of sharpness-aware minimiza- tion algorithms, 2024. URL https: : htttpp://aaxiv.ort/ abs / 2406 03682.
URL hwwpt::tp////
Pittorino, F., Lucibello, C., Feinauer, C., Malatesta, R. M., Perugini, G., Baldassi, C., Negri, M., Demyanenko, E., and Zecchina, R. Entropic gradient descent algorithms and wide flat minima. CoRR, abs/2006.078897 2020. URL https:/ / aaxii.iivvvv//vv//aaaaaaiiaaavva wwwpa.wwgggaaaaaaaat.......w.att htgg/bb//at /tss//2000 2..6. 07897.
Vapnik, V. Principles of risk minimization for learning the- ory. Advances neural information processing systems, 4, 1991.
Vaswani, A. Attention all you need. Advances Neural Information Processing Systems, 2017.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Wu, D., tao Xia, S., and Wang, Y. Adversarial weight perturbation helps robust generalization, 2020. URL 05884.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, J. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21 (140):1-67, 2020.
https://arxin wtwp:/orrgaabb/////////
Xie, W., Pethick, T., and Cevher, V. Sampa: Sharpness- aware minimization parallelized, 2024. URL https:
/ //aaxii.or//ab////2///... 1068.
Zhang, H. mixup: Beyond empirical risk minimization. arXiv preprint arXiv: 1710v01710.094122 2017.
Ridnik, T., Ben--Bnnccle E., Noy, A., and Zelnik-Mani L. Imagenet-2Ik pretraining for the masses, 2021. URL https: //axxiv. org/ hbs//ab///// 2214 10972.
Zhang, Y., Chen, C., Ding, T., Li, Z., Sun, R., and Luo, Z.-Q. Why transformers need adam: A hessian perspec- tive, 2024. https: /aaxiv..//iivve//a org/ abs/ /2022. 16798.
Rissanen, J. Modeling by shortest data description. Auto- matica, 14(5):465-47, 1978.
Sagun, L., Evei, U., Guney, V. U., Dauphin, Y., and Bottou, L. Empirical analysis of the hessian of over-parametrized neural networks, 2018. URL https::///// wttaaaixvvo........a..a abs/1706. 04454.
Zhang, Y., Li, B., and Giannakis, B. Precondi- tioned sharpness-aware minimization: Unifying anal- ysis and novel learning algorithm. arXiv preprint arXiv:2501.06603 2025.
Schraudolph, N. N. Fast curvature matrix-vector products for second-order gradient descent. Neural Computation, 14:1723-1738, 2002.
Sherborne, T., Saphra, N., Dasigi, P., and Peng, H. Tram: Bridging trust regions and sharpness aware minimiza- tion, 2024. URL htL.htts://arxiii http://aaaa....o////////3/ 22310 03646.
Singh, S. P., Bachmann, G., and Hofmann, T. Analytic insights into structure and rank of neural network hes- sian maps. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, W. (eds.), Advances in Neural In- formation Processing Systems, 2021. URL ht https: //openreview.oo wtt.eeww.ee.orrreeee ddd.ddoooo.7M00
12

---

Avoiding spurious sharpness minimization broadens applicability of SAM
A. Additional Results
A.I. Sharpness Composition Plots
ViT, ImageNettk
VIT ImageNet21K
109
109
Sharpness composition
Sharpness composition
10-1-
10-1-
-10-1-
-10-1
-10
-10
7540
15080
22620
30160
37700
11580
23160
34740
46320
57900
Training Steps
Training Steps
ViT JFT
23.9 M, Nanodo, C4
109
109
Sharpness composition
Sharpness composition
10
10-1
Mmmnrmmn mm@m
-10-1
-10-1
-10º
-10º
10680
21360
32040
42720
53400
1600
3200
4800
6400
Training Steps
Training Steps
1208.01 01 M, Nanodo, C4
10"
Sharpness composition
10-1
M M
10-2
M
M "ymmaa m Mww
-10-2
Nogst
-10-1
Trunc
Teross
-100
41150
82300
123450
164600
Training Steps
Figure 4: Sharpness contributions Tlogit, Trune and Teross for various datasets. Teross tends to be negative for most of training.
13

---

Avoiding spurious sharpness minimization broadens applicability of SAM
A.2. Detailed Table Results
Below in Table 5, include the standard deviation for each of the baseline method well, in addition to that listed for PRECOND FUNCTIONAL-SAM in Table 3.
Table 5: Evaluation loss comparison of different methods in Chinchilla like training (Everell et al., 2024), where the optimal
learning rate is found for each model size separately. All results have been averaged over 5 seeds, except for the more expensive billion parameter runs which are based on single seed. Lower is better.
PRECOND
PRECOND
SIZE
SAM
ADAMW
FUNC-SAM
SAM
23.9 M
3.492 +2.9x10-3
3.498
#3.6x10-3
3.516
12.9x10-3
3.497
=2.9x10-3
42.5 M
3.340
12.0x10-3
3.348
11.7x10-3
3.359
13.4x10-a
3.344
_2.5x10
117.9
3.070
1-22x10-
3.074
+10010-
3.094
#3.1x10-8
3.079
=2.7x10-3
1208 M
2.557
2.562
NaN
2.561
A.3. Results across Scales common learning rate
Table 6: Evaluation loss comparison of different methods in fixed-length (10K steps) training setup at common learning rate of 0.001. Lower is better.
PRECOND FUNC-SAM
PRECOND SAM
SIZE
SAM
ADAMW
23,9 M
3.527
3.547
3,587
3,565
42.5 M
3,414
3,439
3,462
3,451
117.9 M
3.252
3.266
3.288
3.280
1208 M
3,054
NaN
3.081
Table 7: Evaluation loss comparison of different methods in Chinchilla like training (Everett al., 2024) common learning rate of 0.001. Lower is better.
PRECOND FUNC-SAM
PRECOND SAM
SIZE
SAM
ADAMW
23.9 M
3.631
3.663
3.706
3.688
42.5 M
3.414
3,435
3,460
3.445
117.9 M
3.096
3.108
3.126
3.120
1208 M
2.612
NaN
NaN
2.627
14

---

Avoiding spurious sharpness minimization broadens applicability of SAM
A.4. Results across non-linearities
Table 8: Comparison of SAM variants across non- non-linearities for Nanodo model with 2M (non-embedding parameters) on C4 dataset.
NON-LINEARITY
METHOD
EVAL Loss
PRECOND
GeLU
3.8614
FUNCTIONAL-SAN
PRECOND
GeLU
3.8894
SAM
GeLU
ADAMW
3,9069
PRECOND
ReLU
3.8777
FUNCTIONAL-SAM
PRECOND
ReLU
3.8937
SAM
ReLU
ADAMW
3,9145
A.5. Hessian spectra results at another model size
Table 9: Comparison of different methods based on Hessian H and GGN HG metries for the 117.9M model trained as per Chinchilla like training setup.
METHOD
EVAL LOss
Amax (H&)
tr(Hz)
tr(He)
PRECOND
3.096
8.884
2790.650
2746.559
FUNC-SAM
PRECOND
3.108
7.415
2920.445
2060.214
SAM
SAM
3.126
3.381
2235.759
2222.779
ADAMW
3.120
9.259
3299.497
3285.881
A.6. Architectural Details.
The 23.9M 33.9M, 42.5M, 117.9M, and 1208M models have the same depth of 6, and whose width has been scaled together with the number of heads. In particular, these correspond to h = 9, 12, 20, and 64 heads per block and the width m scales as x m. The 2M model used for prototyping has depth 3, 4 heads per block, width m 256 and MLP dimension 1024.
m 64 x h, and the MLP dimension
15

---

Avoiding spurious sharpness minimization broadens applicability of SAM
A.7. Downstream benefits of sharpness minimization
159.73 M model, Nanodo
159.73 M model, Nanodo
3.15
10
eval loss
eval loss
3.13
3.12
tno sam, precond loss satossammppprrrree
tnc sam, precond
8.11
lass_sam pros_sam_rreeee precood adamm
3.10
adarm
loss_som
loss_som
0.00
0.05
0.10
0:15
0.20
0.3
0.4
0.5
0.8
0.7
0.0
0.9
pruning ratio
pruning ratio
Figure 5: Effeect of one-shot (unstructured) global magnitude pruning: We see that sharpness minimization methods tend to degrade more gracefully as increasing number of parameters are pruned. Also, from this figure we can see that the performance gained imparted by FUNCTIONAL-SAM over ADAMW equivalent setting about 25% parameters zero, and is thus significant.
B. Additional theory
B.I. Argument for reducing matrix-vector products with inverse preconditioning
We will use a random matrix model to reason about the effects of preconditioning with a matrix inverse. Random matrices have been used to model the spectra of the Hessian of the large models found machine learning, and in particular treating the Gauss-Newton and the functional Hessian/NME as independent yields quantitative insights about the structure of the overall Hessian (Pennington & Bahri, 2017).
Consider two N X N symmetric invertible matrices A and B. Suppose A and B are random and freely independent. Free independence is the non non-commutative analog to classical independence, implied by classical independence of entries in the limit of large matrix size (Pennington & Bahri, 2017). The key feature it induces is
E Ei B El(r | [(rr(B
(13)
for any j. k where tr is the trace.
Given a random unit vector v, the expected squared lengths of matrix vector products with A and B are given by
B|||Av| E[tr[A3]] 31,
B Bv
[[]
(14)
The ratio ri of magnitudes is given by
Bv 12
B
r1 = E A
(15)
A
Now consider the norm of w A L'y passed through each matrix:
E A[][Aw/]
Bw] ElrAA-
- 'B'A-11
(16)
N
The new ratio of magnitudes r2 is given by
E Bw Aw
E(tr B
72
E|tr[A 21
(17)
From Jensen's 's inequality, tr[A?] 2 tr(A -21-1 Therefore r2 > r1: preconditioning by the inverse of A causes matrix-vector products upweigh products with B relative to products with A, relative to the unpreconditioned product.
In neural network settings, there are non-trivial relationships between the gradient, the functional Hessian HF; however, the eigenvectors and eigenvalues of HG and Hp are only weakly related. Therefore even though the exact calculations in the example above don't 't hold, we suspect that generically preconditioning by H will downweigh Hg compared to Hye".
matrix HG and the
16

---

Avoiding spurious sharpness minimization broadens applicability of SAM
C. Code Snippets for SAM and FUNCTIONAL-SAM
from jax import grad, vip
from mox.tree_utee import tree_map
from utils import normalize_grad
def sam_gradients (params, fn, rho):
compute usual loss gradient dL_atheta grad fn) (params)
normalize the gradients dL_atheta normalize_gran (dL dtheta)
perturb the parameters
perturbed_params - tree_map (lambda a, b: + rho b, params, dl, dl__theta)
compute gradient as by
sam_grad grad(loss_fn) (perturbed__arrms)
return sam grad
def Eunctional_sangraadee (params, loss_fn, network_fa, rho):
# compute the usual loss gradient, but also extract dL dlogits (di_dlogits), dt_dtheta-
bax_aux=Tr (params)
normalize the gradients
dt_atheta = normalize_grad (dL dtheta)
perturb the parameters perturbed_parant dree_map (1ambda 8, b: + rho b, params, di__theeaa)
set the VJP at the perturbed parameters
oF fn vip (lambda theta: network fn (theta), perturbed params)
do the We with the (unperturbed) di_dlogits functicnal_sam = dF ff (dI_dlogits) [o]
return functicnal_sst _a______________________________ saa__ra_a__aaaada_aa
Listing 1: Illustration of how to get the gradients in the methods. Functional SAM differs from the SAM implementation only in the last couple lines
17

---

Avoiding spurious sharpness minimization broadens applicability of SAM
D. ANGLE-SAM
In this section, we present general variant SAM, which includes both FUNCTIONAL-SAN and SAM as its particular instantiations. The core idea that once we have been able to decompose the sharpness gradient into those arising from logit and functional paths, we can design our bespoke or custom versions of SAM which lean more less towards one path than another.
To do so, let's assume that out custom path is at an angle & with the functional path. We weigh the functional path by factor of cos( while we weigh the logit path by sin(Ø). Then the gradient update in ANGLE-SAM can be described as:
ANGLE == VoL(O) + sin( Blogit + cos(o) -itune
(18)
VoC(O) + sin( (VF(()) Vr + pE*) - V)L(()) + cos(o) . [Vo + pe*) VFL(0) - VoC(O)) (19)
VoC(O) + po) Hg e" + pcos(o) Hp E" + 0(p2)
(20)
We see that Ø recovers SAM upto first order in while Φ = 0 would yield FUNCTIONAL-SAM and Φ - & would result a variant that optimizes solely along the logit path and which can thus be called LOGIT- All in all, this shows how ANGLE-SAM is a clean generalization of SAM, that equips it with a perturbation angle in addition to the usual perturbation radius p.
At some level, this above formulation could be viewed as making assumption that these two paths orthogonal', On another level, one can simply think of these weights mere strategy to obtain convenient weight settings that have their sum of squares as unity.
We expect that this approach might pay dividends in different pondel-dataset-optmnee triples, and we expect this to be an interesting direction for future work.
"This is not far-fetched. As noted in our experiments, these two paths tend he anti-correlated, but often the correlation is quite small magnitude and approaches zemo
18


Reviews:
[{"review_id": "ZyqBuxViBf", "reviewer": "Reviewer_X2he", "summary": "This paper investigates why Sharpness Aware Minimization (SAM), effective in vision tasks, underperforms in natural language processing (NLP). The authors identify that SAM in NLP overly focuses on reducing sharpness via logit manipulation rather than improving the model's functional geometry, leading to spurious optimization. To address this, they propose Functional-SAM, which prioritizes functional sharpness reduction, and preconditioned SAM, aligning perturbations with optimizer geometry, demonstrating superior generalization across NLP tasks and model scales compared to SAM and AdamW.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed comments and their positive view of our paper. We address their concerns below \n\n\n----\n\n&nbsp;\n\n\n> ### Theoretical justification of the decomposition\n\n- 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*. \n\\\n&nbsp;\n- 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.  \n\n&nbsp;\n\n> ### Lack of theoretical guarantees \n\n- **Convergence Analysis:** We appreciate the comment. However, rigorous convergence analysis for SAM-type methods is notoriously complex:\n  - SAM itself has lacked a general convergence proof in the non-convex case, with initial results only recently published [8]. \n  - Older convergence analyses [1-5] required significant modification of the algorithm itself or strong assumptions to make progress. \n   - 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]. \n\\\n&nbsp;\n-  **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.\n\n&nbsp;\n\n> ### Concerns about efficiency and deployment of methods\n\n**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. \n\n- **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.\n\n&nbsp;\n\n> ### Reliance on C4\n\n- C4 is widely used for benchmarking LLMs and using it facilitates easier comparison to those works. \n\n- 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.*\n\n- 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.\n\n&nbsp;\n\n> ### Decoder-only Transformers\n\nWe 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. \n\n&nbsp;\n\n> ### Practical NLP benefits underdeveloped\n\nWe 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. \n\n&nbsp;\n\n> ### Relation to \"Penalizing Gradient Norm for Efficiently Improving Generalization in Deep Learning\"\n\nThe mentioned paper only uses gradient-norm penalty as a regularizer, and does not analyze it, let alone use the decomposition of the sharpness.\n\n&nbsp;\n\n> ### Complexity of combining Functional-SAM with preconditioning \n\nThe 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).   \n\n&nbsp;\n\n> ### Code \n\nYes, we aim to release code by the camera-ready stage.\n\n&nbsp;\n\n\n---\n\n&nbsp;\n\n*Let us know if you have any further questions; if we have addressed your concerns we hope you will consider **revisiting your review score**.*\n\n&nbsp;\n\n\n\n&nbsp;\n\n[1] M. Andriushchenko and N. Flammarion. Towards understanding sharpness-aware minimization. ICML 2022.\n\n[2] P. D. Khanh et al. Fundamental Convergence Analysis of Sharpness-Aware Minimization. NeurIPS 2024\n\n[3] P. L. Bartlett et al. The dynamics of sharpness-aware minimization JMLR 2023.\n\n[4] Y. Dai et al. The crucial role of normalization in sharpness-aware minimization. NeurIPS 2023.\n\n[5] K. Ahn et al. How to escape sharp minima with random perturbations. ICML 2024.\n\n[6] https://proceedings.mlr.press/v202/agarwala23a\n\n[7] https://arxiv.org/abs/2203.02714\n\n[8] https://arxiv.org/abs/2503.02225", "claims_and_evidence": "While the logit vs. functional sharpness decomposition is intuitive, the theoretical justification relies heavily on empirical observations and simplified assumptions (e.g., free independence of Hessian components). A more rigorous mathematical foundation for the decomposition’s validity across architectures and loss landscapes is lacking.", "methods_and_evaluation_criteria": "Downstream utility (e.g., fine-tuning, robustness) is only briefly explored (via pruning), leaving practical NLP benefits underdeveloped.", "theoretical_claims": "lacks of theoretical guarantees.\n1.No convergence analysis for Functional-SAM or preconditioned SAM.\n2.No theoretical bounds on how much functional sharpness reduction improves generalization.", "experimental_designs_or_analyses": "1.The evaluation primarily focuses on language modeling using the C4 dataset and decoder-only Transformers. The paper does not validate the proposed methods (Functional-SAM and preconditioned SAM) on other NLP tasks (e.g., text classification, machine translation) or diverse datasets, raising questions about broader applicability.\n\n2.Reliance on C4 dataset alone limits insight into performance on noisy or domain-specific corpora.", "relation_to_broader_scientific_literature": "1.This paper provide empirical insights into sharpness regularization.\n\n\n2. Based on their findings, authors propose Functional-SAM and Preconditioned SAM", "essential_references_not_discussed": "The paper references [1], which analyzes gradient norm penalties in the context of SAM sharpness. To clarify the novelty and distinctions of this work, could the authors explicitly discuss how the decomposition of the sharpness gradient in their approach differs from that in [1]?\n\n[1]Penalizing Gradient Norm for Efficiently Improving Generalization in Deep Learning. ICML2022", "other_strengths_and_weaknesses": "The methods inherit SAM’s 2× computational overhead, and the perturbation radius requires careful tuning, especially for large models. The paper notes that tuning \\rho becomes coarser for billion-parameter models, potentially limiting real-world adoption where hyperparameter optimization is costly. Additionally, combining Functional-SAM with preconditioning introduces more complexity, which may hinder ease of use.", "other_comments_or_suggestions": "To ensure reproducibility and facilitate further research, could the authors provide access to the implementation code?", "questions_for_authors": "see questions before.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "24KnB4J9ZF", "reviewer": "Reviewer_TL8s", "summary": "The paper introduces Functionnal-SAM (F-SAM), an alternative to Sharpness-Aware Minimization (SAM) that aims to address its poor performance in NLP tasks. The authors argue that SAM's failure in language modeling is due to its focus on regularizing logit statistics rather than modifying the functional properties of the neural network. They propose F-SAM, which modifies sharpness through the functional path, and PRECONDITIONED-SAM (Pre-SAM), which improves SAM’s perturbation by adapting it to the optimizer’s preconditioning scheme. Their empirical results show very slightly improved performance over both SAM and ADAMW across multiple model scales and training settings.", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your thorough review. We address the primary concerns below:\n\n&nbsp;\n\n----\n\n&nbsp;\n\n> ### 1. Significance of Performance Improvements:\n\n- 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].  \n\n- 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. \n\n-  As clarified in our response to Reviewer Eee3, these gains are statistically significant and indeed represent genuinely meaningful improvements. \n\n&nbsp;\n\n> ### 2. Computational efficiency concerns\n\n - *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.\n\n- *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.\n\n- *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.\n\n&nbsp;\n\n> ### 3. Comparison to Other Methods:\n\n- *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*.\n\n- *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. \n\n&nbsp;\n\n> ### 4. Theoretical Claims:\n\n - **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.\n\n- **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.\n\n&nbsp;\n\n> ### 5. Experimental Details (Hessian Variance):\n\nHessian 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. \n\n&nbsp;\n\n> ### 6. Specific Questions:\n\n- *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.\n- *PEFT Compatibility:*  Likely compatible, but the interaction needs study.\n- *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) \n\n&nbsp;\n\n\nWe 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. \n\n----\n\n\n*Let us know if you have additional questions; if we have answered your concerns, we hope you will consider revisiting your review score.*", "claims_and_evidence": "The main claims of the paper are:\n- SAM performs poorly in NLP tasks because it minimizes sharpness primarily by modifying logits statistics rather than the network’s functional properties.\n- F-SAM improves sharpness regularization by emphasizing functional modifications over logit-based adjustments.\n- Pre-SAM further improves sharpness minimization by adapting perturbations to the optimizer’s preconditioning scheme.\n \nThe combination of F-SAM and Pre-SAM shows slight performance improvement (max improvement displayed is 0.06 loss points on values around 3.5) in large-scale language modeling.\n \nThe evidence includes:\n- Theoretical decomposition of sharpness minimization into logit and functional paths.\n- Empirical validation of the proposed algorithms on multiple model scales (from 2M to 1.2B parameters) in both fixed-length and Chinchilla-style training regimes.\n- Hessian eigenvalue analysis showing F-SAM reduces sharpness more effectively than SAM.\n \nHowever, key claims remain heuristic rather than rigorously proven, and the cost-benefit trade-off is not discussed in sufficient depth.\nFurther, the empirical results show only *very* marginal improvements over SAM and more critically ADAMW, raising questions about the practical significance of the proposed method.", "methods_and_evaluation_criteria": "The proposed methods are evaluated using:\n- Validation loss on language modeling tasks with multiple model scales.\n- Hessian eigenvalue analysis to assess sharpness reduction.\n- Performance comparisons with SAM and ADAMW under equivalent computational budgets.\n \nThe evaluation is generally well-structured but has critical weaknesses:\n- The efficiency trade-offs of F-SAM and Pre-SAM are not analyzed and are merely discussed in Section 7.\n- No training time comparisons or FLOP analysis to assess whether F-SAM justifies its additional computational cost.\n- No comparisons to alternative sharpness minimization methods beyond SAM, leaving open the question of whether F-SAM is the best solution for this problem. Further, the claim about improved versions of SAM (Kwon et al., 2021; Tahmasebi et al.,2024; Li & Giannakis, 2024) being an orthogonal line of work is not substantiated and would benefit from either a more detailed discussion or empirical comparison.", "theoretical_claims": "The paper presents a decomposition of sharpness minimization into logit and functional paths but does not provide a rigorous proof that F-SAM leads to better generalization. Instead, the claims are supported by empirical observations and qualitative reasoning which are unfortunately not sufficiently backed up by experimental results in my opinion.\n \nUnfortunately, there is no formal proof that logit-path minimization is suboptimal for NLP. Furthemore, the paper does not provide a theoretical justification for why F-SAM gives better convergence guarantees than SAM.\n \nThe appendix introduces ANGLE-SAM, which generalizes SAM by parameterizing perturbations using an angle $\\phi$, showing that F-SAM and SAM are special cases. However, this remains an intuitive generalization rather than a rigorous theoretical result. I believe exploring such a theoretical generalization and studying its properties with respect to $\\phi$ could be a strong contribution to the paper.", "experimental_designs_or_analyses": "The experimental setup is very clear and sufficient in terms of dataset and model choices, but it could be improved in several ways:\n- Computational cost is not analyzed, making it unclear whether F-SAM is worth the additional cost.\n- No training time comparisons between F-SAM, SAM, and ADAMW.\n- Limited discussion on efficiency: if F-SAM is significantly more expensive while offering small improvements, it is not practically useful.\n- The results are underwhelming in terms of performance improvements, with the best improvement being 0.06 loss points on values around 3.5. This raises questions about the practical significance of the proposed method.", "relation_to_broader_scientific_literature": "The paper is very well-situated within the sharpness regularization literature, and even in the preconditionned optimization literature, even though one could regret the absence of preconditioning-based optimization methods like Shampoo (a precursor of SOAP by Vyas et al.) or [K-FAC](https://arxiv.org/abs/1503.05671).", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "**Strengths:**\n- Identifies a limitation of SAM in NLP and proposes an alternative.\n- Provides a novel decomposition of SAM into logit and functional paths.\n- Coherent empirical framework across a wide range of model scales.\n- Hessian eigenvalue analysis to support empirical claims is given\n \n**Weaknesses:**\n- No formal proof of F-SAM’s theoretical advantages.\n- The computational efficiency is not analyzed, making practical applicability uncertain.\n- The empirical results are very weak and would not justify the additional cost of F-SAM.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "- Are the proposed methods adapted for fine-tuning settings? Furthermore, do the given pre-trained models lead to better zero/few-shot performance on downstream tasks?\n- Would the proposed methods be compatible with parameter-efficient fine-tuning techniques such as LoRA or Adapters? Could they be combined in a fine-tuning pipeline?\n- SAM is often seen as a regularization technique, however, it is not explained in the paper how F-SAM and Pre-SAM relate to explicit regularization techniques. Could they be combined with L1/L2 regularization or other explicit regularization methods to improve performance further?\n- The Hessian metrics reported are not associated with their variance but are known to be empirically noisy. Could you provide more details on the Hessian analysis and how it was conducted?\n \nIf the overall weaknesses and questions are addressed, I would be happy to raise my score, although I believe the paper would benefit from more substantial improvements (Angle-SAM seems to be a very promising theoretical approach) to justify a higher rating.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}, {"review_id": "oSodCycOgv", "reviewer": "Reviewer_ewUR", "summary": "The paper investigates the limitations of SAM in NLP tasks, where it often degrades performance despite its success in vision tasks. The authors find that SAM's effectiveness varies across domains due to differences in sharpness minimization pathways: the logit path and the functional path. In NLP, the logit path dominates, leading to spurious sharpness minimization. \nThe paper proposes two alternative algorithms: Functional SAM and Preconditioned SAM. Empirical evaluations demonstrate improved performance over AdamW and SAM in NLP tasks.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed comments and address some of their concerns here.\n\n&nbsp;\n\n----\n\n&nbsp;\n\n> ### 1. Computational Overhead:\n\nWe agree that in its current form, Functional SAM is not as FLOPs efficient as Adam. \n\n- 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.\n\\\n&nbsp;\n- 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. \n\\\n&nbsp;\n- 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**).\n\n[1] https://arxiv.org/abs/2203.02714\n\n&nbsp;\n\n> ### 2. Concern regarding long-term performance:\n\nPlease 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. \n\n*Thus, we can be confident that functional SAM does yield long-term performance benefits as well.*\n\n&nbsp;\n\n> ### 3. Demonstrating Enhanced Functional Path: \n\nThis 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.\n\n&nbsp;\n\n> ### 4. Theoretical Support and Preconditioned SAM Motivation:\n\n  - **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.\n\\\n&nbsp;\n - **Preconditioned SAM:** The motivation is twofold: \n    - (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. \n    \\\n&nbsp;\n     - (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.\n\n&nbsp;\n\n> ### 5. Missing References [1-3]:\n\n- [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. \n\\\n&nbsp;\n- Although ref [3] shares similar motivations as SAM and is quite interesting, they in their own words say that the “specific approaches differ significantly”. \n\\\n&nbsp;\n- 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. \n\nRegardless these will be good works to discuss in our paper, and we thank you for suggesting them.\n\n----\n\n&nbsp;\n\n*Let us know if we can clarify any further points. If we have answered your concerns, **please consider giving your score a second thought**.*", "claims_and_evidence": "See Weakness.", "methods_and_evaluation_criteria": "See Weakness.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "See Weakness.", "relation_to_broader_scientific_literature": "See Summary.", "essential_references_not_discussed": "Some improved algorithms of SAM:\n\n[1] Du et al. Efficient sharpness-aware minimization for improved training of neural networks. (ICLR 2022)\n\n[2] Mueller et al. Normalization layers are all that sharpness-aware minimization needs. (NeurIPS 2023)\n\n[3] Wang et al. Improving generalization and convergence by enhancing implicit regularization. (NeurIPS 2024)", "other_strengths_and_weaknesses": "**Strengths**\n- The paper presents a novel decomposition of SAM’s sharpness minimization update into logit and functional paths, revealing that the logit path dominates in NLP tasks.\n- The proposed algorithms, Functional SAM and Preconditioned SAM, empirically outperform SAM in certain NLP tasks.\n\n\n**Weaknesses**\n- **Computational Overhead**: Similar to SAM, Functional SAM and Preconditioned SAM require **twice** the gradient computation per step, making them computationally expensive. Consequently, while the proposed algorithms slightly outperforms Adam *given the same number of iterations*, Adam may still perform better when compared *under equal computational cost*, which is a fairer comparison.\n\n- **Concern regarding long-term performance.** *It is unclear whether the proposed algorithms will be surpassed by AdamW or SAM given sufficient training time.* The experiments on C4 are conducted for a relatively small number of steps. Even with a 1.2B model, the final validation loss remains above 3, which is **significantly higher** than established baselines. For instance, in [4], a 1.2B model achieves a final validation loss of 2.56.\n\n- **Insufficient experimental evidence** to explain how the proposed algorithms work. Although the proposed algorithms are designed to enhance the functional path, no experiments demonstrate whether they indeed result in a larger functional path than SAM.\n\n- **Insufficient theoretical support.**  While the proposed algorithms are motivated by addressing the domination of the functional path,their formulation relies on several approximations. Theoretical support is needed to substantiate that they indeed lead to a larger functional path, particularly for Preconditioned SAM (, whose connection to the main motivation remains unclear).\n\n\n\n[4] Zhao et al. Deconstructing What Makes a Good Optimizer for Autoregressive Language Models. (ICLR 2025)", "other_comments_or_suggestions": "See Weaknesses.\n\nTypo: (Line 377) \"We also that\"", "questions_for_authors": "See Weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FJKuUyPjiV", "reviewer": "Reviewer_Eee3", "summary": "This paper presents an intriguing exploration of the distinction between logit-space and functional-space perturbations within the context of Sharpness-Aware Minimization (SAM). The authors' identification of this subtle difference is interestimng, and while the observed effects might appear minor, the potential ramifications for model training and generalization are substantial. Understanding how these perturbation spaces impact optimization could lead to more robust and efficient training methodologies.\n\n## update after rebuttal\nthanks for addressing some of the concerns. I think this work has merit and I will keep to my original score [weak accept]", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.  \n\n&nbsp;\n----\n\n> ### 1. Significance of Empirical Gains (0.03 loss):\n\nWe understand the concern about seemingly small improvements in loss. However,\n\n- **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.\n\\\n&nbsp;\n- **Statistical Significance:** \n   - 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.\n\\\n&nbsp;\n   - 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. \n\\\n&nbsp;\n\nHence, 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.*\n\n&nbsp;\n\n> ### 2. Preconditioning Novelty (vs. Optax, Granziol, etc.):\n\nWe appreciate the reviewer pointing out related work and common practices. \n\n- 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.\n\\\n&nbsp;\n- 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).\n\\\n&nbsp;\n- 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. \n\n&nbsp;\n\n----\n\nHopefully, this addresses your pending concerns, but please let us know if you have any more questions or comments.", "claims_and_evidence": "The authors propose P-SAM to address issues related to \"un-preconditioned geometry.\" However, it appears that the concept of pre-conditioning the inner adversarial optimizer with the outer optimizer's state, specifically using ADAM, is already a common practice, notably within the Jax's Optax library - and in other works (like Granziol, JMLR; Gordon-Wilson and others). \n\nIf the authors' P-SAM simply replicates this existing approach, then the novelty and contribution are limited. If, on the other hand, P-SAM introduces \"further preconditioning\" beyond established techniques, the justifications provided in the paper are insufficient. The explanations regarding the need for additional pre-conditioning lack the necessary depth and clarity to convince of necessity or effectiveness. More concrete theoretical or empirical evidence is needed to substantiate this claim and differentiate P-SAM from existing implementations.\n\nWhile the authors suggest that F-SAM should deliver superior performance, the empirical evidence provided is underwhelming. A mere 0.03 improvement in loss, based on what appears to be a single seed and a fixed training budget, raises serious questions about statistical significance. Without a clear understanding of the variance across multiple seeds under identical training conditions, it's difficult to ascertain whether this improvement is genuinely meaningful or simply a by-product of experimental fluctuation. For typical NLP problems, where variability between runs can often be considerable, a 0.03 difference might well fall within the noise - this should be commented on. To validate the effectiveness of F-SAM, surely a more rigorous experimental setup, including multiple seeds and a thorough analysis of variance, is essential? Furthermore, it would be beneficial to benchmark these improvements against typical NLP problem improvements to provide more context.\n\nAdding to my concern is the absence of F-SAM results in Figure 2. This would have been more informative had they included F-SAM.", "methods_and_evaluation_criteria": "The general methods suggested make sense for the application - but, as per my comments, there are significance questions about the improvements.", "theoretical_claims": "all proofs checked (to the best of my ability) and all seem to work out.", "experimental_designs_or_analyses": "coming back to the issue of the headline result, which seems to be a small improvement with lack of clarity about its significance. as per previous comments, some multi-run analysis would be good to see if this [slight] improvement is truly significant.", "relation_to_broader_scientific_literature": "there is a other of work in pre-conditioning, some of which is referenced in the paper - but I'd argue that's not the headline of the submission [as pre-conditioning by itself is not novel]. There are some prior works that might be of interest.\nA Random Matrix Theory Approach to Damping in Deep Learning, Diego Granziol, Nicholas Baskerville [arXiv]. Granziol's JMLR paper also looks at a Hessian pre-conditioning [using RMT] and there is related work from Gordon-Wilson and Izmailov and Das's arXiv paper Towards Quantifying the Preconditioning Effect of Adam.", "essential_references_not_discussed": "The key innovation of the paper is the exploration of the distinction between logit-space and functional-space perturbations within the context of Sharpness-Aware Minimization. I could not find prior work that replicates this", "other_strengths_and_weaknesses": "The core novelty of the paper is nice -  logit-space and functional-space SAM. This could have some neat practical implications. While the observation of the logit/functional difference is an interesting contribution, the paper's central claims regarding F-SAM and P-SAM are weakened by apparent methodological shortcomings and a lack of empirical support. To strengthen the paper, the authors should address the (seemingly weak) statistical significance of their results, provide a more comprehensive experimental evaluation and offer a clearer justification for the proposed pre-conditioning method and how it differs from existing approaches.", "other_comments_or_suggestions": "The paper is well-written and an enjoyable read with no obvious errors or typos. The references need Capital letters protecting [minor issue].", "questions_for_authors": "all in previous comments", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DiQx7eB4Jv", "reviewer": "Reviewer_8e58", "summary": "The authors investigate the Sharpness-Aware Minimization (SAM) algorithm for language tasks and find deteriorated performance compared to vision tasks. They explain this by re-writing the SAM update as a gradient norm penalty, and decompose the gradient of the gradient norm into a functional part and a logit part. Through empirical analysis they demonstrate that in language tasks, SAM is biased towards minimizing the logit part - different to vision tasks. They suggest a modified variant of SAM, which explicitly minimizes the functional part, and additionally contains a preconditioning for the perturbation. The authors report improved performance over baselines on language tasks with model sizes over three orders of magnitude (24M-1.2B).", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\n&nbsp;\n\n----\n\n&nbsp;\n\n\n> ###  Same compute budget comparisons\n\nAs 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:\n\\\n&nbsp;\n- **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. \n\\\n&nbsp;\n-  **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. \n\\\n&nbsp;\n- **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.\n\n&nbsp;\n\n> ### Comparison to other SAM variants: \n\n- While interesting, comparing against the plethora of (effectively vision-focused) SAM variants was beyond the scope of this work.\n\\\n&nbsp;\n- 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.\n\\\n&nbsp;\n- 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. \n\\\n&nbsp;\n- 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. \n\n\n&nbsp;\n\n> ### Non-spurious sharpness measures that might correlate better with generalization\n\nThis 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. \n\n&nbsp;\n\n> ### Vision vs Language Intuition: \n\nAnother great question. Our current hypothesis relates to the **nature of the typical output distributions $p(y | x; \\theta)$ in these domains**.\n\n- 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*.\n\\\n&nbsp;\n- 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*.\n\n&nbsp;\n\n> ### Suggestions on Experimental Designs Or Analyses\n\n- Thanks for these great suggestions, we will definitely incorporate them in the camera ready version. \n\n&nbsp;\n\n----\n\n*We hope we have able to address your concerns. We remain at your disposal should you have more questions or comments.*", "claims_and_evidence": "- according to the title, the authors claim to “broaden the applicability of SAM”, implying that SAM becomes practical for (potentially large) language model optimization. However, the results presented in the experiments were conducted for a fixed number of steps, giving SAM twice the compute budget of base optimizers. The authors acknowledge this in their discussion, and point to efficient SAM implementations that could be combined with their work, but this is not explicitly shown. Thus, whether SAM's applicability is indeed broadened is unclear from this work. \n- “Here, we confirm that the generalization benefits imparted by PRECOND FUNCTIONAL-SAM are brought about convergence to a solution with lower curvature, as shown in the Table 4 for the 23.9M model” This only holds when comparing SAM-variants to AdamW. In Table 6, the lowest curvature does not imply lowest eval loss (but still lower curvature compared to AdamW). \n-  “This further highlights how SAM, by default, in language modeling tasks is set up to minimize sharpness spuriously”. The notion of spurious is a bit unclear here, as there is no way of assessing the non-spurious sharpness via the provided numbers, and PRECOND FUNC-SAM shows the lowest value for $tr(H_G)$. It could also just imply that overall the sharpness quantities investigated do not correlate well with generalization.", "methods_and_evaluation_criteria": "Yes, except for the difference in compute budget between SAM and baselines.", "theoretical_claims": "I checked Appendix B.", "experimental_designs_or_analyses": "- like explained above, the difference in compute is problematic\n- Since the authors aim at minimizing the functional path, it would be good to demonstrate that this actually happens with PRECOND FUNC-SAM, e.g. via repeating Figure 4 for _FUNC SAM (precond)_, and potentially also for the other variants\n- Including _Func SAM_ and _SAM (precond)_ in Figure 3 would allow to disentangle the effects of preconditioning and  the functional formulation\n- Since rho is tuned, I suggest reporting the full results in the Appendix (ideally like in Figure 3) for all experiments\n- I recommend also reporting $\\lambda_{max}(H_G)$ in Table 4", "relation_to_broader_scientific_literature": "SAM has mostly been applied to vision tasks, and to a lesser extent for fine-tuning in the language domain. Why it is not used more for language modelling in practice has not been investigated thoroughly. The authors show that training from scratch leads to deteriorated results compared to standard optimizers. This has not been demonstrated in published research. They connect their findings to previous work that highlighted the relevance of the Nonlinear Modeling Error (NME), a component of the Hessian of the loss, for Sharpness minimization. They derive a SAM-variant that explicitly minimizes the functional part, and additionally employs a preconditioning on the SAM perturbation. The preconditioning of the SAM perturbation is conceptually similar to the plethora of SAM variants that exists already. Those variants are discussed, but not compared against in the experiments. The paper is also the first one I am aware of that scales SAM to models bigger than 1B params.", "essential_references_not_discussed": "The authors provide a comprehensive discussion of the relevant literature.", "other_strengths_and_weaknesses": "- the presentation of the paper is nice. I appreciate the effort in communicating the central ideas and results clearly (both through good writing and appropriate use of colors and markers)\n\n- Given that there is a plethora of SAM papers and SAM variants, where each paper claims improved performance over baselines (standard SAM and standard optimizers), it is natural to ask why SAM is - to the best of my knowledge - not used for training LLMs. Investigating and improving the practicality of SAM for language tasks is therefore an important task, and the authors have made a good effort in this direction by finding a difference in sharpness-minimization between vision and language tasks, and proposing a modified SAM variant to mitigate the problems. \n\n- The scale of the experiments (>1B training from scratch) is novel for SAM \n\n- As discussed by the authors, other studies have also proposed preconditioning of SAM or adjusting its perturbation model, and the authors “believe our decomposition approach is orthogonal to this line of work”. While I agree that the perspective of decomposing the sharpness term is novel for SAM, it would still be interesting to see if some of the other SAM variants show the same behaviour as SAM or PRECOND FUNCTIONAL-SAM in language tasks. Perhaps some of the problems are already implicitly mitigated by those variants.", "other_comments_or_suggestions": "- grammar mistake in line 186: “we see that a simple but spurious way to decrease it is by make the network…”\n- typo in line 376: “interestingly, we also that”", "questions_for_authors": "I'll repeat points from above, in decreasing relevance\n1. (most relevant) Do the authors have evidence that their method improves over baselines for the same compute budget? \n2. Have the authors compared to other SAM variants? Could be for small-scale setups\n3. Does the decomposition of the Hessian imply a non-spurious sharpness measure that might correlate better with generalization (e.g. via $\\delta_{func}$)? If so, could the authors report it alongside the metrics in Table 4?\n4. (least relevant) Do the authors have any intuition or evidence about where the difference between the language and vision setup and the corresponding spurious sharpness minimization comes from? I.e., why would one setup be biased towards the respective path? Can it be connected to the loss function, or the number of classes, or not training until convergence? Is there a connection to label smoothing? I'm just curious, and I don’t expect a comprehensive answer here; it will likely not affect my score, so the authors can feel free to not answer this.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: Sharpness-aware minimization (SAM) is known to give great results when training vision models on small datasets but often performs poorly and even degrades performance when training LLMs. The work proposes a new method called Functional-SAM (along with preconditioned variants) which only perturbs the model and not the logits. Unlike regular SAM, the method shows consistent improvements when training language models at different scales. 

Despite lower scores, reviewers appreciated several points in the paper. It is the first work to scale SAM to >1B models, considers the under explored SAM+LLM setting and even improves SAM to work for LLMs for the first time. 

Several concerns were pointed out by reviewers, such as 1) lack of comparison to other SAM baselines, 2) additional overhead of SAM and 3) lack of theoretical advantages of functional SAM.  While reviewers did not change their score, I found most of these points to be addressed well in the rebuttal. There are ways to overcome the additional overheads, and even more computationally expensive methods could be interesting in view of scaling laws that also take the amount of data used into account.

Overall I believe the strong points outweigh the concerns raised by the reviewers. It is important that SAM research moves past convolutional networks on CIFAR and this work makes an important step in that direction and will be of great interest to the community. For the final version, I encourage the authors to take the reviewers suggestions into account, and a comparison of the Zoo of existing SAM variants on (perhaps a small) language model will be very valuable.
Decision: Accept (poster)