Efficient Sharpness-Aware Minimization for Molecular Graph Transformer Models
Authors: Yili Wang
Kaixiong Zhou
Ninghao Liu
Ying Wang
Xin Wang
Venue: ICLR 2024 poster (2024)
Keywords: Sharpness-Aware Minimization
Molecular Graph
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: Sharpness-aware minimization (SAM) has received increasing attention in computer vision since it can effectively eliminate the sharp local minima from the training trajectory and mitigate generalization degradation. However, SAM requires two sequential gradient computations during the optimization of each step: one to obtain the perturbation gradient and the other to obtain the updating gradient. 
Compared with the base optimizer (e.g., Adam), SAM doubles the time overhead due to the additional perturbation gradient. By dissecting the theory of SAM and observing the training gradient of the molecular graph transformer, we propose a new algorithm named GraphSAM,  which reduces the training cost of SAM and improves the generalization performance of graph transformer models. 
There are two key factors that contribute to this result: (i) \textit{gradient approximation}: we use the updating gradient of the previous step to approximate the perturbation gradient at the intermediate steps smoothly (\textbf{increases efficiency}); (ii) \textit{loss landscape approximation}: we theoretically prove that the loss landscape of GraphSAM is limited to a small range centered on the expected loss of SAM (\textbf{guarantees generalization performance}). The extensive experiments on six datasets with different tasks demonstrate the superiority of GraphSAM, especially in optimizing the model update process.

Published as a conference paper at ICLR 2024
EFFICIENT SHARPNESS- AWARE MINIMIZATION FOR MOLECULAR GRAPH TRANSFORMER MODELS
Yili Wang', Kaixiong Zhou?, Ninghao Liu", Ying Wang', Xin Wang
School of Artificial Intelligence, Jilin University, China
2 PInstitute for Medical Engineering & Science, Massachusetts Institute of Technology, USA School of Computing, University of Georgia, USA
*College of Computer Science and Technology. Jilin University, China
wangy121@man jlu. edu. cn, xz34@mit edu, ninghao. 1 1iuuuge uga. edu, (wangying2010, xinwang) @ [ iel.luee. edu. en
ABSTRACT
Sharpness-awar minimization (SAM) has received increasing attention in com- puter vision since can effectively eliminate the sharp local minima from the training trajectory and mitigate generalization degradation. However, SAM re- quires two sequential gradient computations during the optimization of each step: one to obtain the perturbation gradient and the other to obtain the updating gradient. Compared with the base optimizer (e.g., Adam). SAM doubles the time overhead due to the additional perturbation gradient. By dissecting the theory of SAM and observing the training gradient of the molecular graph transformer, we propose a algorithm named GraphSAM, which reduces the training cost of SAM and improves the generalization performance of graph transformer models. There are key factors that contribute to this result: (i) gradient approximation: updating gradient of the previous step to approximate the perturbation gradient at the intermediate steps smoothly (increases efficiency); (ii) loss landscape approxi- mation: we theoretically prove that the loss landscape of GraphSAM is limited to a small range centered on the expected loss of SAM (guarantees generalization per- formance). The extensive experiments on six datasets with different tasks demon- strate the superiority of GraphSAM, especially optimizing the model update process. The code is in: https://githbb.com Www.t////araaSS///reeeee//////
use the
INTRODUCTION
Biochemical molecular property prediction is one of the essential tasks for many applications, including drug discovery (Withnall et al.] 22000] Wu et al. 2018: Ye al. 2022 and molecular fingerprint design (Honda et al. 2019 Kearnes et al. 2016). Motivated by the recent progress in natural language processing and computer vision, graph transformer models have shown promising results for molecular property prediction by treating molecular structures as graphs (Rong et al., 2020: Chen al., 2021a). The transformers learn the global interaction of every node to capture the underlying structure, improving the molccular property classification performance.
Without the hand-crafted features inductive biases encoded in the neural architecture, it is widely observed that the transformer is prone to converge to sharp local minima, where the loss value changes quickly in the neighborhood around model weights (Du et al. 2021 Zhou al.) [2022b] Chen et al.) 2021b Andriushchenko & Flammarion 2022, Zhou et al. 2021b). The sharp local minima are highly correlated with significant generalization errors in various domains. To mitigate the sharpness, existing studies on transformers for molecular property prediction use large-scale pre-training, while the pre-training requires extensive resources and expert knowledge in constructing the informative self-supervised learning tasks (Rong et al. 2020 Chithrananda et al. 2020 Ying et al. 2021; Liu et al. 2021b). These laborious requirements prevent engineers from directly applying the transformer for any downstream application.
Recently, sharpness aware minimization (SAM) (Foret et al. 2020 has been proposed to explicitly smooth the sharp local minima during model training like pre-training. Nevertheless, SAM requires two forward and backward propagations at each step: one to obtain the worst-case adversarial gradient
"oorrespnndence Xin Wang

---

Published as a conference paper at ICLR 2024
for perturbing weights named perturbation gradient and the other to obtain the updating gradient for the training model. Compared with the base optimizer, SAM doubles the time overhead due to the extra computation of the perturbation gradient. In order improve the efficiency of model training, some efficient SAM variants (AE-SAM, RST, ESAM, etc.) (Jiang et al., 2023] /Zao al., 2022a Du al. 2021) have been invented. They seek to enhance efficiency through "generic" techniques (c.g., Stochastic judgment gradient calculation) that can be employed with "any" model and have been successful in Computer Vision. Whether they still maintain a high standard when they step out of their comfort area, on the contrary, the answer is no. As shown in Fig. [ com- pared to the base optimizer of Adam (the second bar), the performance of the pre-training-fres Grover model with SAM has a significant improvement (the third bar). How- ever, other efficient SAM variants have minimal even counterproductive performance. The reason is that they overlooked omitted the principles of SAM's success: gradient direction and gradient size. This is the core aspect that motivates the current work.
Negittraingg
92-
+AE-SAM
+LackSAN
86
80
B0BP
ClinTox
Tex21
Figure 1: The results of GROVER with
different strategies three datasets.
In this paper, we aim to improve the efficiency of SAM while maintaining generalization perfor- mance of graph transformer models on molecular property prediction tasks. One of the existing straightforward solutions to speed up SAM is to compute the perturbation gradient periodically, c.g., RST (Zhao al. 2022a). which unfortunately leads to poor empirical results. In the experiments, we observe the perturbation gradient direction is similar to the indispensable updating gradient direction from the last step. Based on this discovery, we develop GraphSAM to reuse the updating gradient of the previous step and approximate the perturbation gradient with free lunch. Specifically, GraphS contains two sequential gradient computations at each step: one to obtain the updating gradient by one forward and backward propagation and the other use it to prepare the perturbation gradient for the next step. The perturbation gradient is periodically re-anchored and gradually smoothed with the updating gradient at the intermediate steps. By relieving the perturbation gradient computation overhead, GraphSA.M has comparable efficiency with the traditional optimizers. More importantly, we verify through formula derivation and experiments that GraphSAM has the ability to compete with SAM. GraphSAM can converge to the flat minima, as shown in Fig. 2 Our contribution is in the following four aspects:
We observe the sharp local minima in a convergence of the graph transformer model for the molecular property prediction problem. Similar to other domains, the sharp local minima result in poor generalization performance.
We propose GraphSAM to relieve the computation overhead of SAM. In particular, instead of computing the perturbation gradient at every step, GraphSAM computes it periodically and uses the updating gradient of the previous step to approximate it smoothly.
We experimentally observe the approximated perturbation gradient is close to the ground- truth, and theoretically prove the loss landscape is constrained within a small bound centered at the desired loss of SAM.
The extensive experiments on the benchmark molecular datasets show that GraphSAM achieves comparable or even outperforming results, and frees the expensive pre-training. The time overhead is marginal compared with SAM.
2 RELATED WORK
Sharp Local Minima. Sharp local minima can primarily affect the generalization performance of deep neural networks (Izmailov et al.) 2018 Jastrzebski et al. 2017 Guo et al. (022)) Zhou et al.) 2022al
2023 Wang et al./ 2022b 2020). Recently, many studies have attempted to explore how to solve the optimization problem by locating the parameters in flat minima instead of sharp local minima (Zhou et al. 2020; Jiang et al. 2019: Moosavi-Dezfooe
et al. 2019 Zhou et al. 202 Liu
et al. 2020 Wen et al.] 2018). At the same time, it is shown through a large number of experiments that there 18 strong correlation between sharpness and generalization error at various hyperparameter
settings (Jiang et al. 2019). This motivates the idea of minimizing sharpness during training to
2

---

Published as a conference paper at ICLR 2024
improve standard generalization, resulting in sharpness-aware minimization(SAMM) Foret al. 2020).
Sharpness-Aware Minimization. SAM has been successfully used in areas such image classifica- tion (Foret et al. 2020 Du et al. 2021 Liu et al. 2022 Zhao et al., 2022by and natural language processing (Bahri al. 2021). Especially, LookSAM improves computational efficiency by elimi- nating the calculation of the updating gradient (Liu et al.
2022). SAF (Du et al.) 2022) suggests an innovative trajectory loss that mitigates abrupt decreases in the loss sharp local minima during the weight update trajectory to reduce the time loss. Nevertheless, most of the work still ignores the
fact of SAM's double overhead (Damian et al., 202T, Kwon et al. 2021 Wang et al.) 2022a; Li & Giannakis (023) and no studies of SAM are available in the graph domain. This forces us to propose the GraphSAM algorithm in the field of molecular graphs, which retains the generalization ability of SAM while improving computational efficiency.
Molecular Graph Representation Learning. In recent years, Steven et Kearnes et al. (2016);
Xiong et al. (2019) have tried applying GNNs to molecular characterization learning to better utilize the structural information of molecules. Nevertheless, they have poor generalization ability to new- synthesized molecules. Recently, a great deal of pre-training has been done on graphs to capture the rich information in molecular graphs (Hu et al., 2019 2020; Rong et al., 2020; Honda et al., 2019). In contrast, the pre-training-free graph transformer models are prone to overfit and have poor generalization performance, thus some recent studies utilize contrastive learning and data augmentation methods solve the problem (Xia al. 2022 Zhang et al. 2020). However, these methods have a considerable training time and are complex to operate.
PROBLEM STATEMENT AND SAM
3.1 MOLECULAR PROPERTY PREDICTION
A molecular structure can be represented by an attributed graph G (V,E), where V and E denote the atoms (nodes) and chemical bonds (edges) within the molecule, respectively. The initial feature of node v is denoted as Zy. and the initial feature of edge (u, v) is Euv- The set of neighbors of node v is denoted as N₁.. We consider two categories of molecular property prediction tasks: graph classification and graph regression. Both of them are given set of molecular graphs G = [G1,'.. ... Gr } and their labels/targets {v/1. .y J. where the molecular property prediction task is to infer the label/target of new graph.
3.2 SHARP LOCAL MINIMA IN TRANSFORMER
Weight Loss Landscape. Existing studies have demonstrated that converging to a flat region in the loss landscape can improve the generalization performance
of DNNs (Li et al. 2018 Du et al.) 202T Xue et al.) [02T
0.50
+Adam
Juan et al. 2023 Yang et al. 2022). To understand model
45AM
0.45
convergence, we visualize the loss landscape of GROVER
4GraphSAM
Training Loss
0.400
trained on the BBBP dataset in Fig. 2 The training loss
0.35
(y-axis) is defined as C(0 + DD) (c.g., cross-entropy loss), where 0 denotes the best-trained model parameters at the convergence of GROVER. Do is the random gradient direc- tions sampled from Gaussian distribution, and ó (x-axis) controls the scalar size of the step that moves by Dd to get the perturbation parameters, which is called the best parameter's 's neighborhood. In Fig. 2 we can observe that the loss landscape of Adam is much sharper, i.e., its neigh- borhood parameters' training loss is much larger than SAM and GraphSAM. The over-parameterized pre-training-fir graph transformer model located in the sharp local minima.
0.30
0.25
0.20
0.15
-1.06-075-0.0002200.000
0.25 0.50 1.00
Figure 2: The loss landscape of GROVER with different optimizers.
Generalization Performance. Transformer models rely on massive pre-training (Rong et al. 2020 Yuan et al. 2 Liu et al. 2021a: Touvron et al. [202 Sun et al. 2022) to mitigate the sharp local minima problem, such as GROVER (Rong et al. 2020). They are demanding data and computation resources, and require laborious tuning and expert knowledge in designing the pre- pre-training tasks. When the pre- pre-training process is removed, the models often fall into the sharp local minima, which significantly affects the generalization performance of graph transformer models. As shown in Fig.
3

---

Published as a conference paper at ICLR 2024
use the test accuracy value to express the models' generalization performance, and the GROVER performance significantly deteriorates once pre-training is ablated.
3.3 SHARPNESS- AWARE MINIMIZATION
In order to improve generalization and free model construction from large-scale pre-training, Foret et al. (Foret et al. 2020) propose the SAM algorithm to seek parameter values whose entire neigh- borhoods have both low loss and low curvature. Formally, SAM trains transformer by solving the following min-max optimization problem:
min max Co(0 + +), |ellssS
(1)
where p is the size of the gradient stepping ball, G is the training dataset, and 0 denotes the model weight parameters. Here we omit the regularization term for simplicity. At training step t. SAM solves the min-max problem by the following iterative process:
G
E = VoCc(O). ē =
Wi = VoCc + dc+1 = -7 Wr.
||e|112"
For the inner optimization, SAM calculates the perturbation gradient e, with a complete forward and backward propagation and then normalizes e, within the p-ball to get &, The normalized gradient & is applied to update the model as shown by the inner maximization problem in equation [ For the outer optimization, SAM obtains the updating gradient W, with another complete forward and backward propagation. Then is used to update the model with the learning rate γ) towards the expected smooth minima. SAM consumes double overhead due to the extra computation of perturbation gradient compared with the base optimizer.
4 GRAPHSAM
4.1
UNDERSTANDING OF THE GRADIENTS
To optimize the computational efficiency of SAM, we provide deeper insights into the perturbation gradient and the updating gradient wi. Particularly, introduce two key observations as motivation from experiments conducted on the BBBP dataset with two graph transformer models.
3.0
|leus oll
Ho Gelle
Accurary Value
96
SAM
Adam
Number of directions
400
Incondistent
Consistent
2.0
SAM-K
Gradient
1.5
93
300
1.0
92
200
91
0.5
100
0
100
200
300 400 500 600
700
Grover
COMPT
SAM GraphSAMSAM- SAM-One
Step
Model
Algorithm
Figure 3: Illustration on the observation of gradient on GROVER and CoMPT with BBBP dataset. (a) Gradient variation during training. (b) Accuracy of Adam, SAM, SAM-One, and SAM-k. (c) The similarity direction numbers of eS and wi in SAM, or the similarity direction numbers of e, and cñ in GraphSAM, SAM-K, and SAM-One.
Observation (The change of et and we during the iterative training process). As shown in Fig. ((a)
we use ||e:+1 eell2 and | well2 to describe the changing degrees of perturbation gradient e and the updating gradient w/. respectively. It is observed that W/ changes more drastically and has a broader range during the training process. In contrast, the variation of e is relatively minor, which can lead us to focus on the perturbation gradient to improve the computational loss of the model. We also observe |wel|2 > ||<:||2 in Fig. 9 which is useful for the subsequent proof of the theory
According to Observation 1 where the perturbation gradient changes gently, an intuitive solution to reducing the time complexity of SAM is compute 6, only once at the beginning, and then reuse
'hee observation is unique the domain of molecular graphs.
4

---

Published as a conference paper at ICLR 2024
the initial c, at the following training steps. Mathematically, we name such efficient solution as SAM-One, where the perturbation gradient is given by 6 = cu. for t 2 1. Notably, SAM-One can be regarded computed every k step. Although the computational efficiency is improved, the performances of both SAM-One and SAM-k deteriorate significantly as shown in Fig. ((b) For example, SAM-One consistently delivers the worst generalization performance. Compared with SAM on the CoMPT model, the molecule property classification accuracy of SAM-One drops from 96.2% to 92.2%. One possible reason is that the influence of c fluctuation is non-negligible, as shown in the small window of Fig. 3(a) The old perturbation gradient obtained many steps before, fails accurately reflect the scale of current éς- This experimental observation demonstrates that the vanilla algorithm of SAM-One or SAM-k does not adapt well to the molecular property prediction problem. Thus, we further study the correlation between the perturbation gradient and the updating gradient as follows.
special case of SAM-k (Liu et al. (0222, where the perturbation gradient is
Definition 1 (Similarity metric). We use the cosine function to measure the similarity between two vectors, denoted cos(:.)). The vector intersection angle is denoted L, e.g., if cos(a, b) 1. then Lab = -%.
Definition 2 (Gradient direction similarity). Given the periurbution grudient e, and the updating gradient Wr. if 0 I cos(e,, wi) )) $ 1, we call 6, and we are consistent: and if -1 S cos(er,w)) < 0, then e, and Wt inconsistent.. That is, the angle L between consistent vectors satisfies 0 < v 90%.
Observation 2 (Similarity between a-11 and wa). As shown in Fig. 3(c) divide the simi- larity scores between the (c+41,W) pair of SAM into two categories: consistent pairs with 0 < Cos(er+1, < 1, and inconsistent pairs with -1 cos(Er+ wi) < O. We observe that the gradient direction of w; is prone to be close to that of e++1 during the training process. In particular, consistent pairs account for 67. 75.5% in the overall training process.
The above observations motivate us to approximate the perturbation gradient with the updating gradient from the last step, in order to save time for the model optimization computation.
4.2 AN EFFICIENT SOLUTION
In the following, we propose novel optimization algorithm, GraphSAM, to address the above challenges. The innovation of GraphSAM has the following two aspects: perturbation gradient approximation and gradient ball's size (p) scheduler.
Perturbation gradient approximation. GraphSAM contains two sequential gradient computations at each step: one to obtain the updating gradient by forward and backward propagation, and the other to use the updating gradient to prepare the perturbation gradient for the next step. Indeed, at training step f, we have:
C
cu = VoLa(Oo),
& pp
We = VeLo(Ai + &t),
||e 12"
0t+1 0, - n7 We, c+1 = 3 E + (1 3) - wt/ We ll2.
Like the SAM algorithm, at training step t = 0, we need an additional forward and backward propagation compute the perturbation gradient eg. Then we project the perturbation gradient within the p-ball to get é, and the neighboring parameters (0,+ of 0,. With another forward and backward propagation from (0, + ét), we can get the updating gradient wi to update the model. Unlike the SAM algorithm, at training step + 1, GraphSAM uses the idea of moving average to prepare 6 by combining the e, with W, of the previous step, and the hyperparameter 3 E [0, 1) controls the exponential decay rate of the moving average. The moving average method can effectively retain the information of the previous steps of wt, and the initial information of 6g- The 6++1 acquired by this method can be approximated with the ground truth. Meanwhile, it prevents computing the perturbation gradient at each step with additional forward and backward propagation, which increases the efficiency of SAM for model optimization.
Compared with SAM in equation 2 instead of exactly calculating perturbation gradient with extra forwarding and backward, GraphSAM approximates Et+1 for the next step based on the above two observations: &, changes slowly and 6+1 positively relates to w,. As the training step increases, 6+1 obtained by moving average may gradually deviate from the ground-truth exactly computed by
5

---

Published as a conference paper at ICLR 2024
SAM. Therefore, reduce the error, the perturbation gradient c1+1 is periodically re-anchored, e.g.. the perturbation gradient is recalculated at the first-step of each epoch. This ultimately makes the direction of E++1 obtained by the moving average is keeping close to the ground-truth.
To evaluate the approximation performance, we compare the approximated perturbation gradients obtained by SAM-One, SAM-K, and GraphSAM in Fig. [(c] Specifically, we measure their co- sine similarity scores with the ground-truth perturbation gradient obtained by SAM. As shown in Fig. 3(c) consistenlpairs s accounts for 72.46% of the total training process in GraphSAM. The consislenlpuirs statistics of SAM-One (40.05%) and SAM-k (56.74%) 74%) algorithms are much worse perturbation gradient.
than GraphSAM. The results show that our GraphSAM provides a better solution to approximate the
Gradient ball's size (p) scheduler. As seen from equation [ the projected perturbation gradient size & is determined not only by the perturbation gradient <t. but also by the size of gradient ball p. The models' generalization ability is correlated intimately with p (Kwon et al. 2021). as shown in Table in the experiment. In general, the size of the perturbation gradient should change with the training process. The early stages of training require a larger perturbation gradient size to enable graph transformer models to handle the much sharper minima cases. When the model converges at the later stages of training, the perturbation gradient should decrease or even converge to O, allowing the model to fall into flatter region. Therefore, motivated by the learning rate scheduling (e.g.. StepLR (Brownlee] ), we propose a simple but effective gradient hall's s size scheduler as below:
Pnew Pinitial * opoch//
where the hyperparameter À controls the update rate of p. i.e., how many epochs for one update. The hyperparameter y controls the modification scale in p. Summarizing the above, the GraphSAM algorithm is shown in Appendix.
4.3 GRADIENT APPROXIMATION ANALYSIS
In this section, we verify that the perturbation gradient of GraphSAM approximates the ground truth. Meanwhile, we prove the loss landscape of GraphSAM is limited to a small range centered on the expected loss of SAM.
Before going to the mathematical proof, we first illustrate GraphSAM and SAM in Fig. Considering the model parameter learning from 6g to 0.+1, GraphSAM first approximates the perturbation gradient and then ad- versarially perturbs the graph transformer model to the neighborhood, where the model is finally updated to 0.+1 with a complete forward and backward propagation. In particular, the perturbation gradient efe of GraphSAM is obtained from the moving average between c--1 and W1-1 as shown by the last equation in equation 3 The projected perturbation gradient Et practically used to perturb model to neighborhood gadv, is obtained by mapping E over the gradient ball. Based upon 0gdv, we compute the updat- ing gradient wi by a forward and backward propagation to update O2 -> 0.+1. Therefore, each training step of GraphSAM only requires the same tensor flowing as the base optimizer. In contrast, as illustrated in Fig. 4 SAM takes another complete propagation to obtain e and its projected one Eे. Since è is exactly estimated, we treat it as the ground truth compared with the approximated perturbation gradient eG used in GraphSAM.
(1-
(3)
l[w.
-1j W
B ep
6
0,
WE
Figure 4: Visualization of GraphSAM. The black dashed line is obtained from GraphSAM's updating gradient which targets the model's parameters @ moving to a flatter region.
Conjecture 1. Let és and êg denote the perturbation weights of SAM and GraphSAM, respectively, where we ignore the subscript of for the simple representation. Suppose that Tw] e as empirically discussed in Observation and ||es||2 < [|écllafor p > designating és as the ground-truth. We have:
max Le(O + ēs) S
max Le(0 + εG).
kellaS
"This assumption is obtained from Fig.

---

Published as a conference paper at ICLR 2024
We list the detailed proof in Appendix A.2 From the generalization capacity perspective, it is observed that the adversarially perturbed loss of SAM is upper bounded by that of GraphSAM. Recalling the min-max optimization problem of SAM in equation[]] if we replace the inner maximum objective from (a + ēs) to Ls + (G), the graph transformer is motivated to smooth a worse neighborhood loss in the loss landscape. In other words, the proposed GraphSAM aims minimize a rougher neighborhood loss, whose value is theoretically larger than that of SAM, and obtain a smoother landscape associated with the desired generalization.
Let ||Cs(0 + éG.) Lo (0 + és)ll denote the gap of generalization performance between SAM and GraphSAM. is positively correlated with the gradient approximation error as below:
Conjecture 2.
La + ēG) - La + Es) || &G -
We list the detailed proof in Appendix A.I
Given the above-constrainee loss bias and the limited training epochs, the optimization algorithm of GraphSAM is prone to converge to the area neighboring that of SAM, where the loss landscape is smooth and desired for good generalization. Indeed, we can compute the loss bias by:
cG - Es || a ēG |= a- l
6G
11.
(5)
EG
where || a êG || denotes the are length between gradients éG and éş. When the intersection angle (or the size of gradient ball p) close to 0, & is infinitely approximated to ?². Conjecture 2 with equation 5 shows GraphSAM's loss landscape is constrained within small bound centered at the desired loss of SAM.
As the training step t increases, the intersection angle between Es and è gets wider, and the generalization performance of GraphSAM may reduce significantly. As seen from equation [5] in addition a, the size of gradient ball plays another key role in determining the gradient approximation error. Therefore, to ensure GraphSAM has a similar generalization performance as SAM, we periodically reduce the size of p by using equation 4 In summary, we use the moving average method to approximate the ground-truth and control the size p to reduce the gap between é and cp. Both of them are the keys keeping GraphSAM with good generalization performance.
5 EXPERIMENTS RESULTS
In this section, we evaluate the effectiveness of GraphSAM on two graph transformer models. Overall, we aim to answer three research questions follows. Q1: Can GraphSAM improve the generalization performance of the graph transformer models? Q2: How effective is GraphSAM compared to SAM and other variants of SAM? Q3: How do each module and key hyperparameters of GraphSAM affect its efficiency and performance?
5.1 EXPERIMENT SETUP
Datasets. Following the settings of previous molecular graph tasks, consider six public benchmark datasets: BBBP, Tox21, Sider, and ClinTox for the classification task, and ESOL and Lipophilicity for the regression task. We evaluate all models
random split suggested by MoleculeNet (Wul et al. 2018), and split the datasets into training, validation, and testing with a 0.8/0.1/0 ratio. More detailed statistics are provided in Appendix A.3 In addition, the Backbone frameworks and baselines and Implementations are described in detail Appendix A.4
5.2 PERFORMANCE ANALYSIS
5.2.1 GENERALIZATION PERFORMANCE IMPROVEMENT
To answer the research question Q1, we apply GraphSAM to GROVER and CoMPT in comprehen- sive comparison with the baseline approaches on different datasets. We make two key observations.
>> Lacking the large-scale pre-training process, an over-parameterized graph transformer model does not necessarily work better than GNNs based model. As shown in Table I compared with GNNs-based models, GROVER (no pre-training) and CoMPT have even worse performance on most datasets than the GNNs-based CMPNN. fully illustrates that the over- parameterized transformer model tends to fall into the sharp local minima during the training process like Fig. 22, The sharp local minima significantly affect the generalization performance the graph transformer model.
7

---

Published as a conference paper at ICLR 2024
Table 1: Prediction results of GraphSAM and baselines six datasets. We used 5-fold cross- validation with random split and replicated experiments on each task five times. The mean and standard deviation of AUC or RMSE values are reported. Bold is the best optimizer, and underlined is the best model.
Task
Graph Classification (ROC-AUCT)
Graph Regression (RMSE.) )
Dataset
BBBP
Tox21
Sider
ClinTox
ESOL
Lipophilicity
GCN
0.690
0.81940.081
0.62340.022
0.807 +.807+0.044 044
0.970
1.313 149
MPNN
0.901 20.022
0.83 0.804
0.63440.0014
0.881 037
0.702
1.21210.22
DMPNN
0.912
0. 0.845-g0.022
0, 0.64660..20
20.020
0.663+0.000
1.15940.207
CMPNN
0.92540.011 037
0.837
0.64040.01
0.91840.0
0.582 055
0.63340
GROVER
0.917+0.028
0.822+0.010
0. 0.649+0..05
0.853 0..43
0.639+0..08
0.671+0.00
SAM
0.926+00022
0.840+0.005
0.66040.043
0.872=0.084
0.619 20.19
U. 0.662+0.002
GraphSAM
0.928 10.000
0.848_..00
0.665 0.0888
0.86611.00 051
U.62511000
0.654
CoMPT
0.9 0.948
0. 1.828go.oo oos
0. 0.621
0.914
0.062
0 0.618g0.012
SAM
0.962 23.00a
0 0.839 2...0000 oos
0 0.643
0.927 025
0.517:9.002
0.61110.015
-GaphSAM
0.961 1.. 01
0.841_0.004
0.645 20.01
0.937 00.008
0.511
0.608 20.007
>> GraphSAM and SAM can improve the generalization performance of transformer models that fall into sharp local minima. Compared with the vanilla model, adapting GraphSAM each transformer model, it delivers the average improvements of 1.52% (GROVER), and 2.16% (CoMPT), respectively. SAM is 1.55% (GROVER) and 1.97% (CoMPT). We show the loss curves of GROVER and CoMPT with dif- ferent optimizers in Fig. 5 SAM has the best generalization loss (the purple line), Graph- SAM is the second best (the green line), and Adam is the worst (the orange line). In par- ticular, there a large gap between Adam's 's training loss and testing loss. This is because the over-parameterizer pre-training-fee graph transformer model is always prone to overfitting.
COMPT
GROVER
0.8
Aday Juct
1.2
Adary Lind
0.6
1.0
0.8
3..4
0.6
0.4
0.2
5
10
15 Epochs
20
25
30
Epochs
Figure 5: Loss curves of differrnt optimizers on the BBBP dataset.
In addition, shown in Fig. 2 we display the loss landscape and demonstrate that both GraphSAM and SAM have significantly improved the sharpness compared to Adam. GraphSAM and SAM optimize the model update using the perturbation gradient to seek parameter values with low loss and curvature throughout the neighborhood. It also prevents the parameters from falling into sharp local minima. Both GraphSAM and SAM can improve the model's 's generalization.
5.2.2 THROUGHPUT AND ACCURACY OF THE OPTIMIZERS
To answer the research question Q2, we summarize the model efficiency and generalization perfor- mance of different optimization algorithms in Table In summary, we observe that:
> GraphSAM has a similar generalization performance to SAM while having high throughput. We compare the throughput and generalization performance of GraphSAM with other optimization algorithms in Table 2 and Table 7 Where throughput denotes the computational overhead which quantified by graphs processed per second (Graphs/s), and the generalization performance is measured by test ROC-AUC/RMSE The experimental results show that GraphS AM can improve the training speed up to 155.4% compared SAM. The reason for this is that we perform only one forward and backward propagation in most of the training steps, which will greatly improve the model's efficiency. In terms of performance, GraphSAM is close to SAM and even better than SAM on some datasets in Table
The other optimization algorithms also can improve efficiency but lose some generalization perfor- mance. For instance, although SAM-k and SAM-One have a high throughput, their generalization performance is poor. This is because though they also periodically compute the perturbation gradient, it differs significantly from the ground-trute -ruth gradient at other training times, eventually affecting the model's generalization performance. LookSAM is an optimization algorithm designed for the
Under the optimal parameters, SAM theoretically has the best performance
8

---

Published as a conference paper at ICLR 2024
transformer model in the computer vision do- main. Unlike our algorithm, it maintains the
Table 2: Classification accuracy and training speed.
BRBP
perturbation gradient by the updating gradi- ent. Due to the design of LookSAM, it does not conform the properties of the molecu- lar graph model. Therefore, it is inferior to train the model based on the original hyper- parameter. LookSAM has been tuned by us (p 0.0001, a 0.2, k 8). but its per- formances are only the same as Adam's ss in all datasets, and LookSAM takes more time. ad- dition, AE-SAM and RST compute the update gradient by different strategies, respectively. Be- cause of the determination problem of the strate- gies, their computational consumption is very dissimilar in different models and data. There is also no guarantee that the generalization perfor- mance will not be affected. Please see Appendix A.4 and A.5 for details.
GROVER
ROC-AUCT
Throughput
Adam
0.917
362
SAM
0.926
201(100.0%)
SAM-One
0.899
330(164.2%)
SAM-K
0.915
287(142.7%)
LookSAM
0.919
266(132.3%
AE-SAM
0.923
291(144.8%
RST
0.912
312(
GraphSAM
0,928
272
CoMPT
ROC-AUCT
Throughput
Adam
0.948
218
SAM
0.962
112
SAM-One
0.922
199(177.69)
SAM-K
0.941
183(163.4%)
LookSAM
0.952
161(143.8%)
AE-SAM
0.955
178(158.9%)
RST
0.940
186(166.1%)
GraphSAM
0.961
174(155.4%)
5.2.3 ROLES OF GRAPUSAM MODULES,
To answer the research question Q3, we conduct experiments on multiple graph transformer models and datasets for Adam, SAM, and GraphSAM to verify the role of each module. In summary, have the following two observations.
>> The size of directly impacts the generalization of SAM and GraphSAM for different datasets. In particular, we investigate the impact of different fixed-value p and gradient ball's size p schedulers on the performance of various datasets. We list their performances in Table and conduct that: (1) The p of the gradient ball helps SAM and GraphSAM to seek the model parameters that lie in neighborhoods having uniformly low loss. We need manually adjust the size of p to get better performance for different datasets (e.g.. the best performance of BBBP is p = 0.05 but ESOL p 0.001 in SAM). Due to the nature of GraphSAM, which needs a small p to approximate the ground-truth, is unsuitable for over-scaling. (2) Compared to fixed p. the scheduler can improve the model's ' generalization performance on different datasets. The reason for the above conduct is that parameter re- re-scaling can cause a difference in sharpness values so the size of p may weaken the correlation between sharpness and generalization gap (Dinh et al., (017). this phenomenon is named the scale-dependency problem. Our gradient ball's size p scheduler module, which changes the size of rho periodically, can remedy the scale-dependency problem of sharpness. Experimentally, GraphSAM and SAM have relatively stable generalization performance in different datasets under the action of the scheduler.
Accuracy-efficienne trade-off. The secret of GraphSAM's ability to maintain similar perfor- mance to SAM the timely re-anchor of 6, The effectiveness of another module of GraphSAM, moving average, has been shown in Table [ In this part, mainly verify the correlation between accuracy-efficiency. We propose the GraphSAM-K to investigate the impact of the rate of re-anchor to the perturbation gradient & on generalization performance. The K means that we perform an additional forward and backward propagation to re-anchor from to Eo for every K epoch. As in Fig. 6 we analyze GraphSAM-K for different values of K. The N means that only once re-anchor of Ge. When K = GraphSA.M is comparable to SAM. When K = 2, the performance of GraphSAM slightly higher than Adam's. When K 2, the efficiency of GraphSAM is increasing rapidly, but its performance drops sharply. The reason is that as the training steps increase, the error between & obtained from equation 3 and the ground-truth és becomes increasingly large. To reduce the error, we need to re-anchor the é, by forward and backward propagation in each epoch repeatedly.
CONCLUSION
In this paper, we perform a series of analytical experiments to show that the graph transformer model suffers from sharp local minima when the pre-training process is removed. SAM can solve this problem, but its computational loss is double that of the traditional optimizer. Then we propose an efficient algorithm GraphSAM which reduces the training cost of SAM and improves the generaliza- tion performance of graph transformer models. The experiments show the superiority of GraphSAM, especially in optimizing the model update process.

---

Published as a conference paper at ICLR 2024
7 ACKNOWLEDGMENTS
This work was supported by a grant from the National Natural Science Foundation of China under grants NNo.62372211 62272191), the Foundation of the National Key Research and Development of China 2No.2021ZDDN 12500), and the International Science and Technology Cooperation Program of Jilin Province (No. NN..20230402076GH) and the Science and Technology Development Program of Jilin Province (No. 20220201 153GX).
REFERENCES
Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware mini- mization. In International Conference on Machine Learning, pp. 639-668. PMLR, 2022.
Dara Bahri, Hossein Mobahi, and Yi Tay. Sharpness-aware minimization improves language model generalization. arXiv preprint arXiv:2110.08529, 2021.
Lorenz C Blum and Jean-Louis Reymond. 970 million druglike small molecules for virtual screening in the chemical universe database gdb- 13. Journal of the American Chemical Society, 131(25): 8732-8733, 2009.
J Brownlee. Using learning rate schedules for deep learning models in python with keras. Machine learning mastery, June, 21, 2016.
Stephen J Capuzzi, Regina Politi, Olexandr Isayev, Sherif Farag, and Alexander Tropsha. Qsar mod- eling of tox21 challenge stress response and nuclear receptor signaling toxicity assays. Frontiers in Environmental Science, 4:3, 2016.
Jianwen Chen, Shuangjia Zheng, Ying Song, Jiahua Rao, and Yuedong Yang. Learning at- tributed graph representations with communicative message passing transformer. arXiv preprint arXiv:2107.08773, 2021a.
Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform resnets without pre-training or strong data augmentations. arXiv preprint arXiv:2106.01548, 2021b.
Seyone Chithrananda, Gabriel Grand, and Bharath Ramsundar. Chemberta: large-scale self- supervised pretraining for molecular property prediction. arXiv preprint arXiv:2010.09885, 2020.
Alex Damian, Tengyu Ma, and Jason Lee. Label noise sgd provably prefers flat global minimizers. Advances in Neural Information Processing Systems, 34:27449 27461,771111 2021.
John S Delaney, Esol: estimating aqueous solubility directly from molecular structure. Journal of chemical information and computer sciences, 44(3):1000-1005 2004.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In International Conference on Machine Learning, pp. 1019-1028, PMLR, 2017.
Jiawei Du, Hanshu Yan, Jiashi Feng, Joey Tianyi Zhou, Liangli Zhen, Rick Siow Mong Goh, and Vincent YF Tan. Efficient sharpness-aware minimization for improved training of neural networks. arXiv preprint arXiv:2110.0314| 2021.
Jiawei Du, Daquan Zhou, Jiashi Feng. Vincent Tan, and Jocy Tianyi Zhou. aharee training for free. Advances in Neural Information Processing Systems, 35:23439-23451, 2022.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. arXiv preprint arXiv:2010.01412 2020,
Anna Gaulton, Louisa J Bellis, A Patricia Bento, Jon Chambers, Mark Davies, Anne Hersey, Yvonne Light, Shaun McGlinchey, David Michalovich, Bissan Al-Lazikani, et al. Chembl: a large-scale bioactivity database for drug discovery. Nucleic acids research, (NDD)):DI10-D1)
2012.
Kaitlyn M Gayvert, Neel Madhukar, and Olivier Elemento, A data-driven approach to predicting successes and failures of clinical trials. Cell chemical biology, 23(10): 1244:1301 2016.
10

---

Published as a conference paper at ICLR 2024
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pp. 1263-1272. PMLR, 2017.
Kai Guo, Kaixiong Zhou, Xia Hu, Yu Li, Yi Chang, and Xin Wang. Orthogonal graph neural networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp. 3996-4004, 2022.
Shion Honda, Shoi Shi, and Hiroki R Ueda. Smiles transformer: Pre-trained molecular fingerprint for low data drug discovery. arXiv preprint (e.Xix:9111.047555 2019.
Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovec. Strategies for pre-training graph neural networks. arXiv preprint
12265, 2019.
Ziniu Hu, Yuxiao Dong, Kuansan Wang. Kai-Wei Chang, and Yizhou Sun. Gpt-gnn: Generative pre-training of graph neural networks. In Proceedings of the 26th ACM SIGKDD International Conference Knowledge Discovery & Data Mining, pp. 1857-1867, 2020.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson, Av- eraging weights leads to wider optima and better generalization. arXiv preprint arXiv:1803.05400) 2018.
Stanislaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey Three factors influencing minima in sgd. arXiv preprint 1rXiv:1711.0460 2017.
Weisen Jiang, Hansi Yang. Yu Zhang, and James Kwok. An adaptive policy to employ sharpness- aware minimization. arXiv preprint arXiv:2304. 14647, 2023.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. arXiv preprint / 2019.
Xin Juan, Fengfeng Zhou, Wentao Wang, Wei Jin, Jiliang Tang, and Xin Wang. Ins-gnn: Improving graph imbalance learning with self-supervision. Information Sciences, 637:1 337:118935, 2023.
Steven Kearnes, Kevin McCloskey, Mare Berndl, Vijay Pande, and Patrick Riley. Molecular graph convolutions: moving beyond fingerprints. Journal of computer-aided molecular design. 30(8): 595-608, 2016.
Thomas N Kipf and Max Welling. Semi- supervised classification with graph convolutional networks arXiv preprint arXiv: 1600::660.00000 (600002007 2016.
Michael Kuhn, Ivica Letunic, Lars Juhl Jensen, and Peer Bork. The sider database of drugs and side effects. Nucleic acids research, 44(DI):D0075-DI0T) 2016
Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. Asam: Adaptive sharpness-awwe aware minimization for scale-invariant learning of deep neural networks. In International Conference Machine Learning, pp. 5905-5914. PMLR, 2021.
Bingcong Li and Georgios B Giannakis. Enhancing sharpness-aware optiiiiiai through variance suppression. arXiv preprint arXiv:2309. /5639, 2023.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape neural nets. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018
Chen Liu, Mathieu Salzmann, Tao Lin, Ryota Tomioka, and Sabine Süssttrunk. On the loss landscape of adversarial training: Identifying challenges and how to overcome them. Advances in Neural Information Processing Systems, 33:21476-21488 2020.
Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, and Yang You. Towards efficient and scalable sharpness-awe aware minimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12360-12370, 2022.

---

Published as a conference paper at ICLR 2024
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 10012-10022, 2021a.
Zirui Liu, Kaixiong Zhou, Fan Yang, Li Li, Rui Chen, and Xia Hu. Exact: Scalable graph neural networks training via extreme activation compression. In International Conference on Learning Representations, 2021b.
Ines Filipa Martins, Ana L Teixeira, Luis Pinheiro, and Andre O Falcao. A bayesian approach to in silico blood- brain barrier penetration modeling. Journal of chemical information and modeling, 52 (6):1686-1697, 2012.
Seyed-Mohsen Moosavi-Dezfooli, Alhussein Fawzi, Jonathan Uesato, and Pascal Frossard. Robust- ness via curvature regularization, and vice versa. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 9078-9086, 2019.
Raghunathan Ramakrishnan, Mia Hartmann, Enrico Tapavicza, and O Anatole Von Lilienfeld. Electronic spectra from tddft and machine learning in chemical space. The Journal of chemical physics, 143(8). 2015.
Yu Rong, Yatao Bian, Tingyang Xu, Weiyang Xie, Ying Wei, Wenbing Huang, and Junzhou Huang. Sele-super graph transformer on large-scale molecular data. Advances in Neural Information Processing Systems, 33: 12559- 2020.
Ying Song, Shuangjia Zheng, Zhangming Niu, Zhang-Hua Fu, Yutong Lu, and Yuedong Yang. Communicative representation learning on attributed molecular graphs. In IJCAI, volume 2020, pp. 2831-2838, 2020.
Govindan Subramanian, Bharath Ramsundar, Vijay Pande, and Rajiah Aldrin Denny. Computational modeling of 3-secretase (bace-1) 1) inhibitors using ligand based approaches. Journal of chemical information and modeling, 56(10):1936-199) 1936-199, 2016.
Mingchen Sun, Kaixiong Zhou, Xin He, Ying Wang, and Xin Wang. Gppt: Graph pre-training and prompt tuning to generalize graph neural networks. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 1717-1727, 2022.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International Conference on Machine Learning, pp. 10347-10357. PMLR, 2021.
Pichao Wang. Xue Wang, Hao Luo, Jingkai Zhou, Zhipeng Zhou, Fan Wang, Hao Li, and Rong Jin. Scaled relu matters for training vision transformers. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp. 2495-2503, 2022a.
Xiaoyang Wang, Yao Ma, Yiqi Wang, Wei Jin, Xin Wang, Jiliang Tang, Caiyan Jia, and Jian Yu. Traffic flow prediction via spatial temporal graph neural network. In Proceedings of the web conference 2020, pp. 1082-1092. 2020.
Yili Wang, Kaixiong Zhou, Rui Miao, Ninghao Liu, and Xin Wang. Adagel: Adaptive subgraph contrastive learning to generalize large-scale graph training. In Proceedings of the 31st ACM International Conference on Information Knowledge Management, pp. 2055,2055, 2022b.
Wei Wen, Yandan Wang. Feng Yan, Cong Xu, Chunpeng Wu, Yiran Chen, and Hai Li. Smoothout: Smoothing out sharp minima to improve generalization in deep learning. arXiv preprint arXiv: 2018
Michael Withnall, Edvard Lindelöf, Ola Engkvist, and Hongming Chen. Building attention and edge message passing neural networks for bioactivity and physical-chemice property prediction. Journal of cheminformatics, 1-18, 2020.
Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh Pappu, Karl Leswing, and Vijay Pande. Moleculenet: a benchmark for molecular machine learning. Chemical science, 9():5133-300 2018
12

---

Published as a conference paper at ICLR 2024
Jun Xia, Jiangbin Zheng, Cheng Tan, Ge Wang, and Stan Z. Li. Towards effective and generalizable fin-tuning for pre-trained molecular graph models. bioRxiv, 2022.
Zhaoping Xiong, Dingyan Wang. Xiaohong Liu, Feisheng Zhong, Xiaozhe Wan, Xutong Li, Zhaojun Li, Xiaomin Luo, Kaixian Chen, Hualiang Jiang, et al. Pushing the boundaries of molecular representation for drug discovery with the graph attention mechanism. Journal of medicinal chemistry, 63(16))8749-876) 2019.
Haotian Xue, Kaixiong Zhou, Tianlong Chen, Kai Guo, Xia Hu, Yi Chang, and Xin Wang. Cap: Co-adversarial perturbation on weights and features for improving generalization of graph neural networks. arXiv preprint arXiv:2110. (tl1:2111..4855, 2021.
Kevin Yang, Kyle Swanson, Wengong Jin, Connor Coley, Hua Gao, Angel Guzman- Perez, Timothy Hopper, Brian P Kelley, Andrew Palmer, Volker Settels, et al. Are learned molecular representations ready for prime time? 2019.
Yintao Yang, Rui Miao, Yili Wang, and Xin Wang. Contrastive graph convolutional networks with adaptive augmentation for text classification. Information Processing & Management, 59(4): 102946, 2022.
Xian-bin Ye, Quanlong Guan, Weiqi Luo, Liangda Fang, Zhao-Rong Lai, and Jun Wang. Molecular substructure graph attention network for molecular property identification in drug discovery. Pattern Recognition, 128: 108659, 2022.
Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie Yan Liu. Do transformers really perform badly for graph representation? Advances Neural Information Processing Systems, 34:28877-28888, 2021.
Li Yuan, Yunpeng Chen, Tao Wang, Weihao Yu, Yujun Shi, Zi-Hang Jiang, Francis EH Tay, Jiashi Feng. and Shuicheng Yan. Tokens-to-token vit: Training vision transformers from scratch on imagenet. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 558-567, 2021.
Shichang Zhang, Ziniu Hu, Arjun Subramonian, and Yizhou Sun. Motif-driven contrastive learning of graph representations, 2020. URL
12533
Yang Zhao, Hao Zhang, and Xiuyuan Hu. Randomized sharpness-aware training for boosting computational efficiency in deep learning. arXiv preprint MXXXix222330099663 2022a.
Yang Zhao, Hao Zhang, and Xiuyuan Hu. Ss-sam: Stochastic scheduled sharpness-aware minimiza- tion for efficiently training deep neural networks. arXiv preprint arXiv:2203.0998 2022b.
Kaixiong Zhou, Xiao Huang, Yuening Li, Daochen Zha, Rui Chen, and Xia Hu. Towards deeper graph neural networks with differentiable group normalization. Advances in neural information processing systems, 33:4917-4928, 2020.
Kaixiong Zhou, Xiao Huang, Daochen Zha, Rui Chen, Li Li, Soo-Hyun Choi, and Xia Hu. Dirichlet energy constrained learning for deep graph neural networks. Advances in Neural Information Processing Systems, 34:21834-21 21844, 2021a.
Kaixiong Zhou, Qingquan Song, Xiao Huang, Daochen Zha, Na Zou, and Xia Hu. Multi-channel graph neural networks. In Proceedings of the Twenty-Ninth International Conference on Interna- tional Joint Conferences on Artificial Intelligence, pp. 1352 1358, 2021b.
Kaixiong Zhou, Xiao Huang, Qingquan Song, Rui Chen, and Xia Hu. Auto-gnn: Neural architecture search of graph neural networks. Frontiers in big Data, 5:1029307, 2022a.
Kaixiong Zhou, Zirui Liu, Rui Chen, Li Li S Choi, and Xia Hu. Table2graph: Transforming tabular data to unified weighted graph. In Proceedings of the Thirty-Fire International Joint Conference on Artificial Intelligence, IJCAL, pp. 2420-2426 2022b.
Kaixiong Zhou, Soo-Hyun Choi, Zirui Liu, Ninghao Liu, Fan Yang, Rui Chen, Li and Xia Hu. Adaptive label smoothing to regularize large-scale graph training. In Proceedings of the 2023 SIAM International Conference Data Mining (SDM), pp. 55-63. SIAM, 2023.
13

---

Published as a conference paper at ICLR 2024
A APPENDIX
A.1 APPENDIX 1: PROOF FOR CONJECTURE 2
Conjecture 2. Let || L3 + &G) - Ls + ēg) denote the gap of generalization performance
between SAM and GraphSAM. is positively correlated with the gradient approximation error as below:
(0 + ēG) (e + Es) ||||| ēG is ll
Proof 1. For SAM loss function Lo + 2), based on Taylor Expansion (Liu et al. 2022). we can obtain:
(6)
Therefore, have:
Lo(I+ + (G) - Cs(
(Co + &G VoLc(()) (Co + VoCc(0))
(ēG - Volc
According equation 7 the gap in generalization performance between SAM and GraphSAM is proportional to the gradient approximation error.
A.2 APPENDIX 2: PROOF FOR CONJECTURE
Conjecture 1. Let éş and &G be the projected perturbation gradients of SAM and GraphSAM, respectively. We ignore the subscript for the simple representation. For any p > we assume that ||εs|12 < [|cal and êş is the ground-truth, we have
max Cc + ēs) |aa
max Lc(0 + εG).
leallasp
Proof 2. For SAM loss function Combining equation of the main paper, that the size of & is determined by et. Then G is expressed as follows:
G B - (3) |
(1) When = 0, e. = E,, because both of them are computed by the same forward and backward propagation.
(2) When / 1,, c = Beg + (1 (1-B)-
wallz - win
> E
& E = -(C -G) + :3((
WD
lwwlll w/
(10)
B ? + (1 B3
|wollz
as shown in Fig. 3(a) the main paper, the change of perturbation gradient is small, and after experimental analysis, the updating gradient > e.. So e > Es. fort=1 =
(3) When / 2, because e > e.. we have
G - 5 = Bee - ( + (1 - (
W1
| 12
> B(e + + +(1 (1- B)
W1
5)
1
G 8 for =
14

---

Published as a conference paper at ICLR 2024
(4) When & > 2. analogously to the above conclusions, have
G - c = B(c - (1- {3(
W/
5)
w 1|
3(e < + +1 (
WL-1
5)
W/ 1|
In summary, we can conclude thate? G > E,. for t > 0. Following the above, the equation
can be rewritten as follows:
Le + 2 Ls + +
(ce - VoLo
(11)
2 (cr c V (O:) > O.
Therefore, can be proved that
max |
Lo + Es) <
max 1FalleS5
Lo + EG)
Also when p 0 or = 0, then
max 1elll
C6 +
Cs + ēG) ||Gal1.5p
max
Although we prove the correctness Theorem equation Mathematical Induction and experi- ment, still want to find out the main factors affecting the perturbation gradient of GraphSAM so we expand equation as follows:
G = + +(1 -1- B)
Wt 1
W 11 12
B + +(1 31)
Wi-2
+ (1 (8)
Wi
1
W 1|
= B + 3(1 - 8)
Wi-2
+ (1 9))
W/-1
(12)
| 1|2
= BCe
G + 101 (1- (3
wo
+... 30 9)
W4
|woll2
|
As equation 12 we find that the parameter 33 of the moving average plays a key role. So we perform some experiments on it. As shown in Table it is known experimentally that GraphSAM performs best when 3 = 0.99. This result is particularly evident in the ESOL dataset. So equation I0 holds for suitable 3. In the paper, we experimentally observe that the perturbation gradient and the updating gradient differ by two orders of magnitude, i.c., 100 times. Therefore, the optimal result is 3 = 0.99, so 1 3=0.01.
Table 3: The influence of smoothing parameter 3 of moving average on GraphSAM algorithm on CoMPT.
Algorithm
,
BBBP (ROC-AUCT)
ClinTox/ROC-AUEE
ESOL (RMSEL)
Lipophilicity (RMSE.)
0.9
0.958 1 0.014
0.928 1 0.015
0.357 1 0.020
0.614 1 0.021
GraphSAM
0.99
0.961 0.012
0.937 + 0.008
0.511 0.018
0.608 + 0.007
0.999
0.959 + 0.003
0.931 t 0.006
0.549 + 0.023
0.612 + 0.010
15

---

Published as a conference paper at ICLR 2024
A.3
APPENDIX 3: THE STATISTICS OF DATASETS
GraphSAM and SAM is evaluated on six molecular graph datasets, as described below:
Molecular Classification Datasets.
BBBP (Martins et al. 2012): This is Blood-brain barrier penetration (BBBP) dataset, which involves recording whether a compound has the permeability to penetrate the blood-brain barrier. has binary labels for 2.039 molecules.
Tox21 (Capuzzi et al. 2016): This is public database for measuring the toxicy of compounds, which contains 7,831 compounds against 12 different targets.
Sider (Kuhn al. 2016): The Side Effect Resource (Sider) a database of marketed drugs and adverse drug reactions, which contains 1,427 approved drugs with 27 system organ branches.
Clintox (Gayvert et al. 2016): has 1,478 drugs and compares them approved through FDA and drugs eliminated due to the toxicity during clinical trials.
Molecular Regression Datasets.
ESOL (Delaney 2004): is a small dataset with 1,128 molecules to document the solubility of compounds.
Lipophilicity (Gaulton et al., 2012). it was selected from the ChEMBL database, and contains 4, 1,198 molecules determine the lipophilicity of the molecule.
The statistics of the molecular graph datasets used in the graph classification and graph regression summarized in Table
Table 4: Statistics of datasets.
Category
Dataset
#Tasks
Task Type
#Molecule
BBBP
Classification
2,039
Physiology
Tox21
12
Classification
7,831
Sider
27
Classification
1,427
Clintox
2
Classification
1,478
Physical chemistry
ESOL
Regression 1,128
Lipophilicity
Regression 4,198
In addition to the above datasets, we have performed experiments on three additional datasets from biophysics (BACE) and quantum mechanics (QM7, QM8, QM9), as shown in Table 5 and Table 6 This demonstrates the generalization capability of GraphSAM.
Quantum Mechanics Datasets.
QM7 (Blum& & Reymond 2009): The QM7 dataset is subset of the GDB-13 database, which contains 7615 stable, synthesizable organic molecules.
QM8 (Ramakrishnan et al. 2015): The QM8 dataset comes from quantum mechanical com- putational modeling study of the electronic energy spectra and excited state energies of small molecules. It contains 21 786 samples.
Biophysics Dataset.
AAbrammannnnn et al. (016)): The BACE dataset provides combined quantitative (IC50) and qualitative (binary labeling) results for of human f-secretase (BACE-1) inhibitors and contains data for 1,522 molecules. All data are experimental values reported in the scientific literature over the last decade, some of which also provide detailed crystal structures.
16

---

Published as a conference paper at ICLR 2024
Table 5: Statistics of three additional datasets.
Category
Dataset
#Tasks
Task Type
#Molecule
Quantum mechanics
QM7
I
Regression
7.165
QM8
12
Regression
21,786
Biophysics
BACE
1
Classification
1,522
Table 6: Prediction results of GraphSAM and baselines.
BACE
QM7
QM8
GROVER
ROC-AUC
MAE
MAE +
Adam
0.871
78.9
0.0203
SAM
0.886
76.4
0.0189
GraphSAM
0.882
75.8
0.0185
CoMPT
ROC-AUC 1
MAE
MAE J
Adam
0.863
66.1
0.0159
SAM
0.876
64.3
0.0145
GraphSAM
0.880
64.1
0.0141
A.4 APPENDIX 4: TRAINING DETAILS AND ALGORITHM DETAILS
Backbone frameworks and baselines. To implement SAM and GraphSAM on molecular graph data with two widely-used graph transformer backbones, GROVER (Rong et al. (2020) and CoMPT (Chen et al. 2021a). We comprehensively compare them against four baselines in the graph-level task.
GCN (Kipf & Welling 2016) is the most classical graph convolutional neural network. Compared
with GCN, MPNN (Gilmer et al. 2017) and its two variants (DMPNN (Yang et al. 2019 and CMPNN (Song al. 2020)) integrate the edge features into the message message-passing process.
SAM and Efficient Variants of SAM. Here we analyze the advantages and disadvantages of SAM and the efficient variants of SAM, and verify them through experiments. As shown in Table 7
SAM (Foret et al., 2020): First work on Sharpness-ANe Aware Minimization(SAMM type methods. SAM computes the perturbation gradient and the updating gradient by two forward propagation and backpropagation, respectively. This prevents the model parameters from falling into sharp local minima. The disadvantage is twice the computational loss of the base optimizer.
LookSAM (Liu et al., 2022): An efficient approach to SAM variants. It saves computational loss by observing the training gradient of the CV domain model and periodically computing two gradients. The disadvantage is that it only cares about the gradient direction but ignores the gradient size. This leads to poor results when applied other domains if the gradient direction is not similar to the CV domain. When the is tuned down, the performance of LookSAM is only close to the base optimizer.
AE-SAM (Jiang et al., 2023): A generic hot-swappable component AE-SAM utilizes Squared stochastic gradient norms to decide whether to compute only the perturbation gradient and not the updating gradient. When Squared stochastic gradient norms exceed certain threshold, quadratic gradient computation is performed and vice versa. The trend of its Squared stochastic gradient norms in the CV domain not the same in the molecular graph domain. Therefore its performance varies greatly when encountering different models and data while ignoring gradient direction and gradient size.
RST (Zhao et al. 2022a): A method that utilizes the Bernoulli trial periodically compute two gradients. It is more stochastic and faster to calculate than AE-SAM. However, its drawbacks are also apparent, aimlessly deciding whether update the gradient or ends up with poorer performance. Its performance is only slightly higher than SAM-One and lower than SAM-k.
GraphSAM: An efficient SAM-typp method based on gradient descent trend specifically designed for transformer models in the molecular graph domain. In contrast to the single improvement of other SAM-type methods, we design the gradient approximation method for the gradient direction. We design adjustable gradient ball's size (p) scheduler for the gradient size. In addition,
17

---

Published as a conference paper at ICLR 2024
have designed method to compute the gradient periodically that GraphSAM maintains high level of performance improvement. Last but not least, we prove the effectiveness of GraphSAM through experimentation and derivation.
Implementations. The backhone implementations and their hyperparameter settings are provided by the publicly released repositoriesRong et al. (2020); Chen et al. (2021a). GROVER and CoMPT utilize Adam the base optimizer, and neither uses the pre-training strategy. We only adjust the specific hyperparameters introduced by GraphSAM: (1) smoothing parameters of moving average B tuned within [0.9, 0.99, 0.999), (2) the initial size of the gradient ball p selected from {0.05, 0.01, 0.005, 0,001), (3) the p's update rate À is searched over {1, 3. 5), (4) and the scheduler's 's modification scale {0.5, 0.2, 0.1 0. Due to space limitations, we place our experiments on hyperparameters in
Appendix A.6 All the experiments are implemented by Py PyTorch, and run on an NVIDIA TITAN-RTX
(24G) GPU.
GraphSAM uses uniform hyperparameters for different models and datasets to train with the following settings: (1) initial_ _ 0.05; (2) y 0.5; (3) 3 0.99; (4) X = 1. For the other the model parameter settings of the original paper. LookSAM uses the hyperparameters (p- = 0.0001, a = 0.2, k = 8) for training. SAM-k uses k=8, where k stands for step instead of epoch. And GraphSAM algorithm is shown in Algorithm 1
model parameters,
Algorithm 1 GraphSAM
Require: Network fo: Training set G; Batch size b; Neighborhood size p 0: Moving average
hyperparameter 3: Learning rate >> 0; p's update rate À > 0; The modification scale for epoch - 0.1,2. E do
2: Sample a mini-batch in B C G with size b.
3: for < 0. N do
4: if == 0: then
5:
= VoCx(())
6:
if epochh%>= = O: then
7:
Pnew Pinitial y opach/A end if
8:
9:
else
10:
c Bc (1 B) we-1 1|2
11:
end if
= p p- sign(es) TTelllz
leal
13:
Wi = VoLz (0:) b,tt
14:
l.+1 Os 17 We
15:
end for
end for
17: Return: 0r
Ensure: Model trained with GraphSAM.
"hats://githuh.com/// http:///ghhbbbmm///c/ccmaaaaattt Sttt::///tthaa..cm///cchhccc///t
18

---

Published as a conference paper at ICLR 2024
A.5 APPENDIX 5: SUPPLEMENTARY EXPERIMENTS.
In this section, we mainly add some experiments sections 5.2.2 (Throughput and accuracy of the optimizers) and 5.2.3 (Roles of GraphSAM modules) of the main paper.
WW compare the throughput and generalization performance of GraphSAM with other optimization algorithms in Table 7 Where throughput denotes the computational overhead which is quantified by graphs processed per second (Graphs/s), and the generalization performance is measured by test ROC-AUC/RMSE.
We can see that AE-SAM more stable compared RST and LookSAM. This depends on how high its throughput is, i.e., whether it computes the two gradients multiple times. When its throughput high, its performance will be relatively lower, and vice versa. GraphSAM, on the other hand, performs very consistently, improving both the throughput and the generalization performance of the model, which attributed to the three improvement parts of this paper. In addition, due to LookSAM and RST's own shortcomings for molecular graphs, their performance improvement is not sufficient and may even be counterproductive. For example, RST only close to the simple SAM-k method in most cases.
Table 7: Classification accuracy and training speed. The numbers in parentheses (-) indicate the ratio of GraphSAM's training speed w.r.t SAM.
Tox21
Sider
Clin
GROVER
ROC-AUCT
Throughput
ROC-AUCT
Throughput
ROC-AUCT
Throughput
Adam
0.822
131
0.641
130
0.853
161
SAM
0.840
81(100.0%)
0.660
79(00.00%
0.872
110(100.0%)
SAM-One
0.801
118(145.6%
0.610
111(141.3%)
0.818
150(136.
SAM-k
0.815
109(134.6%)
0.631
105(132.9%)
0.836
141((28.1%)
LookSAM
0.829
92(113.6%)
0.653
95(120.3%)
0.847
128(116.45)
AE-SAM
0.834
90(111.1%)
0.649
101
0.858
133(121.0%) 11.%%
RST
0.818
100(123.49%)
0.625
103(130.3%)
0.845
143(130.0%)
GraphSAM
0.846
98(121.0%)
0.665
98(124.1%%
0.866
134(121.8%)
CoMPT
ROC-AUCT
Throughput
ROC-AUCT
Throughput
ROC-AUCT
Throughput
Adam
0.828
119
0.621
219
0.914
211
SAM
0.839
60(100.0%)
0.643
114(100.0%)
0.927
115(100.0%)
SAM-One
0.788
102(170.0%
0.595
198(173.6%
0.879
201(174.7%)
SAM-k
0.819
94(156.6%)
0.610
180(157.8%%
0.895
187(162.6%)
LookSAM
0.825
78(100.0%)
0.625
141(123.6%)
0.916
(127.88.8%)
AE-SAM
0.833
(1118.33%
0.631
150(131.6%)
0.909
155(134.8%)
RST
0.815
94( (A(156.6%)
0.615
(1981148)))
0.890
180(156.5%)
GraphSAM
0.841
(5(11..7%
0.645
(13(124.2%)
0.937
160(139.1%)
ESOL
Lipophiliiity
GROVER
RMSE
Throughput 360
RMSEL
Throughput
Adam
0.639
0.671
160
SAM
0.619
223(100.0%)
0.662
68(100.0%)
SAM-One
0.671
332(148.9%)
0.715
143(210.3%)
SAM-k
0.650
298(133.69)
0.698
129(189.7%)
LookSAM
0.633
278(126. 1%)
0.680
90(132.4%)
AE-SAM
0.631
289(129.6%)
0.674
100(147.0%)
RST
0.638
305(130.9%)
0.690
116(170.1%)
GraphSAM
0.625
282(126.55)
0.654
95
CoMPT
RMSEJ
Throughput
RMSEL
Throughput
Adam
0.562
502
0.618
432
SAM
0.517
282(100.0%)
0.611
234(100.0%)
SAM-One
0.601
457(162.1%)
0.638
392(167.59)
SAM-K
0.575
413(146.5%)
0.630
367(156.8%)
LookSAM
0.543
338(119.9%)
0.621
318(135.9%)
AE-SAM
0.527
329(116.7%)
0.615
308(131.6%)
RST
0.588
391(138.6%)
0.633
(71(158.5%
GraphSAM
0.511
347 (123.000%%
0.608
(31(141.5%)
19

---

Published as a conference paper at ICLR 2024
We investigate the impact of different fixed-value P and gradient ball's size p schedulers the performance of various datasets. We list their performances in Table g
We find that proposed Gradient ball's size (p) schedule is more stable than the fixed (p), and climinates the need for extensive tuning to find the optimal rho.
Table 8: The influence of gradient ball's size on SAM and GraphSAM on the CoMPT.
Algorithm
BBBP/ROC-AUCT)
Tox21ROC-AUCC)
Sider(ROC-AUCT
0.05
0.957 + 0.010
0.822 = 0.011
0.621 0.018
SAM
0.005
0.952 0.022
0.842 0.008
0.635 0.021
0.001
0.953 0.021
0.835 0.013
0.641 0.029
scheduler
0.962 0.033
0.839 0.006
0.643 0.049
0.05
0.933 + 0.025
0.810 0.033
0.618 0.029
GraphSAM
0.005
0.949 + 0.020
0.826 + 0.020
0.625 + 0.018
0.001
0.964 + 0.008
0.838 + 0.006
0.636 + 0.016
scheduler
0.961 + 0.012
0.841 + 0.004
0.645 + 0.013
Algorithm
ClinTo((RO--UUCT)
(SOL(RMSEL)
T Lipophilicity(RMMFE)
0.05
0.910 + 0.031
0.539 + 0.031
0.622 1: 0.016
SAM
0.005
0.918 + 0.016
0.534 + 0.021
0.601 + 0.009
0.001
0.924 0.018
0.527 + 0.014
0.613 + 0.018
scheduler
0.927 + 0.025
0.517 + 0.025
0.611 = 0.015
0.05
0.902 0.045
0.946 0.103
0.691 0.053
GraphSAM
0.005
0.924 0.020
0.598 0.212
0.625 = 0.018
0.001
0.931 0.008
0.528 0.016
0.610 0.016
scheduler
0.937 + 0.008
0.511 + 0.014
0.608 + 0.007
>> Accuracy-efficienoe trade-off. The secret of GraphSAM's ability to maintain similar perfor- mance to SAM the timely re-anchor of 6₁. The effectiveness of another module of GraphSAM, moving average, has been shown in Table In this part, we mainly verify the correlation between accuracy-efficieney We propose the GraphSAM-K to investigate the impact of the rate of re-anchor to the perturbation gradient e, generalization performance. The K means that we perform additional forward and backward propagation to re-anchor from 6, to Eo for every K epoch. As in Fig. 66] we analyze GraphSAM-K K for different values of K. The N means that only once re-anchor of Gt- When K = GraphSAM is comparable to SAM. When = 2, the performance GraphSAM slightly higher than Adam's. When K 2, the efficiency of GraphSA. is increasing rapidly, but its performance drops sharply. The reason is that as the training steps increase, the error between é obtained from equation and the ground-truth és becomes increasingly large. To reduce the error, we need to re-anchor the E by forward and backward propagation in each epoch repeatedly.
97
Accuracy-Training CoMPT on
10.0
93.0
Accuracy-Training Time Grover on
6.2
96
accuracy
9.5
92.5
accuracy
6.0
95
time
9.0
Average fimets 0
92.0
time
5.8
Accuracy
91.5
94
8.5
5.4
93
8.0
90.5
5.2
92
7.5
5.0
90.0
91
7.0
89.5
4.8
4.6
90
6.5
89.0
Figure 6: Accuracy-Training Time of different models for GraphSAM-K. Average Time (s/epoch) represents the average time consumed for each epoch.
20

---

Published as a conference paper at ICLR 2024
A.6 APPENDIX 6: HYPERPARAMETER STUDIES.
We conduct experiments on some hyperparameters introduced by GraphSAM in the training process.
@ The appropriate size of in GraphSAM is essential for model training. We conduct the experiment and observe that the size of p is related to the initial size p. p's update rate X, and the scheduler's s modification scale 7. To effectively evaluate the impact of the size of p. the default values of the nyperparameters this paper are initial p 0.05, À 1. and ? = 0.5. In Fig. Zh. we find that if the initial p is small, the weight perturbation has smaller range, which eventually affects the training performance of GraphSAM. A is big (i.e., p updated more slowly), the perturbation range of GraphSAM deviates severely from the SAM, causing the model generalization performance to drop sharply as shown in Fig. 17p. Similarly, when the modification scale of p is faster i.e. ? 0.1, the weight perturbation range spans a larger extent, making the performance of GraphSAM scaled down, especially in the ESOL dataset When the modification scale of p faster i.e. y 0.1, the weight perturbation range
addition, if the initial / = 0.05 large and the update rate
in Table 9 spans a larger extent, making the performance of GraphS AM scaled down, especially in the ESOL dataset in Table 9
Table 9; The influence of the gradient ball's size p scheduler's modification scale y GraphSAM algorithm with COMPT.
Algorithm y
BBBP (ROC-AUC1)
ClinTox(RO-AUUC)
ESOL (RMSE))
Lipophilicity (RMSE))
0.5
0.961 0.012
0.937 0.008
0.511 0.018
0.608 0.007
GraphSAM
0.2
0.959 = 0.002
0.923 + 0.015
0.560 + 0.010
0.618 = 0.019
0.1
0.956 = 0.005
0.911 + 0.011
0.608 + 0.012
0.634 = 0.010
fil sze lpl the gradient ball
update rate (N
85.6
65.6
25.4
45.4
95.2
95.2
95.0
05.0
94.
0.25
0.01
0.005
0.001
Figure 7: Hyperparameter studies of p, and À of CoMPT+ GraphSAM on BBBP dataset.
A.7 APPENDIX 7: SUPPLEMENT TO THE PRE-TRAINING EXPERIMENT.
We have also conducted experiments on the pre-training method, and both SAM and GraphSAM have certain performance improvements. However, compared to non-pre con-pre-training experiments, the improvement is limited. As shown in table TO, we provide the experimental results of SAM and GraphSAM on the pre-trained model GROVER with six datasets.
Table 10: Prediction results of GraphSAM on pretraining- GROVER on six datasets. We used 5-fold voss-validotion with random split and replicated experiments on each task five times. The mean and standard deviation of AUC or RMSE values are reported.
Task
Graph Classification (ROC-AUC1)
Graph Regression (RMSE.
Dataset
BBBP
Tox21
Sider
Clin'Tox
ESOL
Lipophilicity
GROVER
0.930
0.83340.028
0.661+0.030
0.93140.011
0.627 00.07
0.544g0.039
SAM
0.935
0.
0.669
0.940+0.018
0.623
0.537
GraphSAM
0.93940.011
0.838 00003
0.673 0.006
0.93844.0020
0.61940.111
0.529
21

---

Published as a conference paper at ICLR 2024
A.8 APPENDIX 8: ADDITIONS THE OBSERVATION OF GRADIENT VARIATION.
In this section, we address the phenomenon detailed as Observation I in the main text, as illustrated in Fig. This addition robustly affirms that the variations observed in the perturbation gradient and the updating gradient in Observation I are not isolated instances, but rather general trend. Furthermore, we incorporate the trajectory of changes in the magnitudes of ||ee|12 and ||w (wxll2 as depicted in Fig. 9 This inclusion substantiates our finding ||ell2 ||</|12 confirming the initial observation and providing deeper insight into the underlying dynamics play.
Tex22
Slee
Ohtss
Bosr <dl
1.5
1..2
3..00
0.0
100
200
320
603
300
950
105
120
200
300
430
350
803
700
102
205
200
400
500
600
700
Figure 8: Illustration on the observation of gradient variation during training on GROVER with three datasets.
Tox22
Grvdent
Stop
Stpp
Figure 9: Illustration on the observation | and ||well2 during training on GROVER with three datasets.
Regarding the attention matrix the CoMPT model on BBBP dataset, observe that models optimized with Adam, Gra and SAM obtain average entropies 2.4130, 2.7239, and 2.8382, respectively. That means the application of SAM approaches can smooth the attention scores avoid overfitting. The attention heat map is shown in Fig. 10
Alzention Heatmap (SAN)
(GraphSAM
28
6/40
43
16
0.32
14
A1
Figure 10: Attention heat map.
In CV, the model's perturbation gradient variations are often consistent with updating gradients, which is not line with the phenomena on the molecular graphs. To further address the concern and provide holistic discussion, add the following experiments to compare the performance of lookSAM, AE-SAM, and GraphSAM in the CV and molecule datasets, respectively. LookSAM and GraphSAM efficient SAMs designed based the gradient varification patterns of models in their respective domains, and AE-SAM is an adaptively-updattee
efficient method according to the comparison evaluation between gradient norm and a pre-defined threshold.
22

---

Published as a conference paper at ICLR 2024
Table 11: Performance of various SAM methods in the CV domain.
ResNet-18 (CV)
CIFAR-10
CIFAR-100
+SGD
95.41
78.21
+SAM
96.53
80.18
*LooKSAM
96.25
79.91
+AE-SAM
96.49
80.31
+GraphSAM
95.86
78.69
As shown in Fig 11 and Fig. 7 it is observed that LookSAM has better results on the image datasets instead 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 across the different domains. is challenging 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 monitored inform the necessary gradient re-computation. But is not good as GraphSAM where we accurately the perturbation gradients at each step. In summary, GraphS is optimized particularly for graph transformers based the empirical observations of gradient variations.
23


Reviews:
[{"review_id": "psYtMk8JtI", "reviewer": "Reviewer_QvPb", "summary": "The authors study Sharpness Aware Minimization (SAM), an effective training objective in the image domain to avoid \"sharp minima\", in the context of graph transformers for molecular data. The proposed efficient SAM-alternative, termed GraphSAM, empirically performs on par with SAM while being more efficient.", "questions": "1. Is GraphSAM only a good approximation for graph transformers? How is GraphSAM working, e.g., in the image domain?\n1. 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)?\n1. Is it possible to compare GraphSAM with other efficient SAM derivates besides SAM-k (e.g., see related work section)?\n1. 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?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. More efficient approach for SAM (roughly 30%-50% faster training)\n1. GraphSAM consistently improves the performance of the base model\n1. Approach is ablated and rich explanations are provided\n1. The paper is well organized and it is easy to follow.", "weaknesses": "1. 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).\n1. The theoretical statements could be better motivated and more decisively embedded into the story. It remains a bit vague what the actual implications are. For example, the conclusion of Theorem 1 is speculative \"by replacing the inner maximum of LG(θ + εˆS) with the maximum of LG(θ + εˆG), we tend to smooth the worse neighborhood loss in the loss landscape. In other words, if one could minimize the upper bound given by LG(θ + εˆG), the final loss landscape is as smooth as that obtained from SAM.\" Alternatively, the wording may be improved. It is not clear why optimizing the upper bound necessarily results in a \"loss landscape is as smooth as that obtained from SAM\".\n1. 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.\n\nMinor:\n1. Missing space in the first line of page 2\n1. The bold highlighting in Table 1 is counter-intuitive. Perhaps add an additional marker to highlight the best model per task.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QwAx1VYbUd", "reviewer": "Reviewer_sT1K", "summary": "The study introduces Sharpness-aware minimization (SAM) in computer vision, a technique effective in eliminating sharp local minima in training trajectories and countering generalization degradation. However, its dual gradient computations during optimization increases time costs. To address this, a novel algorithm named GraphSAM is proposed. It lowers the training expenses of SAM while enhancing the generalization performance of graph transformer models. The approach is underpinned by two main strategies: gradient approximation and loss landscape approximation. Empirical tests across six datasets validate GraphSAM's superiority, particularly in refining the model update process. Anonymized code is also provided for further reference.", "questions": "see weakness", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "(1) The paper is exceptionally coherent in its writing. The content is presented seamlessly, and the expression is clear.\n\n(2) This paper is both conceptually and technically innovative. It uniquely reutilizes the updating gradient from the previous step, approximating the perturbation gradient in an efficient manner.\n\n(3) The experimental design of this study is well-conceived. It systematically validates the representational capabilities of GraphSAM. Additionally, the appendix offers an extensive set of supplementary experiments.", "weaknesses": "(1) 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?\n\n(2) The results presented in Table 1 seem to show limited improvement on certain datasets. I suggest the author consider incorporating an efficiency aspect into the table. Efficiency is a crucial contribution, yet it hasn't been emphasized adequately in the experimental section, which seems inconsistent.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UMiLyK5F8c", "reviewer": "Reviewer_mSQh", "summary": "The paper proposes GraphSAM to reduce the training cost of sharpness-aware minimization (SAM) and improve the generalization performance of graph transformer models. GraphSAM uses the updating gradient of the previous step to approximate the perturbation gradient at the intermediate steps smoothly and theoretically proves that the loss landscape of GraphSAM is limited to a small range centered on the expected loss of SAM. Extensive experiments on six datasets with different tasks demonstrate the effectiveness and efficiency of GraphSAM.", "questions": "see the questions in weakness.\n\n======\n\nI have read the rebuttal and other reviews. Most of my concerns are resolved, so I maintain my score of 6.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- **Motivation is clear**: SAM is a powerful optimizer but doubles the computational cost compared with the base optimizer. Thus, **improving its efficiency** (the focus of this paper) is an important problem\n- GraphSAM improves the efficiency of SAM by **approximating the perturbation gradient** $\\epsilon_{t+1}$ as the moving average of $\\epsilon_t$ and $\\omega_{t}$ based on two observations: (i) $\\epsilon_t$ is close to $\\epsilon_{t+1}$ (Figure 3(a)) and (ii) $\\epsilon_{t+1}$ is close to $\\omega_t$ (Figure 3(c))\n- By approximating the perturbation gradient, GraphSAM almost **does not need to compute the perturbation gradient** (only once per epoch). Thus, GraphSAM is efficient\n- Experimental results on the graph dataset show that GraphSAM is comparable with SAM (Tables 1 and 2) and is more efficient (Table 2)", "weaknesses": "- **writing**:\n  - 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)?\n  - Figure 1, it is better to include the proposed GraphSAM as well\n  - 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\n  - 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 why they cannot be used for graph datasets, this is also the motivation for the proposed GraphSAM\n  - Figure 3, sub-caption for each subfigure\n  - Table 1: also compare with GROVER+AE-SAM/SS-SAM\n  - \"AE-SAM and RST periodically compute the update gradient by different strategies\": AE-SAM is an adaptive strategy, not a periodic strategy\n- 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.\n- ablation study for hyperparameters $\\gamma, \\lambda, \\beta$\n- Theorem 2, based on the proof, the proportion ratio depends on the $\\theta$ (Eq(7)), it is trivial. if we assume the loss is Lipschitz, then the constant can be independent of $\\theta$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8DZ37Unp5y", "reviewer": "Reviewer_1va4", "summary": "This paper designs a computation efficient variant of sharpness aware minimization (SAM) for molecular graph transformers.", "questions": "Q1. Table 1 does not fit into the paper properly.\n\nQ2. Some of the expressions are unnecessarily complicated. For example, in equation (2) and (3), $sign(\\epsilon_t)|\\epsilon_t|$ can be replaced by $\\epsilon_t$.\n\nQ3. Missing references on SAM variants. \n\n[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).\n\n[2] Li, Bingcong, and Georgios B. Giannakis. \"Enhancing Sharpness-Aware Optimization Through Variance Suppression.\" arXiv preprint arXiv:2309.15639 (2023).", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "S1. This paper observes that flat minima generalize better for molecular graph transformers. This extends the boundary of sharpness awareness to a new domain.\n\nS2. To overcome the computational burden, this paper leverages domain specific knowledge to design GraphSAM. Supporting theories are derived to explain the reasons behind design.\n\nS3. GraphSAM improves throughput with comparable performance with SAM. Compared to other variants of computation efficient SAM, GraphSAM performs better in the considered setting.", "weaknesses": "W1. 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$?\n\nW2. 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)