Towards Better Out-of-Distribution Generalization of Neural Algorithmic Reasoning Tasks

This is the code to "Towards Better Out-of-Distribution Generalization of Neural Algorithmic Reasoning Tasks".

Dependencies

Please refer to CLRS benchmark's documentation for instructions on how to install this library and dependencies. Note that a separate virtualenv is needed for installing package from this repository if you have already installed CLRS.

Generating Datasets

To generate the datasets reported in our paper, you should run the following command:

CLRS_VENV_PATH=/path/to/venv \
CLRS_ROOT=/path/to/clrs \
CLRS_DATASET_PATH=/path/to/clrs/dataset \
  ./scripts/make_datasets.sh

where CLRS_VENV_PATH is the path to the virtual environment containing CRLS, CLRS_ROOT is the path to the CLRS code root, and CLRS_DATASET_PATH is the desired directory for creating the CLRS dataset.

Reproducing the Results

First, generate the datasets using the above command. Then, run the following script:

CLRS_DATASET_PATH=/path/to/saved/datasets \
CLRS_LOG_PATH=/tmp/clrs_logs \
CLRS_CHECKPOINT_PATH=/tmp/clrs_checkpoints \
  ./scripts/run_experiments.sh

where CLRS_DATASET_PATH is the path to generated datasets (created in above command), CLRS_LOG_PATH is the path for saving run logs, and CLRS_CHECKPOINT_PATH is the path for saving checkpoints.

The raw logs will be saved in the mentioned logging directory. In addition to that, WandB summaries will also be saved locally.

Colab

We also provide a simplified notebook in Google Colab, which contains data generation, training, and visualization for BFS algorithm.

Notes

Datasets are generated without any hints. This saves 10X space and allows for larger number of samples to be generated. In case you would like the hints to be included, you need to disable CLRS_DISABLE_HINTS flag inside the dataset generation script. However, there are no mechanisms for efficient dataset generation at the moment, and you need a high amount of RAM for dataset generation.
Although the seed variable controls the seed for neural network initialization, the order in which batches are created is not controlled. As a result, exact reproduction of numbers is not guaranteed in the CLRS benchmark.
The experiments script runs all the experiments in a sequential order. You might want to make them parallel depending on your cluster system.
The code base also contains the ideas mentioned in appendix of the paper, but are deactivated by default.
The logs are by default written to their corresponding log directory and WandB locally. In case of debugging, you can activate the debug flag to see the logs in stdout and disable WandB logging.
The values of parameters in Google Colab notebook are changed and simplified to facilitate readability and quick training. For complete set of parameters used in our paper please refer to out GitHub code base.

Cite

Please consider citing our paper if you use this code in your research work:

@article{
  mahdavi2023towards,
  title={Towards Better Out-of-Distribution Generalization of Neural Algorithmic Reasoning Tasks},
  author={Sadegh Mahdavi and Kevin Swersky and Thomas Kipf and Milad Hashemi and Christos Thrampoulidis and Renjie Liao},
  journal={Transactions on Machine Learning Research},
  issn={2835-8856},
  year={2023},
  url={https://openreview.net/forum?id=xkrtvHlp3P},
  note={}
}

Questions/Bugs

Please submit a Github issue or contact smahdavi@ece.ubc.ca if you have any questions or find any bugs.

DSL-Lab/clrs