Original: E IN S TEIN VI: G ENERAL AND I NTEGRATED S TEIN
Revised: E IN S TEIN VI: G ENERAL AND I NTEGRATED S TEIN

Original: V ARIATIONAL I NFERENCE
Revised: V ARIATIONAL I NFERENCE

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: 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 ﬂexibility of non-parametric inference methods.
Revised: 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 ﬂexibility of non-parametric inference methods.

Original: 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.
Revised: 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.

Original: EinSteinVI is a lightweight composable library that integrates the latest Stein VI methods with the PPL NumPyro (Phan et al., 2019).
Revised: EinSteinVI is a lightweight composable library that integrates the latest Stein VI methods with the PPL NumPyro (Phan et al., 2019).

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Interest in Bayesian deep learning has surged due to the need for quantifying the uncertainty of predictions obtained from machine learning algorithms.
Revised: 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).

Original: The idea behind Bayesian inference is to describe observed data x using a model with latent variables z .
Revised: The idea behind Bayesian inference is to describe observed data x using a model with latent variables z .

Original: 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 )
Revised: 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 )

Original: =
Revised: =

Original: p ( x | z ) p ( z ) .
Revised: p ( x | z ) p ( z ) .

Original: We obtain the posterior by following the rules of Bayesian inference:
Revised: We obtain the posterior by following the rules of Bayesian inference:

Original: [Equation]
Revised: [Equation]

Original: where Z = (cid:82) z p ( x | z ) p ( z ) d z is the normalization constant.
Revised: where Z = (cid:82) z p ( x | z ) p ( z ) d z is the normalization constant.

Original: For most practical models, the normalizing constant lacks an analytic solution or requires an infeasible number of computations, complicating the Bayesian inference problem.
Revised: For most practical models, the normalization constant lacks an analytic solution or requires an infeasible number of computations, complicating the Bayesian inference problem.

Original: Variational Inference (VI) techniques (Blei et al., 2017; Hoffman et al., 2013; Ranganath et al., 2014) aim to approximate the posterior distribution.
Revised: Variational Inference (VI) techniques (Blei et al., 2017; Hoffman et al., 2013; Ranganath et al., 2014) aim to approximate the posterior distribution.

Original: VI poses a family of distributions over the latent variables q ( z ) ∈ Q and ﬁnd
Revised: VI poses a family of distributions over the latent variables q ( z ) ∈ Q and ﬁnds

Original: [Equation]
Revised: [Equation]

Original: where D is a divergence 1 between the variational distribution q , also called the guide, and the true posterior distribution p ( z | x ) .
Revised: where D is a divergence 1 between the variational distribution q , also called the guide, and the true posterior distribution p ( z | x ) .

Original: The Kullback-Leibler divergence is a typical choice.
Revised: The Kullback-Leibler divergence is a typical choice.

Original: VI often provides good approximations that capture uncertainty and scales to millions of data points by a suitable choice of Q and inference method.
Revised: VI often provides good approximations that capture uncertainty and scales to millions of data points by a suitable choice of Q and inference method.

Original: Stein VI is a family of VI techniques for approximate Bayesian inference based on Stein’s method (see
Revised: Stein VI is a family of VI techniques for approximate Bayesian inference based on Stein’s method (see

Original: Anastasiou et al.
Revised: Anastasiou et al.

Original: (2021) for an overview) that is gaining popularity since it combines the scalability of traditional VI with the ﬂexibility of non-parametric particle-based methods.
Revised: (2021) for an overview) that is gaining popularity since it combines the scalability of traditional VI with the ﬂexibility of non-parametric particle-based methods.

Original: 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.

Original: Stein variational gradient descent (SVGD) (Liu & Wang, 2016) is a recent Stein VI technique which uses a set of particles { z i } Ni =1 as the approximating distribution q ( z ) .
Revised: Stein variational gradient descent (SVGD) (Liu & Wang, 2016) is a recent Stein VI technique which uses a set of particles { z i } Ni =1 as the approximating distribution q ( z ) .

Original: As a particle-based method, SVGD is well suited for capturing correlations between latent variables.
Revised: As a particle-based method, SVGD is well suited for capturing correlations between latent variables.

Original: The technique preserves the scalability of traditional VI approaches while offering the ﬂexibility and modeling scope of techniques such as Markov chain Monte Carlo (MCMC).
Revised: The technique preserves the scalability of traditional VI approaches while offering the ﬂexibility and modeling scope of techniques such as Markov chain Monte Carlo (MCMC).

Original: 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 ﬂow (Liu, 2017) or in terms of the properties of kernels (Liu & Wang, 2018).
Revised: 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 ﬂow (Liu, 2017) or in terms of the properties of kernels (Liu & Wang, 2018).

Original: Many advanced inference methods based on SVGD have been recently developed, including Stein mixtures (Nalisnick & Smyth, 2017), non-linear 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).
Revised: 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).

Original: These techniques have extended the power of SVGD, allowing more ﬂexible and accurate approximations of thetrue posterior distribution.
Revised: These techniques have extended the scope of SVGD, allowing more ﬂexible and accurate approximations of the posterior distribution.

Original: While algorithmic power is growing, a distinct lack of integration with a general probabilistic programming language (PPL) framework remains.
Revised: While algorithmic power is growing, a distinct lack of integration with a general probabilistic programming language (PPL) framework remains.

Original: Such integration would solve one of the most prominent limitations of  traditional VI: its lack of ﬂexibility in capturing rich correlations in the approximated posterior.
Revised: Such integration would solve one of the most prominent limitations of PPLs with traditional VI: their lack of ﬂexibility in capturing rich correlations in the approximated posterior.

Original: This paper presents the EinSteinVI library that extends the NumPyro PPL (Bingham et al., 2019; Phan et al., 2019) with support for SVGD and its extensions mentioned above.
Revised: 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.

Revised: 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.

Revised: We validate our algorithm experimentally, showing signiﬁcant improvements in accuracy.

Original: By embedding EinSteinVI in NumPyro, the  Stein VI methods can utilize the many capabilities of NumPyro.
Revised: 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

Original: These include universal probabilistic programming  (van de Meent et al., 2018), integration with deep learning using
Revised: 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

Original:  Jax (Frostig et al., 2018), automatically generated and custom guides, and automatic optimization.
Revised: NumPyro is Jax (Frostig et al., 2018), which gives access to the powerful program optimization and parallelizability provided by the Jax compiler.

Revised: As EinSteinVI works with arbitrary guides, NumPyro is a well suited language for embedding EinSteinVI.

Revised: This is because NumPyro i) is embedded in

Revised: 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) beneﬁts computationally from Jax.

Revised: Our extensions include SVGD,

Revised: Stein mixtures formulated as ELBO-within-Stein, and the advanced methods mentioned above.

Original: Concretely, our contributions are:
Revised: Concretely, our contributions are:

Original: • EinSteinVI, a  general library that extends NumPyro and allows SVGD to work with custom guide programs based on E LBO-with in - Stein (EinSteinVI) optimization (Nalisnick & Smyth, 2017).
Revised: • E LBO-with in - 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.

Original: 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.
Revised: 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.

Original: • 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).
Revised: • 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).

Original: • A series of examples  demonstrate EinSteinVI and the synergy between different Stein VI techniques.
Revised: • A series of examples that demonstrate EinSteinVI and the synergy between different Stein VI techniques.

Original: The examples include a novel Stein-mixture version of the deep Markovmodel (SM-DMM), Bayesian neural networks for regression, and examples that explore kernels and higher-order optimization.
Revised: The examples include a novel Stein-mixture version of the deep Markov

Original: The paper proceeds as follows.
Revised: The paper proceeds as follows.

Original: We ﬁrst present the background of SVGD in the context of our integrated implementation in Section 2.
Revised: We ﬁrst present the background of SVGD in the context of our integrated implementation in Section 2.

Original: We discuss the general details of the implementation of EinSteinVI in NumPyro in Section 3. In Section 4 we present various examples using EinSteinVI, andﬁnally, we summarize our results and discuss future work in Section 5.
Revised: In Section 3 we introduce Stein mixtures and our algorithm

Original: There has been a proliferation of deep probabilistic programming languages based on tensor frameworks featuring automatic differentiation and supporting various inference techniques. Pyro (Bingham et al., 2019) is a universal PPL based on PyTorch (Paszke et al., 2019).
Revised: ELBO-within-Stein.

Original: The main mode of inference in Pyro is black-box stochastic VI (SVI) (Ranganath et al., 2014) with guides .
Revised: We discuss the general details of the implementation of EinSteinVI in Section 4.

Original: Guides are programs that approximate the posterior distribution by repeating the sample statements of the latent variables in the probabilistic model. They can contain deep neural networks allowing amortization of inference (Kingma & Welling, 2013; Gershman & Goodman, 2014). Pyro also supports various sampling algorithms like Hamiltonian Monte Carlo/NUTS (Duane et al., 1987; Neal, 2011; Hoffman & Gelman, 2014) and sample-adaptive MCMC (Zhu, 2019) which provide more accurate posterior approximations and is well suited for GPU, but lack the scalability of VI techniques. NumPyro is a version of Pyro that runs on the Jax framework (Frostig et al., 2018), which gives access to the powerful program optimization and parallelizability provided by the Jax compiler.
Revised: Next, we discuss related work in Section 5.

Original: Our library, EinSteinVI, adds SVGD to NumPyro and allows approximating the optimizable parameters in the guide by a set of particles instead of a point estimate. Other PPLs with similar functionality to EinStein are PyMC3 (Salvatier et al., 2016) , which includes Operator Variational Inference (OVI) (Ranganath et al., 2016), amortized SVGD (Wang et al., 2016;
Revised: In Section 6 we present various examples using

Original: Feng et al., 2017), and SVGD with a RBF-kernel, and Edward (Tran et al., 2016; 2019) which implements OVI in a repository separate from the main software.
Revised: EinSteinVI, and ﬁnally, we summarize our results and discuss future work in Section 7.

Original:  S TEIN V ARIATIONAL G RADIENT D ESCENT
Revised:  S TEIN V ARIATIONAL G RADIENT D ESCENT

Original: The core idea of SVGD is to perform inference by approximating the target posterior distribution p ( z | x ) by an empirical distribution q Z ( z ) =
Revised: The core idea of SVGD is to perform inference by approximating the target posterior distribution p ( z | x ) by an empirical distribution q Z ( z ) =

Original: N − 1 (cid:80) i
Revised: N − 1 (cid:80) i

Original: δ z i ( z ) based on a set of particles Z = { z i } Ni =1 .
Revised: δ z i ( z ) based on a set of particles Z = { z i } Ni =1 .

Original: Here, δ x ( y ) represents the Dirac delta measure, which is equal to 1 if x = y and 0 otherwise.
Revised: Here, δ x ( y ) represents the Dirac delta measure, which is equal to 1 if x = y and 0 otherwise.

Original: One could thus see the approximating distribution q Z ( z ) as a mixture of point estimates, each represented by a particle
Revised: One could thus see the approximating distribution q Z ( z ) as a mixture of point estimates, each represented by a particle

Original: z i ∈
Revised: z i ∈

Original: Z .
Revised: Z .

Original: The idea is to minimize the Kullback-Leibler divergence D KL ( q Z ( z ) (cid:107) p ( z | x )) between the approximation and the true posterior by iteratively updating the particles using the Stein forces:
Revised: The idea is to minimize the Kullback-Leibler divergence D KL ( q Z ( z ) (cid:107) p ( z | x )) between the approximation and the true posterior by iteratively updating the particles using the Stein forces:

Original: [Equation]
Revised: [Equation]

Original: where (cid:15) is the learning rate and S Z denotes the Stein forces.
Revised: where (cid:15) is the learning rate and S Z denotes the Stein forces.

Original: The Two Forces of SVGD Stein VI consists of two forces which work additively under the form
Revised: The Two Forces of SVGD Stein VI consists of two forces which work additively under the form

Original: S Z = S + Z + S −Z , where the attractive force is given by
Revised: S Z = S + Z + S −Z , where the attractive force is given by

Original: [Equation]
Revised: [Equation]

Original: and the repulsive force by
Revised: and the repulsive force by

Original: [Equation]
Revised: [Equation]

Original: Here k : R d × R d → R is a kernel.
Revised: Here k : R d × R d → R is a kernel.

Original: The attractive force can be seen as pushing the particles towards the modes of the true posterior distribution, smoothed by some kernel.
Revised: The attractive force can be seen as pushing the particles towards the modes of the true posterior distribution, smoothed by some kernel.

Original: For an example of a kernel, consider the radial basis function (RBF)
Revised: For an example of a kernel, consider the radial basis function (RBF)

Original: kernel k
Revised: kernel k

Original: ( z i , z j ) =
Revised: ( z i , z j ) =

Original: exp (cid:0) − 1 h (cid:107)
Revised: exp (cid:0) − 1 h (cid:107)

Original: z
Revised: z

Original: i − z j (cid:107) 22 (cid:1) with bandwidth parameter h , chosen as 1 log n med ( z ) .
Revised: i − z j (cid:107) 22 (cid:1) with bandwidth parameter h , chosen as 1 log n med ( z ) .

Original: The repulsive force moves particles away from each other, ensuring that they do not collapse to the same mode.
Revised: The repulsive force moves particles away from each other, ensuring that they do not collapse to the same mode.

Original: For example with the RBF kernel, the repulsive force becomes E 
Revised: For example, with the RBF kernel, the repulsive force becomes E z

Original: [Equation]
Revised: [Equation]

Original: ] , which has high kernel values for particles that are close, causing them to repel.
Revised: ] , which has high kernel values for particles that are close, causing them to repel.

Original: SVGD works with unnormalized distributions p as the normalization constant becomes additive in the log-posterior log p ( z i | x ) = − log Z + log p ( x | z ) + log p ( z ) and is thus not required for the calculation of the gradient.
Revised: SVGD works with unnormalized distributions p as the normalization constant becomes additive in the log-posterior log p ( z i | x ) = − log Z + log p ( x | z ) + log p ( z ) and is thus not required for the calculation of the gradient.

Original: This property is desirable as normalizing p is often computationally expensive.
Revised: This property is desirable as normalizing p is often computationally expensive.

Original: Non-linear Stein In non-linear Stein (Wang & Liu, 2019a)
Revised: Non-linear Stein In non-linear Stein (Wang & Liu, 2019a)

Original: , the repulsive force can be scaled by a factor λ , resulting in S Z = S + Z + λS −Z .
Revised: , the repulsive force can be scaled by a factor λ , resulting in S Z = S + Z + λS −Z .

Original: This approach is useful when dealing with multimodal distributions.
Revised: This approach is useful when dealing with multimodal distributions.

Original: It is also useful in cases where the repulsive force vanishes compared to the likelihood, which happens for large datasets ( X ).
Revised: It is also useful in cases where the repulsive force vanishes compared to the likelihood, which happens for large datasets ( X ).

Original: Scaling the repulsive force by a constant λ = c ( |X| ) proportional (e.g. c = 0 .
Revised: Scaling the repulsive force by a constant λ = c ( |X| ) proportional (e.g. c = 0 .

Original: 1 or c = 0 .
Revised: 1 or c = 0 .

Original: to the size of the dataset |X| addresses this issue.
Revised: to the size of the dataset |X| addresses this issue.

Original: Matrix-valued kernels
Revised: Matrix-valued kernels

Original: The choice of kernels can be extended to matrix-valued ones (Wang et al., 2019a), K : R d × R d → R d × d , in which case the Stein forces become
Revised: The choice of kernels can be extended to matrix-valued ones (Wang et al., 2019a), K : R d × R d → R d × d , in which case the Stein forces become

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: where the standalone del ∇ z i in the repulsive force represents the vector (cid:16) ∂ ∂z i, 1 , . . 
Revised: where the standalone del ∇ z j in the repulsive force represents the vector (cid:16) ∂ ∂z j, 1 , . . .

Original: , ∂ ∂z i,d [Equation]
Revised: , ∂ ∂z j,d (cid:17) .

Original: This results in
Revised: This results in

Original: [Equation]
Revised: [Equation]

Original: The advantage of matrix-valued kernels is that they allow preconditioning  using the Hessian orFisher
Revised: The advantage of matrix-valued kernels is that they allow preconditioning 2 using the Hessian or

Original:  Information matrix (second-order derivatives), which can capture local curvature and thus achieve better optima and convergence rate than standard SVGD.
Revised: Fisher Information matrix, which can capture local curvature and thus achieve better optima and convergence rate than standard SVGD.

Original: Furthermore, it is easy to capture graphical kernels (Wang et al., 2018a) 
Revised: Furthermore, it is easy to represent graphical kernels (Wang et al., 2018a) using matrix kernels,

Original:  K = diag( { K ( (cid:96) ) } (cid:96) ) where the set of variables are partitioned with each their own local kernel K ( (cid:96) ) .  
Revised: e.g. K = diag( { K ( (cid:96) ) } (cid:96) ) where the set of variables are partitioned with each their own local kernel K ( (cid:96) ) .  S TEIN M IXTURES

Revised: The Stein-mixture was proposed by Nalisnick & Smyth (2017) to resolve the representation issue for SVGD when the target distribution has high dimensionality.

Revised: For SVGD, the number of particles needed to represent a distribution adequately grows exponentially with its dimensionality.

Revised: As the computational complexity for the update rule in SVGD is quadratic in the number of particles, the exponential growth quickly becomes computationally intractable.

Original: ELBO-within-Stein Nalisnick & Smyth (2017) showed that the Stein force for a mixture approximation of the posterior, called a Stein-mixture, can be computed as
Revised: Stein-mixtures are hierarchical variational models (Ranganath et al., 2016b) such that the posterior of the second tier variational distribution is q ( z | X ) = 1 N (cid:80) Nk =1 δ ( z k ) . The guide is the joint model q ( θ , z ) =

Revised: q ( θ | z ) q ( z ) , with q ( θ | z ) the ﬁrst tier distribution and q ( z ) its prior.

Revised: The second tier posterior q ( z | X ) is optimized using SVGD, so that the marginal posterior is q ( θ ; z )

Revised: = 1 N (cid:80) Nk

Revised: =1 q ( θ | z ) , a restricted mixture.

Revised: Nalisnick & Smyth (2017) showed that the Stein force for a mixture approximation of the posterior p ( z | Z ) is given by

Original: [Equation]
Revised: [Equation]

Original: wherewe assume log ∇ z j q ( z j )  is sufﬁciently small to be dropped.
Revised: where log ∇ z j q ( z j ) acts as a regularizer on the variational parameters which they assume is sufﬁciently small to be dropped.

Original: To evaluate E z j ∼ q Z ( z ) [ p ( X, θ ) /q
Revised: To evaluate E z j ∼ q Z ( z ) [ p ( X, θ ) /q

Original: ( θ | z j )] , Nalisnick & Smyth (2017) employ differentiable non-centered parameterizations (Kingma & Welling, 2014) and importance weighted Monte Carlo gradients(Burda et al., 2015) to obtain the black-box update,
Revised: ( θ | z j )] , Nalisnick & Smyth (2017) use DNCP and importance weighted Monte Carlo gradients to obtain the black-box update,

Original: [Equation]
Revised: [Equation]

Original: where ˜ w s is an importance weight for sample ˆ θ s = q ( z j , ξ ) , ξ ∼ p 0 (i.e. the guide).
Revised: where ˜ w s is an importance weight for sample ˆ θ s = q ( z j , ξ ) , ξ ∼ p 0 (i.e. the guide).

Original: Note that (cid:80) 1 s =1 ˜ w s ∇ z j log ( p ( X, θ s ) /q

Original: ( θ s | z j )) is the gradient of the one sample evidence lower bound (ELBO). In EinStein we replace the weighted average over S samples by a single loss ( L ).
Revised: ELBO-within-Stein In ELBO-within-Stein we replace the weighted average over S samples by a single loss ( L ).

Original: We can estimate the loss by averaging over multiple samples to reduce bias.
Revised: We can reduce bias in the loss by averaging over multiple samples.

Original: Note this is different from Nalisnick & Smyth (2017) which estimates E θ l ∼ q ( θ | z j ) [ p ( X, θ l ) /q
Revised: However, we only compute one gradient of L per particle. This is different from Nalisnick & Smyth (2017) who estimate E θ l ∼ q ( θ | z j ) [ p ( X, θ l ) /q

Original: ( θ l | z j )] by averaging over gradients.
Revised: ( θ l | z j )] by averaging over gradients, therefore computes a gradient and an importance weighting for each DNCP sample.

Original: The update thus becomes
Revised: The Stein force in ELBO-within-Stein is given by

Original: [Equation]
Revised: [Equation]

Original: where L is the ELBO.
Revised: where L is the ELBO.

Original:  C OMPOSITIONAL I MPLEMENTATION USING N UM P YRO
Revised:  C OMPOSITIONAL I MPLEMENTATION USING N UM P YRO

Original: EinSteinVI integrates with the existing NumPyro API by adding the Stein VI interface, which closely mimicks NumPyro’s SVI interface.
Revised: EinSteinVI integrates with the existing NumPyro API by adding the Stein VI interface, which closely mimicks NumPyro’s SVI interface.

Original: Mimicking the SVI interface makes programs that use SVI in
Revised: Mimicking the SVI interface makes programs that use SVI in

Original: NumPyro trivial to convert to EinSteinVI.
Revised: NumPyro trivial to convert to EinSteinVI (see Figure 1).

Original: Below, we discuss the key features of EinSteinVI, which include re-initializable guides, EinSteinVI’s core algorithm, and the new kernel interface.
Revised: Below, we discuss the key features of EinSteinVI, which include re-initializable guides, EinSteinVI’s core algorithm, and the new kernel interface.

Original: 3. R E - INITIALIZABLE G UIDES
Revised: 4. R E - INITIALIZABLE G UIDES

Original: The Stein VI interface requires that the initialization is different for each parameter in an inference program.
Revised: The Stein VI interface requires that the initialization is different for each parameter in an inference program.

Original: 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.
Revised: 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.

Original: To support re-initializable guides, we provide the ReinitGuide interface, which requires implementingof 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.
Revised: 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.

Original: The WrappedGuide class provides a guide written as a function.
Revised: The WrappedGuide class provides a guide written as a function.

Original: WrappedGuide makes a callable guide re-initializable.
Revised: WrappedGuide makes a callable guide re-initializable.

Original: It works by running the provided guide multiple times and reinitializing the parameters using NumPyro’s interface as follows:
Revised: It works by running the provided guide multiple times and reinitializing the parameters using NumPyro’s interface as follows:

Original: • WrappedGuide runs the guide transforming each parameter to unconstrained space.
Revised: • WrappedGuide runs the guide transforming each parameter to unconstrained space.

Original: • 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 ] .
Revised: • 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 ] .

Original: • It saves the parameter values for each particle and the required inverse transformations to constrained space to run the model correctly.
Revised: • It saves the parameter values for each particle and the required inverse transformations to constrained space to run the model correctly.

Original: We also allow parameters without reinitialization in order to support neural network libraries like stax 2 that have their own initializers.
Revised: We also allow parameters without reinitialization in order to support neural network libraries like stax 3 that have their own initializers.

Original: The Stein VI interface will correctly wrap the guide during initialization, so from a user perspective the syntax for guides is identical to that of the SVI api.
Revised: 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.

Original: 3. S TEIN VI IN N UM P YRO
Revised: 4. S TEIN VI IN N UM P YRO

Original: The integration of Stein VI into NumPyro requires handling transformations between the parameter representation of NumPyro 3 and the vectorized Stein particles that EinSteinVI operates on.
Revised: The integration of Stein VI into NumPyro requires handling transformations between the parameter representation of NumPyro 4 and the vectorized Stein particles that EinSteinVI operates on.

Original: For this, we rely on Jax PyTrees 4 which converts back and forth between Python collections and a ﬂattened vectorized representation.
Revised: For this, we rely on Jax PyTrees 5 which converts back and forth between Python collections and a ﬂattened vectorized representation.

Original: Algorithm 1 shows the core algorithm of EinSteinVI.
Revised: Algorithm 1 shows the core algorithm of EinSteinVI.

Original: EinSteinVI updates the standard variational model parameters φ and guide parameters φ (cid:48) by averaging the loss over the Stein particles.
Revised: EinSteinVI updates the standard variational model parameters φ and guide parameters ψ by averaging the loss over the Stein particles.

Original: For the
Revised: For the

Original: Stein parameters, the process is more elaborate.
Revised: Stein parameters, the process is more elaborate.

Original: First, we convert the set of individual parameters to a monolithic vector-encoded particle using Jax PyTrees .
Revised: First, we convert the set of individual parameters to a monolithic vector-encoded particle using Jax PyTrees .

Revised: The monolithic particle represents the particles as a ﬂattened and stacked Jax array.

Original: Then we compute a kernel based on the vector-encoded Stein particle; this is delegated to the kernel interface as the computation is kerneldependent.
Revised: 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.

Original: We apply Jax ’s vmap operator (Frostig et al., 2018; Phan et al., 2019) to compute the 
Revised: 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.

Original:  Stein forces for each particle  in  a vectorized manner.
Revised: 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 compute the desired Jacobians.

Original: Note that Algorithm 1 presents how EinSteinVI works with scalar kernels and does not account for the different features presented in Section 2.
Revised: Note that Algorithm 1 presents how EinSteinVI works with scalar kernels and does not account for the different features presented in Section 2.

Original: Finally, we convert the monolithic Stein particle to their non-vectorized dictionary-based form and return the expected changes for standard- and Stein-parameters. 
Revised: Finally, we convert the monolithic Stein particle to their non-vectorized dictionary-based form and return the expected changes for standard- and Stein-parameters. 4. K ERNEL INTERFACE

Original: [Equation]
Revised: 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

Original:  ( z ) , loss L , kernel interface KI .
Revised: Input: Classical VI parameters φ and ψ , Stein parameters { θ i } i , model p φ ( z , x ) , guide q θ , ψ ( z ) , loss L , kernel interface KI .

Original: Output: Parameter changes based on classical VI ( ∆ φ , ∆ φ (cid:48) ) and Stein VI forces ( { ∆ θ i } i ).
Revised: Output: Parameter changes based on classical VI ( ∆ φ , ∆ ψ ) and Stein VI forces ( { ∆ θ i } i ).

Original: procedure E INSTEIN ( φ , φ (cid:48) , { θ i } i , p φ , q θ , φ (cid:48) ) ∆ φ ← E θ [ ∇ φ L ( p φ , q θ , φ (cid:48) )]∆ φ (cid:48) ← E θ [ ∇ φ (cid:48)
Revised: procedure E INSTEIN ( φ , ψ , { θ i } i , p φ , q θ , ψ ) ∆ φ ← E θ [ ∇ φ L ( p φ , q θ , ψ )]

Original:  L ( p φ , q θ , φ (cid:48) )] { a i } i ← PyTreeFlatten( { θ i } i ) k ← KI( { a i } i ) procedure E INSTEIN F ORCES ( a i ) (cid:46)
Revised: ∆ ψ ← E θ [ ∇ ψ L ( p φ , q θ , ψ )] { a i } i ← PyTreeFlatten( { θ i } i ) k ← KI( { a i } i ) procedure E INSTEIN F ORCES ( a i ) (cid:46)

Original: Calculate forces per particle for higher-order vmap function.
Revised: Calculate forces per particle for higher-order vmap function.

Original: θ
Revised: θ

Original: i ← P Y T REE R ESTORE ( a i ) ∆
Revised: i ← P Y T REE R ESTORE ( a i ) ∆

Original: a i ← (cid:80) a j k ( a j , a i ) ∇ a i L ( p φ , q θ i , φ (cid:48) ) + ∇ a i k ( a j , a i ) return ∆ a i end procedure { ∆ a i } i ← VMap ( { a i } i , E INSTEIN F ORCES ) { ∆ θ i } i ← P Y T REE R ESTORE ( { ∆ a i } i ) return ∆ φ , ∆ φ (cid:48) , { ∆ θ i } i end procedure Algorithm 2 https://jax.readthedocs.io
Revised: a i ← (cid:80) a j k ( a j , a i ) ∇ a i L ( p φ , q θ i , ψ ) + ∇ a i k ( a j , a i ) return ∆ a i end procedure { ∆ a i } i ← VMap ( { a i } i , E INSTEIN F ORCES ) { ∆ θ i } i ← P Y T REE R ESTORE ( { ∆ a i } i ) return ∆ φ , ∆ ψ , { ∆ θ i } i end procedure and particles, and the loss function L and returns a differentiable kernel k .

Original: [Table]

Original: en/latest

Original: [Table]

Original: jax.experimental stax.html

Original: .html

Original: 3. K ERNEL INTERFACE

Original: The Kernel interface is straightforward.

Original: 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 .

Original: All kernels are currently static, but the interface could be extended to stateful kernels, allowing conjugate gradients or quasi-Newton optimization.
Revised: All kernels are currently static, but the interface could be extended to stateful kernels, allowing conjugate gradients or quasiNewton optimization.

Original: Appendix A describes the complete list of kernels in EinSteinVI.
Revised: Appendix A gives the complete list of kernels in EinSteinVI.

Original: The implemented kernels are well justiﬁed for SVGD, but they do not consider information geometry and may have erratic behavior with the Stein-mixture method, as it does not use point estimates.
Revised: We are planning to extend EinSteinVI with probability product kernels Nalisnick & Smyth (2017), which take into account the information geometry of the ﬁrst tier variational distributions in Stein mixtures.  R ELATED W ORK 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.

Original: Nalisnick & Smyth (2017) used a probability product kernels to account for this; however, these are not currently part of EinSteinVI.
Revised: 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).

Original: E XAMPLES
Revised: 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 .

Revised: Amortized SVGD (Feng et al., 2017) trains a stochastic network to draw samples from a target distribution.

Revised: The network is iteratively adjusted so that the output changes in the direction of the

Revised: Stein variational gradient (the same gradient as used in SVGD).

Revised: In comparison, EinSteinVI transports a ﬁxed set of particles (each parameterizing a guide) to the target distribution.

Revised: Amortized SVGD is not in EinSteinVI as its extension to arbitrary guides is an open problem.

Revised: OVI optimizes operator objectives, which take functions of functions to a non-negative number.

Revised: Ranganath et al.

Revised: (2016a) include an operator objective based on the Langevin-Stein operator (Anastasiou et al., 2021).

Revised: This is the same operator used for the kernelized Stein discrepancy (Liu et al., 2016) which also underlies SVGD.

Revised: Unlike EinSteinVI, Amortized SVDG and SVGD, OPV is not a particle-based method.

Revised: [Figure]

Revised: ( ' x ' , NormalMixture(jnp

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: SVI( model,AutoNormal(model),

Revised: [Figure]

Revised: [Figure]

Revised: =

Original: We illustrate the features of EinSteinVI on two toy examples, namely 1D and 2D mixtures of
Revised: [Figure] _ELBO() ) [Figure] svi

Revised: [Figure]

Original: Gaussian’s.
Revised: iterations) [Figure]  E XAMPLES 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.

Original: We demonstrate EinSteinVI on real-world examples where we ﬁnd that EinSteinVI tends to outperform alternative methods.
Revised: We demonstrate EinSteinVI on real-world examples and show that EinSteinVI tends to outperform alternative methods.

Original: These examples include regression with Bayesian neural networks and deep Markov models.
Revised: These examples include regression with Bayesian neural networks and deep Markov models.

Original: All experimental code, notebook tutorials and the EinSteinVI package itself are available through an Anonymized URL.
Revised: All experimental code, notebook tutorials and the EinSteinVI package itself are available through an Anonymized URL.

Original: 1D Gaussian mixture To demonstrate the two modes of VI with EinSteinVI, we consider the 1D
Revised: 1D Gaussian mixture To demonstrate the two modes of VI with EinSteinVI, we consider the 1D

Original: Gaussian mixture 1 / 3 N ( − 2 , 1)+2 / 3 N (2 , 1) (see Figure 1).
Revised: Gaussian mixture 1 / 3 N ( − 2 , 1) + 2 / 3 N (2 , 1) (see Figure 1 and Figure 2).

Original: The Gaussian target mixture is bi-modal and well-suited for the nonparametric nature of SVGD and Stein mixtures.
Revised: The Gaussian target mixture is bi-modal and well-suited for the nonparametric nature of SVGD and Stein mixtures.

Original: Figure 1 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.
Revised: 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.

Original: Note the reduction in particles required to estimate the target when using Stein mixtures compared to SVGD.
Revised: Note the reduction in particles required to estimate the target when using Stein mixtures compared to SVGD.

Original: Also, note that the Stein-mixture overestimates the variance and slightly pertubates the locations.
Revised: Also, note that the Stein-mixture overestimates the variance and slightly pertubates the locations.

Original: 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.
Revised: 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.

Original: The Stein-mixture will therefore not be able to exactly capture the mixing components of a target mixture model with one particle per component.
Revised: The Stein-mixture will therefore not be able to exactly capture the mixing components of a target mixture model with one particle per component.

Original: However, with more particles the mixture can be approximated better as demonstrated using three particles.
Revised: However, with more particles the mixture can be approximated better as demonstrated using three particles.

Original: [Figure]
Revised: [Figure]

Original: [Table] [Table] Bayesian Neural Networks We compare SVGD  in EinSteinVI with the implementation of SVGD by Liu & Wang (2016) 5 without model selection, and in PyMC3 (using the Theano backend 6 ,
Revised: 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 [Table]

Original:  Al-Rfou et al.
Revised: [Table] (using the Theano backend 7 , Al-Rfou et al.

Original: (2016)), using Bayesian neural networks (BNN) for regression.
Revised: on Bayesian neural networks (BNN) for regression.

Original:  We use an RBF-kernel in 
Revised: 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 ﬁndings reﬂect the warning in the PyMC3 documentation.

Original: EinStein for fair comparison with Liu & Wang (2016) and PyMC3 .
Revised: We use an RBF-kernel in EinSteinVI for fair comparison with Liu & Wang (2016) and PyMC3 .

Revised: For non-linear Stein we determined the best repulsion factor λ ∗ by a grid search for λ ∈ { 10 − 2 , 10 − 1 , 10 1 , 10 2 } .

Revised: For SVGD PyMC3 we use a temperature 8 of ten.

Original: Like Liu & Wang (2016)
Revised: Like Liu & Wang (2016)

Original: we use a BNN with one hidden layer of size ﬁfty and RELU activation.
Revised: we use a BNN with one hidden layer of size ﬁfty and RELU activation.

Original: We put a Gamma (1 , 0 . 1) prior on the precision of the neurons and the likelihood.
Revised: We put a Gamma (1 , 0 . 1) prior on the precision of the neurons and the likelihood.

Original: For allthree versions, we use 100 particles anditerations.
Revised: For all versions, we use 100 particles and 2000 iterations.

Original: We use a mini-batch of 1000 for Year and 100 for the rest.
Revised: We use a mini-batch of 1000 for Year and 100 for the rest.

Original: All measurements are repeated twenty times and obtained on a GPU 7 for EinSteinVI and PyMC3 , except for Year with PyMC3 which was only  once.
Revised: All measurements are repeated ten times and obtained on a GPU 9 for EinSteinVI and PyMC3 , except for Year with PyMC3 which was only run once.

Original: We do not report times for Liu & Wang (2016) because only the CPU version of their code could be executed without irresolvable issues.
Revised: We do not report times for Liu & Wang (2016) because only the CPU version of their code could be executed without irresolvable issues.

Original: Table 1 shows the performance in terms of the average root mean squared error (RMSE) on the test set.
Revised: Table 1 shows the performance in terms of the average root mean squared error (RMSE) on the test set.

Original: We ﬁnd that EinSteinVI achieves signiﬁcantly better RMSE than Liu & Wang (2016) and that both systems outperform SVGD in PyMC3 .
Revised: We ﬁnd that EinSteinVI achieves signiﬁcantly better RMSE than Liu & Wang (2016) and that both systems outperform SVGD in PyMC3 .

Original: The times in Table 1 measure how long it takes to infer parameters.
Revised: The times in Table 1 measure how long it takes to infer parameters.

Original: Table 1 excludes the ﬁrst run for two reasons: i) that run will ﬁll the GPU caches, and ii)
Revised: Table 1 excludes the ﬁrst run for two reasons: i) that run will ﬁll the GPU caches, and ii)

Original: Jax will trace the programs.
Revised: Jax will trace the programs.

Original: As a result, this run is more costly than subsequent ones.
Revised: As a result, this run is more costly than subsequent ones.

Original: The times for the ﬁrst run are given in Table 2.
Revised: The times for the ﬁrst run are given in Table 2.

Original: By running EinSteinVI for more iterations, we can amortize the initial cost.
Revised: By running EinSteinVI for more iterations, we can amortize the initial cost.

Original: Star Distribution To illustrate the kernels included in EinSteinVI, we approximate the "star distribution", follow Wang et al.

Original: (2019b)

Original: (see Figure 2).

Original: The star distribution is constructed as a 2D

Original: Gaussian mixture,

Original: [Equation]

Original: where U k is a rotation matrix given by

Original: [Equation]

Original: We use ﬁfty particles and Adagrad with the learning rate that yields the best Maximum Mean

Original: Discrepancy (MMD) (Gretton et al., 2012) with a RBF kernel.

Original: To compute the MMD, we use 10, samples from the 2D gaussian mixture.

Original: Keeping these choices ﬁxed, we vary the type of kernel, using a point mass distribution as guide.

Original: [Figure]

Original: We consider the scalar ( R d → R ) linear kernel, the inverse multiquadric (IMQ) kernel, the RBF kernel and the random feature kernel; details are given in appendix A.

Original: In addition, we use a mixture kernel, which is a uniform mixture of the linear and the random feature kernels.

Original: We also use the matrix version of the RBF kernel with a constant preconditioning, using the Hessian matrix for preconditioning.

Original: 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.

Original: Stein Mixture Deep Markov Model Music generation requires a model to learn complex temporal dependencies to achieve local consistency between notes.
Revised: Stein Mixture Deep Markov Model Music generation requires a model to learn complex temporal dependencies to achieve local consistency between notes.

Original: 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.
Revised: 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.

Original: We consider a vectorized version of the DMM (Jankowiak & Karaletsos, 2019) for the generation of polyphonic music using the JSB chorales data set.
Revised: We consider a vectorized version of the DMM (Jankowiak & Karaletsos, 2019) for the generation of polyphonic music using the JSB chorales data set.

Original: The SM-DMM model consists of two feed-forward neural  networks.
Revised: The SM-DMM model consists of two feed-forward neural (FNN) networks.

Original: The Transition network takes care of the conditional dependencies between subsequent latent states in the Markov chain, while the
Revised: The Transition network takes care of the conditional dependencies between subsequent latent states in the Markov chain.

Revised: It consist of two layers with hidden dimension 200 and ReLU activation on the ﬁrst layer and sigmoid on the second.

Original:  Emitter network  produces a likelihood at each time step using the current latent state.
Revised: The Emitter network is a three layers FNN which produces a likelihood at each time step using the current latent state.

Revised: The layers have hidden dimensions 100, 100, and 88, respectively.

Original: The variational distribution is of the form (cid:81)
Revised: The variational distribution is of the form (cid:81)

Original: Nn =1 q ( z n 1: T n | f ( X 1: T n )) where the parametrized feature function f 1: T n is a  gated recurrent unit (Chung et al., 2014), with hidden dimension 400.
Revised: Nn =1 q ( z n 1: T n | f ( X 1: T n )) where the parametrized feature function f 1: T n is a one layered gated recurrent unit (Chung et al., 2014), with hidden dimension 150.

Original: In the JSB chorales dataset, each time step is an eighth of a note.
Revised: The total number of parameters in the DMM is 128,654, so that with ﬁve particles EinSteinVI is optimizing 643,270 parameters for the SM-DMM model. [Table]

Original: We run the SM-DMM using the Adam optimizer (Kingma & Ba, 2014) with a learning rate of 10 − 5 , using an RBF kernel and ﬁve Stein particles for one thousand epochs.
Revised: We train the SM-DMM using the Adam optimizer (Kingma & Ba, 2014) with a learning rate of 10 − 5 , using an RBF kernel and ﬁve Stein particles for four thousand epochs on the polyphonic music generation dataset JSB chorals (Boulanger-Lewandowski et al., 2012). We follow Krishnan et al. and report two version NLL of a) [Table]

Original: The trained SM-DMM achieves a negative log likelihood of6.67 on the  test set, comparable to the 6.85 originally reported in Krishnan et al.
Revised: , where T i is the length of the i th 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 signiﬁcant improvement using Stein mixtures.

Original: (2016) and 6.82 reported in the improved version in Jankowiak & Karaletsos (2019).
Revised: 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).

Original:  S UMMARY
Revised:  S UMMARY

Original: EinSteinVI provides the latest techniques for Stein VI as an extension to Numpyro.
Revised: EinSteinVI provides the latest techniques for Stein VI as an extension to Numpyro.

Original: Our results indicate that the library is substantially faster and more expressive than other available libraries for Stein VI.
Revised: Our results indicate that the library is substantially faster and more expressive than other available libraries for Stein VI.

Original: EinSteinVI provides a familiar and efﬁcient interface for practitioners working with the Pyro/NumPyro PPL and provides a uniﬁed code base to researchers for benchmarking new developments in Stein VI.
Revised: EinSteinVI provides a familiar and efﬁcient interface for practitioners working with the Pyro/NumPyro PPL and provides a uniﬁed code base to researchers for benchmarking new developments in Stein VI.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: A A PPENDIX
Revised: A A PPENDIX

Original: Kernel
Revised: Kernel

Original: [Table]
Revised: [Table]

Revised: Detail

Original: Type
Revised: Type

Original: Reference
Revised: Reference

Original: Radial Basis
Revised: Radial Basis

Original: Function (RBF)
Revised: Function (RBF)

Original: [Equation]
Revised: [Equation]

Original: scalar
Revised: scalar

Original: Liu & Wang (2016)
Revised: Liu & Wang (2016)

Original: [Equation]
Revised: [Equation]

Original: vector
Revised: vector

Original: Pyro
Revised: Pyro

Original: Inverse MultiQuadratic (IMQ)(
Revised: Inverse MultiQuadratic (IMQ)

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: β ∈ ( − 1 , 0) and c >scalar
Revised: β ∈ ( − 1 , 0) and c >scalar

Original: Gorham &
Revised: Gorham &

Original: Mackey (2017)
Revised: Mackey (2017)

Original: RandomFeatureExpansion
Revised: RandomFeatureExpansion

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: ) where w 0 ∼ Unif(0 , 2 π ) and w 1 ∼ N (0 , 1) scalar
Revised: where w 0 ∼ Unif(0 , 2 π ) and w 1 ∼ N (0 , 1) scalar

Original: Liu & Wang (2018)
Revised: Liu & Wang (2018)

Original: Linear x (cid:62) y
Revised: Linear x (cid:62) y

Original: +scalar
Revised: +scalar

Original: Liu & Wang (2018)
Revised: Liu & Wang (2018)

Original: Mixture (cid:80) i w i k
Revised: Mixture (cid:80) i w i k

Original: i
Revised: i

Original: ( x , y ) { k i } i individual kernels, weights w i scalar, vector, matrix
Revised: ( x , y ) { k i } i individual kernels, weights w i scalar, vector, matrix

Original: Liu & Wang (2018)
Revised: Liu & Wang (2018)

Original: Scalar-based
Revised: Scalar-based

Original: Matrix k
Revised: Matrix k

Original: ( x , y )
Revised: ( x , y )

Original: I k scalar-valued kernel matrix
Revised: I k scalar-valued kernel matrix

Original: Wang et al.
Revised: Wang et al.

Original: (2019a)
Revised: (2019a)

Original: Vector-based
Revised: Vector-based

Original: Matrix diag( k ( x , y ))
Revised: Matrix diag( k ( x , y ))

Original: k vector-valued kernel matrix
Revised: k vector-valued kernel matrix

Original: Wang et al.
Revised: Wang et al.

Original: (2019a)
Revised: (2019a)

Original: Graphical diag( { K ( (cid:96) )
Revised: Graphical diag( { K ( (cid:96) )

Original: ( x , y ) } (cid:96) )
Revised: ( x , y ) } (cid:96) )

Original: { K ( (cid:96) ) } (cid:96) matrix-valued kernels, each for a unique partition of latent variables matrix
Revised: { K ( (cid:96) ) } (cid:96) matrix-valued kernels, each for a unique partition of latent variables matrix

Original: Wang et al.
Revised: Wang et al.

Original: (2019a)
Revised: (2019a)

Original: Constant Preconditioned
Revised: Constant Preconditioned

Original: Q −
Revised: Q −

Original: [Equation]
Revised: [Equation]

Original: K
Revised: K

Original: [Equation]
Revised: [Equation]

Original: , Q 12 y ) Q −
Revised: , Q 12 y ) Q −

Original: K is an inner matrix-valued kernel and Q is a preconditioning matrix like the Hessian −∇ 2¯ z log p (¯ z | x ) or Fischer information [Equation]
Revised: K is an inner matrix-valued kernel and Q is a preconditioning matrix like the Hessian −∇ 2¯ z log p (¯ z | x ) or Fischer information − E z ∼ q Z ( z ) [ ∇ 2 z log p ( z | x )]

Original: matrices matrix
Revised: matrices matrix

Original: Wang et al.
Revised: Wang et al.

Original: (2019a)
Revised: (2019a)

Original: Anchor Point
Revised: Anchor Point

Original: Preconditioned
Revised: Preconditioned

Original: [Equation]
Revised: [Equation]

Original:  { a (cid:96) } m(cid:96) =1 is a set of anchor points, Q (cid:96) =
Revised: [Equation] { a (cid:96) } m(cid:96) =1 is a set of anchor points, Q (cid:96) =

Original: Q ( a (cid:96) ) is a preconditioning matrix for each anchor point, K Q (cid:96) is an inner kernel conditioned using Q (cid:96) , and w (cid:96) ( x ) = softmax [Equation] [Equation] ) matrix
Revised: Q ( a (cid:96) ) is a preconditioning matrix for each anchor point, K Q (cid:96) is an inner kernel conditioned using Q (cid:96) , and w (cid:96) ( x ) = softmax (cid:96) ( {N ( x | a (cid:96) (cid:48) , Q − 1 (cid:96) (cid:48) ) } (cid:96) (cid:48) ) matrix

Original: Wang et al.
Revised: Wang et al.

Original: (2019a)
Revised: (2019a)

Revised: B E XPERIMENTS

Revised: Star Distribution To illustrate the kernels included in EinSteinVI, we approximate the "star distribution", follow Wang et al.

Revised: (2019b)

Revised: (see Figure 3).

Revised: The star distribution is constructed as a 2D

Revised: Gaussian mixture,

Revised: [Equation]

Revised: where U k is a rotation matrix given by

Revised: [Equation]

Revised: We use ﬁfty particles and Adagrad with the learning rate that yields the best Maximum Mean

Revised: Discrepancy (MMD) (Gretton et al., 2012) with a RBF kernel.

Revised: To compute the MMD, we use 10,

Revised: [Figure]

Revised: [Table]

Revised: samples from the 2D gaussian mixture.

Revised: Keeping these choices ﬁxed

Revised: [Table]

Revised: we vary the type of kernel, using a point mass distribution as guide.

Revised: We consider the scalar ( R d → R ) linear kernel, the inverse multiquadric (IMQ) kernel, the RBF kernel and the random feature kernel; details are given in appendix A.

Revised: In addition, we use a mixture kernel, which is a uniform mixture of the linear and the random feature kernels.

Revised: We also use the matrix version of the RBF kernel with a constant preconditioning, using the Hessian matrix for preconditioning.

Revised: 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.

Revised: Bayesian Neural Network The amortized SVGD in PyMC3 is an experimental implementation and the documentation clearly discourages its use.

Revised: For completeness we include the performance here.

Revised: The temperature and learning rate were determined by a grid search on the Boston data set.

Revised: All measurements are repeated ﬁve times.

Revised: In Table 4 we report test RMSE and time the the UCI regression benchmark.

Revised: We did not include the Year data set as it takes over 21 hours to run.

Revised: The

Revised: RMSE for amortized SVGD is poor and highly variable across all data sets, except Naval.

Revised: However, even for Naval the RMSE is two orders of magnitude greater than SVGD in EinSteinVI and Liu &

Revised: Wang (2016).