Stochastic Localization via Iterative Posterior Sampling

Description

This is the official repository for the paper Stochastic Localization via Iterative Posterior Sampling. We consider a general stochastic localization framework and introduce an explicit class of observation processes, associated with flexible denoising schedules. We provide a complete methodology, Stochastic Localization via Iterative Posterior Sampling (SLIPS), to obtain approximate samples of this dynamics, and as a by-product, samples from the unnormalized target distribution. Our scheme is based on a Markov chain Monte Carlo estimation of the denoiser and comes with detailed practical guidelines. We illustrate the benefits and applicability of SLIPS on several benchmarks, including Gaussian mixtures in increasing dimensions, Bayesian logistic regression and a high-dimensional field system from statistical-mechanics. We experiment multiple tasks :

• Toy target distributions (8 Gaussians and Rings) (2 dimensions)
• Funnel (10 dimensions)
• Mixture of two Gaussians (from 8 up to 128 dimensions)
• Bayesian Logisitic Regressions (with 34 and 61 dimensions)
• Phi four field system (with 100 dimensions)

This code also contains implementation of the following papers :

• Reverse Diffusion Monte Carlo (see slips.samplers.rdmc)
• Chain of Log-Concave Markov Chains (see slips.samplers.mnm)
• Annealed Importance Sampling (see slips.samplers.smc)
• Sequential Monte Carlo (see slips.samplers.smc)

Installation

The package can be installed as follows

# Clone the projet
git clone https://github.com/h2o64/slips/
cd slips
# Create a virtual environment
python3 -m venv venv
source venv/bin/active
# Install the dependencies
pip install -r requirements.txt
# Install the package
pip install -e .

You can run the experiments using the following command line

python experiments/launch_benchmark.py --results_path [OUTPUT_FOLDER] -target_type [TARGET_TYPE] --algorithm_name [ALGORITHM_NAME] --seed [SEED]

Notebooks

We provide 4 notebooks to play with RDMC, OAT, SMC/AIS and SLIPS on the toy target 8 Gaussians.

Citations

If this codebase is useful towards other research efforts please consider citing us.

@misc{grenioux2024stochastic,
      title={Stochastic Localization via Iterative Posterior Sampling}, 
      author={Louis Grenioux and Maxence Noble and Marylou Gabrié and Alain Oliviero Durmus},
      year={2024},
      eprint={2402.10758},
      archivePrefix={arXiv},
      primaryClass={stat.ML}
}

Contribute

We welcome issues and pull requests (especially bug fixes) and contributions. We will try our best to improve readability and answer questions!