Unifying and Enhancing Graph Transformers via a Hierarchical Mask Framework

We are immensely grateful for the reviewer's constructive and valuable feedback. We are dedicated to thoroughly examining each issue raised and providing detailed responses to enhance our paper.

More Details

We have provided detailed experimental settings in Appendix G, and an anonymous ZIP archive containing our code is included in the Supplementary Material. You may refer to Appendix G for specific experimental configurations, and reproduce our results by running the runs.sh script included in the provided code.

Regarding the selection of the number of clusters, we preset a range based on the total number of nodes in each dataset to ensure that the cluster sizes remain appropriate. The selected cluster numbers are then treated as a tunable hyperparameter. The specific candidate cluster numbers for each dataset are listed in Appendix G.2. As shown in the parameter analysis in Figure 4 and Appendix H.2, the model’s performance is generally robust with respect to the number of clusters in most scenarios.

For the design of computation schemes in different regions, we follow the theoretical guidance provided by Proposition 4.2. Specifically, we use sparse attention for the local (local masks), intra-cluster (cluster masks), and node-to-global (global mask) regions, where the low density of non-zero entries makes sparse computation more efficient. In contrast, we use dense attention for the global-to-node region (also in the global mask), as this region is fully connected (all ones), making dense computation more efficient in practice. Experimental results in Figure 3 and Appendix H.1 empirically confirm the efficiency benefits of this design.

Sharing parameters across experts

We conducted an experiment where the parameters of the three attention experts were shared. The results are presented in the table below:

These results show that sharing parameters among the three experts leads to a noticeable drop in performance across most datasets. We believe this is because each expert is designed to capture interaction information at a different level. When the same parameters are shared among them, each expert is forced to handle multi-level information simultaneously, which may cause conflicts in the learning process and hinder effective representation learning.

Homophily rate and label consistency

While the homophily rate captures label agreement among immediate neighbors, our definition of label consistency is more general. It measures label agreement among structurally related nodes across multiple levels—including multi-hop neighborhoods, clusters, and global contexts—rather than being limited to local connections. This broader perspective enables a better understanding of complex interaction patterns beyond simple homophily.

Thank you for your valuable and positive feedback on our work. I’m glad to share our observations regarding the role of the Feed-Forward Network (FFN) module in node classification tasks.

The impact of FFN module

Our experiments reveal that, unlike in NLP and CV domains, the commonly used FFN module does not contribute positively to performance in node classification. The following table presents comparative results across seven benchmark datasets:

As shown, adding the FFN module often leads to a noticeable performance drop—especially on datasets such as cora, citeseer, chameleon, and squirrel. We speculate that this is due to the nature of node features in these datasets, which are generally simple and easy to learn. In such cases, capturing complex node interactions becomes more crucial, a task better handled by the Multi-Head Attention (MHA) module. While the MHA is often regarded as the core of Transformer architectures, it is worth noting that the standard FFN module typically contains twice as many parameters as the MHA (e.g., with a projection from dimension $d \rightarrow 4d \rightarrow d$). Since meaningful structural interactions are primarily modeled through attention, a natural design choice is to allocate more capacity to the MHA module instead.

In addition to the observed performance improvement, reducing or simplifying the FFN module can significantly enhance computational and memory efficiency, making the model more lightweight and scalable.

We sincerely thank the reviewer for the precious time spent reading through the paper and giving constructive suggestions. Below, we address the reviewer's concerns one by one, hoping that a better understanding of every point can be delivered.

Theoretical assumptions and generality

We thank the reviewer for the insightful comment regarding the assumptions in our theoretical analysis. We acknowledge that the use of class-conditional isotropic Gaussian distributions with orthogonal class means is a simplification and may not fully capture the complexity of real-world feature distributions. However, such assumptions are widely adopted in theoretical studies of machine learning and representation learning [1,2], as they enable analytically tractable results and offer conceptual clarity. Our goal is not to precisely mimic empirical data distributions, but to distill key factors that influence node classification and to provide theoretical grounding for the design of Graph Transformers.

Importantly, our analysis is conducted under the proposed unified hierarchical mask framework, which offers a feasible way for analyzing multi-level interactions. We believe this framework lays the foundation for more refined theoretical studies in the future, potentially under more relaxed or realistic distributional assumptions.

Extension of our method to graph-level task

The label-related global mask in our model is derived from theoretical analysis tailored for node classification. To extend our method to other tasks, we replace this mask with a general global mask, similar to that used in Exphormer[3]. Specifically, we adapt our model for graph-level tasks by incorporating edge features and Laplacian positional encodings, along with the general global mask.

We evaluate this extension on two graph classification datasets: OGBG-MOLBACE and OGBG-MOLBBBP. Due to time constraints, we compare our method with three widely recognized GNN baselines—GCN, GAT, and GINE—known for their strong performance on graph-level benchmarks [4]. The results are presented below:

These results demonstrate that our method generalizes well to graph-level tasks and outperforms strong baselines, highlighting its flexibility and effectiveness beyond node classification.

More ablation studies

We thank the reviewer for the constructive suggestions. To further validate our design choices, we conducted two additional ablation studies: (1) replacing label-informed global nodes with generic, unlabeled ones, and (2) simplifying the hierarchical mask integration by directly concatenating the three masks and applying standard attention. The results are summarized below:

We observe that both variants lead to consistent performance drops, particularly on heterophilic benchmarks (e.g., chameleon and squirrel). For the global nodes, this degradation aligns with our theoretical analysis: generic global masks lack class-specific semantics, limiting their ability to guide meaningful aggregation. In contrast, label-informed global nodes can better capture class-dependent representations and facilitate more targeted global interactions (see Section 4.2).

Similarly, the mask concatenation variant suffers due to semantic ambiguity—once the masks are merged, the model can no longer distinguish whether an edge reflects local, cluster-level, or global structure. As discussed in Section 3, these interaction levels play distinct roles depending on the graph topology. This highlights the importance of disentangled hierarchical modeling and our bi-level attention MoE, which explicitly separates and selectively leverages structural signals at different scales. Simpler alternatives fail to retain this structural granularity, resulting in reduced effectiveness.

Run time performance

We are pleased to provide a detailed time complexity analysis. For standard multi-head attention (MHA), the computational cost primarily arises from: (1) feature transformations of the QKV matrices ($O(3Nd^2)$), (2) computation of the attention matrix ($O(N^2d)$), and (3) multiplication between the attention matrix and the value matrix ($O(N^2d)$). Thus, the overall time complexity is $O(3Nd^2 + 2N^2d)$. For sparse MHA, the complexity includes: (1) QKV transformations ($O(3Nd^2)$), (2) sparse attention score computation ($O(2md)$, step 4 of Algorithm 2), and (3) output aggregation ($O(2md)$, step 6 of Algorithm 2), leading to a total complexity of $O(3Nd^2 + 4md)$. A direct comparison reveals that sparse MHA is more efficient than standard MHA when $\frac{m}{N^2} < \frac{1}{2}$.

In our "Dual Attention Computation Scheme" (Section 4.2), we apply sparse attention in most regions, except for the attention from global nodes to the origin node, which uses a dense pattern (where $\frac{m}{N^2} = 1$; see Figure 1). This hybrid design enables our method to achieve lower overall time complexity than both standard and sparse MHA, i.e., $\min(O(3Nd^2 + 2N^2d), O(3Nd^2 + 4md))$. To empirically validate efficiency, we report training times (over 200 epochs) on seven datasets, comparing our method, its sparse/dense variants, and PolyNormer [5]:

Our method consistently outperforms its dense and sparse MHA variants in training efficiency. Although slightly slower than Polyformer, the use of three MHA experts leads to faster and more stable convergence, ultimately resulting in higher accuracy. Furthermore, since the model converges within 200 epochs on most datasets, the additional overhead is acceptable in practice.

Clarifying the limitations

We acknowledge that the performance gains on standard homophilic benchmarks (e.g., Pubmed, ogbn-arxiv) are relatively modest. This is consistent with our theoretical analysis: in highly homophilic graphs, local neighborhood information already provides strong predictive signals, making complex hierarchical modeling less critical. Nonetheless, our method still consistently outperforms strong baselines across all nine datasets and exhibits stable training behavior (see Figure 5 and Appendix H.3). More importantly, our framework is designed to address the challenges of heterophilic or structurally diverse graphs, where multi-level dependencies are essential. On such datasets (e.g., Chameleon, Squirrel, Minesweeper), our model shows substantial improvements, highlighting its strength in capturing complex structural signals.

Regarding novelty, we emphasize that our key contribution lies not only in the specific model, but in the unified hierarchical mask framework. This framework provides a principled and general perspective to reinterpret GTs and enables theoretical analysis of node classification. It also facilitates a extensible way to incorporate hierarchical interactions via mask design, avoiding ad hoc architecture engineering. Our method is a natural instantiation of this framework, guided by both theory and empirical findings. In addition, to the best of our knowledge, this is the first work to integrate the MoE paradigm into Graph Transformers. Unlike prior approaches in NLP and GNNs that define experts as distinct parameter sets or layers, we treat each MHA module—coupled with a specific hierarchical mask—as a relational expert. This introduces a new form of structure-aware specialization, representing a conceptual shift in how MoE mechanisms can be utilized in graph representation learning.

[1] Shen, Ruoqi, et al. "Data augmentation as feature manipulation." International conference on machine learning. PMLR, 2022.

[2] Li, Yibo, et al. "Graph fairness learning under distribution shifts." Proceedings of the ACM Web Conference 2024. 2024.

[3] Shirzad, Hamed, et al. "Exphormer: Sparse transformers for graphs." International Conference on Machine Learning. PMLR, 2023.

[4] Luo, Yuankai, et al. "Can Classic GNNs Be Strong Baselines for Graph-level Tasks? Simple Architectures Meet Excellence." arXiv preprint arXiv:2502.09263 (2025).

[5] Deng, Chenhui, et al. "Polynormer: Polynomial-expressive graph transformer in linear time." arXiv preprint arXiv:2403.01232 (2024).

We are truly grateful to the reviewer for the constructive comments and valuable feedback regarding our paper. To thoroughly address all the concerns raised, we will provide detailed responses to each question sequentially.

Compare with recent works

We sincerely thank the reviewer for highlighting these recent works. We have carefully reviewed both papers and acknowledge that they are high-quality contributions to the field.

The first paper[1] proposes a novel approach that integrates large language models (LLMs) with graph neural networks (GNNs) to achieve impressive node classification performance. While the results are indeed strong, we believe this method is not a suitable baseline for our comparison. The approach leverages pre-trained LLMs, which inherently encode extensive world knowledge, and utilizes textual attributes as input features. These text-based features are significantly more informative than the original node features (which are typically simple bag-of-words or structural encodings) used in our setting. Therefore, the problem formulation and input assumptions differ substantially from ours, making a direct comparison less meaningful.

The second paper [2] proposes a compelling approach that leverages Residual VQ-VAE to learn expressive token representations and employs node sampling strategies to generate transformer-compatible sequences. We have acknowledged this work in the Introduction, highlighting it as a representative method of the second line of Graph Transformer research. During the rebuttal phase, we conducted a comprehensive comparison with GQT on seven node classification benchmarks. Our model consistently outperforms GQT across all datasets, as summarized below:

Note: Although the authors have released their code on GitHub, we found the implementation to be incomplete and containing several bugs. We made our best effort to complete the missing parts and fix the bugs. However, their paper does not provide detailed hyperparameter settings per dataset. Given the time constraints, we conducted a reasonable grid search to tune hyperparameters, but were unable to fully reproduce the performance reported in their paper. We plan to contact the authors for further clarification and will update our comparison accordingly in a future version of the manuscript.

Time and space complexity

Thanks for your valuable suggestions. We have conducted comprehensive space complexity analysis for standard MHA and sparse MHA in Appendix B. The conclusion is that the total space complexity of standard MHA is $O(2hN^2)$, where $h$ is the number of heads and $N$ is the number of nodes, and the overall space complexity of sparse MHA is $O(6mhd_h)$, where $m$ is the number of non-zero elements in the mask matrix and $d_h$ is the hidden size of each head. According to our proposed "Dual Attention Computation Scheme" and Proposition 4.2, the space complexity of our method is the lower bound of standard MHA and sparse MHA, i.e., $\min(O(2hN^2), O(6mhd_h))$. We have plotted the memory usage of our method and its variants in Figure 3 and Appendix H. The numerical comparison is as follows:

We find that our "Dual Attention Computation Scheme" is more memory-efficient than both the Sparse MHA variant and the Dense MHA variant, demonstrating the effectiveness of our design.

Now, we analyze the time complexity. For the standard MHA, the time complexity mainly comes from the feature transformations of the QKV matrices ($O(3Nd^2)$), the computation of the attention matrix ($O(N^2d)$), and the matrix multiplication between the attention matrix and the value matrix ($O(N^2d)$). Therefore, the total time complexity of standard MHA is $O(3Nd^2 + 2N^2d)$. For sparse MHA, the time complexity includes the feature transformations ($O(3Nd^2)$), the sparse attention score computation ($O(2md)$, step 4 of Algorithm 2), and the output computation ($O(2md)$, step 6 of Algorithm 2), resulting in a total of $O(3Nd^2 + 4md)$. A straightforward calculation shows that sparse MHA is more efficient than standard MHA when $\frac{m}{N^2} < \frac{1}{2}$.

According to our "Dual Attention Computation Scheme" and Proposition 4.2, we apply sparse attention computation in most regions, except for attention from global nodes to the origin node, which involves a dense attention region (where $\frac{m}{N^2} = 1$, as shown in Figure 1). This hybrid design allows our method to achieve lower overall time complexity than both standard and sparse MHA, i.e., $\min(O(3Nd^2 + 2N^2d), O(3Nd^2 + 4md))$.

We compare the training time (200 epochs) of our method with its variants and PolyNormer[3]. The results are as follows:

We observe that our "Dual Attention Computation Scheme" achieves better training speed than both its sparse and dense MHA variants. Although it is slightly slower than PolyNormer, the extra computation stems from the use of three MHA experts, which are essential for achieving superior accuracy. Moreover, 200 epochs are sufficient for convergence on most datasets, making the time overhead acceptable.

Code

We have included an anonymous ZIP archive of our code in the Supplementary Material. The results can be reproduced by running the provided runs.sh script. Upon acceptance of the paper, we will make the complete codebase publicly available on GitHub.

Performance on graph-level tasks

We extend our method to support graph-level tasks by incorporating edge features and Laplacian positional encodings. To evaluate the effectiveness of this extension, we conduct experiments on two graph classification datasets: OGBG-MOLBACE and OGBG-MOLBBBP. Due to time constraints, we compare our method against three widely-recognized GNN baselines—GCN, GAT, and GINE—which have been shown to be strong baselines on graph-level benchmarks [4]. The results are summarized in the following table:

As shown in the table, our extended model outperforms the baselines on both datasets, demonstrating its strong potential for graph-level tasks.

[1] Xu, Zhe, et al. "How to make LLMs strong node classifiers?." arXiv preprint arXiv:2410.02296 (2024).

[2] Wang, Limei, et al. "Learning graph quantized tokenizers." arXiv preprint arXiv:2410.13798 (2024).

[3] Deng, Chenhui, et al. "Polynormer: Polynomial-expressive graph transformer in linear time." arXiv preprint arXiv:2403.01232 (2024).

[4] Luo, Yuankai, et al. "Can Classic GNNs Be Strong Baselines for Graph-level Tasks? Simple Architectures Meet Excellence." arXiv preprint arXiv:2502.09263 (2025).