ELBOing Stein: Variational Bayes with Stein Mixture Inference
Authors: Ola Rønning
Eric Nalisnick
Christophe Ley
Padhraic Smyth
Thomas Hamelryck
Venue: ICLR 2025 Poster (2025)
Keywords: variational Bayes
particle-based inference
mixture models
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We present Stein mixture inference, a particle-based inference method, that mitigates variance collapse in moderetly sized models.

Abstract: Stein variational gradient descent (SVGD) (Liu & Wang, 2016) performs approximate Bayesian inference by representing the posterior with a set of particles.
However, SVGD suffers from variance collapse, i.e. poor predictions due to underestimating uncertainty (Ba et al., 2021), even for moderately-dimensional models
such as small Bayesian neural networks (BNNs). To address this issue, we generalize SVGD by letting each particle parameterize a component distribution in
a mixture model. Our method, Stein Mixture Inference (SMI), optimizes a lower
bound to the evidence (ELBO) and introduces user-specified guides parameterized
by particles. SMI extends the Nonlinear SVGD framework (Wang & Liu, 2019) to
the case of variational Bayes. SMI effectively avoids variance collapse, judging by
a previously described test developed for this purpose, and performs well on standard data sets. In addition, SMI requires considerably fewer particles than SVGD
to accurately estimate uncertainty for small BNNs. The synergistic combination of
NSVGD, ELBO optimization and user-specified guides establishes a promising
approach towards variational Bayesian inference in the case of tall and wide data.

Published as a conference paper at ICLR 2025
ELBOING STEIN: VARIATIONAL BAYES WITH STEIN MIXTURE INFERENCE
Ola Rønning
Eric Nalisnick
Department of Computer Science University of Copenhagen ola@di, faaddikk.dde ku.dk
Department of Computer Science
Johns Hopkins University halisnicke jhu edu
Christophe Ley
Padhraic Smyth
Department of Mathematics University of Luxembourg christe Shristophe. leyêuni. lu
Department of Computer Science University of California, Irvine smyth cs uci edu
Thomas Hamelryek
Departments of Computer Science / Biology University of Copenhagen
thamelryEbio. ku.dk
ABSTRACT
performs approx- imate Bayesian inference by representing the posterior with of particles. However, SVGD suffers from variance collapse, i.e. poor predictions due to under- estimating uncertainty (Ba et al. 2021), even for moderately-dimensiin such as small Bayesian neural networks (BNNs). To address this issue, we gen- eralize SVGD by letting each particle parameterize a component distribution a mixture model. Our method, Stein Mixture Inference (SMI), optimizes a lower bound to the evidence (ELBO) and introduces user-specified guides parameterized by particles. SMI extends the Nonlinear SVGD framework (Wang & Liu 2019) to the case of variational Bayes. SMI effectively avoids variance collapse, judging by previously described test developed for this purpose, and performs well stan- dard data sets. In addition, SMI requires considerably fewer particles than SVGD accurately estimate uncertainty for small BNNs. The synergistic combination of NSVGD, ELBO optimization and user-specified guides establishes a promising approach towards variational Bayesian inference in the case of tall and wide data.
Stein variational gradient descent (SVGD) Liu & Wang 2016)
models
INTRODUCTION
Accurate and safe machine learning necessitates adequate uncertainty estimation to ensure reliability in critical applications such as autonomous vehicles and medical diagnosis. Even in problems considered solved, such as protein structure prediction, better handling of uncertainty particularly in cases of conformational heterogeneity (Gavalda-Garcia et al. 2025) could lead to further improvements. As current deep methods are known to be overly confident in their predictions (Szegedy et al. 2014 Nguyen al., 2015). a more principled treatment of uncertainty necessary. Bayesian probabilistic models are attractive as they assess model uncertainty through coherent framework of updating data-based beliefs. However, Bayesian inference for complex models is often analytically and computationally intractable. Therefore, variational Bayes methods approximate Bayesian posterior with a tractable variational distribution (Jordan et al. 1999 Blei et al. 2017).
Particle-based inference is attractive approach to variational Bayes because it resides an intermediate between variational and sampled- based methods (Saeedi et al. 2017 Domke 2017). As hybrid method, particle- based inference combines several desirable properties: sample efficiency, deterministic updates and asymptotic unbiasedness. Primary among particle variational inference algorithms is Stein variational gradient descent (SVGD) Liu & Wang (2016) due to its tractable and

---

Published as a conference paper at ICLR 2025
((01)5))
θ
00
p ~ 1/m Er- (0)
p( D) 1/m Ee ((11))))
(a) SVGD
(b) SMI
Figure 1: Variational inference with SVGID-derived particles (Liu & Wang 2016) versus with an SMI-derived probability density, formulated as mixture model (this work). Left: SVGD uses m particles O₂ to approximate the posterior p(0) (D). Right: SMI uses mixture model (with uniform weights) of m guides (((bb))) parameterized by particles 3 to approximate p)0 As a result, SMI approximates Bayesian posterior with richer model that alleviates variance collapse in higher dimensional posteriors.
simple update rule. However, SVGD suffers from underestimating variance, also called variance collapse (Ba et al. 2021 Zhuo al. 2018). Overcoming the collapse with SVGD requires using more particles as the model size grows. We will demonstrate that this quickly becomes computationally infeasible with off-the- shell hardware, even for moderately sized models such as small BNNs.
To address the issue of variance collapse in SVGD, we introduce Stein mixture inference (SMI I SMI lets each particle parameterize a component distribution, which we call a guide, resulting in mixture approximation of the posterior. In contrast, SVGD directly represents approximate samples from the posterior using its particles. Figure schematically distinguishes the two methods. The mixture approximation allows SMI to represent neighborhoods of SVGD particles, thereby scaling better with model size. We show that SMI is novel variant of Nonlinear-SVGD (NSVGD) (Wang et al. 2019) applied to the variational approximation of Bayesian posteriors. SMI combines ordinary mean mean--ield variational inference (OVI) (Jordan et al. 1999. Hoffman et al. [2013] Ranganath et al. 2014) with SVGD through the NSVGD framework. Specifically, our article makes the following three contributions:
We introduce SMI and show that it extends NSVGD to variational Bayes.
2. We empirically demonstrate that SMI is more particle efficient than SVGD.
3. We use synthetic and real-world data to show that SMI does not suffer from variance collapse in small- moderately-sine sized models such small BNNs.
Next, in Section 2
we will motivate SMI by outlining the reasoning behind the method.
STEIN MIXTURE INFERENCE IN A NUTSHELL
We aim to construct richer variational approximation (0) of the posterior (() than the one offered by SVGD, while ensuring we also have means to optimize To achieve this, we express q()) as uniform mixture model of m (user-defined) guides, parameterized by particles {1b, ) m that make up the empirical measure Pmn()) = //E Sg. (.),
pm
d b))
'This article extends our preliminary work presented in the (non-archival) workshop paper Nalisnick & Smyth (207))
2

---

Published as a conference paper at ICLR 2025
The goal is optimize the corresponding mixture ELBO, which measures how well the mixture model approximates the true posterior,
Σ ())))) log p(0,2) D) ((O)pm)) I Ilogp(D).
((
(2)
Now, the mixture ELBO can be interpreted symmetrid fanctional F(pm), mapping the particles of the empirical measure to scalar,
F(Pm) [(m))
This interpretation allows us leverage the NSVGD framework optimize Flem]. along with an additional weighted entropy term [H.] to encourage particle diversity, to find
p(0, D) (e(0)e)) log q(0)pmm pm
Pm arg max F F(pm) + a.H pm arg max Pm
+ a:H pm)):
(3)
Pm
where Hi = log denotes the differential entropy and a 2 0. As we will show, despite the inclusion of the entropy term, we still obtain a proper ELBO, CsMI(Pm) Flpm] + - < ((D), if we choose a = 1. This ensures the mixture model g(0)pm) provides a well-justified, diversified posterior approximation.
BACKGROUND
We state the variational objective that NSVGD maximizes, restate the central result from Wang & Liu (2019) that allows us to move pm in theorem 3.1 and finally, in eq. (7). give the tractable iterative update that is the backbone of NSVGD optimization.
After introducing OVI, we detail NSVGD in section 3.2
Notation Let p(x) denote sample generated from unknown distribution p(x). We observe N independent and identically distributed draws from p(x) that constitute the dataset
D = (x a latent variable. Let p(0) denote the prior and ()) the posterior. We assume that the posterior not analytically available except up to constant of proportionality. i.e., p(0]D)
We denote the likelihood function p(D 0) II p(x, ()) where 0 O C Rd
p(D,0). We denote the differential operator as Ve when differentiating with respect to (wrt.) the l'th (random) variable. For example, /i, b) is the differential wrt. a. However, if this notation becomes ambiguous, we use the symbolic subscript notation, e.g. VoJ(a, b) = (a,
3.1 ORDINARY VARIATIONAL INFERENCE
We can approximate an intractable posterior (D) with tractable variational distributior q D by optimizing a lower bound on the evidence pDD), the aforementioned ELBO (Jordan et al. 1999). The ELBO given by
logp(D) 2 E log 0)) + ((( = L(bb).
Maximizing C(bb) is equivalent to minimizing the Kullback-Leible (KL) divergence between the approximate and intractable posterior, but importantly requires computing the joint density p(D, 0) instead of the intractable conditional p(0 D). & is obtained from maximizing the ELBO by gradient or coordinate ascent.
3.2
NONLINEAR STEIN VARIATIONAL GRADIENT DESCENT
Like Markov chain Monte Carlo (MCMC) methods, particle variational methods (Frank al.
2009
Saeedi et al. 2017) approximate samples from the posterior rather than its density. However,
AA function is symmetric if its evaluation is independent of the order of its parameters.
3 Although the differential entropy is formally undefined for an empirical measure, within the NSVGD framework (L et al. [2017] Wang Liu (2019) pir converges weakly to p m - no. Hi(p*] is well defined. *We use the notation p(cla) for conditioning the random variable and p(c; b) for density with parameters b.

---

Published as a conference paper at ICLR 2025
Table 1: NSVGD generalizes SVGD and includes DivMM and SMI (this work). Like SVGD, SMI approximates general posteriors. But where SVGD represents the posterior directly with particles 0g. SMI addresses variance collapse by using mixture model 1/m Σi=1 g( parameterized by particles br. On the other hand, DivMM specializes NSVGD to diversified maximum likelihood estimation for mixture models and cannot approximate general posteriors; moreover, the number of particles in DivMM ís a hyperparameter of the model, unlike in SVGD and SMI, where it relates to the posterior approximation 's richness.
Method
Posterior approximation
Model
SVGD (iu & Wang 2016
122
_ (0)
(T
:=1
DivMM Waang & Liu 2019
None
Σ P 0z)
SMI (This work)
2E=1 (0lbb))
p 0 (T)
unlike MCMC methods, the number of posterior samples is fixed a priori for particle methods. Particle methods are attractive due to their freedom from strong parametric assumptions and resulting flexibility as an approximation. We will designate the samples as "particles" " to emphasize that they not auto-correlated, as with MCMC methods. However, they retain some correlation in the non-asymptotic case (Gallego & Insua 2018).
Wang & Liu (2019) introduces the Nonlinear SVGD framework which allows funetional optimiza- tion under diversification constraints. Wang & Liu (2019) applies this general framework to the (constrained) maximum likelihood estimation of diversified mixture models (DivMM), i.e., mixture models consisting of spread-out mixture components. The NSGVD framework has previously only been applied to SVGD and this type of maximum likelihood estimation.
The variational objective NSVGD iteratively moves an initially simple distribution such as Gaussian, according to T(p) p + cølp] such that the transported distribution Tpp) maximally increases a variational objective. The variational objectives for NSVGD combine a functional, like the likelihood (DivMM), the negative Kl--divergence to a posterior (SVGD) or an ELBO (SMI and OVI) with a diversification constraint on We need the constraint to avoid mode collapse.
In its most general form, NSVGD solves the maximization given by
= arg max F(p] + aHH(p),
(5)
where F[p] nonlinear functional of p, H[p] is the differential entropy and 2 scales the contribution of the entropy. The entropy acts as a regularizer, forcing P to distribute uniformly. promoting particle diversification and avoiding collapse to the closest mode.
Making the optimal perturbation di tractable Finding the steepest perturbation direction di [p) challenging in the general setting. This is because di [p] requires computing the first variation of F, functional analog to the derivative of a function, which may not exist, let alone be computationally tractable. We must weaken our optimization and add additional structure on P and F' to progress toward tractable algorithm by guaranteeing the first variation is always tractable. Theorem from Wang & Liu (2019) provides this structure. First, use an empirical measure Pm on m particles 19 instead of p because with the particles evaluating wrt. Pm is trivial. However, using Pm for means we only approximate the p" from eq. (5) with the guarantee that the optimum Pmm weakly converges to p" when letting m + (Wang Liu 2019 Liu al. 2017). Second, choose F[pm] such that there exists a symmetric and differentiable map f: 01, 02.00..... Oi + F[pm]. Under these two conditions, the first variation of F' wrt. the l'th partiele reduces to ordinary differentiation of ( .2....................... 0m) .
which is tractable to compute.
The optimal perturbation in RKHS
With the additional structure on p and F, the last essential component that gives Ø² closed form is restricting the candidate perturbations to functions in reproducing kernel Hilbert space Liu & Wangl (2016). With these components, we can restate the theorem that gives the closed form for optimal perturbation as:
Theorem 3.1. The Kernelized Steepest Perturbation Waag & Liu 2019) Let F(p) + aHlpp be the variational objective for a transport T(p) p + eΦpp), with e > and distribution p with
4

---

Published as a conference paper at ICLR 2025
supp p S dom f C Rd. Let Pm (mm (.) be the empirical measure of m particles and let f (o...... mm) + Fl be differentiable and symmetric function. If we choose a reproducing
kernel & on R x Rd with reproducing kernel Hilbert space H (Brllintt & Thomas- Agnan 2011 such that Vik and 22 exist and are both continuous, then the optimal perturbation direction o" E H
such that 6= 1 1 satisfies
(²(+) Be pm [k(0, W.V..OO..... V 0m.) + V 1O0,
(6)
Theorem 3.1 combines Theorem lb and 2 from Wang & Liu 2019) and provides the closed form for the optimal perturbation direction ø used in SMI when replacing f with [(pm). The first and second terms that constitute (-) in theorem 3.1 are commonly referred to as the attractive and repulsive force. This is because the first term pulls particles towards the nearest maximum in F, whereas the second keeps particles from collapsing onto each other. Finally, notice we never need to evaluate H (Pm): instead, di uses the kernel gradient in the repulsive force. This sidesteps the issue of not having (pm) available.
The iterative optimization algorithm The NSVGD iterative optimization starts with particles (0, p m according to
T==1 drawn from the simple initial distribution o p(O), At each iteration, every particle moves
0 +
10,,, V. 8m)
+ a ((),
m
which maximally increases the change in eq. (5) by theorem 3.1 We see that SVGD is an instance of NSVGD by replacing f with the expected log joint density, f o 0m) = Eg [log p (D, 0)]. In the case of SVGD, (o) ,.,,. ..,...,0.. reduces to the scaled score function (I, V i log p(0:, D).
Diversified mixture models In the case of DivMM, Wang & Liu 2019) applies the NSVGD framework the functional Flo] EE=DD (log S0>plp((;; {)] DivMM needs the nonlinear form of theorem 3.1 as its functional is not linear in p. The optimization is a constrained maximum likelihood estimation of the diversified mixture model, Σm p.) Here, the number of particles m is a hyperparameter of the model, whereas. for SVGD and SMI, the model is independent of m. In table [ we contrast SVGD, DivMM, and SMI. which are different applications of the same NSVGD framework. In the following section, we introduce SMI and demonstrate that the NSVGD framework can be adapted infer approximate variational posteriors of general Bayesian models.
STEIN MIXTURE INFERENCE
The key to justifying SMI is showing we can optimize the SMI ELBO (SMI(Pm) given by eq. using NSVGD and that it is indeed an ELBO. To this end, we first show that the mixture ELBO [(pm) given by eq. 2 is a differential symmetric function. That means we can use theorem 3. to find the Pt, that maximizes LsMi (Pm) by iterating eq. 7 Next, show that when a = 1 in eq. the resulting quantity (PM)(Pm) is indeed an ELBO, despite the addition an entropy term to the mixture ELBO [(pm) given by eq. (2).
The SMI function(al) and its gradient The mixture ELBO L(pm) given by eq. is symmetric function wit. Pm due the outer sum. If ((m)) also differentiable, we have the desired mapping required to use parametric differentiation to optimize eq. (2) using NSVGD. To show L(p.) is differentiable wrt. the l'th particle, we can compute the gradient as
p(0, D)
V き, ((pm) = Eg(0
log E, ((0bb,,)
Vv, log ((0bb))
(8)
tb (bbbbb))))))
ΣØ ((0bb)
we choose g(((b)) ( to be differentiable wrt. th, we see that C(pm) is also differentiable wrt. sb. The complete derivation is given in the appendix.
5

---

Published as a conference paper at ICLR 2025
SMIs variational objective is ELBO ((mm) is ELBO, but does SMI indeed maximize ELBO? For this to be the case, (pm)( = + which includes an entropic regulariser, must also be an ELBO. We show this by generalizing Pm to the continuous case because, as noted previously, H(m)) is technically undefined for finite particles. First, consider the continuous version of (Pm) given by,
log D(0,D) p())mm))
p(0, D)
[pp
lim mh- +0-200
Σ R whe)
Byjers) EEp log
- log ((D).
772
(((bbb) (
Note that C(pm) weakly converges to Lig] for m - oo when using NSVGD. Next, we construct an upper bound L of Llp)].
p(0, D)) log
[1 = Bp(()
2 [[p),
(00 2b)
as we show in the appendix. Now, applying SMI's variational objective given by eq. indeed results in ELBO, shown by
3 to [T with
log p(0, D)
log p(0,D) ((0b)))({{)
[p] + Hpp = Ep(()
+ Hppp = Ex(w)
q00 bb)
p D)
< log Bp(w) (0()))
= log P(D).
Lq(0 (((b)b))({}))
Here, the inequality comes from repeatedly applying Jensen's inequality. The final equality results from simply marginalizing. Note that the objective is not an ELBO for a * 1. Finally, we can conclude that [p Lip) + H[[p] is also an ELBO, because
Clp) S ct [pl
Clpl + Hpp S cT Hip) S ID))
Thus, the maximizing empirical measure obtained by NSVGD, Pr weakly converges to correspond- ing ELBO (p") when m - sc. Our experiments indicate that SMI generally insensitive a; therefore, we recommend using a 1 to tie the optimization to ELBO maximization.
The iterative optimization algorithm Optimization starts with the empirical measure Ph on particles { p²) } drawn from a simple initial distribution p". We subsequently iterate the
following gradient ascent-like step on the particles, which by theorem 3.1 maximize eq. (3):
172
14 4-1
Σ
k(
2 V ((p)))
Σ
V ViK( w5)
(9)
We continue the optimization until we reach a fixed particle configuration. In eq. 9.. i the step size, & is reproducing kernel, and IR + is the hyper-parameter inherent NSVGD. a controls the spread of the particles (i.e., by scaling H[[pm] and is, together with the kernel k and step size, the syper-parameters of SMI.
Connection to SVGD, OVI and maximum posteriori SVGD, OVI and maximum posteriori (MAP) estimation are all instances of SMI. particular, where SVGD reduces to MAP estimation when only using one particle, SMI reduces to ordinary variational inference (as in eq. (4) in the
single-particle case for an arbitrary guide (00l)). We can also connect SMI to SVGD, and by
extension MAP estimation, by choosing each guide q(0lzb,) as the point mass, i.e., 15, (0). These connections place SMI as a hybrid between a sample- and density-based method. We attribute SMI's ability mitigate variance collapse to this hybrid nature. In appendix A detail how SMI can be reduced to recover SVGD and OVI.
Library implementation We provide open-source implementation (under an Apache version license) of SMI, called VI. the deep probabilistic programming language NumPyro (Phan al. 2019

---

Published as a conference paper at ICLR 2025
5 RELATED WORK
There has been flurry of work on SVGD, but much of it has concerns that are orthogonal ours. The SVGD algorithm itself has been extended to include second-order information (Detommaso)
2 operate on Riemannian manifolds (Liu & Zhu, 2018) and forgo analytic gradients (Han & Liu 2018). Furthermore, SVGD has been pe-puroseed to perform message passing (Wang
al.,
2018, Zhuo et al., 2018), importance sampling (Han & Liu, 2017), generative modeling (Feng et al. 2017; Pu et al. 2017). and reinforcement learning (Liu et al. (017). Theoretical work on
al.
SVGD has analyzed its behavior in asymptotic (Liu 2017: et al.]|2019 Duncan et al., (023) and
non-asymplotic (Chen et al.)
2018a
Liu & Wang 2018 (Shi & Mackey 2022) regimes as well as in
high dimensions (Zhuo et al.
2018
Ba et al. 2021).
A significant part of particle VI research explores understanding SVGD kernelized gradient flow Liu et al.] 2017); [Chewi et al. (2020). This line of research investigates the kernel and the properties of its associated function space in terms of the quality of the gradient flow approximation. These include broadening the functional regularizer Dong et al. (2022), specializing the acceleration schedule the step size Liu et al.] (2019), and alternatives to RBF kernel such as scalar kernels (Gorham & Mackey 2017 Wang et al., 2018)) and matrix variate kernels Wang et al. (2019). Where SMI focuses on the attractive force, these works focus on the repulsive force. As such, there is significant potential for adapting this body of work to SMI.
Annealing SVGD (ASVGD) Angelo & Fortuin (202aa) lap the only alternative that directly addresses variance collapse in SVGD with viable method. The resampling method introduced by Ba et al. (2021) is computationally impractical for large-scale problems, and we demonstrate bias in the appendix. Our experimental results in section O indicate that, unlike SMI, ASVGD follows the same collapse pattern as SVGD.
This work can also be related to work using hierarchical variational models (HVM) (Ranganath cet al. 2016) and mixture approximations (Jaakkola Jordan 1998; Bishop et al. 1998 Gershman
2012 Salimans & Knowles 2013, Miller et al.] (017). Unlike prior work on MsMs, Stein mixture inference does not require auxilaary models or bounds looser than an ELBO. Like SMI, Morningstar et al. 2021) introduces SIWAE, variational objective for mixture inference. While SMI employs an entropic regularizer on particles, SIWAE uses importance weighting. As an ELBO-based method, SIWAE can readily be incorporated as SMI's attractive force, enabling kernel design within the framework, We are unaware of any work that applies SVGD (Liu & Wang (016) or NSVGD (Wang et al., (2019) to optimize HVMs. Pu al. (017) considers SVGD specialized for VAEs; their method is similar to SMI in that it introduces an encoder. However, unlike SMI, their method only applies VAEs.
et al.
EXPERIMENTS
Because SVGD and ASVGD are prone to variance collapse (Ba et al. 2021). increasing the dimen- sionality of the posterior requires more particles represent uncertainty adequately. On the other hand, SMI can adjust the distribution of the variational components, thus requiring fewer particles. We demonstrate this for small and moderately sized models on synthetic and real-world data.
All experiments are carried out on an NVIDIA Quadro RTX 6000 GPU. For clarity, we only outline the experimental settings the following sections, leaving the details necessary for reproduction to a All experiments use our publicly available Stein VI library, and we provide the source code for the experiments
com/ eattry- sci science// smi te___exxrrmmeee
6.1 GAUSSIAN VARIANCE ESTIMATION
Following Zhuo et al. 2018), we estimate the per-dimension variance, called the dimension marginal variance, of standard multivariate Gaussian of increasing dimension with fixed number of particles. Hereafter, variance refers to the dimension marginal variance. The estimated variance will tend towards zero for method prone to collapse. The right panel of fig, 2 demonstrates variance collapse in SVGD and ASVGD with twenty particles. We see no benefit from annealing SVGD. In contrast, for SMI, when we single particle with mean-field multivariate Gaussian guide (i.e., the Stein
7

---

Published as a conference paper at ICLR 2025
102
dimension merginal variance
13
u=10c.0
101
x=10.0
12
x=1.
109
x- 0.1
11
0.6
x= 0.001
10-2
Actual
10-2
0.4
ASVGD0o
SM
30
0.8
SMI
0.2
SVGD,>
10-1
0.7
0.6
0.0
60
Dimensionality
Dimensionality
Dimensionality
Figure Left and middle (zoom): Variance estimation standard multivariate Gaussian obtained with 20-particle ASVGD, SMI and SVGD. Only SMI does not collapse and is robust to changing Right: SMI with one particle and Gaussian guide exactly recovers the multivariate Gaussian.
(a) SMI
(b) SVGD
(c) ASVGD
(d) OVI
(c) NUTS
Figure 3: Top row: High-density interval (HDI) for the low-dimensional model inferred using SMI, SVGD, ASVGD, OVI and NUTS on the 1D wave dataset (dotted line). SVGD, ASVGD, and SMI use five particles. The posteriors are inferred with data drawn from the In region, highlighted with vertical lines. NUTS serves as reference. Bottom row: HDI for the ASVGD and SVGD display collapse by significant narrowing in HDI between the In regions when comparing the low to moderate dimensions. On the other hand, both OVI and SMI widen the HDI with the richer model for the in-between region. In contrast to SMI, OVI overestimates the variance in the In region, where data is available, for the mid-sized model.
model.
particle represents the mean and variance of the guide), the estimated variance stays close one. This is because, when using one particle, the Stein mixture contains the model.
What happens when the SMI posterior is richer than the model? When the posterior is richer than the model, we risk overestimating the variance. The middle panel of fig. illustrates that we can choose a < to improve the overestimation of variance when using more particles than needed. The model and guide are as before, but now SMI uses twenty particles instead of one. We can alleviate the overestimation with a < 1 because it allows overlapping particle neighborhoods, i.c., the mixture components can collapse onto each other, thereby mimicking single particle. The result for SMI implies that a < 1 significantly reduces variance, we are likely using too many particles. In our experiment, tuning is only viable strategy for SMI, the right panel of fig. shows. This is because a acts on o for SMI. whereas choosing a * changes the target posterior for the other methods.
8

---

Published as a conference paper at ICLR 2025
6.2 1D REGRESSION WITH SYNTHETIC DATA
Previously, we demonstrated that SMI with a single particle and a Gaussian guide recovers a multivariate Gaussian distribution regardless of dimension. To extend beyond this scenario, we use synthetic one-dimensional regression dataset (dotted line in fig. 3 to study the uncertainty estimation of 2 hidden-layer BNNs in data-sparse regions. A well-calibrated model should assign high uncertainty in data-sparse areas and low uncertainty in data-rich ones (Foong et al. 2019 Daxberger al., 202T). We compare tiny BNN with 5 hidden units (46 random variables) to a small BNN with 100 hidden units (10,401 random variables).
Does SMI capture uncertainty better? Figure illustrates this using the high-density interval (HDI) (Gelman al. 1995), shown in gray, which represents the narrowest 90% Bayesian credible interval. With No-U- No-U-Turn Sample (NUTS) (Hoffman et al. 2014) serving as reference, well- calibrated model is expected to produce wide HDIs data-sparse regions and narrow HDIs in data-rich areas. Among the variational methods, only SMI demonstrates this desired behavior for low- and moderate-dimension
models. Closing the SMI-SVGD gap in moderate-sized networks by increasing SVGD particles is infeasible our hardware (appendix (B). Moreover, ASVGD shows improvement over SVGD, which is consistent with the variance experiment.
6.3 UCI REGRESSION BENCHMARK
To investigate the improvement uncertainty quantification of SMI on moderately sized models for real-world data, consider the UCI regression benchmark with Standard and Gapl0 splits hidden-layered BNNs. Standard UCI uses ordinary 10% test splits (Mukhoti et al. 2018). Gapl0 sorts each feature dimension to create splits (Foong et al., 20T9): The middle 10% of data is used for testing, while the tails are used for training. well-calibrated method should perform well on standard and catastrophically deteriorate Gap10. The BNN details, datasets and splits are summarized in appendix CC.4 For comparison, we use SVGD, ASVGD, MAP and OVI as baselines and NUTS as the gold standard.
Table summarizes the UCI results, evaluating their performance using root mean squared error (RMSE) and negative log-likelihood (NLL). NLL the primary metric of interest because we evaluate uncertainty estimation. Here, SMI delivers the best performance on Standard and Gapl0 UCI datasets. Notably, the RMSE is best for MAP and SMI, which means SMI has not sacrificed prediction accuracy to improve the NLL.
6.4 MNIST CLASSIFICATION
Next, we examine multi-class classification by applying 2 and hidden-layer Bayesian Neural Networks (BNNs) to the MNIST dataset (LeCun al.) al.] 2010)) Details about the BNN configurations
are provided in appendix C.3 Our evaluation includes accuracy (Acc), confidence (Conf), NLL, and
several classification reliability metries: the Brier score (Brier) (Brier) 1950), expected calibration error (ECE) (Guo et al. 2017), and maximum calibration error (MCE) (Guo al. 2017). Among the reliability metrics, we highlight the Brier score, as ECE and MCE can be sensitive to the choice of the bin count (100 bins were used in this study).
Table 3 summarizes the results. For both BNNs, SMI generally outperforms other methods across all metrics except for ECE and MCE. Judging by the Brier score, SMI is deemed the best-calibrated method for 2 hidden-layer BNNs, while MAP and SMI exhibit comparable calibration performance in the 3-layer case. When considering all metrics collectively, SMI emerges the preferred approach.
DISCUSSION
Limitations
The main limitation of SMI is that the variational approximation could be misspecified by using too many particles or poor choice of the parametric family. We saw an example of this in Section 6. when using twenty particles to estimate a Gaussian with SMI. Another major limitation
'While NUTS is asymptotically exact and serves reference, it is significantly slower to converge than variational inference methods.

---

Published as a conference paper at ICLR 2025
Table 2: Root mean squared error (RMSE) and negative log- log-likeli (NLL) for the UCI regression benchmark with standard and Gap10 splits. Lower is better for RMSE and NLL. Variational inference methods that are less than or equal in distribution to the lowest mean VI method are underlined. Similarly, the best, or equal in distribution, among all methods is highlighted in bold. We compare methods using
(MWU) test (Mann & Whitney 1947) at a significance level 0.05. For VI methods, SMI and MAP perform comparably in RMSE, with SMI outperforming alternatives on probabilistic calibration measured by NLL. Overall, NUTS outperformance SMI on NLL, However, of the two, only NUTS suffer severe deterioration Gapl0.
Hann-WWitnee
NB
RARNE
Diteet
SKUI
SYGD
ASYGD
OV1
MAP
NUTS
5332
SYCI
ASVUD
NUTS
22=00
29
53
30.504
Cheaves
31 1.1.3
10.00
64.1..44
4.1.....
Eeass
08:13
LUSEL
21 6.
B4511
0.S-00.1
16
-03 0.
1.
-15.10.
-121-44
06108
08:100
0.0.00 8.6
08:00
25:000
42501
Pheni
27:00
3.6.000
Wise
10:011
10:001
11204
10208
11400
06200
85:00
07200
0.580
18200
Boass
56-44
76451
41-12
38-06
58453
24-03
41414
4141A
47411
41:.0
464414
Coeant:
01:2222
84244
1A20S
86254
$2203
82530
$1519
41:15
85:30
K7231
Katien
-14003
01:00
012
01:00
-85-2... 0.3
-R12CH
14241
LOMB21100
482
D0S00
13
10200
20:001
28:03
14 12
47283
63:03
23:3
DALE 16.1
25.01
269 963
03203
13:00
13:03
11300
14302
13588
144109
Table 3: Evaluation of 2 and 3 hidden-layer BNNs for MNIST classification on several metrics: confidence (Conf), negative log- likeliikloood (NLL), accuracy (Acc), Brier score (Brier), expected calibration error (ECE), and maximum calibration error (MCE). Lower is better for NLL, Brier, ECE and MCE. Higher is better for Conf and Acc. All methods less than or equal in distribution to the method with the best mean are highlighted in bold. Methods compared using an MWU test at significance level of 0.05. Overall, SMI stands out as the preferred method.
2 Hidden--Layered BNN Acc (f)
Method
Conf (f)
NLL (()
Brier ())
ECE (U)
MCF (.)
SMI
0.979 0.001
0.039
0.003
0.957 1: 0.003
0.065 1 0.005
0.148 = 0.012
0.631 +: 0.047
ASVGD
0.972 4.0000
0.053 0.004
0.949 0.003
0.074 + 0.005
0.135 = 0.007
0.634 - 0.024
MAP
0.900
0.050
0.002
0.952 0.001
0.068
0.000
0.133 0.000
0.574 0.000
OVI
0.921 + 0.006
0.158
0.012
0.908 + 0.006
0.106 + 0.007
0.085 0.010
0.630 = 0.136
SVGD
0.972 + 0.003
0.054 0.006
0.949 + 0.004
0.074 + 0.007
0.139 + 0.014
0.653 + 0.048
3 Hidden-Layered BNN
SMI
0.979 +: 0.002
0.042
0.005
0.956 + 0.002
0.067 + 0.003
0.150 + 0.014
0.653 + 0.057
ASVGD
0.956 0.004
0.104 0.011
0.936 0.004
0.083 1 0.003
0.132 - 0.011
0.651 1 0.075
MAP
0.976 0.001
0.044 0.003
0.955 1 0.001
0.066 : 0.000
0.126 - 0.000
0.614 0.000
OVI
0.913
0.182 0.012
0.899 0.005
0.116 0.005
0.084 0.009
0.652 0.133
SVGD
0.000
0.091 0.011
0.940 = 0.002
0.081 + 0.004
0.135 = 0.013
0.649 = 0.044
for practitioners is knowing how to choose the kernel. We present SMI using RBF kernel, leaving a study of kernel choice to future work.
Future directions We present SMI as an extension to nonlinear SVGD, anchored the kernelized gradient flows theory (Liu et al. 2017 Chewi et al. 2020). However, such flows are not necessarily the best choice of transport for mixture approximations Chen et al. (2018b); Dong et al.) (2022). open question is identifying which properties make gradient flows well-suited for mixtures.
One of the issues with Bayesian modeling neural networks is their inherent non-identifiabilit; which can lead to degenerate posteriors (Yacoby et al. 2022: Roy et al.] 2024). SMI provides several opportunities for addressing this issue via the choice of its kernel. Promising directions are reparameterization invariant kernels (Roy al.) 2024),, probability product kernels (Jebara et al. 2004) and harnessing the connection between SMI and repulsive deep ensembles (D' (D'Angelo & Fortuin 2021b).
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We acknowledge support from the Independent Research Fund Denmark Technology and Produc- tion Sciences (grant 9131-00025BB and the VILLUM Experiment Programme (grant 50240). We thank Ahmad Salim Al-Sibahi, Martin Jankowiak and Du Phan for their discussions and help in implementing the SMI inference engine in NumPyro.
REFERENCES
Jimmy Ba, Murat A Erdogdu, Marzyeh Ghassemi, Shengyang Sun, Taiji Suzuki, Denny Wu, and Tian- zong Zhang. Understanding the Variance Collapse SVGD in High Dimensions. International Conference Learning Representations, 2021.
Alain Berlinet and Christine Thomas-Agnan. Reproducing Kernel Hilbert Spaces in Probability and Statistics. Springer Science & Business Media, 2011.
Christopher M Bishop, Neil D Lawrence, Tommi Jaakkola, and Michael Jordan. Approximating Posterior Distributions Belief Networks using Mixtures. Advances Neural Information Processing Systems, pp. 416 422, 1998.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational Inference: A Review for Statisticians. Journal of the American Statistical Association, 1 112(118):59-877, 2017.
Glenn W Brier. Verification of Forecasts Expressed in Terms of Probability. Monthly weather review, 78(1):1-3, 1950.
Changyou Chen, Ruiyi Zhang, Wenlin Wang, Bai Li, and Liqun Chen. A Unified Particle- Optimization Framework for Scalable Bayesian Sampling. In Conference on Uncertainty in Artificial Intelligence, 2018a.
Yongxin Chen, Tryphon T Georgiou, and Allen Tannenbaum. Optimal Transport for Gaussian Mixture Models. In IEEE Access, volume 7, pp. 6269-6278. IEEE, 2018b.
Sinho Chewi, Thibaut Le Gouic, Chen Lu, Tyler Maunu, and Philippe Rigollet. SVGD as a Kerneiized Wasserstein Gradient Flow of the Chi-Squared Divergence. In Advances in Neural Information Processing Systems, volume 33, pp. 2098-2109, 2020.
Francesco D' D'Angolo and Vincent Fortuin. Annealed Stein Variational Gradient Descent. arXiv preprint arXiv:2 0tt::2111.00900 09815, 2021a.
Francesco D' Angelo and Vincent Fortuin. Repulsive Deep Ensembles are Bayesian. In Advances in Neural Information Processing Systems, volume 34, pp. 3451-3465, 2021b.
Erik Daxberger, Eric Nalisnick, James U Allingham, Javier Antorán, and José Miguel Hernández- Lobato. Bayesian Deep Learning via Subnetwork Inference. In International Conference Machine Learning, pp. 2510-2521. PMLR, 2021.
Gianluca Detommaso, Tiangang Cui, Youssef Marzouk, Alessio Spantini, and Robert Scheichl. A Stein Variational Newton Method. In Advances in Neural Information Processing Systems, pp. 9169-9179, 2018.
Justin Domke. A Divergence Bound for Hybrids of MCMc and Variational Inference and an Application to Langevin Dynamics and SGVI. In International Conference on Machine Learning, pp. 1029-1038. PMLR, 2017.
Hanze Dong, Xi Wang, Yong Lin, and Tong Zhang. Particle-Based Variational Inference with Preconditioned Functional Gradient Flow. arXiv preprint arXiv:2211. 2022.
Andrew Duncan, Nikolas Nüsken, and Lukasz Szpruch. On the Geometry of Stein Variational Gradient Descent. Journal of Machine Learning Research, volume 24, pp. 1-39, 2023.
Yihao Feng, Dilin Wang, and Qiang Liu. Learning to Draw Samples with Amortized Stein Variational Gradient Descent. In Conference Uncertainty Artificial Intelligence. AUAI Press, 2017.

---

Published as a conference paper at ICLR 2025
Andrew YK Foong, Yingzhen Li, José Miguel Hemández-Lobate and Richard E Turner. 'In- Between' Uncertainty in Bayesian Neural Networks. arXiv preprint arXiv:1906.1133/ 2019.
Andrew Frank, Padhraic Smyth, and Alexander T Ihler. Puttcle-Bassd Variational Inference for Continuous Systems. In Advances Neurul Information Processing Systems, pp. 826-834, 2009.
Victor Gallego and David Rios Insua. Stochastic Gradient MCMC with Repulsive Forces. In Stat, volume 1050, pp. 30, 2018.
Jose Gavalda-Garcia, Bhawna Dixit, Adrián Díaz, An Ghysels, and Wim Vranken. Gradations in Protein Dynamies Captured by Experimental NMR Are Not Well Represented by Alphafold2 Models and Other Computational Metrics. Journal of Molecular Biology, 437:168900,, 2025.
Andrew Gelman, John Carlin, Hal Stern, and Donald Rubin. Bayesian Data Analysis. Chapman and Hall/CRC, 1995.
Samuel J Gershman, Matthew D Hoffman, and David M Blei. Nonparametric Variational Inference. In Proceedings of the 29th International Conference on Machine Learning, pp. 235-242. Omnipress, 2012.
Jackson Gorham and Lester W. Mackey. Measuring Sample Quality with Kernels. International Conference on Machine Learning, 70: 1292-1301, 2017.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On Calibration of Modern Neural Networks. In International conference machine learning, pp. 1321-1330. PMLR, 2017.
Jun Han and Qiang Liu. Stein Variational Adaptive Importance Sampling. In Proceedings of the 33rd Conference Uncertainty in Artificial Intelligence. AUAI Press, 2017.
Jun Han and Qiang Liu. Stein Variational Gradient Descent Without Gradient. Proceedings of the 35th International Conference on Machine Learning, volume 80, pp. 1900-1908, 2018
José Miguel Hernández-Lobato and Ryan Adams. Probabilistic Backpropagation for Scalable Learning of Bayesian Neural Networks. In International Conference on Machine Learning, pp. 1861-1869, PMLR, 2015.
Matthew D Hoffman, David M Blei, Chong Wang, and John Paisley. Stochastic Variational Inference. In Journal of Machine Learning Research, 2013.
Matthew D Hoffman, Andrew Gelman, et al. The no-u-turn sampler: adaptively setting path lengths in hamiltonian monte carlo. J. Mach. Learn. Res., 15(1):1993-1623, 2014.
Tommi 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. In Journal of Machine Learning Research, volume 5, pp. 1998844,2 2004.
Michael Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. Introduction to Variational Methods for Graphical Models. In Machine Learning, volume 37, pp. 183-233. Springer, 1999.
Diederik Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. In International Conference on Learning Representations, 2015.
Donald Ervin Knuth. The Art of Computer Programming, volume 3. Pearson Education, 1997.
Alp Kucukelbir, Dustin Tran, Rajesh Ranganath, Andrew Gelman, and David M Blei. Automatic Differentiation Variational Inference. In Journal of Machine Learning Research, volume 18, pp. 1-45, 2017.
Yann LeCun, Corinna Cortes, and CJ Burges. MNIST Handwritten Digit Database. ATT Labs [Online]. Available: http://yammleccu..om/eeee 2, 2010.
12

---

Published as a conference paper at ICLR 2025
Chang Liu and Jun Zhu. Riemannian Stein Variational Gradient Descent for Bayesian Inference. In 32nd AAAI Conference on Artificial Intelligence. 2018.
Chang Liu, Jingwei Zhuo, Pengyu Cheng, Ruiyi Zhang, and Jun Zhu. Understanding and Accelerating Particle-Based Variational Inference. In International Conference on Machine Learning, pp. 4082- 4092. PMLR, 2019.
Qiang Liu. Stein Variational Gradient Descent as Gradient Flow. In Advances in Neural Information Processing Systems, pp. 3115-3122 3123, 2017.
Qiang Liu and Dilin Wang. Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm. In Advances in Neural Information Processing Systems, volume 29, 2016.
Qiang Liu and Dilin Wang. Stein Variational Gradient Descent as Moment Matching. In Advances in 2018.
Neural Information Processing Systems, pp.
Yang Liu, Prajit Ramachandran, Qiang Liu, and Jian Peng. Stein Variational Policy Gradient. In Proceedings of the 33rd Conference on Uncertainty in Artificial Intelligence. AUAI Press, 2017.
Jianfeng Lu, Yulong Lu, and James Nolen. Scaling Limit of Stein Variational Gradient Descent: The Mean Field Regime. In SIAM Journal on Mathematical Analysis, volume 51, pp. 648-671, 2019.
Henry B Mann and Donald R Whitney. On a Test Whether One of Two Random Variables Stochastically Larger Than the Other. In The Annals of Mathematical Statistics, pp. 50-60. JSTOR, 1947.
Andrew C Miller, Nicholas Foti, and Ryan Adams. Variational Boosting: Iteratively Refining Posterior Approximations. In Proceedings of the 34th International Conference Machine Learning, 2017.
Warren Morningstar, Sharad Vikram, Cusuh Ham, Andrew Gallagher, and Joshua Dillon. Automatic Differentiation Variational Inference with Mixtures. In International Conference on Artificial Intelligence and Statistics, pp. 3250-3258. PMLR, 2021.
Jishnu Mukhoti, Pontus Stenetorp, and Yarin Gal. On the Importance of Strong Baselines in Bayesian Deep Learning. In Bayesian Deep Learning, 2018
Eric Nalisnick and Padhraic Smyth. Variational Inference With Stein Mixtures. In Advances in Approximate Bayesian Inference, NeurIPS 2017 Workshop, 2017.
Yurii Nesterov. Primal-dual subgradient methods for convex problems. Mathematical programming, 120(1):221-253) 2009.
Anh Nguyen, Jason Yosinski, and Jeff Clune. Deep Neural Networks are Easily Fooled: High Confidence Predictions for Unrecognizable Images. In Proceedings of the IEEÉ Conference on Computer Vision and Pattern Recognition, pp. 427-436, 2015.
Du Phan, Neeraj Pradhan, and Martin Jankowiak. Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro. In Program Transformations for Machine Learning, NeurlPS Workshop. 2019.
Yuchen Pu, Zhe Gan, Ricardo Henao, Chumyuan Li, Shaobo Han, and Lawrence Carin. VAE Learning via Stein Variational Gradient Descent. In Advances Neural Information Processing Systems, pp. 2017.
4236-4245, 2017.
Rajesh Ranganath, Sean Gerrish, and David Blei. Black Box Variational Inference. In Conference Artificial Intelligence and Statistics, pp. 814-822. 2014.
Rajesh Ranganath, Dustin Tran, and David Blei. Hicrarchical Variational Models. In International Conference Machine Learning, pp. 324-333. 2016.
Hrittik Roy, Marco Miani, Carl Henrik Ek, Philipp Hennig, Marvin Pförtner, Lukas Tatzel, and Søren Hauberg. Reparameterization Invariance in Approximate Bayesian Inference. arXiv preprint arXiv:2406.0334 033333 2024.
13

---

Published as a conference paper at ICLR 2025
Ardavan Saeedi, Tejas D. Kulkarni, Vikash K. Mansinghka, and Samuel J. Gershman. Variational
Particle Approximations. In Journal of Machine Learning Research, volume 18, pp. 1-29, 2017.
Tim Salimans and David A Knowles. Fixed-Form Variational Posterior Approximation through Stochastic Linear Regression. In Bayesian Analysis, volume 8, pp. 837-882, 2013.
Jiaxin Shi and Lester Mackey. Sinite-Partiche Convergence Rate for Stein Variational Gradient Descent. arXiv preprint (tfi::2211.097222 2022.
Christian Szegedy, Wojeiech Zaremba, Ilya Sutskever, Joan Bruna, Dumitru Erhan, lan Goodfellow, and Rob Fergus. Intriguing Propertics of Neural Networks. In International Conference on Learning Representations, 2014.
Aki Vehtari, Andrew Gelman, and Jonah Gabry. Practical Bayesian Model Evaluation using Leave- One-Out Cross- Validatin and WAIC. In Statistics and computing, volume 27, pp. 1413-1432. Springer, 2017.
Dilin Wang and Qiang Liu. Nonlinear Stein Variational Gradient Descent for Learning Diversified Mixture Models. In International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 6376-6585, Long Beach, California, USA, 09-15 Jun 2019.
PMLR.
Dilin Wang, Zhe Zeng, and Qiang Liu. Stein Variational Message Passing for Continuous Graphical Models. In International Conference on Machine Learning, pp. 5219-5227. PMLR, 2018.
Dilin Wang, Ziyang Tang, Chandrajit Bajaj, and Qiang Liu. Stein Variational Gradient Descent with Matrix-Valued Valued Kernels. In Advances in Neural Information Processing Systems, volume 32, pp. 7834, 2019.
Yaniv Yacoby, Weiwei Pan, and Finale Doshi-Velez Mitigating the Effects of Non-Identifiabilit on Inference for Bayesian Neural Networks with Latent Variables. Journal of Machine Learning Research, 23(244): 2022.
Jingwei Zhuo, Chang Liu, Jiaxin Shi, Jun Zhu, Ning Chen, and Bo Zhang. Message Passing Stein Variational Gradient Descent, In International Conference on Machine Learning, pp. 6013-6022, 2018.
14

---

Published as a conference paper at ICLR 2025
A STEIN MIXTURE INFERENCE DETAILS
This section provides the details missing from section In particular, we show that the functional
C [p] is an upper bound to Cip) in appendix A.1 give the complete derivation of the gradient of
(Pm)) in appendix A.2 and finally, in appendix A.3 demonstrate how to reduce SMI to SVGD and OVI.
A. BOUNDING THE SMI FUNCTIONAL
Our goal is to show that
[*] 2 [/p]
because it allows us to conclude that the SMI variational objective an ELBO when the repulsion is not scaled (i.e., 1). Recall that the above functionals given by
p00, D) ( p)
Clp] Ep(e)
log
and
p(0, D)
[p] Ep(()
(0(0
log
g b)
In both functionals, p is continuous distribution, p(0, D) the joint distribution of latent variables and data, and g0) p) = pp()) [0(pp) is SMI's variational approximation.
We need the following inequality to show that we can bound L(p). To shorten the notation, let p p(tb) and q = (())b). With this notation, it holds that
Ep Eg [log Ep [qll] 2 Ep [Eg [log gll
(10)
as shown by
Ep [Bg [log E, [all] 2 Ep [E, [E [log q]]l EE, [Ep [E,, [log a]]l
Ex [E,, [log d]]
With the inequality established, we can show that [t(p) upper bounds Clp) as follows:
log p(0, 1
Ct [p]
By(0110)
(expand log)
((0lbb)
B
[E v)
p(0 D)]]
E.((v) (Bg(())) [log (((}b))]
(negation of eq. (T0)
2 E
Eg(0 v)
[(0, [O]]
Ep(p)) [g(((()) [log Ep(p) [0(O
(combine logs)
p(0, D)
B wb)
log Ep(p) (g bb]
C(p),
which shows that L [p] 2 Cip) as claimed.
A.2 SMI GRADIENT DERIVATION
With the bound Clp) established,
focus showing that the mapping 1/ bbm
[(pm) is differentiable wrt. the f'th particle. For conciseness, we use L(pm) to mean both the map v1 2.......44 1b.... [(pm) from particles { and the functional L parameterized by Pm-- We can do this because ( 1 completely characterises (.n(.) = 1/m 1m (.)
Demonstrating that C(pm) is differentiable an essential component for using theorem 3.1 to optimize our variational objective. Specifically, theorem 3.1 requires us to have a differentiable symmetric mapping C(pm). We already established the symmetric nature of L(pm) the main article. The closed form of the gradient shows that L(pm) is differentiable wrt. v, if q0 x)) is differentiable wrt. sb. In practice, this restricts us to guides q(0) bb) that are differentiable. However, this restriction is shared with OVI and easy to fulfill.
15

---

Published as a conference paper at ICLR 2025
Recall claimed that the gradient of [(pm) wrt. th particle is given by
p(0.2) D)
Vv, ((pm) = B.(0(v.)
Vv, log g(0lv/:) log
> mm j= {(1bb))
(11)
V (000(0))) 1b,)
{(01b)))
with the SMI functional given by
p(O, D) log I m ΣØ= 9 b,,)
((mm)
10
To show that eq. [11 holds first, notice that because L(pm) is symmetric, the ordering of the particles does not matter. For our derivation, we, therefore, simply pick one. With the ordering of the particles now fixed, we can derive the gradient as follows:
p(O,D) d0
V V. ((pm)
((0)pb,))
1b log m
(expand the log)
>m
((0(b)))
V +/
((Olbb,) log p(0, D)d0
Vv, E ((1(bb)) log - d0
Vv, Σ ((1bbb)) log Σ ((0lbb))d
Now, the second term is zero because
V Σ q (0(bb) log _0 vv. log m
1 - / ((0)+b))00 = V log 0,
which gives us
m vv ((pm) V Σ/ ((0b)) log p(0, D)d0 V., E ((0) log Σ (((bbb,)))
Noting that when i +f we have Vv, (flbb) = O, we can eliminate the sum on the first term have
/ V g(0l)b)))))))))) log p0, D) V E/ ((0(bb)) log Σ ((0)bb))))
m V. ((pm)
From here, we use the product rule and combine like terms, obtain
mV v ((pm) = / V (00 log Σ, {(1bb,))
p(0,D) D) d0
Σ/ q (0,,)) Ve, log Σ (00 vb,
Finally, because V Vlogf= V we have
m mVv mVv.C(pmi ((pm) = / V0 log Σ, {0(bbb,))
p(0,D) d0 - -E (()bb,))
V / ((0)bb)))
(b)
Σ, ((0lbb,) d0
p(0, D)
m
V (1
= -Ey(0(+b))
V log (0 log
1 ((((bb)))
1= (00b)
From the above, we have established that eq. 11 holds and C(pm) ) therefore differentiable and symmetric required for using theorem 3.1 to maximize SMI's variational objective.
16

---

Published as a conference paper at ICLR 2025
A.3 REDUCING SMI OVI AND SVGD
In the following, establish that both OVI and SVGD instances of SMI for particular choice of hyper-parameters, namely a single particle and point-mass guide, respectively.
A.3.1 SVGD AND MAP ARE SPECIAL CASES OF SMI
We can connect SMI to SVGD by choosing each guide component ((0)eb,) as a point-mass, i.e., lv (0). Subsequently, the point-mass can be interpreted as a simple variable renaming. Using the point-mass for each particle, we have that
/ 1 (0) log
p(0, D) (0) d0 log
p(V. D)
Substituting this into L(pm), the gradient wrt. the l'th particle becomes
p(p, D)
v.. f(pm) = V log
-V (bb;(
(12)
m
Substituting eq. 12) for eq. in eq. recovers the SVGD update rule given to a constant factor 1/m. From the connection SVGD, we get the connection to MAP estimation for free as it corresponds to SVGD with one particle (Liu & Wang 2016). To be precise, MAP estimation corresponds to SMI with a point-mass guide and one particle. Naturally, we can also recover MAP estimation by first considering particle and then introducing the point-mass guide. Next, demonstrate that if choose an arbitrary (differential) guide and one particle, then SMI corresponds to OVI.
A.3.2 REDUCING SMI TO OVI
Like SVGD reduces to MAP estimation when only using one particle, SMI reduces to ordinary variational inference (as in eq. (4)) in the single-particle case. To see this, first note that with one particle, the kernel k(bb,)) +) is constant, regardless of bb and thus (bb = O. Starting from eq. (9) and denoting the constant value of ki(bb, *b) by c, we obtain
16² .6¹ + ck(pb" V + ca ca ik(b" *²'
7 log (01bb) log p00 D) 4b' ccB2)1000))
((BBg(0l))) V log ((0(b))
((0(bb) (b)
p00 D)
4b' + cc (p9(0 log
d0 ecE (cEy((11r) V (00
(0({))
10, D) *b' + ecV
3b + <c g(
log
(0(bb)
where e 0 is the step size. This means that with one particle, we are doing gradient ascent on the ELBO as defined in eq. (4) The connections SVGD and ordinary VI are attractive because SMI thus naturally bridges particle methods and OVI.
A.4 MINI-BATCHING
As with SVGD, computing the likelihood can become prohibitively expensive for large data sets (N > 1). To avoid the computational dependence the size of the dataset, we approximate the likelihood by data subsampling with the unbiased estimator
PT (D(0) = II (1/
(13)
VET
where I C. (D) and т is draw from the uniform distribution over index permutations. This follows the standard mini-batching method in NumPyro (Phan al. 2019
17

---

Published as a conference paper at ICLR 2025
Table The median recovery point R (> 5 favors SMI) for BNNs inferred with SMI and SVGD on different regions of the wave dataset. SMI uses five particles. Due to reaching hardware limitations, the moderate-dimensionaal results are lower bounds.
Model size
In
Between
Entire
Low Moderate
8
256
> 256
B RECOVERY POINT EXPERIMENT
To quantify the difference visualized in fig. we use the log point-wise predictive density (LPPD), quantity used for model comparison and model fit in the presence of outliers (Vehtari et al. 2017 The empirical LPPD is given by
s Σ p(u Fz 0.)
LPPD Σ log
where {( U:)}; are data points from an evaluation region, 0, g th. D) and 3; is drawn uniformly from the converged particles. We repeat the experiment ten times estimate the empirical LPPD for each region.
recovery-point experiment compares LPPD from SMI using five particles to SVGD with an increasing number of particles. We call the number of particles such that SVGD produces better LPPD the recovery point, R. Table 4 reports the median R over ten repeated trials over the three regions from Table
Can SVGD become on par with SMI by increasing the number of particles? Table 4 shows that increasing the particle count can only compensate for the difference in LPPD between SMI and SVGD for the tiny BNN, the low-dimensional model. For the moderately dimensional model, SVGD reaches the GPU memory limit before reaching the recovery point. The In region results show that MAP estimate is enough only when the noise level is low and there is enough data.
C EXPERIMENTAL DETAILS
This section provides extra results and the experimental setup needed for reproduction. Our experimen- tal code is available at htttpss: /ttttps:/qqtthb com/ http/aleetoryysclee wwtyyyacccceeeeeeeeeeeeeee htspsssmmeeeeeeeeeeeeeeeeeaeeaae experiments SMI is available inference engine under the Apache V2 license part of the deep probabilistic programming language NumPyro (Phan et al. 2019).
C.I VARIANCE ESTIMATION
In this experiment, we aim to recover the per-dimension variance of a multivariate standard Gaussian (MVG) across increasing dimensions. Specifically, we evaluate MVGs with dimensions 1, 4, 8, 10, 20, 40, 60, 80, and 100. We compare the performance of SVGD and ASVGD, each utilizing 20 particles, against SMI with both and 20 particles in estimating the variance.
we employ a factorized Gaussian guide initialized with a scale of 0.1. SMI's Gaussian guide mean is uniformly initialized within each dimension's [-2, 2). In contrast, the particles for ASVGD and SVGD are uniformly initialized within [-20, 20). This wider initialization is crucial, as ASVGD and SVGD fail converge in lower dimensions without
For SMI,
Optimization is performed using the Adam optimizer for SVGD and ASVGD and Adagrad for SMI, each with learning rate of 0.05. We run the optimization for 60,000 steps, sufficient for all three methods to achieve convergence.
Posterior shape To assess each method's ability to recover the shape of the standard MVG, we calculate the Frobenius norm between the estimated covariance matrix and the identity matrix,
18

---

Published as a conference paper at ICLR 2025
10
100
a=100.0
a=10.0
8
Frobenius Distance
80
0=1.0
Perfect fit
a=0.1
60
a=0.001
ASVGDzo
20
SMI
40
SVGD
20
2
1020 40
100
1020 40 60
80
100
Dimensionality
Dimensionality
Figure 4: Left: Frobenius distance between the estimated and the true covariance matrix in the Gaussian variance estimation experiment, using 20 particles for all methods. Only SMI achieves distances close zero, indicating that it accurately captures the shape of the standard Gaussian, unlike the other methods. Right: Frobenius distance when SMI uses single particle. In this case, SMI perfectly recovers the posterior.
0.100
a=100.0
0.075
G=1.0
0.050
a-0.1
a=0.001
Location
0.025
0.000
-0.025
Actual
ASVGD
-0.050
SMIpa
SMI;
-0.075
SVGD2o
RESVGD0e
0.100
10a
10a
Dimensionality
Figure 5: Mean location estimates of a standard Gaussian distribution across different dimensionalities and repulsion scaling (ax) for SMI (with 1 and 20 particles), ASVGD, SVGD and RESVGD (with 20 particles). RESVGD repulsion is not scaled. The "Actual" line represents the true mean location (zero). Only RESVGD exhibits significant bias, particularly in higher dimensions.
representing the true covariance of the MVG. A perfect recovery corresponds to distance between the matrices. As illustrated in Figure 4 SMI is the only method among the three that successfully captures the shape of the standard MVG.
Estimation of Mean Location Our Gaussian variance estimation experiment reveals that SVGD and ASVGD suffer from variance collapse. However, providing unbiased estimates of the mean location of a standard Gaussian distribution is an equally important requirement for these methods. As shown in fig. [S] SMI, SVGD, and ASVGD successfully achieve this. In contrast, fig. 5 also demonstrates that SVGD with resampling (RESVGD), implemented as Algorithm 1 in Ba et al. (202T), produces biased estimate of the mean. For this reason, excluded it from our experiments.
19

---

Published as a conference paper at ICLR 2025
Table 5: Evaluation interval and data size ((2)) of wave datasets. All data points are drawn uniformly from the evaluation interval. The Between and Entire regions contain points outside the clusters used for inference.
Region
Evaluation Interval
D
In
[-1.5, -0.5 U (13,,111
20
Between
|-0.5, 1.3]
60
Entire
[-2, 2]
120
C.2 SYNTHETIC ID REGRESSION
The data data-generating process combines linear and sine-wave periodic trend given by
ply N (ulw = (1.5 [2r(x: + 2/a)] + + 1) - 0.1).
We estimate tiny and small BNN using twenty observations drawn uniformly from each of two separate clusters at the intervals [-1.5, 0.55] and [1.3, 1.7]. The construction provides data-sparse interval (-0.5, 1.3] in between the two clusters. The idea of Foong al. (2019) is to this in- between region to evaluate the inference methods' ability to capture and assign high uncertainty to data-sparse intervals.
We evaluate the BNNs on the In, Between, and Entire regions specified in table 5 The Between and Entire regions contain points outside the data clusters seen in fig. 6 The In region has separate samples for inference and evaluation.
Bayesian networks The BNNs have two hidden layers with tanh activation for both models. The case has a hidden dimension of 100, and the low-dimensional one has a hidden dimension of 5, yielding 10,401 and 46 parameters, respectively. We use standard Gaussian priors on weights and biases, and Gaussian likelihood. The network determines the likelihood mean, while the standard deviation of the likelihood fixed at the known data noise level of 0.1. The noise level is intentionally kept small to ensure that any observed uncertainty arises primarily from the BNN.
Inference details We use the Adam optimizer (Kingma & Ba 2013) with a learning rate of 0.001. We run SVGD, ASVGD, and SMI with five particles for 15,000 steps and OVI for 50,000 steps, sufficient for converging. We use 5,000 draws to estimate performance metric for OVI and SMI. For both SMI and OVI, use factorized Gaussians guides. We
hundred draws estimate the Stein force for SMI (i.e., eq. ( All methods are initialized 1-0.1, 1, 0.1] and measurements are taken for ten different initialization. We use 500 warmup steps and a chain of 3000 steps for NUTS in NumPyro. All other parameters are default: max depth is 10, during warmup, the mass matrix is adapted using the online Welford scheme (Knuth 1997 and step size calibrated using dual averaging Nesterov 2009).
C.2.1 RECOVERY POINT SETUP
The recovery point experiment uses the same hyper- apper-parameter setup above. Recall that the recovery point is the number of particle SVGD required to get an LPPD below five particle SMI. To reach it, begin with one particle SVGD and subsequently double until reach the recovery point. We repeat the experiment ten times.
C.3 MNIST CLASSIFICATION
This section outlines the details required to reproduce our MNIST classification results.
We utilize both 2 and 3 hidden-layer BNNs of size 100 and tanh activation functions. Input images are flattened before being fed into the BNNs. The likelihood is modeled as a 10-class categorical distribution, parameterized by the logits produced by the BNN.
Bayesian network
20

---

Published as a conference paper at ICLR 2025
5.-
5.0
5.0
0.0
0.0
0.0
-.5
-3.0
-3.0
-2
-2
-2
Taan
0.0
5.0
Figure 6: Top row: The samples were drawn from the data-generating process for evaluating Between, In and Entire regions. The In region used for inferring the BNNs highlighted in grey. Bottom row: The samples drawn from the Sata-generating process infer BNN posteriors.
Inference details We employ the Adam optimizer with learning rate of 10 for both MAP and OVI. For SVGD, SMI) and ASVGD, we use the Adagrad optimizer with learning rate of 0.7 for the 2 hidden-layer BNN and 0.8 for the hidden-layer BNN, utilizing five particles in each case. Specifically, the SMI method estimates the attractive force using 55 draws.
All approaches are trained for 100 epochs with a batch size of 128. The images are scaled to [0,1]. Instead of random subsampling, we implement mini-batching and appropriately scale the likelihood, as described in Equation eq.
C.4 UCI REGRESSION BENCHMARK
In this section, we provide the details for reproducing our UCI regression results.
Bayesian network We use 2 hidden-layer Bayesian neural of size 50 and ReLU activation for all datasets. We use a Gaussian likelihood with the mean given by the BNN and Gamma(shape=l, rate=0.1) prior on the precision (i.e., reciprocal variance). For SMI and OVI, we use the softplus (xr + log (exp(xx) + 1)) transformation on the Gaussian approximation to account for the difference in support of the likelihood precision. This is the transformation recommended in Kucukelbir et al. (2017) for inference using automatic differentiation when transforming a random variable from R to Rt The independent variable (x: is standardized, while the dependent variable (y) is kept as is. We randomly initialize guides uniformly in the interval 0.1.)
Inferring the networks We randomly initialize the tested methods uniformly in unconstraint space within the interval [-0.1, 0.1]. This is lower than the NumPyro default of 2). The initialization strategy mimics the initialization from Liu & Wang (2066) for SVGD and substantially reduces the steps needed for good performance.
We choose the learning rate from [5 10 10-176 -i with a grid search on the first split of each data set. We select the learning rate with the best RMSE on 10% validation split from the training data. Table provides the chosen learning rate for each method and dataset.
We use the Adam optimizer for up to 60.000 steps, inferring the BNNs random subsample of size 100 without replacement. The independent variable (æ) standardized, while the dependent variable (y) is kept is.
We use a convergence criterion on the Euclidean norm of ø² for the particle methods. We compare a slow-moving norm average, calculated over the last 350 steps, against fast-moving norm average, computed over the previous 35 steps. If the fast-moving average exceeds the slow-moving average, conclude that the methods fluctuate around minimum and stop iterating Equation (7). The number of past steps was chosen using the first split of the Boston Housing dataset with a learning rate of 0.5 using 10% validation set from training minimize RMSE.
21

---

Published as a conference paper at ICLR 2025
ASVGD
NUTS
OVI
SMI
SVGD
ASVGD
NUTS
OVI
SMI
SVGD
10
-10
10
hOU
50
10
UQ
10
10
500
0
-10
-20
-2
0
-2
2
-2
0
2
-2
0
2
-2
2
-2
2
-2
2
-2
0
-2 0
X
X
(a) Low-dimensional model
(b) cndderate-dimensionna model
Figure 7: Figure 7a High-density interval (HDI) for the low-dimensional model inferred using SMI. NUTS, SVGD, ASVGD and OVI on the ID wave dataset (dotted line). SVGD, ASVGD, and SMI use five particles. The posteriors are inferred with data drawn from the In region, highlighted with vertical lines. Figure 7b) HDI for the by significant narrowing in HDI between the In regions when comparing the low to moderate dimensions. In low-dimensional models, initialization plays a role in narrowing widening HDI for all methods. In mid-sized models, SMI is robust to initialization.
model. ASVGD and SVGD display collapse
22

---

Published as a conference paper at ICLR 2025
Table Learning rate for the methods used on the standard and Gapl0 splits of UCL.
Standard UCI
Dataset
ASVGD
MAP
OVI
SMI
SVGD
Boston
5 5-10-5
5 5-10-3
5 100
5-10-3
5
10-6
Concrete
5-10-4
5-10-4
5-10-3
5-10-4
5-
Energy
5-10-4
5-10-4
5-10-3
5-10-4
5-10-
Kin8nm
5.10-5
5.10-5
5.10-4
5.10-3
5-10-
Naval
5-10-5
5.10-4
5.10-4
5.10-4
5-10-5
Power
5-10-4
5.10-4
5.10-3
5.10-4
5-10-4
Protein
5.10-b
5.10-3
5.10-
5.10-3
5-10-3
Wine
5-10-6
50
5.
5.10-3
5-10-6
Yacht
5.10-5
50
5 5.10
10
5 5-10-3
Gap10 UCI
Boston
5.10-5
5.10-
5.10-
5.10-3
5. 10-2
Concrete
5.10-4
5.10-4
5.10-3
5.10-4
5-10-3
Energy
5-10-4
5-10-2
5.10-2
5.10-4
5-10-4
Kin8nm
5-10-5
5.10-1
5.10-3
5.10-5
5-10-2
Naval
5-10-
5.10-1
5.10-1
5.10-4
5.10-5
Power
5 5-10-1
5.10-5
5.10-2
5.10-4
5-
Protein
5-10-
5-10-3
5.10-3
5.10-4
5-
Wine
5-
50
5-
5-
1-
Yacht
5 50
10
5.
10
5 10-4
Table 7: Summary statisties for the standard UCI benchmark datasets with train-test splits from Hernández- & Adams (2013) and Gap10 benchmark datasets adapted from Foong et al. (2019) use 10% for testing instead of 33%.
Dataset
Train size
Test size
Features
Std Splits
Gap10 Splits
Boston
455
51
13
20
13
Concrete
927
103
8
20
8
Energy
691
77
8
20
8
Kin8nm
7373
819
8
20
8
Naval
10741
1193
17
20
17
Power
8611
957
4
20
4
Protein
41157
4573
9
9
Wine
1439
160
11
20
11
Yacht
277
31
6
20
6
The standard UCI split We use the train-test splits from Mukhoti et al.) (2018) for our standard UCI results. Table 7 gives summary statistics of the datasets. We treat features and responses (i.e., (x, y)) as real values.
The Gap10 UCI split We use the methodology suggested Foong et al. (2019) to construct the GAP dataset. We sort each feature dimension individually, taking the middle tenth as and leaving the two tails as our training split. Using this procedure will result in as many splits as features. However, where /hong et 2219) allocated the middle third for testing, we use tenth have the same test allocation as standard UCI. Comparing Standard to Gap10 in table 7 the Gap10 generally produces fewer splits than standard UCI.
Time comparison for methods In table we reproduce the per-step average inference time [sec/step] the UCI datasets for SMI, SVGD, ASVGD, OVI and MAP. On UCI datasets, SMI exhibits slower inference compared to the VI-based basclines. A portion this overhead arises from JIT compilation, which we believe can be reduced by optimizations in future releases of SMI.
23

---

Published as a conference paper at ICLR 2025
Table 8: The table shows the average time per step (in seconds per step) for datasets in the UCI regression benchmark. Although SMI demonstrates slower inference times than alternative methods, the recovery point experiment indicates that SMI offers more favorable trade-off when considering the associated performance improvements.
Dataset
SMI
SVGD
ASVGD
OVI
MAP
Boston
0.0014
0.0003
0.0003
0,0002
0.0001
Concrete
0.0015
0.0004
0.0003
0.0002
0.0001
Energy
0.0017
0.0003
0,0003
0.0002
0.0001
Kin8nm
0.0192
0.0004
0.0004
0.0003
0.0002
Naval
0.0103
0.0004
0.0004
0.0004
0.0001
Power
0.0079
0.0004
0.0004
0.0002
0.0002
Protein
0.0468
0.0011
0.0008
0.0004
0.0003
Wine
0.0093
0.0003
0.0003
0.0002
0.0001
Yacht
0.0059
0.0003
0.0003
0.0003
0.0001
10
force
Attractivn force
Posulsivve force
Repuliive force
X
SUGD
08-
X
0.0
(a) Location dimension of particles.
(b) Variance dimension of particles.
Figure 8: Converged two-particle approximation of two-dimensional Gaussian using SMI and SVGD. Each SMI particle, denoted as &6, parameterizes a Gaussian guide with & = (z, y, 2, w). where (xr, v) represent the guide's 's location and (z, w) represent its variances. In contrast, an SVGD particle, denoted as 0, only represents location dimensions, i.e., = (z. y). Left: Location dimensions of SMI and SVGD partieles. Shades show the equiprobability contours of the Gaussian. Right: Variance dimensions of the SMI particles. SVGD particles are absent here as they only represent location. SMI effectively approximates the Gaussian by explicitly incorporating variance as part of its particle dimensions. The SMI forces are scaled for better visibility in the location dimensions. No force arrows are visible in the variance dimensions because the system has converged, making the forces negligible.
When considering the recovery point experiment table 4 SMI demonstrates significantly improved runtime efficiency. On the mid-sized network, SMI achieves inference times 6x faster than SVGD. This observation suggests that while VI methods excel in runtime on UCI datasets, SMI provides better trade-off when factoring in performance gains. Thus, in contexts where accuracy and robustness are critical, SMI is preferable despite its higher initial runtime cost.
D SMI VERSUS SVGD: INSIGHTS FROM A SIMPLE TOY MODEL
To address variance collapse, the key distinction between SVGD and SMI lies in the space the particles occupy. SMI particles operate in a higher-dimensional space than SVGD particles. This allows the repulsive term, Vık(æ, y). in SMI to influence the distribution's 's shape and its parameterized location. In contrast, SVGD particles can only control location.
In SVGD, each particle represents latent parameter sample. Meanwhile, in SMI, each particle parameterizes an entire distribution. For example, if the parameterized distribution is a factorized Gaussian, each SMI particle would represent both the mean (location) and variance of the Gaussian.
24

---

Published as a conference paper at ICLR 2025
While the location component of an SMI particle shares the same space SVGD particle, the variance component has no equivalent in SVGD. As a result, the repulsive force in SMI operates in broader space, encompassing both location and variance.
This distinction becomes evident when comparing SVGD and SMI in a two-particle approximation of standard Gaussian distribution. The SMI approximation forms Gaussian mixture. By breaking SMI particles into their location and variance components, we can visualize the location component within the same space SVGD particles and the target Gaussian density. In fig. Sal SMI particle locations converge toward the center of the target Gaussian, while SVGD particles spread out, maintaining equal distances from the Gaussian center.
At first glance, focusing solely on the location dimension of the SMI particles might suggest they have collapsed. However, this interpretation is incomplete because it overlooks the role of variance. In fig. [b we examine the variance component of the SMI particles. We see that the two-particle SMI system together captures the variance of the target density because each dimension fig. [8b] sums to one. Consequently, the SMI estimation has not collapsed. Instead, it captures the target with all four dimensions of its particles (two for location and two for variance) rather than the two dimensions of the SVGD particle.
25


Reviews:
[{"review_id": "1Nb0lxcWQR", "reviewer": "Reviewer_rYZM", "summary": "The authors improve Stein variational gradient descent (Liu & Wang, 2016) by extending nonlinear SVGD (Wang & Liu, 2019) by learning a density-based mixture model to approximate the posterior, instead of solely relying on particles (i.e., delta-distributions).\nThey evaluate their method on a set of (small) scale regression tasks.", "questions": "- Q1: Can the authors speculate on performance as a function of the parameter count, e.g., sticking to BNNs, at which depth/width would the method start to struggle?\n- Q2: What are the increased runtime costs compared to compared baselines?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The method is a straightforward and effective extension of SVGD/NSVGD\n- The paper is well-written and easy to follow and the same goes for the provided codebase", "weaknesses": "- The experiments are rather small-scale and limited to regression data sets. Their aim seems to be primarily to demonstrate the relative performance of the proposed approach compared to prior SVGD-related approaches rather than, its absolute performance. In the list of baselines, at least a comparison against an HMC performance on the UCI data sets would have been nice to see how close it can come to it (or improve upon it).\n- The paper lacks ablations to evaluate what happens as an underlying BNN gets deeper, i.e., to what extent it can handle the increase in parameters. A deep experiment could be a combination with last-layer BNNs, i.e., learn the mixture not for the whole net, but treat only the penultimate layer in a Bayesian fashion.\n- ~~The experiments are limited to regressions with a homoscedastic, known observation noise.~~ What about classification or heteroscedastic regression tasks? _Edit: The claimed fixed noise was an error on my side. See the answer of the authors on this_\n- Citing Agarap (2018) in l478 as a reference for ReLUs seems rather odd. In their work, they evaluate the usage of a ReLU in place of a softmax for classification, i.e., nothing related to the current work nor has the ReLU been introduced in that paper.\n\n### Typos\n- l233 lacks a second closing bracket", "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": "", "extra_scores": {}}, {"review_id": "rZj15ugk9j", "reviewer": "Reviewer_mFTt", "summary": "The paper introduces Stein Mixture Inference (SMI), which optimizes a lower bound to the Evidence Lower Bound (ELBO). \nSMI extends Nonlinear Stein Variational Gradient Descent (NSVGD) to the variational Bayes setting and addresses the issue of variance collapse. \nThe effectiveness of SMI is demonstrated on both synthetic and real-world datasets.", "questions": "1. It’s challenging to distinguish the lines representing different methods in Figure 2 (e.g. $SMI_{1}$, $SMI_{20}$). \nUsing distinct colors for each method would improve the visualization and make the differences clearer.\n2. The experiments in Section 6.1 demonstrate that SMI overcomes variance collapse. It would also be valuable to assess whether the approximate distribution given by SMI accurately captures the shape of the posterior. \nThis could be evaluated by comparing the estimated covariance matrix with the target covariance matrix.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The problem addressed in this paper is both important and compelling. Traditional approaches like ordinary mean-field variational inference (OVI) and Stein Variational Gradient Descent (SVGD) often experience variance collapse, whereas SMI provides more accurate variance estimates, improving uncertainty quantification.\n\n2. The paper is well-written, providing a clear background and a thorough summary of related work. As someone slightly unfamiliar with the field, I particularly appreciated the authors' effort to re-explain and contextualize prior results, which greatly helped in assessing the paper's contributions.\n\n3. SMI is compared with other methods across a variety of synthetic and real-world datasets.", "weaknesses": "1. Variational inference offers a compelling alternative to sampling methods like MCMC due to its efficiency, especially in high-dimensional settings and with large-scale datasets. \nHowever, the current validation of SMI is limited to small to moderately-sized models, which somewhat limits its appeal and persuasiveness for broader, large-scale applications.\n\n2. The paper lacks theoretical insights or guidance on how SMI’s performance depends on the number of particles $m$.\nProviding recommendations or analysis on selecting an appropriate particle count $m$ would greatly enhance its practical applicability.", "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": "", "extra_scores": {}}, {"review_id": "0FFcorPbKr", "reviewer": "Reviewer_HpKW", "summary": "- The authors propose a method called Stein Mixture Inference (SMI) to address the issue of variance collapse observed in Stein Variational Gradient Descent (SVGD).\n- SMI extends the Nonlinear SVGD framework (Wang & Liu, 2019) to variational inference (VI) by allowing each particle to parameterize a component distribution within a mixture model, thereby introducing ELBO-like objectives.\n- The authors show that SMI offers several advantages over standard SVGD by effectively mitigating variance collapse.", "questions": "- Could you provide additional analysis or intuition on why the combination of an ELBO-like objective from VI and the Non-linear SVGD framework effectively mitigates variance collapse? Specifically, how does modeling a mixture of approximate distributions around the vicinity of SVGD particles help in avoiding variance collapse? A more detailed explanation or visual representation would be appreciated.\n  - For example, could you provide a more detailed explanation or visual representation of how the mixture components interact with the ELBO objective to mitigate variance collapse? For instance, a step-by-step explanation or diagram illustrating how the proposed method addresses the variance collapse problem would be helpful.\n- Why is the integration with the NSVGD framework necessary in your method? Is there evidence that the entropy regularization term alone is insufficient to address variance collapse? Given that Figure 2 shows variance collapse is mitigated even when $\\alpha$ takes small values, does this imply that the regularization component may not be as critical? If so, what is the rationale for including it in the framework?\n- Why were methods such as HMC and MMD descent not included in the comparative analysis, especially given their relevance in approximate inference and their use in experiments in (Ba et al., 2021)?\n  - If possible, could you add comparisons with HMC (Neal, 2011) and MMD descent (Arbel et al., 2019; Ba et al., 2021) in the experimental section, particularly at least on the UCI datasets, to provide a broader context for evaluating SMI’s performance in addressing variance collapse? If a full comparison is not feasible, could you discuss how SMI might be expected to compare to these methods theoretically or empirically, based on existing literature?\n- Could you elaborate on why the resampling method from (Ba et al., 2021) was excluded as a comparative method, despite the computational resources available (e.g., “NVIDIA Quadro RTX 6000 GPU”)? Is this method genuinely computationally infeasible for UCI benchmark datasets, or were there other factors influencing its exclusion?\n  - So, could you include the resampling method from Ba et al. (2021) in your comparisons, particularly on the UCI datasets, to strengthen the evaluation? If this is not feasible, could you provide a more detailed justification for why it is computationally infeasible, even with the available GPU resources? Additionally, if an empirical comparison is truly not possible, could you discuss how SMI theoretically compares to the resampling approach?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- The application of variational inference (VI) concepts to Stein Variational Gradient Descent (SVGD) appears novel and intriguing.\n- The authors validate their VI-based approach through numerical experiments on several UCI benchmark datasets, demonstrating good performance. The results seem to suggest that this approach effectively mitigates the impact of variance collapse.", "weaknesses": "### Insufficient Analysis of the Motivation Behind Extending SVGD with VI for Variance Collapse Mitigation:\n- The main objective of this paper, as I understand it, is to mitigate variance collapse by extending the SVGD objective function through a combination of an ELBO-like objective from VI and the Non-linear SVGD framework. However, it is not entirely clear \"why\" this extension effectively mitigates variance collapse. While Figure 1 provides a conceptual illustration, it does not intuitively explain how the proposed method addresses variance collapse. Additionally, while the third paragraph of the Introduction discusses the motivation for this approach, it remains unclear how using a mixture of approximate distributions around SVGD particles with a VI-inspired objective avoids variance collapse.\n- The authors propose controlling particles through variational distributions, similar to VI, as a solution to the variance collapse issue. However, given the use of the NSVGD framework, the critical role of this aspect remains unclear. The entropy regularization term could potentially affect not only mode collapse but also variance collapse. If the VI-inspired approach is indeed effective, the method should perform well even with $\\alpha=0$. In this context, Figure 2 shows that variance collapse is mitigated even when $\\alpha$ takes small values, suggesting that particle control via variational distributions may be effective. On the other hand, this result implies that regularization may not play a significant role, raising questions about the necessity of combining it with NSVGD. Overall, it remains unclear why the NSVGD framework is essential and which part of the proposed approach effectively addresses variance collapse.\n\n### Concerns Regarding the Limited Number of Comparative Methods:\n- For sample approximations of the posterior distribution, methods such as HMC (Neal, 2011) and MMD descent (Arbel et al., 2019; Ba et al., 2021) are also effective. However, this study only compares performance within the SVGD family of methods and EVI, leaving questions about the extent to which the proposed method mitigates variance collapse in the broader context of approximate inference. Given that (Ba et al., 2021) also includes these methods in numerical experiments addressing variance collapse, this comparison is essential for validating contributions in this research area.\n- Additionally, the absence of a comparison with the resampling method proposed by (Ba et al., 2021) raises concerns regarding the integrity of the performance evaluation. While the authors argue in Section 5 that the resampling method is computationally infeasible, I believe this does not fully justify its exclusion as a comparative method. Given the availability of an “NVIDIA Quadro RTX 6000 GPU,” running such methods may not be computationally prohibitive, at least for datasets like the UCI benchmarks.\n- Furthermore, I find it difficult to agree with the authors’ claim: “Annealing SVGD (ASVGD) D’Angelo & Fortuin (2021a) is the only alternative that directly addresses variance collapse in SVGD with a viable method.” I believe that the resampling method proposed by (Ba et al., 2021) is also aimed at mitigating the variance collapse problem.\n\n### Citation:\n- (Neal, 2011): R. M. Neal. MCMC using Hamiltonian dynamics. Handbook of Markov Chain Monte Carlo, 2(11):2. https://arxiv.org/abs/1206.1901.\n- (Arbel et al., 2019): M. Arbel, A. Korba, A. Salim, and Arthur Gretton. Maximum Mean Discrepancy Gradient Flow. NeurIPS2019. https://arxiv.org/abs/1906.04370.", "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": "", "extra_scores": {}}, {"review_id": "Ry7dZ9iLrP", "reviewer": "Reviewer_dciU", "summary": "The posterior of a Bayesian model is sometimes intractable, calling for approximate inference techniques. This paper focuses on the idea of approximating the Bayesian posterior with a mixture distribution where each component is parameterized separately but still in the same family. By viewing the ELBO as an objective with permutation invariant parameters, this paper incorporates ideas from Nonlinear-SVGD (NSVGD) and develop a Stein-style update rule of the mixture parameters. The resulted method, called Stein mixture inference (SMI), prevents itself from variance collapse. The paper also shows that asymptotically the optimized bound is an ELBO.\n\nIn the experiments, it is first shown that Stein-based methods suffer from variance collapse in synthesized Gaussian models and 1D regression models. In contrast, SMI and vanilla VI (OVI) produce the desired uncertainty. On the UCI regression benchmarks, SMI gives the best NLL in the most cases, comparing with other Stein-based methods and OVI.", "questions": "- How does SMI compare with ADVI with mixtures? \n- How is each component in the mixture distributions parameterized?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I believe the method is novel and the main idea is sound. The claims are clearly presented and supported by empirical evidences. Overall it is a complete work with a few concerns.", "weaknesses": "This paper has the title of ELBOing Stein, but I would rather call it Steining ELBO, which seems a bit unnecessary. To convince me to increase my score, I would like to see a discussion of [1], which also uses a mixture distribution to approximate the posterior. \n\nIf one is using a mixture distribution to target the Bayesian posterior, the most direct approach would be to try VI, instead of deriving a complicated Stein-based method. One pro of VI is that the mixture weights can be adjusted while one con is that it does not fully make use of the exchangeability of parameters. However, this work only considers mean-field VI, which is a really weak baseline. I would like to see how the permutation invariance helps the optimization of the mixture distribution.\n\nIt is not surprising that optimizing an VI objective prevents the approximate posterior from the pitfalls of SVGD. As shown in the paper, OVI does not have the issue. The argument that \"SMI is more particle-efficient than SVGD\" is translated to me as \"VI with mixtures is more particle-efficient than SVGD\". Then what is the point of using Stein?\n\nLine 150 says that \"Particle methods are attractive due to their freedom from strong parametric assumptions\". Mixture distribution in this paper seems to be a strong parametric assumption, especially when it uses fewer particles than SVGD, which further drags this work away from Stein. \n\nThe experiment section is also rather weak. The benchmark models all have very low dimensions. I expect a Bayesian inference algorithm in 2024 to be tested on more recent benchmarks, like models in posteriorDB, or larger BNN problems.\n\n[1] Morningstar, W., Vikram, S., Ham, C., Gallagher, A., & Dillon, J. (2021, March). Automatic differentiation variational inference with mixtures. In International Conference on Artificial Intelligence and Statistics (pp. 3250-3258). PMLR.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)