A Benchmark for Variational Sampling Methods

This repository accompanies the paper "Beyond ELBOs: A Large-Scale Evaluation of Variational Methods for Sampling. [ICML'24,BibTeX]" It includes a variety of sampling algorithms, such as sequential importance sampling algorithms and diffusion-based methods. A complete list of the available algorithms is provided below. Additionally, we offer several target densities, both low- and high-dimensional, with varying levels of complexity.

Disclaimer

The codebase is currently under development. We are in the process of unifying various parts of the code, including migrating algorithms from the Haiku library to Flax. As a result, some outcomes may differ from those reported in the paper. These will be updated once the majority of the codebase is complete.

Available Algorithms

The table below provides a overview of all available algorithms.

Acronym	Method	Reference
MFVI	Gaussian Mean-Field VI	Bishop, 2006
GMMVI	Gaussian Mixture Model VI	Arenz et al., 2022
NFVI	Normalizing Flow VI	Rezende & Mohamed, 2015
SMC	Sequential Monte Carlo	Del Moral et al., 2006
AFT	Annealed Flow Transport	Arbel et al., 2021
CRAFT	Continual Repeated AFT	Matthews et al., 2022
FAB	Flow Annealed IS Bootstrap	Midgley et al., 2022
ULA	Uncorrected Langevin Annealing	Thin et al., 2021
MCD	Monte Carlo Diffusion	Doucet et al., 2022
UHA	Uncorrected Hamiltonian Annealing	Geffner et al., 2021
LDVI	Langevin Diffusion VI	Geffner et al., 2022
CMCD	Controlled MCD	Vargas et al., 2023
PIS	Path Integral Sampler	Zhang et al., 2021
DIS	Time-Reversed Diffusion Sampler	Berner et al., 2022
DDS	Denoising Diffusion Sampler	Vargas et al., 2023
GFN	Generative Flow Networks	Lahlou et al., 2023
GBS	General Bridge Sampler	Richter et al., 2023
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
Sonar	61	❌	❌	sonar
Digits	196/784	✔️	✔️	nice_digits
Fashion	784	✔️	✔️	nice_fashion
LGCP	1600	❌	❌	lgcp
MoG	any	✔️	✔️	gaussian_mixture40
MoS	any	✔️	✔️	student_t_mixture

The respective configuration files can be found here.

Installation

First, clone the repo:

git clone git@github.com:DenisBless/variational_sampling_methods.git
cd variational_sampling_methods

We recommend using Conda to set up the codebase:

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

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.

@misc{blessing2024elbos,
      title={Beyond ELBOs: A Large-Scale Evaluation of Variational Methods for Sampling}, 
      author={Denis Blessing and Xiaogang Jia and Johannes Esslinger and Francisco Vargas and Gerhard Neumann},
      year={2024},
      eprint={2406.07423},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

Acknowledgements

Portions of the project are adapted from other repositories (as mentioned in the code):

• https://github.com/shreyaspadhy/CMCD is licensed under MIT,
• https://github.com/OlegArenz/gmmvi is licensed under MIT,
• https://github.com/lollcat/fab-jax is licensed under MIT,
• https://github.com/tomsons22/LDVI is licensed under MIT,
• https://github.com/juliusberner/sde_sampler is licensed under MIT,
• https://github.com/franciscovargas/denoising_diffusion_samplers is licensed under MIT,
• https://github.com/antoninschrab/mmdfuse is licensed under MIT,
• https://github.com/google-deepmind/annealed_flow_transport is licensed under Apache-2.0.