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 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).
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 also shows EinSteinVI scales to large models with more than 125,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 (Wilson & Izmailov, 2020; Wilson, 2020).
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 normalization 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 ﬁnds
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: 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
Revised: The technique preserves the scalability of traditional

Original: VI approaches while offering the ﬂexibility and modeling scope of techniques such as Markov chain
Revised: VI approaches while offering the ﬂexibility and modeling scope of techniques such as Markov chain

Original: Monte Carlo (MCMC).
Revised: 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; 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).
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 scope of SVGD, allowing more ﬂexible and accurate approximations of the 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 PPLs with traditional
Revised: Such integration would solve one of the most prominent limitations of PPLs with traditional

Original: VI: their lack of ﬂexibility in capturing rich correlations in the approximated posterior.
Revised: VI: their lack of ﬂexibility in capturing rich correlations in the approximated posterior.

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

Original: Nalisnick & Smyth (2017) suggest resolving this by using Stein mixtures and propose an inference algorithm that uses differentiable non-centered parameterization (DNCP)
Revised: Nalisnick & Smyth (2017) suggest resolving this by using Stein mixtures and propose an inference algorithm that uses differentiable non-centered parameterization (DNCP)

Original: (Kingma & Welling, 2014) and importance weighted Monte Carlo gradients (Burda et al., 2015).
Revised: (Kingma & Welling, 2014) and importance weighted Monte Carlo gradients (Burda et al., 2015).

Original: In this article we propose a novel algorithm for inference of Stein mixtures called ELBO-within-Stein.
Revised: In this article we propose a novel algorithm for inference of Stein mixtures called ELBO-within-Stein.

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

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

Original: ELBO-within-Stein is the core algorithm of our Stein VI library called EinSteinVI library.
Revised: ELBO-within-Stein is the core algorithm of our Stein VI library called EinSteinVI library.

Original: Our library further includes SVGD as a special case of ELBO-within-Stein and all the advanced methods mentioned above.
Revised: Our library further includes SVGD as a special case of ELBO-within-Stein and all the advanced methods mentioned above.

Original: The only other PPL with Stein VI methods is PyMC3 (Salvatier et al., 2016).
Revised: The only other PPL with Stein VI methods is PyMC3 (Salvatier et al., 2016).

Original: However, PyMC3 does not include any of the advanced Stein VI methods mentioned above or
Revised: However, PyMC3 does not include any of the advanced Stein VI methods mentioned above or

Original: Stein mixture methods.
Revised: Stein mixture methods.

Original: The EinSteinVI library extends the NumPyro PPL (Bingham et al., 2019; Phan et al., 2019).
Revised: The EinSteinVI library extends the NumPyro PPL (Bingham et al., 2019; Phan et al., 2019).

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

Original: The computational backend of
Revised: The computational backend of

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

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

Original: This is because NumPyro i) is embedded in
Revised: This is because NumPyro i) is embedded in

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

Original: Our extensions include SVGD,
Revised: Our extensions include SVGD,

Original: Stein mixtures formulated as ELBO-within-Stein, and the advanced methods mentioned above.
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: • E LBO-with in - Stein (EinSteinVI).
Revised: • E LBO-with in - Stein (EinSteinVI).

Original: A novel algorithm for Stein mixtures that only requires a single gradient evaluation of the ELBO per particle.
Revised: A novel algorithm for Stein mixtures that only requires a single gradient evaluation of the ELBO per particle.

Original: • A general library extension to NumPyro, called EinSteinVI.
Revised: • A general library extension to NumPyro, called EinSteinVI.

Original: EinSteinVI allows SVGD to work with custom guide programs based on ELBO-within-Stein optimization.
Revised: 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 that 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 Markov
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: In Section 3 we introduce Stein mixtures and our algorithm
Revised: In Section 3 we introduce Stein mixtures and our algorithm

Original: ELBO-within-Stein.
Revised: ELBO-within-Stein.

Original: We discuss the general details of the implementation of EinSteinVI in Section 4.
Revised: We discuss the general details of the implementation of EinSteinVI in Section 4.

Original: Next, we discuss related work in Section 5.
Revised: Next, we discuss related work in Section 5.

Original: In Section 6 we present various examples using
Revised: In Section 6 we present various examples using

Original: EinSteinVI, and ﬁnally, we summarize our results and discuss future work in Section 7.
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 Ez
Revised: For example, with the RBF kernel, the repulsive force becomes E

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 and can be chosen by cross-validation on a subset of the data.

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 j in the repulsive force represents the vector (cid:16) ∂ ∂z j, 1 , . . .
Revised: where the standalone del ∇ z j in the repulsive force represents the vector (cid:16) ∂ ∂z j, 1 , . . .

Original: , ∂ ∂z j,d (cid:17) .
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 2 using the Hessian or
Revised: The advantage of matrix-valued kernels is that they allow preconditioning 2 using the Hessian or

Original: Fisher Information matrix, 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 represent graphical kernels (Wang et al., 2018a) using matrix kernels,
Revised: Furthermore, it is easy to represent graphical kernels (Wang et al., 2018a) using matrix kernels,

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

Original:  S TEIN M IXTURES
Revised:  S TEIN M IXTURES

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

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

Original: As the computational complexity for the update rule in SVGD is quadratic in the number of particles, the exponential growth quickly becomes computationally intractable.
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: Stein-mixtures are hierarchical variational models (Ranganath et al., 2016b) such that the posterior of the second tier variational distribution is q ( z | X )
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 )

Original: = 1 N (cid:80) Nk
Revised: = 1 N (cid:80) Nk

Original: =1 δ ( z k ) .
Revised: =1 δ ( z k ) .

Original: The guide is the joint model q ( θ , z ) =
Revised: The guide is the joint model q ( θ , z ) =

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

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

Original: = 1 N (cid:80) Nk
Revised: = 1 N (cid:80) Nk

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

Original: Nalisnick & Smyth (2017) showed that the Stein force for a mixture approximation of the posterior p ( z | Z ) is given by
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: 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.
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
Revised: To evaluate E z

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

Original: ( θ | z j )] , Nalisnick & Smyth (2017) use
Revised: ( θ | z j )] , Nalisnick & Smyth (2017) use

Original: DNCP and importance weighted Monte Carlo gradients to obtain the black-box update,
Revised: 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: ELBO-within-Stein In ELBO-within-Stein 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 in Equation (1) by a single loss ( L ).

Original: We can reducebias in the loss  by averagingover multiple samples.
Revised: If we choose L to be an f -divergence, such as the ELBO, we can reduce the loss bias by averaging multiple samples.

Original: However, we only compute one gradient of L per particle.
Revised: However, ELBO-within-Stein only computes one gradient of L per particle regardless of the number of samples we use to estimate L .

Original: This is different from Nalisnick & Smyth (2017) who estimate E θ l ∼ q ( θ | z j ) [ p ( X, θ l ) /q
Revised: 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, therefore computes a gradient and an importance weighting for each DNCP sample.
Revised: ( θ l | z j )] by averaging over gradients, therefore computes a gradient and an importance weighting for each DNCP sample rather than per particle.

Revised: The difference makes ELBO-within-Stein computationally cheaper as the complexity of adding a gradient is O ( n ) whereas adding a sample to estimate L is O (1) .

Original: The Stein force in ELBO-within-Stein is given by
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 (see Figure 1).
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: 4. 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 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.
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 3 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 follows the API of SVI.
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: 4. 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 4 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 5 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 ψ 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 .

Original: The monolithic particle represents the particles as a ﬂattened and stacked Jax array.
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
Revised: Then we compute a kernel based on the vector-encoded

Original: Stein particle; this is delegated to the kernel interface as the computation is kernel-dependent.
Revised: 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 Stein forces for each particle in a vectorized manner.
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: This is done in unconstrained space so the Stein force must the corrected by the Jacobian of the bijection between constrained and unconstrained space.
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.

Original: Doing this directly on the Jax on the monolithic particle incurs a massive memory overhead in the adjoint.
Revised: Doing this directly on the Jax on the monolithic particle incurs a massive memory overhead in the adjoint.

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

Original: The operation is embarrassingly parallel and so we again use a vmap operator with a nested tree_map to compute the desired Jacobians.
Revised: 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.

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

Original: Input: Classical VI parameters φ and ψ , Stein parameters { θ i } i , model p φ ( z , x ) , guide q θ , ψ ( 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 ( ∆ φ , ∆ ψ ) and Stein VI forces ( { ∆ θ i } i ).
Revised: Output: Parameter changes based on classical VI ( ∆ φ , ∆ ψ ) and Stein VI forces ( { ∆ θ i } i ).

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

Original: ∆ ψ ← E θ [ ∇ ψ L ( p φ , q θ , ψ )] { 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 , ψ ) + ∇ 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 .
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 4. K ERNEL 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 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 quasiNewton 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 gives the complete list of kernels in EinSteinVI.
Revised: Appendix A gives the complete list of kernels in EinSteinVI.

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

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: 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.
Revised: 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: However, only PyMC3 (Salvatier et al., 2016) includes inference using SteinVI.
Revised: However, only PyMC3 (Salvatier et al., 2016) includes inference using SteinVI.

Original: In PyMC3 the inference techniques that use Stein’s method include SVGD with a scalar RBF-kernel, amortized SVGD (Wang et al., 2016;
Revised: In PyMC3 the inference techniques that use Stein’s method include SVGD with a scalar RBF-kernel, amortized SVGD (Wang et al., 2016;

Original: Feng et al., 2017), and Operator Variational Inference (OVI) (Ranganath et al., 2016a).
Revised: Feng et al., 2017), and Operator Variational Inference (OVI) (Ranganath et al., 2016a).

Original: SVGD in PyMC3 manipulates particles to directly capture the target distribution (Liu & Wang, 2016).
Revised: SVGD in PyMC3 manipulates particles to directly capture the target distribution (Liu & Wang, 2016).

Original: In EinSteinVI, SVGD is a special case of Stein mixtures where the guides are point mass distributions.
Revised: In EinSteinVI, SVGD is a special case of Stein mixtures where the guides are point mass distributions.

Original: Because we have guide programs, EinSteinVI allows arbitrary computations to transform random variables in the guide.
Revised: Because we have guide programs, EinSteinVI allows arbitrary computations to transform random variables in the guide.

Original: This is not possible with SVGD in PyMC3 .
Revised: This is not possible with SVGD in PyMC3 .

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

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

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

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

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

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

Original: Ranganath et al.
Revised: Ranganath et al.

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

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

Original: Unlike EinSteinVI, Amortized SVDG and SVGD, OPV is not a particle-based method.[Figure] ( ' x ' , NormalMixture(jnp [Figure]
Revised: Unlike EinSteinVI, Amortized SVDG and SVGD, OPV is not a particle-based method.

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: SVI( model,AutoNormal(model),

Original: [Figure]

Original: [Figure]

Original: =

Original: [Figure]

Original: _ELBO() )

Original: [Figure]

Original: svi

Original: [Figure]

Original: iterations)

Original: [Figure]

Original:  E XAMPLES
Revised:  E XAMPLES

Original: 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.
Revised: 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 and show 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 and Figure 2).
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 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.
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: [Table]

Revised: ± 0 .

Revised: 596 ± 0 .

Revised: 769 ± 0 .0 .

Revised: 604 ± 0 .5 .

Revised: 614 s ± 0 .

Revised: 052 s 1 m 7 .

Revised: 767 s ± 0 .

Revised: 286 s Yacht 0 .

Revised: 531 ± 0 .

Revised: 246 ± 0 .

Revised: 341 ± 3 .0 .

Revised: 94 ± 0 .5 .

Revised: 519 s ± 0 .

Original: 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
Revised: 07 s 1 m 4 . 748 s ± 0 . 392 s Year 8 . 662 ± 0 .

Original: [Table]
Revised: 701 ± 0 .135 .8 . 895 ± 0 .1 m 37 . 569 s ± 0 . 07 s 4 h 19 m 36 s

Original: [Table]
Revised: [Figure]

Original:  (using the Theano backend 7 , Al-Rfou et al.
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 (using the Theano backend 7 , Al-Rfou et al.

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

Original: The PyMC3 documentation clearly states amortized SVGD is experimental and is not suggested to be used.
Revised: The PyMC3 documentation clearly states amortized SVGD is experimental and is not suggested to be used.

Original: We include its experimental results in Appendix B for completeness and note that our ﬁndings reﬂect the warning in the PyMC3 documentation.
Revised: We include its experimental results in Appendix B for completeness and note that our ﬁndings reﬂect the warning in the PyMC3 documentation.

Original: We use an RBF-kernel in EinSteinVI 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 .

Original: 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 non-linear Stein we determined the best repulsion factor λ ∗ by a grid search for λ ∈ { 10 − 2 , 10 − 1 , 10 1 , 10 2 } .

Original: For SVGD PyMC3 we use a temperature 8 of ten.
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
Revised: we use a BNN with one hidden layer of size ﬁfty and

Original: RELU activation.
Revised: 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 all versions, we use 100 particles and 2000 iterations.
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 ten times and obtained on a GPU 9 for EinSteinVI and
Revised: All measurements are repeated ten times and obtained on a GPU 9 for EinSteinVI and

Original: PyMC3 , except for Year with PyMC3 which was only run once.
Revised: 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: 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 [Table] [Table] 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 (FNN) 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.
Revised: The Transition network takes care of the conditional dependencies between subsequent latent states in the Markov chain.

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

Original: The Emitter network is a three layers FNN which 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.

Original: The layers have hidden dimensions 100, 100, and 88, respectively.
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 one layered gated recurrent unit (Chung et al., 2014), with hidden dimension 150.
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: 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.
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.

Original: [Table] 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).
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).

Original: We follow Krishnan et al.
Revised: We follow Krishnan et al.

Original: and report two version NLL of a) 
Revised: and report two version NLL of a) (cid:80) Ni

Original: [Table]
Revised: =1 − p ( x i | θ ) (cid:80) Ni =1 T i and b) 1 N (cid:80) Ni

Original:  , where T i is the length of the i th sequence.
Revised: =1 − p ( x i | θ ) T i , where T i is the length of the i th sequence.

Original: In Table 3 we report NLL(a) and NLL(b) on a held-out test set of JSB.
Revised: In Table 3 we report NLL(a) and NLL(b) on a held-out test set of JSB.

Original: Compared to baseline methods SM-DMM achieves a signiﬁcant improvement using Stein mixtures.
Revised: Compared to baseline methods SM-DMM achieves a signiﬁcant improvement using Stein mixtures.

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

Original: ELBO of 45.10 (higher is better).
Revised: 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]

Original: Detail
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]
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]
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 − E z ∼ q Z ( z ) [ ∇ 2 z log p ( z | x )]
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] [Equation]
Revised: [Equation] [Equation]

Original: { a (cid:96) } m(cid:96) =1 is a set of anchor points, Q (cid:96) =
Revised: { 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 (cid:96) ( {N ( x | a (cid:96) (cid:48) , Q − 1 (cid:96) (cid:48) ) } (cid:96) (cid:48) ) 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)

Original: B E XPERIMENTS
Revised: B E XPERIMENTS

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

Original: (2019b)
Revised: (2019b)

Original: (see Figure 3).
Revised: (see Figure 3).

Original: The star distribution is constructed as a 2D
Revised: The star distribution is constructed as a 2D

Original: Gaussian mixture,
Revised: Gaussian mixture,

Original: [Equation]
Revised: [Equation]

Original: where U k is a rotation matrix given by
Revised: where U k is a rotation matrix given by

Original: [Equation]
Revised: [Equation]

Original: We use ﬁfty particles and Adagrad with the learning rate that yields the best Maximum Mean
Revised: 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.
Revised: Discrepancy (MMD) (Gretton et al., 2012) with a RBF kernel.

Original: To compute the MMD, we use 10,
Revised: To compute the MMD, we use 10,

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: samples from the 2D gaussian mixture.
Revised: samples from the 2D gaussian mixture.

Original: Keeping these choices ﬁxed
Revised: Keeping these choices ﬁxed

Original: [Table]
Revised: [Table]

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

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

Original: In addition, we use a mixture kernel, which is a uniform mixture of the linear and the random feature kernels.
Revised: 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.
Revised: 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.
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.

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

Original: For completeness we include the performance here.
Revised: For completeness we include the performance here.

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

Original: All measurements are repeated ﬁve times.
Revised: All measurements are repeated ﬁve times.

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

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

Original: The
Revised: The

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

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

Original: Wang (2016).
Revised: Wang (2016).