Stein Variational Gradient Descent (SVGD) - Tutorial

This repository provides a collection of tutorials demonstrating how to implement SVGD using Python and PyTorch. It includes toy examples of vanilla SVGD and annealed SVGD applied to problems ranging from low-dimensional to high-dimensional and multimodal distributions.

Corrales M.

📊 Project structure

This repository is organized as follows:

• 📂 svgdtutorial: python library containing routines for the totorials provided (for now empty).
• 📂 data: folder containing data (or instructions on how to retrieve the data).
• 📂 notebooks: set of jupyter notebooks reproducing the tutorials (see below for more details).
• 📂 scripts: set of python scripts used to run multiple experiments (for now empty).

🔥 Notebooks

The following notebooks are provided:

• 📙 00_Theory_Recap.ipynb: Brief recap about Deterministic Inverse problems and Variational inference.
• 📙 00_SVGD_Notes.ipynb: Brief derivation of SVGD.
• 📙 01_SVGD_gaussian.ipynb: SVGD tutorial for one simple gaussian distribution.
• 📙 02_ASVGD_multimodal.ipynb: Annealed SVGD tutorial for multi-modal distributions.
• 📙 03_SVGD_high_dimensions.ipynb: SVGD tutorial for high-dimensional distributions.

The theory notebooks have also been shared as slides in HTML format. You can find them in the following folder 📂 notebooks/html/

🎮 Getting started

To ensure reproducibility of the results, we suggest using the environment.yml file when creating an environment. (Use environment_cpu.yml or environment_gpu.yml depending on your needs)

For example, simply run:

./install_env_cpu.sh

It will take some time, if at the end you see the word Done! on your terminal you are ready to go. Remember to always activate the environment by typing:

conda activate svgdtutorial_cpu

After that you can simply install your package:

pip install .

or in developer mode:

pip install -e .

Note

All experiments have been carried on a Intel(R) Xeon(R) CPU @ 2.10GHz equipped with a single NVIDIA GEForce RTX 3090 GPU, and an Apple M3 Max. Different environment configurations may be required for different combinations of workstation and GPU.

Important

This implementation was inspired on:

• https://github.com/DartMLStein-Variational-Gradient-Descent
• https://github.com/chenxy99/Stein-Variational-Gradient-Descent
• https://github.com/activatedgeek/svgd
• https://github.com/ChunyuanLI/SVGD
• https://github.com/shwangtangjun/SVGD-PyTorch

📖 References

$^1$ Liu, Q., & Wang, D. (2016). Stein variational gradient descent: A general purpose bayesian inference algorithm. Advances in neural information processing systems, 29.

$^2$ D'Angelo, F., & Fortuin, V. (2021). Annealed stein variational gradient descent. arXiv preprint arXiv:2101.09815.

$^3$ Ba, J., Erdogdu, M. A., Ghassemi, M., Sun, S., Suzuki, T., Wu, D., & Zhang, T. (2021). Understanding the variance collapse of SVGD in high dimensions. In International Conference on Learning Representations.

$^4$ Duncan, A., Nüsken, N., & Szpruch, L. (2023). On the geometry of Stein variational gradient descent. Journal of Machine Learning Research, 24(56), 1-39.

$^5$ El Moselhy, T. A., & Marzouk, Y. M. (2012). Bayesian inference with optimal maps. Journal of Computational Physics, 231(23), 7815-7850.

$^6$ Liu, Q., & Wang, D. (2016). Stein variational gradient descent: A general purpose bayesian inference algorithm. Advances in neural information processing systems, 29.

$^7$ Oliver, D. S., Reynolds, A. C., & Liu, N. (2008). Inverse theory for petroleum reservoir characterization and history matching.

$^8$ Ravasi, M., Romero, J., Corrales, M., Luiken, N., & Birnie, C. (2023). Striking a balance: Seismic inversion with model-and data-driven priors. In Developments in structural geology and tectonics (Vol. 6, pp. 153-200). Elsevier.

$^9$ Jordan, M. I., Ghahramani, Z., Jaakkola, T. S., & Saul, L. K. (1999). An introduction to variational methods for graphical models. Machine learning, 37, 183-233.

$^{10}$ Blei, D. M., Kucukelbir, A., & McAuliffe, J. D. (2017). Variational inference: A review for statisticians. Journal of the American statistical Association, 112(518), 859-877.