LightSAM: Parameter-Agnostic Sharpness-Aware Minimization

The mismatch between theory and empirics: the theoretical results focus on optimization convergence for differentiable loss landscapes, while the empirical results demonstrate generalization performance.

The goal of [1] is to find a minima with flat loss landscape, thereby increasing the generalization ablibity of the model. By forming the above goal as a minimax optimization problem and solving it, [1] proposes the SAM optimizer with an extra perturbation step. Since we aim to make SAM parameter-agnostic, we present the generalization performance in the experiments. This organization is adopted other SAM-related works. The convergence performances are listed in section "B. Experiment Illustration" in the Appendix.

[1] Foret P, Kleiner A, Mobahi H, et al. Sharpness-aware minimization for efficiently improving generalization[J]. arXiv preprint arXiv:2010.01412, 2020.

In Theorem 1, we can see that $A_1$ is related to the term $\eta-2\rho^2 \log(\epsilon)$ which would be negative if $\rho$ is too large.

In fact, $\eta-2\rho^2 \log(\epsilon)$ would always be positive since $\epsilon$ is a real number less than 1 (commonly set to $10^{-3}$ or $10^{-4}$ in the experiments).

While the perturbation radius $\rho$ does not affect the convergence result, how does it influence the model's generalization ability in this paper? Any empirical or theoretical insights provided?

We vary the value of $\rho$ to evaluate the stability of LightSAM, and the results are listed in Table 3,5 and global comment above. We could observe that with different values of $\rho$, LightSAM always performs well and achieves the lowest standard deviation.

In line 228, the authors mentioned "Technical Challenge", but what's the key technique to address this challenge?

We could summarize two points of "Technical Challenge", the first is to decouple the randomness in the numerator and denominator of one term and the second is to bound the terms concerning both $\mathbb{E}||\nabla f(x_t)||^2$ and $\mathbb{E}||\nabla f(w_t)||^2$. To solve the first challenge, we turn to bound term $T_2$ in the proof of Theorem 4 (in the Appendix). The term $T_2$ concerning $\nabla f(x_t)$ and $\nabla f(w_t)$, thus corresponding to the second challenge. To solve it, we turn to bound $\sum ||\nabla f(w_t)||^2(\frac{1}{\sqrt{v_{t-1}}}-\frac{1}{\sqrt{v_t}})$ as shown in equation (23) and propose Lemmas 6 and 7 to bound it. Finally, Lemma 8 establishes the inequality relationship between $\nabla f(x_t,\xi_t)$ and $\nabla f(w_t,\xi_t)$ and finishes the proof.

In Section 3.3, the authors provide the analysis of LightSAM-II (AdaGrad), but without explanation.

The difference between LightSAM-II and LightSAM-I is that the multiplication and division becomes element-wise for vectors, thus, the forms and explanations of the theoretical results for these two approaches are similar. Thus, we omit the further explanation of LightSAM-II to prevent the duplication.

The inequality such as Line 836-837 does not hold. $||\nabla f(w_t,\xi_t)||$ could not be directly bounded through Assumption 2.

Firstly, for the equality at Line 836-837, since the term $||\nabla f(w_t)||$ and $v_{t-1}$ do not have randomness at time $t$, we pull them out of the expectation. Secondly, we could obtain that $\mathbb{E}||\nabla f(w_t,\xi_t)||^2 \leq D_0 + D_1||\nabla f(w_t)||^2$ by Assumption 2, here the expectation operation removes the randomness caused by the stochastic sampling $\xi_t$. In other words, this expectation is taken conditioned on the historical stochastic gradients. Thus, when we take global expectation on the LHS of equation (13), we are actually taking expectation on a conditioned expectation, i.e. $\mathbb{E}[\mathbb{E}[a|b]] = \mathbb{E}[a]$, thereby we finally add the expectation notation on $||\nabla f(w_t)||^2$.

Both two ways could be regarded as adaptive perturbation radius since $\rho_1=\rho/||u_t||$ and $\rho_2=\rho/||g_t||$ are all determined by the gradient information. Whether it's necessary to incorporate the $||u_t||$ or are there any essential differences?

Firstly, the perturbation radius adopted in original SAM $\rho_2=\rho/||g_t||$ uses the gradient normalization (discussed in [1]), not the adaptive perturbation radius. Our adopted perturbation radius $\rho_1=\rho/||g_t||$ is adaptive, the same as previous adaptive optimizers (Adagrad-Norm, Adagrad and Adam). Secondly, AdaSAM [2] uses the perturbation radius $\rho/||g_t||$ and the adaptive learning rate. However, AdaSAM is not proved to be parameter-agnostic as our approach. Thus, we think it is necessary to incorporate $||u_t||$ to make SAM parameter-agnostic.

The theoretical contributions in this work differ fundamentally from existing proofs.

Our approach involves two gradients $\nabla f(x_t)$ and $\nabla f(w_t)$, thus the proof is more complex than [3]. For example, the proof of [3] needs to bound the term $\mathbb{E}||\nabla f(x_t)||^2(\frac{1}{\sqrt{v_{t-1}}}-\frac{1}{\sqrt{v_t}})$, while we need to bound the term $\mathbb{E}||\nabla f(w_t)||^2(\frac{1}{\sqrt{v_{t-1}}}-\frac{1}{\sqrt{v_t}})$. As a consequence, we propose Lemmas 1 and 2 to bound it. The analysis in Lemmas 1 and 2 is non-trivial, please refer to the Appendix for more details.

Including additional experiments on tasks beyond classification, such as natural language processing or reinforcement learning, could provide a more comprehensive evaluation.

Thanks for your suggestion. We supplement the experiment on GLUE benchmark and list the results in the global comment above, please refer to it.

There are some typographical errors in the proofs that affect readability and accuracy.

Thanks for your comment. We have corrected the error.

[1] Dai, Yan, Kwangjun Ahn, and Suvrit Sra. "The crucial role of normalization in sharpness-aware minimization." Advances in Neural Information Processing Systems 36 (2024).

[2] Sun, Hao, et al. "Adasam: Boosting sharpness-aware minimization with adaptive learning rate and momentum for training deep neural networks." Neural Networks 169 (2024): 506-519.

[3] Wang, Bohan, et al. "Convergence of adagrad for non-convex objectives: Simple proofs and relaxed assumptions." The Thirty Sixth Annual Conference on Learning Theory. PMLR, 2023.

Can it be understood that the proposed method only gives marginal improvements in table 4 and 5?

Yes, the improvements in the experiments of this work is marginal. However, our goal is to propose a SAM-type algorithm that achieves the parameter-agnostic property, as mentioned on the top of page 2 "Can we make SAM parameter-agnostic?", not to propose an approach that outperforms previous algorithms in the experiments. The theoretical and experimental results have verified that we realize this goal.

The convergence plots for a range of $\rho$ or $\eta$ values can be helpful to provide a more comprehensive ablation study. Compare the three proposed versions with other optimizers across different hyperparameter settings, rather than a single setup of hyperparameters in table 5.

In fact, we have already evaluated our proposed algorithm across different hyper-parameter settings. In Table 3, we adopt four hyper-parameter settings ($\rho$,2$\rho$)*($\eta$,2$\eta$), where $\rho$ and $\eta$ are the values adopted to lead to results in Table 2. In Table 5, the nine results for one baselines are obtained under nine different hyper-parameter settings (0.5$\rho$,$\rho$,2$\rho$)*(0.5$\eta$,$\eta$,2$\eta$), where $\rho$ and $\eta$ are the values adopted to lead to results in Table 4. A clearer presentation could be referred to the Table in the global comment above. We have mentioned these points in the paragraph "Sensitivity to hyper-parameters".

Can other optimizers also get such parameter agnostic properties, aside of the presented 3? What’s essential characteristics of those optimizers contributing to such agnostic properties in convergence?

We have listed some of the existed parameter-agnostic optimizers in the section "Related Work". Among them, the techniques of Normalized SGDM and adaptive optimizers (Adagrad-Norm, Adagrad and Adam) are classical. Since the technique of adaptive learning rate is similar to normalization (upper bounding the gradient norm), we think the gradient normalization or bounding the weight update vector may be essential for the agnostic property.

Why name is as LightSAM?

LightSAM is proved to converge under any hyper-parameter settings, therby reducing the heavy workload of tuning parameters. Thus, we use "light" to name it.

From the experiments, the proposed algorithm performs on par with the AdaSam baseline, while the AdaSam baseline is claimed to have a better convergence rate.

The goal of this work is to propose a SAM-type algorithm with adaptive hyper-parameters and verify the parameter-agnostic property of it, not to propose an approach that outperforms previous algorithms in the experiments. The theoretical analysis and experimental results have already veried that LightSAM is parameter-agnostic. The theoretical contribution of this work is to prove the convergence rate of LightSAM under nearly the weakest assumptions, and do not obtain a better convergence rate than previous algorithms. Since LightSAM and AdaSAM both adopt the adaptive learning rate in the gradient descent step, it is trivial that they achieve similar performances.

How should one decide the initial parameter of the LightSAM?

Since LightSAM adopts the adaptive learning rate in the gradient descent step, we set $\eta$ the same as the value commonly used by Adam optimizer, such as $1e-4$ in ImageNet experiment and $1e-5$ in GLUE experiment (these values are suggested in the experiments in previous studies). Then, the weight perturbation step adopts the same type of hyper-parameter as the gradient descent step, thus we set $\rho$ the same as $\eta$.