ELBO-ing Stein Mixtures
Authors: Ola Rønning
Christophe Ley
Ahmad Salim Al-Sibahi
Thomas Hamelryck
Venue: Submitted to ICLR 2023 (2023)
Keywords: Particle-based variational inference
alpha-indexed Stein mixtures
ELBO-within-Stein

Abstract: Stein variational gradient descent (SVGD) \citep{DBLP:conf/nips/LiuW16} is a particle-based technique for Bayesian inference. 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. Unfortunately, the number of particles required to represent a model adequately grows exponentially with the dimensionality of the model. Stein mixtures \citep{nalisnick2017variational} alleviate the exponential growth in particles by letting each particle parameterize a distribution. However, the inference algorithm proposed by \cite{nalisnick2017variational} can be numerically unstable.  We show that their algorithm corresponds to inference with the R\'enyi $\alpha$-divergence for $\alpha=0$ and that using other values for $\alpha$ can lead to more stable inference. We empirically study the performance of Stein mixtures inferred with different $\alpha$ values on various real-world problems, demonstrating significantly improved results when using $\alpha=1$, which coincides with using the evidence lower bound (ELBO). We call this instance of our algorithm ELBO-within-Stein. A black-box version of the inference algorithm (for arbitrary $\alpha\in \sR$) is available in the deep probabilistic programming language NumPyro \citep{phan2019}.

Under review as a conference paper at ICLR 2023
ELBO-ING STEIN MIXTURES
Anonymous authors Paper under double-blind review
ABSTRACT
Stein variational gradient descent (SVGD) Liu Wang 2016) is particle- based technique for Bayesian inference. SVGD has recently gamedd popularity because it combines the ability of variational inference to handle tall data with the modeling power of non-parametric inference. Unfortunately, variance estimation scales inversely with the dimensionality of the model leading to underestimation, meaning more particles are required to represent high-dimensional models adequately. Stein mixtures (Nalismick & Smyth 2017) alleviate the exponential growth in particles by letting each particle parameterize distribution. However, the inference algorithm proposed by Nalisnick & Smyth 2017) can be numerically unstable. We show that their algorithm corresponds inference with the Rényi a-divergence for 0 and that using other values for a can lead to more stable inference. We empirically study the performance of Stein mixtures inferred with different a values on various real-world problems, demonstrating significantly and consistently improved results when using a 1, which corresponds to using the evidence lower bound (ELBO). We call this instance of our algorithm of the inference algorithm (for arbitrary a E I)) is available in the deep probabilistic programming language NumPyro (Phan et al. 2019).
An easy-to-use version
INTRODUCTION
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 al. (02T). Bayesian inference aims to describe i.i.d. data D = }x using a model with latent a variable z. Bayesian inference does this by computing posterior distribution p(z[D) over the latent variable given a model describing the joint distribution p(z, D) p (Dla (zz). We obtain the posterior by following Bayes' theorem,
p(z 1 p(x, p(x: ()))()/(pp))
where - - Je IT= I p(x, p( (z))zz is the normalization constant. For most practical models, the normalization constant lacks an analytic solution or poses a computability problem, complicating the Bayesian inference problem.
Stein variational gradient descent (SVGD) [Liu Wang 2016) a recent technique for Bayesian inference that uses a set of particles 2 2 F.v to approximate the posterior p(z)D). The idea
behind SVGD is to iteratively transport Z according to a force field Sz, called the Stein force. The Stein force is given by
Sz(z.) = Feg (M; 2,) Vo, log ((2 + Vz, k(z,, 2;)1;
(1)
where k(---) reproducing kernel (Berlinet & Thomas Ahomass-Anan 2011). 92 =N 82, is the empirical measure on the set of particles 2, 0x (y) represents the Dirac delta measure, which is equal to if x = y and 0 otherwise, and V,, log p(2, [D) the gradient of the posterior with respect to the j-th particle. The technique is scalable to tall data (i.c. datasets with many data points) and offers the flexibility and scope of techniques such as Markov chain Monte Carlo (MCMC), SVGD is good at capturing molti-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).
The main problem is that SVGD suffers from the curse dimensionality: variance estimation scales inversely with dimensionality (Ba et al. 2021). Nalisnick & Smyth (2017) suggest resolving this by

---

Under review as a conference paper at ICLR 2023
using Stein mixture (SM). SMs lift each particle to the parameters of variational distribution q. also called a guide. 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 posterior. The introduces guides by replacing each posterior gradient Va, logp(z;))) (D)
Nalisnick & [iqqatioon (1) with the corresponding gradient of the marginal log-variational likelihood given by
p(D,z)
log pDl log ,.),)
[(((D,
Here, we denote the particles by Φ = [6 instead of Z = {z to emphasize they parameterize guide components q<]),, : D). The change in gradient corresponds to minimizing DKL 40 (Ø) p(ø) D)] rather than 1 (z) || p(z] as in SVGD. Note that the line between the model and guide becomes blurred, as pl (D is random in both data (D), as is usually the case, but also in the guide pappee-rraeeters & (Ranganath
al. 2016 Nalisnick & Smyth 2017). To distinguish the two we subsequently refer to p(D) as the evidence and p(D))) as the hierarchical likelihood.
The Stein force using the log hierarchical likelihood, which we call the hierarchical Stein force SH. becomes
( EE, ~u k(,, 0,) log Bgte/
p(D,z) ,,
+ Vo, ()
(3)
where qa is an empirical measure analogous to 12.
Inference converges (i.e. reaches fixed point) when ($)) for all particles, meaning all gradients in S must cancel (i.e. sum to zero). However, computing the gradient of the log-variational likelihood requires numerical estimation as analytical solutions do not exist for most models. Hence, cannot expect the inference converges with noisy gradient estimations compensate for the error in the gradient by a counterforce in the next iteration. Therefore, SMs require good (i.e. low relative variance) gradient approximations; otherwise, the particles will fluctuate around a fixed point without reaching it. We demonstrate that replacing the log hierarchical likelihood with the evidence lower bound (ELBO) can provide better (lower relative variance) gradient approximations. We call the new algorithm ELBO HL.BO-withi--Seee
the Stein force will
(EoS). 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] [0144. Similarly to the ELBO, the VR bound is a lower bound of the evidence, p(D), also called the normalization constant, and given by
1 a
log (x((DD)))) [
((D,z) (2)
DDD) 2
where a 2 is known as its order" Understanding the inference of SMs in terms of the VR bound now understand controlling the
yields insight into the behavior of the two algorithms, variance of each component of the guide. Furthermore, presuming accurate gradient approximation for all (viable) values of a; the connection leads to family of inference algorithms indexed by the VR bound order.
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:
We demonstrate that inaccurate gradient estimates can lead issues with convergence for SMs.
We introduce a new family of inference algorithms for SMs indexed by the parameter a. The family results from connecting inference with SMs to the Rényi a-divergence and includes the inference algorithm by Nalisnick & Smyth (201)) as a special case for O.
Unlike previous work, our algorithm allows for investigating a range of values for for model of interest. This allows to investigate the convergence stability for different a's by measuring the SN-ratio. We find that
is optimal for models with latent variable
'The VR bound can be extended to E R. We presume is finite, but we allow to be less than equal to
zero Van Erven & Harremos) 2014)
2

---

Under review as a conference paper at ICLR 2023
for each data point (local latent variables), resulting in better SN-ratios than all other values. For models where all datapoints share latent variable (global latent variables), using 0.5 (corresponding to the Hellinger distance) is par with Nalismick & Smyth (2017)ss 's
al algorithm (which corresponds to a = 0). Other values for a result in worse SN-ratios.
We evaluate our inference algorithm for different values of a 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.
We describe black-box inference algorithm for our proposed family of inference algorithms and provide a software library, called EinSteing in NumPyro.
In Section 4 we discuss related work. We benchmark our algorithm in Section summarize results in Seectio 6
Finally, we
2 BACKGROUND
Let be latent variable of interest taking values in space Z Rd (up to diffeommrppismm and D = Jee be a set of i.i.d. observations. For many models, exact Bayesian inference is computa- tionally impracticable due to the cost of evaluating the evidence p(D). Therefore, practitioners turn to tractable approximate variational inference (VI).
VI aims to bring computationally cheap variational distribution ((z(D) close to the model posterior. Typically, we measure closeness by the Kullback-Leibler divergence (DRL). ic. DEL (x(z(D) |l p(z[D)). However, generally avoid directly evaluating DKL [g(z(D) || p(z[D)] as this requires evaluating the evidence, p(D). We will concern ourselves with two types of VI.
The first type of VI searches for parameterization r/* of g in family of distributions Q that minimizes the divergence to the posterior. When the divergence is measured by DKI., this type of VI made tractable by maximizing the evidence lower bound (ELBO), that is
p(D, z)
2 = arg max (logp(D) (kllg(z/D:e)) p(z[D))) arg max Ey(z) log q(z 2b)
The second type of VI we consider relies on particle- particle-based methods and the focus of this article. This type of VI relies transporting finite set of particles such that their empirical measure is close to the posterior. We will discuss this method in detail below.
2.1 STEIN VARIATIONAL GRADIENT DESCENT
The core idea of SVGD is to perform inference by approximating the target posterior distribution p(z[D) by an empirical distribution qz (z) = . -N- 8z, (z) based on a set of particles Z, where
2 = | One could thus see the approximating distribution q2 (z) as (uniform) mixture of point estimates, each represented by particle z; E Z. The SVGD algorithm minimizes the Kullback-Leible divergence DKL[9z (z) p(z [D] between the approximated and the true posterior by iteratively updating the particles using the following expression:
2,+1 + (Z,)
where is the learning rate and Sz denotes the Stein force.
The two forces of SVGD The Stein force Sz consists of two underlying forces that work additively, with 5z = 5t + +Z. The attractive force is given by
(t(2;) = Exz-aq [k(z;, 23)) V log p(z, [D)]
and the repulsive force by
(z)): Bz V k(z,, z,)].
(5)
Here k Rd x Rd - R is kernel. The attractive force be seen as pushing the particles towards the modes of the true posterior distribution. smoothed by some kernel. The repulsive force stops particles with high kernel values from collapsing each other. In Anpppediix C. we demonstrate the
3

---

Under review as a conference paper at ICLR 2023
repulsive behavior for radial basis function (RBF) kernel. The computational cost of S2 is quadratic in the size 2, i.c. ((V3), which makes SVGD computationally burdensome for high-dimensionas posteriors. For a particle method such SVGD, the number of particles required to represent posterior distribution adequately is exponential in its dimensionality.
SVGD suffers from the curse of dimensionality (Ba al. 2021). which results in variance collapse (i.e. variance is underestimated), Wang et al. (2018) demonstrates the problem with simple factorized Gaussian, suggesting the (RBF) kernel introduces global statistical dependence driving the need for particles up for accurate representation. [a et al. (02T)) demonstrate that the collapse is due to the deterministic update of the attractive force. They do this by showing that re-sampling the particles at each iteration eliminates the underestimation of variance. Note that their particle re-sampling scheme by Ba et (202T) not generally tractable; hence it does not suffice as a practical solution.
2.2 RÉNYT DIVERGENCE AND THE VARIATIONAL RÊNYT BOUND
The Rényi divergence (Rényi
1961) is family of divergences between distributions p and g indexed by the order parameter a E /[0,1,, [Dal < oo). The divergence is given by
log / p(z) (
Da p ll
The Rényi divergence can be extended to a E {0, 1, oo) } by continuity. addition, we allow for D&lp 9] 0, the order can be further extended orders correspond to known divergences (see (Van Erven & Harremos) 2014) and (L.i & Turner 2016)
I (Van Erven & Harremos 2014). Several
for overview). In particular,
corresponds to DKL-
Analogous to the use the DKL in the ELBO, Da leads to a variational Rényi bound (Li & Turner 2016) which, when formulated used with SMs, is given by
p(z, D) g(((D,))
1-a
logp(D) - Da [g(z| 11 p(z = 1- log E (m(())
(6)
Note that model parem--prameeers (Ø) in the variational posterior, (z) Φ). are lifted to random
variables when doing inference with SMs. See Appendix A for the derivation of Equation (6)
Assuming reparameterization of z possible, we can approximate the gradient A(ő) of Equation (6)
using Monte Carlo integration by
p(Zk, D) ((Zy|D,y)
Ax(
Σ wi (Z 1m1
with Z₁ ())
where E N number of draws used to compute the VR bound and
(2, = 1 p(D, 1-a
K
p(D, 2i)
1-a
with C =
(8)
g(zx (D,
((,;
We provide the derivation in Appendix
2.3 THE SIGNAL-TO-NOISE RATIO
The signal-to-noiss (SN) ratio was introduced by Rainforth et al. (2018 to study the effect of tighter variational bounds on gradient estimation. The SNR given by
E M (Φ)
SNRM.K(0)
(9)
A (Ø)
where a ( is the standard deviation, M, K N are the number of Monte Carlo draws, and A Mr (Φ)
derives from rewriting Equation (7) in the form
1 M
K Σ
p(Zm,k; D) 1-07
A (xx 1 - M m
Σ
V log
(10)
K
g(
km

---

Under review as a conference paper at ICLR 2023
Here, separate tightening the bound (by increasing K) from reducing the noise in the gradient estimation (by increasing M). 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. The counter-intuitive implication is that a tighter bound can worsen the gradient estimation.
2.4 THE STEIN MIXTURE
Variational inference with SMs (Nlisnicc & Smyth 2017) approximates the target posterior distribu-
tion D(z by letting the Stein particles = (0: i= parameterize guide programs, q(ziØ,: D).
SM yields mixture marginal variational posterior, p(z)D) = 1/10) LeEo g(z) Φ. D), from which it takes its name. Formally, SM is a hierarchical variational model (HVM) (Ranganath et al. 2016) with an empirical measure of particles da (defined the same way as (z) as its variational poste- rior, uniform variational prior, and variational likelihood Px(2|D.e) (p(P.zl))/or (x(x)))) Similarly to SVGD, SM minimizes (x |I P())D)) by iteratively transporting the particles according to the following expression
Φ Φ, + (0)
where is the learning and 5 is the hierarchical Stein force.
The attractive force of SM Like SVGD, SM also makes use of two additive forces, S SH+ + +Sp- The repulsive force Sh is the same in SVGD, given by Equation (5) The attractive force is given by
p(D,z) qz(
sul+
(0:) = Espng k( V log (z(0) +
where l: Rd x Rd + R is a kernel. From the construction of SVGD, we require that the kernel has the reproducing property, the kernel is dense in the space of continuous functions. If we choose 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, Φ,) and reduces to the RBF kernel for fixed variance, which is a reproducing kernel. The EL kernel is given by
k 2,), =
,,
where L2 is an inner product and a positive definite kernel.
3 Q-INDEXED STEIN MIXTURES INFERENCE AND
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 and the VR bound given in EEquation (6) for = 0. Presuming the support of the variational likelihood q(z(Ø) is a subset of the support of the prior of p. supp(g(zio)) C supp(p(z)), the gradient of the log hierarchical likelihood is given by
p (DD,2) 2)
V log logp(Dl)) V log E,
(ax
Eq. (6)
V ((D) D (a=o(qgZ\DDe Φ) p(z[D)))
(zD)=oll(2
[()[D)]
(11)
From EEquation (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 a = 0) between the variational posterior qz () and the model posterior p(z|[D),
the other hand. Thus, Equation (11) shows that the attractive hierarchical force (S pushes the components of the variational posterior, q(z) D. towards the model posterior, PDz(D). see see Apelndpxi D for details. The equivalence in Equation (11) suggests whole class of hierarchical attractive forces indexed by the order a of the VR bound. Note that choosing a + means we lose the interpretation of the attractive
force moving the particles towards the nearest peak of the conditional evidence. Assuming
5

---

Under review as a conference paper at ICLR 2023
(a) Low accuracy gra- dient approximation
(b) High accuracy gra- dient approximation
Figure I: Two particle system at theoretical fixed point. The blue arrows indicate the magnitude and direction of the attractive force, the red arrows show the repulsive force, and the black arrows the Stein force. Note that Figure Ib has no Stein force as expected for a converged system.
marginal variational posterior g(2))) ()) is reparameterizable, we can approximate the attractive force for any as
(0.) = Expiga () () (Ø)],
(12)
where (Ø) is given by Equation (7) We call inference with Equation (12) a-indexed Stein mixture inference. There are two special cases
that are worth highlighting. The first is a = 1/2, for which the Rényi divergence corresponds to the Hellinger divergence (Van Erven & Harremos 2014)(Li & Turner 2016). The second is a = 1, corresponding to the DkL-divergence. In this case, the VR bound recovers the ELBO. We call this instance of our c-indexed SM inference algorithm ELBOO-ithinnSteiin In (nppppexxxx E] we show that we can also recover the a = 1 case directly by applying Jensen's inequality to the conditional evidence.
3.1 INVESTIGATING THE TO-SSE RATIO
Estimation of a Stein mixture converges when S which means that the repulsive and attractive forces must be equal and oppose for them to cancel. Hence, convergence requires A x (Φ) and Vo, k( Φ22 to be accurate. In Figure we demonstrate the effect of inaccurate gradient approxima- tions. To study the sensitivity of gradient approximations the choice a, we measure the SN-ratio (see (eceEEtonoon (9) of the VR bound gradients (see Equation (6)). We simulate data {x, )= from a simple latent variable model given by N (D|z,, La) u, Ia). where E Rd is unknown and la is the d-dimensional identity matrix. To approximate its posterior use Stein mixture of the form
(/2(N( 3/2la) + *(d2.))
and expected likelihood kernel. We choose (computationally convenient) fixed variance such that the Stein mixture cannot exactly recover the posterior. We can see this as the posterior is unimodal which is only the case for the Stein mixture |6) Φ2 (Behboodian 1970). but in this interval the variance of the Stein mixture will be greater than or equal to 3/2. With the expected likelihood kernel we can analytically characterize all fixed points for the Stein particles as
1 C
( plzz D) ((2(),)
1-a
((2, D)
1-a
log Eg(z(o,)
V
log Ey(zlo.)
(2(())))
See the Appendix for the derivation. measure the effect of gradient approximation on the system we use Equation to estimate the gradients.
To conduct our experiment, we sample the location μ from a 20-dimensional standard Gaussian and use this μ to simulate n 64 data points D. We then approximate the gradients random point close to a fixed point
1 nD DD
M+nD
(1)
+ V AM. (0,) +e, +1 Vo, (...AM..(()))
+1
where is the data average, u+nD/n+1 is the posterior mean and offsets each dimension by Gaussian with mean zero and variance 0.01.

---

Under review as a conference paper at ICLR 2023
NAZU-000) Ialegelom.5 E.80 11
10
EL8O 11
parrts
(b) SN-ratio convergence with a local latent variable.
(c) SN-ratio convergence with a global latent variable.
(a) Experimental setup.
Figure 2b and Figure 2c show the convergence of the SN-ratio (see Equation (9) as we tighten the VR bound by increasing either K or M. We only show the SN-ratio for the first particle (Φ1) as the
behaviour for the second particle is the same. For the ELBO (a 1), we fix K = and increase M reduce the gradient approximation variance, while for the rest (a E {0, 0.5, 2, 10}), we fix M 1 and increase K to tighten the VR bound. In [Fiigure 26 there is a latent variable for each data point.
Note how the SN-ratio only improves with tightness for a (green line). In Figure 2c all data points share latent variable.
Figure 2a illustrates the experimental setup in two dimensions. For legibility in Figure 2a we do not include the perturbation (i.e. added noise) Φ; to the visualization. The contours correspond the exact posterior. As the particles are placed equidistant from the posterior mean (marked with a blue As would expect (see blue arrows in Figure 2a the
cross), in this setting the Stein forces gradient estimations of V VaAi M, are equal and opposite for the two particles.
For a * 1 we fix M and vary K, while for a == we fix K = 1 and vary M. We do not need to consider # when K = 1 as the associated gradient scaling (1 a)) cancels in the SN-ratio. We empirically estimate the SN-ratio by estimating the expectation and standard deviation from 10, 000 gradient samples. In [Figure 26 we show local variate of the model where there is latent variable z; for each X; E D. We see that for a f 1 the SN-ratio does not depend on the particular choice of and that the growth in SN-ratio superseded by that of
1. 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). In Figure 2c we evaluate a global latent variable variate of the model, so that there is one z for all datapoints in D. As with the local version, we fix either M or K. For this model, we see that a = 0 and a = 1/2 achieve the highest SN-ratio. The result aligns with the BNN example, where = is not dominating in performance over the (0, 1/2) on all datasets.
High precision is desirable to avoid fluctuation at convergence. From the above results, we see that = is not necessarily the best choice for precise (high SN-ratio) gradient estimation. In particular, for local latent variable models, a 1 is better choice, and for that, for global latent models, a = 1/2 is on par with a = in our experiment.
3.2 BLACK-BOX INFERENCE FOR ELBO-WITHN-STEEN
We provide a mini-batch version of ELBO-within-Stein called VI, in NumPyro. To compute the VR bound exactly requires all the datapoints, that is, we cannot represent the bound as point- wise expectation, except for a = 1. Therefore. in order to make EinSteinVI scalable to tall data, we provide an approximation of the VR bound which replaces the likelihood by Pr(D|z, Φ) Ikez p(a,) iz, h) /izz,, where I is a subset of a permutation of the data indices. The approximate attractive force s (Ø) used in EinStein is given by
1
PT (D ))((2)
1-a
5 (b.) = Exp
k(
V
log F
(13)
qz (z)
which recovers the exact VR bound when I provide example programs in Appendix 1
DJ. We describe the NumPyro integration and

---

Under review as a conference paper at ICLR 2023
Table I: Average RMSE (lower is better) test results for BNNs on UCI benchmark. EoS (ours)
corresponds 1, Hell (ours) to = 0.5, and SM (Nalisnick & Smyth 2017 to 0. Parentheses mean we use Dirac delta guides. EoS and Hell gives the best results.
Root Mest Squared Deviation (RMSE) SM
Hell
SVGD
MFVI
Laplace
Boston
[2.5 0.4)
3.5 0.82
(2.71 + 0.26)
3,92 =
(2 (2:715- + 0. 0.263)
3.9 1.29
2,86 1 0.23
3.25 1. 0.0.1
3.68 1. 0.33
Concrese
5.76 + 0.55
0.53 0.34)
6.56 + 0.53
0.3)
(5.23
6.51 0.09
5.54 4 0.33
5.6
5.22 0.43
Energy
(0.45 0.03)
(0.67 + 0.033
0.94 0.15
(0.74 00.05)
1,30 + 0.08
1.75 0.15
0:46
Naval
(0.00 =: 0.D0)
0.044 = 0.04
(0.001 2: 0.000)
0.004 0.001
0.004 0.0002
0.0011 + 0.000
0.007 1: 0.000
0.000 0.00
0.00 1: 0.00
Wine
(0.07 0.00)
0.035
0.61 +-035
0.61 + 0.03
0.62 0.04
0.09 = 0.04
0.61
Yacht
(0.45 0.03)
1.76
(0.87 1: 0.03)
1.66 0.85
(0.74 &: 0.05)
1.61
1.11 0.3
4.09 L 0.31
2.16 0.37
Pawer
(3.91 0.18)
4.04 0.16
(3.98 + 0.19)
4.15 + 0.21
(3.97 +
4.16 0.21
4.06 à: 0.17
3.94 A 0.18
3.99 0.11
4 RELATED WORK
Nalisnick & Smyth (2077) first suggested Stein mixtures as an alternative to HVMs (Ranganath et al. 2016). Using SVGD allows Stein mixtures to side-step the need of HVMs for an auxiliary distribution to keep the bound (learning objective) tight. 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 kernel controls the tightness and we have theoretical understanding of kernels (Wang et al. [0199] Gorham & Mackey| 2017: Liu & Wang, 2018). Mixture approximations have long history of work (Jaakkola & Jordan 1998; Bishop et 1997 Gershman et al., 2012, Miller the intractable mixture ELBO. Van Erven & Harremos (014)) unifies a number of variational techniques by considering them as optimizing different orders of the VR bound. They further demonstrate that two different variants of mini-batch training with the VR bound recover Stochastic EP (Li et al.) (015)) and Black-box a (Hernandez) Lobato et al. 2016), respectively. The Rényi divergence has been studied in other forms under the name ax-divergence (Amaril 2012 Tsallis 1988). Hlernandez-Lobato et al. (2016) introduced a black- box algorithm for variational inference based on the a-divergence using automatic differentiation. Unlike our algorithm, their algorithm is not for HVMs. Rainforth et al. (2088) demonstrated that for VAEs the gradient estimation degrades for salti-sample approximations when using the importance weighted variational autoencoder (IWAE) bound (Burda et al.) 2015). Furthermore, Rainforth et al. (2018) showed that this is not the case when using the ELBO. Rainforth et al. (2018) differs from our work in that the VAEs estimated are with point mass guide, as their inference algorithm is not for HVMs. Le et al. (2020) investigates the deterioration experimentally, providing evidence for it several real world tasks. Tucker et al. (2018) show that by double reparameterizing the gradient estimator, they can eliminate the degrading SN-ratio for multi sauti-sample estimation of the IWAE gradient, among others.
Liu
al. 2017) focusing approximating or lower-boundigg
5 EXAMPLES
We evaluate a-indexed Stein mixture inference by inferring Bayesian neural networks (BNN) and variational autoencoders (VAE) on standard datasets. We use the BNNs for regression on the UCI regression benehmark (the same as HerrHnddezz- Lobato & Adams (2015)) and VAE for unsupervised
learning on MNIST (Salakhutdinov & Murray 2008 LeCun et al. 1998) and OMNIGLOT (Lake
ct al. 2013).
Bayesian neural networks For brevity, we present BNNs for the subset of the UCI regression
benchmark detailed in Appenddi G All datasets use real-valued features. We use a 90-10 split for training and test datasets. We compare a-indexed SM inference for a {0, 0.5, 1} on BNN regression. We test with two guides: factorized Gaussian guides with an EL kernel and point mass (Dirac delta) guides with an RBF kernel. Like Liu & Wang (016),, we layer of size fifty and a RELU activation. We put a Gamma(1, 0.1) prior on the precision of the neurons and the likelihood. We use five particles for all experiments. We run all datasets for 35,000
BNN with one hidden
8

---

Under review as a conference paper at ICLR 2023
Table 2: Average log likelihood (higher is better) test results for BNNs on UCI benchmark. EoS (ours) corresponds a 1, Hell (ours) to a 0.5, and SM (Nalisnick & Smyth 2017) to = 0. Parentheses mean we use Dirac delta guides. EoS generally outperforms Hell and SM.
Average log-likelihood SM
Laplace
Roston
(-0.61
[=1.00 # 1,34)
0.32
(-1.48 1.32)
-0.79
-2.5 + 0.08
-0.74 0.01
-0.50 + 0.03
-0.58 0.25
0.37
0.37
Concrete
(-0.5 1: 0.08)
(-1 0:27)
(-1 (-190.29)
-8.1 1 0.06
-0.54 0.11
00,96 = 0.49
Energy
(-0.25 (.81)
0.022 +
(-202 1.64)
-0.10 + 0.63
(-1611 1.50)
-0.13 0.70
-1.70 1 0.00
0.05 1: 0.12
0.31 d: 0.70
Naval
(-21.00 19.71)
1.72 +..54
(-5.32 0.996)
-0.54 + +0.40
(-4.80 1.044
0.90 + 0.78
3.46 1: 0.04
1,48 1: 1.57
2.05 2: 0.03
--1:23-4-0/06)
-1.38
(-1-43::0.11)
(-1.45
-1.85 0.04
-0.96
-1:35 0.07
-1.02
(-0.20++-0.3 +-.50+
0.46
(-283 2.130
-0.26
(-2,80
-0.27
-2.1 0.63
-0.83 0.05
-0.12
Power
( (-0.20 (::000:::000 (158g
0.04 0.06
[- 0.20 2 0.188
-0.08 1.11.
(- 0.25 (.13)
00.08 0.11
-283 1 0.00
4: 0.00
0.02 2:: 0.07
Table 3: Log likelihood (higher is better) test results for VAE. EoS (ours) corresponds a Hell (ours) to a = 0.5, and SM (Nalisnick & Smyth 2017) to a O.
Hell
MNIST 101.874
100.411
77.401
OMNIGLOT
114
144.257
-18.428
epochs with a subsample size of 32, the Adam optimizer (Kingma & Ba 2014) and a step size of 0.002. All measurements are repeated three times and obtained on a GPL
We compare against the SVGD implantation from Liu et al. (2016) with 20 particles, mean field variational Bayes (MFVI) with factorized Gaussian guide (Hoffman et al.] 2013) and Laplace approximation. For the latter two we inference engines from NumPyro (Phan et al. 2019).
Table shows the root mean squared error (RMSE) on test sets. EoS with delta Guides out perfor- mance baselines and other c-orders, except on Boston. SM and Hell perform similarly with factorized Gaussian guides, which aligns with our SN-ratio experiment that shows the gradient approximations similar for these two cases. Note the Stein mixtures use only five particles, whereas SVGD uses twenty. Table 2 gives the log- log-likelihood on the same test sets. FoS achieves better average log-likelihood for all datasets with factorized Gaussian guides than other o-orders. We see that = and a = _5 performs similarly with a factorized Gaussian prior, which aligns with our SN-ratio experiment in that the quality of gradient approximations is similar.
Variational autoencoder We evaluate Stein mixtures and SVGD for VAEs two datasets with a E {0, 0.5, 1}. We use binarized MNIST (Salakhutdinov & Murray 2008 LeCun et al.) 1998). a dataset of 28 x 28 pixel images of handwritten single digit numbers, and variate of OMNIGLOT Lake et al. 2013), which contains 28 x 28 pixel images of characters from fifty different alphabets. We use the same VAE architecture as Burda et al. (011)), detailed in Appendix H| For both datasets we optimize using the Adam optimizer and learning rate of 5- 10-4 We optimize with a batch size of 20 and use 20 draws to approximate the gradients. For OMNIGLOT we use 20 epochs and for MNIST for E {0.0, 0.5, 1.03. We find that EI.BO-withn-Stere with a = achieves better log-likelshoods on MNIST datasets. On OMNIGLOT 0.5 and a 0 achieve comparable loglliili with a0 the slightly outperforming 0.5.
we use 50. Table 3 show the performance of EL.BO-witin--S5
SUMMARY
We introduce a new algorithm called Se---witt---eiii (EoS) based on new connection between the inference of Stein mixtures and the Rényi variational bound. We demonstrate that Eos provides better gradient approximations than alternative algorithms, which results in better performance for standard benchmark problems. EoS is integrated as black box library in the NumPyro PPL which is distributed freely.
Quadro RTX 6000 with Cuda V11.4.120
9

---

Under review as a conference paper at ICLR 2023
REFERENCES
Shun-ichi Amari. Entrrental-Geomettiiee Methods in Statistics. Springer Science & Business Media, 28, 2012.
Jimmy Ba, Murat A Erdogdu, Marzych Ghassemi, Shengyang Sun, Taiji Suzuki, Denny Wu, and Tianzong Zhang. Understanding the variance collapse of svgd in high dimensions. In International Conference on Learning Representations, 2021.
Javad Behboodian. On the modes of a mixture of two normal distributions. Technometrics, 12(1): 131-139, 1970.
Alain Berlinet and Christine Thomas Agnan. Reproducing Kernel Hilbert Spaces in Probability and Statistics. 2011.
Thierry Bertin-Mahicen Mahieux, Daniel P.W. Ellis, Brian Whitman, and Paul Lamere. The Million Song Dataset. International Conference on Music Information Retrieval, 2011.
Eli Bingham, Jonathan P. Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul A. Szerlip, Paul Horsfall, and Noah Goodman. Pyro: Deep Universal Probabilistic Programming. Journal Machine Learning Research,
2019.
Christopher Bishop. Neil Lawrence, Tommi Jaakkola, and Michael Jordan. Approximating Posterior Distributions in Belief Networks Using Mixtures. Advances in Neural Information Processing Systems, 10, 1997.
Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance Weighted Autoencoders. arXiv preprint:1509.005/1 2015.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive Subgradient Methods for Online Learning and Stochastic Optimization. Journal of Machine Learning Research, 12(7), 2011.
Roy Frostig, Matthew Johnson, and Chris Leary. Compiling Machine Learning Programs via High-Level Tracing. Systems Modeling Language, 2018.
Samuel Gershman, Matt Hoffman, and David Blei. Nonparametric Variational Inference. arXiv preprint: 1206. 4665, 2012.
Jackson Gorham and Lester W. Mackey. Measuring Sample Quality with Kernlls, International Conference on Machine Learning, 70: 10:1292-1301, 2017.
David Harrison and Daniel L Rubinfeld. Hedonic Housing Prices and the Demand for Clean Air. Journal of Environmental Economics and Management, 5-1):81-102, 1978.
Jose Hernandez-Lobato, Yingzhen Li, Mark Rowland, Thang Bui, Daniel Hernández-Lobato. and Richard Turner. Black-Box Alpha Divergence Minimization. In International Conference on Machine Learning, pp. 1511-1520. PMLR, 2016.
José Miguel flemández-Lobbte and Ryan Adams. Probabilistic Backpropagation for Scalable Learning of Bayesian Neural Networks, In International Conference Machine Learning, pp. 1869.1869. PMLR, 2015.
Matthew D Hoffman, David M Blei, Chong Wang, and John Paisley. Stochastic Variational Inference. Journal of Machine Learning Research, 2013.
Pavel Izmailov, Sharad Vikram, Matthew D Hoffman, and Andrew Gordon Gordon Wilson. What Are Bayesian Neural Network Posteriors Really Like? International Conference Machine Learning, pp. 4629 4640. 2021.
Tommi S Jaakkola and Michael Jordan. Improving the Mean Field Approximation via the Use of Mixture Distributions. In Learning in Graphical Models, pp. 163-173. Springer, 1998.
Tony Jebara, Risi Kondor, and Andrew Howard. Probability Product Kernels. The Journal of Machine Learning Research, 5:819-844, 844, 2004.
10

---

Under review as a conference paper at ICLR 2023
Diederik P Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. preprint:1412.6986 2014.
arXiv
Diederik P Kingma and Max Welling. Auto-Encoding Variational Bayes. arXiv 2013.
Brenden M Lake, Russ Salakhutdinov, and Josh Tenenbaum. One-Shot Learning by Inverting Compositional Causal Process. Advances in Neural Information Processing Systems, 26, 2013.
Tuan Anh Le, Adam R Kosiorek, N Siddharth, Yee Whye Teh, and Frank Wood. Revisiting Reweighted Wake-Sleep for Models With Stochastic Control Flow. Uncertainty in Artificial Intelligence, pp. 1039-1049, 2020.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-Based Learning Applied to Document Recognition. Proceedings of the IEEE, 86(1 (0:21)):2772225 23244, 1998.
Yingzhen Li and Richard E Turner. Rényi Divergence preprint:1602.02311 2016.
Variational Inference.
arXiv
Yingzhen Li, José Miguel Hermández-Lobate and Richard E Turner. Stochastic Expectation Propaga- tion. Advances in Neural Information Processing Systems, 28, 2015.
Qiang Liu. Stein Variational Gradient Descent as Gradient Flow. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances Neural Information Processing Systems, pp. 3115-3123, 2017.
Qiang Liu and Dilin Wang. Stein Variational Gradient Descent: General Purpose Bayesian Inference Algorithm. In Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg, Isabelle Guyon, and Roman Garnett (eds.), Advances in Neural Information Processing, pp. 2370-2378, 2016.
Qiang Liu and Dilin Wang. Stein Variational Gradient Descent as Moment Matching. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianche and Roman Garnett (eds.), Advances in Neural Information Processing Systems, pp.
2018.
Qiang Liu, Jason Lee, and Michael Jordan. Kernelized Stein Discrepancy for Goodness-Of-Fit Tests. In International Conference Machine Learning, pp. 276-284, PMLR, 2016.
Andrew C Miller, Nicholas J Foti, and Ryan P Adams. Variational Boosting: Iteratively Refining Posterior Approximations. International Conference on Machine Learning, pp. 2420-2429, 2017.
Eric Nalisnick and Padhraic Smyth. Variational Inference With Stein Mixtures. In Advances in Approximate Bayesian Inference, NeurIPS 2017 Workshop, 2017.
Du Phan, Neeraj Pradhan, and Martin Jankowiak. Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro. In Program Transformations for Machine Learning, NeurIPS Workshop. 2019.
Tom Rainforth, Adam Kosiorek, Tuan Anh Le, Chris Maddison, Maximilian Igl, Frank Wood, and Yee Whye Teh. Tighter Variational Bounds Are Not Necessarily Better. In International Conference on Machine Learning, pp. 4277-4285. PMLR, 2018
PS Rana. Physicochemical Properties of Protein Tertiary Structure Data Set. UCI Machine Learning Repository, 2013.
Rajesh Ranganath, Dustin Tran, and David Blei. Hierarchical Variational Models. International Conference on Machine Learning, pp. 324-334-3333 2234333,222 2016.
Alfréd Rényi. On Measures of Entropy and Information. In Proceedings of the Fourth Berkeley Symposium Mathematical Statistics and Probability, Volume 1: Contributions to the Theory of Statistics, volume 4, pp. 547-562. University of California Press, 1961.
Amit Sabne. XLA Compiling Machine Learning for Peak Performance. 2020.
Ruslan Salakhutdinov and lain Murray. On the Quantitative Analysis of Deep Belief Networks. In Proceedings of the 25th International Conference Machine Learning, pp. 872-879. ACM, 2008.
11

---

Under review as a conference paper at ICLR 2023
Rui Shu, Hung H. Bui, Shengjia Zhao, Mykel J Kochenderfer, and Stefano Ermon. Amortized Inference Regularization. Advances in Neural Information Processing Systems, 31, 2018.
Constantino Tsallis. Possible Generalization of Boltzmann-Gibbs Statistics. Journal of Statistical Physics, 52(1))479-4877 1988.
Athanasios Tsanas and Angeliki Xifara. Accurate Quantitative Estimation of Energy Performance of Residential Buildings Using Statistical Machine Learning Tools. Energy and Buildings, 49: 560-567, 2012.
George Tucker, Dieterich Lawson, Shixiang Gu, and Chris J Maddison. Doubly Reparameterized Gradient Estimators for Monte Carlo Objectives. arXiv preprint: 1810.04 2018.
Pinar Tüfekei. Prediction of Full Load Electrical Power Output of Base Load Operated Combined Cycle Power Plant Using Machine Learning Methods. International Journal of Electrical Power & EEerry Systems, 60: 126-140, 2014.
MillWillem van de Meent, Brooks Paige, Hongseok Yang, and Frank Wood. An Introduction to Probabilistic Programming. arXiv preprint: 1809. 10756, 2018.
Tim Van Erven and Peter Harremos. Rényi divergence and Kullback-Leible Divergence. IEEE Transactions on Information Theory, 60(7):3777-3820, 2014.
Dilin Wang and Qiang Liu. Nonlinear Stein Variational Gradient Descent for Learning Diversified Mixture Models. International Conference on Machine Learning, 97:6576-6585, 09-15 Jun 2019.
Dilin Wang, Zhe Zeng, and Qiang Liu. Stein Variational Message Passing For Continuous Graphical Models. International Conference on Machine Learning, pp. 5219-5227, 2018.
Dilin Wang, Ziyang Tang, Chandrajit Bajaj, and Qiang Liu. Stein Variational Gradient Descent With Matrix- Matrix-Valued Kernels. Advances Neural Information Processing Systems, pp. 7834-7844, 2019.
I-C Yeh. Modeling of Strength of High- Concrete Using Artificial Neural Networks. Cement and Concrete research, 28(12):1977-180; 1998.
A VARIATIONAL RÉNYI BOUND
For convenience, we derive the variational Rényi bound (Li & Turner 2016) in the context of inference with our algorithm below. Recall that Stein mixtures lift the set of guide parametas & (optimized in VI) to a random variable Φ Let D be a finite set of observations, z E be a latent variable, Da [g] [)] the Rényi a-divergence (Rényi 1961) between distributions p and g. and & Rd a set of guide hyper- parameerssssers lifted to a random variable. Then we have
p(z,D) D)
1-07
log PDD) Da (o(2)
log g((((d))
(14)
1
(z((D,))
To see this, consider,
Da (2)(((
log (<)lD,
a -
(r(z.D)
log (z|DD. 6)a
dz
1
((D)
log
(z(D) P(z 1-0 -z p (DD)--
a - 1
1
log
1-2
dz a -
log (D)
0
1
log / (2 (11-a)
p(z, D) 1-a + log
0 -
( p(z, 1-a7
1
(
log (
l())
12

---

Under review as a conference paper at ICLR 2023
from which recover the desired equality by rearranging and multiplying both sides by negative one,
p(z, D) ((2(0)
logp(D) - Da(c())1))(((1) 1 log ()(E(DD)))
To shorten notation, let C (D, E (p(e.)) 1-a
where the expectation is with
respect to q(z)D, Φ). The gradient of the variational Rényi bound with respect to Φ is
1
15 (
p(zzD)
1-a7
p(z,
1-a
V
Caa (D.
E
V
1 a
a
p(z, D)
1-a
(((z D) ((z(D,
= =C Φ) -FE
V log
((]/2,00) の))
D(a.D) ((x(D.+))
D(Z,D)
V log
(D, Φ)
((21D,
q(z|D, 1) )
p(z,
E
(Z, V log
where
wa ( (z. D) (pta.))/((H).dd)) 1-a
(p(e.)) 1-a'
B CHARACTERIZING TWO PARTICLE FIXED POINTS
We give the full derivation of stationary points for the Stein mixture that consider in Section 3 Recall that Seection: investigated the SN-ratio for a Stein mixture given by
2 (N 3/2Ta) + (6 3/2Ta))
where 02 E Rdd are two d-dimensional particles. We use the kernel given by
k()1.0)) = exp ( 10 2212
(15)
where h E RR+ is the bandwidth. The kernel has the following properties:
-V V k Φ2) = - V&, (2) k(..)) = 1,
k Φ2) = k(02 #1),
vak(,.)) = O,
((xxD) 1-o
which we will use the derivation. Finally, we introduce Ea (Ø) I-a log (g(2)) as notation short-hand.
9(x d)
Our two particle configuration reaches fixed point when
(Ø) + c (Q Φ2 e. (Φ²)) (611)22)
where 2 is the step size. Therefore, Stt (Φ1) = 0 and 5 (Φz) at any fixed point. S (Φ1) is given by
$ (Φ)
:() V Vo, (6) k( 2 tk( 22) 62) Va,E((b) Vo (()
(o, k(0, 2 V Ea (b)) Vo, ((0 = 0.
13

---

Under review as a conference paper at ICLR 2023
Therefore,
( (0,5a((1) + k(0 Vo.
(16)
at a fixed point. By a similar argument for 5/ ( we have
V k( 02) = - (V 0. S. 2 k(Ø 02) V (a
(17)
at a fixed point. As V 12) = Va, k(0) 22) it follows from Equations (16) and (17) that
V k( 02) Ve, 02)
a. (2)
k( 22 Ve, En (02)) Vo, 1 Vo, (222)
Va ( (1 + k( 22)) (a) (1 x(
Lo (02) Ve, (o)))
Hence, we that at any fixed point for our two particle configuration, the gradients the VR-bound are equal and opposite.
C KERNELS IN SVGD
For an example of kernel, consider the radial basis function (RBF) kernel k(z,, zz) exp (- 1 || z; -zz
z, 123) with bandwidth parameter h, chosen as log N med(2), where med is the
median operator.
The repulsive force moves particles away from each other, ensuring that they do not collapse onto the same mode. For the RBF kernel, the repulsive force becomes
Bxj
V x (2:
Σ
2 2k(z,, 2z) (z: 2z) ld.
where is the (cuclidean) inner product and 1a is d-dimensional one vector. It follows that z is pushed away from z, when k(z,,2 2,) is large.
D CONDITIONAL EVIDENCE AS RÉNYI DIVERGENCE BETWEEN POSTERIORS
That Equation (11) pushes posteriors towards each other follows from properties of Rényi divergence with a E [0, 1] and the negative direction of the gradient on the Rényi divergence. In particular, we have that the divergence
similarity measure of distributions for 2 so Da=00 q= = (ii) that Da(g [] p] is everywhere positive, and (iii) the divergence is jointly convex (i.e. convex both distributions) [Van Erven & Harremos) 2014). Putting it all together, we see from (ii) and (iii) that the extremum at Da=o ( [p] 0 is global and from the (negative) gradient we minimizing D Da=olo [glpp]. So, we have that V Do (z())D(d)) (DD)) p(z[D)) = = Da=olg(z (D
=
p(z
0 which from (i) means q(z D. p(z(D).
E ALTERNATIVE ELBO-WITHIN- STEIN DERIVATION
For we can derive the attractive force of S directly by applying Jensen's inequality to the log conditional evidence, resulting in
PDD 2() (2
V log Eg
V log
(18)
(2((D,
In ELBO-within Stein, the attractive force takes the simple form
SL FL Ex
[k( Φ) (x()))))))))))))))))()))))) poD, k(o,, Φ) V [a(z[D,
(19)
and the repulsive force is given by Equation (5)
14

---

Under review as a conference paper at ICLR 2023
Table 4: Summary statistics of datasets from the UCI regression benchmark.
Dataset
Data points
Feature count
Boston (Harrison Jr & Rubinfeld 1978)
506
Concrete (Ych) 1998)
1030
8
Energy Tsaaa & Xifara 2012)
768
Power Tütekci 2014
9568
4
Protein Rana 2013)
45730
Year (Bertin-Mahreux et al.
2011)
515345
90
Table 5: Variational autoencoder architecture for MNIST and OMNIGLOT. s denotes a stochastic
layer and d denotes deterministic layer. Read the networks left-to-righh for guide description and right-to-left for model description.
Dataset MNIST d200-d200-s50
Architecture
Activation
tanh
OMNIGLOT
6200-000-:100-1110--000000
tanh
F ILLUSTRATING STEIN MIXTURES
To illustrate the use of a SM, consider the variational autoencoder (VAE) (Kingma & Welling 2013). The VAE simultaneously trains a generative model p(D) (z(z))p(z) and variational approximation (z))fa as (D)) of the posterior p(2 D). Here, 0 and r/ are parameters of the generative neural network so(-) and the inference network f respectively. VAE training typically done by stochastic variational inference (SVI) (Hoffman al. 2013) which optimizes I and & to minimize the ELBO. With a SM, the generative model remains the same, that is. we obtain point estimate of 0. However, the marginal posterior approximation changes to 1/10 Esco ((z)f((D) So with Stein mixture, each particle & parameterizes separate inference network, i.e. fo (.), meaning the guide becomes amortized similar to Shu et al. (2018).
G UCI BENCHMARK DETAILS
We compare ELBO-within-Steee for E (0,0.5, 1) on BNNs regression point mass (Dirac delta) guide and RBF kernel. With ELBO-Within-Stee we recover variant of SVGD with VR gradient rather than the score function. Like Liu & Wang 2016).
BNN with one hidden layer of size fifty and a RELU activation. We put a Gamma 1, 0.1) prior on the precision of the neurons and the likelihood. For both versions epochs and 500 epochs for the rest. We use Adagrad (Duchi et al. 201T) with a step size of 0.05 and a subsample size of 100. All measurements are repeated five times and obtained on a GPU
5 partieles and update Year for 40 epochs, Protein for 100
H VAE DETAILS
Following Li & Turner (0066] Burda et al. (2015); Rainforth et al. (2018), we use VAEs with multiple stochastic layers. The idea define the model through ancestral sampling as
p(x(0) Σ p(zz)p(zL-11) fo... (ZL)) p(x foa (zz)),
where x is data-point (which we will also denote 2o). Zi. Zy are the stochastic layers, and 0 parameterizes a neural network fi which takes Zi+1 to the parameters of the distribution pt, i.c. ((zi).(((z4))) We then let the guide factor in the opposite direction, resulting in
q x) q(zz)ff, (x))) (zi)) ... .p(z fe, (ZL-1)).
Quadro RTX 6000 with Cuda V11.4.120
15

---

Under review as a conference paper at ICLR 2023
We the same network architecture as Rainforth et al. (2018) (summarized in Table 5). In Table 5 s denotes a stochastic layer and d denotes a deterministic layer (affine transforms). We use tanh as the activation functions deterministic layers. Stochastic layers distribute according to factorized Gaussian distribution, and for the likelihood we use the Bernoulli distribution (hence the binarization of the datasets).
I THE EINSTEINVI LIBRARY
We provide a library called EinSteinVi for inferring Stein mixtures in the probabilistic program- ming language (PPL) NumPyro (Bingham et al. 2019 Phan al. 2019). EinSteinVI uses a-indexed SM inference as its core algorithm as described in Section 3) NumPyro is a universal PPL (van de Meent et al. 2018) embedded in Python. NumPyro provides specialized constructs for expressing probabilistic models as Python programs and allows executing arbitrary code in its model and guide. The computational backend of NumPyro ro is Jax (Frostig al. (018), which combines XLA (accelerated linear algebra) (Sabne 2020) program transformations with automatic differentiation. As EinSteinVI works with arbitrary guides, NumPyro a well-suited language for embedding EinSteinVI. Further, we chose NumPyro because:
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.
1.1 A EINSTEINVI PROGRAM EXAMPLE
To demonstrate the two modes of VI VI and Stein mixtures) with EinStein VI, we consider the ID Gaussian mixture 1/3N(-2,, 1) + 2/3N(2, 1) (see Figure 3 and Figure 4). The Gaussian mixture is bi-modal and well-suited for the nonparametric nature of SVGD and Stein mixtures. Figure 4 shows that both SVGD and the dieinemixttre naturally capture the bi- bi-modality of the target distribution, compared to SVI with a Gaussian guide. Note the reduction in particles required to estimate the target when using Stein mixtures compared SVGD, Also, note that the Stein-mixture overestimates the variance and slightly perturbs the locations. The еттог seen at the right mode for the Stein-mixture with two partieles due the uniform weighting of the particles in SVGD, and as such is algorithmie. The Stein-mis Stein-mixture will therefore not be able to exactly capture the mixing components of target mixture model with one particle per component. However, with more particles the mixture can be approximated better as demonstrated using three particles.
1.2 INTEGRATION WITH NUMPYRO
Integrating Ei cinSteinvV with NumP requires handling transformations between the parameter representation of NumP yr and the array representation that HLBB--ithhin- operates on. For this, we rely on Jax's PyTrees which converts back and forth between Python dictionaries and array representations.
Algorithm shows the black-box version of a-iindexed SM inference in NumPyro. The algorithm allows SVI to estimate a subset of the parameters and a-indexed SM inference the rest. To differentiate the two, we denote parameters updated by SVI with ob and parameters updated by ELBO-within-tted -Sein with Φ (i.e. the Stein particles $ { N 1). In the model, only SVI can update parameters which we denote by 0. We update 0 and &a by averaging the loss over the Stein particles. For the Stein particles, the process is more elaborate. First, we convert the set of individual distribution parameters the guide to monolithic array using Jax's PyTrees. The array represents the particles as
WWe recover SVGD with point mass (Delta dirac distribution) on all distributions in the guide. A dictionary mapping parameters to their values, which can be of arbitrary Python type https://jax rewdthedocs..eeeee wtt.s/.ooo///////eeeeeee wttp:///////////eeeeeeeee/
htm1
16

---

Under review as a conference paper at ICLR 2023
def model () samplef ( 'x', NormalMixture (jn.. ([1 / 3, 2 31),
jnp. anp.array ([-2.0, 2.01),
jmp.array ([1.0, 1.01)))
(a) ID Gaussian mixture model
svi SVI (
stein SteinVI model,, AutDelta (model),
moael, oNermal (model),
Adagrad setee__iieee_00
Adagrad (step_iiz=11100 Trace_BL30(), RBFKernel
Trace_ELBO()
results- = svi run (ring nun_iterations)
results stein run (rng_key, nun_iterations)
(b) SVI
(c) SVGD with EinSteinVI
Figure 3: ID Gaussian mixture model in NumPyro. We use the deprecated Norma Mi xtu re over the more general (and more verbose) Mtureeaee ami ly for clarity.
0.1
0.0
0.0
(a) SVI
(b) SVGD (RBF kernel)
(c) Two partiele Stein- mixture (linear kernel)
(d) Three particle Stein- mixture (RBF kernel)
Figure The blue dashed line is the target pdf, while the solid green line is the density of the particles. We estimate the particle density for SVGD with Gaussian kernel density estimation. We use 100 particles for SVGD, and two or three particles for the Stein-mixture. SVI uses Gaussian guide.
17

---

Under review as a conference paper at ICLR 2023
flattened and stacked Jax array. Then, compute a kernel on the particles, delegated to the kemell interface (see AApenddix 1.3) as the computation semnel-dependen..
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. As we compute the Stein forces in unconstrained space, must correct them by the Jacobian of the bijection to constrained space. Naively computing the Jacobian on the monolithic array incurs a massive memory overhead. However, as NumPyro registers 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. Like computing the Stein forces, the correction emburrassingly parallel so that we can use a vmap operator again. Inside the vmap we nest tree map to do the appropriate conversion between representations. Finally, we convert the monolithic array to NumPyro representation and return the expected changes for SVI- and Stein- parameters.
Require: SVI parameters 0 and vh Stein parameters { model Ps (z, x), guide ge. (z), loss La kernel interface KI.
Ensure: Parameter changes based on SVI (40, Axp) and hierarchical Stein forces (f 1
procedure UPDATE(0, th, { Pe. 2
AO Eg VV (po. {o.6)
A&b + E (Po. l
{a, ; t PyTreeFlatten( ( (0)))))
< (T((a:)E
procedure HSTEIN-FORCESGE
Calculate forces per particle for higher-order vmap function.
0. + PYTREERESTORL(a))
Aa; EEa k(aj, aa)j,,VVa,, Lo (pp. 90. + Va, k(aj, ac)
return Aa; end procedure
(aa h VMap((a J. HSTEIN-FORCES) A 1. il + (ATREERESTORE
return AO, Asp, { end procedure
Algorithm 1: a-indexed Stein Mixture inference
1.3 KERNEL INTERFACE
The kernel interface straightforward extend the interface, users must implement the compute function, which accepts as input and particles, and the loss function L and returns differentiable kernel Table gives the complete list of kernels in EinStein
the current set of particles, the mapping between model parameters
18

---

Under review as a conference paper at ICLR 2023
Table 6: Kernels included in the EinStein library.
Kernel
Definition
Comments
Type
Reference
Liu & Wang
Radial Basis Function (RBF)
exp(1 x y (2)
scalar
(2016)
exp( 1/2 (x y))
vector
Pyrd Gorham &
Inverse Multi-
Quadratic
(c2+ x - y 113)
8 E (-1,0) and c 0
scalar
Mackey (2017)
(IMQ)
Φ(x, w
Random
22 cos( w + wo) where wo Unif 2x)
Liu & Wang
Feature
E ( w w)Φ(y, w)]
scalar
(2018)
Expansion
and NV(0 La)
Liu & Wang (2018)
Linear
x +1
scalar
scalar, vector,
{k) individual kernels, weights w;
Liu & Wang
Mixture
(2018)
matrix
Scalar-based
Wang et al. (2019)
Ia
h scalar-valued kernel
matrix
Matrix
Vector-based
Wang el al. (2019)
diag(k(xx)))
k vector-valued kernel
matrix
Matrix
matrix, placed on the diago-
{K(5) scalar- valued kernels, each for a unique partition of latent variables
Graphical
(iag((A(00(x))))
Wang et al. (019)
nal
K is inner matrix- kernel and Q is preconditioning matrix like the Hessian - V log p(ž]x) matrix
Constant Pre- conditioned
Q=tK(Qix.Qty)0=+
Wang et al. (2019)
Fisher information
E2z (x) [V2 log p(2]))]
matrices
[ar m is a set of anchor points, Q₂ ((aa) a preconditioning matrix for each anchor point, Ko is an inner kernel conditioned using Qé, and wx(x) softmaxx( ([N((x)))(e))))) ag, QF )e))
Anchor Point Precondi- tioned
Wang et al. (2019)
Σmh
matrix
19


Reviews:
[{"review_id": "oqk_TLZmEs", "reviewer": "Reviewer_k1hY", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper aims to improve the numerical stability of stein mixture in variational inference. To achieve this, the proposed EBLO-within-stein method allows the choice of a class of hierarchical attractives forces indexed by the hyper-parameter $\\alpha$ in the Renyi divergence. Different values of $\\alpha$ acounts for different types of latent variables, which could be measured by the signal-to-noise ratios. The paper includes empirical results to show various effects of $\\alpha$ on real data analysis.", "summary_of_the_review": "Although more empirical results can make the paper much stronger, this is an overall well-executed paper. The authors need clarify on a few raised questions.\n\n", "strength_and_weaknesses": "Strengths:\n\n- The paper is overall well-written and the method is clearly explained.  The reviewer really like the writing in section 2 that builds up the knowledge step by step\n- The literature review is thorough\n- The connection with the Renyi divergence and the VR bound seems interesting\n\nWeakness: \n\n- Although the connection with existing works is interesting, the reviewer is not fully convinced why ELBO-within-Stein is preferred over the existing ones.\n\n- given a real data analysis, how should one pick the value of $\\alpha$\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ieaJGQWWmR", "reviewer": "Reviewer_Ra8d", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper investigate a pevious work on stein mixtures and show that it corresponds to inference with the Renyi $\\alpha$-divergence for $\\alpha=0$, and using other values of $\\alpha$ may stablize the inference.", "summary_of_the_review": "My main concern is that the original contribution of this work is not significant enough.", "strength_and_weaknesses": "Strength:\n\n1. This paper is overall well-written and the literature review is thorough.\n\nWeaknesses:\n\n1. The motivation of the proposed method is vague. No comparison to other SVGD variants.\n2. The orginality is low as the major contribution seems to be some empirical investigation of a hyperparameter.\n3. The paper is a little bit hard to follow. What is the main difference between stein mixture and a simple mixture model?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Ln1ZjaU2TR", "reviewer": "Reviewer_r9He", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors generalized the Stein Mixtures approach of Nalisnick and Smyth 2017 to general alpha-divergences.  The paper is mostly nicely written and was a good read.  But it did not have a clear motivation nor impressive empirics.", "summary_of_the_review": "The paper is an enjoyable read about extending Stein mixtures to use alpha divergences.  But I recommend rejection because in my assessment it falls short on:\n* Motivating how or why the proposed method addresses limitations of existing methods (particularly dimension dependence) and\n* Empirical demonstration of improvement relative to existing methods.", "strength_and_weaknesses": "Strengths\n* The explanation and demonstration of the trade-off between the need for a large number of Monte Carlo samples when estimating the gradients without too much bias / variance is nicely laid out and is clear.\n* The introduction of alpha-divergences into Stein mixtures is a nice application of a compelling line of work on generalized variational inference and I was glad to see it explored here.\n* The worked out toy example with the analytic fixed points is very nice and illuminating!\n\nWeaknesses:\n* The authors should better defend a key, undefended premise of their work about which I was unconvinced: that Stein mixtures alleviate an “exponential” dependence of SVGD on the dimensionality on the model.  This premise is not made precise nor is it attributed to previous work that establishes it more precisely.  Why is an exponential number of particles needed for SVGD?  This is certainly not the case for MCMC for example.\n* If there is a situation where the dependence is exponential with dimension, would the authors please explain (ideally with a theoretical result) such a situation?\n* Or if there is no theoretical justification or proof to be made about the better dependence on dimension, empirical validation of this claim is needed.  For example, it would help to have an illustrative example of a case (e.g. in simulation) where the dimension scaling is exponential with the dimension. E.g. perhaps a plot with an increasing dimension on the x axis and the number of particles needed to achieve some level of posterior precision on the y-axis (for all three of the SVGD method, Stein mixtures, and the new proposed method).\n* Or alternatively, if this exponential dependence on dimension is known from prior theory or empirics, even replicating those prior results (with citation) would strengthen the paper.\n* Or if the dependence is not actually exponential, the authors must either be precise about the way in which dimensionality poses a problem that is addressed by their approach or drop this as the premise of the method.\n\n* Given that the paper is about an improvement upon Nalisnick and Smyth 2017 (whose applications and limitations most readers may not be very familiar with), the paper could benefit from more justification as to why this method is important/useful and why it must be improved.  At present, I found it challenging to see what problem the new methodology was trying to solve that Nalisnick and Smyth 2017 had not solved.\n* On problems where the developed methodology is applicable, how does it compare to other Bayesian inference methods?  E.g. HMC, MFVB, Laplace approximations etc., in terms of computational cost and accuracy (experiments are not necessarily needed, but I think commentary at least is).  For the UCI comparisons, previously run baselines could be used.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject