Original: ELBO- ING S TEIN M IXTURES
Revised: ELBO- ING S TEIN M IXTURES

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 gradient descent (SVGD) (Liu & Wang, 2016) is a particle-based technique for Bayesian inference.
Revised: Stein variational gradient descent (SVGD) (Liu & Wang, 2016) is a particle-based technique for Bayesian inference.

Original: SVGD has recently gained popularity because it combines the ability of variational inference to handle tall data with the modeling power of non-parametric inference.
Revised: SVGD has recently gained popularity because it combines the ability of variational inference to handle tall data with the modeling power of non-parametric inference.

Original: Unfortunately, the number of particles required to represent a model adequately grows exponentially with the dimensionality of the model.
Revised: Unfortunately, variance estimation grows inversely with the dimensionality of the model, meaning more particles are required to represent high-dimensional models adequately.

Original: Stein mixtures (Nalisnick & Smyth, 2017) alleviate the exponential growth in particles by letting each particle parameterize a distribution.
Revised: Stein mixtures (Nalisnick & Smyth, 2017) alleviate the exponential growth in particles by letting each particle parameterize a distribution.

Original: However, the inference algorithm proposed by Nalisnick & Smyth (2017) can be numerically unstable.
Revised: However, the inference algorithm proposed by Nalisnick & Smyth (2017) can be numerically unstable.

Original: We show that their algorithm corresponds to inference with the Rényi α -divergence for α = 0 and that using other values for α can lead to  more stable inference.
Revised: We show that their algorithm corresponds to inference with the Rényi α -divergence for α = 0 and that using other values for α can lead to a more stable inference.

Original: We empirically study the performance of Stein mixtures inferred with different α values on various real-world problems, demonstrating significantly and consistently improved results when using α = 1 , which corresponds to using the evidence lower bound (ELBO).
Revised: We empirically study the performance of Stein mixtures inferred with different α values on various real-world problems, demonstrating significantly and consistently improved results when using α = 1 , which corresponds to using the evidence lower bound (ELBO).

Original: We call this instance of our algorithm ELBOwithin-Stein.
Revised: We call this instance of our algorithm ELBO-within-Stein.

Original: An easy-to-use version of the inference algorithm (for arbitrary α ∈ R ) is available in the deep probabilistic programming language
Revised: An easy-to-use version of the inference algorithm (for arbitrary α ∈ R ) is available in the deep probabilistic programming language

Original: NumPyro (Phan et al., 2019).
Revised: NumPyro (Phan et al., 2019).

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: The ability of Bayesian deep learning to quantify the uncertainty of predictions by deep models is causing a surge of interest in using these techniques (Izmailov et al., 2021).
Revised: The ability of Bayesian deep learning to quantify the uncertainty of predictions by deep models is causing a surge of interest in using these techniques (Izmailov et al., 2021).

Original: Bayesian inference aims to describe i.i.d.
Revised: Bayesian inference aims to describe i.i.d.

Original: data D = { x i } ni
Revised: data D = { x i } ni

Original: =
Revised: =

Original: 1 using a model with latent a variable z .
Revised: 1 using a model with latent a variable z .

Original: Bayesian inference does this by computing a posterior distribution p ( z |D ) over the latent variable given a model describing the joint distribution p ( z , D )
Revised: Bayesian inference does this by computing a posterior distribution p ( z |D ) over the latent variable given a model describing the joint distribution p ( z , D )

Original: = p ( D| z ) p ( z ) .
Revised: = p ( D| z ) p ( z ) .

Original: We obtain the posterior by following Bayes’ theorem,
Revised: We obtain the posterior by following Bayes’ theorem,

Original: [Equation]
Revised: [Equation]

Original: where p ( D ) = (cid:82) z (cid:81) ni =1 p ( x i | z ) p ( z )
Revised: where p ( D ) = (cid:82) z (cid:81) ni =1 p ( x i | z ) p ( z )

Original: d z is the normalization constant.
Revised: d z is the normalization constant.

Original: For most practical models, the normalization constant lacks an analytic solution or poses a computability problem, complicating the
Revised: For most practical models, the normalization constant lacks an analytic solution or poses a computability problem, complicating the

Original: Bayesian inference problem.
Revised: Bayesian inference problem.

Original: Stein variational gradient descent (SVGD) (Liu & Wang, 2016) is a recent technique for Bayesian inference that uses a set of particles
Revised: Stein variational gradient descent (SVGD) (Liu & Wang, 2016) is a recent technique for Bayesian inference that uses a set of particles

Original: Z = { z i } Ni =1 to approximate the posterior p ( z |D ) .
Revised: Z = { z i } Ni =1 to approximate the posterior p ( z |D ) .

Original: The idea behind SVGD is to iteratively transport Z according to a force field S Z , called the Stein force.
Revised: The idea behind SVGD is to iteratively transport Z according to a force field S Z , called the Stein force.

Original: The
Revised: The

Original: Stein force is given by
Revised: Stein force is given by

Original: [Equation]
Revised: [Equation]

Original: where k ( · , · ) is a reproducing kernel (Berlinet & Thomas-Agnan, 2011), q Z = N − 1 (cid:80) i
Revised: where k ( · , · ) is a reproducing kernel (Berlinet & Thomas-Agnan, 2011), q Z = N − 1 (cid:80) i

Original: δ z i is the empirical measure on the set of particles Z , δ x ( y ) represents the Dirac delta measure, which is equal to 1 if x = y and 0 otherwise, and ∇ z j log p ( z j |D ) is the gradient of the posterior with respect to the j -th particle.
Revised: δ z i is the empirical measure on the set of particles Z , δ x ( y ) represents the Dirac delta measure, which is equal to 1 if x = y and 0 otherwise, and ∇ z j log p ( z j |D ) is the gradient of the posterior with respect to the j -th particle.

Original: The technique is scalable to tall data (i.e. datasets with many data points) and offers the flexibility and scope of techniques such as Markov chain Monte Carlo (MCMC).
Revised: The technique is scalable to tall data (i.e. datasets with many data points) and offers the flexibility and scope of techniques such as Markov chain Monte Carlo (MCMC).

Original: SVGD is good at capturing multi-modality (Liu & Wang, 2016; Wang & Liu, 2019), and has useful theoretical interpretations such as a set of particles following a gradient flow (Liu, 2017) or in terms of the properties of kernels (Liu & Wang, 2018).
Revised: SVGD is good at capturing multi-modality (Liu & Wang, 2016; Wang & Liu, 2019), and has useful theoretical interpretations such as a set of particles following a gradient flow (Liu, 2017) or in terms of the properties of kernels (Liu & Wang, 2018).

Original: The main problem is that SVGD suffers from the curse of dimensionality: the number of particles required to represent a posterior distribution adequately is exponential in its dimensionality.
Revised: The main problem is that SVGD suffers from the curse of dimensionality: variance estimation scales inversely with dimensionality (Ba et al., 2021).

Original: Nalisnick  Smyth (2017) suggest resolving this by using a Stein mixture (SM).
Revised: Nalisnick & Smyth (2017) suggest resolving this by using a Stein mixture (SM).

Original: SMs lift each particle to the parameters of a variational distribution q , also called a guide.
Revised: SMs lift each particle to the parameters of a variational distribution q , also called a guide.

Original: The idea is that each guide in the Stein mixture represents the density of multiple particles in SVGD, thereby reducing the number of particles needed to represent a posterior.
Revised: The idea is that each guide in the Stein mixture represents the density of multiple particles in SVGD, thereby reducing the number of particles needed to represent a posterior.

Original: The Nalisnick & Smyth algorithm introduces guides by replacing each posterior gradient ∇ z j log p ( z j |D ) in Equation (1) with the corresponding gradient of the marginal log-variational likelihood given by
Revised: The Nalisnick & Smyth algorithm introduces guides by replacing each posterior gradient ∇ z j log p ( z j |D ) in Equation (1) with the corresponding gradient of the marginal log-variational likelihood given by

Original: [Equation]
Revised: [Equation]

Original: Here, we denote the particles by Φ = { ϕϕϕ j } Ni =1 instead of Z = { z i } Ni =1 to emphasize they parameterize guide components q ( z | ϕϕϕ i , D ) .
Revised: Here, we denote the particles by Φ = { ϕϕϕ j } Ni =1 instead of Z = { z i } Ni =1 to emphasize they parameterize guide components q ( z | ϕϕϕ i , D ) .

Original: The change in gradient corresponds to minimizing
Revised: The change in gradient corresponds to minimizing

Original: D KL [ q Φ ( ϕϕϕ )
Revised: D KL [ q Φ ( ϕϕϕ )

Original: ∥ p ( ϕϕϕ |D )] rather than D KL [ q Z ( z ) ∥ p ( z |D )] , as in SVGD.
Revised: ∥ p ( ϕϕϕ |D )] rather than D KL [ q Z ( z ) ∥ p ( z |D )] , as in SVGD.

Original: Note that the line between the model p and guide q becomes blurred, as p ( D| ϕϕϕ ) is random in both data ( D ), as is usually the case, but also in the guide hyper-parameters ϕϕϕ (Ranganath et al., 2016; Nalisnick & Smyth, 2017).
Revised: Note that the line between the model p and guide q becomes blurred, as p ( D| ϕϕϕ ) is random in both data ( D ), as is usually the case, but also in the guide hyper-parameters ϕϕϕ (Ranganath et al., 2016; Nalisnick & Smyth, 2017).

Original: To distinguish the two we subsequently refer to p ( D ) as the evidence and p ( D| ϕϕϕ ) as the hierarchical likelihood.
Revised: To distinguish the two we subsequently refer to p ( D ) as the evidence and p ( D| ϕϕϕ ) as the hierarchical likelihood.

Original: The Stein force using the log hierarchical likelihood, which we call the hierarchical Stein force S H Φ , becomes
Revised: The Stein force using the log hierarchical likelihood, which we call the hierarchical Stein force S H Φ , becomes

Original: [Equation]
Revised: [Equation]

Original: where q Φ is an empirical measure analogous to q Z .
Revised: where q Φ is an empirical measure analogous to q Z .

Original: Inference converges (i.e. reaches a fixed point) when S H Φ ( ϕϕϕ i ) = 0 for all particles, meaning all gradients in S H Φ must cancel (i.e. sum to zero).
Revised: Inference converges (i.e. reaches a fixed point) when S H Φ ( ϕϕϕ i ) = 0 for all particles, meaning all gradients in S H Φ must cancel (i.e. sum to zero).

Original: However, computing the gradient of the log-variational likelihood requires numerical estimation as analytical solutions do not exist for most models.
Revised: However, computing the gradient of the log-variational likelihood requires numerical estimation as analytical solutions do not exist for most models.

Original: For such estimates, we cannot expect the gradients to cancel as the Stein force will compensate  thegradient estimation error  by a counter force.
Revised: So we cannot expect the inference converges with noisy gradient estimations as the Stein force will compensate for the error in the gradient by a counterforce in the next iteration.

Original: Therefore, SMs require good (i.e. low relative variance) gradient approximations, as otherwise, the particles will fluctuate around a fixed point without reaching it.
Revised: Therefore, SMs require good (i.e. low relative variance) gradient approximations; otherwise, the particles will fluctuate around a fixed point without reaching it.

Original: We demonstrate that replacing the log hierarchical likelihood with the evidence lower bound (ELBO) can provide better (lower relative variance) gradient approximations.
Revised: We demonstrate that replacing the log hierarchical likelihood with the evidence lower bound (ELBO) can provide better (lower relative variance) gradient approximations.

Original: We call the new algorithm ELBO-within-Stein.
Revised: We call the new algorithm ELBO-within-Stein (EoS).

Original: We connect ELBO-within-Stein with the algorithm proposed by Nalisnick & Smyth (2017) in terms of computing the gradient of different orders of the variational Rényi (VR) bound (Van Erven & Harremos, 2014).
Revised: We connect EoS with the algorithm proposed by Nalisnick & Smyth (2017) in terms of computing the gradient of different orders of the variational Rényi (VR) bound (Van Erven & Harremos, 2014).

Original: Similarly to the ELBO, the VR bound is a lower bound of the evidence, p ( D ) , also called the normalization constant, and is given by
Revised: Similarly to the ELBO, the VR bound is a lower bound of the evidence, p ( D ) , also called the normalization constant, and is given by

Original: [Equation]
Revised: [Equation]

Original: where α ≥ 0 is known as its order 1 .
Revised: where α ≥ 0 is known as its order 1 .

Original: Understanding  inference of SMs in terms of the VR bound yields insight into the behavior of the two algorithms, as we can now understand α as controlling the variance of each component of the guide.
Revised: Understanding the inference of SMs in terms of the VR bound yields insight into the behavior of the two algorithms, as we can now understand α as controlling the variance of each component of the guide.

Original: Furthermore, presuming accurate gradient approximation for all (viable) values of α , the connection leads to a family of inference algorithms indexed by the
Revised: Furthermore, presuming accurate gradient approximation for all (viable) values of α , the connection leads to a family of inference algorithms indexed by the

Original: VR bound order.
Revised: VR bound order.

Original: After reviewing SVGD, the Rényi divergence, and the signal-to-noise ratio (SN-ratio) that is used to estimate the relative variance in Section 2, we make the following contributions:
Revised: After reviewing SVGD, the Rényi divergence, and the signal-to-noise ratio (SN-ratio) that is used to estimate the relative variance in Section 2, we make the following contributions:

Original: • We demonstrate that inaccurate gradient estimates can lead to issues with convergence for SMs.
Revised: • We demonstrate that inaccurate gradient estimates can lead to issues with convergence for SMs.

Original: • We introduce a new family of inference algorithms for SMs indexed by the parameter α .
Revised: • We introduce a new family of inference algorithms for SMs indexed by the parameter α .

Original: The family results from connecting inference with SMs to the Rényi α -divergence and includes the inference algorithm by Nalisnick & Smyth (2017) as  special case for α = 0 .
Revised: The family results from connecting inference with SMs to the Rényi α -divergence and includes the inference algorithm by Nalisnick & Smyth (2017) as a special case for α = 0 .

Original: • Unlike previous work, our algorithm allows for investigating a range of values for α for a model of interest.
Revised: • Unlike previous work, our algorithm allows for investigating a range of values for α for a model of interest.

Original: This allows us to investigate the convergence stability for different α ’s by measuring the SN-ratio.
Revised: This allows us to investigate the convergence stability for different α ’s by measuring the SN-ratio.

Original: We find that α = 1 is optimal for models with local latent variables, resulting in better SN-ratios than all other α values.
Revised: We find that α = 1 is optimal for models with a latent variable for each data point (local latent variables), resulting in better SN-ratios than all other α values.

Original: For models with global latent variables, using α = 0 . 5 (corresponding to the Hellinger distance) is on par with Nalisnick & Smyth (2017)’s algorithm (which corresponds to α = 0 ).
Revised: For models where all datapoints share a latent variable (global latent variables), using α = 0 . 5 (corresponding to the Hellinger distance) is on par with Nalisnick & Smyth (2017)’s algorithm (which corresponds to α = 0 ).

Original: Other values for α result in worse SN-ratios.
Revised: Other values for α result in worse SN-ratios.

Original: • We evaluate our inference algorithm for different values of α on Bayesian neural networks (BNNs) and variational autoencoders (VAEs), showing that the α that results in the highest performance varies depending on both model and data set.
Revised: • We evaluate our inference algorithm for different values of α on Bayesian neural networks (BNNs) and variational autoencoders (VAEs), showing that the α that results in the highest performance varies depending on both model and data set.

Original: • We describe a black-box inference algorithm for our proposed family of inference algorithms and provide a software library, called EinSteinVI , in NumPyro.
Revised: • We describe a black-box inference algorithm for our proposed family of inference algorithms and provide a software library, called EinSteinVI , in NumPyro.

Original: In Section 4 we discuss related work.
Revised: In Section 4 we discuss related work.

Original: We benchmark our algorithm in Section 5.
Revised: We benchmark our algorithm in Section 5.

Original: Finally, we summarize our results in Section 6.
Revised: Finally, we summarize our results in Section 6.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: Let z be a latent variable of interest taking values in a space Z ⊆ R d (up to a diffeomorphism) and
Revised: Let z be a latent variable of interest taking values in a space Z ⊆ R d (up to a diffeomorphism) and

Original: D = { x i } n ∈ N be a set of i.i.d.
Revised: D = { x i } n ∈ N be a set of i.i.d.

Original: observations.
Revised: observations.

Original: For many models, exact Bayesian inference is computationally impracticable due to the cost of evaluating the evidence p ( D ) .
Revised: For many models, exact Bayesian inference is computationally impracticable due to the cost of evaluating the evidence p ( D ) .

Original: Therefore, practitioners turn to tractable approximate variational inference (VI).
Revised: Therefore, practitioners turn to tractable approximate variational inference (VI).

Original: VI aims to bring a computationally cheap variational distribution q ( z |D ) close to the model posterior.
Revised: VI aims to bring a computationally cheap variational distribution q ( z |D ) close to the model posterior.

Original: Typically, we measure closeness by the Kullback-Leibler divergence ( D KL ), i.e.
Revised: Typically, we measure closeness by the Kullback-Leibler divergence ( D KL ), i.e.

Original: D KL [ q ( z |D )
Revised: D KL [ q ( z |D )

Original: ∥ p ( z |D )] .
Revised: ∥ p ( z |D )] .

Original: However, we generally avoid directly evaluating D KL [ q ( z |D )
Revised: However, we generally avoid directly evaluating D KL [ q ( z |D )

Original: ∥ p ( z |D )] as this requires evaluating the evidence, p ( D ) .
Revised: ∥ p ( z |D )] as this requires evaluating the evidence, p ( D ) .

Original: We will concern ourselves with two types of VI.
Revised: We will concern ourselves with two types of VI.

Original: The first type of VI searches for a parameterizatio ψ ∗ of q in a family of distributions Q that minimizes the divergence to the posterior.
Revised: The first type of VI searches for a parameterization ψ ∗ of q in a family of distributions Q that minimizes the divergence to the posterior.

Original: When the divergence is measured by D KL , this type of VI is made tractable by maximizing the evidence lower bound (ELBO), that is
Revised: When the divergence is measured by D KL , this type of VI is made tractable by maximizing the evidence lower bound (ELBO), that is

Original: [Equation]
Revised: [Equation]

Original: The second type of VI we consider relies on particle-based methods and is the focus of this article.
Revised: The second type of VI we consider relies on particle-based methods and is the focus of this article.

Original: This type of VI relies on transporting a finite set of particles such that their empirical measure is close to the posterior.
Revised: This type of VI relies on transporting a finite set of particles such that their empirical measure is close to the posterior.

Original: We will discuss this method in detail below.
Revised: We will discuss this method in detail below.

Original: 2. S TEIN VARIATIONAL GRADIENT DESCENT
Revised: 2. S TEIN VARIATIONAL GRADIENT DESCENT

Original: The core idea of SVGD is to perform inference by approximating the target posterior distribution p ( z |D ) 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 |D ) 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 , where
Revised: δ z i ( z ) based on a set of particles Z , where

Original: Z = { z i } Ni =1 .
Revised: Z = { z i } Ni =1 .

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

Original: Z .
Revised: Z .

Original: The SVGD algorithm minimizes the
Revised: The SVGD algorithm minimizes the

Original: Kullback-Leibler divergence D KL [ q Z ( z ) ∥
Revised: Kullback-Leibler divergence D KL [ q Z ( z ) ∥

Original: p ( z |D )] between the approximated and the true posterior by iteratively updating the particles using the following expression:
Revised: p ( z |D )] between the approximated and the true posterior by iteratively updating the particles using the following expression:

Original: [Equation]
Revised: [Equation]

Original: where ϵ is the learning rate and S Z denotes the Stein force.
Revised: where ϵ is the learning rate and S Z denotes the Stein force.

Original: The two forces of SVGD
Revised: The two forces of SVGD

Original: The Stein force S Z consists of two underlying forces that work additively, with S Z = S + Z + S −Z .
Revised: The Stein force S Z consists of two underlying forces that work additively, with S Z = S + Z + S −Z .

Original: The attractive force is given by
Revised: 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: The repulsive force stops particles with high kernel values from collapsing onto each other.
Revised: The repulsive force stops particles with high kernel values from collapsing onto each other.

Original: In Appendix C we demonstrate the repulsive behaviour for a radial basis function (RBF) kernel.
Revised: In Appendix C, we demonstrate the repulsive behavior for a radial basis function (RBF) kernel.

Original: The computational cost of S Z is quadratic in the size of Z , i.e. O ( N 2 ) , which makes SVGD computationally burdensome for high-dimensional posteriors.
Revised: The computational cost of S Z is quadratic in the size of Z , i.e. O ( N 2 ) , which makes SVGD computationally burdensome for high-dimensional posteriors.

Original: For a particle method such as SVGD, the number of particles required to represent a posterior distribution adequately is exponential in its dimensionality.
Revised: For a particle method such as SVGD, the number of particles required to represent a posterior distribution adequately is exponential in its dimensionality.

Original: SVGD does not require normalized distributions p as the normalization constant becomes additive in  the log-posterior, i.e. log p ( z |D )
Revised: SVGD suffers from the curse of dimensionality (Ba et al., 2021), which results in variance collapse (where variance is underestimated). Wang et al. (2018) demonstrates the problem with a simple factorized Gaussian, suggesting the (RBF) kernel introduces global statistical dependence driving the need for particles up for accurate representation. Ba et al.

Original: = log p ( D| z ) + log p ( z ) − log p ( D ) , and therefore cancels when computing the gradient with respect to z .
Revised: (2021) demonstrate that the collapse is due to the deterministic update of attractive force and that re-sampling the particles at each iteration eliminates the issue. Note that algorithmic change by Ba et al. (2021) is not generally tractable, so it does not suffice as a practical solution.

Original: 2. R ÉNYI DIVERGENCE AND THE VARIATIONAL R ÉNYI BOUND
Revised: 2. R ÉNYI DIVERGENCE AND THE VARIATIONAL R ÉNYI BOUND

Original: The Rényi divergence (Rényi, 1961) is a family of divergences between distributions p and q indexed by the order parameter { α | α ∈ R + / { 0 , 1 } , | D α | < ∞} .
Revised: The Rényi divergence (Rényi, 1961) is a family of divergences between distributions p and q indexed by the order parameter { α | α ∈ R + / { 0 , 1 } , | D α | < ∞} .

Original: The divergence is given by
Revised: The divergence is given by

Original: [Equation]
Revised: [Equation]

Original: The Rényi divergence can be extended to α ∈ { 0 , 1 , ∞} by continuity.
Revised: The Rényi divergence can be extended to α ∈ { 0 , 1 , ∞} by continuity.

Original: In addition, if we allow for
Revised: In addition, if we allow for

Original: D α
Revised: D α

Original: [ p ∥ q ] ≤ 0 , the order can be further extended to α ∈ R (Van Erven & Harremos, 2014).
Revised: [ p ∥ q ] ≤ 0 , the order can be further extended to α ∈ R (Van Erven & Harremos, 2014).

Original: Several orders correspond to known divergences (see (Van Erven & Harremos, 2014) and (Li & Turner, 2016) for an overview).
Revised: Several orders correspond to known divergences (see (Van Erven & Harremos, 2014) and (Li & Turner, 2016) for an overview).

Original: In particular, α = 1 corresponds to D KL .
Revised: In particular, α = 1 corresponds to D KL .

Original: Analogous to the use of the D KL in the ELBO, D α leads to a variational Rényi bound (Li & Turner, 2016) which, when formulated as used with SMs, is given by
Revised: Analogous to the use of the D KL in the ELBO, D α leads to a variational Rényi bound (Li & Turner, 2016) which, when formulated as used with SMs, is given by

Original: [Equation]
Revised: [Equation]

Original: Note that model hyper-parameters ( ϕ ) in the variational posterior, q ( z |D ,ϕϕϕ ) , are lifted to random variables when doing inference with SMs.
Revised: Note that model hyper-parameters ( ϕ ) in the variational posterior, q ( z |D ,ϕϕϕ ) , are lifted to random variables when doing inference with SMs.

Original: See Appendix A for the derivation of Equation (6).
Revised: See Appendix A for the derivation of Equation (6).

Original: Assuming reparameterization of z is possible, we can approximate the gradient Λ( ϕϕϕ ) of Equation (6) using Monte Carlo integration by 
Revised: Assuming reparameterization of z is possible, we can approximate the gradient Λ( ϕϕϕ ) of Equation (6) using Monte Carlo integration by [Equation] Z [Equation]

Revised: ∼ q

Original: [Equation]
Revised: [Equation]

Original:  where K ∈ N number of draws used to compute the VR bound and
Revised: z |D [Equation] [Equation] 7) where K ∈ N number of draws used to compute the VR bound and

Original: [Equation]
Revised: [Equation]

Original: We provide the derivation in Appendix A.
Revised: We provide the derivation in Appendix A.

Original: 2. T HE SIGNAL - TO - NOISE RATIO
Revised: 2. T HE SIGNAL - TO - NOISE RATIO

Original: The signal-to-noise (SN) ratio was introduced by Rainforth et al.
Revised: The signal-to-noise (SN) ratio was introduced by Rainforth et al.

Original: (2018) to study the effect of tighter variational bounds on gradient estimation.
Revised: (2018) to study the effect of tighter variational bounds on gradient estimation.

Original: The SNR is given by
Revised: The SNR is given by

Original: [Equation]
Revised: [Equation]

Original:  where σ [ · ] is the standard deviation, M, K ∈ N are the number of Monte Carlo draws, and ∆ αM,K ( ϕϕϕ ) derives from rewriting Equation (7) in the form 
Revised: (cid:12)(cid:12)(cid:12)(cid:12)(cid:12) (cid:12) , (9) where σ [ · ] is the standard deviation, M, K ∈ N are the number of Monte Carlo draws, and ∆ αM,K ( ϕϕϕ ) derives from rewriting Equation (7) in the form [Equation]

Original: [Equation]
Revised: [Equation]

Original: Here, we separate tightening the bound (by increasing K ) from reducing the noise in the gradient estimation (by increasing M ).
Revised: Here, we separate tightening the bound (by increasing K ) from reducing the noise in the gradient estimation (by increasing M ).

Original: If the rate at which the expected gradient decreases is faster than the rate of decrease of the variance, the gradient estimates worsen as K increases.
Revised: If the rate at which the expected gradient decreases is faster than the rate of decrease of the variance, the gradient estimates worsen as K increases.

Original: The counter-intuitive implication is that a tighter bound can worsen the gradient estimation.
Revised: The counter-intuitive implication is that a tighter bound can worsen the gradient estimation.

Original: 2. T HE S TEIN MIXTURE
Revised: 2. T HE S TEIN MIXTURE

Original: Variational inference with SMs (Nalisnick & Smyth, 2017) approximates the target posterior distribution p ( z |D ) by letting the Stein particles Φ = { ϕϕϕ i } Ni =1 parameterize guide programs, q ( z | ϕϕϕ
Revised: Variational inference with SMs (Nalisnick & Smyth, 2017) approximates the target posterior distribution p ( z |D ) by letting the Stein particles Φ = { ϕϕϕ i } Ni =1 parameterize guide programs, q ( z | ϕϕϕ

Original: i , D ) .
Revised: i , D ) .

Original: A
Revised: A

Original: SM yields a mixture marginal variational posterior, p ( z |D )
Revised: SM yields a mixture marginal variational posterior, p ( z |D )

Original: ≈ 1 / | Φ | (cid:80) ϕϕϕ ∈
Revised: ≈ 1 / | Φ | (cid:80) ϕϕϕ ∈

Original: Φ q ( z | ϕϕϕ, D ) , from which it takes its name.
Revised: Φ q ( z | ϕϕϕ, D ) , from which it takes its name.

Original: Formally, SM is a hierarchical variational model (HVM) (Ranganath et al., 2016) with an empirical measure of particles q Φ (defined in the same way as q Z ) as its variational posterior, a uniform variational prior, and variational likelihood E q ( z |D ,ϕϕϕ )
Revised: Formally, SM is a hierarchical variational model (HVM) (Ranganath et al., 2016) with an empirical measure of particles q Φ (defined in the same way as q Z ) as its variational posterior, a uniform variational prior, and variational likelihood E q ( z |D ,ϕϕϕ )

Original: [ p ( D , z | ϕϕϕ ) / q ( z |D ,ϕϕϕ ) ] .
Revised: [ p ( D , z | ϕϕϕ ) / q ( z |D ,ϕϕϕ ) ] .

Original: Similarly to
Revised: Similarly to

Original: SVGD, SM minimizes D KL ( q ( ϕϕϕ )
Revised: SVGD, SM minimizes D KL ( q ( ϕϕϕ )

Original: ∥ p ( ϕϕϕ |D )) by iteratively transporting the particles according to the following expression
Revised: ∥ p ( ϕϕϕ |D )) by iteratively transporting the particles according to the following expression

Original: [Equation]
Revised: [Equation]

Original: where ϵ ≥ 0 is the learning rate and S H Φ is the hierarchical Stein force.
Revised: where ϵ ≥ 0 is the learning rate and S H Φ is the hierarchical Stein force.

Original: The attractive force of SM Like SVGD, SM also makes use of two additive forces, S H Φ =
Revised: The attractive force of SM Like SVGD, SM also makes use of two additive forces, S H Φ =

Original: S H +Φ + S − Φ .
Revised: S H +Φ + S − Φ .

Original: The repulsive force S − Φ is the same as in SVGD, given by Equation (5).
Revised: The repulsive force S − Φ is the same as in SVGD, given by Equation (5).

Original: The attractive force is given by
Revised: The attractive force is given by

Original: [Equation]
Revised: [Equation]

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

Original: From the construction of SVGD, we require that the kernel has the reproducing property, so the kernel is dense in the space of continuous functions.
Revised: From the construction of SVGD, we require that the kernel has the reproducing property, so the kernel is dense in the space of continuous functions.

Original: If we choose
Revised: If we choose

Original: Gaussian guides, the expected likelihood (EL) kernel (Jebara et al., 2004) is a natural choice because it accounts for the geometry of q ( z |D ,ϕϕϕ j ) and reduces to the RBF kernel for fixed variance, which is a reproducing kernel.
Revised: Gaussian guides, the expected likelihood (EL) kernel (Jebara et al., 2004) is a natural choice because it accounts for the geometry of q ( z |D ,ϕϕϕ j ) and reduces to the RBF kernel for fixed variance, which is a reproducing kernel.

Original: The EL kernel is given by
Revised: The EL kernel is given by

Original: [Equation]
Revised: [Equation]

Original: where L 2 is an inner product and k is a positive definite kernel.
Revised: where L 2 is an inner product and k is a positive definite kernel.

Original:  α - INDEXED S TEIN MIXTURES INFERENCE AND ELBO- WITHIN -S TEIN
Revised:  α - INDEXED S TEIN MIXTURES INFERENCE AND ELBO- WITHIN -S TEIN

Original: To see the connection between the hierarchical Stein force given in Equation (3) and the Rényi divergence, consider the gradient of the log hierarchical likelihood (that occurs in S H +Φ ) and
Revised: To see the connection between the hierarchical Stein force given in Equation (3) and the Rényi divergence, consider the gradient of the log hierarchical likelihood (that occurs in S H +Φ ) and

Original: the
Revised: the

Original: VR bound given in Equation (6) for α = 0 .
Revised: VR bound given in Equation (6) for α = 0 .

Original: Presuming the support of the variational likelihood q ( z | ϕϕϕ ) is a subset of the support of the prior of p , supp ( q ( z | ϕϕϕ ))
Revised: Presuming the support of the variational likelihood q ( z | ϕϕϕ ) is a subset of the support of the prior of p , supp ( q ( z | ϕϕϕ ))

Original: ⊆ supp ( p ( z )) , the gradient of the log hierarchical likelihood is given by
Revised: ⊆ supp ( p ( z )) , the gradient of the log hierarchical likelihood is given by

Original: [Equation]
Revised: [Equation]

Original: From Equation (11), we see that the gradient of the log marginal likelihood is exactly the gradient of the difference between the score log p ( D ) , on the one hand, and the Rényi divergence (at α = 0 ) between the variational posterior q ( z | ϕϕϕ ) and the model posterior p ( z ∥D ) , on the other hand.
Revised: From Equation (11), we see that the gradient of the log marginal likelihood is exactly the gradient of the difference between the score log p ( D ) , on the one hand, and the Rényi divergence (at α = 0 ) between the variational posterior q ( z | ϕϕϕ ) and the model posterior p ( z ∥D ) , on the other hand.

Original: Thus,
Revised: Thus,

Original: Equation (11) shows that the attractive hierarchical force ( S H +Φ ) pushes the components of the variational posterior, q ( z |D ,ϕϕϕ ) , towards the model posterior, p ( z |D ) , see Appendix D for details.
Revised: Equation (11) shows that the attractive hierarchical force ( S H +Φ ) pushes the components of the variational posterior, q ( z |D ,ϕϕϕ ) , towards the model posterior, p ( z |D ) , see Appendix D for details.

Original: The equivalence in Equation (11) suggeststhere is a whole class of hierarchical attractive forces indexed by the order α of the VR bound.
Revised: The equivalence in Equation (11) suggests a whole class of hierarchical attractive forces indexed by the order α of the VR bound.

Original: Note that choosing α ̸ = 0 means we loose the interpretation of the attractive force as moving the particles towards the nearest peak of the conditional evidence.
Revised: Note that choosing α ̸ = 0 means we loose the interpretation of the attractive force as moving the particles towards the nearest peak of the conditional evidence.

Revised: [Figure]

Original: Assuming our marginal variational posterior q ( z | ϕϕϕ ) is reparameterizable, we can approximate the attractive force for any α ≥ 0 as
Revised: Assuming our marginal variational posterior q ( z | ϕϕϕ ) is reparameterizable, we can approximate the attractive force for any α ≥ 0 as

Original: [Equation]
Revised: [Equation]

Original: (12) where Λ K ( ϕ ) is given by Equation (7).
Revised: where Λ K ( ϕ ) is given by Equation (7).

Original: We call inference with Equation (12) α -indexed Stein mixture inference.
Revised: We call inference with Equation (12) α -indexed Stein mixture inference.

Original: There are two special cases of α that are worth highlighting.
Revised: There are two special cases of α that are worth highlighting.

Original: The first is α = 1 / 2 , for which[Figure] the Rényi divergence corresponds to the Hellinger divergence (Van Erven & Harremos, 2014)(Li & Turner, 2016).
Revised: The first is α = 1 / 2 , for which the Rényi divergence corresponds to the Hellinger divergence (Van Erven & Harremos, 2014)(Li & Turner, 2016).

Original: The second is α = 1 , which corresponds to the D KL -divergence.
Revised: The second is α = 1 , corresponding to the D KL -divergence.

Original: In this case, the VR bound recovers the ELBO.
Revised: In this case, the VR bound recovers the ELBO.

Original: We call this instance of our α -indexed SM inference algorithm
Revised: We call this instance of our α -indexed SM inference algorithm

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

Original: In Appendix E we show that we can also recover the α = 1 case directly by applying Jensen’s inequality to the conditional evidence.
Revised: In Appendix E we show that we can also recover the α = 1 case directly by applying Jensen’s inequality to the conditional evidence.

Original: 3. I NVESTIGATING THE SIGNAL - TO - NOISE RATIO
Revised: 3. I NVESTIGATING THE SIGNAL - TO - NOISE RATIO

Original: Estimation of a Stein mixture converges when S H Φ = 0 , which means that the repulsive and attractive forces must be equal and oppositein order for them to cancel.
Revised: Estimation of a Stein mixture converges when S H Φ = 0 , which means that the repulsive and attractive forces must be equal and opposite for them to cancel.

Original: Hence, convergence requires ∆ αM,K ( ϕϕϕ ) and ∇ ϕϕϕ 1 k ( ϕϕϕ 1 ,ϕϕϕ 2 ) to be accurate.
Revised: Hence, convergence requires ∆ αM,K ( ϕϕϕ ) and ∇ ϕϕϕ 1 k ( ϕϕϕ 1 ,ϕϕϕ 2 ) to be accurate.

Original: In Figure 1 we demonstrate the effect of inaccurate gradient approximations.
Revised: In Figure 1 we demonstrate the effect of inaccurate gradient approximations.

Original: To study the sensitivity of gradient approximations to the choice of α , we measure the SN-ratio (see Equation (9)) of the VR bound gradients (see Equation (6)).
Revised: To study the sensitivity of gradient approximations to the choice of α , we measure the SN-ratio (see Equation (9)) of the VR bound gradients (see Equation (6)).

Original: We simulate data { x i } ni =1 from a simple latent variable model given by N ( D| z , I d ) N ( z | µµµ, I d ) , where µµµ ∈ R d is unknown and I d is the d -dimensional identity matrix.
Revised: We simulate data { x i } ni =1 from a simple latent variable model given by N ( D| z , I d ) N ( z | µµµ, I d ) , where µµµ ∈ R d is unknown and I d is the d -dimensional identity matrix.

Original: To approximate its posterior we use a Stein mixture of the form 1 / 2 ( N ( ϕϕϕ 1 , 3 / 2 I d )+ N ( ϕϕϕ 2 , 3 / 2 I d )) , and an expected likelihood kernel.
Revised: To approximate its posterior we use a Stein mixture of the form [Equation] [Equation] I d )

Revised: [Equation]

Revised: 2 I d )) , and an expected likelihood kernel.

Original: We choose a (computationally convenient) fixed variance such that the Stein mixture cannot exactly recover the posterior.
Revised: We choose a (computationally convenient) fixed variance such that the Stein mixture cannot exactly recover the posterior.

Original: We can see this as the posterior is unimodal, which is only the case for the Stein mixture if | ϕϕϕ 1 − ϕϕϕ 2 | < 3 (Behboodian, 1970), but in this interval the variance of the Stein mixture will greater than or equal to 3 / 2 .
Revised: We can see this as the posterior is unimodal which is only the case for the Stein mixture if | ϕϕϕ 1 − ϕϕϕ 2 | < 3 (Behboodian, 1970), but in this interval the variance of the Stein mixture will greater than or equal to 3 / 2 .

Original: With the expected likelihood kernel we can analytically characterize all fixed points for the Stein particles as−∇
Revised: With the expected likelihood kernel we can analytically characterize all fixed points for the Stein particles as

Original: [Equation]
Revised: [Equation]

Original: See the Appendix B for the derivation.
Revised: See the Appendix B for the derivation.

Original: To measure the effect of gradient approximation on the system we use Equation (10) to estimate the gradients.
Revised: To measure the effect of gradient approximation on the system we use Equation (10) to estimate the gradients.

Original: To conduct our experiment, we sample the location µµµ from a 20-dimensional standard Gaussian and use this µµµ to simulate n = 64 data points D .
Revised: To conduct our experiment, we sample the location µµµ from a 20-dimensional standard Gaussian and use this µµµ to simulate n = 64 data points D .

Original: We then approximate the gradients at a random point close to a fixed point
Revised: We then approximate the gradients at a random point close to a fixed point

Original: [Equation]
Revised: [Equation]

Original: where D is the data average, µ + n D / n +1 is the posterior mean and ϵ offsets each dimension by a
Revised: where D is the data average, µ + n D / n +1 is the posterior mean and ϵ offsets each dimension by a

Original: Gaussian with mean zero and variance 0.01.
Revised: Gaussian with mean zero and variance 0.01.

Original: Figure 2a illustrates the experimental setup in two dimensions.

Original: For legibility in Figure 2a, we do not include the perturbation (i.e. added ϵ noise) on ϕϕϕ 1 to the visualization.

Original: The contours correspond to the

Original: [Figure]
Revised: [Figure]

Original: Figure 2b and Figure 2c show the convergence of the SN-ratio (see Equation (9)) as we tighten the
Revised: Figure 2b and Figure 2c show the convergence of the SN-ratio (see Equation (9)) as we tighten the

Original: VR bound by increasing either K or M .
Revised: VR bound by increasing either K or M .

Original: We only show the SN-ratio for the first particle ( ϕϕϕ 1 ) as the behaviour for the second particle is the same.
Revised: We only show the SN-ratio for the first particle ( ϕϕϕ 1 ) as the behaviour for the second particle is the same.

Original: For the ELBO ( α = 1 ), we fix K = 1 and increase M to reduce the gradient approximation variance, while for the rest ( α ∈ { 0 , 0 .
Revised: For the ELBO ( α = 1 ), we fix K = 1 and increase M to reduce the gradient approximation variance, while for the rest ( α ∈ { 0 , 0 .

Original: , we fix M =and increase K to tighten the VR bound.
Revised: , we fix M =and increase K to tighten the VR bound.

Original: In Figure 2b, there is a latent variable for each data point.
Revised: In Figure 2b, there is a latent variable for each data point.

Original: Note how the SN-ratio only improves with tightness for α = 1 (green line).
Revised: Note how the SN-ratio only improves with tightness for α = 1 (green line).

Original: In Figure 2c all data points share a latent variable.
Revised: In Figure 2c all data points share a latent variable.

Original:  exact posterior.
Revised: Figure 2a illustrates the experimental setup in two dimensions. For legibility in Figure 2a, we do not include the perturbation (i.e. added ϵ noise) on ϕϕϕ 1 to the visualization. The contours correspond to the exact posterior.

Original: As the particles are placed equidistant from the posterior mean (marked with a blue cross), in this setting the Stein forces are zero.
Revised: As the particles are placed equidistant from the posterior mean (marked with a blue cross), in this setting the Stein forces are zero.

Original: As we would expect (see blue arrows in Figure 2a) the gradient estimations of ∇ ϕϕϕ ∆ αM,K are equal and opposite for the two particles.
Revised: As we would expect (see blue arrows in Figure 2a) the gradient estimations of ∇ ϕϕϕ ∆ αM,K are equal and opposite for the two particles.

Original: For α ̸ = 1 we fix M = 1 and vary K , while for α = 1 we fix K = 1 and vary M .
Revised: For α ̸ = 1 we fix M = 1 and vary K , while for α = 1 we fix K = 1 and vary M .

Original: We do not need to consider α ̸ = 

Revised: We do not need to consider α ̸ =

Original: when K = 1 as the associated gradient scaling ( 1 − α ) cancels in the SN-ratio.
Revised: when K = 1 as the associated gradient scaling ( 1 − α ) cancels in the SN-ratio.

Original: We empirically estimate the SN-ratio by estimating the expectation and standard deviation from 10 ,gradient samples.
Revised: We empirically estimate the SN-ratio by estimating the expectation and standard deviation from 10 ,gradient samples.

Original: In Figure 2b we show a local variate of the model where there is a latent variable z
Revised: In Figure 2b we show a local variate of the model where there is a latent variable z

Original: i for each x i ∈ D .
Revised: i for each x i ∈ D .

Original: We see that for α ̸ = 1 the SN-ratio does not depend on the particular choice of α and that the growth in SN-ratio is superseded by that of α = 1 .
Revised: We see that for α ̸ = 1 the SN-ratio does not depend on the particular choice of α and that the growth in SN-ratio is superseded by that of α = 1 .

Original: This means there is little to no benefit in increasing K beyond K = 1 , for which we recover the ELBO gradient when the guide is reparameterizable (see Section 2).
Revised: This means there is little to no benefit in increasing K beyond K = 1 , for which we recover the ELBO gradient when the guide is reparameterizable (see Section 2).

Original: In Figure 2c we evaluate a global latent variable variate of the model, so that there is one z for all datapoints in D .
Revised: In Figure 2c we evaluate a global latent variable variate of the model, so that there is one z for all datapoints in D .

Original: As with the local version, we fix either M or
Revised: As with the local version, we fix either M or

Original: K .
Revised: K .

Original: For this model, we see that α = 0 and α = 1 / 2 achieve the highest SN-ratio.
Revised: For this model, we see that α = 0 and α = 1 / 2 achieve the highest SN-ratio.

Original: The result aligns with the BNN example, where α = 1 is not dominating in performance over the α ∈ { 0 , 1 / 2 } on all datasets.
Revised: The result aligns with the BNN example, where α = 1 is not dominating in performance over the α ∈ { 0 , 1 / 2 } on all datasets.

Original: High precision is desirable to avoid fluctuation at convergence.
Revised: High precision is desirable to avoid fluctuation at convergence.

Original: From the above results, we see that α = 0 is not necessarily the best choice for precise (high SN-ratio) gradient estimation.
Revised: From the above results, we see that α = 0 is not necessarily the best choice for precise (high SN-ratio) gradient estimation.

Original: In particular, for local latent variable models, α = 1 is a better choice, and for that, for global latent models, α = 1 / 2 is on par with α = 0 in our experiment.
Revised: In particular, for local latent variable models, α = 1 is a better choice, and for that, for global latent models, α = 1 / 2 is on par with α = 0 in our experiment.

Original: 3. B LACK - BOX INFERENCE FOR ELBO- WITHIN -S TEIN
Revised: 3. B LACK - BOX INFERENCE FOR ELBO- WITHIN -S TEIN

Original: We provide a mini-batch version of ELBO-within-Stein, called EinSteinVI, in NumPyro.
Revised: We provide a mini-batch version of ELBO-within-Stein, called EinSteinVI, in NumPyro.

Original: To compute the VR bound exactly requires all the datapoints, that is, we cannot represent the bound as a pointwise expectation, except for α = 1 .
Revised: To compute the VR bound exactly requires all the datapoints, that is, we cannot represent the bound as a pointwise expectation, except for α = 1 .

Original: Therefore, in order to make EinSteinVI scalable to tall data, we provide an approximation of the VR bound which replaces the likelihood by p
Revised: Therefore, in order to make EinSteinVI scalable to tall data, we provide an approximation of the VR bound which replaces the likelihood by p

Original: I ( D| z ,ϕϕϕ ) =
Revised: I ( D| z ,ϕϕϕ ) =

Original: (cid:81) i ∈I p ( x i | z ,ϕϕϕ ) |D| / |I| , where I is a subset of a permutation of the data indices.
Revised: (cid:81) i ∈I p ( x i | z ,ϕϕϕ ) |D| / |I| , where I is a subset of a permutation of the data indices.

Original: The approximate attractive force S α +Φ ( ϕϕϕ ) used in EinSteinVI is given by
Revised: The approximate attractive force S α +Φ ( ϕϕϕ ) used in EinSteinVI is given by

Original: [Equation]
Revised: [Equation]

Original:  which recovers the exact VR bound when |I| = |D| 
Revised: 13) which recovers the exact VR bound when |I| = |D| .

Original: We describe the NumPyro integration and provide example programs in Appendix I.
Revised: We describe the NumPyro integration and provide example programs in Appendix I.

Original: [Table]
Revised: [Table]

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

Original: Nalisnick & Smyth (2017) first suggested Stein mixtures as an alternative to HVMs (Ranganath et al., 2016).
Revised: Nalisnick & Smyth (2017) first suggested Stein mixtures as an alternative to HVMs (Ranganath et al., 2016).

Original: Using SVGD allows Stein mixtures to side-step the need of HVMs for an auxiliary distribution to keep the bound (learning objective) tight.
Revised: Using SVGD allows Stein mixtures to side-step the need of HVMs for an auxiliary distribution to keep the bound (learning objective) tight.

Original: This is an improvement as the effect of the auxiliary distribution on the approximation is implicit and therefore hard to understand, whereas with Stein mixtures the choice of kernel controls the tightness and we have theoretical understanding of kernels (Wang et al., 2019; Gorham & Mackey, 2017; Liu & Wang, 2018).
Revised: This is an improvement as the effect of the auxiliary distribution on the approximation is implicit and therefore hard to understand, whereas with Stein mixtures the choice of kernel controls the tightness and we have theoretical understanding of kernels (Wang et al., 2019; Gorham & Mackey, 2017; Liu & Wang, 2018).

Original: Mixture approximations have a long history of work (Jaakkola & Jordan, 1998; Bishop et al., 1997; Gershman et al., 2012; Miller et al., 2017) focusing on approximating or lower-bounding the intractable mixture ELBO.
Revised: Mixture approximations have a long history of work (Jaakkola & Jordan, 1998; Bishop et al., 1997; Gershman et al., 2012; Miller et al., 2017) focusing on approximating or lower-bounding the intractable mixture ELBO.

Original: Van Erven & Harremos (2014) unifies a number of variational techniques by considering them as optimizing different orders of the VR bound.
Revised: Van Erven & Harremos (2014) unifies a number of variational techniques by considering them as optimizing different orders of the VR bound.

Original: They further demonstrate that two different variants of mini-batch training with the VR bound recover Stochastic EP (Li et al., 2015) and Black-box α (HernandezLobato et al., 2016), respectively.
Revised: They further demonstrate that two different variants of mini-batch training with the VR bound recover Stochastic EP (Li et al., 2015) and Black-box α (HernandezLobato et al., 2016), respectively.

Original: The Rényi divergence has been studied in other forms under the name α -divergence (Amari, 2012; Tsallis, 1988).
Revised: The Rényi divergence has been studied in other forms under the name α -divergence (Amari, 2012; Tsallis, 1988).

Original: Hernandez-Lobato et al.
Revised: Hernandez-Lobato et al.

Original: (2016) introduced a blackbox algorithm for variational inference based on the α -divergence using automatic differentiation.
Revised: (2016) introduced a blackbox algorithm for variational inference based on the α -divergence using automatic differentiation.

Original: Unlike our algorithm, their algorithm is not for HVMs.
Revised: Unlike our algorithm, their algorithm is not for HVMs.

Original: Rainforth et al.
Revised: Rainforth et al.

Original: (2018) demonstrated that for
Revised: (2018) demonstrated that for

Original: VAEs the gradient estimation degrades for multi-sample approximations when using the importance weighted variational autoencoder (IWAE) bound (Burda et al., 2015).
Revised: VAEs the gradient estimation degrades for multi-sample approximations when using the importance weighted variational autoencoder (IWAE) bound (Burda et al., 2015).

Original: Furthermore, Rainforth et al.
Revised: Furthermore, Rainforth et al.

Original: (2018) showed that this is not the case when using the ELBO.
Revised: (2018) showed that this is not the case when using the ELBO.

Original: Rainforth et al.
Revised: Rainforth et al.

Original: (2018) differs from our work in that the VAEs estimated are with a point mass guide, as their inference algorithm is not for HVMs.
Revised: (2018) differs from our work in that the VAEs estimated are with a point mass guide, as their inference algorithm is not for HVMs.

Original: Le et al.
Revised: Le et al.

Original: (2020) investigates the deterioration experimentally, providing evidence for it on several real world tasks.
Revised: (2020) investigates the deterioration experimentally, providing evidence for it on several real world tasks.

Original: Tucker et al.
Revised: Tucker et al.

Original: (2018) show that by double reparameterizing the gradient estimator, they can eliminate the degrading SN-ratio for multi-sample estimation of the
Revised: (2018) show that by double reparameterizing the gradient estimator, they can eliminate the degrading SN-ratio for multi-sample estimation of the

Original: IWAE gradient, among others.
Revised: IWAE gradient, among others.

Original:  E XAMPLES
Revised:  E XAMPLES

Original: We evaluate α -indexed Stein mixture inference by inferring Bayesian neural networks (BNN) and variational autoencoders (VAE) on standard datasets.
Revised: We evaluate α -indexed Stein mixture inference by inferring Bayesian neural networks (BNN) and variational autoencoders (VAE) on standard datasets.

Original: We use the BNNs for regression on the UCI regression benchmark (the same as Hernández-Lobato & Adams (2015)) and VAE for unsupervised learning on MNIST (Salakhutdinov & Murray, 2008; LeCun et al., 1998) and OMNIGLOT (Lake et al., 2013).
Revised: We use the BNNs for regression on the UCI regression benchmark (the same as Hernández-Lobato & Adams (2015)) and VAE for unsupervised learning on MNIST (Salakhutdinov & Murray, 2008; LeCun et al., 1998) and OMNIGLOT (Lake et al., 2013).

Original: Bayesian neural networks For brevity, we present BNNs for the subset of the UCI regression benchmark detailed in Table 1.
Revised: Bayesian neural networks For brevity, we present BNNs for the subset of the UCI regression benchmark detailed in Appendix G. All datasets use real-valued features.

Original: All datasets use real valued features.
Revised: We use a 90-10 split for training and test datasets.

Original: We compare α -indexed SM inference for α ∈ { 0 , 0 . 5 , 1 } on BNN regression with a factorized
Revised: We compare α -indexed SM inference for α ∈ { 0 , 0 . 5 , 1 } on BNN regression.

Original:  Gaussian guides and a EL kernel. In Appendix G we consider point massguides, with similar results.
Revised: We test with two guides: factorized Gaussian guides with an EL kernel and point mass

Revised: (Dirac delta) guides with an RBF kernel.

Original: Like Liu & Wang (2016), we use a BNN with one hidden layer of size fifty and a RELU activation.
Revised: Like Liu & Wang (2016), we use a BNN with one hidden layer of size fifty and a 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: We use five particles for all experiments.
Revised: We use five particles for all experiments.

Original: Due to the size differences of the datasets we run the Year dataset for 40 epochs, the Protein dataset for 100 epochs and 500 epochs for the rest.
Revised: We run all datasets for 35, [Table] [Table]

Original: We use a subsample size of 32 with the Adam optimizer (Kingma & Ba, 2014) and a step size of 0.001.
Revised: epochs with a subsample size of 32, the Adam optimizer (Kingma & Ba, 2014) and a step size of 0.002.

Original: All measurements are repeated five times and obtained on a GPU 2 .
Revised: All measurements are repeated three times and obtained on a GPU 2 .

Original: [Table]
Revised: We compare against the SVGD implantation from Liu et al. with 20 particles, mean field variational Bayes (MFVI) with a factorized Gaussian guide (Hoffman et al., 2013) and Laplace approximation. For the latter two we inference engines from NumPyro (Phan et al., 2019). Table 1 shows the root mean squared error (RMSE) on test sets.

Original: Figure 3 shows the performance of ELBO-within-Stein for α ∈ { 0 . 0 , 0 . 5 , 1 . 0 } witha factorized
Revised: EoS with delta Guides out performance baselines and other α -orders, except on Boston. SM and Hell perform similarly with factorized

Original: Gaussian prior. Performance is measured in terms of the average root mean squared error (RMSE) and log-likelihood (LL) on test sets.
Revised: Gaussian guides, which aligns with our SN-ratio experiment that shows the gradient approximations are similar for these two cases.

Revised: Note the Stein mixtures use only five particles, whereas SVGD uses twenty.

Original: We find that ELBO-within-Stein with α = 1 achievesa better average log-likelihood for all datasets and a better average RMSE for all but the Boston dataset.
Revised: Table 2 gives the log-likelihood on the same test sets. EoS achieves better average log-likelihood for all datasets with factorized Gaussian guides than other α -orders.

Original: We see that α = 0 and α = .
Revised: We see that α =and α = .

Original: 5 perform similarly with a factorized Gaussian prior, which aligns with our SN-ratio experiment  thatshows the gradient approximations are of similar quality for these two cases.
Revised: 5 performs similarly with a factorized Gaussian prior, which aligns with our SN-ratio experiment in that the gradient approximations are similar.

Original: Variational autoencoder We evaluate Stein mixtures and SVGD for VAEs on two datasets with α ∈ { 0 , 0 . 5 , 1 } .
Revised: Variational autoencoder We evaluate Stein mixtures and SVGD for VAEs on two datasets with α ∈ { 0 , 0 . 5 , 1 } .

Original: We use binarized MNIST (Salakhutdinov & Murray, 2008; LeCun et al., 1998), a dataset of 28 × 28 pixel images of handwritten single digit numbers, and a variate of OMNIGLOT (Lake et al., 2013), which contains 28 × 28 pixel images of characters from fifty different alphabets.
Revised: We use binarized MNIST (Salakhutdinov & Murray, 2008; LeCun et al., 1998), a dataset of 28 × 28 pixel images of handwritten single digit numbers, and a variate of OMNIGLOT (Lake et al., 2013), which contains 28 × 28 pixel images of characters from fifty different alphabets.

Original: We use the same VAE architecture as Burda et al.
Revised: We use the same VAE architecture as Burda et al.

Original: (2015), detailed in Appendix H. For both datasets we optimize using the Adam optimizer and learning rate of 5 · 10 − 4 .
Revised: (2015), detailed in Appendix H. For both datasets we optimize using the Adam optimizer and learning rate of 5 · 10 − 4 .

Original: We optimize with a batch size of 20 and use 20 draws to approximate the gradients.
Revised: We optimize with a batch size of 20 and use 20 draws to approximate the gradients.

Original: For OMNIGLOT we use 20 epochs and for MNIST we use 50.
Revised: For OMNIGLOT we use 20 epochs and for MNIST we use 50.

Original: Table 2 show the performance of ELBO-within-Stein for α ∈ { 0 . 0 , 0 . 5 , 1 . 0 } .
Revised: Table 3 show the performance of ELBO-within-Stein for α ∈ { 0 . 0 , 0 . 5 , 1 . 0 } .

Original: We find that
Revised: We find that

Original: ELBO-within-Stein with α = 1 achieves better log-likelihoods on MNIST datasets.
Revised: ELBO-within-Stein with α = 1 achieves better log-likelihoods on MNIST datasets.

Original: On OMNIGLOT α = 0 .
Revised: On OMNIGLOT α = 0 .

Original: 5 and α = 0 achieve comparable log-likelihood with α 0 the slightly outperforming α = 0 .
Revised: 5 and α = 0 achieve comparable log-likelihood with α 0 the slightly outperforming α = 0 .

Original:  S UMMARY
Revised:  S UMMARY

Original: We introduce a new algorithm called ELBO-within-Stein based on a new connection between the inference of Stein mixtures and the Rényi variational bound.
Revised: We introduce a new algorithm called ELBO-within-Stein based on a new connection between the inference of Stein mixtures and the Rényi variational bound.

Original: We demonstrate that ELBO-withinStein provides better gradient approximations than alternative algorithms, which results in better performance for standard benchmark problems.
Revised: We demonstrate that Eos provides better gradient approximations than alternative algorithms, which results in better performance for standard benchmark problems.

Original: ELBO-within-Stein is integrated as a black box library in the NumPyro PPL which is distributed freely with an Apache License 2.0 license.
Revised: EoS is integrated as a black box library in the NumPyro PPL which is distributed freely.

Original: A V ARIATIONAL R ÉNYI BOUND
Revised: A V ARIATIONAL R ÉNYI BOUND

Original: For convenience  we derive the variational Rényi bound (Li & Turner, 2016) in the context of inference with our algorithm below.
Revised: For convenience [Equation] we derive the variational Rényi bound (Li & Turner, 2016) in the context of inference with our algorithm below.

Original: Recall that Stein mixtures lift the set of guide hyper-parameters ϕ (optimized in VI) to a random variable ϕϕϕ .
Revised: Recall that Stein mixtures lift the set of guide hyper-parameters ϕ (optimized in VI) to a random variable ϕϕϕ .

Original: Let D be a finite set of observations,
Revised: Let D be a finite set of observations,

Original: z ∈ R
Revised: z ∈ R

Original: d be a latent variable, D α [ q || p ] the Rényi α -divergence (Rényi, 1961) between distributions p and q , and ϕϕϕ ∈ R d a set of guide hyper-parameters lifted to a random variable.
Revised: d be a latent variable, D α [ q || p ] the Rényi α -divergence (Rényi, 1961) between distributions p and q , and ϕϕϕ ∈ R d a set of guide hyper-parameters lifted to a random variable.

Original: Then we have
Revised: Then we have

Original: [Equation]
Revised: [Equation]

Original: To see this
Revised: To see this

Original: [Equation]
Revised: [Equation]

Original: consider,
Revised: consider,

Original: [Equation]
Revised: [Equation]

Original: from which we recover the desired equality by rearranging and multiplying both sides by negative one,
Revised: from which we recover the desired equality by rearranging and multiplying both sides by negative one,

Original: [Equation]
Revised: [Equation]

Original: To shorten notation
Revised: To shorten notation

Original: [Equation]
Revised: [Equation]

Original: we let C α ( D ,ϕϕϕ )
Revised: we let C α ( D ,ϕϕϕ )

Original: =
Revised: =

Original: E (cid:104) ( p ( z , D ) / q ( z |D ,ϕϕϕ ) )
Revised: E (cid:104) ( p ( z , D ) / q ( z |D ,ϕϕϕ ) )

Original: 1 − α (cid:105) where the expectation is with respect to q ( z |D ,ϕϕϕ ) .
Revised: 1 − α (cid:105) where the expectation is with respect to q ( z |D ,ϕϕϕ ) .

Original: The gradient of the variational Rényi bound with respect to ϕϕϕ is∇ ϕϕϕ
Revised: The gradient of the variational Rényi bound with respect to ϕϕϕ is

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: B C HARACTERIZING TWO PARTICLE FIXED POINTS
Revised: B C HARACTERIZING TWO PARTICLE FIXED POINTS

Original: We give the full derivation of stationary points for the Stein mixture that we consider in Section 3.
Revised: We give the full derivation of stationary points for the Stein mixture that we consider in Section 3.

Original: Recall that Section 3 investigated the SN-ratio for a Stein mixture given by
Revised: Recall that Section 3 investigated the SN-ratio for a Stein mixture given by

Original: [Equation]
Revised: [Equation]

Original: where ϕϕϕ 1 ,ϕϕϕ 2 ∈ R d are two d -dimensional particles.
Revised: where ϕϕϕ 1 ,ϕϕϕ 2 ∈ R d are two d -dimensional particles.

Original: We use the kernel given by
Revised: We use the kernel given by

Original: [Equation]
Revised: [Equation]

Original: where h ∈ R + is the bandwidth.
Revised: where h ∈ R + is the bandwidth.

Original: The kernel has the following properties:
Revised: The kernel has the following properties:

Original: [Equation]
Revised: [Equation]

Original: which we will use in the derivation.
Revised: which we will use in the derivation.

Original: Finally, we introduce ξα ( ϕϕϕ ) = 1 1 − α log E
Revised: Finally [Equation] we introduce ξ

Original: [Equation]
Revised: [Equation]

Original: as notation short-hand.
Revised: as notation short-hand.

Original: Our two particle configuration reaches a fixed point when
Revised: Our two particle configuration reaches a fixed point when

Original: [Equation]
Revised: [Equation]

Original: where ϵ ≥ 0 is the step size.
Revised: where ϵ ≥ 0 is the step size.

Original: Therefore, S H Φ ( ϕϕϕ 1 ) =[-
-]
Revised: Therefore, S H Φ ( ϕ ϕ ϕ 1 ) =[-
-]

Original: and S H Φ ( ϕϕϕ 2 )
Revised: and S H Φ ( ϕ ϕ ϕ 2 )

Original: = 

Revised: =

Original: at any fixed point.
Revised: at any fixed point.

Original: S H Φ ( ϕϕϕ 1 ) is given by
Revised: S H Φ ( ϕ ϕ ϕ 1 ) is given by

Original: [Equation]
Revised: [Equation]

Original:  Therefore, − ∇
Revised: [Equation] Therefore, [Equation]

Original: [Equation]
Revised: [Equation]

Original: 16) at a fixed point.
Revised: at a fixed point.

Original: By a similar argument for S H Φ ( ϕϕϕ 2 ) , we have
Revised: By a similar argument for S H Φ ( ϕϕϕ 2 ) , we have

Original: [Equation]
Revised: [Equation]

Original: at a fixed point.
Revised: at a fixed point.

Original: As ∇ ϕϕϕ 1 k ( ϕϕϕ 1 ,ϕϕϕ 2 ) = −∇ ϕϕϕ 2 k ( ϕϕϕ 1 ,ϕϕϕ 2 ) , it follows from Equations (16) and (17)
Revised: As ∇ ϕϕϕ 1 k ( ϕϕϕ 1 ,ϕϕϕ 2 ) = −∇ ϕϕϕ 2 k ( ϕϕϕ 1 ,ϕϕϕ 2 ) , it follows from Equations (16) and (17)

Original: that
Revised: that

Original: [Equation]
Revised: [Equation]

Original: Hence, we see that at any fixed point for our two particle configuration, the gradients of the VR-bound are equal and opposite.
Revised: Hence, we see that at any fixed point for our two particle configuration, the gradients of the VR-bound are equal and opposite.

Original: C K ERNELS IN SVGD
Revised: C K ERNELS IN SVGD

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 ∥
Revised: exp (cid:0) − 1 h ∥

Original: z i − z j ∥ 22 (cid:1) with bandwidth parameter h , chosen as 1 log N med ( Z ) , where med is the median operator.
Revised: z i − z j ∥ 22 (cid:1) with bandwidth parameter h , chosen as 1 log N med ( Z ) , where med is the median operator.

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

Original: For the RBF kernel, the repulsive force becomes
Revised: For the RBF kernel, the repulsive force becomes

Original: [Equation]
Revised: [Equation]

Original: where · is the (euclidean) inner product and 1 d is a d-dimensional one vector.
Revised: where · is the (euclidean) inner product and 1 d is a d-dimensional one vector.

Original: It follows that z i is pushed away from z j when k ( z i , z j ) is large.
Revised: It follows that z i is pushed away from z j when k ( z i , z j ) is large.

Original: D C ONDITIONAL EVIDENCE AS R ÉNYI DIVERGENCE BETWEEN POSTERIORS
Revised: D C ONDITIONAL EVIDENCE AS R ÉNYI DIVERGENCE BETWEEN POSTERIORS

Original: That Equation (11) pushes posteriors towards each other follows from properties of Rényi divergence with α ∈ [0 , 1] and the negative direction of the gradient on the Rényi divergence.
Revised: That Equation (11) pushes posteriors towards each other follows from properties of Rényi divergence with α ∈ [0 , 1] and the negative direction of the gradient on the Rényi divergence.

Original: In particular, we have (i) that the divergence is a similarity measure of distributions for α ≥ 

Revised: In particular, we have (i) that the divergence is a similarity measure of distributions for α ≥

Original: so D α =

Revised: so D α =

Original: [ q || p ] = 

Revised: [ q || p ] =

Original: = ⇒ q = p , (ii) that D α [ q ∥ p ] is everywhere positive, and (iii) the divergence is jointly convex (i.e. convex in both distributions) (Van Erven & Harremos, 2014).
Revised: = ⇒ q = p , (ii) that D α [ q ∥ p ] is everywhere positive, and (iii) the divergence is jointly convex (i.e. convex in both distributions) (Van Erven & Harremos, 2014).

Original: Putting it all together, we see from (ii) and (iii) that the extremum at D α =

Revised: Putting it all together, we see from (ii) and (iii) that the extremum at D α =

Original: [ q || p ] = 0 is global and from the (negative) gradient we are minimizing
Revised: [ q || p ] = 0 is global and from the (negative) gradient we are minimizing

Original: D α =

Revised: D α =

Original: [ q || p ] .
Revised: [ q || p ] .

Original: So, we have that −∇ ϕϕϕ D α =

Revised: So, we have that −∇ ϕϕϕ D α =

Original: [ q ( z |D ,ϕϕϕ )
Revised: [ q ( z |D ,ϕϕϕ )

Original: ∥ p ( z |D )]
Revised: ∥ p ( z |D )]

Original: ==
Revised: ==

Original: ⇒
Revised: ⇒

Original: D α =

Revised: D α =

Original: [ q ( z |D ,ϕϕϕ )
Revised: [ q ( z |D ,ϕϕϕ )

Original: ∥ p ( z |D )]
Revised: ∥ p ( z |D )]

Original: = 0 which from (i) means q ( z |D ,ϕϕϕ )
Revised: = 0 which from (i) means q ( z |D ,ϕϕϕ )

Original: =
Revised: =

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

Original: E A LTERNATIVE ELBO- WITHIN -S TEIN DERIVATION
Revised: E A LTERNATIVE ELBO- WITHIN -S TEIN DERIVATION

Original: For α = 1 we can derive the attractive force of S H Φ directly by applying Jensen’s inequality to the log conditional evidence, resulting in
Revised: For α = 1 we can derive the attractive force of S H Φ directly by applying Jensen’s inequality to the log conditional evidence, resulting in

Original: [Equation]
Revised: [Equation]

Original: In ELBO-within-Stein, the attractive force takes the simple form
Revised: In ELBO-within-Stein, the attractive force takes the simple form

Original: [Equation]
Revised: [Equation]

Original: [Equation] and the repulsive force is given by Equation (5). 
Revised: and the repulsive force is given by Equation (5). [Table] [Table]

Original: F I LLUSTRATING STEIN MIXTURES
Revised: F I LLUSTRATING STEIN MIXTURES

Original: To illustrate the use of a SM, consider the variational auto encoder (VAE) (Kingma & Welling, 2013).
Revised: To illustrate the use of a SM, consider the variational autoencoder (VAE) (Kingma & Welling, 2013).

Original: The VAE simultaneously trains a generative model p ( D| g θ ( z ))
Revised: The VAE simultaneously trains a generative model p ( D| g θ ( z ))

Original: p ( z ) and a variational approximation q ( z | f ψ ( D )) of the posterior p ( z |D ) .
Revised: p ( z ) and a variational approximation q ( z | f ψ ( D )) of the posterior p ( z |D ) .

Original: Here, θ and ψ are parameters of the generative neural network g θ ( · ) and the inference network f ψ ( · ) , respectively.
Revised: Here, θ and ψ are parameters of the generative neural network g θ ( · ) and the inference network f ψ ( · ) , respectively.

Original: VAE training is typically done by stochastic variational inference (SVI) (Hoffman et al., 2013) which optimizes θθθ and ψψψ to minimize the ELBO.[Table]
Revised: VAE training is typically done by stochastic variational inference (SVI) (Hoffman et al., 2013) which optimizes θθθ and ψψψ to minimize the ELBO.

Original: With a SM, the generative model remains the same, that is, we obtain a point estimate of θ .
Revised: With a SM, the generative model remains the same, that is, we obtain a point estimate of θ .

Original: However, the marginal posterior approximation changes to 1 / | Φ | (cid:80) ϕϕϕ ∈
Revised: However, the marginal posterior approximation changes to 1 / | Φ | (cid:80) ϕϕϕ ∈

Original: Φ q ( z | f ϕϕϕ ( D )) .
Revised: Φ q ( z | f ϕϕϕ ( D )) .

Original: So with a Stein mixture, each particle ϕϕϕ parameterizes a separate inference network, i.e. f ϕϕϕ ( · ) , meaning the guide becomes amortized similar to Shu et al.
Revised: So with a Stein mixture, each particle ϕϕϕ parameterizes a separate inference network, i.e. f ϕϕϕ ( · ) , meaning the guide becomes amortized similar to Shu et al.

Original: G BNN WITH POINT - MASS GUIDE
Revised: G UCI BENCHMARK DETAILS

Original: We compare ELBO-within-Stein for α ∈ { 0 , 0 . 5 , 1 } on BNNs regression point mass (Dirac delta) guide and a RBF kernel.
Revised: We compare ELBO-within-Stein for α ∈ { 0 , 0 . 5 , 1 } on BNNs regression point mass (Dirac delta) guide and a RBF kernel.

Original: With ELBO-within-Stein we recover a variant of SVGD with a VR gradient rather than the score function.
Revised: With ELBO-within-Stein we recover a variant of SVGD with a VR gradient rather than the score function.

Original: Like Liu & Wang (2016), we use a BNN with one hidden layer of size fifty and a RELU activation.
Revised: Like Liu & Wang (2016), we use a BNN with one hidden layer of size fifty and a 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 both versions we use 5 particles and update Year for 40 epochs, Protein forepochs and 500 epochs for the rest.
Revised: For both versions we use 5 particles and update Year for 40 epochs, Protein forepochs and 500 epochs for the rest.

Original: We use Adagrad (Duchi et al., 2011) with a step size of 0.05 and a subsample size of 100.
Revised: We use Adagrad (Duchi et al., 2011) with a step size of 0.05 and a subsample size of 100.

Original: All measurements are repeated five times and obtained on a GPU 3 .
Revised: All measurements are repeated five times and obtained on a GPU 3 .

Original: Figure 4 shows the performance of ELBO-within-Stein for α ∈ { 0 . 0 , 0 . 5 , 1 . 0 } with a point mass guide.

Original: Performance is measured in terms of the average root mean squared error (RMSE) and log-likelihood (LL) on test sets.

Original: We find that ELBO-within-Stein with α = 1 achieves a better average log-likelihood for all datasets and a better average RMSE for all but the Boston dataset.

Original: We see that α = 0 and α = .

Original: 5 perform similarly with a factorized Gaussian prior, which aligns with our

Original: SN-ratio experiment in that the gradient approximations are of similar quality.

Original: H VAE DETAILS
Revised: H VAE DETAILS

Original: Following Li & Turner (2016)
Revised: Following Li & Turner (2016)

Original: ; Burda et al.
Revised: ; Burda et al.

Original: (2015); Rainforth et al.
Revised: (2015); Rainforth et al.

Original: (2018), we use VAEs with multiple stochastic layers.
Revised: (2018), we use VAEs with multiple stochastic layers.

Original: The idea is to define the model through ancestral sampling as
Revised: The idea is to define the model through ancestral sampling as

Original: [Equation]
Revised: [Equation]

Original: where x is a data-point (which we will also denote z 0 ), z 1 , . .
Revised: where x is a data-point (which we will also denote z 0 ), z 1 , . .

Original: z L are the L stochastic layers, and θθθ l parameterizes a neural network f l which takes z l +1 to the parameters of the distribution p l , i.e. p [Equation] z l | f l ( z l [Equation]
Revised: z L are the L stochastic layers, and θθθ l parameterizes a neural network f l which takes z l +1 to the parameters of the distribution p l , i.e. p ( z l | f l ( z l +1 ) .

Original: We then let the guide factor in the opposite direction, resulting in
Revised: We then let the guide factor in the opposite direction, resulting in

Original: [Equation]
Revised: [Equation]

Original: We use the same network architecture as Rainforth et al.
Revised: We use the same network architecture as Rainforth et al.

Original: (2018) (summarized in Table 3).
Revised: (2018) (summarized in Table 5).

Original: In Table 3, s denotes a stochastic layer and d denotes a deterministic layer (affine transforms).
Revised: In Table 5, s denotes a stochastic layer and d denotes a deterministic layer (affine transforms).

Original: We use tanh as the activation functions on deterministic layers.
Revised: We use tanh as the activation functions on deterministic layers.

Original: Stochastic layers distribute according to a factorized
Revised: Stochastic layers distribute according to a factorized

Original: Gaussian distribution, and for the likelihood we use the Bernoulli distribution (hence the binarization of the datasets).
Revised: Gaussian distribution, and for the likelihood we use the Bernoulli distribution (hence the binarization of the datasets).

Original: I T HE E IN S TEIN VI LIBRARY
Revised: I T HE E IN S TEIN VI LIBRARY

Original: We provide a library called EinSteinVI for inferring Stein mixtures in the probabilistic programming language (PPL)
Revised: We provide a library called EinSteinVI for inferring Stein mixtures in the probabilistic programming language (PPL)

Original: NumPyro (Bingham et al., 2019; Phan et al., 2019).
Revised: NumPyro (Bingham et al., 2019; Phan et al., 2019).

Original: EinSteinVI uses α -indexed SM inference as its core algorithm as described in Section 3.
Revised: EinSteinVI uses α -indexed SM inference as its core algorithm as described in Section 3.

Original: NumPyro is a universal
Revised: NumPyro is a universal

Original: PPL (van de Meent et al., 2018) embedded in Python.
Revised: PPL (van de Meent et al., 2018) embedded in Python.

Original: NumPyro provides specialized constructs for expressing probabilistic models as Python programs and allows executing arbitrary code in its model and guide.
Revised: NumPyro provides specialized constructs for expressing probabilistic models as Python programs and allows executing arbitrary code in its model and guide.

Original: The computational backend of NumPyro is Jax (Frostig et al., 2018), which combines XLA (accelerated linear algebra) (Sabne, 2020) program transformations with automatic differentiation.
Revised: The computational backend of NumPyro is Jax (Frostig et al., 2018), which combines XLA (accelerated linear algebra) (Sabne, 2020) program transformations with automatic differentiation.

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: Further, we chose NumPyro because:
Revised: Further, we chose NumPyro because:

Original: • NumPyro is embedded in Python, the de-facto programming language for data science; • NumPyro includes the necessary data structures for tracking random variables in both model and guide; • NumPyro features stochastic variational inference (SVI) with an application programming interface (API) that is well suited for EinSteinVI ; and • NumPyro benefits computationally from Jax .
Revised: • NumPyro is embedded in Python, the de-facto programming language for data science; • NumPyro includes the necessary data structures for tracking random variables in both model and guide; • NumPyro features stochastic variational inference (SVI) with an application programming interface (API) that is well suited for EinSteinVI ; and • NumPyro benefits computationally from Jax .

Original: I. A E IN S TEIN VI PROGRAM EXAMPLE
Revised: I. A E IN S TEIN VI PROGRAM EXAMPLE

Original: To demonstrate the two modes of VI (SVGD and Stein mixtures) with EinSteinVI, we consider the 1D
Revised: To demonstrate the two modes of VI (SVGD and Stein mixtures) with EinSteinVI, we consider the 1D

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

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

Original: Figure 6 shows that both SVGD 4 and the Stein-mixture naturally capture the bi-modality of the target distribution, compared to SVI with a Gaussian guide.
Revised: Figure 4 shows that both SVGD 4 and the Stein-mixture naturally capture the bi-modality of the target distribution, compared to SVI with a 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 perturbs the locations.
Revised: Also, note that the Stein-mixture overestimates the variance and slightly perturbs 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, 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, 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: def model(): sample('x', NormalMixture(jnp.array([1 / 3, 2 / 3]), jnp.array([-2.0, 2.0]), jnp.array([1.0, 1.0])))

Original: [Figure]

Original: 1D Gaussian mixture model svi = SVI( model,AutoNormal(model),

Original: Adagrad(step_size=1.0),Trace_ELBO() ) results = svi.run(rng_key,num_iterations) (

Original: [Figure]

Original: ) SVI stein = SteinVI(model,

Original: AutDelta(model),Adagrad(step_size=1.0),

Original: Trace_ELBO(),RBFKernel(), ) results = stein.run(rng_key, num_iterations)

Original: [Figure]

Original: [Figure]

Original: I. I NTEGRATION WITH N UM P YRO
Revised: I. I NTEGRATION WITH N UM P YRO

Original: Integrating EinSteinVI with NumPyro requires handling transformations between the parameter representation of NumPyro 5 and the array representation that ELBO-within-Stein operates on.
Revised: Integrating EinSteinVI with NumPyro requires handling transformations between the parameter representation of NumPyro 5 and the array representation that ELBO-within-Stein operates on.

Original: For this, we rely on Jax ’s PyTrees 6 which converts back and forth between Python dictionaries and array representations.
Revised: For this, we rely on Jax ’s PyTrees 6 which converts back and forth between Python dictionaries and array representations.

Original: Algorithm 1 shows the black-box version of α -indexed SM inference in NumPyro .
Revised: Algorithm 1 shows the black-box version of α -indexed SM inference in NumPyro .

Original: The algorithm allows SVI to estimate a subset of the parameters and α -indexed SM inference the rest.
Revised: The algorithm allows SVI to estimate a subset of the parameters and α -indexed SM inference the rest.

Original: To differentiate the two, we denote parameters updated by SVI with ψψψ and parameters updated by ELBO-within-Stein with ϕϕϕ i (i.e. the Stein particles Φ = { ϕϕϕ } Ni =1 ).
Revised: To differentiate the two, we denote parameters updated by SVI with ψψψ and parameters updated by ELBO-within-Stein with ϕϕϕ i (i.e. the Stein particles Φ = { ϕϕϕ } Ni =1 ).

Original: In the model, only SVI can update parameters which we denote by θθθ .
Revised: In the model, only SVI can update parameters which we denote by θ θ θ .

Original: We update θθθ and ψψψ by averaging the loss over the Stein particles.
Revised: We update θ θ θ and ψ ψ ψ by averaging the loss over the Stein particles.

Original: For the Stein particles, the process is more elaborate.
Revised: For the Stein particles, the process is more elaborate.

Original: First, we convert the set of individual distribution parameters in the guide to a monolithic array using Jax ’s PyTrees .
Revised: First, we convert the set of individual distribution parameters in the guide to a monolithic array using Jax ’s PyTrees .

Original: The array represents the particles as a flattened and stacked Jax array.
Revised: The array represents the particles as a html def model(): sample('x', NormalMixture(jnp.array([1 / 3, 2 / 3]), jnp.array([-2.0, 2.0]), jnp.array([1.0, 1.0]))) [Figure] 1D Gaussian mixture model svi = SVI( model,AutoNormal(model), Adagrad(step_size=1.0),Trace_ELBO() ) results = svi.run(rng_key,num_iterations) (

Revised: [Figure]

Revised: ) SVI stein = SteinVI(model,

Revised: AutDelta(model),Adagrad(step_size=1.0),

Revised: Trace_ELBO(),RBFKernel(), ) results = stein.run(rng_key, num_iterations)

Revised: [Figure]

Revised: [Figure]

Revised: flattened and stacked Jax array.

Original: Then, we compute a kernel on the particles, delegated to the kernel interface (see Appendix I.3)
Revised: Then, we compute a kernel on the particles, delegated to the kernel interface (see Appendix I.3)

Original: as the computation is kernel-dependent.
Revised: as the computation is kernel-dependent.

Original: We apply Jax ’s vmap (Frostig et al., 2018; Phan et al., 2019) operator to compute the Stein forces for each particle in a vectorized manner.
Revised: We apply Jax ’s vmap (Frostig et al., 2018; Phan et al., 2019) operator to compute the Stein forces for each particle in a vectorized manner.

Original: As we compute the Stein forces in unconstrained space, we must correct them by the Jacobian of the bijection to constrained space.
Revised: As we compute the Stein forces in unconstrained space, we must correct them by the Jacobian of the bijection to constrained space.

Original: Naively computing the
Revised: Naively computing the

Original: Jacobian on the monolithic array incurs a massive memory overhead.
Revised: Jacobian on the monolithic array incurs a massive memory overhead.

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 distribution parameters rather than the monolith.
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 distribution parameters rather than the monolith.

Original: Like computing the Stein forces, the correction is embarrassingly parallel so that we can use a vmap operator again.
Revised: Like computing the Stein forces, the correction is embarrassingly parallel so that we can use a vmap operator again.

Original: Inside the vmap we nest a tree_map to do the appropriate conversion between representations.
Revised: Inside the vmap we nest a tree_map to do the appropriate conversion between representations.

Original: Finally, we convert the monolithic array to its NumPyro representation and return the expected changes for SVI- and Stein- parameters.
Revised: Finally, we convert the monolithic array to its NumPyro representation and return the expected changes for SVI- and Stein- parameters.

Original: Require: SVI parameters θθθ
Revised: Require: SVI parameters θθθ

Original: and ψψψ , Stein parameters { ϕϕϕ i } Ni =1 , model p ϕϕϕ
Revised: and ψψψ , Stein parameters { ϕϕϕ i } Ni =1 , model p ϕϕϕ

Original: ( z , x ) , guide q θθθ,ψψψ ( z ) , loss
Revised: ( z , x ) , guide q θθθ,ψψψ ( z ) , loss

Original: L α , kernel interface KI .
Revised: L α , kernel interface KI .

Original: Ensure:
Revised: Ensure:

Original: Parameter changes based on SVI ( ∆ θθθ , ∆ ψψψ ) and hierarchical Stein forces ( { ∆ ϕϕϕ i } Ni =1 ).
Revised: Parameter changes based on SVI ( ∆ θθθ , ∆ ψψψ ) and hierarchical Stein forces ( { ∆ ϕϕϕ i } Ni =1 ).

Original: procedure U PDATE ( θθθ , ψψψ , { ϕϕϕ i } Ni =1 , p θθθ , q ϕϕϕ,ψψψ )
Revised: procedure U PDATE ( θθθ , ψψψ , { ϕϕϕ i } N i =1 , p θθθ , q ϕϕϕ,ψψψ )

Original: ∆
Revised: ∆

Original: θθθ ← E θθθ [ ∇ θθθ L α ( p θθθ , q ϕϕϕ,ψψψ )]
Revised: θθθ ← E θθθ [ ∇ θθθ L α ( p θθθ , q ϕϕϕ,ψψψ )]

Original: ∆ ψψψ ← E ψψψ [ ∇ ψψψ
Revised: ∆ ψψψ ← E ψψψ [ ∇ ψψψ

Original: L α ( p θθθ , q ϕϕϕ,ψψψ )] { a i } i ← PyTreeFlatten( { ϕϕϕ i } Ni =1 ) k ← KI( { a i } Ni =1 ) procedure HS TEIN -F ORCES ( a i )
Revised: L α ( p θθθ , q ϕϕϕ,ψψψ )] { a i } i ← PyTreeFlatten( { ϕϕϕ i } Ni =1 ) k ← KI( { a i } Ni =1 ) procedure HS TEIN -F ORCES ( a i )

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

Original: θθθ i ← P Y T REE R ESTORE ( a i ) ∆ a i ← (cid:80) a j k ( a j , a i ) ∇ a i L α ( p ϕϕϕ , q θθθ
Revised: θθθ i ← P Y T REE R ESTORE ( a i ) ∆ a i ← (cid:80) a j k ( a j , a i ) ∇ a i L α ( p ϕϕϕ , q θθθ

Original: i ,ψψψ )
Revised: i ,ψψψ )

Original: + ∇ a i k ( a j , a i ) return ∆ a i end procedure { ∆ a i } i ← VMap ( { a i } i , HS TEIN -F ORCES ) { ∆ ϕϕϕ i } Ni =1 ← P Y T REE R ESTORE ( { ∆ a i } Ni =1 ) return ∆ θθθ , ∆ ψψψ , { ∆ ϕϕϕ i } Ni = end procedure 
Revised: + ∇ a i k ( a j , a i ) return ∆ a i end procedure { ∆ a i } i ← VMap ( { a i } i , HS TEIN -F ORCES ) { ∆ ϕϕϕ i } Ni =1 ← P Y T REE R ESTORE ( { ∆ a i } Ni =1 ) return ∆ θθθ , ∆ ψψψ , { ∆ ϕϕϕ i } Ni = end procedure I. K ERNEL INTERFACE

Original: The kernel interface is straightforward.
Revised: 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 
Revised: 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 .

Revised: Table 6 gives the complete list of kernels in EinSteinVI.

Original: [Table]
Revised: [Table]

Original: Type
Revised: Type

Original: Reference
Revised: Reference

Original: [Table]
Revised: Radial Basis

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

Original: [Equation]
Revised: [Equation]

Original: 22 ) scalar
Revised: [Equation] scalar

Original: Liu & Wang (2016) exp [Equation] y ))
Revised: Liu & Wang (2016) exp( 1 h ( x − y ))

Original: vector
Revised: vector

Original: Pyro
Revised: Pyro

Original: [Table] Multi- [Table] (IMQ) ( 
Revised: Inverse MultiQuadratic (IMQ) ( c 2 + ∥ x − y ∥ 22 ) β β ∈ ( − 1 , 0) and c >

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

Original: Gorham &
Revised: Gorham &

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

Original: RandomFeatureExpansion
Revised: Random Feature Expansion

Original: [Equation]
Revised: [Equation]

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

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

Original: Linear x ⊤ y +scalar
Revised: Linear x ⊤ y +scalar

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

Original: Mixture (cid:80) i ω ω ω
Revised: Mixture (cid:80) i ωωω i

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

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

Original: [Table] based
Revised: Scalar-based

Original: Matrix k
Revised: Matrix k

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

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

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

Original: [Table] 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: Graphical diag( { K ( ℓ ) ( x , y ) } ℓ ) { K ( ℓ ) } ℓ scalar-valued kernels, each for a unique partition of latent variables matrix, placed on the diagonal
Revised: Graphical diag( { K ( ℓ ) ( x , y ) } ℓ ) { K ( ℓ ) } ℓ scalar-valued kernels, each for a unique partition of latent variables matrix, placed on the diagonal

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

Original: Constant Preconditioned
Revised: Constant Preconditioned

Original: [Equation]
Revised: [Equation]

Original: K is an inner matrix-valued kernel and Q is a preconditioning matrix like the Hessian −∇ 2¯ z log p (¯ z | x ) or Fisher 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 Fisher 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: Anchor Point
Revised: [Table]

Original: Preconditioned (cid:80) mℓ =1 K Q ℓ ( x , y )
Revised: Preconditioned (cid:80) mℓ =1 K Q ℓ ( x , y )

Original: ωωω ℓ ( x ) ωωω ℓ ( y ) { a ℓ } mℓ =1 is a set of anchor points, Q ℓ = Q ( a ℓ ) is a preconditioning matrix for each anchor point, K Q ℓ is an inner kernel conditioned using Q ℓ , and ωωω ℓ ( x ) = softmax ℓ ( {N ( x | a ℓ ′ , Q − 1 ℓ ′ ) } ℓ ′ ) matrix
Revised: ωωω ℓ ( x ) ωωω ℓ ( y ) { a ℓ } mℓ =1 is a set of anchor points, Q ℓ = Q ( a ℓ ) is a preconditioning matrix for each anchor point, K Q ℓ is an inner kernel conditioned using Q ℓ , and ωωω ℓ ( x ) = softmax ℓ ( {N ( x | a ℓ ′ , Q − 1 ℓ ′ ) } ℓ ′ ) matrix

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

Original: (2019) and particles, and the loss function L and returns a differentiable kernel k .

Original: Table 4 gives the complete list of kernels in EinSteinVI.