Gapformer: Graph Transformer with Graph Pooling for Node Classification

1 Introduction

Graph Neural Networks (GNNs), which are based on local message-passing [Kipf and Welling, 2017], have achieved notable success in a variety of applications, including biology [Xu et al., 2019b] and recommendation [Zhang et al., 2019]. In contrast to GNNs, Graph Transformers (GTs) allow each node in a graph to directly attend to all other nodes, which enables the aggregation of information from arbitrary nodes. Along with normalization and residual connection, GTs overcome the deficiencies of GNNs in dealing with oversmoothing [Rong et al., 2020b], over-squashing [Alon and Yahav, 2021], heterophily [Zhu et al., 2020], and long-range dependencies [Zhang et al., 2022].

∗This work was done when Chuang Liu worked as an intern at JD Explore Academy.

†Corresponding Author

However, existing GTs [Dwivedi and Bresson, 2021; Kreuzer et al., 2021; Ying et al., 2021] are exploited primarily for graph-level tasks (e.g., graph classification) with a small number of nodes in a graph. Developing GTs for node classification, where the number of nodes in a graph is relatively large (up to around one million), remains a challenging proposition for the following two reasons. First, the quadratic computational complexity O(n2 ) of self-attention in vanilla GTs, in regards to the number of nodes, inhibits their application to node classification in real-world scenarios. Second, vanilla GTs calculate the full connected attention and aggregate messages from arbitrary nodes, including numerous irrelevant nodes; this results in ambiguous attention weights and the aggregation of noise information from incorrectly correlated nodes.

Only a few existing works have attempted to consider GTs for node classification. GT-sparse [Dwivedi and Bresson, 2021] and SAN [Kreuzer et al., 2021] confine the receptive field of each node to its 1-hop neighboring nodes; as a result, expressiveness is sacrificed when important interactions are multiple hops away, especially in the large-scale graph correspondingly requiring a large receptive field. LiteGT [Chen et al., 2021a], DGT [Park et al., 2022], and DET [Guo et al., 2022] propose selecting important nodes for attention using certain specific, fixed node sampling strategies, which may result in the selection of uninformative nodes. GraphGPS [Rampasek et al., 2022] and TokenGT [Kim et al., 2022] directly adopt efficient approximations [Choromanski et al., 2021; Beltagy et al., 2020] in Transformers [Vaswani et al., 2017] to improve the efficiency of vanilla GTs; nevertheless, they neglect unique characteristics of graph data and tend to yield dense attention, causing an enormous amount of noise messages to be aggregated from irrelevant nodes.

In light of the above analysis, we propose Gapformer, which combines Graph Transformer with Graph Pooling, to capture long-range dependencies and improve the efficiency of vanilla GTs. In vanilla GTs, self-attention converts nodes into queries and keys/values, after which each query attends to all the keys. Specifically, self-attention involves computing the inner product between the query and key vectors to generate attention scores. These scores are then used to perform a weighted aggregation of value vectors. To reduce the complexity of the dense inner product, Gapformer first utilizes graph pooling to group key and value nodes into a smaller number of pooling nodes. For graph pooling, we propose two types of strategies to compress the original graph efficiently and effectively: 1) global graph pooling, which condenses the entire graph into several significant pooling nodes; 2) local graph pooling, which compresses the k-hop neighboring nodes of each query node into the corresponding pooling nodes. Subsequently, each query node interacts with pooling keys (fewer in number) and generates representation with the pooling values. In conclusion, our Gapformer transforms the fully connected attention in vanilla GTs into a sparse attention schema by decreasing the number of attended tokens via graph pooling.

As shown in Figure 1, Gapformer has the following advantages. 1) Gapformer enables a larger attention field per node and thus allows to compute multi-hop correlations via graph pooling between each node and its corresponding disconnected nodes. 2) Since the number of pooling nodes is significantly smaller than that of nodes in the original graph, the computational complexity of Gapformer only increases linearly with the number of nodes in a graph, hence making Gapformer suitable for processing large-scale datasets in the node classification task.

Our main contributions are summarized as follows:

1. We propose Gapformer, a deeper combination of Transformer and Graph Neural Networks. Specifically, Gapformer utilizes Graph Pooling to group the attended nodes of each node into pooling nodes (fewer in number) and
2. computes its attention using only the pooling nodes. This design mitigates the overwhelming unrelated information and quadratic complexity issues associated with GTs while preserving long-range interactions.
3. We conduct extensive experiments to compare Gapformer with 20 GNN and GT baseline models in the node classification task on 13 real-world graph datasets, including homophilic and heterphilic datasets. Experimental results consistently validate the effectiveness and efficiency of our proposed Gapformer.

2 RelatedWork

GraphPooling. GraphNeural Networks (GNNs) [Kipf and Welling, 2017; Hamilton et al., 2017] are networks that perform on graph domain. As an essential component of GNNs, Graph Pooling condenses the input graph with node representations into a smaller graph or a holistic graph-level representation. There are two main types of designs proposed for graph pooling: flat and hierarchical. Flat pooling directly generates a graph-level representation in one step, mostly by taking the average or sum over all node embeddings as the graph representation [Duvenaud et al., 2015]. On the other hand, hierarchical pooling gradually coarsens a graph into a smaller one using either node clustering pooling [Ying et al., 2018; Bianchi et al., 2020] or node drop pooling [Gao and Ji, 2019; Lee et al., 2019; Liu et al., 2023]. Node clustering requires significant computational resources to cluster nodes into clusters, while node drop selects a subset of nodes from the original graph to construct a coarsened version that is more efficient and suitable for large-scale graphs [Lee et al., 2019]. For further details, please refer to [Liu et al., 2022b].

Graph Transformers. In recent years, many Transformer variants have been successfully applied to graph modeling, displaying competitive or even superior performance on many tasks when compared to GNNs. Dwivedi et al. [2021] were the first to extend the transformer architecture to graphs and propose position encoding [Ding et al., 2020] for nodes in a graph. Subsequently, Kreuzer et al. [2021] enabled the position encoding by making it learnable, and further divided the fully connected edges into true edges and virtual edges. There are many other existing GTs [Rong et al., 2020a; Chen et al., 2021b; Wu et al., 2021; Hussain et al., 2022; Chen et al., 2022; Nguyen et al., 2022] and the applications of GTs [Xu et al., 2019a; Zhu et al., 2021; Zhu et al., 2022; Cai et al., 2022], for a more detailed introduction, please refer to the recent reviews of GTs [Rampasek et al., 2022; Min et al., 2022]. However, the above methods are mostly designed for graph-level tasks due to the time and memory constraints imposed by the self-attention layer, which requires O(n2 ) complexity. Therefore, several works [Zhao et al., 2021; Choromanski et al., 2022; Guo et al., 2022; Park et al., 2022; Wu et al., 2022] have been proposed to make graph transformers more scalable and efficient, but they still suffer from some issues such as long-range information loss or noise aggregation.

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 dimension of the node features, and A[i, j] = 1 if there exits an edge between node vi and node vj , otherwise, A[i, j] = 0.

Graph Pooling. Let a graph pooling operator be defined as any function Pooling that maps a graph G = (V, E) to a new pooled graph G ′ = (V ′ , E ′ ) :

where |V ′ | < |V|. The main objective of graph pooling is to decrease the number of nodes in a graph while maintaining its semantic information.

Transformer. The vanilla Transformer [Vaswani et al., 2017] consists of two essential parts: a multi-head selfattention (MHA) module and a position-wise feed-forward network (FFN). To build a deeper model, a residual connection [He et al., 2016] is employed to each module, followed by a layer normalization (LN) [Ba et al., 2016]. The selfattention mechanism calculates attention scores by taking the inner product of query vectors (Q) and key vectors (K). It then uses these scores to aggregate value vectors (V) in a weighted manner, resulting in contextualized representations, that is,

where WQ ∈ Rd×d′ ,WK ∈ Rd×d′ , and WV ∈ Rd×d′ are projection matrices, H = [h⊤1 , . . . , h⊤n ] ∈ Rn×d denotes the input matrix of node embeddings, H′ ∈ Rn×d′ is the output matrix, and d′ is the output hidden dimension. Note that, Equation (3) denotes the single-head self-attention module, which can straightforwardly generalize to MHA.

3.2 ProposedMethod: Gapformer

In this section, we present the model architecture of Gapformer. First, we introduce the base architecture of Gapformer and its core module; that is, the attention enhanced with graph pooling (AGP). We then provide a comprehensive description of AGP from both global and local perspectives.

Architecture

As shown in Figure 2 (a), self-attention in vanilla GT calculates the dot product between each pair of nodes after projection (QK⊤). Therefore, the computation of full self-attention comes with potential noises from long-distance neighbors and O(n2 ) complexity, limiting its capacity to analyze large-scale graphs. Graph pooling manages to reduce the number of graph nodes while maintaining semantic information. This encourages our employment of graph pooling to overcome the deficiencies of GTs. To our knowledge, no efforts have been made to integrate GTs and graph pooling.

Attention Enhanced with Graph Pooling. In light of the above analysis, Gapformer uses a sparse attention schema based on graph pooling to replace the full self-attention mechanism. Specifically, as shown in Figure 2 (b), Gapformer first produces query, key, and value matrices (Linear Module); that is,

We define the query vector of node vi as ˜ qi, while its corresponding key and value matrices in the vanilla self-attention are ˜K ∈ Rn×d′ and ˜V ∈ Rn×d′ , respectively. Subsequently, we apply graph pooling to compress ˜K and ˜V, which is defined as follows:

where KS(i) ∈ Rn′×d′ and VS(i) ∈ Rn′×d′ are the compressed key and value matrices of node vi respectively, and the size of pooled node sets S (i) for node vi is n′ , significantly smaller than n. Pooling(·) refers to graph pooling, which is discussed in detail in the next section. The attention enhanced with graph pooling is then calculated as:

where α is a constant scalar (α = 1 √d′ ).

Following [Guo et al., 2022; Zhao et al., 2021], we also maintain the message-passing with the neighboring nodes. The process is defined as follows:

where ˜KN(i) and ˜VN(i) are the key and value matrices of neighboring nodes, respectively. Therefore, the node vi's final output of single attention module (AGP) in Gapformer is calculated as follows:h′i = hi + β ∗ zi, (9)

where β is a balanced hyper-parameter which controls the combination for the attention enhanced with graph pooling and the neighboring attention.

Other Modules. In addition to AGP, Gapformer also contains layer normalization (LN(·)) applied after the multihead self-attention (MHA(·)) and the feed-forward blocks (FFN(·)), as illustrated in Figure 2. We formalize the Gapformer layer as below:

As with most GT methods, our Gapformer also adopts the positional encodings (PEs), i.e., Laplacian eigenvectors encodings (LapPE) [Dwivedi and Bresson, 2021; Kreuzer et al., 2021] and random-walk positional encodings (RWPE) [Dwivedi et al., 2022].

Two Types of Attention Enhanced with Graph Pooling

In this section, we discuss how to implement the attention enhanced with graph pooling from the global and local views.

Attention Enhanced with Global Graph Pooling (AGPG). AGP-G reduces the number of attended nodes by compressing the original nodes into new pooling nodes in smaller sizes. Intuitively, all information in a graph is compressed into the new pooling nodes. As shown in Figure 3, given node features H ∈ Rn×d with their adjacency information A ∈ {0, 1}n×n , we first construct new keys and values using graph pooling operations, as follows:

Then, as shown in the left part of Figure 3, each node (Q) in the original graph attends to the pooling nodes (KGlobal S ) in the new sets S to generate attention scores. For the pooling operation (Pooling(·)), we empirically explore several different pooling methods [Liu et al., 2022b] to perform compressions, including flat pooling methods (e.g., the mean and max pooling), and trainable pooling mechanisms (e.g., SetPool [Vinyals et al., 2016] and SAGPool [Lee et al., 2019]).

Attention Enhanced with Local Graph Pooling (AGP-L). Different from AGP-G, AGP-L works by compressing the neighbor information of each node. Specifically, as shown in the right part of Figure 3, for each node (i.e., each query Qi), we execute graph pooling on its neighboring nodes extracted from its k-hop subgraphs. Formally, the new keys and values for each node (i.e., node vi) are generated by

where N (i, k) refers to the neighbor set within k hops of node vi. The nodes in a graph then perform attention with those in the corresponding pooled sets. Formally, the output of the AGP-L for node vi is calculated by Eq. (7). Note that the operation of extracting k-hop subgraphs can be performed in the preprocessing stage without consuming additional resources in the training stage. Finally, as with the global pooling, we empirically explore several different pooling methods for use in performing compression (Pooling(·)), including the flat pooling methods (e.g., the mean and max pooling) and trainable pooling mechanisms (e.g., SetPool [Vinyals et al., 2016] and SAGPool [Lee et al., 2019]).

3.3 Merits of Gapformer

Reducing Computational Complexity. We first analyze the complexity of Gapformer. The computational complexity of the attention enhanced with graph pooling (Eq. (7) ) is O(n′n). Since n′ is a constant and usually much smaller than n, the computational complexity can be simplified as O(n). Moreover, the computational complexity of the neighboring attention (Eq. (8) ) is O(|E|). Therefore, the overall complexity of Gapformer is O(n + |E|). To illustrate the superiority of Gapformer, we conduct experiments on synthetic datasets. The results in Figure 4 demonstrate that the time and memory usages of Gapformer do indeed scale linearly with the number of nodes, unlike the full self-attention mechanism that scales exponentially, which enables the application of Gapformer to extremely large-scale datasets.

Reducing the Ratio of Noisy Connections. In most existing graph transformer models, each node aggregates information from all nodes in a graph, which provides the global receptive field. This approach may pose a challenge for node classification tasks since the aggregated information could potentially contain noise and irrelevant data that is not useful for the target node. Consequently, this can hinder the model's ability to perform effectively. To address this issue, our proposed Gapformer modifies the standard full selfattention to a sparse schema, which helps greatly reduce the ratio of noisy connections, thereby enhancing the capacity of graph transformer-based models in node classification.

HandlingLong-rangeDependency. As shown in Figure 1, the receptive field of Gapformer is flexible and ranges from linear growth to exponential growth. Thus, it requires fewer layers than traditional GNN models to capture long-distance connections. This is a remarkable benefit for the case when significant correlations are multiple hops away.

4 Experiments

4.1 Experimental Settings

Datasets. We employ a total of 13 real-world datasets, including nine homophilic graph datasets (Cora, Citeseer, Pubmed, DBLP, CoraFull, CS, Physics, Photo, and ogbnarxiv) and four heterophilic graph datasets (Cornell, Texas, Wisconsin, and Actor), involving diverse domains (citation, co-authorship, co-purchase, and web pages) and sizes

(ogbn-arxiv is a large-scale dataset). The dataset statistics are summarized in Table 1. Please note that in reference to [Zhu et al., 2020], Homo. refers to the ratio of edges linking nodes with identical labels. We use different training, validation, and test splits for various datasets. Specifically, for Cora, Citeseer, and Pubmed datasets we follow the (48%/32%/20%) split as proposed in [Pei et al., 2020]. The same splits used by [Zhu et al., 2020] and [Liu et al., 2022a] are adopted for the four heterophilic graph datasets. For all other datasets, we randomly split them into 60%/20%/20% training/validation/test sets following [Zhang et al., 2022]. All the adopted graph datasets, except ogbn-arxiv, can be downloaded from PyTorch Geometric (PyG) [Fey and Lenssen, 2019] 1 , and obgn-arxiv can be downloaded from Open Graph Benchmark (OGB) 2 .

Baseline. To demonstrate the effectiveness of our proposed method, we compare Gapformer with the following 20 baselines: (I) 7 standard GCN-based models: GCN [2017], GatedGCN [2016], APPNP [2019], GCNII [2020], GAT [2018], GATv2 [2022], and SuperGAT [2021]; (II) 5 heterophilic-graph-oriented models: MLP [2015], MixHop [2019], FAGCN [2021], H2GCN [2020], and GPRGNN [2021]; (III) 8 transformerbased models for graphs: GT-sparse [Dwivedi and Bresson, 2021], SAN [2021], Graphormer [2021], UniMP [2021], LiteGT [2021a], DET [2022], NAGformer [2023], and ANSGT [2022]. The last five transformer-based models are efficient graph transformer models.

Implementation Details. We assess the effectiveness of our proposed model by measuring its accuracy in node classification. To ensure reliability, we conduct 10 trials for each model using random seeds. We utilize Adam optimizer for GCN-based and heterophily-based methods, while Adamw is adopted for all graph transformer-based models. Each method and dataset are run for 200 epochs, with the test accuracy reported based on the epoch that achieves the highest validation accuracy. For ease of tuning work, we set some hyperparameters: dropout at 0.5, weight decay at 5e-4 , position encoding dimension at 20, and hidden dimension within {64, 128, 256}. Our implementation of Gapformer is developed using Python (3.7.0), Pytorch (1.11.0), and Pytorch Geometric (2.2.0). All experiments are conducted on a Linux server with two NVIDIA A100s.

4.2 Overall Performance

We evaluate the effectiveness of the proposed model in terms of accuracy. For each model and dataset, we conduct 10 tri- Notations: 1) Gapformer (w/o GP) refers to Gapformer without graph pooling, which is also the GT-sparse baseline model [Dwivedi and Bresson, 2021]. 2) The results of Graphormer are taken from NAGphormer [2023] and Specformer [2023]. 3) The full-version of GT [2021] and SAN [2021] is OOM even on the small-scale datasets. 4) Another recent graph Transformer, SAT [Chen et al., 2022], is not considered, as it reports OOM even on the small-scale datasets.

1 https://github.com/pyg-team/pytorch geometric

2 https://ogb.stanford.edu/docs/nodeprop/#ogbn-arxiv

Table 2: Experimental results for the node classification task on eight common datasets (mean accuracy (%) and standard deviation over 10 different runs). Red: the best performance per dataset. Blue: the second best performance per dataset. OOM denotes out-of-memory.

als with random seeds, and then take the mean accuracy and standard deviation, which are reported in Tables 2 and 3.

Performance on Homophilic Graphs. From the results in Table 2, we can observe: 1) Gapformer achieves the stateof-the-art performance on five datasets and competitive performance on three datasets, which demonstrates the effectiveness of our proposed method. Gapformer has a significant advantage over its variant that does not include the graph pooling module, referred to as Gapformer (w/o GP). 2) Compared with GCN-based methods, Gapformer performs better on graphs with more nodes (e.g., Photo, Physics, and CS). This is likely because local message-passing based GCN methods neglect long-range dependencies, but our Gapformer enables to learn more informative node representations from the multi-hop neighborhoods, which is a remarkable benefit for the bigger graphs, where the required receptive field is large [Guo et al., 2022; Park et al., 2022; Wu et al., 2022]. 3) The performance of Gapformer surpasses that of graph transformer-based methods on the small-scale datasets (e.g., Cora and Citeseer). The reason may be that vanilla GTs with full connected attention (e.g., Graphormer) and sampling-based GTs (e.g., LiteGT [Chen et al., 2021a] and DET [Guo et al., 2022] ) both introduce more noises from massive unrelated nodes. However, GCN-based models perform better than GT-based methods on small-scale datasets. This is likely because, on small-scale datasets, local information is more important. Moreover, GTs, including our Gapformer, have more parameters than GCN-based methods, meaning that they may suffer from over-fitting on small datasets. 4) Our Gapformer can be applied to large-scale graphs, such as ogbn-arxiv, while some other transformerbased methods cannot be applied to such graphs due to their poor scalability. We have noticed that Graphormer [Ying et al., 2021] and LiteGT [Chen et al., 2021a] encounter outof-memory errors, even when processing small graphs. This highlights the need for a graph Transformer that can scale effectively to handle large-scale graphs.

Performance on Heterophilic Graphs. Table 3 summarizes the results of models on heterophilic graphs. From these results, we can make the following observations: 1) GCN- Notations: 1) Gapformer (w/o GP) refers to Gapformer without graph pooling, which is also the GT-sparse baseline model [Dwivedi and Bresson, 2021]. 2) The official codes of LiteGT [Chen et al., 2021a] and ANS-GT [2022] fail to handle the above four heterophilic datasets.

Table 3: Experimental results for the node classification task on four heterophilic datasets (mean accuracy (%) and standard deviation over 10 different runs). Red: the best performance per dataset. Blue: the second best performance per dataset.

based models exhibit relatively inferior performance on heterophilic graphs. This is because most GCNs utilize directly connected nodes for aggregation even in heterophilic graphs. 2) Surprisingly, transformer-based models show poor performance, which implies that GTs fail to filter out irrelevant messages. 3) Instead, our proposed Gapformer achieves superior performance on heterophilic graph datasets. In particular, Gapformer significantly outperforms transformer-based baselines. This phenomenon is probably because Gapformer summarizes the graph structure or neighbor information from a global view instead of a similarity view.

4.3 Further Discussions

Efficiency of Gapformer. To validate the efficiency of Gapformer, we compare its training cost in terms of running time (s) and GPU memory (MB) with such representative methods as GT [Dwivedi and Bresson, 2021], SAN [Kreuzer et al., 2021], and ANS-GT [Zhang et al., 2022]. The results are summarized in Table 4. From these results, we can observe that Gapformer with AGP-G shows high efficiency compared with all existing GTs, especially when dealing with large-scale graphs. Moreover, Gapformer with AGP-L incurs lower memory costs compared to vanilla GTs, although its time cost is high, comparable to NAGphormer [Chen et al., 2023] and ANS-GT [Zhang et al., 2022].

Ablation on Transformer Components and Pooling Type. We next study the effects of the three components of the transformer and different pooling types. We conduct experiments with AGP-G on three graph datasets. Please note that, apart from the selected components, all other parts remain identical to the complete model. We can observe in Figure 5 that the performance drops after the LayerNorm and Residual components are removed. However, the performance increases after removing the FFN module, which indicates that this module may cause over-fitting. Moreover, from the results in Figure 6, we can determine that SAGPool [Lee et al., 2019] performs better than other simple pooling methods. This encourages our search for more effective and efficient pooling methods to improve the performance of Gapformer.

Impact of Number of Layers and Balance Parameter. We analyzed the effects of l and β using AGP-G on three graph datasets (Citeseer, Pubmed, and Texas) and presented the results in Figure 7. Specifically, we investigated how the number of layers impacts node classification performance. Our findings indicate that as l increases from low to high values, test accuracy decreases due to over-fitting. Additionally, when β is relatively small, our model's accuracy curve remains smooth indicating less sensitivity to hyper-parameters.

5 Conclusion

We propose Gapformer, which combines Graph Transformers (GTs) with Graph Pooling for efficient node classification. Our Gapformer addresses the two main issues of existing GTs: potential noises from long-distance neighbors and the quadratic computational complexity in regards to the number of nodes. Extensive experiments on 13 graph datasets demonstrate that Gapformer outperforms existing GTs and Graph Neural Networks. Despite its competitive performance, Gapformer still has room for improvement. For instance, 1) devising an effective manner to combine the proposed local pooling enhanced attention and global pooling enhanced attention, and 2) incorporating useful techniques to further enhance the performance on large-scale graph datasets.

Acknowledgments

This work was supported in part by the Natural Science Foundation of China (Nos. 61976162, 82174230, 62002090, 62172354), Artificial Intelligence Innovation Project of Wuhan Science and Technology Bureau (No. 2022010702040070), and Science and Technology Major Project of Hubei Province (Next Generation AI Technologies) (No. 2019AEA170). Dr Wu is partially supported by ARC Projects LP210301259 and DP230100899.