Adjoint-based diffusion samplers is a new class of diffusion models for energy-based Boltzmann distributions that is highly scalable, extremely data-efficient, and achieves strong performance. This repository contains the implementation of adjoint-based diffusion samplers for synthetic energies, e.g., Double Well and Lennard Jones potentials. Currently, the repository contains the following adjoint-based diffusion samplers:

• Adjoint Sampling (ICML 2025)
• Adjoint Schrödinger Bridge Sampler (NeurIPS 2025 Oral)

Note that AS is a special case of ASBS with memoryless condition. For amortized conformer generation, please check here. For improved ASBS with chemically-sounded exploration (WT-ASBS), please check here.

Both Anaconda and micromamba should work. We recommend micromamba for its faster compiling.

micromamba env create -f environment.yml
micromamba activate adjoint_samplers

Run scripts/demo.sh to generate similar demo figure in ASBS paper.

Run the following script to download the necessary reference samples to the folder data for evaluation purposes:

bash scripts/download.sh

Training scripts to generate similar results in the papers can be found under scripts. Checkpoints and figures are saved under the folder results.

python train.py experiment={dw4,lj13,lj55}_{asbs,as} seed=0,1,2 -m

If you find this repository helpful, please consider citing our paper:

@inproceedings{liu2025asbs,
  title={{Adjoint Schr{\"o}dinger bridge sampler}},
  author={Liu, Guan-Horng and Choi, Jaemoo and Chen, Yongxin and Miller, Benjamin Kurt and Chen, Ricky T. Q.},
  booktitle={Advances in Neural Information Processing Systems (NeurIPS)},
  year={2025},
}

This repository is licensed under the CC BY-NC 4.0 License, with some portions of the project subject to separate license terms: DEM, DDS, and bgflow are each licensed under the MIT License. Please refer to the respective repositories for details.