readme.md

Source code of our NeurIPS 2025 paper "Unifying and Enhancing Graph Transformers via a Hierarchical Mask Framework"

Requirements

python==3.10.14
torch==2.3.1+cu121 --index-url https://download.pytorch.org/whl/cu121
ogb==1.3.6
torch-geometric==2.5.3
torch-scatter==2.1.2+pt23cu121 --index-url https://download.pytorch.org/whl/cu121
torch-sparse==0.6.18+pt23cu121 --index-url https://download.pytorch.org/whl/cu121
torch-cluster==1.6.3+pt23cu121 --index-url https://download.pytorch.org/whl/cu121
metis

Run Experiments

Run main.py with the commands from runs.sh, for example:

python main.py --dataset cora --lr 0.001 --wd 0.0001 --h_dim 128 --layers 3 --n_head 2 --dropout 0.5 --attn_dropout 0.3 --norm_type rms --norm_pos pre --local_type GAT --num_clusters 128 --use_res --epochs 200 --runs 5 --display_step 100 --split random --device 0

Run Experiments with Checkpoints

To ensure reproducibility, we provide pre-trained checkpoints for 5 datasets in the models directory. You can evaluate these models using eval.py with the same command-line arguments as in runs.sh, for example:

python eval.py --dataset cora --lr 0.001 --wd 0.0001 --h_dim 128 --layers 3 --n_head 2 --dropout 0.5 --attn_dropout 0.3 --norm_type rms --norm_pos pre --local_type GAT --num_clusters 128 --use_res --epochs 200 --runs 5 --display_step 100 --split random --device 0