EinSteinVI: General and Integrated Stein Variational Inference

1 INTRODUCTION . Interest in Bayesian deep learning has surged due to the need for quantifying the uncertainty of predictions obtained from machine learning algorithms ( Wilson & Izmailov , 2020 ; Wilson , 2020 ) . The idea behind Bayesian inference is to describe observed data x using a model with latent variables z . The goal is to infer a posterior distribution p ( z|x ) over the latent variables given a model describing the joint distribution p ( z , x ) = p ( x|z ) p ( z ) . We obtain the posterior by following the rules of Bayesian inference : p ( z|x ) = Z−1p ( x|z ) p ( z ) where Z = ∫ z p ( x|z ) p ( z ) dz is the normalization constant . For most practical models , the normalization constant lacks an analytic solution or requires an infeasible number of computations , complicating the Bayesian inference problem . Variational Inference ( VI ) techniques ( Blei et al. , 2017 ; Hoffman et al. , 2013 ; Ranganath et al. , 2014 ) aim to approximate the posterior distribution . VI poses a family of distributions over the latent variables q ( z ) ∈ Q and finds arg min q∈Q D ( q ( z ) |p ( z|x ) ) , where D is a divergence1 between the variational distribution q , also called the guide , and the true posterior distribution p ( z|x ) . The Kullback-Leibler divergence is a typical choice . VI often provides good approximations that capture uncertainty and scales to millions of data points by a suitable choice of Q and inference method . Stein VI is a family of VI techniques for approximate Bayesian inference based on Stein ’ s method ( see Anastasiou et al . ( 2021 ) for an overview ) that is gaining popularity since it combines the scalability of traditional VI with the flexibility of non-parametric particle-based methods . Stein variational gradient descent ( SVGD ) ( Liu & Wang , 2016 ) is a recent Stein VI technique which uses a set of particles { zi } Ni=1 as the approximating distribution q ( z ) . As a particle-based method , SVGD is well suited for capturing correlations between latent variables . The technique preserves the scalability of traditional 1An asymmetric distance that might satisfy the triangle inequality . VI approaches while offering the flexibility and modeling scope of techniques such as Markov chain Monte Carlo ( MCMC ) . SVGD is good at capturing multi-modality ( Liu & Wang , 2016 ; Wang & Liu , 2019a ) , and has useful theoretical interpretations such as a set of particles following a gradient flow ( Liu , 2017 ) or in terms of the properties of kernels ( Liu & Wang , 2018 ) . Many advanced inference methods based on SVGD have been recently developed ; these include nonlinear Stein ( Wang & Liu , 2019a ) , factorized graphical models ( Zhuo et al. , 2018 ; Wang et al. , 2018a ) , matrix-valued kernels ( Wang et al. , 2019a ) and support for higher-order gradient-based optimization ( Detommaso et al. , 2018 ) . These techniques have extended the scope of SVGD , allowing more flexible and accurate approximations of the posterior distribution . While algorithmic power is growing , a distinct lack of integration with a general probabilistic programming language ( PPL ) framework remains . Such integration would solve one of the most prominent limitations of PPLs with traditional VI : their lack of flexibility in capturing rich correlations in the approximated posterior . The main problem with SVGD is that it suffers from the curse of dimensionality : the number of particles required to adequately represent a posterior distribution is exponential in its dimensionality . Nalisnick & Smyth ( 2017 ) suggest resolving this by using Stein mixtures and propose an inference algorithm that uses differentiable non-centered parameterization ( DNCP ) ( Kingma & Welling , 2014 ) and importance weighted Monte Carlo gradients ( Burda et al. , 2015 ) . In this article we propose a novel algorithm for inference of Stein mixtures called ELBO-within-Stein . Unlike prior work on Stein mixtures ( Nalisnick & Smyth , 2017 ) , ELBO-within-Stein i ) is simple to implement , ii ) only has the learning rate as a tuneable hyper-parameters , iii ) requires only one gradient evaluation of the ELBO for each particle ( 10 is typical for Stein mixtures ) and iv ) supports the use of a tailored guide . We validate our algorithm experimentally , showing significant improvements in accuracy . ELBO-within-Stein is the core algorithm of our Stein VI library called EinSteinVI library . Our library further includes SVGD as a special case of ELBO-within-Stein and all the advanced methods mentioned above . The only other PPL with Stein VI methods is PyMC3 ( Salvatier et al. , 2016 ) . However , PyMC3 does not include any of the advanced Stein VI methods mentioned above or Stein mixture methods . The EinSteinVI library extends the NumPyro PPL ( Bingham et al. , 2019 ; Phan et al. , 2019 ) . NumPyro is a universal probabilistic programming language ( van de Meent et al. , 2018 ) for Python , which allows arbitrary code to be executed in both its model and guide . The computational backend of NumPyro is Jax ( Frostig et al. , 2018 ) , which gives access to the powerful program optimization and parallelizability provided by the Jax compiler . As EinSteinVI works with arbitrary guides , NumPyro is a well suited language for embedding EinSteinVI . This is because NumPyro i ) is embedded in Python , the de facto programming language for data science , ii ) includes the necessary data-structures for tracking random variables in both model and guide , iii ) features SVI with an API that is highly suitable for EinSteinVI , and iv ) benefits computationally from Jax . Our extensions include SVGD , Stein-mixtures formulated as ELBO-within-Stein , and the advanced methods mentioned above . Concretely , our contributions are : • ELBO-within-Stein ( EinSteinVI ) . A novel algorithm for Stein mixtures that only requires a single gradient evaluation of the ELBO per particle . • A general library extension to NumPyro , called EinSteinVI . EinSteinVI allows SVGD to work with custom guide programs based on ELBO-within-Stein optimization . The library is compositional with NumPyro features , including support for deep learning , loss functions ( ELBO , Rényi ELBO ( Li & Turner , 2016 ) , custom losses ) , and optimization methods , thus making it possible for EinSteinVI to grow organically with NumPyro development . • Integration of recent developments in Stein variational inference which includes : nonlinear optimization ( Wang & Liu , 2019a ) , a wealth of kernels ( Liu & Wang , 2016 ; 2018 ; Gorham & Mackey , 2017 ) , matrix-valued kernels ( Wang et al. , 2019a ) supporting higherorder optimization , an ( experimental ) update based on Stein point MCMC ( Chen et al. , 2019a ) , and factorization based on conditional independence between elements in the model ( graphical kernels ) ( Wang et al. , 2018b ) . • A series of examples that demonstrate EinSteinVI and the synergy between different Stein VI techniques . The examples include a novel Stein-mixture version of the deep Markov model ( SM-DMM ) , Bayesian neural networks for regression . An examples that explore kernels and higher-order optimization is available in the appendix . The paper proceeds as follows . We first present the background of SVGD in the context of our integrated implementation in Section 2 . In Section 3 we introduce Stein mixtures and our algorithm ELBO-within-Stein . We discuss the general details of the implementation of EinSteinVI in Section 4 . Next , we discuss related work in Section 5 . In Section 6 we present various examples using EinSteinVI , and finally , we summarize our results and discuss future work in Section 7 . 2 STEIN VARIATIONAL GRADIENT DESCENT . The core idea of SVGD is to perform inference by approximating the target posterior distribution p ( z|x ) by an empirical 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 zi ∈ Z . The idea is to minimize the Kullback-Leibler divergence DKL ( qZ ( 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 is the learning rate and SZ denotes the Stein forces . The Two Forces of SVGD Stein VI consists of two forces which work additively under the form SZ = S + Z + S − Z , where the attractive force is given by S+Z ( zi ) = Ezj∼qZ ( z ) [ k ( zi , zj ) ∇zj log p ( zj |x ) ] and the repulsive force by S−Z ( zi ) = Ezj∼qZ ( z ) [ ∇zjk ( zi , zj ) ] . Here k : Rd × Rd → R is a kernel . The attractive force can be seen as pushing the particles towards the modes of the true posterior distribution , smoothed by some kernel . For an example of a kernel , consider the radial basis function ( RBF ) kernel k ( zi , zj ) = exp ( − 1h ‖ zi − zj ‖ 2 2 ) with bandwidth parameter h , chosen as 1lognmed ( z ) . The repulsive force moves particles away from each other , ensuring that they do not collapse to the same mode . For example , with the RBF kernel , the repulsive force becomes Ezj∼qZ ( z ) [ − exp ( − 1h ‖ zi − zj ‖ 2 2 ) 2 h ∑ ` ( zi ` − zj ` ) ] , which has high kernel values for particles that are close , causing them to repel . SVGD works with unnormalized distributions p as the normalization constant becomes additive in the log-posterior log p ( zi|x ) = − logZ + log p ( x|z ) + log p ( z ) and is thus not required for the calculation of the gradient . This property is desirable as normalizing p is often computationally expensive . Non-linear Stein In non-linear Stein ( Wang & Liu , 2019a ) , the repulsive force can be scaled by a factor λ , resulting in SZ = S+Z + λS − Z . This approach is useful when dealing with multimodal distributions . It is also useful in cases where the repulsive force vanishes compared to the likelihood , which happens for large datasets ( X ) . Scaling the repulsive force by a constant λ = c ( |X | ) proportional ( e.g . c = 0.1 or c = 0.01 ) to the size of the dataset |X | addresses this issue and can be chosen by cross-validation on a subset of the data . Matrix-valued kernels The choice of kernels can be extended to matrix-valued ones ( Wang et al. , 2019a ) , K : Rd × Rd → Rd×d , in which case the Stein forces become S+Z ( zi ) = Ezj∼qZ ( z ) [ K ( zi , zj ) ∇zj log p ( zj |x ) ] and S−Z ( zi ) = Ezj∼qZ ( z ) [ K ( zi , zj ) ∇zj ] where the standalone del ∇zj in the repulsive force represents the vector ( ∂ ∂zj,1 , . . . , ∂∂zj , d ) . This results in ( K ( zi , zj ) ∇zj ) ` = ∑ k ∇kK ` , k ( zi , zj ) . The advantage of matrix-valued kernels is that they allow preconditioning2 using the Hessian or Fisher Information matrix , which can capture local curvature and thus achieve better optima and convergence rate than standard SVGD . Furthermore , it is easy to represent graphical kernels ( Wang et al. , 2018a ) using matrix kernels , e.g . K = diag ( { K ( ` ) } ` ) where the set of variables are partitioned with each their own local kernel K ( ` ) .

Summary:    This paper describes an implementation of EinsteinVI, an improvement over Stein VI, a method that can be used to apply VI in prob. models assuming a flexible distribution represented in terms of particles. The paper describes how this method has been implemented in NumPyro, a software for prob. programming. The method is evaluated and compare to other alternative prob. languages such as PyMC3 that uses also methods based on Stein VI.