Equivariance Everywhere All At Once: A Recipe for Graph Foundation Models

This repository contains the official code base of the paper Equivariance Everywhere All At Once: A Recipe for Graph Foundation Models

Installation

To reproduce the results please use Python 3.10, PyTorch version 2.3.0, Cuda 11.8, PyG version 2.5.3.

pip3 install torch==2.3.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip3 install torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.3.0+cu118.html
pip3 install torch-geometric==2.5.3
pip3 install neptune
pip3 install triton

hyper-parameters

The script we use to run our frameworks is ./main.py. Note that the script should be run with . as the main directory or source root.

The parameters of the script are:

• --project: name of your neptune project.
• --is_train: whether to train a new model or test an existing one.
• --train_test_setup: If 'trainset1' -- training on cora and testing on the remaining graphs. If 'inc_trainset' -- training on a trainset of size 'train_size' and testing on a heldout set of 20 graphs.
• --train_size: Only effective when train_test_setup is 'inc_trainset' and equals the size of the trainset.
• --gnn_type: the type of the GNN used for the GFM.
• --hidden_dim: the hidden dimension.
• --num_layers: the number of layers.
• --ls_num_layers: the number of least square powers to consider.
• --lp_ratio: the ratio of labels in the trainset that are propagated (the rest are used as part of the loss). Note that in test-time, the whole train-set is used for propagation.
• --max_epochs: the number of epochs used.
• --checkpoints: additional epochs to save the model in.
• --lr: the learn rate.

An example

To perform training for a 2 layered GAT-based architecture with a lr of 0.001, while propagating 30% in train time:

python -u main.py --is_train --gnn_type GAT --lr 0.001 --lp_ratio 0.3

best configs

python -u main.py --is_train --train_test_setup trainset1 --gnn_type MEAN_GNN --hidden_dim 16 --num_layers 2
--lp_ratio 0.4 --max_epochs 2000 --lr 0.01 
python -u main.py --is_train --train_test_setup trainset1 --gnn_type GAT --hidden_dim 16 --num_layers 2
--lp_ratio 0.25 --max_epochs 2000 --lr 0.03 

python -u main.py --is_train --train_test_setup inc_trainset --train_size 5 --gnn_type MEAN_GNN --hidden_dim 16 --num_layers 2
--lp_ratio 0.6 --max_epochs 300 --lr 0.003 
python -u main.py --is_train --train_test_setup inc_trainset --train_size 7 --gnn_type MEAN_GNN --hidden_dim 16 --num_layers 2
--lp_ratio 0.6 --max_epochs 500 --lr 0.003 
python -u main.py --is_train --train_test_setup inc_trainset --train_size 9 --gnn_type MEAN_GNN --hidden_dim 16 --num_layers 2
--lp_ratio 0.7 --max_epochs 100 --lr 0.005 

Note that to test this models you run the same command, this time without the is_train flag.

Cite

If you make use of this code, or its accompanying paper, please cite this work as follows:

@misc{finkelshtein2025equivarianceeverywhere,
      title={Equivariance Everywhere All At Once: A Recipe for Graph Foundation Models}, 
      author={Ben Finkelshtein and İsmail İlkan Ceylan and Michael Bronstein and Ron Levie},
      year={2025},
      archivePrefix={arXiv},
}