Equivariant Subgraph Aggregation Networks (ESAN)

This repository contains the official code of the paper
Equivariant Subgraph Aggregation Networks (ICLR 2022)


First create a conda environment

conda env create -f environment.yml

and activate it.

Prepare the data


python data.py --dataset $DATASET

where $DATASET is one of the following:

  • PTC
  • NCI1
  • NCI109
  • ogbg-molhiv
  • ogbg-moltox21
  • ZINC
  • CSL
  • EXP
  • CEXP

Run the models

To perform hyperparameter tuning, make use of wandb:

  1. In configs/ folder, choose the yaml file corresponding to the dataset and setting (deterministic vs sampling) of interest, say <config-name>. This file contains the hyperparameters grid.

  2. Run

    wandb sweep configs/<config-name>

    to obtain a sweep id <sweep-id>

  3. Run the hyperparameter tuning with

    wandb agent <sweep-id>

    You can run the above command multiple times on each machine you would like to contribute to the grid-search

  4. Open your project in your wandb account on the browser to see the results:

    • For the TUDatasets, the CSL and the EXP/CEXP datasets, refer to Metric/valid_mean and Metric/valid_std to obtain the results.

    • For the ogbg datasets and the ZINC dataset, compute mean and std of Metric/train_mean, Metric/valid_mean, Metric/test_mean over the different seeds of the same configuration.
      Then, take the results corresponding to the configuration obtaining the best validation metric.


For attribution in academic contexts, please cite

title={Equivariant Subgraph Aggregation Networks},
author={Beatrice Bevilacqua and Fabrizio Frasca and Derek Lim and Balasubramaniam Srinivasan and Chen Cai and Gopinath Balamurugan and Michael M. Bronstein and Haggai Maron},
booktitle={International Conference on Learning Representations},


View Github