Einstein VI: General and Integrated Stein Variational Inference in NumPyro

Ahmad Salim Al-Sibahi, Ola Rønning, Christophe Ley, Thomas Wim 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 inferencethat is recently gaining popularity since it combines the scalability of traditionalVariational Inference (VI) with the flexibility of non-parametric particle basedinference methods. While there has been considerable progress in developmentof algorithms, integration in existing probabilistic programming languages (PPLs)with an easy-to-use interface is currently lacking. EinStein VI is a lightweightcomposable library that integrates the latest Stein Variational Inference methodswith the NumPyro PPL. Inference with EinStein VI relies on ELBO-within-Stein tosupport use of custom inference programs (guides), non-linear scaling of repulsionforce, second-order gradient updates using matrix-valued kernels and parametertransforms. We demonstrate the achieved synergy of the different Stein techniquesand the versatility of EinStein VI library by applying it on examples. Comparedto traditional Stochastic VI, EinStein VI is better at capturing uncertainty andrepresenting richer posteriors. We use several applications to show how one canuse Neural Transforms (NeuTra) and second-order optimization to provide betterinference using EinStein VI. We show how EinStein VI can be used to infer theparameters of a Latent Dirichlet Allocation model with a neural guide. The resultsindicate that Einstein VI can be combined with NumPyro’s support for automaticmarginalization to do inference over models with discrete latent variables. Finally,we introduce an example with a novel extension to Deep Markov Models, calledthe Stein Mixture Deep Markov Model (SM-DMM), which shows that EinStein VIcan be scaled to reasonably large models with over 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, adds optimizable parameter transforms and supports ELBO optimization.

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics

Reviewed Version (pdf): https://openreview.net/references/pdf?id=hBhu-Wd3u