Rethinking Graph Transformer Architecture Design for Node Classification

Transformer can alleviate certain limitations encountered by Graph Neural Networks (GNNs) in graph-related tasks, such as over-squashing (Alon and Yahav 2021), long-range dependency (Rampášek et al. 2022; Shirzad et al. 2023; Ma et al. 2023), weak connectivity (Wu et al. 2022; Ying et al. 2021), and limited expressiveness (Kreuzer et al. 2021; Mialon et al. 2021; Rong et al. 2020). These advantages have motivated extensive research on Graph Transformer (GT), demonstrating its strong potential in tasks such as graph classification, link prediction, and node classification.

Despite the superior performance of GT over classical GNNs in certain scenarios, its adaptation for the prevalent node classification task exhibits several limitations, which can be summarized as follows: 1) Self-attention mechanism inevitably introduces noise. Previous studies  (Xie et al. 2023; Zhu et al. 2020a; Abu-El-Haija et al. 2019) suggest that shallow GNNs can only capture local neighborhood information of nodes and cannot effectively access higher-order information. Instead, GTs employ the self-attention mechanism to receive information from all other nodes, thereby capturing higher-order relationships such as long-range dependencies. However, self-attention actually disregards the topological structure of the graph and treats it as fully connected (Ying et al. 2021), which inevitably introduces irrelevant information, or noise, especially in homophilous scenarios. Thus, the adaptability of GTs relying on self-attention in node classification task is still open to question. 2) The self-attention mechanism is unsuitable for large-scale graph computation. For vanilla GT, the computational complexity of self-attention is quadratic with respect to the number of nodes. In real-world scenarios such as social networks and e-commerce recommendations, graph-structured data often contains a massive number of nodes, making the direct application of the self-attention mechanism computationally and memory-intensive. Therefore, GTs relying on self-attention are highly inefficient for node classification on large-scale graphs.

For the former limitation, several studies (Min et al. 2022; Dwivedi and Bresson 2020) utilize topological information as masks or biases to recalibrate the self-attention distribution, thereby reducing the introduction of noise to some extent. Regarding the latter limitation, several studies (Wu et al. 2022, 2024; Zhang et al. 2022b; Liu et al. 2023) attempt to improve the efficiency of GTs on large-scale graphs by reducing the computational complexity of attention computation or compressing the graph size. Despite the observed performance improvement, these studies persist in applying the GT architecture to node classification tasks without thoroughly investigating whether the current GT architecture is genuinely adaptability for this purpose.

In this work, we question this adaptability and obtain the supporting observations, as shown in Figure 1, by substituting the multi-head self-attention module (MHA) in vanilla GT with single graph attention layer (more details will be discussed later). Specifically, we find that: 1) the applicability of the self-attention module in GT architecture to node classification is not guaranteed; 2) the feed-forward network module (FFN) in GT architecture appears to be beneficial for node classification. Based on extensive observations, we further disentangle message passing into propagation $\mathbf{P}$ and transformation $\mathbf{T}$, and design a graph transformer architecture called GNNFormer that allows for flexible combination of $\mathbf{P}$ and $\mathbf{T}$, making it suitable for node classification in both homophilous and heterophilous scenarios. The main differences compared to the vanilla GT architecture are as follows: 1) substituting the multi-head self-attention module with stackable message-passing blocks (PT-block); 2) optimizing the feed-forward network with SwishGLU; 3) improving the residual connections to adaptive initial residual connections. The incorporation of three key design in our GNNFormer enables the mitigation of global noise interference during message passing, facilitates effective fusion of graph attribute and structural features, and ultimately achieves powerful node representation learning. Extensive experiments on 12 homophilous and heterophilous benchmark datasets validate the effectiveness and superiority of our GNNFormer framework in node classification tasks, compared to existing graph transformers.

A graph is denoted as $G=(V,E,\boldsymbol{X},\boldsymbol{Y})$, where $V$ and $E$ are the set of nodes and edges respectively, $\boldsymbol{X}\in\mathbb{R}^{|V|\times d}$ is the node feature matrix, and $\boldsymbol{Y}\in\mathbb{R}^{|V|\times C}$ is the node label matrix. Each node is associated with a feature vector $\boldsymbol{x}\in\mathbb{R}^{d}$ and a one-hot encoded label vector $\boldsymbol{y}\in\mathbb{R}^{C}$. Here we use $|V|$, $d$ and $C$ to denote the number of nodes, the dimension of the node features, and the number of classes, respectively. The graph topology information $(V,E)$ can also be represented by an adjacency matrix $\boldsymbol{A}\in\mathbb{R}^{|V|\times|V|}$, where $\boldsymbol{A}_{ij}=1$ indicates the existence of an edge between $v_{i}$ and $v_{j}$, and $\boldsymbol{A}_{ij}=0$ otherwise. Based on the adjacency matrix, we can define the degree distribution of $G$ as a diagonal degree matrix $\boldsymbol{D}\in\mathbb{R}^{|V|\times|V|}$ with $\boldsymbol{D}_{ii}=\sum_{j=1}^{|V|}\boldsymbol{A}_{ij}$ representing the degree value of $v_{i}$. Node classification is a fundamental task in graph machine learning. It involves assigning labels to the nodes of a graph based on their features and the graph topology structure.

Based on the aforementioned motivating observations, we propose a novel graph transformer architecture named GNNFormer, specifically tailored for node classification. The architecture is illustrated in Figure 3(b). Compared to the vanilla GT block, the modifications are mainly reflected in three aspects: 1) the original multi-head self-attention module is substituted with the stackable P-T combination message passing block (PT-block); 2) the FFN is optimized via SwishGLU; 3) the original residual connection is improved to the initial residual connection. The design details of GNNFormer will be discussed below.

As shown in Figure 3(a), vanilla GT learns node representations by stacking transformer blocks, with the input to the current block being the output of the previous block. Differently, our GNNFormer consists of only a single transformer block, which takes node features and adjacency information as input and outputs the final node representations.

First, the input features $\boldsymbol{X}$ will be transformed into an initial feature embedding through a linear transformation parameterized by $\boldsymbol{W}_{0}\in\mathbb{R}^{d\times d^{\prime}}$ and a ReLU activation:

where $d^{\prime}$ is the hidden dimension. Next, we stack message passing blocks, i.e., PT-blocks, to further learn node representations. Each PT-block consists of a message passing function $f$, an initial residual connection, and a layer normalization operation $\operatorname{LN}\left(\cdot\right)$. For the $l$-th PT-block, it takes the adjacency matrix $\boldsymbol{A}$ and the node representations $\boldsymbol{H}^{(l-1)}$ output by the previous PT-block as input, and outputs new node representations via propagation and transformation:

where $f$ is selected from the candidate message passing set {$\mathbf{PP}$, $\mathbf{PT}$, $\mathbf{TP}$, $\mathbf{TT}$}, and $\alpha_{l}$ is a learnable parameter that controls the adaptive residual connection.

After message passing through $l$ PT-blocks, GNNFormer has effectively fused the attribute information of the nodes with the topological information. Furthermore, motivated by observation 1, which suggests that appending an FFN can enhance the expressive power of vanilla GNNs, we preserve the FFN module in the GNNFormer architecture. Specifically, the node representations $\boldsymbol{H}^{(l)}$ will first be fed into a two-layer FFN that employs the SwishGLU in place of the first linear layer and the activation function:

where $\boldsymbol{W}_{1}$, $\boldsymbol{W}_{2}$, $\boldsymbol{W}_{3}\in\mathbb{R}^{d^{\prime}\times d^{\prime}}$ are the transformation weight matrices, and $\otimes$ is the element-wise multiplication. The SwishGLU (Shazeer 2020) combines the smoothness of the Swish (Ramachandran, Zoph, and Le 2017) activation function with the selectivity of Gated Linear Units (Dauphin et al. 2017). The former ensures more stable gradient updates during back-propagation, while the latter allows the GNNFormer to selectively propagate more important features, especially in heterophilous scenarios. Consequently, the enhanced FFN improves the non-linear expressive power of GNNFormer, enabling it to better capture complex node relationships and feature interactions within various graphs.

Subsequently, an initial residual connection and layer normalization will be applied to the output of the FFN:

where $\beta$ is a learnable parameter that controls the adaptive residual connection. Finally, we further integrate the topology information to obtain the final node representations:

where $\boldsymbol{W}_{4}\in\mathbb{R}^{|V|\times d^{\prime}}$ is the weight matrix that encodes the topology information, and $\gamma$ is a learnable parameter that controls the contribution of topology information.

To achieve node classification, we finally use a prediction head $f_{\text{pred}}$ parameterized by $\boldsymbol{W}_{5}\in\mathbb{R}^{d^{\prime}\times C}$ and Softmax activation to obtain the node predictions:

During model training, binary cross-entropy classification loss is used as the optimization objective:

where the trace operation $\operatorname{trace}\left(\cdot\right)$ is used to compute the sum of the diagonal elements of the matrix.

To demonstrate the inherent misalignment of Graph Transformers in node classification tasks, we start with observational experiments and progressively develop the GNNFormer framework. Extensive experiments on 12 benchmark datasets demonstrate that, compared to existing Graph Transformer methods, our framework not only maintains superior performance but also improves computational efficiency. Additionally, we find that GNNFormer, which employs a transformation-first message-passing manner exhibits greater universality in both homophilous and heterophilous scenarios. This insight also motivates us to explore the automatic search of GNNFormer architectures in future work, i.e., driving GNNFormer to automatically search for the optimal combination of transformation and propagation operations, which will further enhance the generalizability and superiority of our framework.

This work was supported in part by the Key R&D Program of Zhejiang under Grants 2022C01018 and 2024C01025, by the National Natural Science Foundation of China under Grants 62103374 and U21B2001.