Paper: Efficient Sharpness-Aware Minimization for Molecular Graph Transformer Models (ICLR-2024)

Reviewer comment (Theory):
Can the authors elaborate more on equation (3)? In particular, are there any specific reasons for choosing $\epsilon_t$ as the moving average of $w_t/\|w_t \|$? What will happen if $\epsilon_t$ becomes moving average of $w_t$?

Author response:
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.

Reviewer comment (Experiments):
What are the specific reasons behind the choice of StepLR for the $\rho$-schedulers? Does other approaches, such as cosine schedulers or inverse square root schedulers help?

Author response:
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.

|GROVER+GraphSAM | BBBP(ROC-AUC↑) | Tox21 (ROC-AUC↑)| Sider (ROC-AUC↑)| Clintox(ROC-AUC↑)|ESOL(RMSE↓) | Lip(RMSE↓)|
|   :----: |   :---: |    :---: |  :---: |    :---: |    :---: |    :---: |
| Original model| 0.917 | 0.822| 0.649 |0.853|0.639 | 0.671|
| +StepLR| **0.928** | **0.846**| **0.665** |**0.866**|**0.625** | **0.654**|
| +CosineLR | 0.921 | 0.828| 0.651|0.856| 0.633|0.661
| +InverseSquareLR | 0.924 | 0.836 |0.654|0.858|0.628|0.665

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.

Reviewer comment (Presentation):
Table 1 does not fit into the paper properly.

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

Reviewer comment (Novelty):
Missing references on SAM variants. 

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

Author response:
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).

Reviewer comment (Presentation):
The colors in Figure 1 are inconsistent with the observations: "compared to the base optimizer of Adam (the blue bar)", but the blue is LookSAM? also, SAM (cyan)?

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

Reviewer comment (Novelty):
Related works: "most of the work still ignores the fact of SAM’s double overhead (Damian et al., 2021; Kwon et al., 2021;Wang et al., 2022)." I think some methods have attempted to mitigate this issue, e.g., AE-SAM, SS-SAM, better to discuss them here.

Author response:
###**[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.

Reviewer comment (Novelty):
Section 3, "SAM consumes double overhead due to the extra computation of perturbation gradient compared with the base optimizer." Existing methods (AE-SAM, SS-SAM) have tried to mitigate this issue; better to discuss them here, and explain why they cannot be used for graph datasets, as this is also the motivation for the proposed GraphSAM.

Author response:
###**[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.

Reviewer comment (Experiments):
Table 1: also compare with GROVER+AE-SAM/SS-SAM.

Author response:
###**[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.]**

Reviewer comment (Writing):
"AE-SAM and RST periodically compute the update gradient by different strategies": AE-SAM is an adaptive strategy, not a periodic strategy.

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

Reviewer comment (Experiments):
It seems that the proposed GraphSAM is general, not limited to graph transformers, and can also be used in computer vision. Thus, it is better to conduct some experiments on computer vision.

Author response:
###**[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.

|ResNet-18 (CV) | CIFAR-10 | CIFAR-100|
|   :----: |   :---: |    :---: |  
| +SGD| 95.41 | 78.21|
| +SAM| **96.53** | 80.18|
| +LookSAM | 96.28 | 79.91|
| +AE-SAM |96.49 | **80.31**|
| +GraphSAM | 95.86 | 78.69 |


|CoMPT (Graph) | Tox21 | Sider| ClinTox|
|   :----: |   :---: |    :---: |   :---: |
| +Adam| 82.81 | 62.17| 91.41|
| +SAM| 83.96 | 64.33| 92.73|
| +LookSAM | 82.55| 62.54| 91.64|
| +AE-SAM | 83.33| 63.16| 91.92|
| +GraphSAM | **84.11** | **64.58** |**93.78** |


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.

Reviewer comment (Experiments):
Ablation study for hyperparameters $\gamma, \lambda, \beta$.

Author response:
###**[W3 - Ablation study for hyperparameters β, γ,λ ? — We have included experiments for hyperparameter analysis in the Appendix.]**

Thank you for your inquiry regarding the ablation study for hyperparameters. Recognizing the length constraints of the main paper, we have included detailed experiments on hyperparameter ablation in Appendix with Table 3 for β, Table 9 for γ and Fig.7 for λ. These sections in the appendix offer a comprehensive view of how various hyperparameters influence the performance and behavior of our proposed method.

Reviewer comment (Experiments):
In Observation 1, the author notes the differing changes of the perturbation gradient and the updating gradient throughout the training process. I wonder, is this a general phenomenon? Could more case studies be provided to illustrate this further?

Author response:
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.

Reviewer comment (Novelty):
The relation to graph machine learning is a bit obscure: is GraphSAM only a good approximation for graph transformers? After all, GraphSAM does not rely on anything graph-related besides the observations made with SAN + graph transformers. Although it is fine to focus on molecular data, the paper would benefit significantly from a holistic discussion (and some rudimentary experiments).

Author response:
###**[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.

|ResNet-18 (CV) | CIFAR-10 | CIFAR-100|
|   :----: |   :---: |    :---: |  
| +SGD| 95.41 | 78.21|
| +SAM| **96.53** | 80.18|
| +LookSAM | 96.28 | 79.91|
| +AE-SAM | 96.49 | **80.31**|
| +GraphSAM | 95.86 | 78.69 |


|CoMPT (Graph) | Tox21 | Sider| ClinTox|
|   :----: |   :---: |    :---: |   :---: |
| +Adam| 82.81 | 62.17| 91.41|
| +SAM| 83.96 | 64.33| 92.73|
| +LookSAM | 82.55| 62.54| 91.64|
| +AE-SAM | 83.33| 63.16| 91.92|
| +GraphSAM | **84.11** | **64.58** |**93.78** |


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.

Reviewer comment (Theory):
Proof of Theorem 1 relies on empirical observations and (from my perspective) strong assumptions. However, the authors do sufficiently not discuss this in the main part and, thus, the current presentation is misleading. In other words, these assumptions, etc. should be made explicit in a rather prominent way in the main part (e.g. "and after ||ω0 ||2 experimental analysis, the updating gradient ω ≫ ε " is not stated clearly in main part). I think it would be better to drop the term "Theorem" and rather give some intuitive, mathematically motivated explanation.

Author response:
###**[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.

Reviewer comment (Presentation):
The bold highlighting in Table 1 is counter-intuitive. Perhaps add an additional marker to highlight the best model per task.

Author response:
###[**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.

Reviewer comment (Evaluation):
Are there any observations of how the behavior of the graph transformer changes if trained with GraphSAM? For example, do the attention scores then align better with the graph connectivity, or do they become smoother (higher entropy)?

Author response:
###[**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.

Reviewer comment (Experiments):
Is it possible to compare GraphSAM with other efficient SAM derivates besides SAM-k (e.g., see related work section)?

Author response:
###[**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.

Reviewer comment (Theory):
Is the proportionality in Theorem 2 not merely caused by the first-order Taylor approximation (Eq 6 in A1.1)? How do the authors know that the linear term is a sufficient approximation?

Author response:
###[**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.