mattiascu/pr-mpnn

Official implementation for paper: Probabilistically Rewired Message-Passing Neural Networks, accepted at ICLR 2024.

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  pytorch-cuda=11.7 -c pytorch-nightly -c nvidia
conda install openbabel fsspec rdkit -c conda-forge
pip install cmake
pip install --verbose git+https://github.com/pyg-team/pyg-lib.git
pip install --verbose torch_scatter
pip install --verbose torch_sparse
pip install --verbose torch_geometric

pip install ogb
pip install ml-collections
pip install numba
pip install sacred
pip install PyYAML
pip install wandb
pip install matplotlib
pip install seaborn
pip install GraphRicciCurvature
pip install gdown

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.

• In-place / not in-place: if in-place, will add the edges based on the original edges, otherwise will return a graph with only the added edges.

Samplers that we use

• 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 configs/zinc/edge_candidate/best.yaml