Efficient Sharpness-Aware Minimization for Molecular Graph Transformer Models

We appreciate a lot for your careful review, and would like to provide responses to your mentioned questions and weaknesses one by one.

###[W1 - What will happen if $ε_{t+1}$ becomes moving average of $w_{t}$? — It can cause gradient explosion.]

Thank you for your insightful observation. Utilizing $|w_{t}|$ without normalization could make $ε_{t+1}$ grow excessively large, which might trigger a gradient explosion. Such an occurrence would critically disrupt the training process, rendering it not only ineffective but also at risk of divergence. In Eq. (3), the choice to use $w_{t}/|w_{t}|$ for updating $ε_{t+1}$ is deliberate and crucial for the stability and effectiveness of the update process. This formulation ensures that the update direction of $ε_{t+1}$ is determined by $w_{t}$ while its magnitude is controlled by the factor (1-β). By normalizing $w_{t}$ with its magnitude $|w_{t}|$ , we prevent the scale of $w_{t}$ from disproportionately influencing the update of $ε_{t+1}$. If we were to use $w_{t}$ directly for updating $ε_{t+1}$, the absence of normalization would lead to significant issues. Our experiments have demonstrated that $|w_{t}|$ >> $|ε_{t}|$, further underscoring the necessity of this normalization step.

###[W2 - Does Cosine schedulers or inverse square root schedulers help? — They have some effect, but StepLR is the best.]

Thank you for your question regarding the choice of the StepLR scheduler for the ρ-schedulers in our model. The decision to use StepLR was primarily guided by its simplicity and effectiveness in controlling the learning rate in discrete steps. This allows for a straightforward reduction of the ρ size at predetermined epochs, which we found to be beneficial for the steady convergence to a flat region of our model's parameters. In response to your suggestion, we have added the following two sets of experiments:

CosineLR: $ρ_{new}$=$\frac{ρ_{initial}}{2}*(1+cos(\frac{epoch·π}{EPOCH}))$

InverseSquareLR: $ρ_{new}$=$\frac{ρ_{initial}}{\sqrt{max(epoch,k)}}$

where epoch is the current training epoch, EPOCH is total training epoch, and k=1, $ρ_{initial}$ = 0.05.

We acknowledge that other approaches, such as cosine schedulers or inverse square root schedulers, could also be effective. While these alternative schedulers present intriguing possibilities, our initial experiments with StepLR showed the most promising results, which led to its selection for the current study. However, we recognize the value in exploring these alternative scheduling methods in future work, as they may offer different advantages in specific training contexts or with certain types of datasets.

###[Q1 - Table 1 does not fit into the paper properly. — Fixed, and thank you for the close read!]

###[Q2 - $sign(ε_{t})|ε_{t}|$ can be replaced by $ε_{t}$! — Yes, you are right and we fixed it.]

###[Q3 - Missing references on SAM variants. — We have added them.]

Thank you for highlighting the omission of references related to SAM variants in our manuscript. In our revised manuscript, we will include and discuss the two SAM variants [1, 2] in the related work.

References:

[1] Du, Jiawei, Daquan Zhou, Jiashi Feng, Vincent YF Tan, and Joey Tianyi Zhou. "Sharpness-Aware Training for Free." arXiv preprint arXiv:2205.14083 (2022).

[2] Li, Bingcong, and Georgios B. Giannakis. "Enhancing Sharpness-Aware Optimization Through Variance Suppression." arXiv preprint arXiv:2309.15639 (2023).

Thank you for the many valuable comments on writing, we have also been corrected! Notably, the SS-SAM you mention is the same paper as the RST compared in our paper.

###[W1.1 - The colors in Fig.1 are inconsistent — Fixed, and thank you for the close read! ]

###[W1.2 - Add GraphSAM in Fig.1 — We have added GraphSAM in Fig.1.]

###[W1.3 and W1.4 - Discuss existing efficient methods, e.g., AE-SAM and SS-SAM. — We have discussted and compared them in Appendix.]

Indeed, we have provided a detailed comparison of these efficient SAM versions, including AE-SAM and SS-SAM(RST), in Appendix A.4 and A.5. We specifically discuss why these methods are less suited for molecular graph datasets and how GraphSAM is uniquely designed to address the challenges inherent to such data.

For AE-SAM, the trend of its squared stochastic gradient norms in CV domain is not consistent with the molecule data. This leads to a decrease in its sensitivity, which is not as good as its performance in CV, but works better than other methods.

For SS-SAM, the advantage is efficiency; but the disadvantage is obvious: random decisions about whether or not to update the gradient can lead to performance reduction. Its performance in CV literature is also poor.

###[W1.5 - sub-caption for each subfigure. — We put the sub-caption for each subfigure uniformly into the description.]

###[W1.6 - Table 1: also compare with GROVER+AE-SAM/SS-SAM. — We have compared with CoMPT and Grover (+ every efficient SAM) in detail in Table 2 and Table 7.]

###[W1.7 - Description of AE-SAM is incorrect. — Fixed, and thank you for the close read!]

###[W2 - GraphSAM can also be used in CV? — Yes, but it does not perform ideally.]

Sorry, GraphSAM doesn't perform well in the CV domain. It is experimentally found that GraphSAM is a good approximation for graph transformers rather than in CV domain. Although our proposed method is designed based on the empirical observations, actually, Observations 1 and 2 are presented only in the graph transformers, while they do not appear in those of CV domain [1]. In CV, the model's perturbation gradient variations are often consistent with updating gradients, which is not in line with the phenomenons on the molecular graphs.

To further address the concern and provide holistic discussion, we add the following experiments to compare the performance of lookSAM [1], AE-SAM and GraphSAM in the CV and molecule datasets, respectively. LookSAM and GraphSAM are efficient SAMs designed based on the gradient varification patterns of models in their respective domains, and AE-SAM is an adaptively-updating efficient method accoding to the comparison evaluation between gradient norm and a pre-defined threshold.

It is observed LookSAM has better results on the image datasets instead on the molecular graphs investigated in this work. On the contrary, our GraphSAM works on the molecules but leads to the worst performances on the image benchmarks. That is because model gradient variations are diverse accross the different domains. It is challenging to transfer the efficient SAMs designed based on the specific gradient patterns. AE-SAM achieves an acceptable performance on the molecular graphs since the perturbation gradient norms of graph transformers are monitored to inform the necessarity of gradient re-computation. But it is not as good as GraphSAM where we accurately fit the perturbation gradients at each step. In summary, GraphSAM is optimized particularly for the graph transformers based on the empirical observations of gradient variations.

References:

[1] Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, Yang You. Towards efficient and scalable sharpness-aware minimization. CVPR 2022.

We appreciate a lot for your careful review, and would like to provide responses to your mentioned questions and weaknesses one by one.

###[W1 - Is the differing changes of gradient a general phenomenon? — Yes, it is.]

The observation of differing changes in the perturbation gradient and the updating gradient throughout the training process is indeed a general phenomenon. To further substantiate this, we have enriched the Figure 8 in Appendix 8 with additional case studies illustrating these variations. They are consistent with the conclusions in Observation 1.

###[W2 - Efficiency is a crucial contribution, add experiments to illustrate — Agreed, we have compared the efficiency in Table 2 and Table 7.]

We have included comprehensive efficiency comparisons in Table 2 of the main paper and Table 7 of Appendix A.5. These tables specifically highlight the efficiency aspects of GraphSAM, offering a balanced view of both its performance and computational efficiency.

###[M1 - Missing space — Fixed, and thank you for the close read! ]

###[M2 - Counter-intuitive bold highlighting — Sure, we have changed it]

In the revised manuscript, we maintain the bold highlighting for the best optimizer of graph transformer per task, but also add 'underline' as an additional marker to further emphasize the top-performing model in each category.

###[Q2 - How the behavior of the graph transformer changes if trained with GraphSAM? — When the model trained with GraphSAM, nodes' attention score matrix average entropy of same molecular graph is higher.]

This is a wonderful suggestion. Regarding the attention matrix over BBBP dataset, we observe that models optimized with Adam, GraphSAM, and SAM obtain average entropies of 2.4130, 2.7239, and 2.8382, respectively. That means the application of SAM approaches can smooth the attention scores to avoid the overfitting. We have added the attention heat map analysis in appendix. In addition, Figure 5 at main paper shows that incorporating GraphSAM leads to smoother training and testing loss curves and reduces the gap between them. This smoother convergence suggests that GraphSAM helps learning a more robust and stable representation, which indicates a more evenly distributed attention across the graph's features.

###[Q3 - Compare GraphSAM with other efficient SAM derivates? — Yes, we have compared with them, and GraphSAM is both efficient and high performing.]

We have compared with other efficient SAMs (e.g., LookSAM, AE-SAM, and RST) in Table 2 of the main paper and Table 7 of Appendix A.5. We observe GraphSAM is both efficient and high-performaning on molecular graphs, while the other efficient SAMs deliver poor generalization results. The detials can be checked in paper.

###[Q4 - Is The proportionality in Theorem 2 not merely caused by the first-order Taylor approximation? How do the authors know that the linear term is a sufficient approximation? — Yes, we leverage the first-order Taylor approximation, but it is accurate enouhgh.]

The proportionality in Theorem 2 is indeed analyzed mainly according to the first-order Taylor approximation. We are aware that although the higher-order terms could offer a more nuanced view, the first-order term often dominates in practical scenarios, making it a sufficient approximation for the purpose of our analysis. Take the second-order Taylor expansion as an example:

$L_{G}$(θ+ϵ̂) ≈$L_G(θ)$ + ϵ̂$\nabla_\theta$ $L_G(θ)$ +$\frac{1}{2} ϵ̂^{2}H(θ)$,

where $\mathbf{H}(\theta)$ is a Hessian matrix (i.e., a matrix of second-order derivatives) of the function with respect to $\theta$. $\hat{\epsilon}$ is a mapped perturbation gradient obtained by Eq. (2) and $||\hat{\epsilon}||_{2}$ is pretty small, which is often at the scale $10^{-3}$. Therefore the second-order expansion term as well as the laters can be omitted. Taylor's first-order expansion can be sufficient approximation.

We appreciate a lot for your careful review, and would like to provide responses to your mentioned questions and weaknesses one by one.

###[W1 & Q1 - Is GraphSAM only a good approximation for graph transformers? How is GraphSAM working, e.g., in the image domain? — Yes, GraphSAM only works good at molecular graph domain.]

It is experimentally found that GraphSAM is a good approximation for graph transformers rather than in CV domain. Although our proposed method is designed based on the empirical observations, actually, Observations 1 and 2 are presented only in the graph transformers, while they do not appear in those of CV domain [1]. In CV, the model's perturbation gradient variations are often consistent with updating gradients, which is not in line with the phenomenons on the molecular graphs.

To further address the concern and provide holistic discussion, we add the following experiments to compare the performance of lookSAM [1], AE-SAM and GraphSAM in the CV and molecule datasets, respectively. LookSAM and GraphSAM are efficient SAMs designed based on the gradient varification patterns of models in their respective domains, and AE-SAM is an adaptively-updating efficient method accoding to the comparison evaluation between gradient norm and a pre-defined threshold.

It is observed LookSAM has better results on the image datasets instead on the molecular graphs investigated in this work. On the contrary, our GraphSAM works on the molecules but leads to the worst performances on the image benchmarks. That is because model gradient variations are diverse accross the different domains. It is challenging to transfer the efficient SAMs designed based on the specific gradient patterns. AE-SAM achieves an acceptable performance on the molecular graphs since the perturbation gradient norms of graph transformers are monitored to inform the necessarity of gradient re-computation. But it is not as good as GraphSAM where we accurately fit the perturbation gradients at each step. In summary, GraphSAM is optimized particularly for the graph transformers based on the empirical observations of gradient variations.

###[W2 - The conclustion of Theorem 1 is not clear. — Will add to the following new explanations.]

Thank you for your constructive feedback regarding the theoretical aspects of our paper. We hope the following revisement can mitigate your concerns. The revised sentence is: "Recalling the min-max optimization problem of SAM in Eq. (1), if we replace the inner maximum objective from $L_{G}(θ + ε^{S})$ to $L_{G}(θ + ε^{G})$, the graph transformer is motivated to smooth a worse neighborhood loss in the loss landscape. In other words, the proposed GraphSAM aims to minimize a rougher neighborhood loss, whose value is theoretically larger than that of SAM, and obtain a smoother landscape associated with the desired generalization."

###[W3 - Proof of Theorem 1 relies on empirical observations and strong assumptions. — Agreed, we have now changed it.]

Thanks for your suggestions, and we have uniformly changed 'Theorem' to 'Conjecture'. Conjecture is a statement or conclusion based on insufficient evidence, rather than through rigorous proof. We rewrite Theorem 1 to include the experimental assumption that supports our conclusion.

The new Conjecture 1 is: Let $ϵ̂_{S}$ and $ϵ̂_{G}$ denote the perturbation weights of SAM and GraphSAM, respectively, where we ignore the subscript of t for the simple representation. Suppose that ω/|ω| >> ϵ ,as empirically discussed in Observation 1, and $|ϵ̂_{S}|$ < $|ϵ̂_{G}|$ for ρ> 0, designating $ϵ̂_{S}$ as the ground-truth. We have:

L_{G}(θ+$ϵ̂_{S}$) ≤ L_{G}(θ+$ϵ̂_{G}$)

References:

[1] Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, Yang You. Towards efficient and scalable sharpness-aware minimization. CVPR 2022.

Dear Reviewer QvPb,

We sincerely appreciate your thoughtful feedback. In response to your suggestion, we elaborate in Observation 1 of the main text on why GraphSAM is only a good approximation specifically for molecular tasks. Meanwhile, we compare GraphSAM with other efficient SAMs designed for the CV domain and analyze their limitations in the molecular graph domain in detail as shown in Appendix A.4.

Best regards,

The Authors of ’Efficient Sharpness-Aware Minimization for Molecular Graph Transformer Models‘