EinSteinVI: General and Integrated Stein Variational Inference
Authors: Ola Rønning
Ahmad Salim Al-Sibahi
Christophe Ley
Thomas Hamelryck
Venue: ICLR 2022 Submitted (2022)
Keywords: Stein variational inference
variational inference
probabilistic programming
Pyro
deep probabilistic programming
deep learning

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 variational inference,  integration in existing probabilistic programming languages (PPLs) with an easy-to-use interface is currently lacking.  EinSteinVI is a lightweight compostable library that integrates the latest Stein variational inference method with the PPL NumPyro (Phan et al., 2019).  EinSteinVI provides ELBO-within-Stein to support the use of custom inference programs (guides), implementations of a wide range of kernels, non-linear scaling of the repulsion force (Wang & Liu,2019b), and second-order 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 shows EinSteinVI scales to large models with more than 500,000 parameters.

Under review as a conference paper at ICLR 2022
EINSTEINVI: GENERAL AND INTEGRATED STEIN VARIATIONAL INFERENCE
Anonymous authors Paper under double-blind review
ABSTRACT
Stein variational inference is a technique for approximate Bayesian inference that has recently gained popularity because it combines the scalability of variational inference (V1) with the flexibility of non-parametrie inference methods. While there has been considerable progress in developing algorithms for Stein VI, integration in existing probabilistic programming languages (PPLs) with an casy-to-use interface is currently lacking. Vi a lightweight composable library that integrates the latest Stein VI methods with the PPL NumPyro (Phan et al. (019). EinStein VI also provides our novel algorithm ELBO-within-Sie Stein to support the use of custom inference programs (guides), in addition to implementations of a wide range of kernels, non-linear scaling the repulsion force (Wang Liu 2019b) and second- order gradient updates using matrix-values kernels (Wang et al. 2019b). We illustrate EinStein VI using toy examples and show results on par with or better than existing state-of fet-hheertt methods for real- real-world problems. These include Bayesian neural networks for regression and a Stein-mixture deep Markov model, which also shows EinStein VI scales to large models with more than 125,000 parameters.
INTRODUCTION
Interest 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 The goal infer posterior distribution p(z) x) over the latent variables given model describing the joint distribution p(z, x) = p(x)))p(Z). We obtain the posterior by following the rules of Bayesian inference:
x(zx) px(z(x)))
where = f (tpp((())(((ee))) is the normalization constant. For most practical models, the nor- malization constant lacks analytic solution requires an infeasible number of computations, complicating the Bayesian inference problem.
Variational Inference (VI) techniques (Blii al. 2017 Hoffman al. 2013 Ranganath al. 2014) aim to approximate the posterior distribution. VI poses a family of distributions over the latent variables (z)) Q and finds
arg min D(1(()) qE
where is divergencd between the variational distribution g, also called the guide, and the true posterior distribution p(z[x). The Kullback-Leibe 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 a family of techniques for approximate Bayesian inference based on Stein's method (see Anastasiou et al. (202T) for an overview) that is gaining popularity since it combines the scalability of traditional VI with the flexibility of non- parametric particle-base methods. Stein variational gradient descent (SVGD) (Liu & Wang 2016) is recent Stein VI technique which uses set of particles {z, J as the approximating distribution (z). As a particle-based method, SVGD is well suited for capturing correlations between latent variables. The technique preserves the scalability of traditional
An asymmetric distance that might satisfy the triangle inequality.

---

Under review as a conference paper at ICLR 2022
VI approaches while offering the flexibility and modeling scope of techniques such 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 or in terms of the properties of kernels (Liu & Wang 2018).
Liu 2017)
Many advanced inference methods based on SVGD have been recently developed; these include non- linear Stein (Wang & Liu (0119)), factorized graphical models (Zhuo et al.] [2018] Wang et al. [018a] matrix-valued kernels (Wang al. 2019a) and support for higher-order gradient-basee 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 PPL.s with traditional VI: their lack of flexibility in capturing rich correlations in the approximated posterior.
The main problem with SVGD that suffers from the curse of dimensionality: the number of particles required to adequately represent 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 novel algorithm for inference of Stein mixtures called ELBO-withing within-Stein. Unlike prior work on Stein mixtures (Nalismick & Smyth 2017), ELBO-within Stein i) is simple to implement, ii) only has the learning rate 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 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 EinStein VI library. Our library further includes SVGD as special case of LLLBO-wtthi--teee 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.
and all the advanced
The EinStein VI library extends the NumPyro PPL (Bingham al. 2019 Phan et al. 2019). NumPyro is universal probabilistic programming language (van de Meent et al., (2118) 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 EinStein VI works with arbitrary guides, NumPyro is well suited language for embedding EinStein VI. 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 VinSteinVI, and iv) benefits computationally from Jax. Our extensions include SVGD, Stein-mixtures formulated MLLBO-within-StH and the advanced methods mentioned above.
Concretely, our contributions are:
ELBO-within-Stee (EinStein VI). 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 Ein EinStein allows SVGD to work with custom guide programs based on ELBO-within-Steer 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 possible for EinStein grow organically with NumPyro development.
Integration of recent developments in Stein variational inference which includes: non- linear optimization (Wang & Liu 2019a), wealth of kernels (Liu & Wang] 2016 2018 Gorham & Mackey 2017). matrrxxvvllud kernels (Wang et al., 2019a) supporting higher- order optimization, an (experimental) update based on Stein point MCMC (Chen al.) 2019a), and factorization based on conditional independence between elements in the model (graphical kernels) (Wang al. 2018b)
A series of examples that demonstrate EinStein VI and the synergy between different Stein VI techniques. The examples include novel Stein- mixture version of the deep Markoy
2

---

Under review as a conference paper at ICLR 2022
model (SM-DMM), Bayesian neural networks for regression. An examples that explore kernels and bigher-order optimization 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 we introduce Stein mixtures and our algorithm ELBO-within-Steinn We discuss the general details of the implementation of EinStein in Section 4 Next, we discuss related work in Section [S] In Section 6 we present various examples using
Vin and finally, we summarize our results and discuss future work in Section
STEIN VARIATIONAL GRADIENT DESCENT
The core idea of SVGD to perform inference by approximating the target posterior distribution
p(z|x) by an empirical distribution (z) = N E, dz, (z) based on set of particles 2 = { Here, (x()) represents the Dirac delta measure, which is equal to 1 x = y and 0 otherwise. One could thus by a particle 2; E Z. The idea is to minimize the Kullback-Leible divergence DKL (2z (z) p(z|x))
the approximating distribution gz (z) a mixture of point estimates, each represented
between the approximation and the true posterior by iteratively updating the particles using the Stein forces:
2: - Z; + eSz
where is the learning rate and Sg denotes the Stein forces.
The Two Forees SVGD Stein VI consists of two forces which work additively under the form Sz S + Sz. where the attractive force is given by
(z.) = [k(z V logp(z,(x))
and the repulsive force by
Sz(z.) Be V k(z;, 2;)]-
Here k: Rd x Rd -> R is kernel. The attractive force 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(z,, 2j) exp (- A z; Zq (1ž) with bandwidth parameter chosen as logn med(z).
The repulsive force moves particles away from each other, ensuring that they not col- lapse to the same mode. For example, with the RBF kernel, the repulsive force becomes (2) exp - z, ny [1:] E Er - 2,z)], which has high kernel values for particles that are close. causing them to repel.
SVGD works with unnormalized distributions as the normalization constant becomes additive in the log-posterior log xx) = log Z + log p(x/z) + log and is thus not required for the calculation of the gradient. This property is desirable as normalizing is often computationally
expensive.
Non-linear Stein In non-linear Stein (Wang & Liu| 2019a), the repulsive force can be scaled by factor X, resulting in Sz 5 + A52- This approach is useful when dealing with multi- modal distributions. It is also useful cases where the repulsive force vanishes compared to the likelihood, which happens for large datasets (X'). Scaling the repulsive force by constant À = c(I2') proportional (e.g. e 0.1 or e 0.01) to the size of the dataset |.Y' addresses this issue and can be chosen by cross-vas volid-aaliottion on subset of the data.
Matrix-valued kernels 2019a), Rd Rd + idd in which case the Stein forces become
The choice of kernels can be extended to matrix-valuee ones (Wang et al.
(z,) = Ez Enj (x) (z, z, Vz, log p(zy [x]
and
S (2.) Ez S (K) (K(z, Vz,)
3

---

Under review as a conference paper at ICLR 2022
where the standalone del 2z, in the repulsive force represents the vector Øzy. ..... Fzy. This results in
Σ VeKee
The advantage of matrix-valued kernels is that they allow preconditioning using the Hessian or Fisher Information matrix, which can capture local curvature and thus achieve better optima and convergence rate than standard SVGD. Furthermore, 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(t)
STEIN MIXTURES
The Stein-mixture was proposed by Nalisnick & Smyth (017 to resolve the representation issue for SVGD when the target distribution has high dimensionality. For SVGD, the number particles needed to represent 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 (z) Σ k= (. The guide is the joint model q(, z) q( g(0lz)g(z) with (0)) (z) the first tier distribution and g its prior. The second tier posterior
X) is optimized using SVGD, so that the marginal posterior is q(0; z) = N Σ ( z). a restricted mixture.
Nalisnick & Smyth
2017) showed that the Stein force for a mixture approximation of the posterior
plz|2 ()) given by
(2,),) E,, (2)
[kzi, 2z) Bo (or (p(X, (0)(0)))((0))3,)]]
+ Fo (z)
[k((2; 2;) Va, log logn(z;))
+ S (z:),
where log V2, (Zzz) acts regularizer on the variational parameters which they assume is suffi- ciently small to be dropped. To evaluate Bzj (x) [p(X, ())((()); (())z;)), DNCP and importance weighted Monte Carlo gradients to obtain the black-box update,
Nalisnick & Smyth (2017) use
S
) q(0,) 2z)
0()
(z(),) (z)00g(())
k(zi, 2j)
W. Vz, log
+ 5Z(z))
where W, is an importance weight for sample 0, = u(2j, E), po (i.e. the guide).
EittO-withnn-Stein In ELBO-witthnnSteem we replace the weighted average over samples in Equation T by single loss (C). If we choose L to be f-divergence, such the ELBO, reduce the loss bias by averaging multiple samples. However, ELBO-within-S Stein only computes gradient of L per particle regardless of the number of samples different from Nalisnick & Smyth (217) who estimate Eg.-v(01,,) [p(X, (,) by averaging over gradients, therelore computes gradient and an importance weighting for each DNCP sample rather than per particle. The difference makes ELBO-within-Seeii computationally cheaper as the complexity of adding gradient Onn) whereas adding a sample estimate O(1). The Stein force in w-tthin--ttein is given by
to estimate L. This is
Sz(z:) = Ee [k(zs, V,, (2))))
(2)
where L is the ELBO.
Using preconditioner matrix Q, such that Q M has lower condition number than M.
4

---

Under review as a conference paper at ICLR 2022
4 COMPOSITIONAL IMPLEMENTATION USING NUMPYRO
EinStein VI 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 EinStein VI (see Figure I
Below, we discuss the key features EinSteinVI, which include re-initializable guides. EinSteinVT's core algorithm, and the new kernel interface.
4.1 RE-INITIALIZABLE GUIDES
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, provide the Reini tGiide interface, which requires implement- ing a function find_paraa ams that accepts a list of random number generator (RNG) keys in addition the arguments for the guide and returns a set of freshly initialized parameters for each RNG key.
The WrappedGuide class provides guide written as function. WrappedGui de makes callable guide re-initializable. works by running the provided guide multiple times and reinitializing the parameters using NumPyro's '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 initializa- which initializes each parameter with uniform
tion strategy, e.g., tostt_tt_uniirrree random value the range [-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 P have their own initializers.
The Stein VI interface will correctly wrap the guide during initialization, so from user perspective the syntax for guides follows the API of SVI.
4.2 STEIN VI IN NUMPYRO
The integration of Stein VI into NumPyro requires handling transformations between the parameter representation of NumPyrd and the vectorized Stein particles that FinStein VI operates on. For this, we rely on Jax PyTrees which converts back and forth between Python collections and a flattened vectorized representation.
Algorithm shows the core algorithm of EinStein VI. EinStein updates the standard variational model parameters and guide parameters &p 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 monolithic vector-encoded particle using Jax PyTrees. The monolithic particle represents the particles as a flattened and stacked Jax array. Then we compute kernel based on the vector-encoded Stein partiele; 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 vectorized manner. This is done unconstrained space 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 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 vmap operator with a nested tree_mant to
readthedocs.ioo/e//lateeee httpss//jia.rraaaoooss.oooee
http://eeetttpeeeeepernentttalla stax. html A dictionary mapping parameters to their values, which can he arbitrary Python type https: http:///..eaattttddsss/////////s
https://jax
5

---

Under review as a conference paper at ICLR 2022
compute the desired Jacobians. Note that Algorithm presents how EinSteinVI VI 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 oonterctrr pactionary-based form and return the expected changes for standard- and Stein Stein-parammn parameters.
Input: Classical VI parameters Φ and sb. Stein parameters {0.).. model pæ(z, guide qe.vi v (z), loss C, kernel interface KI.
Output: Parameter changes based on classical VI (A Axp) and Stein VI forces (f)0,):).
procedure EINSTEIN(Ø), sb, {0. his Pæ. 90,4)
AΦ < Ep V a (pp., (
AΦb - Eg [ (p qo.0)
{a.): ; f FyyTreeFlater (1)):)
< kI((a:)
procedure FINSTEINFORCESSaa))
> Calculate forces per particle for higher-order vmap function.
0, <- PYTREERESTORE((2,
Aa; + --n,, k(aj, a, V ((ps. 20 V k(a,
return Aa, end procedure
Aa; : + VMap((a:): EINSTEINFORCES) A0; J: + PYTREIRISTORE(Aa.))
return Ad. Awb, (A0:): end procedure
Algorithm 1: EinSteinVI
4.3 KERNEL INTERFACE
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 C and returns a differentiable kernel k. All kernels are currently static, but the interface could be extended to stateful kernels, allowing conjugate gradients or quasi- Newton optimization. Appendix Aives the complete list of kernels EinStein VI. We are planning to extend EinStein VI with probability product kernels Nalisnick & Smyth (201)), which take into account the information geometry of the first tier variational distributions in Stein mixtures.
RELATED WORK
There has been proliferation of deep probabilistic programming languages (Ge 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 St.inVI.
In PyMC3 the inference techniques that use Stein's 's method include SVGD with a scalar RBF-kernel, amortized SVGD (Wang et al. 2016: Feng et al.) (017). and Operator Variational Inference (OVT)
(Ranganath et al. 2016a).
SVGD in PyMC3 manipulates particles to directly capture the target distribution (Liu & Wang 2016). In FinStein VI, SVGD is special case of Stein mixtures where the guides are point mass distributions. Because we have guide programs, EinStein VI allows arbitrary computations to transform random variables in the guide. This is not possible with SVGD in PyMC3.
Amortized SVGD (Feng al. 2017) trains stochastic network draw samples from target distribution. The network is iteratively adjusted so that the output changes in the direction of the Stein variational gradient (the same gradient used in SVGD). In comparison, EinStein VI transports a fixed set of particles (each parameterizing guide) to the target distribution. Amortized SVGD is not in EinStein VI its extension arbitrary guides is open problem.

---

Under review as a conference paper at ICLR 2022
def model() sample(" x', Normal (arr 31).
jnp apraar[( ([-2.0, 2.0]), imp.array(Cn) 01.0, 1.0])))
(a) ID Gaussian mixture model
svi SVI(
stein SteinVI( model, (atDella(model)
model, AutoNormal (model),
Astepa (ttepdsieee__ 1ize=1100,
@tterad(ste__sizer www.ssi.eeio. Trace_ELBO(), RBFKermel)),
Trace_ELBO()
results = svi. svi.run(rng_kos num num_iterat
results = stein. raalnngunnnnnneae key,
num_iterations)
SVI
(c) SVGD with Ei
Figure 1: 1D Gaussian mixture model in NumPyro.
OVI optimizes operator objectives, which take functions of functions to non-negative number. Ranganath et al. (2016a) include operator objective based the Langevin-Stein operator (Anas- tasiou et | 2021). I. This is the same operator used for the kernelized Stein discrepancy (Liu et al. 2016) which also underlies SVGD. Unlike EinStein VI, Amortized SVDG and SVGD, OPV is not particle-based method.
6 EXAMPLES
We illustrate the features of FinStein on two toy examples, namely 1D) mixture of Gaussians below, and 2D mixtures of Gaussian in Appendix B We demonstrate EinStein VI 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 EinStein VI package itself are available through an Anonymized URL.
1D Gaussian mixture To demonstrate the modes of VI with EinStein VI, consider the 1D Gaussian mixture 1/3 1) 22, 1) (see Figure 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 SVI with 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- mixinee 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 [}, and as such is algorithmic. The Stein-mixture will therefore not be able to exactly capture the mixing components of target mixture model with one particle per component. However, with more particles the mixture can be approximated better demonstrated using three particles.
Bayesian Neural Networks We compare SVGD and non-linear Stein in EinStein with the implementation of SVGD by Liu & Wang 2016 (without model selection), and amortized SVGD (using the Theano backend Al-Rfou al. (2016) Bayesian neural networks (BNN) for regression. The PyMC3 documentation clearly states amortized SVGD 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 EinStein VI for fair comparison with Liu & Wang (016)) and PyMC3, For non- nin-linear Stein we determined the best repulsion factor X by grid search for {10 2 10 10 102}. For SVGD PyMC3 we use
'https://github com/ http//iiiinnggn///eeeeeeee htw.://///tteiiaaaaaatiiiiaaaaa Variational Gradient DDscentt wettantt/booo//mmtteeeee python/ http:/baaesiiiann..aoa PPyMC4, which uses Jax its backend, does not include SVGD the time of writing.

---

Under review as a conference paper at ICLR 2022
0.3
0.3
0.3
0.
0.2
0.2
0.2
0.2
0.3
0.1
0.3
0.1
0.0
0.01
0.01
(a) SVI
(b) SVGD (RBF kernel)
Two particle Stein-
(d) Three particle Stein- mixture (RBF kernel)
mixture (linear kernel)
Figure 2:
The blue dashed line is the target pdf, while the solid green line is the density of the particles. We estimate the particle density for SVGD with Gaussian kernel density estimation. use 100 particles for SVGD, and two three particles for the Stein-mixture. SVI uses Gaussian guide.
Table 1: Average test RMSE and time for inference for the UCI regression benchmarks.
Test RMSE
Time
EinStein
PyMC3
Liu & Wang 2016
EinStein
PyMC3
Dataset
SVGD
NL-Stein
SVGD
SVGD
SVGD
SVGD
2.610= 0.044
2.492
0.028
0.244
2,990 0.013
5.758s 0.048:
166.396s 0.301s
Concrete
1.175 = 0.025
4.782 + 0.017
6.255
0.390
5.927 = 0.021
3.602s + 0.048s
1m5.815s 0.278s
Energy
0.321 + 0.004
0.375 + 0.007
3.494
1.500
1.251 + 0.017
5.562s + 0.053s
1m6. 1m6.175 175x + 0.41s
Kin8nm
0.073 =
0.076
0.000
0.164
0.040
0.115 = 0.001
6.141s + 0.04s
2m40.106 0.298s
Naval
0.001 + 01.000
0.001
+ 0.000
0.029 0.015
0.008 + 0.001
6.51s + 0.048s
3m51.05 0.385s
Power
3,952 0.004
3.907 0.006
2:219 5.389
4.215 + 0.002
6.314s - 0.054s
2m51 654s 0.566s
Protein
4.509 01.066
4.654 1: 0.006
5.921
2.539
4,846 + 0.002
9.436s 1 0.032s
11m11.9s 0.399s
Wine
0.599 + 0.006
0.596 +: 0.005
0.212 0.769
0.604 + 0.001
5.614 + 0.052s
1m7.767s 0.286s
Yacht
8.341
5.519s
1m4.748s 0.392s
Year
8.662 + 0.002
8.701 + 0.002
135.316
8.896 + 0.001
1m37.569s 0.07s
4h Jh19336s
temperatur ten. Like Liu & Wang (2016) we use a BNN with one hidden layer of size fifty and RELU activation. We put Gamma 0.1) prior 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
GP 19 for EinStein VI and PyMC3, except for Year with PyMC3 which was only run once. We do not report times for Liu (2016) because only the CPU version of their code could be executed without irresolvable issues.
repeated ten times and obtained
Liu & Wang
Table shows the performance terms of the average root mean squared error (RMSE) on the test set. We find that EinStein VI achieves significantly better RMSE than Liu & Wang (016) and that both systems outperform SVGD in PyMC3. The times in Table measure how long it takes to infer parameters. exxlluds the first run for reasons: i) that run will fill the GPU caches, and ii) Jax will trace the programs. As result, this run is more costly than subsequent ones. The times for the first run are given Table [ initial cost.
By running EinStein VI for more iterations, we can amortize the
Stein Mixture Deep Markoy Model Music generation requires 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 model parameters. We consider 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. consist of two layers with hidden dimension 200 and ReLU activation on the first layer and sigmoid
'We found the performance he very sensitive the choice of temperature.
QQadro RTX 6000 with Cuda V11.4.120
Table
Time for first repetition with EinStein for UCI regression benchmarks.
Datasct
Boston
Concrete
Energy
Kin8nm 43.592s
Naval
Power
Protein
Wine
Yacht
Year
Time
41.665s
41.642s
41.591s
44.2570s
44.058
47.87s
41.9s
41.409k
2m18. 191
8

---

Under review as a conference paper at ICLR 2022
Table 3: Test negative log- likelihood (lower is better) on Polyphonic Music Generation (JSB) dataset. Baseline results from Krishnan et al. (2016). ISN-DMM and ISN-DMM-Aug (Krishnan et al. 2016
TSBN and HMSBN GGn al. 2015)
ISN-DMM
ISN-DMM-Aug
HMSBN
TSBN
SM-DMM
NLL (a)
6.926
6.773
8.0473
-45.983
NLL (b)
6.856
6.692
7,9970
7.48
-46.066
on the second. The Emitter network is a three layers FNN which produces likelihood at each time step using the current latent state. The layers have hidden dimensions 100, 100, and 88, respectively. The variational distribution of the form II g(=".r.. ( function SI:T, is one layered gated recurrent unit (Chung al. (2014), with hidden dimension 150. The total number of parameters in the DMM is 128,654, so that with five particles EinStein is optimizing 643,270 parameters for the SM-DMM model.
where the parametrized feature
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
Lewandowski et al. (012). We follow Krishnan et al.
(0066) and report two version NLL of a)
0(x.00) where T; is the 2 p(x. and b) N -== 7.
EN
length of the ith sequence. In Table we report NLL(a) and NLL(b) on a held-out test set of JSB. Compared to baseline methods SM-D MMDDMM achieves a significant improvement using Stein mixtures. We see similar improvements in the test ELBO, Jankowiak Karaletsos (2019) reports test ELBO of -6.82 nats on the JSB dataset for their approach, SM-DMM with EinStein VI achieves an test ELBO of 45. 10 (higher is better).
SUMMARY
EinStein provides the latest techniques for Stein VI as an extension Numpyro. Our results indicate that the library is substantially faster and more expressive than other available libraries for Stein VI. EinStein provides familiar and efficient interface for practitioners working with the Pyro/NumPyro PPL and provides 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 Num NumPyro's automatic enumeration features (Obermeyer et al. 2019) when guides are used, and (d) stateful matrix kemels for conjugate gradient 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 get better mode hopping properties.
ACKNOWLEDGEMENTS
REFERENCES
Rami AI-Rfou, Guillaume Alain, Amjad Almahairi, Christof Angermueller, Dzmitry Bahdanau, Nicolas Ballas, Frédéric Bastien, Justin Bayer, Anatoly Belikov, Alexander Belopolsky, al. Theano: A python framework for fast computation of mathematical expressions. arXiv e-prints, pp. ar0iv-1605, 2016. URL /trL:ttt///aiii.orr////// abs/1//5/ 02688
Andreas Anastasiou, Alessandro Barp, François-Xavier Briol, Bruno Ebner, Robert Gaunt, Fatemeh Ghaderinezhad, Jackson Gorham, Arthur Gretton, Christophe Ley, Qiang Liu, Lester Mackey, Chris. J. Oates, Gesine Reinert, and Yvik Swan. Stein's method meets statistics: A review of some recent developments. arXiv preprint arXiv:2105.034: 2021. https::: org. abb// 2205. 03481
Eli Bingham, Jonathan P. Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul A. Szerlip, Paul Horsfall, and Noah D. Goodman. Pyro: Deep

---

Under review as a conference paper at ICLR 2022
universal probabilistic programming. J. Mach. Learn. Res., 20:28:1-288. 2019. URL http: //jmlr www.rr..oggppppr///2//////
html
David M. Blei, Alp Kucukelbir, and Jon D. McAuliffe. Variational inference: A review for statisticians.
Journal of the American Statisticaal Association, 112(518):55--877 877, 2017. doi: 10. 1080/01621459. 2017.1285773. htt.sttps:///o..orrr
1080 0011245961 2017. 1285773
Nicolas Boalanger-Lewandowssss Yoshua Bengio, and Pascal Vincent. Modeling temporal dependen- cies in high-dimensional sequences: Application to polyphonic music generation and transcription. arXiv preprini arXiv: /206.6392, 2012.
Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance weightte autoencoders. arXiv preprint 1XXi::15090005 2015. URL BKL..ttpss//aarx..orr//bb////// 00519
Wilson Ye Chen, Lester W. Mackey, Jackson Gorham, François-Xavier Briol, and Chris Oates. Stein points. In Jennifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmüssuun, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 843-852. PMLR, 2018. URL http://proceediinn / chen18f.. html
mlr, prwss/v80
Wilson Ye Chen, Alessandro Barp, François- Xavier Briol, Jackson Gorham, Mark Girolami, Lester Mackey, and Chris Oates. Stein point markov chain monte carlo. In International Conference Machine Learning, pp. 101 1021-1 PMLR, 2019a. URL btttps://axxiv.org/// 1905
03673
Wilson Ye Chen, Alessandro Barp, François-Xave Briol, Jackson Gorham, Mark A. Girolami, Lester W. Mackey, and Chris Oates. Stein Point Markov Chain Monte Carlo. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 1011-1021. PMLR, 2019b. URL http //proceedings.mil mmr. http./eerees////cee110000 html
Junyoung Chung. Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling, arXiv preprint arXiv:1412.3555, 2014. URL
www.gg////b///////// .555
Marco Cusumano-Townes Feras Saad, Alexander K. Lew, and Vikash K. Mansinghka. Gen: general-purpose probabilistic programming system with programmable inference. In Proceedings of the 40th ACM SIGPLAN Conference on Programming Language Design and Implementation. PLDI 2019, pp. 221-236 New York, NY. USA, 2019. ACM. ISBN 978 00044550377222 doi: 10.1145/3314221.11 4642.. URLpttp: acm. org /00 11t..1111/33322200 3314642
Gianluca Detommaso, Tiangang Cui, Youssef M. Marzouk, Alessio Spantini, and Robert Scheichl. A Stein variational Newton method. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicold Cesa-Bianchi, and Roman Garnett (eds.), Advunces in Neural Information Processing Systems 31: Anmual Conference on Neural Information Processing Systems 2018, NeurlPS 2018, 3-8 December 2018, Montréal, Canada, pp. 9187-9197, 2018. URL http: \/papers. cc/paper 8130-a--1 www.iinaaaiaiinnaalllaaaaaa newtonn method
Joshua V Dillon, Ian Langmore, Dustin Tran, Eugene Brevdo, Srinivas Vasudevan, Dave Moore, Brian Patton, Alex Alemi, Matt Hoffman, and Rif A Saurous. Tensorflow distributions. arXiv 2017.
preprint ar arXiv:1711.
Yihao Feng, Dilin Wang, and Qiang Liu. Learning draw samples with amortized stein variational gradient descent. arXiv preprint arXiv 2017. URL. 1707. 06626
Roy Frostig, Matthew Johnson, and Chris Leary. Compiling machine learning programs via high-level tracing. In SysML 2018, 2018. URL URL http:://////w cc/ym/llcc/oo///01///4/444 pdf
Zhe Gan, Chunyuan Li, Ricardo Henao, David E Carlson, and Lawrence Carin. Deep temporal sigmoid belief networks for sequence modeling. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama,
10

---

Under review as a conference paper at ICLR 2022
and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 28. Cur- ran Associates, Inc., 2015. URL https:: otl.ap/pw./pppoeettr.....too..rp 95 351110bbbbbf//brttaaaaaaaaaaaaa. Papeerr----ppe.... .dff
Hong Ge, Kai Xu, and Zoubin Ghahramani. Turing: a language for flexible probabilistic inference. In International Conference on Artificial Intelligence and Statistics, AISTATS 2018, 9-11 April 2018, Playa Blanca, Lanzarote, Canary Islands, Spain, pp. 1682-1690, 2018. URL http:
//proceedings. mlr arees//v4/ge188 html
Jackson Gorham and Lester W. Mackey. Measuring sample quality with kernels. In Doina Precup and Yee Whye Teh (eds.), Proceedingy of the 34th International Conference Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pp. 1292- 1301. PMLR, 2017. URL http::/ppooeeeiiiiii mwr.smmmlppp.... mlr.press/ v70 770/gorhaml7a.hee html
Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. kernel two-sample test. The Journal of Machine Learming Research, 13(1):723-773, 2012. 2012 URL httpp:://wet jmlr
Matthew Hoffman, David M. Blei, Chong Wang, and John W. Paisley. Stochastic variational inference. J. Mach. Learn. Res., 14(1):1303-138) 2013. URL http://il..cc.oor//iiaat...o orgg/itatiintio
cfm?id=2502622.
Martin Jankowiak and Theofanis Karaletsos. Pathwise derivatives for multivariate distributions. In Kamalika Chaudhuri and Masashi Sugiyama (eds.), The 22nd International Conference on Artificial Intelligence and Statistics, AISTATS 2019, 16-18 April 2019, Naha, Okinawa, Japan, volume 89 of Proceedings of Machine Learning Research, pp. 333-342. PMLR, 2019. URL http://proceedings mmriiinnnsm.......mmmmmmm.m wwt.pppwssssvjjjeeee jttp://innnn......rrerr//////aaa jaaess//8///akkooia11....../
Tony Jebara, Risi Kondor, and Andrew Howard. Probability product kernels. The Journal of Ma- 2004. URL hethh:ttp//:w. jmw....jmlr...mmm.m..m.......... jmlr. www.oogg/ppp/ee. ctoperrsvvoluees
chine Learning Research,
jebara04a/jebaraA44 .d.f
Diederik Kingma and Max Welling, Efficient gradient-based inference through transformations between bayes nets and neural nets. In International Conference on Machine Learning, pp. 1782-1790. PMLR, 2014. URL http:://proceeeingggg////o/o mlr prwss//v32//iingmaaaaa pdf
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
6980
Rahul G Krishnan, Uri Shalit, and David Sontag. Structured inference networks for nonlincar state space models. arXiv preprint arXiv: 2016. URL [K..httpss//axxivorr///5s 09869.
Yingzhen Li and Richard E Turner. Rényi divergence variational inference. arXiv preprint arXiv:1602.023311 2016. ItLp
02311
Qiang Liu. Stein variational gradient descent as gradient flow. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio. Hanna M. Wallach, Rob Fergus, S. V. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference Information Processing Systems 2017, 4-9 December 2017, Long Beach, CA, USA, pp. 3115- 3123, 2017. URL httpthtttppppss.....e/ nips. wwtpwcccpppppeeeeccc stein-tein--ie www.sstewaaaiinaalaoooooooo gradiennt descent - as- gradient- -srattaiiggrrdnnttnn-a flow
Neural
Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose Bayesian inference algorithm. In Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg, Isabelle Guyon, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 29: Annual Conference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain, pp. 2370- 2378, 2016. URL http: htt/hattp//ppppes..iipppp nipse http://ssssninccccpppeeee/////a wttper//----sseeeeaaaaiiiiioooo -radiieng descent www.geneallpppm purpose-- bayeoooseebbbbeiaanna.aa
paceeeaaaaaaallpppppeeeeeaaaiiaa www.riiiieeieeiee algorithm

---

Under review as a conference paper at ICLR 2022
Qiang Liu and Dilin Wang. Stein variational gradient descent moment matching. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò desa-Bianchi, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 31: Annual Conference Neural Information Processing Systems 2018, NewrlPS 2018, 3-8 December 2018, Montréal, Canada, pp.
2018. URL httt:://appeeeee gradient-descene as moment matching
Qiang Liu, Jason Lee, and Michael Jordan. A kernelized stein discrepancy for goodness-of-ln tests. In International conference on machine learning, pp. 276-284. PMLR, 2016.
Eric Nalisnick and Padhraic Smyth. Variational inference with stein mixtures. In Advances in Approx- imate Bayesian Inference, NIPS 2017 Workshop, 2017. http//ttpp///prroxi/mttiieeeeeee .dff
org/2017 hhgp:/277/cccppeee/aaaass/////a/
Fritz Obermeyer, Eli Bingham, Martin Jankowiak, Du Phan, and Jonathan P Chen. Functional tensors for probabilistic programming. arXiv preprint arXiv:1910.10775, 2019. URL https
wwwp:ii..rrg/bbb/19990 10775
Du Phan, Neeraj Pradhan, and Martin Jankowiak. Composable effects for flexible and accelerated probabilistic programming numpyro. In Program Transformations for Machine Learning, NeurlPS 2019 Workshop, 2019.
Rajesh Ranganath, Sean Gerrish, and David M. Blei. Black box variational inference. Proceedings of the Seventeenth International Conference Artificial Intelligence and Statistics, AISTATS 2014, Reykjavik, Iceland, April 22-25, 2014, volume 33 of JMLR Workshop and Conference Proceedings, pp. 814 822. JMLR.org. 2014. URL http::p:p/prrocceiingss... .mr...mm.......... wtt........sss////aaaaaaaatt.aa html
Rajesh Ranganath, Dustin Tran, Jaan Altosaar, and David Blei. Operator variational inference. Advances in Neural Information Processing Systems, 29:496-504, 2016a. URL https?
//proceedings. newrips..ppp... cwttrips.cc/pppppre.ee pc/ppppe//211/ffl/////2222222210 httpp////..//...........//////0. Paper. .pff
Rajesh Ranganath, Dustin Tran, and David Blei. Hierarchical variational models. In International
Conference on Machine Learning, pp. 324-333. PMLR, 2016b.
John Salvatier, Thomas V. Wiecki, and Christopher Fonnesbeck. Probabilistic programming in python using pymc3. PeerJ Comput. Sci., 2:e55, 2016. doi: 10.7717/peerj-cs.55. URL httpps //doi. org. 10.
Dustin Tran, Alp Kucukelbir, Adji B. Dieng, Maja Rudolph, Dawen Liang, and David M. Blei. Edward: A library for probabilistic modeling, inference, and criticism. arXiv preprint 16Xiv: 09787, 2016.
Wil-We van de Meent, Brooks Paige, Hongseok Yang, and Frank Wood. An introduction to probabilistic programming. CoRR, abs/1809. 10756, 2018. URL http:///://axx org. arg89 10756
Dilin Wang and Qiang Liu. Nonlinear stein variational gradient descent for learning diversified mixture models. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 6576-6585, Long Beach, California, USA, 09-15 Jun 2019a. PMLR. URL http://proceeding m wwwsss/svvi https//////wwggg9/// html
Dilin Wang and Qiang Liu. Nonlinear stein variational gradient descent for learning diversified mixture models. In International Conference on Machine Learning, pp. 6576-6585. PMLR, 2019b. URL httphhtt://prrccednnngg /ttpeedddn.g......prse////////a/ press//97/wwng119h.t html
Dilin Wang, Yihao Feng, and Qiang Liu. Learning to sample using stein discrepancy. In NIPS workshop on bayesian deep learning, volume 192, 2016. hut.tppppp/aayeiimmepeeemmmig org/2016/ 2ttp://////pppr//BB111111
12

---

Under review as a conference paper at ICLR 2022
Dilin Wang. Zhe Zeng, and Qiang Liu. Stein variational message passing for continuous graphical models. In Jennifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference Machine Learning. ICML 2018, Stockholmsmäässan. Stockholm, Sweden, July 10-15 2018, volume 80 of Proceedings of Machine Learning Research, pp. 5206-5214, PMLR, 2018a. URL htt.htttpppproceeiggg.. mir vesss//80/wang181..
Dilin Wang, Zhe Zeng, and Qiang Liu. Stein variational message passing for continuous graphical models. In International Conference on Machine Learning, pp. 5219-5227. PMLR, 2018b. URL http://proceedinng. www.igs..m...ppp.... wil..ppressvv8)) www//wanggnl.l... html
Dilin Wang, Ziyang Tang, Chandrajit Bajaj, and Qiang Liu. Stein variational gradient de- scent with matrix-vale vallud kernels. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelz- imer, Florence d' d'Alcéé-Bue Buc, Emily B. Fox, and Roman Garnett (eds.). Advances in Neu- ral Information Processing Systems 32: Annual Conference on Neural Information Pro- cessing Systems 2019, NeurIPS 2019, 8-14 December 2019, Vancouver, BC, Canada, pp. 7834-7844, 2019a. URL http://papers www..ins..ccccc....m wwwp.nccpp.pp.r......c 8998- http//898888tttnenne ste----triiiiionnallll medieent-descent with- valtriix-vvlue--eeeellll
Dilin Wang, Ziyang Tang, Chandrajit Bajaj, and Qiang Liu. Stein variational gradient descent with matrix-valued kernels. Advances in Neural Information Processing Sys- lems, 32:7834, 2019b. URL https: https://proceed: http://gw...aaiesscc////////////
JaccooddBB301167773333322222242 Paper. pdf
Andrew Gordon Wilson. The case for bayesian deep learning. arXiv preprint arX.v.2001.10995. 2020.
Andrew Gordon Wilson and Pavel Izmailov. Bayesian deep learning and probabilistic perspective of generalization. arXiv preprint 2rX1v:2002.0879), 2020.
Jingwei Zhuo, Chang Liu, Jiaxin Shi, Jun Zhu, Ning Chen, and Bo Zhang. Message passing Stein variational gradient descent. In Jennifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference
Machine Learning, ICML 2018, Stockholmsmüssan, Stockholm, Sweden, July 10-15, 15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 6013- 6022. PMLR. 2018. htLp:ttppproceedinggss www.sme.r.rrrree0eeem htte//huolll...a httll
13

---

Under review as a conference paper at ICLR 2022
A APPENDIX
Kernel
Definition
Detail
Type
Reference
Liu & Wang
Radial Basis Function (RBF)
exp( hh2 x - y (12)
scalar
(2016)
exp( :/j (x yy)
vector
Pyrc
Inverse Multi-
Gorham &
Quadratic
(c2+ y 1
8 (-1,0) and >
scalar
Mackey
(IMQ)
(2017)
Φ(x, x, w) =
Random
V2 cos( h w x + wo) where wo ~ Unif(0, 2T) and W ~ (0,
Liu & Wang
Feature
Ewlo(x, w))(y,w))
scalar
(2018)
Expansion
Liu & Wang
Linear
x y
scalar
(2018)
scalar,
(ke): individual kernels, weights wi
Liu & Wang
Mixture
Σ, wx,k(xx)))
vector,
(2018)
matrix
Ccalar-based
Wang et al. (209aa)
k(x,y)I
k scalar-valued kernel
matrix
Matrix
Vector-based
ding(k(x,y))
1: vector-valued kernel
matrix
Wang et al. (2019a)
Matrix
{K(()) matrix-valued kernels, each for unique partition of latent variables K' is inner matrix-valued
Wang et al.
Graphical
matrix
(019a)
kernel and Q is a preconditioning matrix like
Constant Pre- conditioned
Q Q=
the Hessian - V log p(z x)
matrix
Wang et al. 2019a)
Fischer information
E Ex (x) [V2 p(z x)]
matrices
a Te is a set of anchor
points, Q₂ Q(a+)
Anchor Point Precondi-
preconditioning matrix for each anchor point, Kq, is an inner kernel conditioned using Qr. and wi wx(x) sox lag, QF )e)
Wang et al.
Σm
matrix
2019a)
tioned
B EXPERIMENTS
Star Distribution distribution". follow Wang al. (2019b) (see Figure 3. The distribution is constructed as 2D Gaussian mixture,
To illustrate the kernels included in EinStein we approximate the "star
K Σ
p(x) K 1=1
N(<lp,, Ex) P1 = [0, 1.5), uk U341 Σi = diag([,, j 1): Σx = UKE
where Us is a rotation matrix given by
cos(O) sin
sin( cos(
UR =
(k E [L,.., K])0 2kT K.
We use fifty particles and Adagrad with the learning rate that yields the best Maximum Mean Discrepancy (MMD) (Gretton al. 2012) with RBF kernel. To compute the MMD, 10,000
14

---

Under review as a conference paper at ICLR 2022
Initial
Unear
IM0
RBF
MQ
Fanduy facle
Minto-oons
Berot
Figure 3:
Particle positions for EinStein VI with different kernels after 1000 iterations with a point mass guide, starting from particle positions given in the upper left frame labelled Initial. The MMD, bottom right frame, evaluated using the RBF kernel.
Table 4: Average test RMSE and time for inference for the UCI regression benchmarks with amortized SVGD.
RMSE
Time
Boston
87230.469 + 54219.534
5m466 = 3s
Concrete
2458.250 +: 754.653
6m3s +
Energy
2458.250 + 754.653
5m45s 3s
Kin8nm
302.218 + 43.818
19m54s 5s
Naval
0.386 = 0.128
28m39s + 3s
Power
3333.796 + 517.241
21m38s 68
Protein
64.791 + 12.181
1h32m28s + 1mls
Wine
12.929 + 6.202
7m21s = 3s
Yacht
2822.983 + 944.449
5m35s 4s
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 In addition, we use 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 EinStein and that the matrix RBF with preconditioning yields the lowest MMD.
experimental implementation and the documentation clearly discourages its use. For completeness we include the performance here. The temperature and learning rate were determined by grid search 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
Bayesian Neural Network The amortized SVGD in PyMC3
it takes over 21 hours 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 EinStein VI and Liu Wang (2016).
15


Reviews:
[{"review_id": "z7L0B82gYxc", "reviewer": "Reviewer_Hnfk", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "[Strength]\n\nThe paper is overall well-written and the method is clearly explained. The literature review is thorough.\n\n\n[Weakness]\n- Motivation is unclear to me. Why implementing the Stein Variational inference methods in PPL framework is preferable to non-PPL framework? And also why NumPyro is preferable to other PPLs? In the introduction, it is briefly mentioned that \"Such integration would solve one of the most prominent limitations of traditional VI: its lack of flexibility in capturing rich correlations in the approximated posterior\", which I consider as one of the motivations for using PPL framework. Still, theoretical or empirical evidence for this statement is missing and this work would be more appealing if the authors elaborate on it.\n\n- The originality is low since all the Stein variational inference algorithms implemented in NumPyro in this work are proposed by prior work and the authors do not propose any new algorithm.", "summary_of_the_paper": "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.\n", "summary_of_the_review": "My main concern is that the original contribution of this work is not significant enough for being published as a conference paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "bvInj3TaQfq", "reviewer": "Reviewer_KQPj", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "It is not clear what is the difference between EinSteinVI and the Stein VI methods included in Edward and PyMC3.\n\nTh experiments are extensive, considering different prob. models and problems. The computational time of each method is recorded comparing favorably for the implemented one.\n\nThe paper compares favorably with competing software for prob. programing languages such as PyMC3. However, it does not compare results with Edward. This questions the significance of the results.\n\nThe novelty of the paper is low since it does not describe any new method. It only consists in an implementation of an already known method to be used within NumPyro for prob. programming.\n\nMy main concern with the paper hence the lack of novelty, although some practitioners may consider the implementation interesting and useful.\n\n", "summary_of_the_paper": "  \nThis 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.\n", "summary_of_the_review": "A nice implementation of an already known method that could be used for machine learners using prob. programs.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "YB0lfLAEg4-", "reviewer": "Reviewer_uwM1", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Although I am fully supportive of having source code presented in academia, I do not think this paper brings anything of particular interest, in particular for this venue.\nAlthough the introduction to the different Stein-VI methods is nicely written and very understandable, this paper does not aim to be a review which is already done in Anastasiou 21' as the authors point out.\nWhat the paper could bring, would be novel techniques or implementations to improve speed, but unfortunately only using `vmap` to solve things seems a bit light.\nAnother major weakness is that no code is available, for a submission on a software this is unacceptable.\n\nThe positive point I found was the generalization of these different algorithms through a unified algorithm, but this alone is not enough to make the work sufficiently significant to be accepted.\n", "summary_of_the_paper": "The paper presents the implementation of a framework to work with different \"Stein VI\" methods in NumPyro.\nA quick introduction to the different algorithms is given as well as an evaluation of a series of examples.", "summary_of_the_review": "The paper does not bring anything interesting enough research-wise and such a paper should be submitted to a software venue/journal like [JOSS](https://joss.theoj.org/).\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "55H_543ScTY", "reviewer": "Reviewer_N78P", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "# Strength\n- Integration of SteinVI into numpyro seems very useful. Users can easily take advantage of the state-of-the-art SteinVI algorithms for their own Bayesian modelings. Moreover, experimental results show that it is faster than the existing implementation.\n- The concept of ELBO-within-Stein is innovative and seems interesting to be explored.\n- Extending the stein mixture method to deep Markov models is a novel application.\n\n# Weakness\n- No code example is presented. Thus I could not tell how EinSteinVI is easier or better than the standard Numpyro and PyMC3.\n- Related to the above, I was not convinced by the motivation of this work in the Introduction. The authors explained that \n> Currently Stein VI methods are available in PyMC3 (Salvatier et al., 2016) and Edward (Tran et al., 2016), neither of which include the advanced methods available in EinSteinVI. $\\ldots$ 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 traditional VI: its lack of flexibility in capturing rich correlations in the approximated posterior\n\n-  As far as I understood, these phrases correspond to the motivation of this work. But after reading this paper, I thought that why PyMC3 or naive Jax implementation is not enough to resolve the above issues and why proposed EinSteinVI resolves them differently from PyMC3 or Jax implementation. Add an explanation or show additional experimental results for that.\n\n- Some important Stein VI methods seem lacking, e.g,\n    - Function space particle optimization for bayesian neural network (https://arxiv.org/abs/1902.09754)\n    - Understanding and Accelerating Particle-Based Variational Inference (https://arxiv.org/abs/1807.01750)\n    - Learning Equivariant Energy Based Models with Equivariant Stein Variational Gradient Descent (https://arxiv.org/abs/2106.07832)\n- The presentation of the paper is poor. I think the concept of ELBO-within-Stein and its application to deep Markov models are the technical contributions of this work, but they are poorly explained. \n- No experiments to support the usefulness of EinSteinVI for Non-linear Stein VI, Matrix-valued kernel stein VI, and message passing stein VI.\n\n\n# The reason for the score\nAlthough the developed library seems useful for the community and the concept of ELBO-within-Stein seems interesting, I think the contribution is not enough experimentally and theoretically as I described in Weakness above. \n\n\n# Comments\n- The bandwidth tuning is critical for Stein VI. I recommend adding the tuning method other than the median method, e.g, the heat kernel approach introduced in Understanding and Accelerating Particle-Based Variational Inference (https://arxiv.org/abs/1807.01750)\n- I think the presentation can be improved to make this work more attractive.  The technical contribution, that is, the concept of ELBO-within-Stein and its application to deep Markov models, should be more emphasized. The authors should explain how novel they are in more detail. Also, the authors should add more experiments to support the usefulness of the concept of ELBO-within-Stein and stein mixtures in deep Markov models.", "summary_of_the_paper": "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.\n", "summary_of_the_review": "The contribution is not enough experimentally and theoretically to be published. Although the proposed framework seems beneficial for the community, it lacks a detailed comparison regarding the usability or performance compared to existing frameworks, e.g., PyMC3. Regarding the technical contribution, the extension of stein mixtures seems promising, but it lacks the explanation, numerical experiments, and theoretical justification.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper aims to integrate Stein variational inference methods into the existing probabilistic programming language NumPyro. The implemented methods include variantions of Stein variational gradient descent with different types of kernel functions, non-linear scaling of update terms, and matrix-valued kernels. The paper includes empirical results with a comparsion with existing baselines in real-world problems. Using this framework, the authors developed a new Stein mixture algorithm for deep Markov models, which shows better performance than existing methods.

Strengths:

- The paper is overall well-written and the method is clearly explained. 
- The literature review is thorough.
- Integration of SteinVI into numpyro seems useful. Users can easily take advantage of the state-of-the-art SteinVI algorithms for their own Bayesian modelings.
- Extending the stein mixture method to deep Markov models is a novel application.

Weaknesses:

- The originality is low the authors propose algorithms that are very similar to previous work and there is a lack of experiments to verify the usefulness of the proposed method, for example,
  to verify the decreased variance of the gradient estimates claimed by the authors.
- Efforts are required to illustrate why ELBO-within-Stein is preferred over the existing work.
- Some important Stein VI methods seem lacking.
- No experiments to support the usefulness of EinSteinVI for Non-linear Stein VI, Matrix-valued kernel stein VI, and message passing stein VI.

All reviewers vote for rejection. I recommend the authors to addrss the limitatoins mentioned above and improve the paper before its resubmission to another venue.
Decision: Reject