Implicit Regularization of Sharpness-Aware Minimization for Scale-Invariant Problems

Summary

This paper analyzes the implicit regularization effect of Sharpness-Aware Minimization (SAM), focusing specifically on scale-invariant problems. While existing research emphasizes sharpness, this study introduces a new concept called Balancedness, demonstrating both theoretically and empirically that SAM promotes Balancedness. Additionally, the proposed Balancedness-Aware Regularization (BAR) is shown to significantly improve computational efficiency while achieving superior performance in fine-tuning language models using LoRA.

Strengths

Originality - Introducing a new concept of Balancedness and reinterpreting the implicit regularization effect of SAM from a new perspective. Quality - Consistent theoretical analysis and empirical validation with clear results. Clarity - Clear presentation with a structure that is easy for readers to understand. Significance - Revealing a new aspect of SAM's adaptability to data anomalies and proposing a computationally efficient BAR. Soundness - 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. Contribution - 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. - Experiments using various models confirm that the proposed BAR exhibits superior performance in fine-tuning with LoRA compared to traditional SAM and SGD. - 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. [1] https://arxiv.org/abs/2210.14199

Weaknesses

- While LoRA training is performed with AdamW in the original paper, demonstrating that SAM or BAR is more effective than these would make the paper more solid. - Code such as `eval.py` was not included in the supplementary materials, and its release is expected.

Questions

- As cited in the Appendix, it is known that SAM promotes the acquisition of low-rank features [2]. However, it was pointed out that explicit regularization for acquiring low-rank features did not lead to improved generalization ability. My curiousity is about the relationship between Balancedness and low-rank features, and whether there is any connection to sharpness. - Another question is how optimization methods commonly used in LLM training, such as AdamW, affect Balancedness. [2] https://arxiv.org/abs/2305.16292

Rating

8

Confidence

3

Soundness

4

Presentation

4

Contribution

4

Limitations

The authors sufficiently discuss the limitations of this study, particularly regarding the application range and computational efficiency of BAR.

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?

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

4

Limitations

The limitation is clearly addressed.

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$?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes, the authors adequately addressed the limitations.

Summary

### Summary: This 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. Moreover, they show that outliers strongly affect balancedness, which helps explain SAM behavior. Finally, they propose the BAR algorithm and conclude with several experiments.

Strengths

### Pros: - the studied problem is very important, and the result can be impactful - the balancedness is completely new and well-motivated for SAM (previously some works studied related quantities but just for SGD) - having several experiments

Weaknesses

### Cons: - 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 - some ambiguities about Theorem 1

Questions

### Questions/Comments: This 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. Some comments: - 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. - 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: - A Universal Class of Sharpness-Aware Minimization Algorithms, ICML 2024 - Do authors claim that Theorem 1 is the original result of this paper? I think this is known from previous works on the implicit biases of SGD, but it looks like the paper is claiming the results without references. See, for example, Equation 2.5 in this paper (already cited): - Ji, Ziwei, and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. ICLR. Please explain whether the theorem follows from the previous results or if there is a major difference between them. - 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. - 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?

Rating

6

Confidence

5

Soundness

3

Presentation

2

Contribution

4

Limitations

N/A

Official Comment by Reviewer t2ze

I thank the authors for their detailed response. I am satisfied with the answers provided, and I would like to keep my positive score.