νSAM: Memory-Efficient Sharpness-Aware Minimization via Nuclear Norm Constraints
Authors: Thomas Pethick
Parameswaran Raman
Lenon Minorics
Mingyi Hong
Shoham Sabach
Volkan Cevher
Venue: Accepted by TMLR ()

Abstract: Sharpness-aware minimization (SAM) has been shown to improve the generalization of neural networks. However, the method comes at the expense of storing a perturbation of the model parameters, which can be restrictive when memory bound. We design a variant of SAM, called $\nu$SAM, which obtains a low-rank perturbation by modifying the perturbation constraint. The update almost entirely removes the memory footprint of the perturbation without increasing the computational complexity, thus achieving close to a $1/3$ memory saving regarding the parameters when using SGD as the base optimizer. We demonstrate comparable performance of $\nu$SAM with SAM on vision transformers both when training models from scratch and for fine-tuning. Interestingly, $\nu$SAM seems to significantly improve performance for MLP-Mixer architectures across both settings. The results are corroborated theoretically, where we show that SAM with an \emph{arbitrary} norm choice (which includes $\nu$SAM) can converge even with fixed perturbation radius.

Published in Transactions on Machine Learning Research (01/2025)
vSAM: Memory-Efficeent Starpness-Aware Minimization via Nuclear Norm Constraints
Thomas Pethick
hommaspethickklepff.cc
LIONS, IEM, STI, Ecole Polytechnique Fédérale de Lausonné"
Parameswaran Raman
prooman&aamczon.com
Amazon Web Services
Lenon Minorics
minoricss@mmczon.com
Amazon Web Services
Mingyi Hong University of Minnesota
mhong@umn edu
Shoham Sabach
Faculty of Data and Decision Sciences, Technion Israel Institute of Technology'
Volkan Cevher
wolkn..ceuhr@eepf..cd
LIONS, IEM, STI, Ecole Polytechnique Fédérale de Lausanne
Reviewed on OpenReview: https: openreview. net/ forum? id- V61a5NWIMD
Abstract
Sharpness-aware minimization (SAM) has been shown to improve the generalization of neu- ral networks. However, the method comes at the expense of storing a perturbation of the model parameters, which be restrictive when memory bound. We design variant of SAM, called vSAM, which obtains low-rank perturbation by modifying the perturbation constraint. The update almost entirely removes the memory footprint of the perturbation without increasing the computational complexity, thus achieving close to 1/3 memory saving regarding the parameters when using SGD as the base optimizer. We demonstrate comparable performance of vSAM with SAM on vision transformers both when training models from scratch and for fine-tuning. Interestingly, vSAM seems to significantly improve performance for MLP-Mixer architectures across both settings. The results corrobo- rated theoretically, where show that SAM with an arbitnary norm choice (which includes vSAM) can converge even with fixed perturbation radius.
1 Introduction
Sharpness-aware minimization (SAM) (Foret et al., 2020) has seen rising popularity due to increasing the generalization capability across a wide range of tasks. The method consistently improves classification error (Foret et al., 2020), replaces heavy data augmentation otherwise used in pretraining of vision transformers (ViTs) (Chen al., 2021), and improves fino-tuning of large language models (LLMs) (Bahri al., 2021; Zhong el. al., 2022). SAM has also been shown to be more robust. (Foret al., 2020), interpretable (Chen et al., 2021) and reproducible (Somepalli et al., 2022).
"The work of TP was done while interning Amazon Web Services. 'Concurrent positions as an Amazon Scholar and as faculty corresponding institutes. This paper representa the work performed Amazon.

---

Published in Transactions on Machine Learning Research (01/2025)
Table 1: (Left) The number of parameters stored for the perturbation. The vSAM method uses almost no memory for storing the perturbation in comparison with SAM. (Right) Removing the perturbation entirely would lead an overall memory saving regarding the model of up to 1/3. Memory is provided in multiples of the model size.
SAM
vSAM
Saving
Method
Memory
Saving
VIT-S/16
22M
75k
99.66%
SGD
1/3
VIT-B/16
87M
149k
99.83%
SGD momentum
3
1/1
VIT-L/ 16
304M
395k
99.87%
Adam(W)
4
1/5
BERT (base)
110M
201k
99.82%
Adam(W) bf16 momentum
1/4
BERT (large)
335M
479k
99.86%
Adafactor
~2
1/3
The update rule of SAM proceeds by finding a (norm constrained) perturbation vector I2 and subse- quently updating the weights x Rd using gradient ff computed at the perturbed set of weights,
Despite its popularity, one subtle problem is the increased memory footprint SAM as compared con- ventional first-order methods such as stochastic gradient descent (SGD). Specifically, at any given time the SAM method needs to store: the weights, gradient and additionally the perturbation. This is 50% more memory demanding in terms of the weights than SGD, which only needs to store the weights and gradient.
This naturally raises the following research questions:
it possible to obtain a SAM formulation that induces a memory-efficient algorithm without intro- ducing additional computational overhead?
In this work, we answer the above question in the affirmative. Concretely we make the following contributions:
By revisiting the original SAM formulation we notice that the structure of the network is lost in the computation of the perturbation & since the parameters are treated as a vector. This motivates us replace the original vector 4-norm constraint with a particular matrix norm constraint (specifically the muclear norm) that allows to obtain layerwise low-rank e. The modification leads to substantial memory-savings regarding the perturbation, saving more than 99.8% of the memory required by the perturbation in the original SAM method on ViT-B/16 B/16 (sce Table 1).
We extensively evaluate our method =SAM on vision transformers (ViTs) and MLP-Mixer models when both fine-tuning and training from scratch, and additionally fine-tune BERT on set of language tasks. We find that vSAM consistently outperforms the bascline AdamW in all cases and achieves comparable performance with SAM. Surprisingly, this is the case even when the low-rank decomposition of the perturbation E is only coarsely approximated through single power iteration, which avoid adding any wall-clock time as compared with SAM. Interestingly, we find that vSAM enjoys a substantial improvement over SAM for MLP-Mixer models on both fine-tuning and when training from scratch.
We provide a strong bascline for ViTs on relatively small datascts like CIFAR, which might be of independent interest, increasing the baseline in (Mueller et al., 2023) by more than 3pp (see Sections 5.1 and 5.2.1). We interestingly find that, in this setting, both z/SA and SAM substantially improve if the perturbation is not activated until after several epochs. The delay raises the percentage points improvement over AdamW by striking factor of ~4. It appears that the perturbation can be taken larger if delayed, partially explaining the benefit.
Thcoretically, we find that SAM-type methods with arbitrary norm choice for the perturbation (which includes vSAM) converges even with fixed perturbation radius p, although only for a re- strictive class of certain convex quadratics. This is in stark contrast. with SAM variant using
2

---

Published in Transactions on Machine Learning Research (01/2025)
(normalized) random perturbation, which is not guaranteed to converge, demonstrate. This observation possibly sheds some light on the practically observed advantage of using gradient infor- mation in the perturbation.
Limitations Our work focuses on transformers and MLP-MMxer models, since our memory-eficiint refor- mulation of SAM heavily relies on the presence of matrix structures in the network.
2 Related work
Memorr-efficient SAM There has been other works trying to lower the memory consumption of SAM. FSAM (Zhong et al., 2022) and SSAM (Mi et al., 2022) use the Fisher information to identify subset of the parameters to perturb. Tn practice SSAM save 50%, and FSAM roughly 90%, of the perturbation (sce Zhong et al. (2022, Sec. 6.1)) whereas =SAM saves 99.8% (sce Table 1). Very recently, similar memory saving as vSAM was made possible with SAM-ON, which only perturbs the normalization layers (Mueller et al., 2023). Very surprisingly, this minor modification leads to improvements over SAM in many scttings. We compare in detail in the experimental section where we find that our method has an advantage for specifically MLP-Mixer architectures and for fine-tuning tasks across architectures.
Convergence of SAM The analysis in GSAM (Zhuang et al., 2022), SSAM (Mi et al., 2022) and An- driushchenko & Flammarion (2022, Thm. 2) takes decreasing perturbation radius Pk. Almost sure conver- gence of SAM was shown in Nam et al. (2023) using similar construction. The decreasing Pk avoided in (Andriushchenko & Flammarion, 2022, Thm. 6) by ignoring the normalization in SAM. Very recently, SAM with fixed perturbation radius was studied in Si & Yun (2023) with convergence guarantees for (strongly)-convex objectives, while providing negative results for stochastic and nonconvex cases.
Power iteration An important method for eigenvalue and singular value computation the power itera- tion method (Mises & Pollaczek- Peitiezel-Geiringer,, 1929) (see e.g. Golub & Van der Vorst (2000)). The method has been extensively used in the machine learning community. It is, for instance, the backbone of the PageRank algorithm (Page al., 1999), and used for computing the spectral norm in training of generative adversarial networks (Miyato et al., 2018).
Low-rank approximations Thore have been several work exploiting low rank structures machine learning application: Vogels al. (2019) uses a low rank approximation to save on communication costs in distributed settings, recently low-rank fine-tuning was popularized by LoRA (Hu et al., 2021), and other work explores the effect of initialization on low-rank pretraining (Kamalakara et al., 2022).
3 Algorithmic derivation
The starting point of SAM is the following saddle point problem for given loss f : + R
min max f(x +), FEX SS
(2)
where p E [0, oo) is the perturbation radius and ||-I| is some norm be defined. To obtain computationally efficient method, SAM (Foret et al., 2020) linearizes the maximization problem with first-order Taylor expansion as follows
Ex E Argmax f(x + sEx:||el|so
Argmax (f(x) + (e. vf(x))) Arg max (e, vf(x)) sEX:||<||s>
(3)
<EXii|<||SP
For the particular choice of fg-norm constraints, the argmax has a closed form solution, which leads to the following update rule:
* = |V/ V((2*),
p
(SAM)
x A+1 = th V/( + )
3

---

Published in Transactions on Machine Learning Research (01/2025)
Algorithm Nuclear norm based sharpness-aware minimization (vSAM)
Require: Parameter initialization x=1 = "MT''...,,,,"" SVD initialization u, 12 Vie
Repeat for k = 0,1,.... K -1
1: Draw a sample & ~ Compute the adversarial perturbation 5*+1 (M,++))E
uute of SVor ( (V M. , ux),u, vi M Mf put (ut)
Update the weights xk1 (MM+t))EE))
M M put (o+) 7 VM (x) e
Return x² = (M ..... k²
For simplicity the update is only for matrix shaped parameters. Update non-matrix parameters using gradient descent.
Algorithm Top singular value decomposition top1)
Require: Matrix A E R initialization u-1 E R", 10-12 E (0,
Repeat for / = 0, I....,T----......... -1
uft1
n
ut 2
+1 A 12+1
Au
Return u. U
One potential limitation of SAM is the requirement to store one entire additional set of model parameters. To economize on memory we revisit (3) and notice that the RHS can be interpreted as a linear minimization oracle (LMO), i.e.,
Imoz(s) = Arg min (2,
(4)
2E2
To be precise, (3) can be written as
Arg max (c, Vf())) Imog SEXi||E||ss
(5)
where & {EE X |||| SP). It is well-known that for certain choices of norms, such the nuclear norm, the LMO leads to sparse solution, which will exploit in what follows.
Abstractly, with slepsize 7 > 0, our update takes the following form
* Imog (x)*))) k+1 = = +
(6)
VV
Instead of flattening the model parameters, will consider the where model is parametrized by matrix in X Raxm and we choose the muclear norm constraint &, { X ||=||~ <p}. In that case update (6) reduces to only computing
ck = (T a l+1 x ~Vfla" +
(7a)
(7b)
where & and * are the singular vectors associated with the top singular value of - f((*) (see Section 3.1 for dcrivations).
We crucially rearranged the update to avoid ever storing sk and a simultaneously:
4

---

Published in Transactions on Machine Learning Research (01/2025)
I² = x² + +p po" ( xk+1 = -* pu (o+)T ff*
(8a)
(8b)
In other words, I² can be completely discarded after step in vSAM (Algorithm 1), which is critical achieving the memory saving. In (8a), only need to store one matrix x* and vectors w and v by performing the update in-place. Similarly, in (8b), it suffice to store the matrices z and vf((*t) and two vectors u and v²_ In contrast, (SAM) requires storing three matrices simultancously.
The top singular vectors can efficiently be computed using (warmstarted) power iterations (Algorithm 2). We simply ignore the update for other parameters inchuding the vector-valued bias terms in the MLP layers. The resulting algorithm described in =SAM (pronounced "newSAM") where the weight parameterization consists of I weight matrices M, E RT i.c., x (M....., M)) and the update uses stochastic gradicnts.
As in vSAM (Algorithm 1) the LMO is applied independently to each layer matrix, which is beneficial for distributed settings (see Section 3.2). We note that the resulting update still be interpreted norm-constrained LMO on the joint parameters namely as the max-norm over muclear norms,
| == max [| || [Mall.......MMMl))
Consequently, the convergence in Theorem 4.3 applies. We also develop and compare with normalized variant of vSAM in Appendix C.I, where layer updates are no longer independent.
The memory savings are large if the embedding dimensions (i.e., layer width) are large (c.f. Table 1). The low- rank perturbation allows vSAM to reduce the memory consumption associated with storing the perturbation by almost an order 103, e.g. the 304M parameter ViT-L/16 model only requires 395k parameters for the perturbation.
3.1 The low-rank update
To arrive at (8a) we seek a closed form solution to the linear minimization oracle (LMO) under a nuclear norm constraint,
Imo (W) arg min (W,X)
(9)
We include the closed form solution below for completeness (see e.g. Jaggi (2013) for similar result). Lemma 3.1. The LMO under wuclear constraind in (9) can be efficienlly compuled
Imo (W) = puu where u and v are the left and right singular vectors associated with the top singular value of W.
Lemma 3.1 yields the update rule (7), which, when approximated with power iterations, leads Algo- rithm 1.
3.2 Implementation
In practice we avoid perturbing the bias term and the (single) convolutional layer ViTs if not otherwise specified.
See Algorithm 3 in Appendix C for pseudo-code for distributed setting with distributed data parallel (DDP). We note that vSAM is also compatible with fully sharded data parallel (FSDP) as long the neural network structure is preserved when sharded, such that the matrix structure is not lost. Compatibility with FSDP opens up the possibility to apply =SAM to models that do not. fit. on a single GPU. As of writing, the FairScale implementation (FairScale authors, 2021) allows preserving the structure whereas the PyTorch implementation flattens the network parameters (v2.0.0). One favorable property of vSAM under FSDP is that no global normalization needs be computed and distributed between the layer computations, which is otherwise the case for SAM.
5

---

Published in Transactions on Machine Learning Research (01/2025)
10"²
10a
10.4
10:4
30
Gaus:SAM
12
VSAM (exact)
10"
10a
Itaration
Figure Demonstration of Theorem 4.3. The direction of the perturbation is important: Both vSAM and SAM with decrcasing stepsize 7kk = 1/k enjoys 0(12)) rate on min.g[k--1] [V' while perturbing with normalized Gaussian noise (GaussSAM) exhibits nonconvergence. Gradient descent (GD) with the same stepsize only converges O(1ss).
The vSAM method almost entirely removes the memory overhead of the perturbation. It important, in conjunction, to also reduce the memory footprint of the optimizer states. This work uses the popular Adam (Kingma, 2014) and AdamW (Loshchilov & Hutter, 2017) as the base optimizer, but Adam(W) can be readily replaced by memory efficient variant, such Adafactor (Shazeer & Stern, 2018), which is commonly used in the vision transformer literature (see e.g., (Zhai al., 2022, sec. 3.4) and Beyer et al. (2022)). For fine-tuning specifically, we note that vSAM is also compatible with using low-rank updates for the minimizer as in e.g,, LoRA (Hu el al., 2021). We provide an overview of methods in Table 1, which shows that the relative memory saving is higher when the base optimizer has a small memory footprint.
4 Convergence analysis
The gradient direction in the perturbation appearing through the LMO of both SAM and vSAM turns out to be important for convergence under fixed perturbation radius p. If a random direction is used instead, convergence cannot be guaranteed as illustrated in Figure 1. We will now make this claim precise.
In this soction, || will more generally denote the dual norm of some norm || |l. instead of specifically referring to the nuclear norm. For the matrix case || ||2 refers to the entry-wise 2-norm (i.e., the Frobenius norm).
Abstractly our update takes the following form:
* x + p
(10)
k k 7h (T(**)
where the particular choice e E A' is to be defined.
Interestingly, we can show convergence without decreasing p for certain convex quadratics, if we restrict the E-perturbation the LMO over an (arbitrary) norm-ball.
We make the following assumptions and define f* == infzex f(x).
Assumption 4.1. The function f X + is strongly-conver with parameter μ > 0, i.c.,
(Vf(x) Vf(x'), x²) 2 2'11
Yr, X.

---

Published in Transactions on Machine Learning Research (01/2025)
Assumption 4.2. The gradient Vf: X + X is L-Lipschitz with L 0, i.c.,
fV) V f (2 2 - L||x - r'12 Vx, x' E EX.
Theorem 4.3. Suppose Assumptions 4.1 and 4.2 hold with == p. Then, (10) with the (arbitrary) norm perturbation e E Arg Arggmx max s:1lell. (V](x*)), satisfies the following descent inequality with p 2 and 7 > 0,
1(ah+2)- 8- s - |K (f - 8*) - -x - (L)L))/V((("" 22
Furthermore, for 7 min ا log(a), # with a= - A02n2K²
it follows that
D'P'E!E"
min f(x") - = =Ö(exp(- 2/2 ) + Ce RT
ke[k]
with Do = C( = max,ex ||=||//||=||||| and &a max /
Remark 4.4. Theorem 4.3 holds for any norm-bal ball constrained perturbation including both SAM and vSAM. In this sense, Theorem 4.3 can be seen as a generalization of the recent result of Si & Yun (2023). Interestingly, the fact that the proof uses the particular direction of the perturbation appears to not only be an artifact of the proof technique. In particular, if the perturbation is instead replaced with normalized Gaussian noise, the scheme does not converge (ef. Figure 1).
The assumptions of Theorem 4.3 are rather restrictive, but they are sufficient for showing the importance of including the gradient direction in the perturbation. Theorem 4.3 applies to quadratic minimization problem of the form fx) 4 || 2 + e for some b E X and e E R, which we demonstrate it on in Figure 1. The parameter p is kept fixed and the stepsize is taken as Tik (1/2)) as suggested by the theory.
The particular direction of the perturbation turns out to be important not only for the proof. If normalized Gaussian noise is used (i.e., ex = ''/le=laa with et ~ I), denoted as GaussSAM) we observe that the iterates only converges within a neighborhood of the solution as also suggested by theory (cf. e.g. (Li et al., 2024, Thm. 4)). This negative result is particularly interesting in the light of the generalization bound of SAM in Foret et al. (2020, Thm. 1), whose statement would also hold for GaussS A.M It is therefore not obvi- ous from the generalization bound alone that the gradient direction in both vSAM and SAM is preferred over GaussSAM. The optimization perspective in Theorem 4.3 provides a possible explanation, by showing that only the former enjoys convergence guarantees. In addition, we observe the fast O(11²) rate for vSAM and SAM, whereas the (unperturbed) gradient descent (GD) with the same stepsize has the slower ((1/x) rate.
If the perturbation radius can be taken docreasing it is possible to show convergence without convexity assumptions, as long as the gradients are bounded. No structural assumptions are needed on the perturbation ** (apart from being bounded) and the result. thus also applies to vSAM.
Assumption 4.5. The gradient Vf X + X is bounded, i.e., [V]((=) S S Va X. Theorem 4.6. Suppose Assumplions 4.2 and 4.5 hold and || S for all [K]. Then (10) satisfies the following convergence guarontee for Pk 2 0 and Th E LL)),
BotC. ΣK +2p
En K-1
TAPE
<
E 71(1-v64/2)
with Ao = f1 and C2 LGE. In particular, for constant stepsize 7 = = 1/LVK and p VR
1 the following rate is obtained
min
0
d VK
VK
LK37
Remark 4.7. To get convergence one ncods to take Pk decreasing. Specifically, it suffice to take Σ TkPk < Do and Ex- 7 co. We note that decreasing p is used for SOTA experimental SAM results as in Zhuang et al. (2022). Theorem 4.6 also implies a fallback guarantee for when p is taken constant, since the theorem in that case ensures convergence to p-dependent neighborhood.
7

---

Published in Transactions on Machine Learning Research (01/2025)
Training from scratch on CIFAR10/100 where vSAM shows strong performance on especially Mixer models where test accuracy is increased over SAM by between 0.28pp and 0.55pp, while memory usage associated with the perturbation is reduced by almost orders of magnitude.
Table 2:
Method Model
AdamW
SAM
SAM-ON
vSAM
Dataset
CIFAR10
Mixer-B, B/4
91.00
91.70
91.73
(+0.03)
92.23
(+0.53)
Mixer-S/ //4
91.08
91.96
92.00
(+0.04)
92.24
(+0.28)
VIT-B/4
93.45
94.62
94.90
(+0.28)
94.44
(-0.18)
ViT-S/4
93.37
94.21
94.65
(+0.44)
94.39
(+0.18)
CIFAR100
Mixer-Be B/4
66.90
69.54
68.25
(-1.29)
70.09
(+0.55)
Mixer-S/ /4
68.03
70.35
69.35
(-1.00)
70.79
(+0.44)
VIT-B/4
68.25
70.26
71.25
(+0.99)
70.15
(-0.11)
ViT-S/4
68.17
71.71
72.70
(+0.99)
70.60
(-1.11)
CIFAR10 VIT-S/4
CIFAR100 Mixer-S/4
Adamw
Test accuracy
0.94
ISAM
SAM
/dam
Test
SAM ON
0.05
0.30
2000000
2,00
10.00
20.00
0.2
2.0
10.0
Figure Sweep over p for ppretraining results in Table 2. See Figure 4 in Appendix D for remaining
configurations.
5 Experiments
We evaluate vSAM both when training from scratch and for fine-tuning against several baselines, namely Adam(W) (Kingma, 2014; Loshchilov & Hutter, 2017), SAM (Foret et al., 2020) and SAM-ON (Mueller et al., 2023). The relative (colored) number in the tables captures the difference in performance from SAM.
Note that we use a single power iteration (T= in Algorithm 2) for the approximation in =SAM if not otherwise stated and only the weight matrices are perturbed. Specifically, we simply ignore the (single) convolutional layer in the ViT and the bias terms when applying vSAM. This leads to method that has wall-clock time matching that of SAM (c.f. Table 14 in Appendix D.1). We experiment with using more power iterations in Table 10 of Appendix D, but conclude that T performs sufficiently well, whilc remaining compulationally cheap.
5.1 Training from scratch
We train multiple sizes of ViTs and MLP-Mixer on CIFAR10/100 from scratch. We specifically take the original architecture configuration of ViIs and MLP-Mixer models in Dosovitskiy al. (2020); Tolstikhin al. (2021) and downsize the patch size to to fit the smaller image size of 32 32. See Table 6 in Appendix D for details.
Baseline & hyperparameters Since there is a lack of good hyperparameter defaults for these architec- tures small datasets, we first find good configuration for the base optimizer AdamW on CIFAR10. We provide a substantially better bascline (AdamW) than e.g. Mueller et al. (2023, Table 3) (93.37% instcad of 90.34% on ViT-S/4). The baseline is comparable with SOTA ViTs for small datasets (Gani et al., 2022; Lee et al., 2021; Liu et al., 2021) without optimizing the VIT structure and without using CutMix, MixUp, repeated augment, stochastic depth and random crase. The final hyperparameters can be found in Table 7 of Appendix D. We use standard augmentations (random cropping and flipping) and AutoAugment in Gani et al. (2022) and cosine learning rate schedule with linear warmup.
8

---

Published in Transactions on Machine Learning Research (01/2025)
Advrtii :SAM SAN
Test accuracy
0.04
0.94
0.942
0.92
.SAM eparh-d
spoch-100
Adamw
.SAM epoch-30
15AM epoch-200
0.01
0.10
0.20
2,00
0.938
Test accuracy
0.94
0.000
0.92
SAM epoch-0
SAM spoch-100
AdamW
0.934
SAM opach-30
SAM epoch-200
0.01
0.10
0.20
0.30
0.50
1.00
2,00
01052
58
105
150
200
250
330
Perturbattin epoch
(b) The optimal p for =SAM increases with the perturba-
(a) Performance of the best p for given perturbation de- tion start epoch. Both SAM and vSAM suffer if activated immidately.
lay.
Figure 3: Both SAM and =SAM benefit. from delaying perturbing until after the learning rate warmup phase.
Table 3: Pretraining without AutoAugment. The gap comparison with AdamW larger and =SAM provides substantial improvement over SAM, while significantly reducing memory usage.
Method
AdamW
SAM
SAM-ON
SSAM
Model
Dataset
Mixer-S/4
CIFAR10
85.79
87.92
86.18
(-1.74)
88.55
(- 0.63)
CIFAR100
59.09
60.73
60.52
(-0.21)
61.09
(-336)
VIT-S/4
CIFARIO
88.49
89.90
90.74
(+0.84)
90.42
(+0.52)
CIFAR100
61.22
65.11
66.01
(+0.90)
65.29
(+0.18)
Comparison of SAM variants In order to fairly compare variants of SAM, we sweep over the hyperpa- rameter p and pick the best one for each combination of model type, dataset and method (c.f. Figure 4 in Appendix D). We discover that p can be significantly larger for both SAM and =SAM if the perturbation is only started after the learning rate warmup period. Delaying the perturbation increases the test accuracy for both mothods, so we use this as the default throughout our experiments. We further investigate this curious phenomenon in Section 5.1.1.
Results The test accuracies of the best iterate are shown in Table 2. The vSAM method yields a large improvement on the MLP-Mixer architecture, which is maybe not surprising since there are no convolutional layers in ViTs. What is maybe surprising, is that uSAM performs substantially better than all the baselines, including SAM, on MLP-Mixer across both datasets. We additionally provide experiments without AutoAugment on VIT-S/4 in Table 3 where vSAM is found to consistently outperform SAM.
5.1.1 Delaying the perturbation
We notice that delaying the perturbation in SAM and vSAM until after the learning rate warmup phase greatly improves the test accuracy in our training setup. This in contrast with earlier observations (Agar- wala & Dauphin, 2023). To investigate the effect of the delay sweep over multiple radii, P, for each delay to ensure optimality. The results can be found in Figure 3 which uses a horizon of 300 epochs on CIFAR10 and a ViT-S/4. We compute a mean and standard deviation over 3 independent runs for each configuration.
Interestingly, the best p for =SAM increases as the delay is increased. appears that perturbing in the early training phase is problematic for both vSAM and SAM. Therefore, we delay perturbation for both method until after 10% of the epochs (after the learning rate warmup). Curiously, SAM-ON (Mueller et al., 2023) the other hand does not seem to benefit from similar delay (c.f. Table of Appendix D).
9

---

Published in Transactions on Machine Learning Research (01/2025)
Table 4: Fine-tuning on CIFAR1O/100 where vSAM consistently outperforms SAM-ON and provides a significant performance boost over SAM for the MLP-Mixer model on CIFAR10 (0.21pp) and CIFAR100 (0.68pp).
Method Model
AdamW
SAM
SAM-ON
=SAM
Dataset
CIFAR10
Mixer-B/16
97.16
1 0.16
97.52 1 0.08
97.58
(+0.06) 1 0.10
97.72
(+0.21) 1 0.05
ViT-B/16
98.36
+ 0.19
98.87 + 0.05
98.74
(-0.13) 0.10
98.89
(+0.02) + 0.03
ViT-S/16
98.38 + 0.02
98.85 + 0.00
98.73
(-0.12) 0.01
98.81
(-0.04) + 0.02
CIFARI00
Mixer-B/16
85.87 1 0.05
0.19
86.39 (+0.19) 0.17
86.88
(+0.68) 1 0.25
VIT-B/16
90.54 + 0.03
92.12 0.13
91.47
(-0.65) + 0.05
91.89
(-0.23) + 0.13
VIT-S/16
91.10 + 0.12
91.74 + 0.23
91.37 (-0.37) # 0.13
91.60 (-0.14) + 0.05
Table Fine-tuning of BERT-base (uncased) on GLUE
Adam
SAM
SAM-ON
vSAM
Length
CoLA
57.31
+ 1.58
57.02 + 1.41
57.84 (++0.2) + 0.89
57.38 (+0.37) + 1.12
8.5k
MNLI-m
83.82
1 0.04
84.23 + 0.11
84.08 (-0.14) + 0.27
84.03 (-0.20) + 0.23
393k
MNLI-mm
84.02
0.29
84.37 1 0.16
83.98 (-0.39) 0.29
84.16 (-0.21) 0.29
393k
MRPC
85.11 +: 0.76
86.13 + 1.12
85.64 (-0.49) + 1.17
86.42 (+0.29) + 0.96
3.7k
MRPC (FI)
89.56 + 0.58
90.31 + 0.77
89.95 (-0.36) + 0.74
90.46 (+0.115 + 0.69
3.7k
QNLI
90.20 + 0.82
90.87 1: 0.47
90.90 (+0.03) +:2 0.48
90.11 (-0.76) + 0.27
105k
QQP
91.03 0.11
91.42 1 0.23
91.29 (-0.13) 0.05
91.29 (-0.13) 0.09
364k
QQP (F1)
87.94 + 0.11
88.47 + 0.33
88.30 (-0.17) + 0.08
88.32 (-0.14) 0.11
364k
RTE
58.12 + 0.98
61.01 + 2.35
60.58 (-0.43) + 2.05
59.93 (-1.08) + 0.57
2.5k
SST-2
92.63 + 0.83
92.78 1 0.69
92.27
(-0.50) + 0.52
92.68 (-0.09) + 0.53
67k
STS-B (Pearson)
87.34 0.18
87.43 1 0.20
87.41 (-0.02) 0.32
87.43 (-0.01) 0.43
7k
STS-B (Spearman)
87.06 +: 0.21
87.
+ 0.19
87.16 (-0.02) 0.31
87.26 (+-0.88) + 0.37
7k
Avg
81.22 + 0.35
81.83 + 0.17
81.70 (-0.13) 0.37
81.61
(-0.22) + 0.13
5.2
Fine-tuning
We fine-tune on both vision tasks and language tasks across multiple datasets.
5.2.1
Vision task
Setup We take VIT and MLP-Mixer architectures pretrained on ImageNet and fine-tune the models on CIFAR10/100. Since there is no pretrained Mixer-S/16 release from Tolstikhin et al. (2021); Steiner (2022) we restrict the experiments concerning MLP-Mixer architectures to Mixer-B/16. We optimize the baseline AdamW, whose final hyperparameters can be found in Table 11 of Appendix D. We sweep over the pertur- bation radius p for all perturbation- based methods for fair comparison (see Figure in Appendix D for p sweep). Each configuration is run 3 times to computes mean and standard deviation.
Our baseline performs substantially better than similar experiments in Zhuang et al. (2022, Table 3). The improvement upon the baseline is primarily due to the use of a smaller learning rate and no weight decay. The configuration additionally allows us to run for only 10 epochs (20 times less itcrations) while still achicving better performance.
Results The results are shown in Table 4. Maybe surprisingly, =SAM seems to especially provide an improvement for the larger base model size. Similar training from scratch in Section 5.1 we also see strong performance on the MLP-Mixer architecture. Even if AdamW optimized further and given a larger computational budget it does not close the performance gap (see Table 13 in Appendix D).
10

---

Published in Transactions on Machine Learning Research (01/2025)
5.2.2 Language task
We fine-tune a pretrained BERT- basc (uncased) (Devlin et al., 2018) on the GLUE benchmark (Wang et al., 2018) following the selup and hyperparameters Geiping & Goldstein (2022). We use a 10% perturbation delay for SAM and vSAM (c.f. Section 5.1.1) and sweep over p E {0.01, 0.02, 0.053. The best configuration is picked and run 5 times provide a mean and standard deviation. Hyperparametens can be found in Table 12 of Appendix D. The results shown in Table 5 where find that all SAM variants improve upon the Adam baseline. The vSAM method appears to exhibit the smallest variance across all the methods on average.
6 Conclusion
We developed vSAM, sharpness-aware minimization algorithm that almost entirely removes the additional memory otherwise required for storing the perturbation. The method uses low-rank approximation that only needs to be very coarsely approximated in practice, resulting in no runtime overhead compared with the original SAM method. We observe strong performance on particularly MLP-Mixe models across both fine- tuning and when training from scratch. We additionally find that both vSAM and SAM benefit from delaying activating perturbation in certain settings. Interesting future work involves better understanding the effect of the architecture and the benefit of perturbation delay as well as exploring the use of other norm choices.
7 Acknowledgments
S. Sabach, M. Hong, and V. Cevher hold concurrent appointments as an Amazon Scholar and as faculty at Technion, University of Minnesota, and EPFL, respectively.
References
Atish Agarwala and Yann Dauphin. Sam operates far from home: eigenvalue regularization dynamical phenomenon. In International Conference on Machine Learning, pp. 152-168. PMLR, 2023.
Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware minimization. International Conference on Machine Learning, pp. 639-668. PMLR, 2022.
Dara Bahri, Hossein Mobahi, and Yi Tay. Sharpness-aware minimization improves language model general- ization. arxiv preprint Xrv::22110. 08529, 2021.
Lucas Beyer, Xiaohua Zhai, and Alexander Kolesnikov. google-research/bigg vision, 2022.
Big
vision.
https://github.com
Xiangning Chen, Cho-Jui Hsich, and Boqing Gong. When vision transformers outperform RosNets without
pre-training or strong data augmentations. arXiv preprinl arXiv:2106.01548. 2021.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirec- tional transformers for language understanding. CoRR, abs/1810.04805, 2018. URL http://arxiv.om /rrg/ abs/1810. 04805.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa. Dehghani, Maithias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words:
Transformers for image recognition at scale. arXiv preprint arXic:2010.11 11929, 2020.
FairScale authors. Fairscale: A general purpose modular pytorch library for high performance and large scale training. https://ggithe http:h/www.mcccccoceearrrrrraaaa 2021.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalivation. arXiv preprint arXin:2010.0141 2020.
11

---

Published in Transactions on Machine Learning Research (01/2025)
Hanan Gani, Muzammal Nascer, and Mohammad Yaqub. How to train vision transformer on small-scale datasets? arviv preprint arXi::2210.0770 072400 2022.
Jonas Geiping and Tom Goldstein. Cramming: Training a Language Model on a Single GPU in One Day, December 2022. doi: 24.5550/arXi.222222144424 URL https://arxii //apxii http://wwwvvv...a.a/.///e /a/aabbb//222/ 14034.
Gene H Golub and Henk A Van der Vorst. Eigenvalue computation in the 20th century. Journal of Compu- tational and Applied Mothemalics, (1-()33) 65, 2000.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. arXiv preprint arXiv:2106. 09685, 2021.
Martin Jaggi. Revisiting Frank- Frank-Wolfe: Wolfe: Projection-free sparse convex optimization. In International confer- ence on machine learning, pp. 427-435. PMLR, 2013.
Siddhartha Rao Kamalakara, Acyr Locatelli, Bharat Venkitesh, Jimmy Ba, Yarin Gal, and Aidan N Gomez. Exploring low rank training of deep neural norworks. arXiv preprinl Xiv: 13569, 2022.
Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.600 2014.
Seung Hoon Lee, Seunghyun Lee, and Byung Cheol Song. Vision transformer for small-size datasets, arXiv preprint Xrx:v221 12. 19492, 2021.
Bingcong Li and Georgios Giannakis. Enhancing sharpness-awart optimization through variance suppression. Aduances in Newral Information Processing Systems, 36, 2024.
Tao Li, Qinghua Tao, Weihao Yan, Zehao Lei, Yingwen Wu, Kun Fang, Mingzhen Hc, and Xiaolin Huang. Revisiting random weight perturbation for efficiently improving generalization. arXiv preprinl 0all:2444.003577 2024.
Yahui Liu, Enver Sangineto, Wei Bi, Nicu Sebe, Bruno Lepri, and Marco Nadai. Efficient training of visual transformers with small datasets. Advances in Neural Information Processing Systems, 11:23818-23838) 2021.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXie preprint orXhv:1711.05101, 2017.
Peng Mi, Li Shen, Tianhe Ren, Yiyi Zhou, Xiaoshuai Sun, Rongrong Ji, and Dacheng Tao. Make sharpness- aware minimization stronger: A sparsified perturbation approach. Advances in Neural Information Pro- cessing Systems, 35:30950-30962, 2022.
RV Mises and Hilda Pollaczek-Geiringee Praktische verfahren der gleichungsauftssung ZAMM-Journal of Ayplied Mathemalics and SichanicssZZiisschenn für Angewandle Mathemalík und Mechanaák, 9(1):58 77, 1929.
Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for gener- ative adversarial networks. arXiv preprint arXio:1802.0500 2018.
Maximilian Mueller, Tiffany Joyce Vlaar, David Rolnick, and Matthias Hein. Normalization layers are all that sharpness-aware minimization needs. In Shirth- Conference on Neural Information Processing Systems, 2023. URL https: /ttper///ppnrrview .et//foouurrit--1rra.3......
Kyunghun Nam, Jinseok Chung, and Namhoon Tee. Almost sure last iterate convergence of sharpness-aware minimization. 2023.
Lawrence Page, Sergey Brin, Rajoev Motwani, and Terry Winograd. The PageRank citation ranking Bringing order the web. In The Web Conference, 1999. URL attps://api semanticscholar. org/ Io
12

---

Published in Transactions on Machine Learning Research (01/2025)
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp. 4596-4604. PMTR, 2018.
Dongkuk Si and Chulhee Yun. Practical sharpness-awart aware minimization cannot converge all the way optima. arx preprint (xx::3366008800 2023.
Gowthami Somepalli, Liam Fowl, Arpit Bansal, Ping Yeh-Chiang, Yehuda Dar, Richard Baraniuk, Micah Goldblum, and Tom Goldstein. Can neural nets learn the same model twice? Investigating reproducibility and double descent from the decision boundary perspective. In Proceedings of the IEEE/CVE Conference on Compuler Vision and Pallern Recognition, pp. 186699-13700 2022.
Andreas Steiner. Vision transformer. btttps://githen cona/poogeeerrssearcchiiiii.... -transforerr 2022.
Ila Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jes- sica Yung, Andreas Steiner, Daniel Keysers, Jakob Usækoreit, Mario Lucic, and Alexey Dosovitskiy. ML.P- Mixer: An all-MLP architecture for vision. arXiv preprint arXiv:2205.0106 2021.
Thijs Vogels, Sai Praneeth Karimireddy, and Martin Jaggi. PowerSGD: Practical low-rank gradient com- prossion for distributed optimization. Advances in Neural Information Processing Systems, 32, 2019.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv:1804.074. 2018.
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lueas Beyer. Scaling vision transformers. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pp. 2104-12113, 2022.
Qihuang Zhong, Liang Ding, Li Shen, Peng Mi, Juhua Liu, Bo Du, and Dacheng Tao. Improving sharpness- aware minimization with Fisher mask for better generalization on language models. arXiv preprint orXiv:2210.05.p9 2022.
Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha Dvornek, Sekhar Tatikonda, James Duncan, and Ting Liu. Surrogate gap minimization improves sharpness-aware training. arXiv preprint Xiv:22003 08065, 2022.
13

---

Published in Transactions on Machine Learning Research (01/2025)
Appendix
Table of Contents
A Proofs for Section 3 (Algorithmic derivation)
15
B Proofs for Section 4 (Convergence analysis)
15
C Implementation
18
C.1 vSAM variant with normalization
18
D Experiments
18
D.1 Computational resources
14

---

Published in Transactions on Machine Learning Research (01/2025)
A Proofs for Section 3 (Algorithmic derivation)
Lemma 3.1. The LMO under nuclear constraint in (9) can be efficiently computed as
Imo (W) pun
where u and v are the left and right singular vectors associated with the top singular value of W.
Proof. We are asked to find X satisfying mingxx (W,X), which relates to the definition of the dual norm,
1 min (W,X)= max ( X) =: ||]|| where X x==-xx. |X
(11)
e SS
P
where ||-Il is the spectral norm. We can rewrite (11) into a vector problem that we know how to solve. Note that the spectral norm unitary invariant since it Schatten norm, i.e., ||||| = ||UAV So consider instead the SVD, W diag(o)) for which the dual norm simplifies,
Wi || = I|U diag(o) [T = ||al = max (0,
(12)
||
In other words, need ||0||00 (o,x) which is attained by
1 j er 0i
By otherwise
Vj e [d].
(13)
This solution also trivially satisfics the ||
constraaai.
Tt remains to translate x (13) into a matrix solution X, which can be done by rewriting the objective and constraint in (12) back into the problem in (11),
||>||] = diag(x))l. = IU diag(=)V ||. = ||
(o, .) = ur(V V diag(o)) T diag(x))
(14)
= =tr(V diag(o)U UU diag(r)V = (U diag(o)VT UU diag(z)V = (W,U diag(x)VT)
So it must be that X = -pX = di)g (=) For our choice the solution reduces to X ppuv" T. 0
B Proofs for Section (Convergence analysis)
Theorem 4.3. Suppose Assumptions 4.1 and 4.2 hold with L p. Then, (10) with the (arbilrary) norm
perturbation c AAgg magl:ll: (x*) e)) satisfies the following descent inequality with p 2 0 and 7
[( S* pm (f(x*) - -x11- -=+)))/\\VV[+)))) 2222
Furthermore, for i = ? = min
log(a); 1 with a
it follows that
kk
-f = (00eep(p555) K )0 +
f(a kE[K]
min
Cp²
with =0 f(0) c= , = maxxex ||=||//1|=|| and &. maxrex /==1|//1/=|1..
Proof. By using smoothness have
1(a*+1) S {(x*) + (VS(a*), 2 k+1 ~ 111 *11
= =f(x²) - 7k (V)
= =s(x²) - 7 (1- - 2ل 2 (V}
V
1]6(1 (78
vf(x²))
(15)
15

---

Published in Transactions on Machine Learning Research (01/2025)
Let us develop the last term of (15) with some 0:
- vf(=")) (f(E")
c (ie) v vf((*))
= - ( V/(x") vs(xa))
(Assumption 4.1) 2 / -opt, v)((*))
- - -V cpe
(Assumption 4.2) 2 (Jic 2 cº 12 -B
/ 113 1 (***
cpe 11
[t - ccppl 1 +
]1 + I
cpe 113
(16)
where we have used the definition of dual norm in the third equality. The first term of (16) is positive by assuming μ 2 2 + 1² which is feasible when μ (pick = L).
2
c c= VV/LL
Plugging (16) into (15) (and ignoring the last good terms of (16)) we get
f(* S f(x²) 7 |I f 113 7x(1 - (( 7. 11 v
(Assumption 4.2) S f(a"t) - 3 IIV /
112 (x - I. V/ "'''0" 1|
S f(x²) - 3 IIV 112 74(1 - -XLL) Vf |+ TE
where the last line uses E, == max
|اا| and [[ S From the PI, condition, f(x) f S [ল].
#1 [ Vz X, (implied by strong convexity and Lipschitz continuity) it follows that
f(a -f s (1 - (m) f) --) (1 1 (L)) 2.0'00 (17)
2
Set ?/k = y and consider two cases.
Case I: When |[] 2 "L' Pe" for all up to K, the second last term 7(1 ((1= dominates, 21-74) leading to geometric decay. We have
f(x*) f 141-py (1-))
(1-14))
Case II: On the other hand, when |f(x>))|1 -
2(1-*L) we have from the PL condition that
f(x²) f- < V 11 < E [V][=>)
PL'PE 8p(
At any k, either of the two cases hold, so we can upper bound by their sum as with g = max,ex [
|اا|
follows
min kE[K]
,K) Do (1
8p(1-21)²
$ (1 )n7)* K Do
72 pf
< exp(-1uyk3 Ao
where for simplicity we have picked 7 S 21 in the second last inequality and the last inequality uses (1 x exp(-xK) for a my
16

---

Published in Transactions on Machine Learning Research (01/2025)
Consider the stepsize ? min pk 27. where the scalar is to be decided. There are two cases.
First, 27. > ak then = nk and consequently
L'pestea 2
min f(x*) - fr S exp(-c))oo kE[K]
(1 (1++
where we have picked e = log( 2g2 E$p= to optimize the bound.
On the other hand, if 1 v c 22 ak
then 7 = 21 S uk
and consequently
d'''p''s"" 2μ²k2
min f f(x²) - f v exp( P. K)4o kE[K]
where the second term uses that S μk
In any case we can upper bound using the sum of the two bounds
min KE[K]
f(ah) - f s =xp((= + (2 22) 2p
L'p'ste
+ K))
L'PESte" 12
where hides logarithmic factors. This completes the proof.
Theorem 4.6. Suppose Assumplions 4.2 and 4.5 hold and || S for all [K]. Then (10) salisfies the following convergence guarontee for Pk 2 0 and "ik (0,2/L).
Ag+C; Σ KK +Pt+C, K-1 TAPE
min
Ex (1(1-11//3)
C with = f6o) 1332 and C2 I.G.
In porticular, for constant stepsize 7he 7 = 1/LVK and p 1
VR the following rate is oblained
[V = 00 3 C2
VK
LK372
Proof.
Using smoothness we have
s f(x*) )(
2
=((x²) - 74 (V +
-(x*) - 71 -
76(1 - L)))) (Vf(z=)) ()= vf(z))
(Assumption 4.2) < f f2** - (1 - 2 Vf ) 2
vf(*t) vf(2tt)
(18)
What remains is the last term of (18):
(Vi((*)
(Assumption 4.2) < Pn L|e [/<>>)
(Assumption 4.5) < LG
(19)
Plugging (19) back into (18) and using that || 1|2 < E we have
f(zk+1) < f(x))( - - V Li0''s + 7x (LL
Subtracting f both sides, telescoping and rearranging completes the proof.
17

---

Published in Transactions on Machine Learning Research (01/2025)
C Implementation
Algorithm 3 Pseudo-code for distributed implementation of vSAM.
//Compute the gradient each device distinct batch: batch)
grad &- compute_ petpautte_grdiinttsiiieeee
3: if not m-SAM then
4: //Average and synchronize the gradients:
5: grad + all_ aal_reduce(grad operation "mean")
end if
//Perturb the model (while maintaining a sparse representation the perturbation):
sparse_epa eps compute perpppte__eeruurbatttiommmddddd model f add_ aaparss_perturbatiiono
sparse_appppppee _eps)
10: //Compute the gradient on the perturbed model on each device and synchronize:
11: grad + compute, lerel_rrrdieee endddidelmmod,,, batch)
12: grad + all_ all_reduce( (grad, operation = "mean")
13: //Remove the (possibly distinct) perturbation on each device:
14: model + subtract_ subtract_sparse_ee.m wwtpsseppettrrrteeee
petseeeurrurbbtiiooooeeee sparse_ eps)
15: //At this point model and grad identical across all devices:
16: model + peeiiiizer_stepmoddeee grad)
C.1 sSAM variant with normalization
In this section, we derive for completeness variant of vSAM which normalizes the layerwise perturbations, that we refer to as vSAM-Norm.
It follows from the Eckart- Bockart-Young-s Young-Mirsky theorem that a rank-1 approximation to a matrix E guxm is given by its top singular value and associated singular vectors u, follows
arg min ||A X||2 aue
(20)
Xmann(X)g1
Suppose for simplicity that the network parameters composed only of matrices, i.e., (M Mi) such that fzxx) fV,, V V,, (z)). Define a; to be the top singular value and u,, 2 the associated singular vectors of V fxr). Then the SAM perturbation update on the rank-1 approximation instead of the gradient reduces to the following
E po /c with c= [| ..... onnne 1l2 1122 VE,
(21)
In other words, the resulting ascent update for vSAM-Norm is identical to Step in vSAM but with a scaling factor of o VVE Ei=r for the 7th mairix. Similar to uSAM, the perturbation update be seen raineed LMO on H, so convergence gaurantees of Theorem 4.3 also applies to sSAM-Norm. Experimental results can be found in Table 8 of Appendix D, where we find that vSAM-Norm has improved performance in certain cases, but is generally dominated by vSAM.
D Experiments
For training from scratch see Tables 6 to 10 and Figure 4.
For fine-tuning Tables 11 13 and Figure 5.
Computational resources and wall-clock time are specified Appendix D.1.
18

---

Published in Transactions on Machine Learning Research (01/2025)
Table 6: The architectures used for CIFARI0/100 training from scratch. We take the original model con- figuralions used in Dosovitskiy et al. (2020); Tolstikhin et. al. (2021) but changes the patch size to 4 to accommodate for the smaller image size 32 x 32.
Model
Embedding dim.
Depth
Heads
Patch size
ViT-S/4
384
12
ViT-B/4
768
12
12
Mixer-S/4
512
8
4
Mixer-B/4
768
12
4
Table 7: Baseline (AdamW) hyperparameters for training from scratch. We use cosine learning rate schedule with linear warmup. We use standard augnentations (random cropping and flipping) and AutoAngment.
Hyperparameter
Value
Learning rate
0.0005
Label smoothing
0.1
Weight decay
0.05
Warmup epoch
10%
Epochs
300
Dropout rate
0.0
Drop path rate
0.1
Gradient clipping
Disabled
Batch size
128
CIFAR10 VIT-5/4
CIFAR10D VIT-5/4
1est accuracy
Test accuract
0.94
AdtamtWW
0.65
Adammti
0.32
.SAM
:SAM
SAM
0.60
SAM
SAN-ON
SAMAON
0.05
0.10
0.200.30
0.50
1.00
2.00
5.00
10.00
20.00
00
0.05
0.10
0.20
0.50
1.00
2.00
5.00
10.00
20.00
CIFAR10 VIT-B/4
CIFAR100 VIT-B/4
0.96
0.70
lezz accuracy
accuracy
0.65
SAM ON
Adammii
0.92
.SAM
0.90
0.50
10.0
CIFAR1O Mixer-5/4
CIFAR100 Mixer-S/4
Test accurasy
Adamm
.SAM
Test accuracy
0.70
AdamTW
+5AM
0.86
0.2
0.5
1.0
2.0
5.0
10.0
20.0
50.0
0.1
0.2
0.5
1.0
2.0
5.0
10.0
20.0
50.0
CIFAR10 Mixer-B/4
CIFAR100 Mixer-B/4
Test accuracy
Test accuracy
SAM
SAM ON
0.03
SAM-ON
0.50
6x-001
Figure 4: Sweep over for pretraining results in Table 2.
19

---

Published in Transactions on Machine Learning Research (01/2025)
Table 8: vSAM-Norm variant on ViT-S/4 and Mixer-S/4 (see Appendix C.1 for the formulation).
Method
AdamW
vSAM
uSAM-Norm
Perturbation ion start epoch
30
30
Dataset
Model
CIFARIO
Mixer-S/4
91.08
92.24
(+1.16)
92.10
(+1.02)
ViT-S/4
93.37
94.39
(+1.02)
93.95 (+0.58)
CIFAR100
Mixer-S/4
68.03
70.79
(+2.76)
70.17 (+2.14)
ViT-S/4
68. 17
70.60
(+2.43)
71.13 (+2.96)
CIFARID VIT-5/4
CIFAR100 VIT-S/4
0.35
That arcurocy
0.90
0.45
0.85
Adamily
Adamid
0.80
(SAM-4tmm upoch 30
JSAM epoch
.SAN Norn sooch 30
JAN eporn
0.05
CIFARIO Mixer-S/4
CIFAR100 Mixer-S/4
0.50
sccurace
0.69
0.55
Adamni
Tel.
0.50
0.1
0.2
0.5
1.0
1.0
5.0
01
0.2
0.5
1.0
2.0
5.0
Table 9: SAM-ON with perturbation delay interestingly does not improve the test accuracy.
Perturbat ion start epoch
30
Dataset
Model
CIFAR100
VIT-S/4
72.70
71.85
CIFAR100 VIT-S/4
AstariWW
Text accurecy
0.72
epach
10
20
Table 10: Ablation of the number of power iterations. Maybe surprisingly, increasing the number of power iterations does not necessarily lead to higher test accuracies. One possible explanation that the warmstart in the power iterations (Algorithm 2) might act as a variance reduction like mechanism by implicitly incorpo- rating previous stochastic gradients in the approximation. Applying variance reduction to the perturbation direction has shown to empirically be beneficial (Li & Giannakis, 2024). Investigating this connection is interesting future work.
Method
v/AM 50
Power iterations
Model
Dataset
VIT-B/4
CIFAR10
94.44
94.35
CIFAR100
70.15
70.27
VIT-S/4
CIFARI0
94.39
94.47
CIFAR100
70.60
70.24
20

---

Published in Transactions on Machine Learning Research (01/2025)
CIFAR10 VIT-B/16
CIFARI00 VIT-B/16
0.99
Text accuricc
6,920
Adamiiv
Adamil
6.410
1560
SAN
SAM
0,95
SAN-ON
5.005
0.01
0.32
0.55
0.10
0.0
0.50
1.00
2.00
5.00
10,00
20.00
0.01
0.02
0.05
0.00
0.50
100
2.00
5,00
10
20.00
CIFARID VIT-S/16
CIFAR100 VIT-S/16
0.0700
0.000
0.975
Adammi
Adamm:
SAISONN
0.9100
0.01
0.02
0.00
6.10
0.25
1.00
1.00
2.00
CIFARIO Mixer-B/16
CIFARICO Mixer-B/16
Test stcurs3
0.07
Adamiiii
Adaniw
Teat
SAN
.CAM
SAUI
I5AM
0.05
SAM-OON
6.2
0.00
5.0
0.20
0.50
1.00
2.00
5.00
0.02
0.05
0.10
0.20
0.50
1.00
2:50
5.00
Figure 5: Sweep over for fine-tuning results in Table 4. Mean and standard deviation for each configuration is computed over 3 independent runs.
Table 11: Bascline (AdamW) hyperparamoters for finc-tuning. We cosine learning rate schedule.
Hyperparameter
Value
Learning rate
0.0001
Label smoothing
0.0
Weight decay
0.0
Warmup epoch
Epochs
10
Dropout rate
0.0
Drop path rate
0.0
Gradient clipping
Disabled
Batch size
96
Table 12: Hyperparameters for fine-tuning on GLUE.
Hyperparameter
Adam
SAM-
DSAM
SAM-ON
Model
Protrained BERT base (uncased) (Devlin et al., 2018)
Epochs
5
Batch size
32
Learning rate
4 10-b (cosine schedule)
Learning rate warmup
10%
P
{0.01, 0.02, 0.05 }
0001 0.02, 0.05 }
(..01, 0.02, 0.05
Perturbation delay
10%
10%
0%
All methods uses Adam as the base optimizer.
21

---

Published in Transactions on Machine Learning Research (01/2025)
Table 13: Even if AdamW is given a more refined budget to optimize the learning rate (Ir = 5e- - 05), (V)SAM with the same learning rate improves even further when using the default paramcter p = 0.05 (p 0.02). Providing AdamW with twice the computational budget also does not close the gap.
Test accuracy
Ir 0.0005
AdamW
97.19
1 0.0001
AdamW (2x epochs)
98.49 (+0.11)
AdamW
98.38
SAM
98.76 (+0.38)
1r=0. 0.00005
AdamW (2x epochs)
98.49 (-0.13)
AdamW
98.62
SAM
98.89 (+0.27)
#SAM
98.81 (+0.19)
1=0000001
AdamW
98.44
SAM
98.72 (+00288)
D.1 Computational resources
All experiments are run on either single NVIDIA V100 GPU, Table 14 shows the wall-clock time.
Table 14: Wall-clock time for training VIIs and MLP-Mixer models from scratch on a single V100 GPU.
Both vSAM (with one power iteration) and SAM require the same computational complexity.
minutes
epoch
Dataset
Model
Method
CTFAR10
Mixer-B/4
=SAM
4.4
SAM
4.4
Mixer-S/4
sSAM
1.6
SAM
1.6
VIT-B/4
VSAM
4.3
SAM
4.2
VIT-S/4
SSAM
1.6
SAM
1.7
CIFAR100 Mixer-B/ B/4
SAM
4.4
SAM
4.3
Mixer-SI S/4
sSAM
1.6
SAM
1.6
VIT-B/4
SSAM
4.3
SAM
4.2
VIT-S/4
SSAM
1.7
SAM
1.7
22


Reviews:
[{"review_id": "MSmjWgxpx1", "reviewer": "Reviewer_MzVp", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- The paper is well-written, with little need for improvement in writing. \n- the empirical results are sound and adequate. \n- some of the additional insights such as delay in perturbation are interesting. \n\nWeaknesses\n- there is almost no new theoretical results. the presented theorems are a rehash of previously known theorems applied to the low rank setting.\n- given the sam-on paper, the results are mixed at best. i am not sure if the method itself is very useful unless further combined, investigated vs sam-on.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes low rank perturbations instead of the inf norm based pertubations done previously for SAM. this results in substantial savings on memory.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- can u elaborate on table 1. how are there 22M params stored for perturbation instead of only 75k. i understand per iteration costs as explained on page4,5 but those still dont account for such a large discrepancy. \n\n- some of the competitors need more details within the paper when discussing the related sections. especially sam-on. \n\n- some of the writing needs cleaned up. please expand on how the nuclear norm constraint as claimed in the front matter of the paper is implied when using rank-one updates. \n\n- have u tried combining sam-on and vsam? Why not just do low rank perturbations on the normalization layers?", "broader_impact_concerns": "none", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "fnMVRVJZMA", "reviewer": "Reviewer_WxYQ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* The paper demonstrates excellent writing quality and organization, making complex concepts accessible and maintaining a clear logical flow throughout. The presentation of ideas is systematic and well-structured.\n* The experimental validation is comprehensive.\n* The theoretical foundations are rigorously developed, with formal proofs for convergence guarantees under arbitrary norm choice. \n\nWeaknesses:\n* The notation choices deviate from common conventions in optimization literature, particularly: Using $x$ to denote model weights instead of the more standard $w$ or $\\theta$\n* Computational efficiency trade-offs need clearer discussion: While the method achieves significant memory savings, the time complexity implications from the additional SVD computation are not thoroughly analyzed.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work proposes νSAM, a variant of Sharpness-aware Minimization that significantly reduces memory requirements by modifying the perturbation constraint. This achieves nearly 1/3 of the memory footprint compared to standard SAM while maintaining comparable performance. They demonstrate the effectiveness of νSAM through comprehensive experiments. They provide theoretical foundations by proving convergence guarantees for SAM with arbitrary norm choice, even with fixed perturbation radius. This extends the theoretical understanding of SAM-based optimization.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* A formal analysis of the method's time complexity compared to standard SAM would strengthen the paper\n* It's better to use some commonly-used notations", "broader_impact_concerns": "No", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "JsA40p1XYf", "reviewer": "Reviewer_Utmq", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "***Strength***\n\nThe paper is generally well-written. Convergence theory is provided, although seems quite restrictive. Experimental results including training from scratch, as well as fine-tuning for language and vision tasks are included.\n\n\n***Weaknesses/questions***\n- Why is SAM framed as injecting noise to the gradient?… Claim about perturbing with Gaussian noise and not leading to convergence is quite misleading. SAM does NOT do that; it takes a gradient ascent step.\n- Nuclear norm constraint is defined as $\\sum_i \\sigma_i(A)$. How does Eq (7a) and (7b) follow? It just adds a rank-1 matrix.\n- I’m also a bit confused about the claim of Theorem 4.3. First, where do I see that $f$ is quadratic? Second, why is it reasonable to assume $L = \\mu$? Third, there already exists much more general theory: Theorem 2 in [2]. I find it a bit misleading to compare with the generalization bound from Foret et al (2020, Thm1).\n- If $f$ is strongly convex quadratic with $\\mu = L$, then the Hessian has uniform eigenvalues, or the hessian has condition number 1. Am I missing something? This is an extremely restrictive setting.\n- Delaying the perturbation also has been considered in [3], with theoretically principled explanation.\n- Experimental results are not very convincing. SAM-ON, which has similar compression effect as the proposed method, works better in quite many settings.\n- I believe the practical concern about SAM is having to evaluate gradient twice, more than the storage issue. In this regard, how does this work compare with [1]?\n\n[1] J. Du et al. (2022) “Sharpness-Aware Training for Free”\n\n[2] M. Andriushchenko et al (2022) “Towards Understanding Sharpness-Aware Minimization”\n\n[3] A. Agarwala “SAM operates far from home”", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a variant of sharpness-aware minimization (SAM) based on nuclear norm constraints on the ascent step, which uses 1/3 memory compared to the vanilla SAM, while achieving comparable performance.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I find Theorem 4.3 to be too restrictive. Even without assuming $f$ is quadratic, having $\\mu=L$ as an assumption is not convincing. Experimental results are also somewhat inconclusive—it’s hard for me to find a reason to use $\\nu$SAM compared to SAM-ON, for example. I think addressing these points are important.", "broader_impact_concerns": "Not applicable.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "ibAewlkslN", "reviewer": "Reviewer_ffMj", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n1. The paper tackles an important problem on the additional memory burden of SAM, particularly relevant when fine-tuning large models such as large language models (LLMs).\n2. The writing is clear, and the paper is well-structured and easy to follow.\n3. The experiments are thorough and provide convincing evidence for the effectiveness of the approach.\n4. The proposed method is both sound and reasonable. Interpreting SAM's perturbation from a norm perspective is particularly interesting.\n\n**Weaknesses:**\n1. $\\nu$SAM does not demonstrate a significant advantage over SAM-ON, which is the primary limitation of this approach. Note that SAM-ON also does not impose a significant memory burden compared to SGD and is much simpler.\n2. The assumptions in the convergence analysis, such as Assumption 4.1, are quite strong. For example, the strongly-convex assumption is rarely satisfied in practice, particularly for deep neural networks (DNNs).\n3. The paper lacks experiments with more commonly used architectures, such as ResNet, VGG, and Wide-ResNet.\n4. It appears that the current approach can only be applied to matrix weights in DNNs. A broader explanation of how the method could handle matrices with more than one dimensionality would be beneficial.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes $\\nu$SAM, a method designed to address the increased memory footprint associated with SAM training. $\\nu$SAM employs a layer-wise low-rank perturbation that leads to substantial memory savings. The paper provides convergence analysis and extensive experiments on various datasets.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. On Page 4, it would be helpful to provide references for the statement: \"since including the parameters did not improve performance and could even sometimes lead to performance degradation.\"\n2. Include experiments on more commonly used architectures, such as ResNet, VGG, and Wide-ResNet, particularly for CIFAR-10/100.\n3. It would be beneficial to provide convergence analysis without relying on the strongly-convex assumption, which is rarely valid in DNNs.\n4. Additionally, many other works have applied random weight perturbation to improve generalization and have shown promising performance (e.g., [1], [2]). It would be helpful to discuss these works when claiming that \"normalized Gaussian noise (GaussSAM) exhibits non-convergence.\"\n\nReferences:\n\n[1] Bisla et al. \"Low-pass filtering SGD for recovering flat optima in the deep learning optimization landscape.\" AISTATS 2022.\n\n[2] Li et al. \"Revisiting Random Weight Perturbation for Efficiently Improving Generalization.\" TMLR 2024.", "broader_impact_concerns": "No ethical concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper studies an important topic of optimization with the goal of keeping the model generalization at a good level. While the proposed approach has a clear limitation due to its design based on the matrix structure, as implied by the low-rank approach, it shows promising results on practical benchmarks and allows for significant memory savings. The paper is very relevant to TMLR and its claims are supported by sufficient empirical evidence, which is why I believe it should be accepted.
Decision: Accept (poster)