conda create -n rehub python=3.10
conda activate rehub
conda install pytorch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 pytorch-cuda=12.4 -c pytorch -c nvidia
pip install torch_geometric
pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.4.0+cu124.html
conda install lightning -c conda-forge
pip install pymetis yacs torchmetrics wandb ogb pandas einops performer-pytorch
conda install openbabel fsspec rdkit -c conda-forge
conda clean --allconda activate rehub
# Running ReHub for LRGB Datasets
python main.py --cfg configs/ReHub/peptides-func-ReHub.yaml wandb.use False
python main.py --cfg configs/ReHub/peptides-struct-ReHub.yaml wandb.use False
python main.py --cfg configs/ReHub/pcqm-contact-ReHub.yaml wandb.use False
# Running ReHub for OGBN-Arxiv
python main.py --cfg configs/Rehubs/ogbn-arxiv-ReHub.yaml wandb.use FalseThe config files structure are the same as in GraphGPS. For Rehub you can change the following config parameters:
rehub:
prep: True # Set True to enable preprocessing of partitioning the graph
# into clusters.
hubs_per_spoke: 3 # Set the number of hubs each spoke would connect to.
num_hubs_type: 'D' # Options are 'D' for dynamically ajusting the number of hubs
# according the graph size, or 'S' for a static amount of hubs
# for all graphs.
num_hubs: 1.0 # For num_hubs_type='D' this value is the Hub Ratio, meaning
# that for each graph there would be num_hubs*sqrt(#Spokes) hubs.
# For num_hubs_type='S' this value is the number of hubs.
spokes_mlp_before_hub_agg: True # Adds an MLP for the spokes features before they
# are aggregated to the hub features
reassignment_strategy: 'k_closest_by_attention' # For the reassignment strategy
# introduced in the paper choose 'k_closest_by_attention'.
# To disable reassignment replace with None.
The most substanial part of the model can be found in:
graphgps/network/rehub_model.pywhere hubs are initialized in theHubCreatorclass.graphgps/layer/rehub_layer.pywhere theReHubclass holds the global part of "Long-Range Spoke Update Layer". i.e. from "Spokes to Hubs" to "Hub (Re)Assignment".graphgps/transform/spoke_hub_clusters.pywhich clusters each graph according to the number of hubs.
We also provide code for the "Large Random Regular Graphs" custom dataset at graphgps/loader/dataset/large_random_graph.py.
Additional util files are available on:
graphgps/layer/rehub_utils.pyfor hub reassignment and metrics.graphgps/train/metric_utils.pyfor metrics.
Our work can be cited using the following bibtex:
@article{
borreda2025rehub,
title={ReHub: Linear Complexity Graph Transformers with Adaptive Hub-Spoke Reassignment},
author={Tomer Borreda and Daniel Freedman and Or Litany},
journal={Transactions on Machine Learning Research},
issn={2835-8856},
year={2025},
url={https://openreview.net/forum?id=L4S54TUOQR}
}