Einstein VI: General and Integrated Stein Variational Inference in NumPyro

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

Tasks

Reproductions