Probabilistically Rewired Message-Passing Neural Networks

Reference implementation of our rewiring method as proposed in

Probabilistically Rewired Message-Passing Neural Networks
Chendi Qian*, Andrei Manolache*, Kareem Ahmed, Zhe Zeng, Guy Van den Broeck, Mathias Niepert^†, Christopher Morris^†

*These authors contributed equally.
^†Co-senior authorship.

Environment setup

conda create -y -n prmpnn python=3.10
conda activate prmpnn

conda install pytorch==2.1.2 pytorch-cuda=11.8 -c pytorch -c nvidia
pip install torch_geometric==2.4.0  # maybe latest also works
pip install https://data.pyg.org/whl/torch-2.1.0%2Bcu118/torch_scatter-2.1.2%2Bpt21cu118-cp310-cp310-linux_x86_64.whl
pip install https://data.pyg.org/whl/torch-2.1.0%2Bcu118/torch_sparse-0.6.18%2Bpt21cu118-cp310-cp310-linux_x86_64.whl

pip install ogb
pip install ml-collections
pip install sacred
pip install wandb
pip install gdown

# maybe need to downgrade numpy
pip install numpy==1.26.4

Datasets

We empirically evaluate our rewiring method on multiple datasets.

Real-world datasets

TUDatasets: PyG class, paper

• ZINC
• Alchemy
• MUTAG
• PRC_MR
• PROTEINS
• NCI1
• NCI109
• IMDB-B
• IMDB-M

OGB: website, paper

• ogbg-molhiv

WebKB: PyG class

• Cornell
• Texas
• Wisconsin

LRGB: code, paper

• peptides-func
• peptides-struct

QM9 used in DRew and SP-MPNN. Note there are different versions of QM9, e.g., PPGN

Synthetic datasets

EXP: code, paper

CSL: code, paper

Trees-NeighborsMatch: code, paper

Trees-LeafColor: Our own ⭐ ⭐ ⭐

Rewire options

We provide rewiring options as following:

• Add edges / remove edges

• Directed / undirected: meaning adding or deleting edges in a directed way or not. If not, will add and remove undirected edges.

• Separated / merged: if separated, will sample 2 graphs, one with edges added and the other with edges removed. If merged, will merge the 2 graphs as one.

Sampler candidates

• SIMPLE, code, paper
• I-MLE, code, paper
• Gumbel softmax for subset sampling

To replicate experiments

We provide yaml files under configs, run e.g. python run.py with PATH_TO_CONFIG

Note that this repo provides a taste of how PR-MPNN works, with examples given by GIN network. For replicating the results in our paper, please see to the backup branch, or contact chendi.qian@log.rwth-aachen.de

Known issue

If you are using a different version of PyTorch, you might have some error in the gradients produced by SIMPLE gradient estimator. You might want to check if SIMPLE gives nonzero gradients.