Efficient Training of Boltzmann Generators Using Off-Policy Log-Dispersion Regularization

This repository contains the code to reproduce the experiments of our paper Efficient Training of Boltzmann Generators Using Off-Policy Log-Dispersion Regularization, accepted at ICML 2026.

This project studies off-policy log-dispersion regularization (LDR) for training Boltzmann generators. LDR uses target energy labels from fixed datasets as an additional regularization signal, improving final performance and data efficiency across unbiased data, biased data, and variational training settings. For methodological details, please refer to the paper.

Note

This repository builds on code bases developed for our previous publications. The implementation was first introduced in TA-BG, then extended for CMT, and has now been further extended for the LDR experiments in this work.

Requirements

An environment with all dependencies can be installed in the following way:

conda env create -f environment.yaml

Since we use weights and biases to track experiments, you first need to login to your account:

wandb login

Downloading ground truth datasets

Ground truth datasets can be downloaded from Zenodo: https://doi.org/10.5281/zenodo.18822445

The Zenodo entry contains the archives data.zip and data_ldr.zip. Extract both archives and place the contents of the contained data/ and data_ldr/ directories in:

./annealed_bg/data/

Running the experiments

All commands to reproduce the experiments, including the required hyperparameters, are listed in experiments.md.

Using AutoSlurm

This repository supports AutoSlurm, which provides automatic resuming. This feature is essential for larger systems that often require more time than the maximum job limit on GPU clusters. Using AutoSlurm is recommended for starting training runs on the larger molecular systems.

Questions and issues

If anything is unclear or does not work as expected, please open a GitHub issue.