Underdamped Diffusion Bridges with Applications to Sampling

This repository accompanies the paper "Underdamped Diffusion Bridges with Applications to Sampling" [ICLR'25,BibTeX].

Available Algorithms

The table below provides a overview of all available algorithms.

Acronym	Method	Reference	ID
ULA	Uncorrected Langevin Annealing	Thin et al., 2021	ula
ULA-UD	Uncorrected Hamiltonian Annealing	Geffner et al., 2021	ula_ud
MCD	Monte Carlo Diffusion	Doucet et al., 2022	mcd
MCD-UD	Langevin Diffusion VI	Geffner et al., 2022	mcd_ud
DIS	Time-Reversed Diffusion Sampler	Berner et al., 2022	dis
DIS-UD	Underdamped Time-Reversed Diffusion Sampler	Ours	dis_ud
DBS	Diffusion Bridge Sampler	Richter et al., 2023	dbs
DBS-UD	Underdamped Diffusion Bridge Sampler	Ours	dbs_ud
The respective configuration files can be found here.

Available Target Densities

The table below provides a overview of available target densities. The 'ID' column provides identifier for running experiments via comand line. Further details in the Running Experiments section.

	dim	True log Z	Target Samples	ID
Funnel	10	✔️	✔️	funnel
Credit	25	❌	❌	credit
Seeds	26	❌	❌	seeds
Cancer	31	❌	❌	cancer
Brownian	32	❌	❌	brownian
Ionosphere	35	❌	❌	ionosphere
ManyWell	50	✔️	✔️	many_well
Sonar	61	❌	❌	sonar
LGCP	1600	❌	❌	lgcp

The respective configuration files can be found here.

Installation

First, clone the repo. For installation we recommend using Conda to set up the codebase:

conda create -n underdamped_sampling python==3.10.14 pip --yes
conda activate underdamped_sampling

Install the required packages using

pip install -r requirements.txt

Finally, we use wandb for experiment tracking. Login to your wandb account:

wandb login

You can also omit this step and add the use_wandb=False command line arg to your runs.

Running Experiments

Configuration

We use hydra for config management. The base configuration file sets parameters that are agnostic to the specific choice of algorithm and target density. The wandb entity can be set in the setup config file.

Running a single Experiment

In the simplest case, a single run can be started using

python run.py algorithm=<algorithm ID> target=<target ID>

The algorithm ID is identical to the Acronym in the algorithm table. The target ID can be found in the ID column of the target table.

Running multiple Experiments

Running multiple experiments can be done by using the hydra multi-run flag -m/--multirun flag. For instance, running multiple seeds can be done via

python run.py -m seed=0,1,2,3  algorithm=<algorithm ID> target=<target ID>

Using comma separation can also be used for running multiple algorithms/targets.

Running Experiments on a Cluster via Slurm

Running experiments on a cluster using Slurm can be done via

python run.py +launcher=slurm algorithm=<algorithm ID> target=<target ID>

which uses the slurm config config. Please make sure that you adapt the default settings to your slurm configuration.

References

If you use parts of this codebase in your research, please cite us using the following BibTeX entries.

@inproceedings{
blessing2025underdamped,
title={Underdamped Diffusion Bridges with Applications to Sampling},
author={Denis Blessing and Julius Berner and Lorenz Richter and Gerhard Neumann},
booktitle={The Thirteenth International Conference on Learning Representations},
year={2025},
url={https://openreview.net/forum?id=Q1QTxFm0Is}
}