Einstein VI:   General and Integrated Stein Variational Inference in NumPyro

1 INTRODUCTION . Interest in Bayesian deep learning has surged due to the need for quantifying the uncertainty of predictions provided by machine learning algorithms . The idea behind Bayesian learning is to describe observed data x using a model with latent variable z ( representing model parameters and nuisance variables , see e.g. , Fig . 4a ) . The goal is then to infer a posterior distribution p ( z|x ) over latent variables given a model describing the joint distribution p ( z , x ) = p ( x|z ) p ( z ) following the rules of Bayesian inference : p ( z|x ) = Z−1p ( x|z ) p ( z ) where the normalization constant Z = ∫ z p ( x|z ) p ( z ) dz is intractable for most practical models including deep neural networks : an analytic solution is lacking or may require an infeasible number of calculations . Variational Inference ( VI ) techniques ( Blei et al. , 2017 ; Hoffman et al. , 2013 ; Ranganath et al. , 2014 ) provide a way to find an approximation of the posterior distribution . VI poses a family of distributions over latent variables q ( z ) ∈ Q ( e.g. , Fig . 4b ) and chooses the one that minimizes a chosen divergence1 D ( q ( z ) ‖ p ( z|x ) ) ( e.g. , Kullback-Leibler ) to the true posterior distribution . VI often provides good approximations that can capture uncertainty , scaling to millions of data points using mini-batch training . 1Asymmetric distance Stein Variational Inference ( Liu and Wang , 2016 ) is a recent non-parametric approach to VI which uses a set of particles { zi } Ni=1 as the approximating distribution q ( z ) to provide better flexibility in capturing correlations between latent variables . The technique preserves the scalability of traditional VI approaches while offering the flexibility and modelling scope of techniques such as Markov Chain Monte Carlo ( MCMC ) . Stein VI has been shown to be good at capturing multi-modality ( Liu and Wang , 2016 ; Wang and Liu , 2019a ) , and has useful theoretical interpretations as particles following a gradient flow ( Liu , 2017 ) and as a moment matching optimization system ( Liu and Wang , 2018 ) . Many advanced inference methods based on Stein VI have been recently developed , including Stein mixtures ( Nalisnick , 2019 ) , non-linear Stein ( Wang and Liu , 2019b ) , factorized graphical models ( Zhuo et al. , 2018 ; Wang et al. , 2018b ) , matrix-valued kernels ( Wang et al. , 2019 ) and support for higher-order gradient-based optimization ( Detommaso et al. , 2018 ) . These techniques have been shown to significantly extend the power of Stein VI , allowing more flexible and effective approximations of the true posterior distribution . While algorithmic power is growing , there remains a distinct lack of integration of these techniques into a general probabilistic programming language ( PPL ) framework . Such an integration would solve one of the most prominent limitations of traditional VI , which lacks flexibility to capture rich correlations in the approximated posterior . This paper presents the EinStein VI library that extends the NumPyro PPL ( Bingham et al. , 2019 ; Phan et al. , 2019 ) with support for the recent developments in Stein Variational Inference in an integrated and compositional fashion ( see Fig . 1c and Fig . 1d ) . The library takes advantage of the capabilities of NumPyro— including universal probabilistic programming ( van de Meent et al. , 2018 ) , integration with deep learning using JAX ( Frostig et al. , 2018 ) , and automatic optimization and marginalization of discrete latent variables ( Obermeyer et al. , 2019 ) —to provide capabilities that work synergetically with the Stein algorithms . Concretely , our contributions are : • Einstein VI , a general library that extends NumPyro and allows Stein Variational Inference to work with custom guide programs based on ELBO-within-Stein optimization ( Nalisnick , 2019 ) . The library is compositional with NumPyro features , including supported deep learning and automatic marginalization , loss functions ( ELBO , Rényi ELBO , custom losses ) and optimization methods , allowing it to grow organically with NumPyro development . • Integration of recent developments in Stein variational inference within the EinStein library . This includes support for non-linear optimization ( Wang and Liu , 2019b ) , a wealth of kernels ( Liu and Wang , 2016 ; 2018 ; Gorham and Mackey , 2017 ) , matrix-valued kernels ( Wang et al. , 2019 ) supporting higher-order optimization , and factorization based on conditional independence between elements in the model ( graphical kernels ) . • We support the application of transforms on the parameter space , including triangular ( Parno and Marzouk , 2018 ) and neural transforms ( Hoffman et al. , 2018 ) to improve the gradient geometry of the inference problem . • A series of examples demonstrate the power of an integrated library such as EinStein VI and the synergy between different Stein VI techniques . The examples include a novel Stein Mixture version of Deep Markov Models ( SM-DMM ) , Stein Mixture Latent Dirichlet Allocation ( SM-LDA ) , and several examples using neural transforms and higher-order optimization . The paper proceeds as follows . We first present a primer on the theory of Stein VI in Section 2 relating it to our integrated implementation in EinStein VI . We discuss the general details of the implementation of EinStein VI in NumPyro in Section 3 . We present the various examples using EinStein VI in Section 4 and finally summarize our results and future work in Section 5 . 1.1 RELATED WORK . There has been a proliferation of deep probabilistic programming languages based on tensor frameworks with automatic differentiation , supporting various inference techniques . Pyro ( Bingham et al. , 2019 ) is a universal PPL based on PyTorch ( Paszke et al. , 2019 ) . The main mode of inference in Pyro is black-box Stochastic Variational Inference ( Ranganath et al. , 2014 ) with guides , which are flexible programs that approximate the posterior distribution by repeating the ( non-observed ) sample statements as the probabilistic model and can contain deep neural networks allowing amortization of inference ( Kingma and Welling , 2014 ; Gershman and Goodman , 2014 ) . Pyro also supports various sampling algorithms like Hamiltonian Monte Carlo/NUTS ( Neal , 2011 ; Hoffman and Gelman , 2014 ) and sample-adaptive MCMC ( Zhu , 2019 ) which provide more accurate posterior approximations but lack the scalability of VI techniques . NumPyro is a version of Pyro that runs on the JAX framework ( Frostig et al. , 2018 ) , which allows it to exploit the powerful program optimization and parallelizability of the JAX compiler . Our library , EinStein VI , extends the Stochastic VI mode of inference by adding Stein Variational Inference on top of NumPyro , allowing the optimizable parameters in the guide to be approximated by a set of particles instead of a point estimate . Other languages with similar feature set include PyMC3 ( Salvatier et al. , 2016 ) , Edward ( Tran et al. , 2016 ; 2019 ) and HackPPL ( Ai et al. , 2019 ) . 2 PRIMER ON STEIN VI . The core idea of Stein VI ( Liu and Wang , 2016 ) is to perform inference by approximating the target posterior distribution p ( z|x ) by an approximate distribution qZ ( z ) = N−1 ∑ i δzi ( z ) based on a set of particles Z = { zi } Ni=1 . Here , δx ( y ) represents the Dirac delta measure , which is equal to 1 if x = y and 0 otherwise . One could thus see the approximating distribution qZ ( z ) as a mixture of point estimates , each represented by a particle z ∈ Z . The idea is then that we minimize the Kullback-Leibler divergence DKL ( q ( z ) ‖ p ( z|x ) ) between the approximation and the true posterior by iteratively updating the particles using the Stein forces : zi ← zi + SZ ( zi ) where represents the learning rate and SZ the Stein forces . The Two Forces of Stein VI Stein VI consists of two forces which work additively under the form SZ = S + Z + S − Z , where the attractive force S+Z ( zi ) = Ezj∼qZ ( z ) [ k ( zj , zi ) ∇zi log p ( zi|x ) ] and the repulsive force S−Z ( zi ) = Ezj∼qZ ( z ) [ ∇zik ( zj , zi ) ] . Here k : Rd × Rd → R is a kernel . The attractive force can be seen as pushing the particles towards the direction that maximizes the true posterior distribution , smoothed by some kernel . For an example of a kernel , consider the standard RBF kernel k ( zi , zj ) = exp ( − 1h ‖ zi − zj ‖ 2 2 ) with bandwidth parameter h , usually chosen as 1lognmed ( z ) . The normalization constant becomes additive in the log-posterior log p ( zi|x ) = − logZ + log p ( x|z ) + log p ( z ) and so does not need to be computed for the gradient . The repulsive force can be seen as moving particles away from each other , ensuring that they do not all collapse to the same mode . For the RBF kernel , the repulsive force becomes Ezj∼qZ ( z ) [ k ( zj , zi ) 2 h ∑ ` ( zi ` − zj ` ) ] and so particles that are close together ( thus having a high kernel value ) will be pushed away from each other . Non-linear Stein In non-linear Stein ( Wang and Liu , 2019a ) , the repulsive force can be scaled by a factor λ , so SZ = S+Z + λS − Z which is often useful when dealing with multi-modal distributions . It is also useful in our framework since the repulsive force often vanishes compared to the likelihood for large datasets X and therefore scaling the repulsive force by a constant λ = λ ( |X | ) proportional to the size of the dataset |X | appears logical , e.g . 0.1 or 0.01 . Matrix-valued kernels The choice of kernels can be extended to matrix-valued ones ( Wang et al. , 2019 ) K : Rd × Rd → Rd×d in which case the Stein forces become S+Z ( zi ) = Ezj∼qZ ( z ) [ K ( zj , zi ) ∇zi log p ( zi|x ) ] and S−Z ( zi ) = Ezj∼qZ ( z ) [ K ( zj , zi ) ∇zi ] where the standalone del∇zi in the repulsive force represents the vector ( ∂ ∂zi,1 , . . . , ∂∂zi , d ) and so ( K ( zj , zi ) ∇zi ) ` = ∑ k∇kK ` , k ( zj , zi ) . The advantage of matrix-valued kernels is that they allow preconditioning using the Hessian or Fisher Information matrix ( second-order derivatives ) , which can capture local curvature and thus achieve better optima than standard Stein VI . Furthermore , it is easy to capture graphical kernels ( Wang et al. , 2018b ) K = diag ( { K ( ` ) } ` ) where the set of variables are partitioned with each their own local kernel K ( ` ) . ELBO-within-Stein The standard form of Stein VI guides the particles toward maximizing the log-posterior log p ( z|x ) but in principle we could replace this with another objective ( negative loss ) Lθ which we want to maximize . Assume our guide qθ ( z ) comes from a parametric distribution with parameters θ , as in traditional VI . We can then for example maximize the ELBO ( Kingma and Welling , 2014 ) Lθ = Ez∼qθ ( z ) [ log p ( z , x ) − log qθ ( z ) ] as in traditional VI but using a set of particles for the parameters { θi } i guided by the Stein forces . One can see this as a generalization of traditional VI using a Stein-based mixture-model ( Nalisnick , 2019 ) , giving the parameters the flexibility to better capture correlations between distributions and to cover multiple modes ( Wang and Liu , 2019a ) . Another perspective is that ELBO-within-Stein adds uncertainty to the approximating distribution qθ ( z ) , allowing it to work with distributions over latent variables that are richer than simply point masses . The objective to be maximized in our framework is compositional , which means that ELBO can be replaced by other objectives , e.g. , Rényi ELBO ( Li and Turner , 2016 ) , Tail-adaptive f-divergence ( Wang et al. , 2018a ) or Wasserstein pseudo-divergence ( Ambrogioni et al. , 2018 ) when these are implemented in NumPyro . Einstein VI makes it possible to integrate all these developments within Stein VI in a coherent fashion . Parameter Transforms Deterministic transport maps have been suggested to be integrated in MCMC methods by Parno and Marzouk ( 2018 ) and Hoffman et al . ( 2018 ) , to allow constructing a better geometry for the posterior approximation and thus improve inference . We provide a novel adaptation of the idea to work with Stein-based methods , in the EinStein VI framework . Consider the general objective we optimize in EinStein VI , which has the form : SΘ ( θi ) = Eθj∼q ( θ ) [ k ( θj , θi ) ∇θiLθi +∇θik ( θj , θi ) ] where Lθ is our negative loss objective ( e.g. , ELBO ) parameterised by θ which we update using a set of Stein particles Θ = { θi } i . The core idea is then that we would like to reparametrize the calculation so each Stein particle θi is expressed as an ( invertible ) deterministic transformation of another particle φi , so θi = T ( φi ) ( thus Θ = { T ( φ ) |φ ∈ Φ } ) . We do this reparametrization by first declaring the set of particles to be optimized Φ = { φi } i and the transformation to be used T , and then define the reparametrized Stein force as follows : SΦ ( φi ) = SΘ ( T ( φi ) ) ( ∇φiT ( φi ) ) > where ∇φiT ( φi ) : Rd×e is the reparametrization Jacobian , that adjusts for the change of variables ( from θi : Re to φi : Rd ) . The power of parameter transforms is amplified in our EinStein VI framework , in that we allow for its parameters to be learnable . We can thus use e.g. , sparse triangular maps , invertible neural networks or normalizing flows , which have been shown to be particularly effective for Bayesian inference .

Summary: This paper introduces EinStein VI: a lightweight composable library for Stein Variational Inference (Stein VI). The library is built on top of NumPyro and can take advantage of many of NumPyro's capabilities. It supports recent techniques associated with Stein VI, as well as novel features. The paper provides examples of using the EinStein VI library on different probabilistic models.