Sharpness-aware Zeroth-order Optimization for Graph Transformers

1 Introduction

Graph Transformers (GTs) represent a significant advancement in machine learning, introducing an emerging class of network architectures purposefully designed to handle graphstructured data with considerable efficiency and flexibility [Dwivedi and Bresson, 2020; Liu et al., 2024a]. In contrast to conventional graph neural networks (GNNs), GTs have demonstrated a strong capacity to address key issues such as over-smoothing and over-squashing [Topping et al.,

∗Corresponding author

2021]. The incorporation of transformer-based architectures into graph models has proven to be highly adaptable and powerful. Notably, this approach has yielded remarkable empirical results in domains as varied as drug discovery [Liu et al., 2022a], where intricate molecular interactions must be accurately modeled, and molecular structure prediction [Rong et al., 2020], which demands precise representation of large, interconnected systems. The adaptability, robustness, and improved representational power of GTs drive ongoing research and offer significant potential for advancing various machine learning applications [Liu et al., 2024b; Chang et al., 2025].

Motivation. Despite extensive efforts to develop advanced optimizers for Transformer-based architectures, a fundamental challenge arises from the presence of certain complex operators within GTs, such as shortest-path distance computations, which are inherently non-differentiable. These operators preclude the straightforward application of conventional gradient-based optimization techniques (e.g., ADAM) [Ying et al., 2021]. In response to this challenge, Zeroth-Order Optimization (ZOO) has emerged as a promising alternative, providing a means of gradient estimation without direct differentiation. Such approaches have recently gained prominence in addressing optimization challenges in large-scale Transformer-based language models [Zhang et al., 2024]. However, the naive integration of ZOO methods in GTs faces substantial challenges in inferior and non-robustness, particularly as model complexity and dimensionality increase [Kunstner et al., 2023]. These challenges arise from the steep gradients in the parameter space of GTs and the sharp loss landscape encountered during ZOO training [He et al., 2022], where convergence to sharp local minima can significantly impair generalization [Foret et al., 2020] and overall performance. Furthermore, we conduct a demonstration experiment which illustrates the poor performance of vanilla ZOO (see Figure 2). Consequently, novel strategies are in urgent need to leverage ZOO principles in GTs. To address the challenges associated with the loss landscape in ZOO, we consider incorporating Sharpness-Aware Minimization (SAM) [Foret et al., 2020], which aims to find local minima with a flatter neighborhood and enhance the robustness of the optimizer. In technique, SAM updates the model parameters using two gradients (one from the neighborhood and one from the perturbed parameters), minimizing the average loss over the perturbed parameters, thereby enhancing generalization by reducing sensitivity to small variations in the parameter space.

In this work, we propose a Sharpness-aware Zeroth-order Optimizer (SZO) for graph transformers which basically incorporates SAM technique into the ZOO training process. By regulating the curvature of the loss landscape, SAM effectively steers the optimization process toward flatter minima, which in turn promotes improved generalization performance and heightened resilience against adversarial perturbations. Besides, we introduce a novel parallelized variant of the Randomized/Coordinate-wise Gradient Estimator which can reduce the computational burden. For the whole optimization framework, we present a thorough theoretical analysis from both the convergence properties and generalization guarantees. Empirically, we conduct extensive experiments involving a wide range of publicly available datasets and evaluate performance across various representative graphoriented tasks. The results consistently demonstrate the effectiveness and versatility of the proposed methods, emphasizing their potential to drive advancements in the state-of-the-art in graph transformer optimization.

Our contributions are summarized as follows:

1. Sharpness-aware Zeroth-order Optimization. We are the first to investigate the training performance of zerothorder optimization for graph transformers and propose a sharpness-aware zeroth-order optimization (SZO) to enhance the robustness and generalization capabilities of the model.
2. Comprehensive Theoretical Analysis. Our analysis provides an in-depth theoretical examination, addressing aspects of both generalization and convergence. We prove the convergence of this combined training approach and demonstrate that the generalization bound of SZO is confined to a neighborhood-wise loss.
3. Experimental Evaluation. We evaluate SZO on various classical graph transformers and a wide range of public benchmark datasets. Extensive experimental results demonstrate the effectiveness in achieving superior performance

compared to several modern optimizers and SAM variants. 1

2 Mathematical Background

Notation. Given a training dataset S ≜ ⋃ni=1 {(xi, yi)} drawn i.i.d. from distribution D = X × Y, we aim to learn a graph transformer model that generalizes well. In particular, consider a family of models parameterized by θ ∈ Θ ⊆ Rd . Given a per-data-point loss function l : Θ×X ×Y → R+, we define the in-sample (empirical) training loss LS(θ) ≜ Lin(θ) = 1n ∑ni=1 l(w, xi, yi) and the out-of-sample (theoretical) loss LD(θ) ≜ Lout(θ) = E(x,y)∼D[l(θ, x, y)]. Having observed only S, the goal of training is to learn model parameters θ with lower loss LD(θ). For graph-level task, the input training data includes graphs {G(i) }Ni=1 (each graph G(i) = (X(i) , A(i) )) and associated labels {y(i) }Ni=1 .

2.1 Graph Transformer

In this section, we present our focused framework, Graph Transformers (GTs). GTs enable each node in a graph to attend to all other nodes through a global attention mechanism. The architecture of GTs consists of two primary components: a self-attention module and a subsequent feed-forward neural network (FFN). In self-attention module, the input feature matrix X ∈ RN×d is originally projected into three spaces: query matrix Q ∈ Rd×dq , key matrix K ∈ Rd×dk and value matrix V ∈ Rd×dv (where dq = dk):

Then we can compute the self-attention score as:

The output of the self-attention module is followed by a skipconnection and a feed-forward network (FFN), which jointly compose a Transformer block:

where W1 and W2 denote the parameters in FFN.

2.2 Zeroth-order Optimization

Zeroth-order optimization (ZOO) has gained significant attention due to its applicability in scenarios where the gradient of the objective function is not readily available or difficult to compute. In this section, we fisrt introduce two classical gradient estimators which are commonly employed to approximate the gradient of a target function. The first one, known as the Coordinate-wise Gradient Estimator (CGE), evaluates the partial derivatives of the function with respect to each separate coordinate. More formally, consider a differentiable function f : Rn → R, and let θ ∈ Rn be the point where we aim to approximate the gradient ∇f(θ). The CGE method approximates each partial derivative ∂f ∂θi by perturbing θ along the i-th coordinate direction. Specifically, it computes gi(θ) ≈ f(θ+ϵei)-f(θ) ϵ , where ϵ > 0 is a small constant and ei represents the i-th standard basis vector in Rn . By performing this operation for each coordinate i of θ, we obtain an approximation of the gradient ∇f(θ). The formal CGE update procedure is thus defined by evaluating the above finite difference estimates for all coordinates and then assembling them into an approximate gradient vector as:

1 https://github.com/liu-yang-maker/SZO

where ei denotes the standard basis vector. Another zerothorder gradient estimation approach is the Randomized Gradient Estimator (RGE), which offers a more flexible alternative to coordinate-wise perturbations by sampling random directions in the parameter space as:

where ui represents a random direction vector, generally sampled from the standard Gaussian distribution N(0, I). The parameter q denotes the number of function queries, while ϵ > 0 signifies a small perturbation step-size, also referred to as the smoothing parameter.

2.3 Sharpness-aware Minimization

The sharpness of the loss landscape in GTs under ZOO [Wang et al., 2024] necessitates the introduction of an alternative optimization technique. To address this issue, we aim to flat the neighborhood of local minima, thereby enabling the development of a more robust and effective model. In this work, Sharpness-Aware Minimization (SAM) is utilized as a proposed solution [Foret et al., 2020].

SAM Problem. Inspired by bounding generalization ability in terms of neighborhood-wise training loss, [Foret et al., 2020] proposed Sharpness-Aware Minimization (SAM) problem as follows:

where LSAM S (θ) ≜max∥ϵ∥p≤ρ LS(θ+ϵ), ρ ≥ 0 is a hyperparameter and p ∈ [1,∞].

Minimize LSAM S (θ). An efficient approximation of its gradient has been determined in as (details in Appendix)

For each update step t, we seek to find the best update ∆ ∈ Rd for θt+1 = θt - η∆, where η ∈ R+ denotes the learning rate. Then, we aim to solve the following optimization problem respect to ∆:

where δ represents the maximum step size and f refers to the original parametric model under consideration.

3 Proposed Approach

In this section, we describe SZO in details. First we give an overview procedure for SZO. We then describe how to design in technical view from the original version.

3.1 Proposed Optimization Procedure: SZO

Building on the SAM framework and the optimization challenges outlined earlier (see Eq.2 and Eq.3), we propose the Sharpness-aware Zeroth-order Optimization (SZO) scheme (here we take CGE as an example):

where ρ is the radius of the gradient stepping ball and η is the step size. A similar scheme with RGE can be derived by replacing the calculation of ϵt and ˆ∆t above.

The training procedure of the SZO algorithm is outlined in Algorithm 1. The algorithm begins by initializing the model parameters θ (line 1). For each epoch t from 1 to MaxIter (line 2), a batch of graphs {(x1, y1), . . . , (xb, yb)} is sampled from the training dataset Dtrain (line 3). In the First Differentiation Step, the approximate gradient ϵt is computed using a coordinate-wise (or randomized) gradient estimator, achieved by perturbing each coordinate independently (line 4). The perturbed model parameters θϵ t are then calculated (line 5). In the Second Differentiation Step, an additional gradient approximation ˆ∆t is computed (line 7). At the end of each epoch, the model parameters are updated by applying the computed gradient approximation via gradient descent (line 11).

3.2 Optimization Design

Given the huge complexity and parameter dimensions of graph transformers, traditional ZOO methods face efficiency bottlenecks. We implement algorithm acceleration using forward parallelization [Chen et al., 2023]. Forward parallelization decouples parameter perturbations and leverages distributed computation, significantly reducing training time while maintaining performance. This acceleration is essential to meet the practical demands of large-scale graph data processing.

Decoupling Parameter Perturbations. For each coordinate in ZOO Perturbation, we select the active parameter to conduct the corresponding independent forward pass to allow for parallel computation. A demo procedure of parameters divided by θQ, θK, θV (refer toWQ,WK,WV ) can be described as

where I ∈ {Q,K,V } and SI denote the active parameters assigned to module I. In this case, only |SI| forward passes are taken in each case.

Distributed and Sparse Implementation. The computation is distributed across multiple processes (or GPUs). Each process is assigned a subset of the parameters to perturb and compute the corresponding forward passes. By distributing the parameter perturbations and forward passes across multiple machines, the training process is significantly accelerated. This approach avoids the limitations of traditional data parallelization, such as the potential for performance loss with large batch sizes. Furthermore, we use a top-K operator before each variable matrix in attention as Y = (X⊙M) ·WT , where M ∈ {0, 1}N×D denotes Topk (|X|) which is the mask tensor that indicates the top-K activations in the input tensor X with absolute values and ⊙ is the Hadamard product. This enhances the efficiency of SZO by ensuring high computational efficiency while maintaining performance.

3.3 SAM-based Design

We employ the Sharpness-Aware Minimization scheme (SAM, see Section 2.3) [Foret et al., 2020; Wang et al., 2024] for zeroth-order optimization.

SAMscheme. We introduce a SAMscheme in our method for better generalization [Wang et al., 2024]. As discussed in Section 2.3, we perform the gradient estimation operation twice, regardless of whether using CGE, RGE, or a combination of both. This training scheme improves the model's generalization by facilitating a flatter loss landscape.

Perturbation size scheduler. ϵi+1 = ϵ0 · γ epoch/λ where the hyperparameter λ determines the update frequency of ϵ, specifying the number of epochs per update, while γ regulates the scale of modifications in ϵ.

3.4 Theoretical Analysis

The training pipeline of SZO is illustrated in Algorithm 1. In this section, we provide a justification for the convergence and generalization ability of SZO to ensure a reliable training process. Note that the theorems only depend on the vanilla setting and ignore the sparsity improvement.

Convergence Justification

Theorem 3.1. Considering an α-strongly convex and βsmooth loss function ℓ, if we run SZO optimizer starting at θ0 with perturbation radius ϵ > 0 and any learning rate η < 2β to minimize ℓ, we have:

Generalization Justification

The original SAM [Foret et al., 2020] approach was inspired by the established connection between the sharpness of the loss landscape and the model's generalization capabilities. To further elucidate this motivation, the following analogous theorem demonstrates how generalization ability can be bounded by neighborhood-wise training loss:

Theorem 3.2 (informal). For any ϵ > 0 and any distribution D, with probability 1 - δ over the choice of the training set S ∼ D, we have

where h : R+ → R+ is a strictly increasing function under some technical conditions on ℓD and we assumed ℓD(θ) ≤ Eϵ∼N (0,δ′ )[ℓD(θ + ϵei)].

Remark 1. Theorem 3.1 establishes a solid foundation for the reliability and effectiveness of the proposed SZO method. The convergence results demonstrate that SZO achieves guaranteed progress under both strongly convex and smooth loss functions, ensuring stable optimization with appropriate learning rates. The generalization analysis (see Theorem 3.2) highlights the role of the SAM framework in achieving flatter minima, which contributes to better robustness and improved generalization capabilities. Together, these theoretical results validate the efficiency of SZO for optimizing complex graph transformer models.

4 Experiment

4.1 Experimental Setup

Dataset. Following established settings inmolecular graph tasks, we utilize nine public benchmark datasets: BBBP, Tox21, Sider, ClinTox and BACE for classification, and ESOL, Lipophilicity, QM7 and QM8 for regression. We assess all models using a random split methodology as recommended by MoleculeNet [Wu et al., 2018], dividing the datasets into training, validation, and testing sets with an 80%/10%/10% ratio. In Table 1, we summarize the statistics of the molecular graph datasets utilized in graph classification and graph regression tasks. Details see Appendix.

Baseline. We evaluate the effectiveness of SZO on two graph transformer models including GROVER and CoMPT [Rong et al., 2020; Chen et al., 2021] with different modern optimizers [Wang et al., 2024]. Also, we add several competitive GNN-based models as baselines.

• GROVER [Rong et al., 2020], which stands for Graph Representation frOm self-superVised mEssage passing tRansformer, is a novel framework designed to address the challenges in molecular representation learning. It leverages self-supervised tasks at the node, edge, and graph levels to learn rich structural and semantic information from unlabelled molecular data. GROVER integrates Message Passing Networks into a Transformerstyle architecture, enhancing its capability to encode complex molecular information.
• CoMPT [Chen et al., 2021]. The Communicative Message Passing Transformer (CoMPT) is a neural network designed to improve molecular graph representation. It enhances traditional Graph Neural Networks by strengthening interactions between nodes and edges using Transformer architecture. CoMPT includes a mes-

sage diffusion mechanism to manage graph connectivity and prevent information overload, capturing both local and global structure.

Implementations. The implementations of the backbone models and their respective hyperparameter configurations are sourced from publicly available repositories as detailed in [Rong et al., 2020] and [Chen et al., 2021]. Both GROVER and CoMPT employ Adam as the base optimizer without employing any pre-training strategies. In our experiments, we solely adjust the hyperparameters introduced by SZO.

We implemented several optimizers (i.e. SGD, ADAM and SZO) on molecular graph data using two widely-adopted graph transformer backbones: GROVER [Rong et al., 2020] and CoMPT [Chen et al., 2021]. A comprehensive comparison was conducted against several baseline models including GCN [Kipf and Welling, 2016], MPNN [Gilmer et al., 2017], DMPNN [Yang et al., 2019] CMPNN [Song et al., 2020], SPMM [Chang and Ye, 2024], PretrainGNN [Hu et al., 2020], and MolXPT [Liu et al., 2023b].

SAM and its Variants. Recent advancements in Sharpness-Aware Minimization (SAM) have introduced several efficient variants. The original SAM [Foret et al., 2020] prevents sharp minima convergence but doubles computational costs. LookSAM [Liu et al., 2022b] improves efficiency by intermittently computing gradients but neglects gradient magnitude, limiting its generalizability. AE-SAM [Jiang et al., 2023] uses squared gradient norms to skip updates, though its performance varies across domains. RST [Zhao et al., 2022] employs Bernoulli trials for stochastic updates, yielding suboptimal results. GraphSAM [Wang et al., 2024] enhances efficiency for molecular graph transformers by reusing prior gradients, maintaining strong generalization.

4.2 Performance on Graph Classification Task

We employed SZO on GROVER and CoMPT for an extensive evaluation against baseline methods across various datasets. The comparative results are presented in Table 2. It demonstrates that the impact of integrating SZO with GROVER and CoMPT models on graph classification tasks across six datasets (BBBP, Tox21, Sider, ClinTox, BACE). The SZOenhanced versions of both models consistently outperform their baseline counterparts, achieving some of the highest ROC-AUC scores, particularly on BBBP, Tox21, and ClinTox datasets. Notable improvements include GROVER with SZO reaching 0.929 in BBBP and CoMPT with SZO achieving 0.940 in ClinTox. This highlights the effectiveness of SZO in boosting model performance.

4.3 Performance on Graph Regression Task

Table 3 displays the performance of various models on graph regression tasks, evaluated using RMSE and MAE across four datasets (ESOL, Lipophilicity, QM7, QM8). The integration of SZO with GROVER and CoMPT models generally leads to improved results. For instance, CoMPT with SZO achieves the lowest RMSE in ESOL (0.505) and the best MAE in QM8 (0.0139). Similarly, GROVER with SZO shows significant improvements, notably achieving the best RMSE in Lipophilicity (0.588). These results highlight the effectiveness of SZO in enhancing model accuracy for graph regression tasks.

4.4 Additional Verification

Ablation Study. The performance results of the ZO and SZO optimizer across four datasets are presented in Figure 2. It is evident that the ZO optimizer alone yields subpar outcomes, demonstrating a lack of effectiveness in achieving optimal performance on the datasets. However, when SAM is integrated into the optimization process, there is a substantial enhancement in performance. The addition of SAM helps to address the limitations of the ZO optimizer, significantly boosting the model's accuracy and stability across all datasets, thus showcasing the considerable benefits of combining SAM with the ZO optimization approach.

Different Optimizers. Table 4 shows the performance of GROVER and CoMPT models with various optimizers across different datasets for both graph classification (ROC-AUC)

and graph regression (RMSE). For GROVER, the RSZO optimizer consistently achieves high scores, such as 0.929 in BBBP and 0.598 in ESOL. Similarly, for CoMPT, RSZO yields excellent results, including 0.964 in BBBP and 0.505 in ESOL. Figure 2 visualizes the total ROC-AUC scores for these models and datasets, clearly indicating that RSZO (purple) and SZO (pink) optimizers lead to superior performance, particularly in the datasets of BBBP, ClinTox, and Tox21. This demonstrates the effectiveness of SZO in optimizing model performance across various tasks.

5 Related works

Graph Transformer. The remarkable achievements of Transformers [Vaswani et al., 2017] in natural language processing (NLP), and more recently in computer vision [Han et al., 2021] and biological information [Bai et al., 2023; Niu et al., 2022; Wu et al., 2022], have generated substantial interest in adapting Transformers for graph data analysis. One of the widely recognized benefits of graph Transformers compared to message-passing neural networks (MPNNs) is their capacity to capture long-range interactions, effectively addressing challenges such as over-smoothing and over-squashing. The introduction of the fully-connected Graph Transformer [Dwivedi and Bresson, 2020] utilized the eigenvectors of the graph Laplacian for node positional encoding (PE). Subsequent research has introduced numerous PE methods to advance Graph Transformers. These include the invariant aggregation of Laplacian eigenvectors in SAN [Ranzato et al., 2021], pair-wise graph distances in Graphormer [Ying et al., 2021] and more. Future research on Graph Transformers includes integratingMPNNs, improving substructure encoding, and optimizing for directed graphs.

Zeroth-order Optimization. A subclass of gradient-free optimization known as zeroth-order (ZO) optimization [Liu et al., 2020; Chen et al., 2023; Liu et al., 2023a] is used in many machine learning domains such as large language model. It is employed to handle optimization problems similar to gradient-based techniques. However, it only uses function values and does not require a gradient. The potential uses of ZO optimization are showcased, including but not limited to the assessment of the resilience of deep learning models [Chen et al., 2017; Wang et al., 2020] and the production of justifications from implicit systems, as well as the effective management of sensors in real-time scenarios.

SAM Optimization. Sharpness-Aware Minimization (SAM) has demonstrated significant effectiveness across diverse domains, including image classification [Foret et al., 2020; Kwon et al., 2021; Du et al., 2022; Wang et al., 2024; Ilbert et al., 2024] and natural language processing [Bahri et al., 2021]. [Foret et al., 2020] proposed SAM during training to improve standard generalization. In particular, LookSAM [Liu et al., 2022b] improves computational efficiency by eliminating the requirement to compute the updating gradient. [Du et al., 2022] proposed SAF, an innovative trajectory loss designed to mitigate sudden decreases in loss at sharp local minima during weight updates, thereby reducing the associated time loss. Additionally, the lack of SAM research in the graph domain necessitates the development of GraphSAM [Wang et al., 2024], designed to preserve SAM's generalization abilities while simultaneously improving computational efficiency.

6 Conclusion

In this work, our exploration into Sharpness-aware Zerothorder Optimization (SZO) for graph transformers has revealed significant advancements in model performance and robustness. By integrating SAM and parallelized gradient estimation techniques, SZO has demonstrated robust generalization capabilities without sharp loss landscape and better performance across various GT models. This study underscores the potential of SZO to transform optimization practices in graphbased neural networks.

Acknowledgements

The AI-driven experiments, simulations and model training were performed on the robotic AI-Scientist platform of Chinese Academy of Science. This work was partially supported by the Strategic Priority Research Program of the Chinese Academy of Sciences (No. XDB0680101), the National Natural Science Foundation of China (No. 62472416), and the CAS Project for Young Scientists in Basic Research (No. YSBR-008).