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: The paper shows how a particle-based nonparameteric Variational Inference methodology known as Stein Variational Inference is integrated in a full-featured Probabilistic Programming Language, NumPyro. The paper goes into a fair amount detail describing a number of enhancements that have been made into numpyro using the general technique of particle-based representation of non-parameteric approximating distributions. They describe how geometric transforms of the parameter space can fit into their scheme, how matrix-valued kernels can be integrated. Also, they describe a new variant of Stein VI which they call ELBO-within-Stein. This introduces a new line of research for Stein VI. They also describe a Stein Mixture extension to Deep Markov Models (SM-DMM) and demonstrate on a very large dataset for the latter method.