Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Unifying and Enhancing Graph Transformers via a Hierarchical Mask Framework

Yujie Xing, Xiao Wang, Bin Wu, Hai Huang, Chuan Shi

TL;DR

This work tackles rigid Graph Transformer designs by introducing a unified hierarchical mask framework that ties architectural choices to attention-mask construction, enabling flexible modeling of local, cluster, and global interactions. It derives a design principle showing the probability of correct classification improves with a larger receptive field and greater label consistency, formalized through bounds involving $k$, $\rho_c$, and $\alpha_c$. Building on this, the authors propose M$^{3}$Dphormer, a Mixture-of-Experts GT with three theoretically grounded masks ($\mathbf{M}^{l2},\mathbf{M}^{c4},\mathbf{M}^{g3}$), a bi-level expert routing mechanism, and a dual attention computation scheme for efficiency. Empirical results across nine datasets demonstrate state-of-the-art performance and reveal the practical benefits of multi-level masking and adaptive routing, while also highlighting memory-efficiency advantages of the dual-attention approach. Overall, the framework provides a principled, scalable path for designing versatile Graph Transformers that capture diverse hierarchical interactions.

Abstract

Graph Transformers (GTs) have emerged as a powerful paradigm for graph representation learning due to their ability to model diverse node interactions. However, existing GTs often rely on intricate architectural designs tailored to specific interactions, limiting their flexibility. To address this, we propose a unified hierarchical mask framework that reveals an underlying equivalence between model architecture and attention mask construction. This framework enables a consistent modeling paradigm by capturing diverse interactions through carefully designed attention masks. Theoretical analysis under this framework demonstrates that the probability of correct classification positively correlates with the receptive field size and label consistency, leading to a fundamental design principle: an effective attention mask should ensure both a sufficiently large receptive field and a high level of label consistency. While no single existing mask satisfies this principle across all scenarios, our analysis reveals that hierarchical masks offer complementary strengths, motivating their effective integration. Then, we introduce M3Dphormer, a Mixture-of-Experts-based Graph Transformer with Multi-Level Masking and Dual Attention Computation. M3Dphormer incorporates three theoretically grounded hierarchical masks and employs a bi-level expert routing mechanism to adaptively integrate multi-level interaction information. To ensure scalability, we further introduce a dual attention computation scheme that dynamically switches between dense and sparse modes based on local mask sparsity. Extensive experiments across multiple benchmarks demonstrate that M3Dphormer achieves state-of-the-art performance, validating the effectiveness of our unified framework and model design.

Unifying and Enhancing Graph Transformers via a Hierarchical Mask Framework

TL;DR

This work tackles rigid Graph Transformer designs by introducing a unified hierarchical mask framework that ties architectural choices to attention-mask construction, enabling flexible modeling of local, cluster, and global interactions. It derives a design principle showing the probability of correct classification improves with a larger receptive field and greater label consistency, formalized through bounds involving , , and . Building on this, the authors propose MDphormer, a Mixture-of-Experts GT with three theoretically grounded masks (), a bi-level expert routing mechanism, and a dual attention computation scheme for efficiency. Empirical results across nine datasets demonstrate state-of-the-art performance and reveal the practical benefits of multi-level masking and adaptive routing, while also highlighting memory-efficiency advantages of the dual-attention approach. Overall, the framework provides a principled, scalable path for designing versatile Graph Transformers that capture diverse hierarchical interactions.

Abstract

Graph Transformers (GTs) have emerged as a powerful paradigm for graph representation learning due to their ability to model diverse node interactions. However, existing GTs often rely on intricate architectural designs tailored to specific interactions, limiting their flexibility. To address this, we propose a unified hierarchical mask framework that reveals an underlying equivalence between model architecture and attention mask construction. This framework enables a consistent modeling paradigm by capturing diverse interactions through carefully designed attention masks. Theoretical analysis under this framework demonstrates that the probability of correct classification positively correlates with the receptive field size and label consistency, leading to a fundamental design principle: an effective attention mask should ensure both a sufficiently large receptive field and a high level of label consistency. While no single existing mask satisfies this principle across all scenarios, our analysis reveals that hierarchical masks offer complementary strengths, motivating their effective integration. Then, we introduce M3Dphormer, a Mixture-of-Experts-based Graph Transformer with Multi-Level Masking and Dual Attention Computation. M3Dphormer incorporates three theoretically grounded hierarchical masks and employs a bi-level expert routing mechanism to adaptively integrate multi-level interaction information. To ensure scalability, we further introduce a dual attention computation scheme that dynamically switches between dense and sparse modes based on local mask sparsity. Extensive experiments across multiple benchmarks demonstrate that M3Dphormer achieves state-of-the-art performance, validating the effectiveness of our unified framework and model design.
Paper Structure (36 sections, 6 theorems, 58 equations, 7 figures, 8 tables, 2 algorithms)

This paper contains 36 sections, 6 theorems, 58 equations, 7 figures, 8 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Preliminary
  3. Revisiting GTs through a unified hierarchical mask framework
  4. The unified hierarchical mask framework for Graph Transformers
  5. Theoretical analysis
  6. Experimental analysis
  7. Method
  8. Overall architecture of M$^3$Dphormer
  9. Theorem-guided hierarchical mask design strategy
  10. Bi-level attention expert routing mechanism
  11. Dual attention computation scheme
  12. Experiments
  13. Related Work
  14. Conclusion
  15. Theorem Proofs
  16. ...and 21 more sections

Key Result

Theorem 3.1

The updated representation of node $u$ follows a Gaussian distribution: Assume the classifier is well-trained such that $\delta_c - k \rho_c \alpha_c \le 0$, and the attention weights satisfy the constraint $0 \le \alpha_{c'} \le \frac{1}{k} \le \alpha_c \le \frac{1}{k \rho_c}$. Then, the probability that node $u$ is correctly classified by a similarity-based classifier where $\Phi(\cdot)$ denote

Figures (7)

  • Figure 1: Hierarchical interactions.
  • Figure 2: Overview of the M$^3$Dphormer with Pre-RMSNorm rmsnorm19 and ReLUrelu18: a) The overall network architecture. b) The bi-level expert routing mechanism. c) The theorem-guided hierarchical mask design strategy. Self-connections are added for nodes in local and cluster masks. Nodes 2, 4, 5 and 8 are included in the training set. d) The dual attention computation scheme.
  • Figure 3: Memory usage of different schemes.
  • Figure 4: Test accuracy vs. $P$
  • Figure 5: Comparison of accuracy and loss curves between M$^3$DFormer and GCN* on Ogbn-Arxiv.
  • ...and 2 more figures

Theorems & Definitions (11)

  • Theorem 3.1
  • Proposition 4.1
  • Proposition 4.2
  • Theorem A.1
  • Theorem A.2
  • Proposition A.3
  • proof
  • proof
  • proof
  • proof
  • ...and 1 more