Implicit Regularization of Sharpness-Aware Minimization for Scale-Invariant Problems
Authors: Bingcong Li
Liang Zhang
Niao He
Venue: NeurIPS 2024 poster (2024)
Keywords: sharpness-aware minimization
LoRA
implicit regularization
finetuning
language models
Primary area: optimization_for_deep_networks
TL;DR: We leverage implicit regularization to enhance the computational efficiency of sharpness-aware minimization (SAM).

Abstract: Sharpness-aware minimization (SAM) improves generalization of various deep learning tasks. Motivated by popular architectures such as LoRA, we explore the implicit regularization of SAM for scale-invariant problems involving two groups of variables. Instead of focusing on commonly used sharpness, this work introduces a concept termed *balancedness*, defined as the difference between the squared norm of two variables. This allows us to depict richer global behaviors of SAM. In particular, our theoretical and empirical findings reveal that i) SAM promotes balancedness; and ii) the regularization on balancedness is *data-responsive* -- outliers have stronger impact. 
The latter coincides with empirical observations that SAM outperforms SGD in the presence of outliers. 
Leveraging the implicit regularization, we develop a resource-efficient SAM variant, balancedness-aware regularization (BAR), tailored for scale-invariant problems such as finetuning language models with LoRA. BAR saves 95% computational overhead of SAM, with enhanced test performance across various tasks on RoBERTa, GPT2, and OPT-1.3B.

Implicit Regularization of Sharpness-Aware Minimization for Scale-Invariant Problems
Bingcong Li
Liang Zhang
Niao He
Department of Computer Science ETH Zurich, Switzerland
{bingeong. 1i, liang. zhang,
ethz. ch
Abstract
Sharpness-aware minimization (SAM) improves generalization of various deep learning tasks. Motivated by popular architectures such as LoRA, we explore the implicit regularization of SAM for scale-invariant problems involving two groups of variables. Instead of focusing on commonly used sharpness, this work introduces concept termed balancedness, defined the difference between the squared norm of two variables. This allows us to depict richer global behaviors of SAM. particular, theoretical and empirical findings reveal that i) SAM promotes bal- ancedness; and ii) the regularization on balancedness is data-responsive - outliers have stronger impact. The latter coincides with empirical observations that SAM outperforms SGD in the presence of outliers. Leveraging the implicit regularization, we develop resource-efficient SAM variant, balancedness-aware regularization (BAR), tailored for scale-invariant problems such as finetuning language models with LoRA. BAR saves 95% computational overhead of SAM, with enhanced test performance across various tasks on RoBERTa, GPT2, and OPT-1.3B.
Introduction
Sharpness-aware minimization (SAM) is emerging as an appealing optimizer, because it enhances generalization performance on various downstream tasks across vision and language applications (Foret et al., 2021; Chen et al., 2022; Bahri et al., 2022). The success of SAM is typically explained using its implicit regularization (IR) toward a flat solution (Wen al., 2023a).
However, existing results only characterize sharpness/latness near local minima (Wen et al., 2023a). Little is known about early convergence, despite its crucial role in SAM's implicit regularization (Agarwala and Dauphin, 2023). In addition, thcoretical understanding of SAM highly hinges upon the existence of positive eigenvalues of Hessians (Wen et al., 2023a), leaving gaps in nonconvex scenarios where the Hessian can be negative definite. The limitations above lead to our first question (QI): can we broaden the scope of implicit regularization depict global behaviors in SAM?
Moreover, scenarios where SAM popularizes often involve certain form of data anomalies, such as outliers and large data variance. SAM has provable generalization benefits sparse coding problems in the small signal-to-noise ratio (SNR) regime (Chen et al., 2023). Remarkable performance SAM also observed under distributional shifts, e.g., domain adaptation (Wang al., 2023), meta-learning (Abbas et al., 2022), and transfer learning in language models (Bahri et al., 2022; Sherborne et al., 2023). Evidences above motivate second question (Q2): implicit regularization of SAM reflect its enhanced performance under data anomalies?
This work answers both QI and Q2 within a class of scale-invariant problems. The focus on scale- invariance is motivated by its prominence in deep learning architectures. Consider variables x Rd
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

abs||| 1
1lo.1l lloy1l
abs[[]/ll - -11
||o.11 ||o,1l
30
1.2
20
10
0.4
0
50000
100000
50000
100000
50000
100000
50000
100000
iter
iter
iter
iter
SGD a
a
a 0.5
SGD g
or
(a) non-overpprramelrizee (NOP)
(b) overparametrized (OP)
Figure 1: Implicit regularization of SAM on balancedness. The losses for NOP and OP are B|||xy
(A + [l and B[[|x y y- - (a + am)ll respectively. Here, A is the ground truth matrix, N is the Gaussian noise, and a controls the SNR. Left of (a) and (b): |xxx|* l || vs. iteration. Right of (a) and (b): |||xx || ||gy, | vs. iteration, where (gx, Sy,) denotes stochastic gradients.
and y R both in high- digenniimnsional space. The problems of interest can be categorized into aanooverparameetizattonn (NOP) and overparametrization (OP), based on whether the dimension of variables (d1 + d2) is greater than dimension of dom f.
NOP: min f (xy EE (xy ')]
(la)
x.y
OP:
min fo(x y) - Be-DD_Leee [f6(x y)]
(Ib)
x.y
Here, d1 = d2 is assumed for OP, and D denotes the training data. For both cases, the losses are nonconvex in y). Scale-invariance refers to that yax,y/a) share the same objective value Va o 0. It naturally calls for implicit regularization from optimization algorithms to determine the value of a. We focus on two-variable problems in the main text for simplicity and generalize the results to multi-layer cases in the appendix. Problems (1a) and (1b) are inspired by widely-adopted modules in deep learning, where low rank adapters (LoRA) for finetuning language models NOP, and softmax in attention falls in OP framework (Hu et al., 2022; Vaswani et al., 2017).
This work studies SAM's implicit regularization on balancedness, defined as B, [1 -Iyy lI?). Balancedness is a useful alternative to sharpness for (1) because: i) it enables us to go beyond local minima and describe the behavior over SAM's entire trajectory; ii) analyses and assumptions can be significantly simplified when working with B,; and, iii) it enables a data-driven perspective for understanding SAM. Building on balancedness, answer our major questions.
For QI, we prove that even with imbalanced initialization, SAM drives |B. -> for OP, while ensuring small [B,] in NOP. In contrast, also prove that balancedness of SGD is unchanged over iterations. This clear distinction between SAM and SGD is illustrated in Fig. 1. Thanks to the results implicit regularization have requirement the batchsize compared to (Wen et al., 2023a) and can be extended to explain m-sharpness (Foret et al., 2021).
adoption of balancedness,
Regarding Q2, we present analytical and empirical evidences that data anomalies (e.g., samples with large noise) have stronger impact on balancedness for both NOP and OP. Fig. showcases an example where SAM is applied on the same problem with different SNRs. Smaller SNR (i.e., larger a) promotes balancedness faster. Being more balanced with noisy data also aligns well with previous studies (Chen et al., 2023; Wang et al., 2023), which show that SAM performs better than SGD under data anomalics. This data-driven behavior of SAM is well depicted through balancedness.
Our theoretical understanding on balancedness also cultivates practical tools. In particular, we explicify the implicit regularization of SAM data-driven regularizer. When applied on top of, e.g., SGD, it enables computationally efficient variant of SAM, balancedness-awwe -wware regularization (BAR), suited for ceale-invarian problems such as finetuning language models with LoRA (Hu et al., 2022). BAR eliminates the need to compute the second gradient in SAM, thereby significantly reducing overhead in large-scale settings. BAR improves the test performance of LoRA on three representative downstream tasks on RoBERTa, GPT2, and OPT, while saving 95% computational overhead of SAM. Moreover, this is the first efficient SAM approach derived from SAM's implicit regularization. In nutshell, our contribution can be summarized as:

---

Theories. Balancedness is introduced as a new metric for implicit regularization in SAM. Compared to sharpness, balancedness enables to depict richer behaviors- SAM favors balanced solutions for both NOP and OP, and data anomalies have stronger regularization on balancedness. Practice. Implicit regularization of SAM made explicit for practical merits. The resulting approach, balancedness- regularization (BAR), improves accuracy for finetuning language models with LoRA, while significantly saving computational overhead of SAM.
Notation. Bold lowercase (capital) letters denote column vectors (matrices); || stands for & (Frobenius) norm of a vector (matrix), and (.) refers to transpose.
1.1 Related Work
Related topics are streamlined here, with comprehensive discussions deferred to Apdx. A.2.
Scale-invariance in deep learning, Scale-invariant modules prevalent in modern neural networks, such as LoRA, ReLU networks, and softmax in attention. However, not yet fully understood, especially from theoretical perspective. Neyshabur et al. (2018) develop scale-invariant PAC-Bayesie bounds for ReI.U networks. A scale-invariant SGD is developed in (Neyshabur et al., 2015), and this approach becomes more practical recently in (Gonon et al., 2024). Linear neural networks entail scale-invariancs and overparametrization simultaneously, and IR of GD)GD on quadratic loss is established in (Arora al., 2018; Du al., 2018; Gidel al., 2019). IR of GD for softmax attention in transformers is studied in (Sheen et al., 2024) assuming linearly separable data. It is pointed out (Dinh et al., 2017) that sharpness is sensitive scaling, while our results indicate that when taking the training trajectory into account, SAM excludes extreme scaling,
problems are
Mechanism behind SAM. To theoretically explain the success of SAM, Bartlett et al. (2023) analyze sharpness on quadratic losses. Wen et al. (2023a) focus on sharpness of SAM near the solution manifold on smooth loss functions, requiring batchsize be 1 in the stochastic case. Andriushchenko and Flammarion (2022) consider sparsity of SAM on (overparametrized) diagonal linear networks on regression problem. Chen et al. (2023) study the benign overfitting of SAM two-layer ReLU network. In general, existing studies on SAM's implicit regularization focus more on sharpness and do not fully capture ccale-invariance In comparison, our results i) are Hessian-free and hence sharpness-free; ii) have no constraint on batchsize; and iii) hold for both NOP and OP.
SAM variants. Approaches in (Kim al., 2022; Kwon et al., 2021) modify SAM for efficiency under coordinate- wise ill-scaling, while our results suggest that SAM favors balancedness between layers. Computationally efficient SAM variants are developed through reusing sparsifying gradients (Liu et al., 2022; Mi et al., 2022); stochastic perturbation (Du et al., 2022a); switching to SGD) (Jiang al., 2023); and connecting with distillation (Du al., 2022b). Our BAR be viewed as resource-effictte SAM applied specifically for scale-invariant problems such as LoRA. Different from existing works, BAR is the first to take inspiration from the implicit regularization of SAM.
2 Preliminaries
This section briefly reviews SAM and then compares sharpness with balancedness. For a smoother presentation, our main numerical benchmark, LoRA (Hu et al., 2022), is revisited in Sec. 5.
2.1 Recap of SAM
Sharpness-aware minimization (SAM) is de- signed originally to seek for solutions in flat
Algorithm SAM (Foret et al., 2021)
1: Initialize: Wo. p. T,
basins. The idea is formalized by enforcing small loss around the entire neighborhood in parameter space, i.e., minw max]jells h( ww+ E), where p is the radius of considered neigh-
2:
for 1 0..... T do
3: Sample & to get minibatch Mr
4: Define stochastic gradient on Me as Vh+(-)
5: Find c, (w:,(w://lll
borhood, and h(w): = Eg [he (w)). Practical implementation of SAM summarized un- der Alg. 1. It is proved in (Wen et al., 2023a)
6: Update via W = Wt - 7/, +c)
7: end for
that [[ (w] A 0 (in line 5) holds for any P under most initialization. Based on this result and similar to (Dai et al., 2023), we assume that SAM iterates are well-defined.
3

---

Limitation of sharpness. Coming naturally with SAM the so-termed sharpness, given by S(w) ma h(w + E) h(w). When [Vh(w)] + 0, S(w) can be approximated using (scaled) largest eigenvalue of Hessian (Zhuang et al., 2022). This approximation is widely exploited in literature to study the implicit regularization of SAM. Consequently, most results only hold locally behaviors near [V](w)] -> 0 are studied. In addition, sharpness (the largest eigenvalue) is not always informative for scale-invariant problems (I), Consider h(z, y) = zy for example. The sharpness is for any (z, y) these points are not distinguishable in terms of sharpness.
2.2 Prelude on Balancedness
Balancedness :;:= 1 (|xx ||= -- turns out to be intriguing alternative to sharpness on the scale-invariant problem (1). Being global metric, balancedness is capable of describing the entire trajectory of an algorithm, regardless of proximity to critical points or definiteness of Hessian,
How does B, evolve in different algorithms? To set a comparing benchmark of SAM, we first borrow results from previous works on SGD. Following implicit regularization literature such (Arora et al., 2018, 2019b; Wen et al., 2023a), we consider SGD with infinitesimally small learning rate - 0 for the NOP problem (1a)
x+1 = X, -/xx: Y1+1 = Y: - '/8y.-
Theorem ((Arora et al.. 2018, 2019a; Ji and Telgarsky, 2019; Ahn et al., 2023)). When applying SGD on the NOP (Ia), the limiting flow with r7 + 0 satisfies ||x;||2 1lyal1 2 |xxl| ll for all > O. In other words, dB, dt
holds.
Theorem I shows that B, Bo given 7 + A graphical illustration can be found in Fig. 1 (a). Another interesting observation is that given the same initialization, B, is fixed for SGD regardless of training datasets. This suggests that SGD is less adaptive data. A similar result of Theorem can be established for SGD on OP. The full statement is deferred to Apdx. C.I; see also Fig. (b).
Merits of being balance. Because B, is preserved, SGD is sensitive initialization. For example, (xo. yo) and (2x;, 0.5yp) can result in extremely different trajectories, although the same objective value is shared at initialization. Most of existing works initialize B6 ss 0 to promote optimization benefits, because the variance of stochastic gradient is small and the local curvature is harmonized around a balanced solution. Take the stochastic gradient of NOP minibatch M for example
gx
M
gy
I
Σ SEM
V }] x.
(3)
SEM
Assuming bounded variance [|||| M ΣEEM ()5( (x 1)1 S 0², it can be seen that
the variance of [gx, gyy is bounded by [1|x|22 11)1))). In other words, among ((x, |> -)) gradient variance minimized if ||x|| ||y||, i.e., being balance. Moreover, block smoothness parameters L and LX also hint upon the difficulties for optimization, where large values typically correspond slow convergence (Bottou et al., 2018; Nesterov, 2004). With the help of Assumption I (in the next subsection), it can be seen that L Ln lly and LX = L,, |x|| In other words, large [B, implies difficulty for optimizing one variable than the other. For these reasons, balancedness well-appreciated in domains such as matrix factorization/sensing 2016; Bartlett al., 2018; Du et al., 2018; Ge et al., 2017). It is also observed that balanced neural networks are easier to optimize relative to unbalanced ones (Neyshabur et al., 2015).
a special case of (Ia) (Tu et al.,
2.3 Assumptions and Prerequisites
To gain theoretical insights of acale-invariant problems in (1), assume that the loss has Lipschitz continuous gradient on dom f following common nonconvex optimization and SAM analyses (Bottou al., 2018; Andriushchenko and Flammarion, 2022; Wen et al., 2023a).
Assumption 1. Let W E RRd xdz, and w E R. For each &, (S()) and 15 (w) in (1) have Lmn and L. Lipschitz continuous gradient, respectively.
Definition of L.: for tixed ||gx. Exzll S L [|]] X2

---

Scale-invariant problems are challenging to solve even simple problems in Fig. 1. Even GD can diverge on some manually crafted initialization (De Sa et al., 2015; Arora et al., 2019a). With proper hyperparameters this rarely happens in practice; hence, we focus on scenarios where SGD and SAM do not diverge. This assumption is weaker than the global convergence needed in (Andriushchenko and Flammarion, 2022), and is similar to the assumption on existence (Wen et al., 2023a).
3 SAM for Non-Overparamettrizd Problems
This section tackles the implicit regularization of SAM on NOP (la). Motivated by practical scenarios such LoRA, we focus on cases initialized with large [Bo].
When ambiguity absent, the subseript in fn and L, ignored in this section for convenience. Applying Alg. 1 on NOP, the update of SAM can be written as
x Xx + px
ye = Yt + pulgy.
(4a)
gx. (
gy: ( X;
(4b)
X X 7gxx
V++1 Y+ 788.
(4c)
where p > 0 is the radius of SAM perturbation; de == 1/v |1x, 2 gy, 12 and ft, V denote the loss, stochastic gradient minibatch Mr, respectively.
Theorem 2. (Dynamics of SAM.) Suppose that Assumption holds. Consider SAM for NOP in (4)
with sufficiently small p. Let B, == 2 (11x,11 - l1y.ll".. For some [Ae] = (pp2 L.) and n -> 0, the limiting flow of SAM guarantees that
dB,
[xx 2
By
+ A..
(5)
dt
1xx 2 gy.
2
Moreover. the change B, depends on the difference of stochastic gradients on x, and yı, i.e.
dB,
p |x.
By. |I| - (pp²))
S PV gx 1| |gy 11 1 O
(6)
dt
Unlike SGD for which d8, O. Theorem 2 states that the balancedness for SAM is driven by gradient
dt
ly. 2 To gain some intuition, if we estimate ||gx, 2 ggy 21 < ||y: 1| ||x,|| || difference lgg, 2 By.
based on (3) and ignore A., it can be seen that dE -pB, This indicates the contraction on [8,]- A graphical illustration on decreasing [B,], and its relation with gradient difference can be found in Figs. (a) and 2 (a). Moreover, this implicit regularization on balancedness is global as it holds for all t regardless of whether (x,y), is close to local optima. Thanks to adopting balancedness as the metric, Theorem also poses no requirement on the batchsize.
SAM promotes balancedness. As discussed in Section 2.2, unbalancedness is burdensome for optimization. SAM overcomes this by implicitly favoring relatively balanced solutions.
Corollary 1. (Informal.) Under some regularity conditions, there exists B > 0 such that whenever [B, By the magnitude of B, shrinks, where Bi can be found in (21) at appendix.
Corollary 1 shows that SAM promotes balancedness until [Be] reaches lower bounds B Because B depends on SAM's trajectory, we plot + f-" B dt using dotted lines for better visualization in Fig. (a). It can be seen that our calculation on B3 almost matches the balancedness of SAM after sufficient convergence. Being balance also reveals that the benefit of SAM come from optimization, which a perspective typically ignored in literature.
Noisy data have stronger impact on balancedness. Although our discussions extend to more general problems, for simplicity we consider the example in Fig. (a), i.c., B[]xxy T - (A + [N)]] where A is ground truth; N is data noise; and determines SNR. For this problem, noisy data directly lead to noisy gradients. It can seen in Fig. 2 (a) that smaller SNR coincides with faster decreasing of [B,)) To explain such data-responsive behavior in implicit regularization, Theorem 2 states that balancedness changes largely when the difference of ||gy. and |gx. is large. Since Elilgy. I|2 ||gx, [|2] x a² if assuming elements of N to be iid unit Gaussian variables, thus implies that small SNR (large a) offers large regularization on balancedness.
5

---

abs|]] [x|| (Н))))
119.11 - 119.11
|bs[[[ -|>|PP)
||.11l 1|o.1l
1.2
20-
0.8
0.4
10
0.0
0-
50000
100000
0.0
-10
50000
100000
50000
100000
SAM a 0.5
SAM c=1
SAM o 15
Iter
thres
thres
thres
SGD
SAM
oBAR
nBAR
(a) threshold of balancedness B,
(b) relation with regularization
Figure 2: Implicit regularization of SAM on NOP Eijlxy" (A + aN) [?]]. where a controls SNR.
(a) the threshold of balancedness Bţ in Corollary I; (b) implicit vs. explicit regularization.
Extension to LoRA (multi-layer two-variable NOP), For LoRA, the objective is minimize
D blocks of variables simultancously. i.c., min lEg [f H- It is established in Theorem 5 appendix that SAM cultivates balancedness in a layer-wise fashion, i.c., the magnitude of By == 1 [|x,,|l - lyr,)l?)) cannot be large for each l. However, the |dB,,//dt can be O(VD) times smaller than Theorem 2 in the worst case because of the additional variables.
Validation of IR modern architectures. Going beyond the infinitesimally small step size, adopt 1 = 0.1 on modern language models to validate our theoretical findings. We consider finetuning RoBERTa-large with LoRA for few-shot learning tasks. More details can be found later in Section 6.1. Balancedness of SAM and SGD on different layers in various datasets are plotted in Fig. 3. SAM has clear trend of promoting balancedness, aligning well with our theoretical predictions.
4 SAM for Overparametrizee Problems
Next, we focus on SAM's implicit regularization on OP (1b). Overparametrization enables SAM to have stronger regularization balancedness. Subscripts in fo and L, are omitted for convenience. SAM's per iteration update for OP can be summarized as
X, x, + puyy:
ÿ: ye Xu
(7a)
gx, fx )r.
gy. (
8
(7b)
X = X 7x.
Y++1 = Yt ngs.
(7c)
where ii := (o.[)((xTT)))))) x 12 |2 Si and S denote the loss, stochastic gradient on
minibatch Mr, respectively. Different from NOP, SAM has stronger regularization balancedness, where [B,] decreases whenever the norm of stochastic gradient is large. To see this, it is convenient to
define C2 := |||x||| llyell.. Note that C, S /2)
Theorem 3. Consider y + 0 for (7). The limiting flow of SAM on OP ensures a decreasing magnitude of B, whenever (x, y+)l C. O(pL(B.)). Moreover, the speed of decrease can be lower- and upper- bounded as
dB.
plf((xx] -C, - O(p2L(BB < dt S Pl vol + O.)
Given P - 0 and sufficiently noisy data, Theorem 3 implies that |B.| O. Moreover, Theorem 3 also states that the regularization power on balancedness is related to both gradient norm and balancedness itself. The elbow-shaped curve of [B, in Fig. 1 (b) demonstrates that the regularization power is reducing, as both gradient norm and balancedness shrink over time.
Noisy data have stronger impact on balancedness. As shown in Fig. 1 (b), balancedness is promoted faster on problems with lower SNR. This data-responsive behavior can be already seen from Theorem 3, because |dB,/dt is directly related with [x, y.)l. and (x,, [] is clearly larger when data are more noisy. other words, SAM exploits noisy data for possible optimization merits from balancedness (see discussions in Sec. 2.2). Overall, the implicit regularization on balancedness aligns well with the empirical observations in presence of data anomalies (Wang et al., 2023; Sherborne et al., 2023), where SAM outperforms SGD by large margin.

---

RTE
SST-5
MNLI
2.68
2,70
2.7
2.66
2.65
2.64
2,6
2.5
2.60
SAM
2.55
SGD
2.58
2.4
2.50
2000
4000 6000 8000 10000
0 2000 4000 6000 8000 10000
0
2000 4000 6000 8000 10000
Figure 3: Implicit regularization of SAM on LoRA. We consider few shot learning with LoRA RoBERTa-large. For datasets RTE, SST-5, and MNLI, 1st, 12th and 24th query layers' 2| are plotted, respectively. The layers are chosen to represent early, middle, and final stages of RoBERTa. The averaged B in Corollary is 0.37, 0.21, and 0.29, respectively.
Extension to m--sharpness. m--sharpness is variant of SAM suitable for distributed training. It is observed to empirically improve SAM's performance (Foret et al., 2021). ). mn-sharpness evenly divides minibatch Mr into disjoint subsets, i.e., (/t,)/=11 and perform SAM update independently on each subset; see (38) in appendix. It turns out that m-sharpness can also be explained using balanced- ness. With formal proofs in Apdx. C.3, the IR of m-sharpness amounts to substitute (x y:)l in Theorem 3 with 1 ΣØ= Iifk (x, y.)l. This means that the regularization on balancedness from m-sharpness is more profound than vanilla SAM, because m Σド= [x,, v.)l (x]
Finally, we connect balancedness with sharpness local minima of OP.
Lemma 1. Let W = {(x,y) = w. f'(ww 0, f" > 0} be non-empty. For the OP W
problem (Ib), minimizing sharpness within W is equivalent finding
This link showcases that by studying balancedness we can also obtain the implicit regularization on sharpness for free. A concurrent work also links balancedness with sharpness (the largest eigenvalue) for some one-hidden layer neural networks (Singh and Hofmann, 2024). Compared with (Wen al., 2023a), this is achieved with less assumptions and simplified analyses. More importantly, balancedness enables us to cope with arbitrary batchsize, to explain SAM's stronger regularization with noisy data, and extend results m-sharpness.
5 Implicit Regularization Made Explicit
Next, insights from our theoretical understanding of SAM are leveraged to build practical tools. We adopt LoRA (Hu al., 2022) as our major numerical benchmark for scale-invariant problems given its prevalence in practice. More diverse examples on both OP and NOP can be found in Apdx. A.3. Compared to full parameter-tuning LoRA is more economical in terms of memory not only for finetuning, but also for serving multiple downstream tasks. LoRA and its variants are actively developed and well welcomed by the community; see e.g., HuggingFace's PEFT codebase."
5.1 Overview of LoRA
Given a pretrained model with frozen weight W, E Rdiids on particular layer the objective of LoRA to find low rank matrices X; dR xr and Y, E Rd with & min(di, d2) such that the loss is minimized for a downstream task, i.c.,
C((W++XXX) + X,Y,))) (X,,Yi)
min
LoRA enjoys parameter efficiency for finetuning thanks to the low-rank matrices X; and Y, For instance, only requires 0.8M trainable parameters finetune 355M -arrameer RoBERTa-large (Hu et al., 2022). The outer product of X, and Y, induces scale-invariance, and the number of variables renders it NOP. The downside of LoRA, on the other hand, is the drop on test performance due to the parsimony trainable parameters. Unbalancedness is also unavoidable for LoRA, due
7

---

Algorithm 2 nBAR
Algorithm 3 oBAR
1: Initialize: learning rate { regularization coefficient {az)
1: Initialize: learning rate (nt) regularization coefficient {a+}
2: for t- 0..... --.....T111. do
2: for t 0,.. 0...., T do
3: Get stochastic gradient Bx, and gy,
3: Get stochastic gradient gx. and gy,
4:
if [|8x.|] || 2 |gy Il then
4: if |xx|| 2 || then
5:
X₁ f [1 Ocels [x,
5:
X + (1 oxeT)e
6:
Yr -- (1 Or (yt
6:
Yr +- (1 + or
7:
else
7:
else
8:
X₁ f (1 ai7 \x Yr. + (1 + )yt
8:
Xt + (1 + ac77 Xx
9:
9: Yt + (1 ay
10: end if'
10:
end if
Optimizer update (via Adam or SGD)
11: Optimizer update (via Adam or SGD)
12: end for
12:
end for
to the need of initializing X NNO, o²), Y, = O; see an example of RoBERTa-large in Fig. The unbalancedness leads to instability of LoRA when finetuning RoBERTa datasets SST-2 and MNLI; see more details in Apdx. D.4.
Integrating SAM with LoRA is a case with mutual benefits - LoRA reduces the additional memory requirement of SAM, while SAM not only overcomes the distributional shift in finetuning (Zhou et al., 2022), but also mitigates the possible inefficiency associated with LoRA's unbalancedness.
5.2 Balancedness-Aware Regularization (BAR)
However, directly applying SAM variants on LoRA exhibits two concerns: i) SAM doubles computa- tional cost due to the need of two gradients; and ii) additional efforts are required to integrate SAM with gradient accumulation and low-precision training (HuggingFace), which are common techniques for memory and runtime efficiency large-scale finetuning. Note that concern i) is annoying given the size of language models, especially in setups involving model parallelism.
regularization (BAR) a highly efficient approach to address both concerns, and it fixes the accuracy drop of LoRA relative to full-parameter finetuning, BAR is also the first efficient SAM variant derived from implicit regularization. The key observation for our algorithm design is that SAM's implicit regularization on balancedness can be achieved with an explicit regularizer 0x) x _ y y This regularizer originates from matrix sensing; see e.g., (Tu al., 2016; Ge et al., 2017). For OP, choosing a == (x y.)l/ [xx||² + ||y+ll²) recovers SAM's dynamic on B, up to an error of O(p2); ef. Lemma appendix. By ignoring this error, it can be seen that B3 decreases when ||x,|| 2 l|y=|l. Following this dynamic, we regulate balancedness based on whether ||x;I| 2 ||y.ll- The resultant approach termed as overparamterized BAR (oBAR) to reflect its source in OP.
Our balancedness-
On the other hand, because LoRA is NOP inherently, we take inspiration from Theorem 2 dropping the term Ar and mimicking dynamics of SAM. In particular, we regulate the objective with ar (x x y y) if |gx ||2 < ||gy,, 1| otherwise Or (y y - x x). The resultant approach is termed as nBAR. graphical illustration
be found in Fig. (b). It can be observed that nBAR shares similar performance as SAM on NOP. Both nBAR and oBAR can be implemented in the same manner as weight decay, and their detailed steps are summarized in Algs. and 3, respectively.
Another benefit of BAR, in additional to the lightweight computation, is that it can be applied individually each LoRA layer. As previously discussed (cf. Theorem 5), the number of layers has a negative impact on balancedness. By overcoming this "curse of multi-layer", BAR can induce better test performance over SAM.
Schedule of a,, In both nBAR and oBAR, one can employ a decreasing scheduler for a, for algorithmic flexibility. This is motivated by the fact that for both NOP and OP problems, the implicit regularization of SAM less powerful after sufficient balancedness or near optimal. Commonly adopted cosine and linear schedules work smoothly.

---

Table 1: Few shot learning RoBERTa (355M). i denotes results reported by (Malladi al., 2023)
RoBERTa
SST-2
SST-5
SNLI
MNLI
RTE
TREC
avg (1)
LoRA
20.1
52.3+2.9
84.3+0.3
18.1
27.
96.6=1.
80.0
LoRA-SAM
+0.2
54.2+2.0
80.5
78.7
80.6=1.3
96.7=0.2
81.3
LoRA-oBAR
91.5+
24.5
84.9 20.5
78.3+2.2
79.7=2.0
96.7.2055
80.9
LoRA-nBAR
91.4
55.0
84.9
78.1
81.0
96.721.0
81.2
Zero-Shot 79.0
35.5
50.2
48.8
51.4
32.0
49.5
Table 2: Runtime of BAR (normalized to LoRA, Ix) on OPT-1.3B. SAM relies on FP32 for stability. LoRA and BAR adopt FP16 training since this is the default choice for large models. nBAR and oBAR share similar runtime, hence reported together.
runtime (4)
SST-2
CB
RTE
COPA
ReCoRD
SQuAD
LoRA-SAM
4.43x
3.34x
4.10x
3.28x
4.35x
3.54x
LoRA-BAR
1.05x
1.03x
1.04x
1.05x
1.04x
1.03x
Numerical Experiments
To demonstrate the effectiveness of BAR, numerical experiments are conducted on various deep learning tasks using language models (LMs). Bold and underlined numbers used to highlight the best and second best performance, respectively, More experimental details can be found in Apdx. D. Code is available https: /g github....... com com/Bingcoonggi/BAA
6.1 Few-shot Learning with RoBERTa-large and OPT-1.3B
The first task to consider is few-shot learning with LoRA (Malladi et al., 2023), where the goal is to finetune language model with small training set. We follow the settings in (Malladi et al., 2023), and choose the backbones as RoBERTa-large. a masked LM with 355M parameters, and OPT-1.3B, an autoregressive LM (Liu et al., 2019; Zhang et al., 2022).
Results of the proposed oBAR and nBAR on RoBERTa-larg are summarized in Table 1. As indicated by the zero-shot performance, the distributional shift between finetuning and pretraining datasets is obvious. This is natural setting suitable for SAM and BAR. The averaged test accuracy is improved by 0.9 and 1.2 via oBAR and nBAR, respectively. The performance of nBAR close SAM. Moreover, BAR saves 74% additional runtime of SAM; see more details in Table 7 in the appendix.
The proposed nBAR and oBAR perform overhead of SAM by more than 95% because of its compatibility with FP16 training; see Table 2. Note that applying FP16 directly with SAM leads to underflow; see more in Apdx. D. This signifies the flexibility of BAR over SAM when scaling to large problems, as FP16 is the default choice for LMs. Prefix tuning (Li and Liang, 2021) is also included performance. We report FI score for SQuAD and accuracy for other datasets Table 3. The averaged improvement over LoRA 0.9 and 1.6 from oBAR and nBAR, respectively. both outperforming SAM. We conjecture that the performance gap between SAM and BAR comes from their different effectiveness regularizing balancedness. Balancedness of a particular layer is decreasing slower in SAM due to multiple layers, shown in Theorem 5, while BAR promotes balancedness faster as it can be applied individually on each LoRA layer. Comparing the absolute improvement for RoBERTa-large (355M) and DPT-1.3B. it is conjectured that BAR has more potential for larger models, and the verification is left for future due to hardware constraints.
better when scaling up OPT-1.3B. BAR reduces the
benchmark for comparisons
test
6.2 Finetuning with RoBERTa-large
Having demonstrated the power of BAR in few-shot learning, then apply it to finetune RoBERTa- large with LoRA. The results can be found Table It can be observed that nBAR and oBAR improve the performance of LoRA and prefix tuning (Li and Liang, 2021) on most of tested datasets.

---

Table 3: Performance of BAR for few shot learning using OPT-1.3B.
OPT-1.3B
SST-2
CB
RTE
COPA
ReCoRD
SQuAD
avg (t)
Prefix
92.9
71.643.0
65.2.2.6
73.0
69.7
82.1
75.8
LoRA
93.1
72.643.7
44.8
78.0
70.8
81.9
77.6
LoRA-SAM
93.5
7..3
70.6
78.0
70.9
83.0
78.4
LoRA-oBAR
93.6+0.6
75.64:4.5
70.4
78.0+0.0
70.9
82.5
78.5
LoRA-nBAR
93.7
79.8
70.5
78.0+0.0
71.0+1.0
82.3+1.8
79.2
Zero-Shot
53.6
39.3
53.1
75.0
70.2
27.2
53.1
Table Finetuning RoBERTa (355M) with BAR. Results marked with are taken from (Hu et al.,
2022), and those with * refer to Adapter" in (Hu et al., 2022).
RoBERTa
# para
STS-B
RTE
MRPC
CoLA
QQP
avg (f)
FT
355M
92.4
86.6
90.9
68.0
90.2
85.6
Adapter*
0.8M
91.9+10.4
80.1=2.9
89.7+112
67.8
91.7++0.2
84.2
LoRA
0.8M
92.4+0.1
88.2=0.6
89.6=005
6..8
91.4+0.1
85.3
LoRA-oBAR
0.8M
92.6+-0.1
88.7-02
90.3+
65.1+1100
91.6+0011
85.7
LoRA-nBAR
0.8M
92.6+0.22
1.3
90.3
65.6
91.6+0.1
85.9
Table 5: Finetuning GPT2 (345M) with BAR on WebNLG. Results prefix tuning and full full-parameter finetuning are obtained from (Hu et al., 2022).
GPT2
FT
Prefix
LoRA
LoRA-oBAR
LoRA-nBAR
# param BLEU (f)
354M
0.35M
0.35M
0.35M
0.35M
46.5
55.1
54.99:10.24
55.15=0.19
55.20+10.16
On average, oBAR leads to gain of 0.4, and nBAR raises the test performance by 0.6. BAR thereby fills the gap test performance between LoRA (0.8M) and full-parameter (355M) finetuning.
6.3 Text Generation on GPT2-medium
Lastly, we consider BAR a text-generation problem using GPT2-medium, model with 345M parameters. Results on WebNLG (Gardent et al., 2017) are reported in Table It can be seen that oBAR matches the performance of prefix tuning, while nBAR achieves the best BLEU score.
7 Discussions
This work provides theoretical and empirical evidence on the implicit regularization of SAM for both scale-invariant NOP and OP problems. Balancedness, an alternative to commonly adopted sharpness, is employed as the metric to capture global and data-responsive behaviors SAM. We find that i) SAM promotes variables to have (relatively) balanced norms; and ii) noisy data have stronger impact on balancedness. Lastly, we explicify the implicit regularization as a data-driven regularizer to foster the design of computationally efficient SAM variant, termed BAR. The effectiveness of BAR demonstrated using various tasks on RoBERTa-large, GPT2 and OPT. BAR saves 95% overhead of SAM and enhances the accuracy of LoRA to the level of full-parameter finetuning.
Limitation and Future directions. Our approach, BAR, is best applied on ocale-invariam modules in neural networks. Finetuning language models with LoRA, as popular option in practice, is setting naturally suitable for our approach. However, our approach does not apply for linear models, e.g., logistic regression. Regarding future directions, interesting one is whether SAM has other forms of implicit regularization beyond balancedness and sharpness. The exploration of other icale-invariant architectures beyond LoRA, e.g., the softmax function in attention, is also deferred to future work.
10

---

Acknowledgements
We thank anonymous reviewers for their suggestions. BL is supported by Swiss National Science Foundation (SNSF) Project Funding No. 200021-207343. LZ gratefully acknowledges funding by the Max Planck ETH Center for Learning Systems (CLS), NH is supported by ETH research grant funded through ETH Zurich Foundations and SNSF Project Funding No. 200021-207343.
References
Momin Abbas, Quan Xiao, Lisha Chen, Pin- Chen, and Tianyi Chen. Sharp-MAML: Sharpness- aware model-agnostic meta learning. In Proc. Int. Conf. Machine Learning, pages 10-32. PMLR, 2022.
Atish Agarwala and Yann Dauphin. SAM operates far from home: cigenvalue regularization as a dynamical phenomenon. In Proc. Int. Conf. Machine Learning, pages 152-168. PMLR, 2023.
Kwangjun Ahn, Sébastien Bubeck, Sinho Chewi, Yin Tat Lee, Felipe Suarez, and Yi Zhang. Learning threshold neurons via edge of stability. Proc. Ady. Neural Info. Processing Systems, volume 36, 2023.
Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-awary mini- mization. In Proc. Int. Conf. Machine Learning, pages 639-668. PMLR, 2022.
Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameterizztion. In Proc. Int. Conf. Machine Learning, pages PMLR, 2018.
Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks. In Proc. Int. Conf. Learning Represention, 2019a.
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. In Proc. Adv. Neural Info. Processing Systems, volume 32, 2019b.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In Proc. Int. Conf. Machine Learning, pages 322-332. PMLR, 2019c.
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent the edge of stability in deep learning. In Proc. Int. Conf. Machine Learning, pages 948-1024. PMLR, 2022.
Dara Bahri, Hossein Mobahi, and Yi Tay. Sharpness-aware minimization improves language model generalization. In Proc. Conf. Assoc. Comput. Linguist. Meet., pages 7360-7371, 2022.
David Barrett and Benoit Dherin. Implicit gradient regularization. In Proc. Int. Conf. Learning Represention, 2021.
Peter Bartlett, Dave Helmbold, and Philip Long. Gradient descent with identity initialization efficiently learns positive definite lincar transformations by deep residual networks. In Proc. Int. Conf. Machine Learning, pages 521-530. PMLR, 2018
Peter Bartlett, Philip Long, and Olivier Bousquet. The dynamics of sharpness-aware minimization: Bouncing across ravines and drifting towards wide minima. J. Mach. Learn. Res., 24(316):1 2023.
Léon Bottou, Frank Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM Review, 60(2):223-311, 2018.
Samuel Bowman, Gabor Angeli, Christopher Potts, and Christopher D Manning. A large annotated corpus for learning natural language inference. In Proc. Conf. Empir. Methods Nat. Lang. Process., pages 632-642, 2015.
Daniel Cer, Mona Diab, Eneko Agirre, Iñigo Lopez-Gazpio, and Lucia Specia. SemEval-2017 task I: Semantic textual similarity and cross-lingual focused evaluation. In Proc. Workshop Semant. Eval., pages 1-14. ACL, 2017.

---

Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-SGD: Biasing gradient descent into wide valleys. In Proc. Int. Conf. Learning Represention, 2017.
Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform ResNets without pre-training or strong data augmentations. In Proc. Int. Conf. Learning Represention. 2022.
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Long- LoRA: Efficient fine-tuning of long-context large language models. In Proc. Int. Conf. Learning Represention, 2024.
Zixiang Chen, Junkai Zhang, Yiwen Kou, Xiangning Chen, Cho-Jui Hsieh, and Quanquan Gu. Why does sharpness-awar minimization generalize better than SGD? In Proc. Adv. Neural Info. Processing Systems, volume 36, 2023.
Yan Dai, Kwangjun Ahn, and Suvrit Sra. The crucial role normalization in sharpness-aware minimization. In Proc. Adv. Neural Info. Processing Systems, volume 36, 2023.
Marie-Catherine De Marneffe, Mandy Simons, and Judith Tonhauser. The CommitmentBank: Investigating projection in naturally occurring discourse. Proc. Sinn und Bedeutung, 23(2):107- 124, 2019.
Christopher De Sa, Christopher Re, and Kunle Olukotun. Global convergence of stochastic gradient descent for some non-convex matrix problems. In Proc. Int. Conf. Machine Learning, pages 2332-2341. PMLR, 2015.
Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. QLoRA: Efficient finetuning of quantized LLMs. In Proc. Adv. Neural Info. Processing Systems, volume 36, 2023.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In Proc. Int. Conf. Machine Learning, pages 1019-1028. PMLR, 2017.
Bill Dolan and Chris Brockett. Automatically constructing a corpus of sentential paraphrases. In Proc. Int. Workshop Paraphrasing, 2005.
Jiawei Du, Hanshu Yan, Jiashi Feng, Joey Tianyi Zhou, Liangli Zhen, Rick Siow Mong Goh, and Vincent Y. F. Tan. Efficient sharpness-aware minimization for improved training of neural networks. In Proc. Int. Conf. Learning Represention, 2022a.
Jiawei Du, Daquan Zhou, Jiashi Feng, Vincent Y. F. Tan, and Jocy Tianyi Zhou. Sharpness- aware training for free. In Proc. Adv. Neural Info. Processing Systems, 2022b.
Simon S Du, Wei Hu, and Jason D Lee. Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced. In Proc. Adv. Neural Info. Processing Systems, volume 31, 2018.
Gintare Karolina Dziugaite and Daniel M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Proc. Conf. Uncerntainty Artif. Intel., 2017.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In Proc. Int. Conf. Learning Represention, 2021.
Claire Gardent, Anastasia Shimorina, Shashi Narayan, and Laura Perez-Beltrachini. The WebNLG challenge: Generating text from RDF data. In Proc. Int. Conf. Nat. Lang. Gener,, pages 124-133. ACL, 2017.
Rong Ge, Chi Jin, and Yi Zheng. No spurious local minima in nonconvex low rank problems: A unified geometric analysis. In Proc. Int. Conf. Machine Learning, pages 1233-1242. PMLR, 2017.
Gauthier Gidel, Francis Bach, and Simon Lacoste-Julien. Implicit regularization of discrete gradient dynamics in linear neural networks. In Proc. Adv. Neural Info. Processing Systems, volume 32, 2019.
12

---

Antoine Gonon, Nicolas Brisebarre, Elisa Riccietti, and Rémi Gribonval. A path nath toolkit for modern networks: consequences, promises and challenges. In Proc. Int. Conf. Learning Represention, 2024.
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan. and Sylvain Gelly. Parameter-efficien transfer learning for NLP. In Proc. Int. Conf. Machine Learning, pages 2790-2799. PMLR, 2019.
Edward 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. In Proc. Int. Conf. Learning
Represention, 2022.
HuggingFace. Gradient accumulation. URL URI.attsss//uugggggaaae http:////.....o////////c/eetttta
usage_guu Llagge_guies//// seriides/graieeee aradmaat..ccmmmlltiin accumulation.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry P. Vetrov, and Andrew Gordon Wilson Averaging weights leads to wider optima and better generalization. In Proc. Conf. Uncerntainty in Artif. Intel., pages 876-885, 2018
Stanislaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in SGD. 171.42171,.0463,, 2017.
Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. In Proc. Int. Conf. Learning Represention, 2019.
Weisen Jiang, Hansi Yang. Yu Zhang, and James Kwok. An adaptive policy to employ sharpness- aware minimization. In Proc. Int. Conf. Learning Represention, 2023.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. In Proc. Ini. Conf. Learning Represention, 2020.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In Proc. Int. Conf. Learning Represention, 2016.
Minyoung Kim, Da Li, Shell Xu Hu, and Timothy M. Hospedales. Fisher SAM: Information geometry and sharpness aware minimisation. In Proc. Int. Conf. Machine Learning, pages 11 11148-11161, 2022.
Dawid Jan Kopiezko, Tijmen Blankevoort, and Yuki M Asano. VeRA: Vector-based random matrix adaptation. In Proc. Int. Conf. Learning Represention, 2024.
Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. ASAM: Adaptive sharpness- aware minimization for scale-invariant learning of deep neural networks. Proc. Int. Conf. Machine Learning, pages 5905-5914, PMLR, 2021.
Bingcong Li and Georgios B Giannakis. Enhancing sharpness-aware optimization through variance suppression. In Proc. Adv. Neural Info. Processing Systems, volume 36, 2023.
Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. In Proc. Conf. Assoc. Comput. Linguist. Meet., pages 4582-4597, 2021.
Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss? A mathematical framework. In Proc. Inf. Conf. Learning Represention, 2022.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. RoBERTa: A robustly optimized BERT pretraining approach. arXiv preprint arXiv: 1907. 1192, 2019.
Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, and Yang You. Towards efficient and scalable sharpness-aware minimization. In Proc. Conf. Computer Vision and Pattern Recognition, pages 12350-12360, 2022.
Iyya Loshchilov and Frank Hutter. Decoupled weight decay regularization. Proc. Int. Conf. Learning Represention, 2019.
13

---

Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In Proc. Int. Conf. Learning Represention, 2020.
Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Alex Damian, Jason D. Lee, Danqi Chen, and Sanjeev Arora. Fine-tuning language models with just forward passes. In Proc. Adv. Neural Info. Processing Systems, volume 36, 2023.
Peng Mi, Li Shen, Tianhe Ren, Yiyi Zhou, Xiaoshuai Sun, Rongrong Ji, and Dacheng Tao. Make sharpness-aware minimization stronger: A sparsified perturbation approach. In Proc. Adv. Neural Info. Processing Systems, volume 35, 2022.
Yurii Nesterov. Introductory lectures convex optimization: A basic course, volume 87. Springer Science & Business Media, 2004.
Behnam Neyshabur, Russ R Salakhutdinov, and Nati Srebro. Path-SGD: Path-normalized optimization in deep neural networks. In Proc. Adv. Neural Info. Processing Systems, volume 28, 2015.
Behnam Neyshabur, Srinadh Bhojanapalli, David Mcallester, Nathan Srebro, and Nati Srebro. Ex- ploring generalization in deep learning. In Proc. Ady. Neural Info. Processing Systems, volume 30, pages 5947-5956, 2017.
Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro. A PAC-bayesian approach to ppectrallynnormalize margin bounds for neural networks. In Proc. Int. Conf. Learning Rep- resention, 2018
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ questions for machine comprehension of text. In Proc. Conf. Empir. Methods Nat. Lang. Process., pages 2383-2392, 2016.
Pranav Rajpurkar, Robin Jia, and Percy Liang. Know what you don't know: Unanswerable questions for SQuAD. In Proc. Conf. Assoc. Comput. Linguist. Meet., pages 784-789, 2018.
Melissa Roemmele, Cosmin Adrian Bejan, and Andrew S Gordon. Choice of plausible alternatives: An evaluation of commonsense causal reasoning. In AAAI Spring Symposium Series, 2011.
Heejune Sheen, Siyu Chen, Tianhao Wang, and Harrison Zhou. Implicit regularization of gradient flow on one-layer softmax attention. arXiv preprint 2xXi:22403.0869, 2024.
Tom Sherborne, Naomi Saphra, Pradeep Dasigi, and Hao Peng. TRAM: Bridging trust regions and sharpness aware minimization. In Proc. Int. Conf. Learning Represention, 2023.
Dongkuk Si and Chulhee Yun. Practical sharpness-aware minimization cannot converge all the way to optima. In Proc. Adv. Neural Info. Processing Systems, volume 36, 2023.
Sidak Pal Singh and Thomas Hofmann. Closed form of the hessian spectrum for some neural networks. In High-dimensional Learning Dynamics 2024 The Emergence of Structure and Reasoning, 2024.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Y Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proc. Conf. Empir. Methods Nat. Lang, Process., pages 1631-16422 2013.
Behrooz Tahmasebi, Ashkan Soleymani, Dara Bahri, Stefanie Jegelka, and Patrick Jaillet. A universal class of sharpness-awann minimization algorithms. arXiv preprint arXiv:2406.03683 2024.
Stephen Tu, Ross Boczar, Max Simchowitz, Mahdi Soltanolkotabi, and Ben Recht. Low-rank solutions of linear matrix equations via procrustes flow. In Proc. Int. Conf. Machine Learning, pages 964 973. PMLR, 2016.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proc. Adv. Neural Info. Processing Systems, volume 30, 2017.
Ellen Voorhees and Dawn M Tice. Building a question answering test collection. In Proc. Annu. Int. ACM SIGIR Conf. Res. Dev. Inf. Retr., pages 200-207, 2000.
14

---

Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. SuperGLUE: A stickier benchmark for ceneral-purposs language understanding systems. In Proc. Adv. Neural Info. Processing Systems, volume 32, 2019a.
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. In Proc. Int. Conf. Learning Represention, 2019b.
Pengfei Wang. Zhaoxiang Zhang, Zhen Lei, and Lei Zhang. Sharpness- aware gradient matching for domain generalization. In Proc. Conf. Computer Vision and Pattern Recognition, pages 3769-3778,
2023.
Ziqiao Wang and Yongyi Mao. On the generalization of models trained with SGD: Information- theoretic bounds and implications. In Proc. Int. Conf. Learning Represention, 2022.
Alex Warstadt, Amanpreet Singh, and Samuel R Bowman. Neural network acceptability judgments. Trans. Assoc. Comput. Linguist., 7:625 643566112222 2019.
Kaiyue Wen, Tengyu Ma, and Z hiyuan Li. How does sharpness sharpness-aware aware minimization minimizes sharpness. In Proc. Int. Conf. Learning Represention, 2023a.
Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. Sharpness minimization algorithms do not only mini- mize sharpness to achieve better generalization. In Proc. Adv. Neural Info. Processing Systems, volume 36, 2023b.
Adina Williams, Nikita Nangia, and Samuel R Bowman, A broad-coverage challenge corpus for sentence understanding through inference. In Proc. Conf. North Am. Chapter Assoc. Comput. Linguist.. pages 11 111221122 112,, 2018
Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Proc. Annual Conf. Learning Theory, pages 3635-3673. PMLR, 2020.
Dongxian Wu, Shu-Tao Xia, and Yisen Wang. Adversarial weight perturbation helps robust general- ization. In Proc. Adv. Neural Info. Processing Systems, volume 33, pages 2958-22888 2558, 2020.
Wenhan Xia, Chengwei Qin, and Elad Hazan. Chain of LoRA: Efficient fine-tuning of language models via residual learning. arXiv preprint trXiv:2401.0415/ 2024.
Qingru Zhang, Minshuo Chen, Alexander Bukharin, Pengeheng He, Yu Cheng, Weizhu Chen, and Tuo Zhao. Adaptive budget allocation for parameter-efficient fine-tuning. In Proc. Int. Conf. Learning Represention, 2023a.
Ruipeng Zhang, Ziqing Fan, Jiangchao Yao, Ya Zhang, and Yanfeng Wang. Domain-inspired sharpness aware minimization under domain shifts. In Proc. Int. Conf. Learning Represention, 2023b.
Sheng Zhang, Xiaodong Liu, Jingjing Liu, Jianfeng Gao, Kevin Duh, and Benjamin Van Durme. ReCoRD: Bridging the gap between human and machine commonsense reading comprehension. arXiv preprini arXiv: 12885, 2018
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. OPT: Open pre-trained transformer language models. arXiv preprint
Ttt.:220500066. 2022.
Yang Zhao, Hao Zhang, and Xiuyuan Hu. Penalizing gradient norm for efficiently improving generalization in deep learning. In Proc. Int. Conf. Machine Learning, pages 26982-26992, 2022.
Wenxuan Zhou, Fangyu Liu, Huan Zhang, and Muhao Chen. Sharpness-aware minimization with dynamic reweighting In Proc. Conf. Empir. Methods Nat. Lang. Process., pages 5686-5699 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. In Proc. Int. Conf. Learning Represention, 2022.
15

---

Supplementary Document for
"Implicit Regularization of Sharpness-Awars Minimization for Scale- Ieali-Invariant Problems"
A
Missing Details
A.1 Broad Impact
The theories and approaches are applicable across various scenarios. The proposed algorithmic tool simplifies finetuning language models, improves performance of downstream tasks, and consumes less resource compared to SAM. For tasks such as sentiment classification, our approach facilitates real world systems such recommendation by improving accuracy. However, caution is advised when the downstream tasks of language models involve generation. For these tasks, users should thoroughly review generated content and consider to implement gating methods to ensure safety and trustworthiness.
A.2 More on Related Work
Sharpness and generalization. Sharpness observed to relate with generalization of SGD in deep learning (Keskar et al., 2016). It is found that sharpness varies with the ratio between learning rate and batchsize in SGD (Jastrzebski et al., 2017). Large scale experiments also indicate sharpness-bar based measures align with generalization in practical scenarios (Jiang et al., 2020; Chen et al., 2022). Theoretical understandings on generalization error using sharpness-related metrics can be found in e.g., (Dziugaite and Roy, 2017; Neyshabur et al., 2017; Wang and Mao, 2022). There is large body of literature exploring sharpness for improved generalization. Entropy SGD leverages local entropy in search of flat valley (Chaudhari al., 2017). A similar approach as SAM is also developed in (Wu et al., 2020) while putting more emphases on adversarial robustness. Stochastic weight averaging is proposed for finding flatter minima in (Izmailov et al., 2018). It is shown later in (Wen et al., 2023b) that the interplay between sharpness and generalization subtly depends on data distributions and model architectures, and there are unveiled reasons beyond sharpness for the benefit of SAM.
SAM variants. Although SAM is successful in various deep learning tasks, it can be improved further by leveraging local geometry in a fine-grained manner. For example, results in (Zhao et al., 2022; Barrett and Dherin, 2021) link SAM with gradient norm penalization. Zhuang et al. (2022) optimize sharpness gap and training loss jointly. A more accurate manner solve inner maximization in SAM is developed in (Li and Giannakis, 2023). SAM and its variants are also widely applied to domain generalization problems; see e.g., (Zhang et al., 2023b; Wang et al., 2023).
Other perspectives for SAM. The convergence of SAM is comprehensively studied in (Si and Yun, 2023). Agarwala and Dauphin (2023) focus the
behavior of unnormalized SAM on quadratic problems. Dai et al. (2023) argue that the normalization in SAM, i.c., line of Alg. 1, is critical. Sharpness measure is generalized to any functions of Hessian in (Tahmasebi et al., 2024). However, even the generalized sharpness cannot provide implicit regularization for simple functions such as h(r, y) zy, because the Hessian is the same for all (æ, v). In addition, when Hessian is negative definite, some of the generalized sharpness measures (e.g., determinate of Hessian) may not be necessarily meaningful.
Implicit regularization. The regularization effect can come from optimization algorithms rather than directly from the regularizer in objective functions. This type of the behavior is termed as implicit regularization or implicit bias of the optimizer. The implicit regularization of (S)GD is studied from multiple perspectives, such margin (Ji and Telgarsky, 2019; Lyu and Li, 2020), kernel (Arora et al., 2019c), and Hessian (Li et al., 2022; Arora et al., 2022). Initialization can also determine the implicit regularization (Woodworth et al., 2020). Most of these works explore the overparametrization regime.
LoRA and parfmeter-efficient finetuning. LoRA (Hu al., 2022), our major numerical benchmark, is an instance of parameter-efficient finetuning (PEFT) approaches. PEFT reduces the resource requirement for large language models on various downstream tasks, at the cost of possible accuracy drops on test performance. The latter, together with the transfer learning setup jointly motivate the adoption of SAM. Other commonly adopted PEFT methods include, e.g., adapters (Houlsby et al., 2019) and prefix tuning (Li and Liang, 2021). There are also various efforts to further improve
16

---

LoRA via adaptivity (Zhang et al., 2023a), chaining (Xia et al., 2024), aggressive parameter saving (Kopiczko al., 2024), low-bit training (Dettmers et al., 2023), and modifications for Long-sequenn (Chen et al., 2024). Most of these efforts are orthogonal to BAR proposed in this work.
A.3 Additional Applications of Scale-Invariant Problems in Deep Learning
Attention in transformers. Attention is one of the backbones of modern neural networks (Vaswani et al., 2017). Given the input D, attention can be written as
min softmax DQK DJ (DQKTT)DD Q.K.V
where {Q, K. }} are query, key, and value matrices to be optimized. This is scale-invariant problem because scaling {Q, K) does not modify the objective function. Considering the number of variables, the optimization of {Q, K is considered as OP.
Two-layer linear neural networks. This problem is simplified version of two-layer ReLU neural nets, and its objective can be defined as
1(WI, W2) - (aa [W, W2a-bl13
(10)
This is usually adopted as an example for overparametrrzation. and can be extended to deeper linear neural networks; see e.g., (Arora et al., 2019a). Moreover, it is known that the optimization for such problem is quite challenging, and GD can fail to converge if W1 and W₂ 2 are not initialized with balancedness (Arora et al., 2019a). An extension of (10) is two-layer ReLU networks, which are widely adopted in theoretical frameworks to understand the behavior of neural networks. ReLU networks are scale-invariant, but only when the scaling factor is positive.
Other examples. For ResNets, two-variable scale-invariant submodules also include affine Batch- Norm and the subsequent convolutional layer. For transformers, scale-invariant submodules besides attention include LayerNorm and its subsequent linear layer.
A.4 SAM Pays More Attention to Difficult Examples
Testing example for NOP. The problem presented below adopted Fig. I (a) and Fig. 2 for visualization of SAM's behavior on NOP. We consider a special case of problem (la), where the goal fit (rank-1) matrices by minimizing
fn yx Ee [|xx] (A oN
(11)
where A 3x33 diag (0.5, 0, 0] and N₂ E 3x33 denote the ground truth and Gaussian noise, respectively; and a controls the SNR. Here we choose N₂- == diag 1.0, 0.8, 0.5) where entries of Ug unit Gaussian random variables.
In our simulation of Fig. (a), we set the step size to be 10 -4 and the total number of iterations as T' 10t for both SGD and SAM. Parameter p chosen as 0.1 for SAM. For both algorithms, initialization is X = [0.2, -0.1, 0.3) and Yo = -3xg- Note that we choose a small step size to mimic the settings of our theorems.
Testing example for OP. The problem presented below is adopted in Fig. I (b) for visualization of SAM on OP. A special case of problem (1b) is considered with objective function
y) EEg [[|x* y (a + [m]
(12)
where a ER and ng IR denote the ground truth and Gaussian noise, respectively. We choose
0.5 and ne as a unit Gaussian random variable. Here, controls the SNR of this problem.
In our simulation of Fig. 1 (b), we set 10-4 and T = 10s for both SGD and SAM. Parameter p is set as 0.2 for SAM. For both algorithms, initialization is xo [0.2, -0.1, 0.3] and yo = -3x--
A.5 Scale-Invariance in OP
Scale-invariance also bothers OP in the same fashion as it burdens NOP. For completeness, the scale-invariance of OP can be verified by
fo(x y) = fo (oxx
T (--)) Va + O.
(13)
a
17

---

An optimizer has to determine a for OP despite it does not influence objective value. Hence, scaling redundant for OP.
Similar to NOP, the (stochastic) gradient of OP is not scale-invariant, In particular, given a minibatch of data M, the stochastic gradient for OP (1b) can be written as
Ex
1
Σ EEM
(
gy
Σ (x y))x
(14)
[
ÊEM
Consequently, being balance also brings optimization benefits for OP as discussed previously in Section 2.2.
A.6 BAR in Detail
BAR is inspired jointly from the balancedness-promoting regularizer |||x|| ||y]|l and the dynamics of SAM both NOP and OP. The implementation of BAR is similar as weight decay in AdamW (Loshchilov and Hutter, 2019).
3.6
3.2
2.6
Here we use nBAR as an example. If ignoring A in Theo- rem 2, it can be seen that B, for NOP decreases whenever
0.02
2.4
0.00
2.0
|gx. < lggy.l ||- other words, the balancedness of SAM driven by the difference between the gradient norms at X₁ and yt. nBAR mimics this and triggers balancedness when stochastic gradients gx, and gy, are not balanced; see Alg. 2.
0.02
1.6
0.04
0.8
0.4
1..005-.0000000255.000 0.025 0.050 0.075
Finally, we illustrate more the reasons for employing regularization in OP rather than posing ||x:|| = lyy|| as hard constraint or initializing in balanced manner, i.e., |xxol = I First, it is quite clear that ||x|| = lyl nonconvex set and how project such set is still debatable. Second, the 'symmetry' associated with the scale-invariant problems does always favor this constraint. For the purpose of graphical illustration, we consider 2-dimensional example f(x, y) 30000(ay 0.005)2. It is quite clear that the objective is symmetric regarding the line -y- which satisfies |a]| ly|: see Fig. 4. However, it is not hard to see that SGD can never leave x = yy once it reaches this line via a hard constraint initialized on this line. In other words, directly adding |x||| = ly|| as constraint can trap the algorithm at saddle points. This symmetric pattern is even more complicated high dimension, i.e., symmetry over multiple lines hyperplanes. Hence, one should be extremely careful about this hard constraint, and regularization is safer and more practical choice.
Figure 4: The value of f(x,)) Once SGD reaches the dotted line, i.c. the hard constraint |æ| = ly/. it can only converge to saddle point (0,0).
B Missing Proofs for NOP
B.1 Proof of Theorem 1
Proof.
For notational convenience, we let G; (x(.y))) Then, we have that
d d||x;|1?
dx:
dt 2x,
2x, Ex. = -2x, Geyt.
dl
Similarly, we have that
y 112
dy:
2y
2y, gy. = -2y G. Xx.
dt Combining these two inequalities, arrive at x 2
di
dI
0.
dt
dt
The proof is thus completed.
18

---

B.2
Extension to Stochastic Normalized Gradient Descent (SNGD)
Next, we extend Theorem SNGD, whose updates can be written as
gx
gy: 8xx
X++1 X
Yt-1 yt
(15)
8x
2
By, 1
By
Theorem 4. When applying SNGD (15) on NOP problem (la), the limiting flow with + guarantees that ||x,]] 11 |yy
0
|
yy 1 for all >> O. In other words, at
dB.
holds.
Proof. For notational convenience, we let G2 == vf((x())
Then, we have that
d|x dt 2x,
dx:
x gx,
X G.y,
-2
=-2
dt Vllxx.1 2 1ly 11
8x, 1 + ggy 12
Similarly, we have that
d y 12
2y, dy:
Yz gy:
yr G. X
= 2 v /llxx Ex. 1
-
dt
dt
+ By:
v gxx 1| + gy. 2
Combining these two inequalities, we arrive at
x|x+11
dllyell?
dt
dt
The proof is thus completed.
B.3 Proof of Theorem
Proof. Denote G, Vi((xyy)) and G. = ((((x) for notational convenience. Following
SAM updates in (4) and setting - O, we have that
dx, dt
dy:
Gy pu, x.), dt
-x, pid Gey:)-
This gives that d( [|x+,|
11y,113)
pud
y, G. G.y. x G.G.
(16a)
dt
pue ||xx, 1 |l 113
G.) + pup y, (G, - Gi) Ex, xx (G: G.)
(16b)
The second term in (16b) is A Theorem 2. Next, we give upper bound on [Acl. Using Assumption 1. we have that
| G.ll S LI x,y.
Ellpuu(x18j) gx,
puu
Bx.gy.
Lp |x1gy, + gx
||gx.gy
Lp²
Vlxx,1 1 + By. 1
|gx
| +
By.
2
_ + [|yell) + Lp² = O(Lp)
2
where (a) uses the definition of us (b) follows from ||ab = ||a||||||| and the finite convergence assumption. To bound At, we also have
ly (G; G.) T gx.
ly. TG. G.)
puely? IyT (G. G.) gx,
<
gx.
V |xx +
gy, 1
[xx
S pll S| Gelllyyll O(Lp²)
(17)
where the last line also uses the finite convergence. We can bound pu, xx_ (G, G.)gy
0(p2L)
in similar manner. Combining (17) with (16b) gives the bound on [Ap O(p²L).
0
19

---

B.4 Proof of Corollary 1
Here, we prove the formal version of Corollary 1.
Corollary 2. Suppose that llgx, || > 0 and |gy ll > 0 and + 0, then there exists Be such that the
magnitude of Be shrinks whenever [B:] > B,-
Proof. Without loss of generality, suppose that B, O, i.e., ||x||| > ||yll| 0. Let x, and y, be the scaled version of and y, such that ||x: |r and x x.y, satisfied. This suggests that X, 0r and y, = V1/an, where or v x |l/ ltpll/llyyl /llll Next, we show that whenever B, is large enough, we have that
dB:
8x, 1 By, 1|
+ L(p2t) < 0.
(18)
dt
gx 2 By, 1
Since + O, we only need to show that for some small = (pL.) 2 0,
gx. 2
gy: 1|
-E.
(19)
8x,
By 1|
By the definition of gx, Sy, and X,, Ve. we have that (19) can be rewritten as
af || G. x 112 |/.,y.1//a
(20)
o G, x, 22 G G.
Note that the function h(2) = (az - b/2) / a2 + b/2 is monotonically increasing in 2 when
a, b > and 2 > 0 as (²()) = (a²2 + 6ab/z / + 1/2==)/(2(az= + b/2)3/2) > 0. This implies that
hx2)) 0 when > Vb/a, and thus the condition in (20) can be satisfied for c = O(L)) - 0 when
af > a², where a² = || x.|-- This condition on ar is equivalent to
B: 2 (||x,|1 - 11y+11
(Ila,x; - 1////,,11
2 (I@xx112 ||5./a/13)-
Combining everything together, we have that dB: < 0 if
1 2
B,
(|ax
(21)
The proof is thus completed. We also note that in the case of p 0. the same condition (21) can be derived by obtaining the inverse function of f(z)) evaluated at O(pL), and the corresponding ap and By can be defindd similarly.
B.5 Extension to LoRA (layer-wise NOP problem)
Let E (1,2...., D} be the layer index. Denote fc as the loss function on minibatch Mr. To simplify the notation, we also let G.,, == Xx..yI, Fe yehl), G.. := V VT se (66,,,,, Fr.h).
and u, 1/ V/VV2 (llgx., 2 ||gv., 112). The update of SAM for layer can be written
x Xx,1 + pu Gr,
Yt, Yt,l pp GT,x,
(22a)
BR. Gr.rt..s
By.. G
(22b)
X(+1J Xt,! 1/x...
Y++1.1 = Yt.1 ngs....
(22c)
Refined assumption for LoRA. Direct translating Assumption to our multi-layer setting gives
[ h) - V h < L
Σ
x - a,bT
(23)
1=1
20

---

However, the above assumption is loose, and our proof only needs block-wits wise smoothness, i.c.,
[Vrf((xxy)) Vifi(a/b)))l 12 arb; 11 14
(24)
where Vi refers to the gradient on x,y;" It can be seen that VDL 2 L. but one can assume that VDL for intuitive understandings.
Theorem 5. Suppose that block smoothness assumption (24) holds. Consider the limiting flow of SAM in (22) with + O and a sufficiently small Let Bu = (1 lIve,ll?) and
B, = Ei=l Br,, For some Arl = O L), SAM guarantees that
dB, Σ 1= Ex,,e 11 EE= By 11?
dt =
+ Ar-
(25)
V |gx., 112 +E 112...12
Furthermore, for per layer balancedness satisfies that for some |A,,1 = O()21))
dB dt
8xx,1 112 - Byr,, 1|
p
Ar,i-
(26)
D 1= |xx,, 112
E1=1 |y.,11?
Understanding Theorem 5. Ar, , and A₁ are the same order because of the possible unbalancedness among gradient norms for different layers. Comparing per layer balancedness B,, with Theorem 2, it can be roughly estimate that the regularization power is O(vD) times smaller in Be,-- This estimation comes from İ =z L/VD, and the first term also O(VD) smaller than the same term in Theorem 2. other words, the regularization on balancedness can be reduced by (vVD) times in LoRA in the worst case, and the worst case comes from gradient unbalancedness among layers.
Proof. Following (22) and setting 7) O, have that
dx,,1 di
dy
Gr.( + pur G.
( + pr Gr.)).)..
di
This gives that
dB.,
= put
y,, GT
G.G.G.G.YY... Gr G.G.G..
(27a)
dt
=
pu
1|2x,,12 ||sr,,1?? 112 + pue yT
8x,, - T
XT, - G.,) (..))
=
(27b)
Proof for (25). Let Ar = E, Att-- start with, we have that
|| - Gr.ill < iIIR,.NTTI x..
(x + Bx..V + &uuggBr, T..
21

---

Next, based on finite convergence assumption, we have that
D
VT. -G.,) 8x.a
(28)
t=1
D
puell puulGGr. -Gual -axxllllxx..ll [lBx.il)
(a)
0
L + Ex..ellll (IEx..ll)
(
D
Σ
0
[1 + (l |18x(ee]
1=1
Σ ||
Σ 13
Ei=1 gx..
@xr.ill Il
= -L' o
1i=1(1|8x |xx, 1 ||gy..: Σ |x.. ll.v.. 112)
where in (a) we use the fact that p is chosen small; (b) uses finite convergence assumption and lab T |a|| |b]l. Using similar arguments, we can bound A = O f).
Proof for (26). Next, we give upper bound on |Ar,y Using similar argument as (28), we have that pue yT,( VV,,GG.... G.., G..)) 8x(,1
(29)
<0 p i( (1l8yaa + |18x, [|
|Bx,,1 1|2
8xx,1
By
= -p-L.
(30)
1= 8xx,1
1 |18y..! (12)
|xx,1
1|2
Using (29), we have that
1gx. + lgy 11
11gx..1
|1llBy... ||
Ard S p.... 0
+
1 (ll 112 ||gy.,
D
(11gx,,1
11
= = L).
The proof is is thus completed.
C Missing Proofs for OP
C.1 Unbalancedness of SGD in OP
Theorem 6. Applied SGD or SNGD on problem (1b), both of them ensure that ||x, lly.
[x
|| for all 0. In other words, B, keeps unchanged.
Proof. We consider SGD and NSGD separately.
SGD. It is straightforward to see that
d Xx 2
d|ly,|12
(x(
y()x
dt
This completes the proof of SGD.
NSGD. The gradient update of NSGD
dx; dt
gx.
dy:
gy
(31)
gx,
2 +
Byy,
dt
gx,
By.
Then we have that for NSGD,
d d||x]||
x, x,y²
dl
(xx
dt
V [x.
+
1y.
1
df
This gives the result for SNGD.
22

---

C.2 Proof of Theorem 3
To prove this theorem, we first focus on the dynamic of SAM.
Lemma 2. Suppose that Assumption I holds. Consider the limiting flow of SAM in (7) with - 0.
Let B == 2 (1|x|1 llyel??) and p small. Then, for some |A (pp [(B.). SAM guarantees
dB. 22
(x yr)
B. At-
(32)
dt
Xx 2 + 1yyll
Proof. For notational convenience, write ft f((x] y.) and ft = [[(8 81). Using similar arguments as Theorem 2, we have that
pu (11x,12-11xxxl)
- l1xel1?)
(33)
2 dt
||x+1 (1|x,|1 11)
V 1x11 11?
(1xx.11 - -1yx.113)
|x,11 112 + 11y.112
sgn sgn( (f))
(11x,11 1y+e1
|x,11
|yxell 1|?
Next bound [Ael- To start with, we have that
& y:
x, yel = 1 u/xx
x
Ye + pue /uuel|x/11 pree ppuelly.+13))
(34)
x Yr
p²
|| PV
X
11
0²
$ +
PV
X 12 12
2
Using Assumption 1 and (34), we arrive at
ft f s T
yr
x7
xTv;
Olpl. / X 2 + 1yy 12).
(35)
Hence, we arrive at The proof is thus completed.
|
||
AN S plft
O(p2LLBB1).
Vxx
2 ly: 2
Next, the proof of Theorem is provided.
Proof. Lemma 2 has already indicated the concentration of B, towards 0, if the magnitude of the first
term is larger than [Ae]. To see this, notice that lower bound // /|xx
lyyl by
Xx 2
1
(I|x + Ilx Iyx 2 |||xx||
||xel
C.-
(36)
X + 1
12 |
Hence, long as (x, y()l C, > O(p²L|B,]) we have the first term dominating the dynamic of SAM, leading contraction of B,, This completes the proof the first part.
Next we prove the second part, which is the lower- and upper- bound on B, The lower bound can be seen from (36). For the upper bound, we have
| 2| - ||y:|2
||x; 1|2 |y
V
(37)
| 2 y
/||x
Plugging (37) into (33) finishes the proof.
23

---

C.3 m-sharpness for OP
m-sharpness is a variant of SAM that is empirically observed to improve generalization, and is especially useful for distributed training on multiple GPUs (Foret al., 2021). However, the reason behind the improved performance is not fully understood. (Andriushchenko and Flammarion, 2022) show that m-sharppess is more pparse-promotinn for diagonal linear neural networks minimized via a quadratic loss. However, diagonal linear networks are not scale-invariant.
For consistent notation with (7), use fe(-) to denote the loss function minibatch Mr. In m-sharpness, the minibatch Mr is divided into m disjoint subsets. Without loss of generality, we also assume that the minibatch evenly divided. We denote the loss function on each subset as Se. = Jt. With these definitions, the update of fr. i E (1, 2, mj. Note that we have 1==1 m-sharpness can be written as
X1, X1 pu
Yt, Yt + p
(38a)
8x.
By. (Z,
(38h)
1
Y4+1 Yv 1
X++1 X
Byn
(38c)
m
where de := ff, (x, y:))
x 12 yvll Comparing with the SAM update for OP in (7),
the difference is that perturbed gradient calculated on each fe. Next, we analyze the dynamic of SAM with m-sharpness.
Lemma 3. Suppose that Assumption holds. Consider the limiting flow of SAM in (38) with rj - O.
Let B, == 1 (||x,|1 - lly.ll?) and be small. Then. for some |Ael = O(p² L). SAM guarantees that
-==1 (x yz) dB. 2 >
B, + Ac. 1
(39)
di
m
Proof. For notational convenience, we write fiy == (x, y:) and fi. := f1, Then, we have that
dt [|x,|| -1lve11
11v:113)
P
Ft. i 21xx|12-1lxxl)) -1xxa11?)
(40)
sgn sgn(Fr,)..
P
(1120112 -1x8112-11ve113))
/|xx| 2 + 11y.112
1=1
P
1=1
-1xa111)
m
+ Ty
sgn
(11x.112 - 11
:=
Next, using (34) and Assumption 1, we have
IL, Pe.l
S L J
x
/ 1 2 lly.13)) 12).
Hence, can bound |Atr,l
|
2
Iy.
Ar,
$ 1f
F.
OpL(BB))
2
|
The proof is thus completed by plugging A into (40).
24

---

C.4 Extension to Layer-wise OP
We start with the notation. Let E 21,2...., D3 be the layer index. Denote fi as the loss on minibatch Mr- Let St. Vufi((x(, i.e., the 1-th entry of gradient (w.r.. the variable
x(,),,), fi,- == S..d) and == 1/v E 18 [1x,,11 + [[ The update of SAM for layer can be written as
x111 X1,1 purfiyyed
Yt,1 yt, putfi
(41a)
8R = f1,
By f1,
(41b)
X++11 X1,1 - 18x V++1, = Yt, 1/y
(41c)
Refined assumption for LoRA. Our proof only needs block wise smoothness, i.c., Vxfi y:) ViSe (a) b:) x, aj bil3; vl,
(42)
where Vi refers the gradient on x yı. can be seen that v vDL L. but one can assume that VDL = for more clear intuition.
Theorem 7. Suppose that block smoothness assumption in (42) holds. Consider the limiting flow of SAM in (41) with 7) + O and a sufficiently small p. Let Be,1 = (1xx,,111 lly(,ll²) and
Br max [Ba,) For some AAl = O0p² ),maax,, SAM guarantees that E D 1 1 (1xx2111 lyea
dB. =
+ A...
(43)
dt
EP= 15t, x 11? 2 + [[]]]
Furthermore, for some Arr ( 221 L(B,,(). per layer balancedness satisfies that
dBr.1 |5 (1 l1v+,11
Ar,s.
(44)
EP 15 xx [1
Proof. Using similar derivation before, have that
d/ | l|yeeel2 -wellfl (1xx,111 11),,.1?²) puffiilfi
Fl.)) - (1x5. (yall? AmA
Next, based on (42), we have that
IL1 f'al S Ż V.. x__,v.ol S plu..Sllll (1 + Ilye,ll??) Combining these two equations, and applying similar argument as Theorem 5, it is not difficult to arrive at |Ar, O(p² (lB,,)) and [Ael O(p² (pp imaa)..
+ + Pul
C.5 Proof of Lemma
Proof. Within W", the Hessian on (x, y) can be calculated as ""(xTy)lyT,xx") J.]. The largest eigenvalue is f"(w)((Ixx + l?)) By the AM-GM inequality, it can be seen that the largest eigenvalue is minimized when ||]|| ||y]], whose balancedness is
D Missing Experimental Details
We mainly focus on finetuning LMs with LoRA. This setting naturally includes distributional shift the finetuning dataset does not usually have the same distribution the pretraining dataset as validated through zero-shot performance. All experiments are performed on server with AMD EPYC 7742 CPUs and NVIDIA GeForce RTX 3090 GPUs each with 24GiB memory. All numerical results from Section report test performance (e.g., accuracy, scores, or BLEU scores) and the standard deviation across multiple runs.
25

---

D.1
Details on Datasets
Our evaluations are carried out on commonly-used datasets in the literature.
GLUE benchmark. GLUE is designed to provide general- purpose evaluation of language under- standing (Wang et al., 2019b). Those adopted our work include MNLI (inference, (Williams et al., 2018)), SST-2 (sentiment analysis, (Socher et al., 2013)), MRPC (paraphrase detection, (Dolan and Brockett, 2005)), CoLA (linguistic acceptability (Warstadt et al., 2019)), QNLI (inference (Rajpurkar al., 2018)) 8)), QQP (guestion-answering)) RTE* (inference), and STS-B (textual similarity (Cer et al., 2017)). These datasets are released under different permissive licenses.
SuperGLUE benchmark. SuperGLUE (Wang et al., 2019a) is another commonly adopted bench- mark for language understanding and more challenging compared with GLUE. The considered datasets include CB (inference, (De Marneffe al., 2019)), ReCoRD (multiple-choice question answering (Zhang et al., 2018)), COPA (question answering (Roemmele et al., 201 1)). These datasets released under different permissive licenses.
WebNLG Challenge. This dataset is commonly used for data-to-text evaluation (Gardent et al., 2017). It has 22K examples in total with 14 distinct categories. Among them, 9 are seen during training, and the unseen training data are used to test the generalization performance. The dataset is released under license CC GY-NC-SA 4.0.
Additional datasets. We also use SQuAD (question answering (Rajpurkar et al.. 2016)) in our experiments, which is released under license CC BY-SA 4.0. Other datasets include TREC (topic classification (Voorhees and Tice, 2000)) and SNLI (inference (Bowman al., 2015)). Both of them are licensed under CC BY-SA 4.0.
D.2 Details Language Models
We summarize the adopted language models in evaluation. All model checkpoints obtained from HuggingFace.
RoBERTa-large. This is 355M parameter model. The model checkpoint is released under the MIT license.
OPT-1.3B. The model checkpoint is released under non-commercial license.
GPT2-medium. This is a 345M parameter model. Its checkpoint is under MIT License.
D.3 Few-shot Learning with RoBERTa and OPT
Experiments on RoBERTa-large. We follow the k-shot learning setup in (Malladi et al., 2023) and focus classification tasks. The training set contains & 512 samples per class while the test set has 1000 samples. We also employ prompts for finetuning; where the adopted prompts are the same those in (Malladi et al., 2023, Table 13). AdamW is adopted the base optimizer, and hyperparameters are tuned from those Table Our experiments are averaged over 3 random trials. The estimated runtime is about 5 minutes per dataset.
The per-iteration runtime on the SST-5 dataset of BAR, SAM, and the baseline optimizer are compared in Table 7. It can be seen that SAM is much more slower than the baseline approach, and BAR reduces 74% additional runtime of SAM, while achieving comparable accuracy. We believe that this runtime saving can be even larger with additional engineering efforts such as kernel fusion, which we leave for future work. This validates the computational efficiency of BAR.
aattps://quoraaaaaa wow.aaa..ommiittttoooraaaaaaaeaa htw.aaaaaaaaatteeeeeeeeeeeaaaaaa
'nttps://huggii httppp///uuggnnng......../////ss
md
'nttps://s3, www.aaams http:///w....ooooooooooo.ooaaaea
bin
26

---

Table 6: Hyperparameters used for few-shot learning with RoBERTa-large.
Hyper-parameters
Values
LoRA (rank)
8
LoRA
16
# iterations batchsize learning rate for SAM Aó for BAR scheduler for BAR
1000
16
1x10-1 11,1,,x10011 1x100--3x10-,55x100
0.05, 0.1, 0.2 0.5, 1.0. 2.0
linear, cosine
Table 7: Per-iteration runtime for finetuning RoBERTa-large on SST5,
SST5
baseline
SAM
BAR
time (s)
0.105
0.265
0.146
Experiments on OPT. For OPT-1.3B, we consider tasks from the SuperGLUE benchmark covering classification and multiple-choice. We also consider generation tasks on SQuAD. Following (Malladi et al., 2023), we randomly sample 1000 data for training and the other 1000 for testing. AdamW is adopted base optimizer. The hyperparameters adopted are searched over values in Table 8. Estimated runtime is less than or around 10 minutes, depending on the dataset.
If we directly apply FP16 training with SAM, underflow can happen if one does not take care of the gradient scaling on the two gradients calculated per iteration. This means that SAM not flexible enough to be integrated with the codebase for large scale training, as FP16 is the default choice for finetuning LMs. We employ FP32 to bypass the issue with SAM. Consequently, the training speed is significantly slowed down; see summary in Table 9, It further demonstrates the effectiveness of BAR for large ccale-training.
Overall, the results for few-shot learning indicate that given limited data, BAR can effectively improve generalization using significantly reduced computational resources relative to SAM.
Table 8: Hyperparameters used for few-shot learning with OPT-1.3B.
Hypar-parameters
Values
LoRA (rank)
8
LoRA # iterations batchsize learning rate P for SAM (4n for BAR scheduler for BAR
16
1000
2, 4, 8
1x10 1x111x011410000 1x1104,,5x0 5x100 5x000
0.05, 0.2
0.2, 0.5, 1.0, 2.0
linear, cosine
D.4 Finetuning with RoBERTa-large
Our implementation is inspired from (Hu et al., 2022)", which is under MIT License. The hyper- parameters are chosen the same as provided in its GitHub Repo. AdamW adopted as the base optimizer. However, we employ single GPU rather than multiple ones and use gradient accumulation rather than parallelism due to memory constraint. We also note that there could be failure cases for LoRA using certain seed, e.g., SST-2 with seed 1 and MNLI with seed 2. These cases ignored when comparing. We consider the GLUE benchmark and report the mismatched accuracy for MNLI, Matthew's 's correlation for CoLA, Pearson correlation for STS-B, and accuracy for other datasets. Larger values indicate better results for all datasets. For LoRA, we employ r = 8 and a = 16.
27

---

Table 9: Per-iteration runtime for finetuning OPT-1.3B on RTE,
RTE
baseline
SAM
BAR
precision
FPI6
FP32
FP16
time (s)
0.1671
0.708
0.1731
Table 10: Experiments on finetuning RoBERTa (355M). Results marked with f are taken from (Hu al., 2022), and those with refer Adapter in (Hu al., 2022).
RoBERTa
para
SST2
STS-B
QQP
QNLI
MRPC
MNLI
CoLA
avg
FT
355M
96.4
92.4
86.6
92.2
94.7
90.9
90.2
68.0
88.9
Adapter*
0.8M
96.6
91.9
80.1
91.7
94.8
89.7
67.8
LoRA
0.8M
95.8
92.4
88.2
91.4
94.7
89.6
90.6
64.8
88.4
LoRA-oBAR
0.8M
96.0
92.6
88.7
91.6
94.8
90.3
90.6
65.1
88.7
LoRA-nBAR
0.8M
96.0
92.6
89.2
91.6
94.7
90.3
90.8
65.6
88.9
Experiments are conducted over three random trials for all datasets, with the exception of QQP, for which only two trials are performed due to its large size. The results of final test performance can be found in Table 10. Estimated runtime varies for different datasets from to 15 hours, except QQP which takes 3 days on our device.
For the hyperparameters of oBAR and nBAR, μo is typically chosen from {0.2, 0.5, 1.0]: however, for QQP, value of 0.05 is used. The scheduler is chosen from linear and constant. We also observe that for datasets such COLA and RTE, setting weight decay as 0 works best for BAR.
D.5
GPT2 medium WebNLG Challenge
A is adopted as base optimizer. The hyperparameters can be found in Table 11. Our results are obtained from three random trials. Each trial takes roughly 8 hours hardware.
Table 11: Hyperparameters used for GPT2.
Hyper-parameters
Values
LoRA (rank)
4
LoRA a # epochs batchsize
32
5
8
learning rate label Smooth
2 10-4 4
0.1
140 for BAR scheduler for BAR
0.1, 0.15, 0.2, 0.25, 0.3
linear, constant
beam size length penalty
10 0.8
28

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our claims are supported by theoretical results in Sections and and numerical experiments in Sections 5 and 6. Due to space limitation, missing proofs and implementation details be found in the appendix.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer:
[Yes]
Justification:
Limitation is discussed in Section 7.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model cell-specification. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
29

---

Answer: [Yes]
Justification: All assumptions are stated along with theorems. The proofs are listed in appendix. All theories and equations are properly referenced.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
The details on proposed algorithms can be found in Section and Appendix A.6. Experimental details on setups, datasets, architectures, and hyperparameters can be found Section 6 and Appendix D. Code can be found in our github repo
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
30

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification:
The code is open-sourced github. The code gathers details for reproducing our experiments. For example, the exact environment is listed in environment ymll. Instructions on preparation (e.g., data, packages, etc) and commands to stated in ReadMe. md.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (attps://nips.cm cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/niiss /ttps.cccpublicc.i wwt.cc/pccccg.ccccccoooo geaiiic/ggee/////eeeeeeeeessssse /Ctps//oooeeeeeeeeeeeee wwwsssssiiiiyyyyyo for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if' applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Experimental details can be found Section 6 and Appendix D.
Guidelines:
The answer NA means that the paper does not include experiments.
the paper level of detail that is necessary to appreciate the results and make sense of them.
The experimental setting should be presented in the
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Mean and standard deviation are obtained by three random trials for most of experiments. See details Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoo bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
31

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for caleulating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (c.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Hardware for our experiments is detailed in Appendix D. Estimated runtime is also provided in Appendix D.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https : /ttt:///e..../eeeeeee cc cc/ Ctublic/Ethicsuutie
Answer:
[Yes]
Justification: NeurlPS Code of Ethics is followed.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Broader impacts can be found in Appendix A.1. It is put in appendix due to space limitation.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
32

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This paper does not release new data or models.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
The datasets and model checkpoints used in this work are widely adopted ones in the field. Their licenses are listed in Appendix D.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
33

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is available online, the authors encouraged to reach out the asset's s creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yss]
Justification: Our code is released under MIT license, unless model checkpoints and datasets are under more restrictive licenses; see more in supplementary materials.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
This work does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
34

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
35


Reviews:
[{"review_id": "bA8yXVpdJ4", "reviewer": "Reviewer_Cq3N", "summary": "### Summary:\n\nThis work introduces balancedness (instead of sharpness) and proposes Balancedness-Aware Regularization (BAR), used for scale-invariant problems (e.g., LoRA). Given an objective such as $f(xy^T)$ or $f(x^Ty)$ with parameters $x,y$ and a fixed function $f$, the balancedness is defined as $B_t = \\frac{1}{2}(\\|x_t\\|^2 - \\|y_t\\|^2)$. They prove that SAM decreases balancedness while SGD preserves it (Theorems 1&2&3) and argue that this new quantity goes beyond local information and relies on the entire SAM trajectory. Indeed, they show that whether $B_t$ goes to zero and gets quite small under SAM. \nMoreover, they show that outliers strongly affect balancedness, which helps explain SAM behavior. Finally, they propose the BAR algorithm and conclude with several experiments.", "questions": "### Questions/Comments:\n\n\nThis is an interesting paper. The authors nicely mention that balancedness changes under SAM and this could be one potential way to explain SAM behavior. Moreover, they propose balancedness minimization, which leads to even better algorithms. Unfortunately, the paper must be modified to reach the audience and become readable since the current version is hard to follow. \n\nSome comments:\n\n- line 37-39 -- this part is not well written and is confusing. What is the role of $d_1+d_2$? I suggest explaining it in a few sentences in the paper. \n\n- Limitation of sharpness (Section 2.1): the max e.v. is not the only sharpness measure considered in the literature. Indeed, there is apparently a class of invariant sharpness measures for SAM; for example, see this:\n        - A Universal Class of Sharpness-Aware Minimization Algorithms, ICML 2024\n\n- Do authors claim that Theorem 1 is the original result of this paper? I think this is known from previous works on the implicit biases of SGD, but it looks like the paper is claiming the results without references. See, for example, Equation 2.5 in this paper (already cited): \n        - Ji, Ziwei, and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. ICLR. \n\nPlease explain whether the theorem follows from the previous results or if there is a major difference between them. \n\n - The BAR algorithm is barely discussed in the main body of the paper—the algorithm itself is provided in the appendix. Because it is one of the main contributions of the paper, authors should consider rearranging stuff in the paper to place it in the main body and discuss it more. \n\n- Moreover, the ideas behind Algorithms 2 and 3 should be discussed in the main body because of their importance. What are the role of the hyperparameters in the algorithms?", "limitations": "N/A", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "### Pros:\n\n- the studied problem is very important, and the result can be impactful\n- the balancedness is completely new and well-motivated for SAM (previously some works studied related quantities but just for SGD)\n- having several experiments", "weaknesses": "### Cons:\n\n- the paper is not well-written, and it is difficult to follow it. The authors need to explain results, ideas, and notation via simple sentences and then deal with formulae/equations\n\n- some ambiguities about Theorem 1", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the time and efforts devoted to this paper. Below please find our point-to-point responses. Please let us know if there are other questions or unclearness during discussion. We are always glad to improve the quality of our submission.\n\n**W1.** *The paper is not well-written, and it is difficult to follow it.*\n\nWe hope a brief summary of this paper helps to clarify the overall logistic.\n\n* We show that SAM promotes balancedness on scale-invariant problems. Unlike many of existing metrics, balancedness is a global metric for implicit regularization. Theoretical and empirical evidences are provided to support SAM's implicit regularization on balancedness.\n\t\n* Since balancedness is computationally tractable, we mimic the dynamic of SAM (e.g., Theorem 2) and explicify balancedness as a regularizer -- BAR. BAR can be applied on top of other optimizers such as SGD or Adam in the same way as weight decay. BAR overcomes the need for the second gradient computation in SAM, yet achieving similar numerical performance as SAM.\n\nWe are also glad to clarify any ambiguities and revise the draft for better presentation -- please let us know during the discussion phase.\n\n\n**W2.** *Some ambiguities about Theorem 1*\n\nSince the confusion is detailed in questions, see our responses in Q3.\n\n\n**Q1.** *line 37-39... What is the role of $d_1 + d_2$?*\n\nTaking NOP problem $\\min_{\\mathbf{x}, \\mathbf{y}}f(\\mathbf{x}\\mathbf{y}^\\top)$ with $\\mathbf{x} \\in R^{d1}$ and $\\mathbf{y} \\in R^{d2}$ as an example. Here, we have that the loss function $f : R^{d1 \\times d2} \\mapsto 1$, in other words, the dimension of $dom f$ is $d1 \\times d2$. The number of variables to be optimized is only $dim(\\mathbf{x}) + dim(\\mathbf{y}) = d1 + d2$ since we parametrize $dom f$ by $\\mathbf{x}\\mathbf{y}^\\top$. Since the number of variables is smaller than the dimension of $dom f$, we are using insufficient variables to perform optimization, and hence its name -- non-overparametrization.\n\n\n**Q2.** *Other sharpness ... There is apparently a class of invariant sharpness measures for SAM; for example, see this: - A Universal Class of Sharpness-Aware Minimization Algorithms, ICML 2024*\n\nThank you for pointing out this work. Unfortunately, it is impossible for us to cite this ICML paper because it appeared later on arXiv (June 6) than NeurIPS submission due date (May 22). However, we have cited and discussed the differences with an earlier workshop version of the same ICML paper; see lines 596 - 597 in Appendix. Per request, the main differences with the ICML paper are re-summarized below.\n\nThe ICML paper introduces generalized sharpness measures -- sharpness as any function of eigenvalues of Hessian. However, even the generalized sharpness cannot provide implicit regularization for function $h(x,y)=xy$, simply because the Hessian is the same for all $(x, y)$. In addition, when Hessian is negative definite, some of the generalized sharpness measures (e.g., determinate of Hessian) may not be necessarily meaningful. Balancedness overcomes these two problems.\n\nIn terms of practical algorithms inspired by theoretical derivations, our goal is in sharp difference with the ICML paper. Our balancedness targets at computational efficiency of SAM. On the other hand, the ICML paper targets at improving over SAM under the same or slightly more computational budget (when the number of samples $n$ is set to be greater than $1$ in their Algorithm 1). \n\n\n**Q3.** *Do authors claim that Theorem 1 is the original result of this paper? I think this is known from previous works on the implicit biases of SGD, but it looks like the paper is claiming the results without references. See, for example, Equation 2.5 in this paper (already cited): - Ji, Ziwei, and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. ICLR.*\n\nThere is a potential misunderstanding. We do not hope to take credits from Theorem 1, as we have already cited several previous works in lines 131 - 134. Even in the ICLR paper of [Ji and Matus 2019], they have cited the same set of papers as we do. We slightly extend the results from these cited papers to the case with stochastic gradients. We will rephrase this paragraph to eliminate possible misunderstanding.\n\n\n**Q4 \\& Q5.** *The BAR algorithm is barely discussed in the main body of the paper—the algorithm itself is provided in the appendix ... Moreover, the ideas behind Algorithms 2 and 3 should be discussed in the main body because of their importance. What are the role of the hyperparameters in the algorithms?*\n\nThe design and the implementation of balancedness-aware regularizer (BAR) are already discussed in lines 280 - 285 in the main context. Since BAR is a regularizer, it can be used the same way as e.g., an $\\ell_2$ regularizer or weight decay. Note that this is quite standard, hence we assume the use of BAR should also be straightforward. Algs. 2 and 3 are just detailed BAR implementation in the same format as weight decay. Given the space limitation, instead of duplicating this weight-decay type implementation in the main context, we provide a short yet informative sentence in line 284 - 285 that convey the same massage. However, we would love to expand our discussions if we had additional space. \n\nThe ideas of Algs. 2 and 3 have been already discussed in lines 274 - 278 and lines 280 - 282, respectively. To summarize, the idea is that since SAM promotes balancedness implicitly (e.g., in Theorem 2),  we can explicify this dynamic and use it to regularize balancedness. \n\nThe only hyperparameter for BAR is the coefficient on this regularizer. It plays the same role as the coefficient for e.g., an $\\ell_2$ regularizer, that is, determining the tradeoff between regularization and loss-fitting.", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "cMAqkJsu4b", "reviewer": "Reviewer_7PUq", "summary": "This paper investigates the implicit regularization effects of sharpness-aware minimization (SAM) on scale-invariant optimization problems, introducing \"balancedness\" as a new metric for analysis. The authors provide theoretical results showing that SAM promotes balanced solutions for both non-overparameterized and overparameterized scale-invariant problems, and demonstrate that SAM's regularization effect is stronger on noisy data. Based on these insights, they propose balancedness-aware regularization (BAR), a computationally efficient variant of SAM. The paper evaluates BAR on language model fine-tuning tasks using LoRA, demonstrating that it can match or exceed SAM's performance while significantly reducing computational overhead.", "questions": "- How to understand the lower bound in Theorem 3, which is not always greater than zero during the optimization. \n\n- How to choose a good parameter $\\rho$?", "limitations": "Yes, the authors adequately addressed the limitations.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The theoretical analysis provides new insights into SAM's implicit regularization effects by introducing \"balancing\".\n- The proposed BAR method offers a more efficient alternative to SAM for scale-invariant problems by adding explicit regularization. \n- Comprehensive experiments on language model fine-tuning demonstrate practical benefits.", "weaknesses": "- The scope is limited to scale-invariant problems, and LoRA types optimization problem.\n- Assumption 1 requires  Lipschitz continuous gradient which may not always hold in practice (e.g., ReLU network)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the time and efforts devoted to this paper. Below please find our responses to weakness and questions.\n\n**W1.** *The scope is limited to scale-invariant problems, and LoRA types optimization problem.*\n\nWe agree with the reviewer that the major application is on scale-invariant problems such as LoRA and its variants. However, we hope to emphasize that LoRA is already one of the most popular approaches for finetuning language models. LoRA is more economical and easier to serve in practice compared to full parameter-tuning. LoRA type approaches are actively developed and well welcomed by the community; see e.g., HuggingFace's PEFT codebase [1]. Given these evidences, we believe that this scenario is important and our results are useful. \n\nMoreover, we definitely hope to inspire following-ups and even entirely new approaches to extend our results to more general settings. We are positive that this will take place in a step-by-step fashion, and we hope that our work is helpful for initiating this. \n\n**W2.** *Assumption 1 requires Lipschitz continuous gradient which may not always hold in practice (e.g., ReLU network).*\n\nLipschitz continuous gradient (or smoothness) is a standard assumption for the theoretical understanding of gradient-based methods; see e.g., SGD [2], SGD with momentum [3], AdaGrad [4], and Adam [5]. While smoothness does not necessarily hold in every practical scenario, the theoretical insights from the analysis still yield robust and meaningful guidance for practice, evidenced by the popularity of [2 - 5] for training neural networks. \n\nMoreover, we also hope to point out that in our LoRA scenario, there is even no ReLU activation for the trainable parameter (activations only apply for the frozen parameters). Thus, smoothness is slightly more reasonable compared to the case of Adam on deep neural networks.\n\n\n**Q1.** *How to understand the lower bound in Theorem 3, which is not always greater than zero during the optimization.*\n\nWhen it is smaller than $0$, the inequality trivially holds. However, when $\\rho$ is chosen small, it is more likely that the RHS is positive given that $\\rho^2 \\ll \\rho$. We will update the lower bound to $\\max$ {0, lower bound in Theorem 3}.\n\n**Q2.** *How to choose a good parameter $\\rho$?*\n\nEmpirical experiences recommend $0.05$ or $0.1$. In practice, a grid search is also helpful.\n\n\n**References.**\n\n[1] https://github.com/huggingface/peft\n\n[2] S Ghadimi, and G Lan. Stochastic First- and Zeroth-order Methods for Nonconvex Stochastic Programming. SIAM J-OPT 2013.\n\n[3] Y Liu, Y Gao, and W Yin. An Improved Analysis of Stochastic Gradient Descent with Momentum. NeurIPS 2020\n\n[4] R Ward, X Wu, and L Bottou. AdaGrad Stepsizes: Sharp Convergence over Nonconvex Landscapes. JMLR 2020\n\n[5] X Chen, S Liu, R Sun, and M Hong. On the Convergence of A Class of Adam-Type Algorithms for Non-Convex Optimization. ICLR 2019", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "WfNEmwPxOm", "reviewer": "Reviewer_9YQw", "summary": "This paper investigates the dynamics of SAM when the loss is of the form f(xy^T) or f(x^Ty). This formulation includes interesting scenarios such as LoRA. This paper shows that SAM will promote balancedness, which is the difference between the squared norms of two variables. Based on this new analysis, this paper proposes to regularize balancedness and show that this method leads to similar or sometime better result than SAM with significantly less compute.", "questions": "1. Because SGD with a small learning rate will make balancedness almost unchanged, does the result imply that initializing LoRA to be completely balanced and using SGD should also improve performance?\n\n2. Connecting to the first problem, as in Figure 3, the balancedness of the weight trained by SAM, while decreasing faster than SGD, still remains at a high level. Is this the case for BAR as well?", "limitations": "The limitation is clearly addressed.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. This paper considers the interesting interplay between SAM and scale-invariance of the loss and discovers that SAM implicitly regularizes balancedness in this case.\n2. The empirical verification of the theoretical prediction is thorough and convincing.", "weaknesses": "1. Arguments including that SGD will make balancedness unchanged are only correct for infinitesimal learning rate and claims like this should be made rigorous (for example at line 137).\n\n2. In section 2.1, when the limitation of sharpness is discussed, the example that h(x,y) = xy is very confusing because there is no local minima for such loss.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer for the nice questions. We will add these discussions to the draft. \n\n**W1.** *Arguments including that SGD will make balancedness unchanged are only correct for infinitesimal learning rate and claims like this should be made rigorous (for example at line 137).*\n\nThank you for pointing this out. We will proofread a few more times and rephrase these sentences. In Figure 3, we have shown that even with a learning rate $0.1$, SGD only slightly changes the balancedness.  \n\t\n**W2.** *In section 2.1, when the limitation of sharpness is discussed, the example that $h(x,y) = xy$ is very confusing because there is no local minima for such loss.*\n\t\nThank you for pointing this out. We will rephrase this sentence. Our intention was that the function $h(x,y) = xy$ has the same Hessian (hence eigenvalues) for any $(x, y)$. Thereby, one cannot provide any implicit bias using the largest eigenvalue. \n\t\n\t\n**Q1.** *Because SGD with a small learning rate will make balancedness almost unchanged, does the result imply that initializing LoRA to be completely balanced and using SGD should also improve performance?*\n\nUnfortunately, it is difficult to initialize LoRA in a balanced fashion. Using the notation of equation (6), one has to set $\\mathbf{Y} = 0$ to ensure that at initialization $\\mathbf{X}\\mathbf{Y}^\\top = \\mathbf{0}$. To achieve exact balancedness, the only choice is to choose $\\mathbf{X}=0$ as well. This is not a good initialization because $\\mathbf{X}=0$ and $\\mathbf{Y}=0$ is a saddle point (as one can see from the gradients). \n\nAn alternative is to initialize $\\mathbf{Y}=0$ and choose a small $\\mathbf{X}$ in magnitude, i.e., initializing close to a saddle point. However, it is known that SGD escapes saddle points quite slow [1]. \n\nThis means that LoRA initialized in a balanced manner can slow down convergence, and it in turn necessities our balancedness-aware regularization (BAR).\n\n**Q2.** *Connecting to the first problem, as in Figure 3, the balancedness of the weight trained by SAM, while decreasing faster than SGD, still remains at a high level. Is this the case for BAR as well?*\n\nIn Figure 3, the balancedness of SAM is decreasing less faster because of the multiple layers of the RoBERTa. As we have discussed in lines 202 - 206 and Theorem 5 in Appendix, multiple layers slowdown the implicit regularization of SAM, and the slowdown is roughly proportional to the square root of (LoRA) layers. In the case of Figure 3, there are 24 transformer layers, and LoRA is applied to every query and value in each transformer layer. This gives a $\\sqrt{24 * 2} \\approx 7$x slowdown. \n\nThe proposed BAR does not suffer from the slowdown from number of layers, because the regularization can be applied individually on each LoRA module for the query and value layers. This gives a better regularization on balancedness and explains why nBAR can improve over SAM in Table 3.\n\n\n**References**\n\n[1] C Fang, Z Lin, and T Zhang. Sharp analysis for nonconvex SGD escaping from saddle points. COLT 2019.", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "zob7EAsp7Q", "reviewer": "Reviewer_t2ze", "summary": "This paper analyzes the implicit regularization effect of Sharpness-Aware Minimization (SAM), focusing specifically on scale-invariant problems. While existing research emphasizes sharpness, this study introduces a new concept called Balancedness, demonstrating both theoretically and empirically that SAM promotes Balancedness. Additionally, the proposed Balancedness-Aware Regularization (BAR) is shown to significantly improve computational efficiency while achieving superior performance in fine-tuning language models using LoRA.", "questions": "- As cited in the Appendix, it is known that SAM promotes the acquisition of low-rank features [2]. However, it was pointed out that explicit regularization for acquiring low-rank features did not lead to improved generalization ability. My curiousity is about the relationship between Balancedness and low-rank features, and whether there is any connection to sharpness.\n\n- Another question is how optimization methods commonly used in LLM training, such as AdamW, affect Balancedness.\n\n[2] https://arxiv.org/abs/2305.16292", "limitations": "The authors sufficiently discuss the limitations of this study, particularly regarding the application range and computational efficiency of BAR.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "Originality\n- Introducing a new concept of Balancedness and reinterpreting the implicit regularization effect of SAM from a new perspective.\n\nQuality\n- Consistent theoretical analysis and empirical validation with clear results.\n\nClarity\n- Clear presentation with a structure that is easy for readers to understand.\n\nSignificance\n- Revealing a new aspect of SAM's adaptability to data anomalies and proposing a computationally efficient BAR.\n\nSoundness\n- The technical claims, experimental methods, and research techniques of this paper are robust, and the central claims are well-supported by ample evidence. Both the theoretical analysis and experimental results are consistent, and the conclusions are clear. Potential concerns, such as those regarding m-sharpness, are addressed.\n\nContribution\n- The paper theoretically analyzes the mechanism by which SAM promotes balance and proves that balance converges to |Bt| \\rightarrow 0. It particularly highlights the strong impact of data anomalies (outliers). This explains why SAM outperforms SGD even in the presence of data anomalies. The implicit regularization of SAM is made explicit, and a new, computationally efficient variant of SAM called Balancedness-Aware Regularization (BAR) is proposed.\n- Experiments using various models confirm that the proposed BAR exhibits superior performance in fine-tuning with LoRA compared to traditional SAM and SGD.\n- While SAM has been noted as a promising optimization method to enhance the performance of large language models [1], its high computational cost and implementation complexity have limited its practical use. However, this study's regularization from the perspective of Balancedness has a potentiall to addresses these issues in conjunction with LoRA.\n\n[1] https://arxiv.org/abs/2210.14199", "weaknesses": "- While LoRA training is performed with AdamW in the original paper, demonstrating that SAM or BAR is more effective than these would make the paper more solid.\n- Code such as `eval.py` was not included in the supplementary materials, and its release is expected.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the time and efforts from the reviewer put into this review. We also want to thank the reviewer for recognizing the strength of our work. We will update our draft and code repo to further improve the quality of this work.\n\n**W1.** *While LoRA training is performed with AdamW in the original paper, demonstrating that SAM or BAR is more effective than these would make the paper more solid.*\n\nThe LoRA baselines in our experiments are indeed trained with AdamW (or Adam for different experiments), and our implementation follows the official LoRA repo. \n\t\n**W2.** *Code such as ```eval.py``` was not included in the supplementary materials, and its release is expected.*\n\nThe code for ```eval.py``` can be found in the official repo of LoRA; see [1]. We will update the ReadMe.md to provide improved instructions.\n\t\n**Q1.** *As cited in the Appendix, it is known that SAM promotes the acquisition of low-rank features [2]. However, it was pointed out that explicit regularization for acquiring low-rank features did not lead to improved generalization ability. My curiousity is about the relationship between balancedness and low-rank features, and whether there is any connection to sharpness.*\n\nThe setting in [2] does not extend to our case. This is because that the feature learning happens in the principle space in [2], while LoRA learns features in the residue space. Recall that low rankness is induced by parameterizing the last layer as $\\mathbf{A} \\mathbf{B}^\\top$; see section 6 of [2]. In other words, the learned feature is $\\mathbf{X}\\mathbf{A} \\mathbf{B}^\\top$, where $\\mathbf{X}$ is data or feature from the previous layer.  However, in our LoRA case, the learned feature is $\\mathbf{X}(\\bf{W} + \\mathbf{A} \\mathbf{B}^\\top ) $, where $\\mathbf{W}$ is the pretrained model. In sum, low rankness can play fundamentally different roles in [2] compared with our LoRA case.\n\nRegarding the relation of balancedness and sharpness, there are indeed links such as what we have shown in Lemma 1. Regarding the relation between balancedness and low rank features, this is beyond the scope of this work. As there are multiple approaches to induce low-rankness (and [2] employs one of them), we believe that a more systematic study should be carried out for answering this question, and we leave it to future work.\n\n\t\n**Q2.** *Another question is how optimization methods commonly used in LLM training, such as AdamW, affect Balancedness.*\n\nIt is not difficult to analytically show that a variant of ADAM (without momentum) decreases balancedness slightly on a simple loss function $f(x,y)=0.5(xy - 1)^2$ when learning rate $\\eta \\rightarrow 0$. However, the balancedness does not decrease after approaching near optimal. This can be seen from Fig. 1 in the additional PDF file (see general response). Note that we use $\\eta = 5e-4$ with $10^4$ iterations. However, this analysis on the simple loss function may not extend to more general cases.\n\nIn practice, we perform the same experiments as Fig. 3 and plot the balancedness of all layers after training using Adam in Fig. 2 of the additional PDF. No specific trend on balancedness (2${\\cal B}_{t,l}$) is observed, because the balancedness on some layers increases compared to initialization; while decreases for other layers. Note that the initialized balancedness is within (2.7 - 3.2) for all layers.\n\n**References**\n\n[1] LoRA official repo: https://github.com/microsoft/LoRA/blob/main/examples/NLG/eval/eval.py\n\n[2] M. Andriushchenko, D. Bahri, H. Mobahi, N. Flammarion. Sharpness-aware minimization leads to low-rank features. NeurIPS 2023", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:
We thank the ACs and reviewers for handling this submission. Your comments are appreciated, and the manuscript will also be updated accordingly. Our point-to-point responses can be found below, and a .pdf file is also attached as graphical illustration to questions from Reviewer t2ze.

Meta-review: The paper presents a novel and well-supported contribution to the understanding and application of SAM in scale-invariant problems, particularly in the context of LoRA. The introduction of BAR (Balancedness-Aware Regularization) offers a practical and computationally efficient alternative to SAM, with potential for significant impact in the field. The reviewers unanimously recommended to accept this paper.

For the camera-ready version, the authors should incorporate suggestions and discussions in the review and rebuttal, including improving the clarity and presentation of the BAR algorithm and its theoretical foundations as well as explicitly addressing and clarifying any assumptions made in the paper.
Decision: Accept (poster)