EinSteinVI: General and Integrated Stein Variational Inference

Ola Rønning, Ahmad Salim Al-Sibahi, Christophe Ley, Thomas Hamelryck

Keywords: Stein variational inference, variational inference, probabilistic programming, Pyro, deep probabilistic programming, deep learning

Abstract: Stein variational inference is a technique for approximate Bayesian inference that has recently gained popularity because it combines the scalability of variational inference (VI) with the flexibility of non-parametric inference methods. While there has been considerable progress in developing algorithms for Stein variational inference, integration in existing probabilistic programming languages (PPLs) with an easy-to-use interface is currently lacking. EinSteinVI is a lightweight compostable library that integrates the latest Stein variational inference method with the PPL NumPyro (Phan et al., 2019). EinSteinVI provides ELBO-within-Stein to support the use of custom inference programs (guides), implementations of a wide range of kernels, non-linear scaling of the repulsion force (Wang & Liu,2019b), and second-order gradient updates using matrix-valued kernels (Wang et al.,2019b). We illustrate EinSteinVI using toy examples and show results on par with or better than existing state-of-the-art methods for real-world problems. These include Bayesian neural networks for regression and a Stein-mixture deep Markov model, which shows EinSteinVI scales to large models with more than 500,000 parameters.

One-sentence Summary: We present EinStein Variational Inference, a technique for inference that integrates all the latest developments within Stein VI into NumPyro and supports ELBO optimization.