title: EinSteinVI: General and Integrated Stein Variational Inference
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.
authors: Ola Rønning Ahmad Salim Al-Sibahi Christophe Ley Thomas Hamelryck
hIndex: None 0