Yilun Du, Shuang Li, Joshua B. Tenenbaum, Igor Mordatch

A link to our paper can be found on arXiv.

Official codebase for Learning Iterative Reasoning through Energy Minimization. Contains scripts to reproduce experiments.

Please install the listed requirements.txt file.

pip install -r requirements.txt

We provide code for running continuous graph reasoning experiments in graph_train.py and code for running continuous matrix experiments in train.py.

To run continuous matrix addition experiments you utilize the following command:

python train.py --exp=addition_experiment --train --num_steps=10 --dataset=addition --train --cuda --infinite

To evaluate the final performance of the model after training, you may use the command:

python train.py --exp=addition_experiment  --num_steps=10 --dataset=addition --cuda --infinite --resume_iter=10000 

and the following command for the OOD test set:

python train.py --exp=addition_experiment  --num_steps=10 --dataset=addition --cuda --infinite --resume_iter=10000  --ood

We may substitute the flag --dataset with other keywords such as inverse or lowrank (as well as additional ones defined in dataset.py).

To run discrete graph reasoning experiments you may utilize the following command:

python graph_train.py --exp=identity_experiment --train --num_steps=10 --dataset=identity --train --cuda --infinite

We may substitute the flag --dataset with other datasets such as shortestpath or connected (as well as additional ones defined in graph_dataset.py).

To evaluate the model, you may then utilize the following command:

python graph_train.py --exp=identity_experiment --num_steps=10 --dataset=identity --cuda --infinite  --resume_iter=10000

Please cite our paper as:

@article{du2022irem,
  title={Learning Iterative Reasoning through Energy Minimization},
  author={Yilun Du and Shuang Li and Joshua B. Tenenbaum and Igor Mordatch},
  booktitle={Proceedings of the 39th International Conference on Machine 
                    Learning (ICML-22)},
  year={2022}
}

Note: this is not an official Google product.

MIT