Title: Implicit Regularization of Sharpness-Aware Minimization for Scale-Invariant Problems
Venue: NeurIPS 2024
Accepted: yes
Keywords: sharpness-aware minimization

LoRA

implicit regularization

finetuning

language models
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

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.

1 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 et al., 2023a).

However, existing results only characterize sharpness/flatness 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, theoretical 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 ( Q1 ): can we broaden the scope of implicit regularization to 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 on sparse coding problems in the small signal-to-noise ratio (SNR) regime (Chen et al., 2023). Remarkable performance of SAM is also observed under distributional shifts, e.g., domain adaptation (Wang et 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 our second question ( Q2 ): can implicit regularization of SAM reflect its enhanced performance under data anomalies?

This work answers both Q1 and Q2 within a class of scale-invariant problems. The focus on scaleinvariance is motivated by its prominence in deep learning architectures. Consider variables x ∈ R d 1

0

50000

100000

iter

0.0

0.4

0.8

1.2

abs(||

x

||

2

||

y

||

2

)

SGD

=1

SAM

=1

SAM

=0.5

0

50000

100000

iter

0

10

20

30

||

gx

||

2

||

gy

||

2

0

50000

100000

iter

0.0

0.4

0.8

1.2

abs(||

x

||

2

||

y

||

2

)

SGD

=1

SAM

=1

SAM

=0.5

0

50000

100000

iter

0

5

10

15

20

||

gx

||

2

||

gy

||

2

(a) non-overparametrized (NOP)

(b) overparametrized (OP)

and y ∈ R d 2 , both in high-dimensional space. The problems of interest can be categorized into non-overparametrization (NOP) and overparametrization (OP), based on whether the dimension of variables ( d 1 + d 2 ) is greater than dimension of dom f ,

glyph[negationslash]

Here, d 1 = d 2 is assumed for OP, and D denotes the training data. For both cases, the losses are nonconvex in ( x , y ) . Scale-invariance refers to that ( α x , y /α ) share the same objective value ∀ α = 0 . It naturally calls for implicit regularization from optimization algorithms to determine the value of α . 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 is 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 t = 1 2 ( ‖ x t ‖ 2 -‖ y t ‖ 2 ) . 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 t ; and, iii) it enables a data-driven perspective for understanding SAM. Building on balancedness, we answer our major questions.

For Q1, we prove that even with imbalanced initialization, SAM drives |B t | → 0 for OP, while ensuring a small |B t | in NOP. In contrast, we 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 adoption of balancedness, our results on implicit regularization have no requirement on the batchsize compared to (Wen et al., 2023a) and can be extended to explain m -sharpness in (Foret et al., 2021).

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. 1 showcases an example where SAM is applied on the same problem with different SNRs. Smaller SNR (i.e., larger α ) 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 anomalies. 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 as a data-driven regularizer. When applied on top of, e.g., SGD, it enables a computationally efficient variant of SAM, balancedness-aware regularization (BAR), suited for scale-invariant 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 a nutshell, our contribution can be summarized as:

v Theories. Balancedness is introduced as a new metric for implicit regularization in SAM. Compared to sharpness, balancedness enables us to depict richer behaviors - SAM favors balanced solutions for both NOP and OP, and data anomalies have stronger regularization on balancedness.

v Practice. Implicit regularization of SAM is made explicit for practical merits. The resulting approach, balancedness-aware 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 glyph[lscript] 2 (Frobenius) norm of a vector (matrix), and ( · ) glyph[latticetop] 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 are prevalent in modern neural networks, such as LoRA, ReLU networks, and softmax in attention. However, scale-invariant problems are not yet fully understood, especially from a theoretical perspective. Neyshabur et al. (2018) develop scale-invariant PAC-Bayesian bounds for ReLU 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-invariance and overparametrization simultaneously, and IR of (S)GD on quadratic loss is established in (Arora et al., 2018; Du et al., 2018; Gidel et 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 in (Dinh et al., 2017) that sharpness is sensitive to scaling, while our results indicate that when taking the training trajectory into account, SAM excludes extreme scaling.

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 to be 1 in the stochastic case. Andriushchenko and Flammarion (2022) consider sparsity of SAM on (overparametrized) diagonal linear networks on a regression problem. Chen et al. (2023) study the benign overfitting of SAM on a two-layer ReLU network. In general, existing studies on SAM's implicit regularization focus more on sharpness and do not fully capture scale-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.

SAMvariants. Approaches in (Kim et 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 or sparsifying gradients (Liu et al., 2022; Mi et al., 2022); stochastic perturbation (Du et al., 2022a); switching to SGD (Jiang et al., 2023); and connecting with distillation (Du et al., 2022b). Our BAR can be viewed as resource-efficient 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 designed originally to seek for solutions in flat basins. The idea is formalized by enforcing small loss around the entire neighborhood in parameter space, i.e., min w max ‖ glyph[epsilon1] ‖≤ ρ h ( w + glyph[epsilon1] ) , where ρ is the radius of considered neighborhood, and h ( w ) := E ξ [ h ξ ( w )] . Practical implementation of SAM is summarized under Alg. 1. It is proved in (Wen et al., 2023a)

Algorithm 1 SAM (Foret et al., 2021)

1: Initialize: w 0 , ρ, T, η

2: for t = 0 , . . . , T -1 do

3: Sample ξ to get a minibatch M t

4: Define stochastic gradient on M t as ∇ h t ( · )

5: Find glyph[epsilon1] t = ρ ∇ h t ( w t ) / ‖∇ h t ( w t ) ‖

6: Update via w t +1 = w t -η ∇ h t ( w t + glyph[epsilon1] t )

7: end for

that ‖∇ h t ( w ) ‖ glyph[negationslash] = 0 (in line 5) holds for any ρ under most initialization. Based on this result and similar to (Dai et al., 2023), we assume that SAM iterates are well-defined.

Limitation of sharpness. Coming naturally with SAM is the so-termed sharpness, given by S ( w ) := max ‖ glyph[epsilon1] ‖≤ ρ h ( w + glyph[epsilon1] ) -h ( w ) . When ‖∇ h ( 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 ‖∇ h ( w ) ‖ → 0 are studied. In addition, sharpness (the largest eigenvalue) is not always informative for scale-invariant problems (1). Consider h ( x, y ) = xy for example. The sharpness is 1 for any ( x, y ) - these points are not distinguishable in terms of sharpness.

2.2 Prelude on Balancedness

Balancedness B t := 1 2 ( ‖ x t ‖ 2 -‖ y t ‖ 2 ) turns out to be an intriguing alternative to sharpness on the scale-invariant problem (1). Being a 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 t 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 as (Arora et al., 2018, 2019b; Wen et al., 2023a), we consider SGD with infinitesimally small learning rate η → 0 for the NOP problem (1a)

Theorem 1 ((Arora et al., 2018, 2019a; Ji and Telgarsky, 2019; Ahn et al., 2023)) . When applying SGD on the NOP (1a) , the limiting flow with η → 0 satisfies ‖ x t ‖ 2 -‖ y t ‖ 2 = ‖ x 0 ‖ 2 -‖ y 0 ‖ 2 for all t > 0 . In other words, d B t d t = 0 holds.

Theorem 1 shows that B t ≡ B 0 given η → 0 . A graphical illustration can be found in Fig. 1 (a). Another interesting observation is that given the same initialization, B t is fixed for SGD regardless of training datasets. This suggests that SGD is less adaptive to data. A similar result of Theorem 1 can be established for SGD on OP. The full statement is deferred to Apdx. C.1; see also Fig. 1 (b).

Merits of being balance. Because B 0 is preserved, SGD is sensitive to initialization. For example, ( x 0 , y 0 ) and (2 x 0 , 0 . 5 y 0 ) can result in extremely different trajectories, although the same objective value is shared at initialization. Most of existing works initialize B 0 ≈ 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 on minibatch M for example

Assuming bounded variance E [ ‖ 1 |M| ∑ ξ ∈M ∇ f ξ n ( xy glyph[latticetop] ) -∇ f n ( xy glyph[latticetop] ) ‖ 2 ] ≤ σ 2 , it can be seen that the variance of [ g x , g y ] is bounded by σ 2 ( ‖ x ‖ 2 + ‖ y ‖ 2 ) . In other words, among { ( x , y ) | xy glyph[latticetop] = W } , gradient variance is minimized if ‖ x ‖ = ‖ y ‖ , i.e., being balance. Moreover, block smoothness parameters L x n and L y n 1 also hint upon the difficulties for optimization, where large values typically correspond to slow convergence (Bottou et al., 2018; Nesterov, 2004). With the help of Assumption 1 (in the next subsection), it can be seen that L x n = L n ‖ y ‖ 2 and L y n = L n ‖ x ‖ 2 . In other words, a large |B t | implies difficulty for optimizing one variable than the other. For these reasons, balancedness is well-appreciated in domains such as matrix factorization/sensing - a special case of (1a) (Tu et al., 2016; Bartlett et 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).

2.3 Assumptions and Prerequisites

To gain theoretical insights of scale-invariant problems in (1), we assume that the loss has Lipschitz continuous gradient on dom f following common nonconvex optimization and SAM analyses (Bottou et al., 2018; Andriushchenko and Flammarion, 2022; Wen et al., 2023a).

Assumption 1. Let W ∈ R d 1 × d 2 , and w ∈ R . For each ξ , f ξ n ( W ) and f ξ o ( w ) in (1) have L n , and L o Lipschitz continuous gradient, respectively.

Scale-invariant problems are challenging to solve even on 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-Overparametrized Problems

This section tackles the implicit regularization of SAM on NOP (1a). Motivated by practical scenarios such as LoRA, we focus on cases initialized with large |B 0 | .

When ambiguity is absent, the subscript in f n and L n is ignored in this section for convenience. Applying Alg. 1 on NOP, the update of SAM can be written as

where ρ > 0 is the radius of SAM perturbation; u t := 1 / √ ‖ g x t ‖ 2 + ‖ g y t ‖ 2 ; and f t , ∇ f t denote the loss, stochastic gradient on minibatch M t , respectively.

Theorem 2. (Dynamics of SAM.) Suppose that Assumption 1 holds. Consider SAM for NOP in (4) with a sufficiently small ρ . Let B t := 1 2 ( ‖ x t ‖ 2 -‖ y t ‖ 2 ) . For some |A t | = O ( ρ 2 L ) and η → 0 , the limiting flow of SAM guarantees that

Moreover, the change on B t depends on the difference of stochastic gradients on x t and y t , i.e.,

Unlike SGD for which d B t d t = 0 , Theorem 2 states that the balancedness for SAM is driven by gradient difference ‖ g x t ‖ 2 -‖ g y t ‖ 2 . To gain some intuition, if we estimate ‖ g x t ‖ 2 -‖ g y t ‖ 2 ∝ ‖ y t ‖ 2 -‖ x t ‖ 2 based on (3) and ignore A t , it can be seen that d B t d t ∝ -ρ B t . This indicates the contraction on |B t | . A graphical illustration on decreasing |B t | , and its relation with gradient difference can be found in Figs. 1 (a) and 2 (a). Moreover, this implicit regularization on balancedness is global as it holds for all t regardless of whether ( x t , y t ) is close to local optima. Thanks to adopting balancedness as the metric, Theorem 2 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 ρ t ≥ 0 such that whenever |B t | > ¯ B ρ t , the magnitude of B t shrinks, where ¯ B ρ t can be found in (21) at appendix.

Corollary 1 shows that SAM promotes balancedness until |B t | reaches lower bounds ¯ B ρ t . Because ¯ B ρ t depends on SAM's trajectory, we plot 1 T ∫ T 0 ¯ B ρ t dt using dotted lines for better visualization in Fig. 2 (a). It can be seen that our calculation on ¯ B ρ t almost matches the balancedness of SAM after sufficient convergence. Being balance also reveals that the benefit of SAM can come from optimization, which is 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. 2 (a), i.e., E [ ‖ xy glyph[latticetop] -( A + α N ) ‖ 2 ] , where A is ground truth; N is data noise; and α determines SNR. For this problem, noisy data directly lead to noisy gradients. It can be seen in Fig. 2 (a) that smaller SNR coincides with faster decreasing of |B t | . To explain such a data-responsive behavior in implicit regularization, Theorem 2 states that balancedness changes largely when the difference of ‖ g y t ‖ and ‖ g x t ‖ is large. Since E [ ‖ g y t ‖ 2 -‖ g x t ‖ 2 ] ∝ α 2 if assuming elements of N to be iid unit Gaussian variables, it thus implies that a small SNR (large α ) offers large regularization on balancedness.

0

50000

100000

iter

0.0

0.4

0.8

1.2

abs(||

x

||

2

||

y

||

2

)

SAM

=0.5

thres

=0.5

SAM

=1

thres

=1

SAM

=1.5

thres

=1.5

0

10

||

gx

||

2

||

gy

||

2

0

10

0

50000

100000

iter

0

10

0

50000

100000

iter

0.0

0.4

0.8

1.2

abs(||

x

||

2

||

y

||

2

)

SGD

SAM

oBAR

nBAR

0

50000

100000

iter

10

0

10

20

||

gx

||

2

||

gy

||

2

(a) threshold of balancedness

¯

B

ρ

t

(b) relation with regularization

Extension to LoRA (multi-layer two-variable NOP). For LoRA, the objective is to minimize D blocks of variables simultaneously, i.e., min E ξ [ f ξ ( { x l y glyph[latticetop] l } D l =1 )] . It is established in Theorem 5 in appendix that SAM cultivates balancedness in a layer-wise fashion, i.e., the magnitude of B t,l := 1 2 ( ‖ x t,l ‖ 2 - ‖ y t,l ‖ 2 ) cannot be large for each l . However, the | d B t,l / d t | can be O ( √ D ) times smaller than Theorem 2 in the worst case because of the additional variables.

Validation of IR on modern architectures. Going beyond the infinitesimally small step size, we adopt η = 0 . 1 on modern language models to validate our theoretical findings. We consider finetuning a 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 a clear trend of promoting balancedness, aligning well with our theoretical predictions.

4 SAM for Overparametrized Problems

Next, we focus on SAM's implicit regularization on OP (1b). Overparametrization enables SAM to have stronger regularization on balancedness. Subscripts in f o and L o are omitted for convenience. SAM's per iteration update for OP can be summarized as

where u t := sgn ( f ′ t ( x glyph[latticetop] t y t )) / √ ‖ x t ‖ 2 + ‖ y t ‖ 2 ; f t and f ′ t denote the loss, stochastic gradient on minibatch M t , respectively. Different from NOP, SAM has stronger regularization on balancedness, where |B t | decreases whenever the norm of stochastic gradient is large. To see this, it is convenient to define C t := |‖ x t ‖ - ‖ y t ‖| . Note that C t ≤ √ 2 |B t | .

Theorem 3. Consider η → 0 for (7) . The limiting flow of SAM on OP ensures a decreasing magnitude of B t whenever | f ′ t ( x glyph[latticetop] t y t ) | · C t > O ( ρL |B t | ) . Moreover, the speed of decrease can be lower- and upper- bounded as

Given ρ → 0 and sufficiently noisy data, Theorem 3 implies that |B t | → 0 . 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 t | 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 | d B t / d t | is directly related with | f ′ t ( x glyph[latticetop] t y t ) | , and E [ | f ′ t ( x glyph[latticetop] t y t ) | ] is clearly larger when data are more noisy. In 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 a large margin.

0

2000 4000 6000 800010000

2.68

2.66

2.64

2.62

2.60

2.58

RTE

SAM

SGD

0

2000 4000 6000 800010000

2.7

2.6

2.5

2.4

SST-5

0

2000 4000 6000 800010000

2.70

2.65

2.60

2.55

2.50

MNLI

Extension to m -sharpness. m -sharpness is a variant of SAM suitable for distributed training. It is observed to empirically improve SAM's performance (Foret et al., 2021). m -sharpness evenly divides minibatch M t into m disjoint subsets, i.e., { f t,j } j m =1 , and perform SAM update independently on each subset; see (38) in appendix. It turns out that m -sharpness can also be explained using balancedness. With formal proofs in Apdx. C.3, the IR of m -sharpness amounts to substitute | f ′ t ( x glyph[latticetop] t y t ) | in Theorem 3 with 1 m ∑ m j =1 | f ′ t,j ( x glyph[latticetop] t y t ) | . This means that the regularization on balancedness from m -sharpness is more profound than vanilla SAM, because 1 m ∑ m j =1 | f ′ t,j ( x glyph[latticetop] t y t ) | ≥ | f ′ t ( x glyph[latticetop] t y t ) | .

Finally, we connect balancedness with sharpness on local minima of OP.

Lemma 1. Let W ∗ = { ( x , y ) | x glyph[latticetop] y = w,f ′ ( w ) = 0 , f ′′ ( w ) > 0 } be non-empty. For the OP problem (1b) , minimizing sharpness within W ∗ is equivalent to finding B = 0 in W ∗ .

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 et 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 to extend results to 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 et 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. 2

5.1 Overview of LoRA

Given a pretrained model with frozen weight W l ∈ R d 1 × d 2 on a particular layer l , the objective of LoRA is to find low rank matrices X l ∈ R d 1 × r , and Y l ∈ R d 2 × r with r glyph[lessmuch] min { d 1 , d 2 } such that the loss is minimized for a downstream task, i.e.,

LoRA enjoys parameter efficiency for finetuning thanks to the low-rank matrices X l and Y l . For instance, it only requires 0.8M trainable parameters to finetune a 355M-parameter RoBERTa-large (Hu et al., 2022). The outer product of X l and Y l 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 on trainable parameters. Unbalancedness is also unavoidable for LoRA, due

Algorithm 2 nBAR

Algorithm 3 oBAR

to the need of initializing at X l ∼ N (0 , σ 2 ) , Y l = 0 ; see an example of RoBERTa-large in Fig. 3. The unbalancedness leads to instability of LoRA when finetuning RoBERTa on 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 computational 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 in large-scale finetuning. Note that concern i) is annoying given the size of language models, especially in setups involving model parallelism.

Our balancedness-aware regularization (BAR) is 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 α t | x glyph[latticetop] x -y glyph[latticetop] y | . This regularizer originates from matrix sensing; see e.g., (Tu et al., 2016; Ge et al., 2017). For OP, choosing α t := O ( | f ′ ( x glyph[latticetop] t y t ) | / √ ‖ x t ‖ 2 + ‖ y t ‖ 2 ) recovers SAM's dynamic on B t up to an error of O ( ρ 2 ) ; cf. Lemma 2 in appendix. By ignoring this error, it can be seen that B t decreases when ‖ x t ‖ ≥ ‖ y t ‖ . Following this dynamic, we regulate balancedness based on whether ‖ x t ‖ ≥ ‖ y t ‖ . The resultant approach is termed as overparamterized BAR (oBAR) to reflect its source in OP.

On the other hand, because LoRA is NOP inherently, we take inspiration from Theorem 2 - dropping the term A t and mimicking dynamics of SAM. In particular, we regulate the objective with α t ( x glyph[latticetop] x -y glyph[latticetop] y ) if ‖ g x t ‖ 2 < ‖ g y t ‖ 2 ; otherwise α t ( y glyph[latticetop] y -x glyph[latticetop] x ) . The resultant approach is termed as nBAR. A graphical illustration can be found in Fig. 2 (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. 2 and 3, respectively.

Another benefit of BAR, in additional to the lightweight computation, is that it can be applied individually on 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 α t . In both nBAR and oBAR, one can employ a decreasing scheduler for α t for algorithmic flexibility. This is motivated by the fact that for both NOP and OP problems, the implicit regularization of SAM is less powerful after sufficient balancedness or near optimal. Commonly adopted cosine and linear schedules work smoothly.

6 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 are used to highlight the best and second best performance, respectively. More experimental details can be found in Apdx. D. Code is available at https://github.com/BingcongLi/BAR .

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 a language model with a 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-large are summarized in Table 1. As indicated by the zero-shot performance, the distributional shift between finetuning and pretraining datasets is obvious. This is a 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 is close to SAM. Moreover, BAR saves 74% additional runtime of SAM; see more details in Table 7 in the appendix.

The proposed nBAR and oBAR perform even better when scaling up to OPT-1.3B. BAR reduces the 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 as a benchmark for comparisons on test performance. We report F1 score for SQuAD and accuracy for other datasets in Table 3. The averaged improvement over LoRA is 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 in regularizing balancedness. Balancedness of a particular layer is decreasing slower in SAM due to multiple layers, as 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 OPT-1.3B, it is conjectured that BAR has more potential for larger models, and the verification is left for future due to hardware constraints.

6.2 Finetuning with RoBERTa-large

Having demonstrated the power of BAR in few-shot learning, we then apply it to finetune RoBERTalarge with LoRA. The results can be found in Table 4. 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.

On average, oBAR leads to a gain of 0 . 4 , and nBAR raises the test performance by 0 . 6 . BAR thereby fills the gap of test performance between LoRA (0.8M) and full-parameter (355M) finetuning.

6.3 Text Generation on GPT2-medium

Lastly, we consider BAR on a text-generation problem using GPT2-medium, a model with 345M parameters. Results on WebNLG (Gardent et al., 2017) are reported in Table 5. 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, as an alternative to commonly adopted sharpness, is employed as the metric to capture global and data-responsive behaviors of 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 a computationally efficient SAM variant, termed BAR. The effectiveness of BAR is 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 scale-invariant modules in neural networks. Finetuning language models with LoRA, as a popular option in practice, is a setting naturally suitable for our approach. However, our approach does not apply for linear models, e.g., logistic regression. Regarding future directions, an interesting one is whether SAM has other forms of implicit regularization beyond balancedness and sharpness. The exploration of other scale-invariant architectures beyond LoRA, e.g., the softmax function in attention, is also deferred to future work.

References

Momin Abbas, Quan Xiao, Lisha Chen, Pin-Yu Chen, and Tianyi Chen. Sharp-MAML: Sharpnessaware 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: eigenvalue 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. In Proc. Adv. Neural Info. Processing Systems , volume 36, 2023.

Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware minimization. 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 overparameterization. In Proc. Int. Conf. Machine Learning , pages 244-253. 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 on 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 linear 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-36, 2023.

Léon Bottou, Frank E 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 1: Semantic textual similarity-multilingual and cross-lingual focused evaluation. In Proc. Int. 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. LongLoRA: 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-aware 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 of 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):107124, 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 Joey 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 in 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.

Antoine Gonon, Nicolas Brisebarre, Elisa Riccietti, and Rémi Gribonval. A path-norm 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-efficient 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 https://huggingface.co/docs/accelerate/en/ usage_guides/gradient_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.

Stanisław Jastrz˛ ebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in SGD. arXiv:1711.04623 , 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 sharpnessaware 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. Int. 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 11148-11161, 2022.

Dawid Jan Kopiczko, 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 sharpnessaware minimization for scale-invariant learning of deep neural networks. In 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. Int. 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.11692 , 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.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In Proc. Int. Conf. Learning Represention , 2019.

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 on 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. Exploring generalization in deep learning. In Proc. Adv. Neural Info. Processing Systems , volume 30, pages 5947-5956, 2017.

Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro. A PAC-bayesian approach to spectrally-normalized margin bounds for neural networks. In Proc. Int. Conf. Learning Represention , 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 H Zhou. Implicit regularization of gradient flow on one-layer softmax attention. arXiv preprint arXiv:2403.08699 , 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-1642, 2013.

Behrooz Tahmasebi, Ashkan Soleymani, Dara Bahri, Stefanie Jegelka, and Patrick Jaillet. A universal class of sharpness-aware minimization algorithms. arXiv preprint arXiv:2406.03682 , 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, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proc. Adv. Neural Info. Processing Systems , volume 30, 2017.

Ellen M 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.

Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. SuperGLUE: A stickier benchmark for general-purpose 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: Informationtheoretic 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-641, 2019.

Kaiyue Wen, Tengyu Ma, and Z hiyuan Li. How does sharpness-aware minimization minimizes sharpness. In Proc. Int. Conf. Learning Represention , 2023a.

Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. Sharpness minimization algorithms do not only minimize 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 1112-1122, 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 generalization. In Proc. Adv. Neural Info. Processing Systems , volume 33, pages 2958-2969, 2020.

Wenhan Xia, Chengwei Qin, and Elad Hazan. Chain of LoRA: Efficient fine-tuning of language models via residual learning. arXiv preprint arXiv:2401.04151 , 2024.

Qingru Zhang, Minshuo Chen, Alexander Bukharin, Pengcheng 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 preprint arXiv:1810.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 arXiv:2205.01068 , 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.

Supplementary Document for

'Implicit Regularization of Sharpness-Aware Minimization for Scale-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 as 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 is 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 (Jastrz˛ ebski et al., 2017). Large scale experiments also indicate sharpness-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 a large body of literature exploring sharpness for improved generalization. Entropy SGD leverages local entropy in search of a flat valley (Chaudhari et 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 to 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 on the edge-of-stability-like behavior of unnormalized SAMon quadratic problems. Dai et al. (2023) argue that the normalization in SAM, i.e., line 5 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 ( x, y ) = xy , 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.

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 as 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 parameter-efficient finetuning. LoRA (Hu et 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

LoRA via adaptivity (Zhang et al., 2023a), chaining (Xia et al., 2024), aggressive parameter saving (Kopiczko et al., 2024), low-bit training (Dettmers et al., 2023), and modifications for long-sequences (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

where { Q , K , V } are query, key, and value matrices to be optimized. This is a 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 a simplified version of two-layer ReLU neural nets, and its objective can be defined as

This is usually adopted as an example for overparametrization, 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 W 1 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 BatchNorm 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 is adopted in Fig. 1 (a) and Fig. 2 for visualization of SAM's behavior on NOP. We consider a special case of problem (1a), where the goal is to fit (rank-1) matrices by minimizing

where A ∈ R 3 × 3 := diag [0 . 5 , 0 , 0] and N ξ ∈ R 3 × 3 denote the ground truth and Gaussian noise, respectively; and α controls the SNR. Here we choose N ξ := diag [1 . 0 , 0 . 8 , 0 . 5] U ξ , where entries of U ξ are unit Gaussian random variables.

In our simulation of Fig. 1 (a), we set the step size to be η = 10 -4 and the total number of iterations as T = 10 5 for both SGD and SAM. Parameter ρ is chosen as 0 . 1 for SAM. For both algorithms, initialization is x 0 = [0 . 2 , -0 . 1 , 0 . 3] glyph[latticetop] and y 0 = -3 x 0 . 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. 1 (b) for visualization of SAM on OP. A special case of problem (1b) is considered with objective function

where a ∈ R and n ξ ∈ R denote the ground truth and Gaussian noise, respectively. We choose a = 0 . 5 and n ξ 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 = 10 5 for both SGD and SAM. Parameter ρ is set as 0 . 2 for SAM. For both algorithms, initialization is x 0 = [0 . 2 , -0 . 1 , 0 . 3] glyph[latticetop] and y 0 = -3 x 0 .

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 glyph[negationslash]

An optimizer has to determine α for OP despite it does not influence objective value. Hence, scaling is 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

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 t ‖ 2 -‖ y t ‖ 2 | and the dynamics of SAM on both NOP and OP. The implementation of BAR is similar as weight decay in AdamW (Loshchilov and Hutter, 2019).

Here we use nBAR as an example. If ignoring A t in Theorem 2, it can be seen that B t for NOP decreases whenever ‖ g x t ‖ < ‖ g y t ‖ . In other words, the balancedness of SAM is driven by the difference between the gradient norms at x t and y t . nBAR mimics this and triggers balancedness when stochastic gradients g x t and g y t are not balanced; see Alg. 2.

Finally, we illustrate more on the reasons for employing regularization in OP rather than posing ‖ x t ‖ = ‖ y t ‖ as a hard constraint or initializing in a balanced manner, i.e., ‖ x 0 ‖ = ‖ y 0 ‖ . First, it is quite clear that ‖ x ‖ = ‖ y ‖ is a nonconvex set and how to project on such a set is still debatable. Second, the 'symmetry' associated with the scale-invariant problems does not always favor this

0.075

0.050

0.025 0.000

0.025

0.050

0.075

x

0.08

0.06

0.04

0.02

0.00

0.02

0.04

0.06

0.08

0.0

0.4

0.8

1.2

1.6

2.0

2.4

2.8

3.2

3.6

y

constraint. For the purpose of graphical illustration, we consider a 2 -dimensional example f ( x, y ) = 30000( xy -0 . 005) 2 . It is quite clear that the objective is symmetric regarding the line x = -y , which satisfies | x | = | y | ; see Fig. 4. However, it is not hard to see that SGD can never leave x = -y once it reaches this line via a hard constraint or initialized on this line. In other words, directly adding ‖ x ‖ = ‖ y ‖ as a constraint can trap the algorithm at saddle points. This symmetric pattern is even more complicated in high dimension, i.e., symmetry over multiple lines or hyperplanes. Hence, one should be extremely careful about this hard constraint, and regularization is a safer and more practical choice.

B Missing Proofs for NOP

B.1 Proof of Theorem 1

Proof. For notational convenience, we let G t := ∇ f t ( x t y glyph[latticetop] t ) . Then, we have that

Combining these two inequalities, we arrive at

Similarly, we have that

The proof is thus completed.

B.2 Extension to Stochastic Normalized Gradient Descent (SNGD)

Next, we extend Theorem 1 to SNGD, whose updates can be written as

Theorem 4. When applying SNGD (15) on NOP problem (1a) , the limiting flow with η → 0 guarantees that ‖ x t ‖ 2 -‖ y t ‖ 2 = ‖ x 0 ‖ 2 -‖ y 0 ‖ 2 for all t > 0 . In other words, d B t d t = 0 holds.

Proof. For notational convenience, we let G t := ∇ f t ( x t y glyph[latticetop] t ) . Then, we have that

Similarly, we have that

Combining these two inequalities, we arrive at

The proof is thus completed.

B.3 Proof of Theorem 2

Proof. Denote G t = ∇ f t ( x t y glyph[latticetop] t ) and ˜ G t = ∇ f t (˜ x t ˜ y glyph[latticetop] t ) for notational convenience. Following SAM updates in (4) and setting η → 0 , we have that

This gives that

The second term in (16b) is A t in Theorem 2. Next, we give upper bound on |A t | . Using Assumption 1, we have that

where (a) uses the definition of u t ; (b) follows from ‖ ab glyph[latticetop] ‖ = ‖ a ‖‖ b ‖ and the finite convergence assumption. To bound A t , we also have

where the last line also uses the finite convergence. We can bound ρu t | x glyph[latticetop] t ( ˜ G t -G t ) g y t | = O ( ρ 2 L ) in a similar manner. Combining (17) with (16b) gives the bound on |A t | = O ( ρ 2 L ) .

B.4 Proof of Corollary 1

Here, we prove the formal version of Corollary 1.

Corollary 2. Suppose that ‖ g x t ‖ > 0 and ‖ g y t ‖ > 0 and ρ → 0 , then there exists ¯ B t such that the magnitude of B t shrinks whenever |B t | > ¯ B t .

Proof. Without loss of generality, we suppose that B t > 0 , i.e., ‖ x t ‖ > ‖ y t ‖ > 0 . Let ¯ x t and ¯ y t be the scaled version of x t and y t such that ‖ ¯ x t ‖ = ‖ ¯ y t ‖ and ¯ x t ¯ y glyph[latticetop] t = x t y glyph[latticetop] t are satisfied. This suggests that x t = α t ¯ x t and y t = ¯ y t /α t , where α t = √ ‖ x t ‖ / ‖ y t ‖ . Next, we show that whenever B t is large enough, we have that

Since ρ → 0 , we only need to show that for some small glyph[epsilon1] = O ( ρL ) ≥ 0 ,

By the definition of g x t , g y t and ¯ x t , ¯ y t , we have that (19) can be rewritten as

Note that the function h ( z ) := ( az -b/z ) / √ az + b/z is monotonically increasing in z when a, b > 0 and z > 0 as h ′ ( z ) = ( a 2 z +6 ab/z + b 2 /z 3 ) / (2( az + b/z ) 3 / 2 ) > 0 . This implies that h ( z ) > 0 when z > √ b/a , and thus the condition in (20) can be satisfied for glyph[epsilon1] = O ( ρL ) → 0 when α 2 t > ¯ α 2 , where ¯ α 2 := ‖ G t ¯ y t ‖ / ‖ G glyph[latticetop] t ¯ x t ‖ . This condition on α t is equivalent to

Combining everything together, we have that d B t d t < 0 if

The proof is thus completed. We also note that in the case of ρ > 0 , the same condition as (21) can be derived by obtaining the inverse function of h ( z ) evaluated at glyph[epsilon1] = O ( ρL ) , and the corresponding ¯ α ρ and ¯ B ρ t can be defined similarly.

B.5 Extension to LoRA (layer-wise NOP problem)

Let l ∈ { 1 , 2 , . . . , D } be the layer index. Denote f t as the loss function on minibatch M t . To simplify the notation, we also let G t,l := ∇ x t,l y glyph[latticetop] t,l f t ( { x t,l , y t,l } l ) , ˜ G t,l := ∇ ˜ x t,l ˜ y glyph[latticetop] t,l f t ( { ˜ x t,l , ˜ y t,l } l ) , and u t := 1 / √ ∑ D l =1 ( ‖ g x t,l ‖ 2 + ‖ g y t,l ‖ 2 ) . The update of SAM for layer l can be written as

Refined assumption for LoRA. Direct translating Assumption 1 to our multi-layer setting gives

However, the above assumption is loose, and our proof only needs block-wise smoothness, i.e.,

where ∇ l refers to the gradient on x l y glyph[latticetop] l . It can be seen that √ D ˆ L ≥ L , but one can assume that √ D ˆ L ≈ L for intuitive understandings.

Theorem 5. Suppose that block smoothness assumption in (24) holds. Consider the limiting flow of SAM in (22) with η → 0 and a sufficiently small ρ . Let B t,l := 1 2 ( ‖ x t,l ‖ 2 - ‖ y t,l ‖ 2 ) and B t = ∑ D l =1 B t,l . For some |A t | = O ( ρ 2 ˆ L ) , SAM guarantees that

Furthermore, for per layer balancedness it satisfies that for some |A t,l | = O ( ρ 2 ˆ L ) .

Understanding Theorem 5. A t,i and A t are at the same order because of the possible unbalancedness among gradient norms for different layers. Comparing per layer balancedness B t,l with Theorem 2, it can be roughly estimate that the regularization power is O ( √ D ) times smaller in B t,l . This estimation comes from ˆ L ≈ L/ √ D , and the first term is also O ( √ D ) smaller than the same term in Theorem 2. In other words, the regularization on balancedness can be reduced by O ( √ D ) times in LoRA in the worst case, and the worst case comes from gradient unbalancedness among layers.

Proof. Following (22) and setting η → 0 , we have that

This gives that

Proof for (25) . Let A t := ∑ l A t,l . To start with, we have that

Next, based on finite convergence assumption, we have that

where in (a) we use the fact that ρ is chosen small; (b) uses finite convergence assumption and ‖ ab glyph[latticetop] ‖ = ‖ a ‖‖ b ‖ . Using similar arguments, we can bound A t = O ( ρ 2 ˆ L ) .

Proof for (26) . Next, we give upper bound on |A t,l | . Using similar argument as (28), we have that

Using (29), we have that

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 t ‖ 2 -‖ y t ‖ 2 = ‖ x 0 ‖ 2 -‖ y 0 ‖ 2 for all t > 0 . In other words, B t keeps unchanged.

Proof. We consider SGD and NSGD separately.

SGD. It is straightforward to see that

This completes the proof of SGD.

NSGD. The gradient update of NSGD is

Then we have that for NSGD,

This gives the result for SNGD.

C.2 Proof of Theorem 3

To prove this theorem, we first focus on the dynamic of SAM.

Lemma 2. Suppose that Assumption 1 holds. Consider the limiting flow of SAM in (7) with η → 0 . Let B t := 1 2 ( ‖ x t ‖ 2 -‖ y t ‖ 2 ) and ρ be small. Then, for some |A t | = O ( ρ 2 L |B t | ) , SAM guarantees

Proof. For notational convenience, we write f ′ t := f ′ t ( x glyph[latticetop] t y t ) and ˜ f ′ t := f ′ t (˜ x glyph[latticetop] t ˜ y t ) . Using similar arguments as Theorem 2, we have that

Next we bound |A t | . To start with, we have that

Using Assumption 1 and (34), we arrive at

Hence, we arrive at

The proof is thus completed.

Next, the proof of Theorem 3 is provided.

Proof. Lemma 2 has already indicated the concentration of B t towards 0 , if the magnitude of the first term is larger than |A t | . To see this, notice that we can lower bound 2 |B t | / √ ‖ x t ‖ 2 + ‖ y t ‖ 2 by

Hence, long as ρ | f ′ t ( x glyph[latticetop] t y t ) | · C t > O ( ρ 2 L |B t | ) , we have the first term dominating the dynamic of SAM, leading to contraction of B t . This completes the proof to the first part.

Next we prove the second part, which is the lower- and upper- bound on B t . The lower bound can be seen from (36). For the upper bound, we have

Plugging (37) into (33) finishes the proof.

C.3 m -sharpness for OP

m -sharpness is a variant of SAM that is empirically observed to improve generalization, and it is especially useful for distributed training on multiple GPUs (Foret et al., 2021). However, the reason behind the improved performance is not fully understood. (Andriushchenko and Flammarion, 2022) show that m -sharpness is more sparse-promoting for diagonal linear neural networks minimized via a quadratic loss. However, diagonal linear networks are not scale-invariant.

For consistent notation with (7), we use f t ( · ) to denote the loss function on minibatch M t . In m -sharpness, the minibatch M t is divided into m disjoint subsets. Without loss of generality, we also assume that the minibatch is evenly divided. We denote the loss function on each subset as f t,i , i ∈ { 1 , 2 , . . . , m } . Note that we have 1 m ∑ m i =1 f t,i = f t . With these definitions, the update of m -sharpness can be written as

where u t,i := sgn ( f ′ t,i ( x glyph[latticetop] t y t )) / √ ‖ x t ‖ 2 + ‖ y t ‖ 2 . Comparing with the SAM update for OP in (7), the difference is that perturbed gradient is calculated on each f t,i . Next, we analyze the dynamic of SAM with m -sharpness.

Lemma 3. Suppose that Assumption 1 holds. Consider the limiting flow of SAM in (38) with η → 0 . Let B t := 1 2 ( ‖ x t ‖ 2 -‖ y t ‖ 2 ) and ρ be small. Then, for some |A t | = O ( ρ 2 L ) , SAM guarantees that

Proof. For notational convenience, we write f ′ t,i := f ′ t,i ( x glyph[latticetop] t y t ) and ˜ f ′ t,i := f ′ t,i (˜ x glyph[latticetop] t,i ˜ y t,i ) . Then, we have that

Next, using (34) and Assumption 1, we have

Hence, we can bound |A t,i | as

The proof is thus completed by plugging |A t,i | into (40).

C.4 Extension to Layer-wise OP

We start with the notation. Let l ∈ { 1 , 2 , . . . , D } be the layer index. Denote f t as the loss on minibatch M t . Let f ′ t,l := ∇ l f t ( { x glyph[latticetop] t,l y t,l } l ) , i.e., the l -th entry of gradient (w.r.t. the variable x glyph[latticetop] t,l y t,l ), ˜ f ′ t,l := ∇ l f t ( { ˜ x glyph[latticetop] t,l ˜ y t,l } l ) , and u t := 1 / √ ∑ D l =1 | f ′ t,l | 2 [ ‖ x t,l ‖ 2 + ‖ y t,l ‖ 2 ] . The update of SAM for layer l can be written as

Refined assumption for LoRA. Our proof only needs block-wise smoothness, i.e.,

where ∇ l refers to the gradient on x glyph[latticetop] l y l . It can be seen that D ˆ L ≥ L , but one can assume that √ D ˆ L ≈ L for more clear intuition.

Theorem 7. Suppose that block smoothness assumption in (42) holds. Consider the limiting flow of SAM in (41) with η → 0 and a sufficiently small ρ . Let B t,l := 1 2 ( ‖ x t,l ‖ 2 - ‖ y t,l ‖ 2 ) and B max t = max l |B t,l | . For some |A t | = O ( ρ 2 ˆ L B max t ) , SAM guarantees that

Furthermore, for some |A t,l | = O ( ρ 2 ˆ L |B t,l | ) , per layer balancedness satisfies that

Proof. Using a similar derivation as before, we have that

Next, based on (42), we have that

| f ′ t,l -˜ f ′ t,l | ≤ ˆ L ∣ ∣ ˜ x glyph[latticetop] t,l ˜ y t,l -x glyph[latticetop] t,l y t,l ∣ ∣ ≤ ρ ˆ Lu t | f ′ t,l | ( ‖ x t,l ‖ 2 + ‖ y t,l ‖ 2 ) + ρ 2 ˆ Lu 2 t | f ′ t,l | 2 | x glyph[latticetop] t,l y t,l | . Combining these two equations, and applying similar argument as Theorem 5, it is not difficult to arrive at |A t,i | = O ( ρ 2 ˆ L |B t,l | ) and |A t | = O ( ρ 2 ˆ L B max t ) .

C.5 Proof of Lemma 1

Proof. Within W ∗ , the Hessian on ( x , y ) can be calculated as f ′′ ( x glyph[latticetop] y )[ y glyph[latticetop] , x glyph[latticetop] ] glyph[latticetop] [ y glyph[latticetop] , x glyph[latticetop] ] . The largest eigenvalue is f ′′ ( w ) ( ‖ x ‖ 2 + ‖ y ‖ 2 ) . By the AM-GM inequality, it can be seen that the largest eigenvalue is minimized when ‖ x ‖ = ‖ y ‖ , whose balancedness is 0 .

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 as the pretraining dataset as validated through zero-shot performance. All experiments are performed on a server with AMD EPYC 7742 CPUs and NVIDIA GeForce RTX 3090 GPUs each with 24GiB memory. All numerical results from Section 6 report test performance (e.g., accuracy, F1 scores, or BLEU scores) and the standard deviation across multiple runs.

D.1 Details on Datasets

Our evaluations are carried out on commonly-used datasets in the literature.

GLUE benchmark. GLUE is designed to provide a general-purpose evaluation of language understanding (Wang et al., 2019b). Those adopted in 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 et al., 2018)), QQP 3 (question-answering), RTE 4 (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 benchmark for language understanding and is more challenging compared with GLUE. The considered datasets include CB (inference, (De Marneffe et al., 2019)), ReCoRD (multiple-choice question answering (Zhang et al., 2018)), COPA (question answering (Roemmele et al., 2011)). These datasets are 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 BY-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 (V oorhees and Tice, 2000)) and SNLI (inference (Bowman et al., 2015)). Both of them are licensed under CC BY-SA 4.0.

D.2 Details on Language Models

We summarize the adopted language models in our evaluation. All model checkpoints are obtained from HuggingFace.

RoBERTa-large. This is a 355 Mparameter model. The model checkpoint 5 is released under the MIT license.

OPT-1.3B. The model checkpoint 6 is released under a non-commercial license. 7

GPT2-medium. This is a 345 Mparameter model. Its checkpoint 8 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 on classification tasks. The training set contains k = 512 samples per class while the test set has 1000 samples. We also employ prompts for finetuning; where the adopted prompts are the same as those in (Malladi et al., 2023, Table 13). AdamW is adopted as the base optimizer, and hyperparameters are tuned from those in Table 6. 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.

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 as 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 is 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 a summary in Table 9. It further demonstrates the effectiveness of BAR for large scale-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.

D.4 Finetuning with RoBERTa-large

Our implementation is inspired from (Hu et al., 2022) 9 , which is under MIT License. The hyperparameters are chosen the same as provided in its GitHub Repo. AdamW is 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 are ignored when comparing. We consider the GLUE benchmark and report the mismatched accuracy for MNLI, Matthew'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 α = 16 .

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 2 to 15 hours, except for QQP which takes 3 days on our device.

For the hyperparameters of oBAR and nBAR, µ 0 is typically chosen from { 0 . 2 , 0 . 5 , 1 . 0 } ; however, for QQP, a value of 0 . 05 is used. The scheduler is chosen from linear and constant. We also observe that for datasets such as COLA and RTE, setting weight decay as 0 works best for BAR.

D.5 GPT2 medium on WebNLG Challenge

AdamW 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 on our hardware.

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 3 and 4 and numerical experiments in Sections 5 and 6. Due to space limitation, missing proofs and implementation details can be found in the appendix.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in 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 or 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.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are 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 no 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 a 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 (e.g., independence assumptions, noiseless settings, model well-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 a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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 fairness.

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 acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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 a complete (and correct) proof?

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 crossreferenced.

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 experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The details on proposed algorithms can be found in Section 5 and Appendix A.6. Experimental details on setups, datasets, architectures, and hyperparameters can be found in Section 6 and Appendix D. Code can be found in our github repo.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to 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 a 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 a 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 to 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 to replicate the results, access to a hosted model (e.g., in the case of a 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 all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a 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 to 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 be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The code is open-sourced on github. The code gathers details for reproducing our experiments. For example, the exact environment is listed in environment.yml . Instructions on preparation (e.g., data, packages, etc) and commands to use are 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 ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to 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 a 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 as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: Experimental details can be found in Section 6 and Appendix D.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

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 in Appendix D.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

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 calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It 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 a 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 of range (e.g. negative error 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 in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

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 or 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 as well as 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 it 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://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: NeurIPS Code of Ethics is followed.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or 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.

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), fairness 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 not tied to particular applications, let alone deployments. However, if there is a 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 as 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 or 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 a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This paper does not release new data or models.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 to 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 we encourage authors to take this into account and make a 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 use 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 is used and, if 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.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

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/model 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. Crowdsourcing 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 as 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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or 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 not involve crowdsourcing nor research with human subjects.

Depending on 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.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nOriginality\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\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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\n\nLIMITATIONS:\nThe authors sufficiently discuss the limitations of this study, particularly regarding the application range and computational efficiency of BAR.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-01"}

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

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.
2. 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).

2. 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.

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?

2. 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-12

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

STRENGTHS:
- The theoretical analysis provides new insights into SAM's implicit regularization effects by introducing "balancing".
- The proposed BAR method offers a more efficient alternative to SAM for scale-invariant problems by adding explicit regularization. 
- Comprehensive experiments on language model fine-tuning demonstrate practical benefits.

WEAKNESSES:
- The scope is limited to scale-invariant problems, and LoRA types optimization problem.
- Assumption 1 requires  Lipschitz continuous gradient which may not always hold in practice (e.g., ReLU network)

QUESTIONS:
- How to understand the lower bound in Theorem 3, which is not always greater than zero during the optimization. 

- How to choose a good parameter $\rho$?

LIMITATIONS:
Yes, the authors adequately addressed the limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n### 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.\n\nSTRENGTHS:\n### 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\n\nWEAKNESSES:\n### 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\n\nQUESTIONS:\n### 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?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-23"}