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: The aim of this work is to integrate the Stein variational inference methods in the existing probabilistic programming language NumPyro. The implemented methods include variants of Stein variational gradient descent algorithms with a range of kernel functions, non-linear scaling of update terms, and matrix-valued kernels. Empirical results with existing baselines for real-world problems are provided.