Gradformer: Graph Transformer with Exponential Decay

1 Introduction

Graph Transformers (GTs) [Dwivedi and Bresson, 2021] have shown remarkable success in achieving state-of-the-art performance across various applications. Unlike the local message-passing in graph neural networks (GNNs), GTs can capture long-range information from distant nodes. Specifically, the self-attention mechanism in GTs allows each node to directly attend to other nodes in a graph, enabling information aggregation from arbitrary nodes.

∗Corresponding Author

The self-attention in GTs offers considerable flexibility and the capacity to aggregate information both globally and adaptively. However, this mechanism significantly neglects intrinsic inductive biases in graphs, which poses challenges in capturing the essential graph structural information. Without considering structural relationships, the indiscriminate attention to all nodes in a graph will render the self-attention mechanism's inadequate focus on key information and the aggregation of redundant information. This comes with the generation of meaningless attention scores and, ultimately, suboptimal learning outcomes. Furthermore, the above issue is particularly evident in scenarios with limited data.

Numerous studies have focused on integrating graph inductive biases into self-attention learning mechanisms. These studies are broadly divided into two main approaches: 1) Injected position encoding: Studies such as GT [Dwivedi and Bresson, 2021] and SAN [Kreuzer et al., 2021] suggest using laplacian eigenfunctions as positional encodings (PEs) to contain the structural characteristics of graphs. 2) Attention bias: Graphormer [Ying et al., 2021] proposes a direct addition of structural encodings into the attention mechanism as a bias, enhancing the GT's capacity to model graph-structured data. Notably, applying PEs, which are typically concatenated with input features, would affect attention scores and could be considered as a way to introduce attention bias, as illustrated in Figure 2. However, concatenated PEs and attention bias fail to directly and effectively guide attention scores to fully capture structural information in graphs. From the Figure 1, it is evident that the attention patterns in GTs augmented with PE or attention bias exhibit minimal deviation from those observed in the vanilla GTs. This observation suggests that the inductive biases introduced by these methods have a limited effect on the model's attention mechanism. Furthermore, the attention in these models appears to be inattentive (characterized by dense patterns), potentially leading to a failure in focusing on key information and aggregating redundant information. These findings highlight a critical need for more sophisticated approaches capable of effectively incorporating structural insights into the self-attention framework of GTs. Therefore, the goal is to develop a method that not only enhances the model's attention mechanism by focusing on structurally significant features but also reduces redundancy in the information aggregation process.

Different from the aforementioned methods, we propose Gradformer, a novel method that innovatively integrates GTs with inductive bias. Specifically, Gradformer integrates an exponential decay mask into the GT self-attention architecture. This mask, multiplied with the attention score, explicitly controls each node's attention weights relative to other nodes, ensuring that the attention weights decay with an increasing node distance. In addition, the exponential decay ensures a gradual reduction in attention weights at the boundary of the full attention zone, avoiding an abrupt truncation. As a result, the introduced decay mask, rooted in the graph's structural bias, effectively guides the learning process within the selfattention framework. Furthermore, Gradformer incorporates a learnable constraint within the decay mask, applied to the attention heads, dynamically adjusting the node distance for full attention in a graph. This constraint amplifies the model's ability to discern local structural nuances and diversifies the attention heads, facilitating more effective assimilation of diverse structural information.

The design elements of Gradformer offer multiple benefits: 1) The decay mask effectively integrates a form of prior knowledge, deeply connected to the graph's structural attributes, into the GT models. 2) The structure-oriented prior knowledge precisely governs each node's interaction radius and hence defines the extent of its attention relative to other nodes, thus preventing the aggregation of redundant information. This is achieved by the design of exponential decay, which enables mask values for distant nodes approaching near-zero levels, thus effectively reducing their influence in the aggregation process. 3) The exponential decay mask endows the enhanced attention mechanism in Gradformer with the ability to function as a unified form of GNNs and GTs, synergizing their strengths: the local processing power of GNNs and global aggregation capabilities of GTs (further expounded in Section 3.3). In summary, Gradformer empowers the self-attention mechanism to effectively concentrate on structural information within the graph and limit unnecessary aggregation from distant nodes.

To evaluate the effectiveness of our Gradformer, we conduct extensive experiments on various commonly-used datasets, including the large-scale Open Graph Benchmark (OGB) [Hu et al., 2020]. The experimental results consistently demonstrate that Gradformer outperforms state-of-theart GT models on most datasets with remarkable stability and accuracy improvements, even as network depth increases. This finding stands in stark contrast to the performance of other GT models, where accuracy tends to decline with increased depth [Zhao et al., 2023]. Additionally, incorporating the decay mask design into Gradformer provides notable advantages for the graph classification task in low-resource settings, underscoring the versatility and practicality of our model. Our main contributions are summarized as follows:

1. We propose a general GT with an exponential decay mask attention mechanism, termed Gradformer, which enhances the self-attention mechanism by fully leveraging the graph's structural information. Therefore, Gradformer maintains the ability to capture long-range information while prioritizing the local information of the graph, guided by an intense inductive bias.
2. We conduct extensive experiments to compare Gradformer with 14 GNN and GT baseline models for the graph classification task on various real-world graph datasets, including OGB. The experimental results consistently validate the effectiveness of the proposed Gradformer.

2 RelatedWork

Graph Transformers. In recent years, many transformer variants have been applied to graph modeling. Unlike GNNs, transformers display competitive or even superior performance across various tasks. Dwivedi et al. [2021] were the first to extend the transformer architecture to graphs and propose PEs [Ding et al., 2020]. Subsequently, numerous variants of GT have been proposed, making significant progress in graph-level tasks [Rong et al., 2020; Chen et al., 2021; Wu et al., 2021; Chen et al., 2022; Yin and Zhong, 2023]. However, these methods are primarily designed for graphlevel tasks due to the time and memory constraints of the selfattention layer, which requires O(n2 ) complexity. Therefore, multiple studies [Zhao et al., 2021; Choromanski et al., 2022; Guo et al., 2022; Park et al., 2022; Wu et al., 2022; Wu et al., 2023a; Liu et al., 2023; Liu et al., 2024] have been conducted to enhance GTs' scalability and efficiency, facilitating their application in node-level tasks. Additionally, GTs have been applied in various fields, including natural language processing [Cai et al., 2022], computer vision [Zhu et al., 2021; Tang et al., 2023], recommendation systems [Xu et al., 2019a; Li et al., 2023], multimodal contexts [Chen and Li, 2022; He and Wang, 2023], and reinforcement learning [Hu et al., 2023]. For a more detailed introduction, please refer to the recent GT reviews [Rampasek et al., 2022; Min et al., 2022].

Prior Knowledge in Graph Transformer. Numerous efforts have been undertaken to integrate prior knowledge into GTs to augment their performance. This integration can be classified into three categories. 1) Proposed position encodings: SAN [Kreuzer et al., 2021] suggests utilizing learned Laplace eigenfunctions as a substitute for the original PE. GraphGPS [Rampasek et al., 2022] summarizes possible positional and structural encodings for GTs. 2) Improved selfattention: Graphormer [Ying et al., 2021] suggests adding structural encodings to the attention metric as a bias to enhance GTs' structured-data modeling. Moreover, RT [Diao and Loynd, 2023] and EGT [Hussain et al., 2022] consider incorporating the edge vectors in their self-attention calculations. 3) Modified architecture: GraphTrans [Wu et al., 2021] introduces a transformer sub-network above a GNN layer. In GraphGPS, the transformer and GNN layers are placed in parallel. In this manuscript, our proposed method lies on the line of improving the self-attention learning.

3 Methodology

3.1 Preliminaries

Notations. A graph G can be represented by an adjacency matrix A ∈ {0, 1}n×n and a node feature matrix X ∈ Rn×d , where n is the number of nodes, d is the node feature dimension, and A[i, j] = 1 if an edge between nodes vi and vj exists, otherwise, A[i, j] = 0.

Graph Transformer. GTs [Vaswani et al., 2017; Ying et al., 2021] consist of two essential parts: a multi-head selfattention (MHA) module and a feed-forward network (FFN). Given the node embedding matrix H(l) ∈ Rn×d(l) in a graph, a single attention head is computed as follows:

where H(l+1) ∈ Rn×d(l+1) is the output matrix, d(l+1) is the output hidden dimension, and Q(l) ∈ Rn×d(l) , K(l) ∈ Rn×d(l) , and V(l) ∈ Rn×d(l) are the query, key, and value vectors, respectively, which are the projection results of H(l) ∈ Rn×d(l) :

Q(l) =H(l)WQ;K(l) =H(l)WK ;V(l) =H(l)WV , (2) where WQ ∈ Rd(l)×d(l+1) ,WK ∈ Rd(l)×d(l+1) , and WV ∈ Rd(l)×d(l+1) are projection matrices. Note that the above single-head self-attention module can be generalized into a MHA via the concatenation operation.

3.2 ProposedMethod: Gradformer

This subsection introduces the Gradformer method. First, we describe the fundamental Gradformer architecture with a primary focus on its central module, the exponential decay mask attention mechanism. The mechanism comprises two components: exponential decay and learnable constraints. Subsequently, we explore the Gradformer method in detail, with a focus on its general form and computational complexity.

Architecture

As depicted in Figure 2, Gradformer extends the original GT architecture by incorporating a decay mask. The decay mask, derived from the graph's structural information, introduces a strong inductive bias into the self-attention learning.

Masking Attention with Exponential Decay. According to Eq. (1), we observe that the attention matrix is comparable to a row-normalized adjacency matrix of a directly weighted complete graph. This dictates the aggregation of node features in a graph, similar to GCN. Unlike the input graph, this dynamically forms graph representations through the attention mechanism. However, the basic transformer does not directly consider the input structure (i.e., existing edges) when forming these weighted graphs (i.e., the attention matrices). Consequently, the attention mechanism may overlook the importance of certain neighbouring nodes due to feature similarity, and hence these dynamic graphs collapse immediately after aggregation. To address this issue, we introduce a decay maskM, related to the graph's structural information, to participate in the aggregation process as follows:

where λ is a scalar and ψ (vi, vj) refers to a function that measures the spatial distance between vi and vj in a graph. The function ψ(·) can be defined by the connectivity between the nodes in the graph. In this paper, we choose ψ (vi, vj) to present the distance of the shortest path (SPH) between vi and vj if the two nodes are connected. If not, we set the output of ψ(·) to a special value (i.e., -1). Based on the study [Sun et al., 2023], our approach configures the attention mask M to decrease monotonically and exponentially, as depicted in the left part of Figure 4. Compared to methods with linear decay, exponential decay eliminates the need for an additional endpoint and effectively filters attention at extended distances, with values for distant nodes diminishing to nearly zero. After obtaining the attention mask, it is multiplied with the attention score S ∈ Rn×n , as illustrated in Figure 3, and formally defined as:

where ˜S(l) denotes the refined attention at layer l, and Attn denotes the self-attention operation with decay mask.

Mask Decay with Learnable Constraints. To enhance the attention mechanism's ability to model local structural information, given the strong locality inherent in graphs, a refined version of Eq. (3) is provided. This version incorporates a distance constraint into the decay operation, effectively tailoring the attention mechanism to be more sensitive to the local information. The refined decay mask is as follows:

where sp is a scalar that defines the starting point of the decay. Specifically, if the spatial distance in ψ (vi, vj) is less than sp, the attention score between vi and vj remains unaffected by the decay. Hence, this design accurately controls each node's interaction radius, delineating the extent of its attention relative to other nodes and preventing the aggregation of redundant information.

Furthermore, to enable various attention heads to capture diverse structural information, we make the constraints sp learnable. Consequently, when the decay mechanism is extended to multi-head attention, a decay mask bearing different constraints is applied to each head, as shown below:

where h represents the total number of attention heads. As illustrated in the right part of Figure 4, the constraints of different attention heads are learned to capture various structural information. Specifically, we observe that the learned constraint for Head 3 is relatively large. This characteristic enables Head 3 to focus on global information of the graph. Conversely, Head 1 exhibits a different learning pattern, which allows model to capture proximate information.

3.3 Discussion

Gradformer is a General Form of GNNs and GTs. In essence, Gradformer represents a more generalized form of the GNN [Veliˇckovi´c et al., 2018] and GT [Dwivedi and Bresson, 2021] models. This generalization becomes evident when considering the behavior of the decay mask M under varying settings of the parameter λ. Specifically, when λ = 1, the decay mask in Eq. (3) essentially transforms into an all-ones matrix. In this configuration, Gradformer effectively mirrors the GT model, as the decay mask exerts no modifying influence on the attention scores. Additionally, when λ = 0 , the decay mask M0 retains the mask value of the 1-hop neighbors as one and sets all others to zero. This configuration parallels the adjacency matrix A, thereby aligning Gradformer closely with the GNN models. The formal Notations: *: DeepGraph utilizes a distinct evaluation metric on the pattern dataset, diverging from other models. Employing this metric, our method achieves an accuracy of 90.88±0.06. **: The result is derived from the fine-tuning of a large pretrained model.

Table 1: Experimental results on eight common datasets (the mean accuracy (Acc.), AUROC, and MAE, and standard deviation over 10 different runs). Bold: the best performance per dataset. Underline: the second best performance per dataset.

derivation of the above process is as follows:

Taking the GCN layer as an example, the message passing mechanism is expressed as follows:

where σ denotes the Sigmoid activation function. From these equations, we can deduce that Gradformer essentially constitutes a generalized GNN form, possessing at least the same expressive capability as traditional GNNs.

In summary, Gradformer surpasses traditional GNNs by broadening its receptive field to encompass more relevant nodes. Furthermore, compared to GTs, Gradformer demonstrates superior capacity in fusing node representations with graph structure, capturing more topological information.

Computational Complexity. This subsection analyzes Gradformer's complexity. Compared to vanilla GTs, Gradformer incurs only the additional computational cost of pointwise multiplication between the decay and attention matrices. This process has significantly less computational complexity than self-attention. Additionally, Gradformer's additional memory requirement is the mask matrix, which shares the same size as the self-attention one. To illustrate Gradformer's complexity, detailed experimental results for time and memory are provided in the following section (See Figure 7).

4 Experiments

4.1 Experimental Settings

Datasets. We utilize nine commonly-used real-world datasets from various sources to ensure diversity, including five graph datasets from the TU database [Morris et al., 2020] (i.e., NCI1, PROTEINS, MUTAG, IMDB-B, and COLLAB), three datasets from Benchmarking GNN [Dwivedi et al., 2023] (i.e., PATTERN, CLUSTER, and ZINC), and one dataset from OGB [Hu et al., 2020] (i.e., OGBG-MOLHIV), involving diverse domains (e.g., synthetic, social, biology, and chemistry), sizes (e.g., ZINC and OGBG-MOLHIV are large datasets), and tasks (e.g., node classification, graph classification and regression). For all datasets, we strictly follow the evaluation metrics and dataset split recommended by the given benchmarks [Ying et al., 2021].

Baseline. To demonstrate the effectiveness of our proposed method, we compare Gradformer with the following 14 baselines: (I) 4 standard GCN-based models: GCN [2017], GAT [2018], GIN [2019b], and GatedGCN [2016]; (II) 10 GT-based graph models: GraphTransformer [2021], SAN [2021], Graphormer [2021], GraphTrans [2021], SAT [2022], EGT [2022], GraphGPS [2022], LGIGT [2023], KDLGT [2023b], and DeepGraph [2023]. For each baseline, we utilize the recommended settings as per the official implementation guidelines.

Implementation Details. We assess our proposed model's effectiveness by measuring its performance in graph classification and regression tasks. To ensure reliability, we conduct 10 trials for each model using different random seeds. Accordingly, we report the average test accuracy/AUROC/MAE based on the epoch when the best validation accuracy/AUROC/MAE is achieved. Furthermore, all the experiments are conducted on a server equipped with 8 NVIDIA A100s.

4.2 Overall Performance

We evaluate the proposed model's effectiveness by comparing it with GT models that are both shallow and deep. For each model and dataset, we conduct 10 trials with random seeds, and then measure the mean accuracy and standard deviation, which are reported in Tables 1 and 2.

Performance of Gradformer. From the results in Table 1, we observe that: 1) Gradformer demonstrates exceptional performance, achieving state-of-the-art results on five datasets while remaining competitive on three others. This highlights the efficacy of the proposed method. 2) Notably, on small-scale datasets such as NCI1 and PROTEINS, Gradformer outperforms all 14 methods, exhibiting improvements of 2.13% and 2.28%, respectively. This demonstrates Gradformer's effective integration of inductive biases into the GT model, a crucial advantage especially in scenarios with limited data availability. To further validate this observation, additional experiments under low-resource conditions for the large-scale datasets have been conducted (see next section). 3) Furthermore, Gradformer delivers competitive results on large-scale datasets (e.g., ZINC), demonstrating its potential applicability across various dataset scales.

Performance of Gradformer with Deep Layers. Table 2 displays the results of Gradformer with deep layers. Based on these results, the following observations can be made. 1) Increasing the number of layers notably decreases the accuracy of deep GTs, such as GraphGPS (e.g., a drop of 14.4% on NCI1 and 9.54% on CLUSTER). 2) In contrast, Gradformer demonstrates a significant improvement in this context. It sustains performance on the NCI1 dataset and surpasses shallow models on CLUSTER and PATTERN datasets. This improvement can be attributed to Gradformer's inherent flexibility in accommodating increased model depth. Specifically, while self-attention in shallow networks is inclined to focus locally, it adopts a more global perspective in deeper networks. Therefore, incorporating residual connections enhances information flow between layers, stabilizing the learning process. The decay mask utilized by Gradformer plays a pivotal role in ensuring that the updated node information maintains higher fidelity to the original data, producing effects similar to those observed in shallower models.

4.3 Further Discussions

Results on Low-resource Settings. We conduct experiments using datasets with a reduced number of labels on NCI1 and ZINC datasets. The results in Table 3 indicate that Gradformer consistently outperforms GraphGPS across all settings, particularly when the number of labels is very low, with improvements of 2.4% and 2.3% for 5% and 10% labeled data, respectively. This finding indicates that our method effectively incorporates the inductive biases into the GTs, enabling them to efficiently assimilate graph information with a limited number of labeled datasets. Furthermore, this discovery highlights the potential applicability of our model under resource-constrained conditions, such as the biomedical field where labeled data are scarce and costly.

Performance without MPNN or Position Encoding. In our subsequent analysis, we examine the impact of two key GT components: the Message Passing Neural Network (MPNN) module and PE. This investigation involves four graph datasets. It is important to note that except for the specific components under study, all other aspects of the models are maintained in line with the complete models. The findings, summarized in Table 4, indicate that: 1) The removal of the MPNN module leads to a slight decline in Gradformer's performance, especially in contrast to GraphGPS. For instance, the performance of GraphGPS without MPNN diminishes by 15.67 on the PATTERN dataset, whereas Gradformer only experiences a reduction of 0.07. 2) The absence of PE seems to have a minimal, and in some cases even beneficial, impact on Gradformer. The model demonstrates improved performance on certain datasets (i.e., PATTERN) when PE is omitted. This finding suggests that Gradformer's architecture, particularly its decay mask design, can compensate for the absence of positional information. In summary, these results imply that our method effectively enables self-attention mechanisms to capture more structural information.

The Impact of Graph Structure Index. As detailed in Section 3.2, our model's decay mask is formulated based on specific graph structural indices. In addition to the previously mentioned SPH, our study also examines another structural index: the discrete Ricci curvature (Curve) [Lai et al., 2023]. This index is the graph distance based on Riemannian manifold. Furthermore, we evaluate the effectiveness of a featurebased index, specifically the feature cosine similarity (FS). The results, illustrated in Figure 5, show that SPH consistently outperforms the other two indexes. Notably, the accuracy achieved with FS is significantly lower than that obtained with the structural indices (i.e., SPH and Curve). These findings further emphasize the importance of incorporating structural information in enhancing model performance.

The Impact of Decay Ratio and Decay Function. In an analysis conducted on three datasets, NCI1, PROTEINS, and COLLAB, the impact of varying the parameter λ is examined. The results are depicted in Figure 6. These findings reveal that as λ increases, the model's accuracy initially rises, reaching optimal performance typically around the values of 0.5 or 0.6. Subsequently, the model's accuracy declines. Furthermore, as elaborated in Section 3.3, specific λ values confer distinct characteristics to Gradformer. For instance, at λ = 1, Gradformer aligns with the conventional GT model, whereas at λ = 0, it nearly regresses to a GNN model. Notably, across varying decay ratios, Gradformer consistently exhibits superior performance over the standard GNN model. In addition, observations indicate that the exponential decay function consistently outperforms the linear one.

Efficiency Analysis. To validate Gradformer's efficiency, its training cost is compared with prominent methods such as SAN [Kreuzer et al., 2021], Graphormer [Ying et al., 2021], and GraphGPS [Rampasek et al., 2022], with a focus on metrics such as running time and GPU memory usage. The comparative results are presented in Figure 7. These findings reveal that Gradformer achieves an optimal balance between efficiency and effectiveness. Notably, Gradformer outperforms SAN and GraphGPS in computational efficiency and accuracy. Although Gradformer exhibits a marginally longer runtime compared to Graphormer, the former significantly surpasses the latter in accuracy, highlighting its superiority in balancing resource usage with high-performance outcomes.

5 Conclusion

In this study, we introduce Gradformer, a novel integration of GT with intrinsic inductive biases, achieved by applying an exponential decay mask with learnable parameters to the attention matrix. Through extensive experimentation across 9 graph datasets, Gradformer has shown its superiority by outperforming 14 contemporary GTs and GNNs. Notably, Gradformer's strength lies in its ability to maintain or even surpass shallow models in accuracy while deepening the network architecture. This is a feat not commonly observed in other GTs where accuracy tends to decline significantly in the same context. Despite its competitive performance, Gradformer still has areas for further improvement, including 1) exploring the feasibility of achieving state-of-the-art structure without the use of MPNN, and 2) investigating the potential for the decay mask operation to significantly improve GT efficiency.

Acknowledgments

The work of Wenbin Hu was supported by the National Key Research and Development Program of China (2023YFC2705700). This work was supported in part by the Natural Science Foundation of China (No. 82174230), Artificial Intelligence Innovation Project of Wuhan Science and Technology Bureau (No. 2022010702040070), Natural Science Foundation of Shenzhen City (No. JCYJ20230807090211021).

Contribution Statement

Chuang Liu and Zelin Yao contributed equally to this work.