Year: 2021

Abstract:
Stein variational inference is a technique for approximate Bayesian inference that has recently gained popularity because it combines the scalability of variational inference (VI) with the flexibility of non-parametric inference methods. While there has been considerable progress in developing algorithms for Stein VI, integration in existing probabilistic programming languages (PPLs) with an easy-to-use interface is currently lacking. EinSteinVI is a lightweight composable library that integrates the latest Stein VI methods with the PPL NumPyro (Phan et al., 2019). EinSteinVI also provides our novel algorithm ELBO-within-Stein to support the use of custom inference programs (guides), in addition to implementations of a wide range of kernels, non-linear scaling of the repulsion force (Wang & Liu, 2019b) and secondorder gradient updates using matrix-valued kernels (Wang et al., 2019b). We illustrate EinSteinVI using toy examples and show results on par with or better than existing state-of-the-art methods for real-world problems. These include Bayesian neural networks for regression and a Stein-mixture deep Markov model, which also shows EinSteinVI scales to large models with more than 125,000 parameters.

heading: 1 INTRODUCTION
text: 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.

heading: 2 STEIN VARIATIONAL GRADIENT DESCENT
text: 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(`).

heading: 3 STEIN MIXTURES
text: The Stein-mixture was proposed by Nalisnick & Smyth (2017) to resolve the representation issue for SVGD when the target distribution has high dimensionality. For SVGD, the number of particles needed to represent a distribution adequately grows exponentially with its dimensionality. As the computational complexity for the update rule in SVGD is quadratic in the number of particles, the exponential growth quickly becomes computationally intractable.
Stein-mixtures are hierarchical variational models (Ranganath et al., 2016b) such that the posterior of the second tier variational distribution is q(z|X) = 1N ∑N k=1 δ(zk). The guide is the joint model q(θ, z) = q(θ|z)q(z), with q(θ|z) the first tier distribution and q(z) its prior. The second tier posterior q(z|X) is optimized using SVGD, so that the marginal posterior is q(θ; z) = 1N ∑N k=1 q(θ|z), a restricted mixture.
Nalisnick & Smyth (2017) showed that the Stein force for a mixture approximation of the posterior p(z|Z) is given by
SZ(zi) = Ezj∼qZ(z) [ k(zi, zj)Eθl∼q(θ|zj) [p(X,θl)/q(θl|zj)] ] +
+ Ezj∼qZ(z) [ k(zi, zj)∇zj log q(zj) ] + + S−Z (zi),
where log∇zjq(zj) acts as a regularizer on the variational parameters which they assume is sufficiently small to be dropped. To evaluate Ezj∼qZ(z) [p(X, θ)/q(θ|zj)], Nalisnick & Smyth (2017) use DNCP and importance weighted Monte Carlo gradients to obtain the black-box update,
SZ(zi) = Ezj∼qZ(z) [ k(zi, zj) S∑ s=1 w̃s∇zj log ( p(X, θ̂s) q(θ̂s|zj) )] + S−Z (zi), (1)
where w̃s is an importance weight for sample θ̂s = q(zj , ξ), ξ ∼ p0 (i.e. the guide).
ELBO-within-Stein In ELBO-within-Stein, we replace the weighted average over S samples in Equation (1) by a single loss (L). If we choose L to be an f -divergence, such as the ELBO, we can reduce the loss bias by averaging multiple samples. However, ELBO-within-Stein only computes one gradient of L per particle regardless of the number of samples we use to estimate L. This is different from Nalisnick & Smyth (2017) who estimate Eθl∼q(θ|zj) [p(X,θl)/q(θl|zj)] by averaging over gradients, therefore computes a gradient and an importance weighting for each DNCP sample rather than per particle. The difference makes ELBO-within-Stein computationally cheaper as the complexity of adding a gradient is O(n) whereas adding a sample to estimate L is O(1). The Stein force in ELBO-within-Stein is given by
SZ(zi) = Ezj∼qZ(z) [ k(zi, zj)∇zjL(zj) ] + S−Z (zi) (2)
where L is the ELBO.
2Using preconditioner matrix Q, such that Q−1M has a lower condition number than M .

heading: 4 COMPOSITIONAL IMPLEMENTATION USING NUMPYRO
text: EinSteinVI integrates with the existing NumPyro API by adding the Stein VI interface, which closely mimicks NumPyro’s SVI interface. Mimicking the SVI interface makes programs that use SVI in NumPyro trivial to convert to EinSteinVI (see Figure 1).
Below, we discuss the key features of EinSteinVI, which include re-initializable guides, EinSteinVI’s core algorithm, and the new kernel interface.

heading: 4.1 RE-INITIALIZABLE GUIDES
text: The Stein VI interface requires that the initialization is different for each parameter in an inference program. The reason is that different Stein particles need to be initialized to different values in order for optimization to work correctly and to avoid that all particles collapse into the posterior mode.
To support re-initializable guides, we provide the ReinitGuide interface, which requires implementing a function find_params that accepts a list of random number generator (RNG) keys in addition to the arguments for the guide and returns a set of freshly initialized parameters for each RNG key.
The WrappedGuide class provides a guide written as a function. WrappedGuide makes a callable guide re-initializable. It works by running the provided guide multiple times and reinitializing the parameters using NumPyro’s interface as follows:
• WrappedGuide runs the guide transforming each parameter to unconstrained space.
• It replaces the values of the parameters with values provided by a NumPyro initialization strategy, e.g., init_to_uniform(r), which initializes each parameter with a uniform random value in the range [−r; r].
• It saves the parameter values for each particle and the required inverse transformations to constrained space to run the model correctly.
We also allow parameters without reinitialization in order to support neural network libraries like stax 3 that have their own initializers.
The Stein VI interface will correctly wrap the guide during initialization, so from a user perspective the syntax for guides follows the API of SVI.

heading: 4.2 STEIN VI IN NUMPYRO
text: The integration of Stein VI into NumPyro requires handling transformations between the parameter representation of NumPyro4 and the vectorized Stein particles that EinSteinVI operates on. For this, we rely on Jax PyTrees5 which converts back and forth between Python collections and a flattened vectorized representation.
Algorithm 1 shows the core algorithm of EinSteinVI. EinSteinVI updates the standard variational model parameters φ and guide parameters ψ by averaging the loss over the Stein particles. For the Stein parameters, the process is more elaborate. First, we convert the set of individual parameters to a monolithic vector-encoded particle using Jax PyTrees. The monolithic particle represents the particles as a flattened and stacked Jax array. Then we compute a kernel based on the vector-encoded Stein particle; this is delegated to the kernel interface as the computation is kernel-dependent.
We apply Jax’s vmap operator (Frostig et al., 2018; Phan et al., 2019) to compute the Stein forces for each particle in a vectorized manner. This is done in unconstrained space so the Stein force must the corrected by the Jacobian of the bijection between constrained and unconstrained space. Doing this directly on the Jax on the monolithic particle incurs a massive memory overhead in the adjoint. However, as NumPyro registers a bijection for each distribution parameter we can eliminate the overhead by computing the Jacobian on the Jax representations of the individual parameters. The operation is embarrassingly parallel and so we again use a vmap operator with a nested tree_map to
3https://jax.readthedocs.io/en/latest/jax.experimental.stax.html 4A dictionary mapping parameters to their values, which can be arbitrary Python type 5https://jax.readthedocs.io/en/latest/pytrees.html
compute the desired Jacobians. Note that Algorithm 1 presents how EinSteinVI works with scalar kernels and does not account for the different features presented in Section 2.
Finally, we convert the monolithic Stein particle to their non-vectorized dictionary-based form and return the expected changes for standard- and Stein-parameters.
Input: Classical VI parameters φ and ψ, Stein parameters {θi}i, model pφ(z,x), guide qθ,ψ(z), loss L, kernel interface KI. Output: Parameter changes based on classical VI (∆φ, ∆ψ) and Stein VI forces ({∆θi}i). procedure EINSTEIN(φ, ψ, {θi}i, pφ, qθ,ψ)
∆φ← Eθ[∇φL(pφ, qθ,ψ)] ∆ψ ← Eθ[∇ψL(pφ, qθ,ψ)] {ai}i ← PyTreeFlatten({θi}i) k ← KI({ai}i)
procedure EINSTEINFORCES(ai) . Calculate forces per particle for higher-order vmap function.
θi ← PYTREERESTORE(ai) ∆ai ← ∑ aj k(aj ,ai)∇aiL(pφ, qθi,ψ) +∇aik(aj ,ai)
return ∆ai end procedure
{∆ai}i ← VMap({ai}i, EINSTEINFORCES) {∆θi}i ← PYTREERESTORE({∆ai}i) return ∆φ, ∆ψ, {∆θi}i
end procedure Algorithm 1: EinSteinVI

heading: 4.3 KERNEL INTERFACE
text: The Kernel interface is straightforward. To extend the interface users must implement the compute function, which accepts as input the current set of particles, the mapping between model parameters and particles, and the loss function L and returns a differentiable kernel k. All kernels are currently static, but the interface could be extended to stateful kernels, allowing conjugate gradients or quasiNewton optimization. Appendix A gives the complete list of kernels in EinSteinVI. We are planning to extend EinSteinVI with probability product kernels Nalisnick & Smyth (2017), which take into account the information geometry of the first tier variational distributions in Stein mixtures.

heading: 5 RELATED WORK
text: There has been a proliferation of deep probabilistic programming languages (Ge et al., 2018; Bingham et al., 2019; Salvatier et al., 2016; Tran et al., 2016; Cusumano-Towner et al., 2019; Dillon et al., 2017) based on tensor frameworks featuring automatic differentiation and supporting various inference techniques. However, only PyMC3 (Salvatier et al., 2016) includes inference using SteinVI.
In PyMC3 the inference techniques that use Stein’s method include SVGD with a scalar RBF-kernel, amortized SVGD (Wang et al., 2016; Feng et al., 2017), and Operator Variational Inference (OVI) (Ranganath et al., 2016a).
SVGD in PyMC3 manipulates particles to directly capture the target distribution (Liu & Wang, 2016). In EinSteinVI, SVGD is a special case of Stein mixtures where the guides are point mass distributions. Because we have guide programs, EinSteinVI allows arbitrary computations to transform random variables in the guide. This is not possible with SVGD in PyMC3.
Amortized SVGD (Feng et al., 2017) trains a stochastic network to draw samples from a target distribution. The network is iteratively adjusted so that the output changes in the direction of the Stein variational gradient (the same gradient as used in SVGD). In comparison, EinSteinVI transports a fixed set of particles (each parameterizing a guide) to the target distribution. Amortized SVGD is not in EinSteinVI as its extension to arbitrary guides is an open problem.
OVI optimizes operator objectives, which take functions of functions to a non-negative number. Ranganath et al. (2016a) include an operator objective based on the Langevin-Stein operator (Anastasiou et al., 2021). This is the same operator used for the kernelized Stein discrepancy (Liu et al., 2016) which also underlies SVGD. Unlike EinSteinVI, Amortized SVDG and SVGD, OPV is not a particle-based method.

heading: 6 EXAMPLES
text: We illustrate the features of EinSteinVI on two toy examples, namely a 1D mixture of Gaussians below, and 2D mixtures of Gaussian in Appendix B. We demonstrate EinSteinVI on real-world examples and show that EinSteinVI tends to outperform alternative methods. These examples include regression with Bayesian neural networks and deep Markov models. All experimental code, notebook tutorials and the EinSteinVI package itself are available through an Anonymized URL.
1D Gaussian mixture To demonstrate the two modes of VI with EinSteinVI, we consider the 1D Gaussian mixture 1/3N (−2, 1) + 2/3N (2, 1) (see Figure 1 and Figure 2). The Gaussian target mixture is bi-modal and well-suited for the nonparametric nature of SVGD and Stein mixtures. Figure 2 shows that both SVGD and the Stein-mixture naturally capture the bi-modality of the target distribution, compared to SVI with a simple Gaussian guide. Note the reduction in particles required to estimate the target when using Stein mixtures compared to SVGD. Also, note that the Stein-mixture overestimates the variance and slightly pertubates the locations. The error seen at the right mode for the Stein-mixture with two particles is due to the uniform weighting of the particles in SVGD (the target posterior is approximated with an empirical distribution, see Section 2), and as such is algorithmic. The Stein-mixture will therefore not be able to exactly capture the mixing components of a target mixture model with one particle per component. However, with more particles the mixture can be approximated better as demonstrated using three particles.
Bayesian Neural Networks We compare SVGD and non-linear Stein in EinSteinVI with the implementation of SVGD by Liu & Wang (2016)6 (without model selection), and amortized SVGD (using the Theano backend7, Al-Rfou et al. (2016)) on Bayesian neural networks (BNN) for regression. The PyMC3 documentation clearly states amortized SVGD is experimental and is not suggested to be used. We include its experimental results in Appendix B for completeness and note that our findings reflect the warning in the PyMC3 documentation. We use an RBF-kernel in EinSteinVI for fair comparison with Liu & Wang (2016) and PyMC3. For non-linear Stein we determined the best repulsion factor λ∗ by a grid search for λ ∈ {10−2, 10−1, 101, 102}. For SVGD PyMC3 we use a
6https://github.com/dilinwang820/Stein-Variational-Gradient-Descent/blob/master/ python/bayesian_nn.pyy
7PyMC4, which uses Jax as its backend, does not include SVGD at the time of writing.
temperature8 of ten. Like Liu & Wang (2016) we use a BNN with one hidden layer of size fifty and RELU activation. We put a Gamma(1, 0.1) prior on the precision of the neurons and the likelihood. For all versions, we use 100 particles and 2000 iterations. We use a mini-batch of 1000 for Year and 100 for the rest. All measurements are repeated ten times and obtained on a GPU9 for EinSteinVI and PyMC3, except for Year with PyMC3 which was only run once. We do not report times for Liu & Wang (2016) because only the CPU version of their code could be executed without irresolvable issues.
Table 1 shows the performance in terms of the average root mean squared error (RMSE) on the test set. We find that EinSteinVI achieves significantly better RMSE than Liu & Wang (2016) and that both systems outperform SVGD in PyMC3. The times in Table 1 measure how long it takes to infer parameters. Table 1 excludes the first run for two reasons: i) that run will fill the GPU caches, and ii) Jax will trace the programs. As a result, this run is more costly than subsequent ones. The times for the first run are given in Table 2. By running EinSteinVI for more iterations, we can amortize the initial cost.
Stein Mixture Deep Markov Model Music generation requires a model to learn complex temporal dependencies to achieve local consistency between notes. The Stein-mixture deep Markov model (SM-DMM) is a deep Markov model that uses a mixture of Stein particles to estimate distributions over model parameters. We consider a vectorized version of the DMM (Jankowiak & Karaletsos, 2019) for the generation of polyphonic music using the JSB chorales data set.
The SM-DMM model consists of two feed-forward neural (FNN) networks. The Transition network takes care of the conditional dependencies between subsequent latent states in the Markov chain. It consist of two layers with hidden dimension 200 and ReLU activation on the first layer and sigmoid
8We found the performance to be very sensitive to the choice of temperature. 9Quadro RTX 6000 with Cuda V11.4.120
on the second. The Emitter network is a three layers FNN which produces a likelihood at each time step using the current latent state. The layers have hidden dimensions 100, 100, and 88, respectively. The variational distribution is of the form ∏N n=1 q(z n 1:Tn |f(X1:Tn)) where the parametrized feature function f1:Tn is a one layered gated recurrent unit (Chung et al., 2014), with hidden dimension 150. The total number of parameters in the DMM is 128,654, so that with five particles EinSteinVI is optimizing 643,270 parameters for the SM-DMM model.
We train the SM-DMM using the Adam optimizer (Kingma & Ba, 2014) with a learning rate of 10−5, using an RBF kernel and five Stein particles for four thousand epochs on the polyphonic music generation dataset JSB chorals (Boulanger-Lewandowski et al., 2012). We follow Krishnan et al. (2016) and report two version NLL of a) ∑N i=1−p(xi|θ)∑N
i=1 Ti and b) 1N
∑N i=1 −p(xi|θ) Ti
, where Ti is the length of the ith sequence. In Table 3 we report NLL(a) and NLL(b) on a held-out test set of JSB. Compared to baseline methods SM-DMM achieves a significant improvement using Stein mixtures. We see similar improvements in the test ELBO, Jankowiak & Karaletsos (2019) reports a test ELBO of -6.82 nats on the JSB dataset for their approach, SM-DMM with EinSteinVI achieves an test ELBO of 45.10 (higher is better).

heading: 7 SUMMARY
text: EinSteinVI provides the latest techniques for Stein VI as an extension to Numpyro. Our results indicate that the library is substantially faster and more expressive than other available libraries for Stein VI. EinSteinVI provides a familiar and efficient interface for practitioners working with the Pyro/NumPyro PPL and provides a unified code base to researchers for benchmarking new developments in Stein VI.
Possible further extensions include (a) supporting updates inspired by Stein Points (Chen et al., 2018; 2019b), (b) extending the kernel interface with probability product kernels (Jebara et al., 2004) to account for information geometry when using Stein-mixtures, (c) adding full support for NumPyro’s automatic enumeration features (Obermeyer et al., 2019) when guides are used, and (d) stateful matrix kernels for conjugate gradient or quasi-Newton optimization. Stein Points updates could handle cases where initialization of particles is sub-optimal by initializing from a Stein point MCMC chain and interleave MCMC updates to get better mode hopping properties.

heading: ACKNOWLEDGEMENTS

heading: A APPENDIX
text: Kernel Definition Detail Type Reference
Radial Basis Function (RBF)
exp( 1h ‖ x− y ‖ 2 2) scalar
Liu & Wang (2016)
exp( 1h (x− y)) vector Pyro10 Inverse Multi-
Quadratic (IMQ)
(c2+ ‖ x− y ‖22)β β ∈ (−1, 0) and c > 0 scalar Gorham &
Mackey (2017)
Random Feature
Expansion Ew[φ(x,w)φ(y,w)]
φ(x,w) =√ 2 cos( 1hw > 1 x + w0)
where w0 ∼ Unif(0, 2π) and w1 ∼ N (0, 1)
scalar Liu & Wang(2018)
Linear x>y + 1 scalar Liu & Wang(2018)
Mixture ∑ i wiki(x,y)
{ki}i individual kernels, weights wi
scalar, vector, matrix
Liu & Wang (2018)
Scalar-based Matrix k(x,y)I k scalar-valued kernel matrix
Wang et al. (2019a)
Vector-based Matrix diag(k(x,y)) k vector-valued kernel matrix
Wang et al. (2019a)
Graphical diag({K(`)(x,y)}`) {K(`)}` matrix-valued
kernels, each for a unique partition of latent variables
matrix Wang et al.(2019a)
Constant Preconditioned Q − 12K(Q 1 2x,Q 1 2y)Q− 1 2
K is an inner matrix-valued kernel andQ is a preconditioning matrix like the Hessian −∇2z̄ log p(z̄|x)
or Fischer information −Ez∼qZ(z)[∇2z log p(z|x)]
matrices
matrix Wang et al.(2019a)
Anchor Point Precondi-
tioned
∑m `=1KQ`(x,y)w`(x)w`(y)
{a`}m`=1 is a set of anchor points,Q` = Q(a`) is a preconditioning matrix for each anchor point, KQ` is an inner kernel conditioned usingQ`, and w`(x) =
softmax`({N (x|a`′ ,Q−1`′ )}`′)
matrix Wang et al.(2019a)

heading: B EXPERIMENTS
text: Star Distribution To illustrate the kernels included in EinSteinVI, we approximate the "star distribution", follow Wang et al. (2019b) (see Figure 3). The star distribution is constructed as a 2D Gaussian mixture,
p(x) = 1
K K∑ k=1 N (x|µk,Σk) µ1 = [0, 1.5], µk = Ukµ1 Σ1 = diag([1, 1 100 ]),Σk = UkΣ1Uk
where Uk is a rotation matrix given by
Uk = [ cos(θ) − sin(θ) sin(θ) cos(θ) ] , (k ∈ [1, ..,K])θ = 2kπ/K.
We use fifty particles and Adagrad with the learning rate that yields the best Maximum Mean Discrepancy (MMD) (Gretton et al., 2012) with a RBF kernel. To compute the MMD, we use 10,000
samples from the 2D gaussian mixture. Keeping these choices fixed, we vary the type of kernel, using a point mass distribution as guide.
We consider the scalar (Rd → R) linear kernel, the inverse multiquadric (IMQ) kernel, the RBF kernel and the random feature kernel; details are given in appendix A. In addition, we use a mixture kernel, which is a uniform mixture of the linear and the random feature kernels. We also use the matrix version of the RBF kernel with a constant preconditioning, using the Hessian matrix for preconditioning. We see that including curvature information by means of the Hessian leads to faster convergence in EinSteinVI and that the matrix RBF with preconditioning yields the lowest MMD.
Bayesian Neural Network The amortized SVGD in PyMC3 is an experimental implementation and the documentation clearly discourages its use. For completeness we include the performance here. The temperature and learning rate were determined by a grid search on the Boston data set. All measurements are repeated five times. In Table 4 we report test RMSE and time the the UCI regression benchmark. We did not include the Year data set as it takes over 21 hours to run. The RMSE for amortized SVGD is poor and highly variable across all data sets, except Naval. However, even for Naval the RMSE is two orders of magnitude greater than SVGD in EinSteinVI and Liu & Wang (2016).

Review summaries:
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. 

This paper proposed a new probabilistic programming framework for stein variational gradient descent and its variants.  The framework put its basis on Numpyro. Experimental results show that this framework shows faster computation for inference in wall clock time. Moreover, using this framework, the authors developed a new Stein mixture algorithm for deep Markov models, which shows better performance than existing methods. 

   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.