N2GNN

This repository is the official implementation of the N2-GNN proposed in the [NeurIPS23]Extending the Design Space of Graph Neural Networks by Rethinking Folklore Weisfeiler-Lehman (previous named "Towards Arbitrarily Expressive GNNs in $O(n^2)$ Space by Rethinking Folklore Weisfeiler-Lehman").

🔥News

• 2024. 01: Fix a bug on ZINC-full dataset. In the previous code, it uses the test set of ZINC-subset as the test set and causes the wrong result. After the correction, the final result for ZINC-full is 0.022 ± 0.002. We update the paper accordingly. We sincerely apologize for the mistake.

• 2023. 10: Update code base to version 2.0. In version 2.0, we:

  1. Improve the implementation of N2GNN. The current implementation is more memory saving than the previous one. The detail is provided in the camera-ready version of the paper.
  2. Add experiment for BREC dataset. We achieve new SOTA among all GNN methods!
  3. Minor bugs fixing and comment polishing.

• 2023. 09: Our paper is accepted at NeurIPS 2023!

Requirements

python=3.8
torch=2.0.1
PyG=2.3.1
pytorch_lightning=2.0.9
wandb=0.13.11

We also provide a docker environment and corresponding docker file if users prefer a docker style setup. Before run the experiments, you may need to set up wandb.

Usage

The result for all experiments will be uploaded to Wandb server with the project named as: [task]_<dataset_name>_<GNN_layer>_<GNN_model>_<num_layer>_<hidden_channels>_<data_type>_<num_hop>_[rd].
task: The number index of the task if dataset has multiple tasks (QM9; substructure counting).
GNN_layer: Name of GNN layer. See detail in models/gnn_conv.py.
GNN_model: Name of GNN model. See detail in models/GNNs.py.
num_layer: Number of layer in GNN model.
hidden_channels: Number of hidden channels for each layer.
data_type: Data preprocessing type. See detail in data_utils.py.
num_hop: Number of hop for overlapping subgraph.
rd: Name ends with rd means model add resistance distance as additional feature.

Reproducibility

ZINC-Subset and ZINC-Full

For ZINC-Subset:

python train_zinc.py --config_file=configs/zinc.yaml

For ZINC-Full:

python train_zinc.py --config_file=configs/zinc_full.yaml

Counting substructure

For cycle counting, set data=count_cycle and task=0, 1, 2, 3 for 3-cycles, 4-cycles, 5-cycles, and 6-cycles, respectively.
For substructure counting, set data=count_graphlet and task=0, 1, 2, 3, 4 for tailed-triangles, chordal cycles, 4-cliques, 4-paths, and triangle-rectangle, respectively. Additionally, set the number of hop h for different tasks. You can find exact setting in the paper.

python train_count.py --config_file=configs/count.yaml --task=$task --dataset_name=$data --num_hops=$h

QM9

Run a single target:

python train_qm9.py --config_file=configs/qm9.yaml --task=7

Run all targets:

python train_qm9.py --config_file=configs/qm9.yaml --search

SR25

python train_sr25.py --config_file=configs/sr25.yaml

CSL

python train_CSL.py --config_file=configs/csl.yaml

EXP

python train_EXP.py --config_file=configs/exp.yaml

BREC

To run the experiment on BREC dataset, you need first download brec_v3_no4v_60cfi.npy from BREC repository and put it in the data/BREC directory. Then, run the following command:

python train_BREC.py --config_file=configs/BREC.yaml

Contact

Jiarui Feng ([email protected])

Citation

If you find this work useful, please kindly cite our paper:

@inproceedings{
Feng2023extending,
title={Extending the Design Space of Graph Neural Networks by Rethinking Folklore Weisfeiler-Lehman},
author={Jiarui Feng and Lecheng Kong and Hao Liu and Dacheng Tao and Fuhai Li and Muhan Zhang and Yixin Chen},
booktitle={Advances in Neural Information Processing Systems},
year={2023}
}