Official implementation of LEAPS

This is the official implementation of LEAPS: A discrete neural sampler via locally equivariant neural networks, a work by Peter Holderrieth*, Michael S. Albergo*, and Tommi S. Jaakkola.

Setup environment

conda env create -f environment.yml

Run experiments for Ising and Potts model

1. Edit configs: Edit the configs in ymls/ising_config.yml or ymls/potts_config.yml to account for the filepaths on your server, i.e. change code_path, ckpts_dir and results_path.

2. Run commands: Run code with the following commands:

python main.py --yml=ymls/ising_config.yml

python main.py --yml=ymls/potts_config.yml

Run experiments for different energy model

To run experiments on a different energy model, you need to define an annealing path $\rho_t$. You can use the examples in src/modules/rhot.py for the Ising model and the Potts model as guiding examples.

Model checkpoints

We provide two model checkpoints for the two Boltzmann distributions we studied:

• Ising model at critical temperature: ckpts/ising_final.ckpt
• Potts model at critical temperature: ckpts/potts_final.ckpt

Analyis of trained models

We have two notebooks that illustrate how to load models and evaluate their ESS:

• notebooks/load_ising_network.ipynb
• notebooks/load_potts_network.ipynb

DISCS benchmark

In our work, we benchmarked against MCMC samplers using the DISCS benchmark. To replicate these experiments, perform the following steps:

1. Go to DISCS repo and clone it.
2. Add the model configs in discs_configs/ that you can find in this repo as a new model config to the DISCS benchmark in discs/models/configs/ within the DISCS repo.
3. Copy the scripts discs_configs/isingcritical_run_discs_benchmark.sh or discs_configs/pottscritical_run_discs_benchmark.sh to the DISCS repo.
4. Run the scripts.