Convexified Message-Passing Graph Neural Networks

By Saar Cohen, Noa Agmon, Uri Shaham [arXiv]

Code for implementation of convexified message-passing graph neural network (CGNN), evaluated through extensive experiments on several benchmark graph classification datasets.

If any part of this code is used, the following paper must be cited:

Saar Cohen, Noa Agmon and Uri Shaham. Convexified Message-Passing Graph Neural Networks. In AISTATS'26: the 29th Annual Conference on Artificial Intelligence and Statistics, 2026 (To Appear).

Experiments

We apply our convexification procedure to the message-passing mechanisms of popular GNNs and hybrid graph transformers, creating convex counterparts. All models share the same architecture design for fair comparison.

• The directory baselines includes implementation of non-convex baselines.
• The file graphClassCxGNN.py implements convexified message-passing GNNs.

Setup and Dependencies

Requirements are in setup.sh. Simply running bash setup.sh will create a conda environment called torch-1-9 that usually works for our experiments, which can then be activated through the command conda activate torch-1-9.

In particular, additional dependencies include:

• PyTorch 1.9
• PyTorch Geometric 2.0.1.
• PyTorch Scatter 2.0.9
• PyTorch Sparse 0.6.12
• PyTorch Spline Conv
• TUDataset
• Open Graph Benchmark (OGB)

Note: You may have to edit the CUDA variable in setup.sh depending on the CUDA version of your GPUs.

Execution

Experiments for Convexified Message-Passing GNNs (CGNNs)

For running an experiment on a CGNN of type CGNN, execute the following command:

python graphClassCxGNN.py --model_name CGNN --dataset dataset_name --nhid nhid --nlayer nlayer --seed seed --saveDirRoot save_dir --nEpochs n_epochs --batchSize batch_size --maxHops max_hops --nystromDimsCxAggGNN nystromDims

• dataset_name - Name of the dataset on which the experiment should be performe (examples include DD/PROTEINS/NCI1/NCI109/Mutagenicity/ENZYMES). For datasets from TUDataset, one can download a .zip file of the desired dataset, upload it to a directory titles datasets and then the dataset will be extracted. Verify whether the dataset`s subdirectories are located in their proper location.
• nhid - Number of hidden units
• nlayer - Number of layers
• seed - Seed value
• save_dir - The name of the directory to which all files related to the experiments will be saved (e.g., the model`s training, validation and test accuracies).
• n_epochs - Number of epochs
• batch_size - Batch size
• max_hops - Maximum number of hops for the GNN`s information aggregation
• nystromDims - The Nystrom dimensions used for the convexification procedure

Experiments for Non-Convex Message-Passing GNNs

For running an experiment on a non-convex GNN of type non-convex GNN, execute the following command (this command can receive arguments similar to the above):

python baselines/main.py --model_name non-convex GNN