Variational Inference with  Mixtures of Isotropic Gaussians
Authors: Marguerite Petit-Talamon
Marc Lambert
Anna Korba
Venue: NeurIPS 2025 poster ()
Keywords: variational inference
optimization
optimal transport
Primary area: probabilistic_methods

Abstract: Variational inference (VI) is a popular approach
in Bayesian inference, that looks for the best approximation of the posterior distribution within a
parametric family, minimizing a loss that is typically the (reverse) Kullback-Leibler (KL) divergence. In this paper, we focus on the following parametric family: mixtures of isotropic Gaussians (i.e., with diagonal covariance matrices proportional to the identity) and uniform weights. 
We develop a variational framework and provide efficient algorithms suited for this family. In contrast with mixtures of Gaussian with generic covariance matrices, this choice presents a balance between accurate approximations of multimodal Bayesian posteriors, while being memory and computationally efficient. 
Our algorithms implement gradient descent on the location of the mixture components (the modes of the Gaussians), and either (an entropic) Mirror or Bures descent on their variance parameters. 
We illustrate the performance of our algorithms on numerical experiments.

Variational Inference with Mixtures of Isotropic Gaussians
Marguerite Petit-Talamon
Mare Lambert INRIA, Ecole Normale Supérieure DGA, French Procurement Agency marc. lambert fr
CREST, ENSAE
Institut Polytechnique de Paris pegggerrtte..etittaaalammmoeeeee
Anna Korba CREST, ENSAE Institut Polytechnique de Paris anna. wwwaaakrrbeeeeeee....m fr
Abstract
Variational inference (VI) is popular approach in Bayesian inference, that looks for the best approximation of the posterior distribution within a parametric family, minimizing loss that is typically the (reverse) Kullback-Leibler (KL) diver- gence. In this paper, we focus on the following parametric family: mixtures isotropic Gaussians (i.e., with diagonal covariance matrices proportional the identity) and uniform weights. We develop a variational framework and provide efficient algorithms suited for this family. In contrast with mixtures of Gaussian with generic covariance matrices, this choice presents balance between accurate approximations of multimodal Bayesian posteriors, while being memory and computationally efficient. Our algorithms implement gradient descent on the location of the mixture components (the modes of the Gaussians), and either (an entropic) Mirror or Bures descent on their variance parameters. We illustrate the performance of our algorithms on numerical experiments.
Introduction
The core problem of Bayesian inference is sample from posterior distribution model parameters, combining prior knowledge with observed data. Unfortunately, the posterior distribution is generally difficull compute due to the presence of an intractable integral (the normalization constant), and density with respect to the Lebesgue measure on Rd also denoted æ, only known in unnormalized form alternative to standard Markov chain Monte Carlo (MCMC, [Roberts and Rosenthal (2044) that approximates the posterior by optimizing
e Variational inference (VI, (Blei et al. 2017 is prominent
family of tractable distributions. While this restriction can introduce bias, VI typically much faster than MCMC, since it reframes the sampling problem the parameters of the variational family.
(generally inite-dimensional) optimization Specifically, VI seeks a distribution a parametric family that minimizes a discrepancy to the target posterior, typically the reverse Kullback-Leibe (KL) divergence:
min (L(H))),
(I)
#EC
where KL(μ(T)= = f log(ds/ds)) ddu if μ is absolutely continuous with respect to n denoting du/dx its Radon Radon-Nikodym density, and +oo else.
There exist various choices of variational families and suited algorithms in the literature of VI. For instance, Mean-field variational inference (MFVI) aims to find an approximate posterior that
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Recent studies by Armese and Lacker [2024) and Lavenant and Zanella [2024) have investigated solving this problem through the lens of coordinate ascent variational inference (CAVI, Bishop (2066), or over polyhedral sets with first-order algorithms (Jaan et al., 2024]. Traditionally, the variational family parameterized by a fnite-dimensional set of parameters, yet, Yao and Yang [202]) - Yao et al. (2024) have extended this framework to MFVI over infiniteddimensionall families. Alternatively, Gaussian variational inference [Barber and Bishop 1997 Opper and Archambeau 2009] has garnered significant attention due to its computational tractability and theoretical appeal Challis and Barber 2013). More recently, a new class of Gaussian
factors as a product of distributions [Lacker 2023
VI algorithms has emerged, based on the discretization of the gradient flow of the KL divergence on the space of Gaussian measures equipped with the Wasserstein-2 metric: Lambert et al. [202]]
proposed a forward (explicit) time discretization approach, while Diao et al.
[2023] Domke et al.
[023] introduced a seward--ackwwar
splitting scheme inspired by [Salim et al.
2020 Other
algorithms were recently introduced for Gaussian VI by optimizing weighted score-based (Fisher)
divergence [Modi et al. [025] Cai et al. 2024b). These methods bring valuable insights for Gaussian
VI, but may provide too crude approximation of the target distribution if the latter is multimodal.
In this setting, deep generative approaches such as normalizing flows (Tabak and Canden-Eijnden 2010 Papamakarios et al., 2021 Kobyzev et al. 2020] offering tractable, normalized densities and efficient sampling, have emerged as flexible and powerful alternatives to traditional Gaussian or factorized variational families. While it offers a flexible and widely applicable approach VI. is not tailored to specific variational family. Moreover, despite their ability to represent multimodal distributions using expressive variational families with neural networks, it can suffer from mode collapse (Soletskyi et al., 20224. especially as the dimension increases, as shown recently in large scale empirical evaluation from Blessing et al. 2024).
A particularly compelling variational family is the one of mixtures of Gaussians. Indeed, the latter can capture complex, multimodal distributions. Also, they lead to tractable approximate posteriors, since sampling mixtures is computationnally cheap, and also marginalizations, or expectations of linear and quadratic functions can be computed in closed-form. For this variational family, several algorithms
have been proposed [Lin et al. [019 Arenz et al.) 2018 2023] based on natural gradient descent over the natural parameters of the Gaussians. An extension of the Wasserstein gradient flow approach for Gaussian mixtures has also been proposed in Lambert et al.] [2022). However, their practical utility is often hindered by the computational challenges of inference, particularly when handling high-dimensional distributions. Indeed, a key challenge in deploying Gaussian mixture models lies in parameterizing the covariance matrices. While full covariance matrices provide maximum flexibility, their quadratic scaling with dimensionality leads to high computational demands, even if some solutions have been proposed to store them efficiently [Challis and Barber 2013] Bonnabel et al., 2024). To address this issue, we restrict the variational family of mixtures of Gaussians to diagonal, more precisely isotropic covariance matrices (i.e., assuming equal variance across all dimensions), with uniform weights. Mode collapse may result from mean alignment, i.c., two components' means may align with the same mode of the target, instead of covering multiple modes; and vanishing do, prevents the latter. With this structure,
weights [Soletskyi et al. 2024]. Using fixed weights, each Gaussian component in the mixture is parameterized by + 1 parameters in dimension d. As result, a mixture of N isotropic Gaussians requires a memory cost of N(d + 1), compared to Ndd2 + d) for a full covariance mixture. Then, optimizing mixtures of isotropic Gaussians incurs a memory cost roughly equivalent to that of optimizing the means only. We show in this study that this choice of variational family balances between accuracy of the variational approximation, i.e., its ability to model multimodal target distributions, and the computational efficiency of the associated algorithms.
Our contributions include the development of variational framework and algorithms tailored to isotropic Gaussian mixtures, as well as an empirical evaluation across synthetic and rorl-world datasets, demonstrating that our approach achieves a compelling balance between modeling accuracy for multimodal targets and computational efficiency. This paper is organized as follows. Section provides the relevant background on the geometry of the space of isotropic Gaussians, and on optimization schemes based on the Bures and entropic mirror descent geometries. In Section we introduce the general setting for optimization over mixtures of isotropic Gaussians with uniform weights. Section 4 presents our algorithms to efficiently optimize over this variational family. In Section [5] we discuss related work in the Variational Inference literature. Our numerical results are to be found in Section 6
2

---

Notation. We write a Gaussian distribution on R with mean m and variance as (mm, ela), where La denotes the d-dimensional identity matrix; and M(x; m, eI,) its density evaluated æ. We denote by P2 (2R")) the set of probability distributions on IR with bounded second moments. Consider μ., v (2(R)) the Wasserstein-2 (W2) distance is defined W v) inf seS J|1 yll2ds(x; y), where S(u, v) is the set of couplings between and w. We will denote k. the normalized Gaussian kernel on IEd with variance 6, i.e. (x) = (2 (ea0-I1a11////
For E (RR()) we denote by ke j its convolution with the Gaussian kernel, that writes ki 11 f ke(- x))dp(x) We denote the trace function.
2 Preliminaries (Gaussian VI)
This section introduces key concepts on the space of isotropic Gaussians, as well as different (time- discretized) gradient flows one can consider through the Bures Wasserstein or entropic mirror descent geometries.
2.1 Isotropic Gaussians (IG)
The space of isotropic Gaussians is defined as IG = {N(m, cla). m E R RR,CER" and is subspace of (2) When equipped with the W2 distance, this space has particularly tractable ge- ometric structure. Indeed, the W2 distance between two isotropic Gaussians N(m, ela), N (mm',, ( rla) takes the form of Bures- Wasserstentnnnn (BW) distance:
BW Na), rTa)) = [mm-m'lf B rTd),
where B denotes the Bures metric [Bhatia et al. 2019] between positive definite matrices and (eld, rId) d(c + T 2VTT). This formula reflects the separable nature (with respect to the means and variances) of the Wasserstein metric on the space of isotropic Gaussians IG. Interestingly. the metric space (IG, BW) of isotropic Gaussians equipped with the BW distance can be seen as submanifold of the space of (all) Gaussian distributions equipped with the same metrie, which can itself be seen as submanifold of the Wasserstein space (P2 (Rd), W2). Indeed, the BW geodesic between μ4, v E IG also lies in IG, see Appendix A.
2.2 Bures- Wasserstein gradient descent on IG
Recall that our goal is to minimize a functional objective KT<(-I=) defined in Eq 1 where и XX e-V, firstly on IG in this section, before being able to tackle mixtures of isotropic Gaussians. In this subsection we explain how to derive a gradient flow with respect to the Bures- wasssessssnttnin geometry, and provide a discrete optimization scheme. To this goal, we first define a minimizing movement scheme on IG. For po E IG and 7 > 0 step-size, define:
1 BW² KT(p)=) + 2ny
Pk+1 arg min
PEIG
which corresponds JKO cchemeJJodan et al. 19988] but where the solution is constrained to lie in IG. In the limit y - 0, we obtain Wasserstein gradient flow of measures projected on IG, i.e. a continuous curve (Pt): 1G decreasing the KL, and which is governed by differential equations for the mean (mz), and variance (Ez); (see Appendix A.3). Such flow can exhibit favorable dynamical behavior under strong log-concavity assumption the target distribution, demonstrated in the following proposition.
Proposition 2.1. Suppose that V z ala some a E R. Then, for any Po E IG, there is a unique solution (pt)z to the flow obtained limit of Eq ( y + 0. Then, for all 2 0 and p" IG,
(L(p.. KL(p'1m) S (L (L](p"m)))
implying that the flow converges linearly when a
The full proof of Proposition 2.1 is provided in Appendix A.2 and is direct application of the one
of [Lambert al. 2022, Corollary 3]. relies on the fact that KT-(-I=) is an a-convex objective functional along W2 geodesics when the target potential V is a-convex, e.g. [Villani 2009
3

---

Theorem 17.15]; and since We (equivalently BW) geodesics between isotropic Gaussians lie in IG (see Appendix [.I].. the objective is also convex on BW). This result indicates that strongly log-coneave targets can be efficiently approximated using isotropic Gaussians.
Here, we propose evaluate explicit dime-discretzzation of this gradient flow, it is computation- ally less expensive than an implicit scheme such as Eq (3). To this end, let F IR x IR+ defined as F(m, e) := (L(N(m, ela)( TT. Starting from some On (mo: 60) E 12d IR t* we consider the following scheme:
(VF(02)) 0x)
1 (No.No. 27
0k+1 argmin XE
(4)
denoting 0 :- (m, e) and Ngg := (m, ela). Note that the scheme above is similar to Eq (3) but where the objective has been linearized. Thanks to the decomposition of the BW distance given in Eq (2) it leads to the following updates the mean and variances:
m mg vV m Fmm cx)
2% F.F (x) Cx.
Ek+1 (1-
(5)
d
where Vn FmF(mmx, ex) = Ep,. [V] and V v.F(m, = 2 deu Eyg [[. VVV
cN ) We
observe that while the first update on the mean is a simple gradient descent, the latter update ensures that the variance remains positive and differs from simple Euler discretization of the associated differential equation (see Appendix (A.3). We provide the details of the computation as well as an interpretation of these updates Riemannian gradient descent in Appendix A.4
2.3 Entropic Mirror Descent on IG
We now turn to an alternative descent scheme on IG, namely mirror descent, which relies on geometry different from the W2 one described in the previous subsection. Mirror descent is an optimization algorithm that was introduced to solve constrained convex problems [Nemirovskij and Yudin 1983]. and that uses in the optimization updates a cost (or "geometry") that is Bregman divergence [Bregman I967), whose definition is given below.
Definition 2.2. Let á X I strietly convex and differentiable functional convex set X, referred to as a Bregman potential. The Ø-Bregman divergence is defined for any x, E X by:
(x(l(z) <(y) <(x) y
Further details mirror descent and its connection to standard algorithms such as gradient descent are provided Appendix A..
Hence, we propose to choose an appropriate Bregman divergence on the space of covariance matrices, namely generalized
BealbacckLeibble divergence between positive definite matrices: KL(A (A]B) = Tr(A(log A log B)) - (er)))((A))) Tr(A)+ Tr(B). The latter object, also called Von Neumann relative entropy, is a Bregman divergence whose Bregman potential is the Von Neumann entropy Φ A + Tr(A log A) (and where (A, B) = Tr(AB)). Note that (T.(<Ta|TTa) = d (c log & - - Then, we can define a descent scheme on IG as follows, starting from some 6u = (mo, cu) E IRd x I
0 argmin x
(VF(0L)) 0x) 24 Il ma
(T((eaa(aAaa))
27
denoting again 0 (m,e) c). Note that compared to the scheme of Eq 4 only the update the
variance differs, and given by:
G++ 6; exp (- 27 V. F(ms, (t).
(6)
see Appendix A.6 for the computations. This update, as the one in Eq (5). also guarantees that the variance parameter e remains strictly positive; and is known as entropic mirror descent [Beck and Teboulle 2003
Beck and Teboulle 2003 used this exponential update followed by renormalization to optimize over the simplex.

---

3 Mixtures of Isotropic Gaussians (MIG)
We now turn the problem of optimizing the KL objective to the target distribution over the family of mixtures of isotropic Gaussians. We will consider the VI problem Eq ①T for a specific setting where the variational family is the set mixtures of isotropic Gaussians, for some N N with equally weighted components:
C2
1 N
Σ
N(m², eTa), e E (R* xRt))
Note that any distribution v E C. writes V = E}= kig * 6mi, for some [m², ed E (R3 x R+* R+*),, where kes is the Gaussian kernel with variance el and smi is the Dirac m². Then, define our loss function F ( (RR-) + R as:
F(mm), el( == KL ( Σ
N(me, e
/=1
The following proposition provides useful formulas regarding the gradients of this objective.
Proposition 3.1. Let μ &V E}= dms and denote Kp TV Σラ=1 his kdms. Assume C (IR.).
The gradients of with respect to m2 et E RRd. x IR write:
V F((m) eNN ( 2N
RK +5
kip 1 ( m (--m²)
()(})
π
V ms ((mm) e _ ,/ Ek, +5
Vn
ki & 1 ())
π
The proof of Proposition 3.1 can be found in Appendix A.7 Note that the means and variances in the mixture interact through the terms V Vnnk @ (1) in the gradients. Remarkably, our computations provide expression of the gradient with respect the variance that only involves scalar product
(- (-m''""Vln K.@u which can be computed efficiently with a computational cost (d). In
practice, the expectations over the Gaussian components k. * 6 for j = 1,.... .N will be estimated with Monte Carlo integration.
4 Algorithms for VI on MIG
4.1 General optimization framework
We propose, for the optimization of the objective F' on (Rd)N x (R+*)N defined in Eq. 7 or equivalently KL(-IT) on C*V to perform joint optimization the means and variances of the mixture. This joint optimization involves gradient descent update on the means, and either a Bures or entropic mirror descent update on the variances. Our approach is summarized in Algorithm [
Algorithm 1 MIG optimization with IBW or MD
Input: initial means and variances (mg. d step-size 7, number of iterations T. for 1 to do
for i 1 to N do
Update of with IBW (Eq. 12 m² F' (mmg kkj=1)
Update mk+t1 = m4 - NNV
(GD)
MD (Eq. 13
end for end for
We now describe the optimization of the variance parameters using either Bures entropic mirror descent updates in the next subsections. These methods rely on careful adaptations the schemes introduced in Section 2 originally defined for single isotropic Gaussian, to the mixture setting.
5

---

4.2 Bures (IBW) update
This section extends the framework of Section 2.2 to Gaussian mixtures and presents a new formula- tion of the JKO scheme adapted to this setting. The We distance between two Gaussian mixtures intractable and does not admit a closed form, in contrast with the BW distance on IG. Then, we cannot obtain direct updates on the mean and variances for JKO scheme restricted to C.. To address this issue, we will represent a mixture with its associated mixing measure. [ambert et al. [022] proposed a similar approach to derive a Wasserstein gradient flow on Gaussian mixtures. However, they considered a mixing measure with an infinite number of components and did not provide a formal derivation of fully explicit discrete (in time and space) scheme.
Any uniform-weighh Gaussian mixture E
C. can be identified a mixing measure ý N j=1 ((ms,e/) on (TR x IR. IR.+) where for any J E TRd we have:
v(z)= NZx; m, ela)dj(m,e)
(8)
Note that the space CN allows for the full identification of a mixture, up to reordering of the indices, since the corresponding mixing measure contains only N particles with equal weights. This avoids the identifiability issues that arise when dealing with mixing measure supported
continuous (infinite) set of components, which constitutes an overparameterized model, see [Chewi et al.) 2024 Section 5.6]. See also Appendix [B.] for more details. Following Chen al. (2019], ,. first consider a Wasserstein distance between mixing measures denoted Whw. where the cost is a squared Bures-
Wasserstein distance, i.e. c((m, <), (m'; -)) = (W2(Mmm, (A(m, W rId)). We then construct the JKO scheme on mixing measures at each step is as:
1 W bio (p. Pb.)
Pk+1 arg min KL(U/T)+ (ml e )=
+ 2yy
where p NV Eド= ) ( ,)), Pk # ) (mg. and is defined in Eq 8 The resulting
Gaussian mixture is Pk+1 ʃ [m,, dla)dpe+1 (m,e). e). Since P and Pu are two discrete measures with an equal number of components, the above Wasserstein distance simplifies as:
W
1 bio (p. = min N
Σ (e.)) ol)) ,c 1a))))) #=1
(10)
where is a permutation of the N indices in the mixture. Solving the JKO scheme Eq 9 is
now tractable and we can compute the limiting flow as y - 0, since at the limit o (i) i. The
continuous-ting time equations of the flow in the isotropic case are given in Appendix B.2 They match
the continuous- time equations for the means and covariances derived in [Lambert et al. al. 2022 Section 5.2] and recalled in Appendix E..
Similarly to Section 2.2 we consider an explicit dine-disceetization of this flow, using a linearization of the objective in Eq This leads us to the scheme:
N
10+ arg min (VP(()))))) 1 10))
Σ
BNs (II)
(
2
assuming that (1) in Eq 10 for 7 small enough. Finally, the variance updates for the Gaussian components are:
IBW update For 1,..., N d+ =
2N V ( d
(12)
d
The update the means takes the form of Eq GD). Details of the computations are deferred to Appendix B.3 Ultimately, we obtain system of Gaussian particles (ml e that interact through the gradient of the objective.
Remark 4.1. Note that can characterize the discrepancy between our JKO scheme on the mixing measure and the original JKO scheme restricted to Gaussian mixtures. Indeed, the Wasserstein distance on the mixing measure is related the W2 distance on P2 (2 as follows: 0 < W (p. fu)- W2 (v, vx) S 2/ 2de where e* is the maximal variance of the mixtures v, Vg. This result is direct consequence of [Delon and Desolneux, 2020 Proposition 6]. see Appendix B.4 for further details.

---

4.3 Entropic mirror descent (MD) update
In this section we provide an alternative way to optimize the variances of the mixture, based on mirror descent ideas introduced in Section 2.3] In particular, generalizing to N components what we have done for Eq (6). and by analogy with the scheme Eq (II). we consider the following scheme:
10 IN arg min (01)
Σ
m I
N
Then, at step h 2 0, the udpate on the variances takes the form:
MD update
For j ...., N: 4.41 og exp
2NT
V. F(mm) <())=n)
(13)
d
while the update on the means remains Eq (G)), see Appendix
A.6 for the detailed computations.
5
Related Work
In this section, we provide an overview of relevant work on VI with mixtures of Gaussians.
Several studies have addressed VI for mixture models, emphasizing computational aspects. Gershman ct al. 2012 optimize an approximate ELBO using L-BFGS (a quasi-Newton method), relying on successive approximations of ELBO terms for mixtures of Gaussians. However, while the original KL objective in VI defines a valid divergence between probability distributions, their optimization objective departs significantly from it.
Lin et al. 2019 Arenz et al. [2018) propose natural gradient descent (NGD) updates on the natural parameters of the Gaussians for each component of the mixture, and on the categorical distribution over weights. These methods are unified and extended [Arenz et al.) 22033, which introduces computational improvements. Natural gradient descent differs from standard gradient descent by performing steepest descent with respect to changes the underlying distribution, measured using Fisher information metric. In other words, the natural gradient is the standard gradient preconditioned by the inverse of the Fisher information matrix. For exponential families, such as Gaussians, the natural gradient of the objective with respect to natural parameters coincides with the standard gradient of the (reparametrized) objective when expressed in terms of expectation parameters (i.c., the moments of the Gaussians). This has some pleasant consequences, including closed-form updates on means and covariances, since the natural parameter admits simple expression in terms of means and variances for Gaussians. The NGD updates (fixing the weights of the mixture) means and variances write:
1-1 h+1 c
1
2N Vel F(mm (.) m -m = N F(mm e.I.) (14)
We refer to Appendix ffor more details and the computations. In particular, the latter update does guarantee that the variances remain positive, and the update on the mean is multiplied by the current covariance. In contrast, our updates on the means and covariances are decoupled (except through the gradient of the objective), and updates variances enforce positivity.
2024 considered the setting where the variances the mixture are shared, equal to ela with gradient descent (GD). In that setting, they proved descent lemma showing that the KL objective functional decreases along the GD iterations, under some conditions including a maximal step-size, boundedness conditions on the second moment of the (distribution) iterations and a finite number of components N. They also provided upper bound on the approximation error, i.c. the minimal KL divergence hetween a N N-component mixture of Gaussians with uniform weights and shared isotropic covariance matrices between components is upper bounded as O (o)((M)), ), when π writes as an infinite mixture of these components. Interestingly, this bound valid for mixtures of isotropic Gaussians with different variances, as we show in Appendix [ Yet, fixing the covariances to a constant factor e of the identity, as in [Huix al. 2024], limits lot the expressiveness of the variational family. Hence, we focus in this work the more general variational family defined by C.N.
Huix et al., I +* that is kept fixed, and only the means (m mird....,.'' .m>) are optimized with
7

---

6 Experiments
In this section, we evaluate our proposed methods summarized in Algorithm I namely IBW and MD, on different experiments on both toy and real data. In practice, in all our experiments. the gradients in Proposition 3.1 are computed with Monte Carlo approximation involving B samples from Gaussian distributions in the components. Our other hyperparameters are the following: N the number of components in the mixture; y the step-size; and we also vary the initial values of the means and covariances of the candidate mixture. Our code is available at htttts://ggiie wttprrrtttttallaa///aaaaaa
com/
We compare our methods with different competitors. The algorithm presented in Lambert et al. [002), based on Bures geometry (as IBW) but updating full covariance matrices in the mixture, is abbreviated as BW (see Appendix [E. [E. for more details on BW). Note that the latter has a computational complexity and memory scaling proportionally to [(d + )" instead of [N + 1) for schemes IBW and MD. The algorithm of [Huix et al. [024)), considering mixtures of isotropic Gaussians with shared variance eld, that only updates the means with Eq. (D) while keeping the variance e fixed, is denoted GD. The Shhared-variance version of our schemes (where = 1, 1,.... .N, we impose the constraint é¹ = c, but the latter is optimized) are denoted IBW-s and MD-s. The natural gradient descent algorithm of et al. (2019). adapted to isotropic Gaussians with fixed weights and given in Eq (14), is referred to as NGD. We also evaluate Normalizing Flow (NF) implementation based on Real NVPs [Dinh et al. 2017]. We also used Hamiltonian Monte Carlo (HMC), an MCMC scheme (hence non parametric) and Automatic Differentiation Variational Inference (ADVI) Kucukelbir et al.
2017). All the experimental details are provided in Appendix E.2
Gaussian-mixture target in two dimensions. We first illustrate the behavior of our methods for two-dimensional target æ that is a Gaussian mixture with 5 components. In Figure [ (top), we evaluate our methods for N = 1, 5, 10, 20. We observe several interesting facts. Choosing a number of components N = 20 leads to the lowest final KL objective. This is in accordance with the fact that the approximation error with a mixture of isotropic Gaussians goes to zero as N tends to infinity, see Section The shared-variance versions of our algorithms (IBWss and MD-s) did not perform as well as the original schemes we propose, namely IBW and MD. This demonstrates the benefit of optimizing each component's variance to better fit given target's 's shape. In Appendix [ we visualize the optimized Gaussian components (represented as circles), highlighting the benefit of allowing each component to have its own variance. The NF method appears slower than ours without improving the approximation. We plot the approximated density with N = 10 for BW, IBW, GD and NF methods together with the target density in Figure (bottom).
N
N = 5
N 10
N 20
10
10-L
10-2
10-3
1805
10
10a
10
10a
10 iterations
10a
10a iterations
iterations
Herations
target
1BW
BW
GD
NF
Figure 1: Illustration of convergence of Algorithm I for a two-dimensional target distribution.
We also evaluated these VI methods on alternative challenging two-dimensional target distributions, these results are deferred to the Appendix. More precisely, we test the performance of these methods on a Funnel distribution and heavy-tailed targets in Appendix E.3 Then we investigate how well our schemes on mixtures of isotropic Gaussians can fit challenging Gaussian mixtures (e.g. with highly unbalanced mode weights) in Appendix E.4
8

---

haussian-mixture target in high dimension. We then consider a Gaussian mixture target with 10 components in dimension d = 20, and variational mixture model with N 20. Figure [ we plotted the marginals along each dimension together with the KT, evolution for the schemes (BW, IBW, MD, NGD).
Jerge
102
density
10 2
IBW
XD
BW
100
10 0 10
10
0
10
10
0
10
10
0
10
10 0
10
10-
10a
iteration
Figure 2: (First 10) Marginals (left) and KL objective (right) for MD, IBW, BW and NGD (d = 20).
We observe that both our schemes IBW and MD provide good approximation, along with BW. In the mean update of NGD, the gradient is rescaled by e, which leads large step for spread-out Gaussian. While this rescaling allows faster convergence, we observe that it makes the algorithm more sensitive to the initial conditions. Then, we compare in Figure the time per iteration and the KL objective value for BW (that updates full covariances matrices), and our isotropic version (IBW) for a similar target over several dimensions, for N = 15. We note that IBW performs comparably to BW, while enjoying a faster time execution, and still with cost in memory linear in the dimension instead of quadratic.
IBW BW
0.015
tims XL
22
0.010
0.005
0.000
10
20
30
40
50
dimension
Figure Time/KL evolution w.r.t
Bayesian posteriors. We evaluate methods two probabilistic inference tasks using classical datasets. The first one is Bayesian logistic regression (BLR) for two UCI datasets: breast_cance (2 labels, d = 30) and wine (3 labels, d 39). The second one aims to compute a Bayesian neural network (BNN) posterior on a regression task on the boston dataset using a single hidden layer neural network of 50 units (d 601), and the MNIST with (d = 203530). In each case, we assume a standard Gaussian prior on the parameters, and compute the posterior distribution given observations. More details given in Appendix [.6]. Note that the first task leads to log-concave posteriors (in contrast to the previous mixture of Gaussians targets which are typically non log-concave) while the second typically leads to a multimodal one [lzmailov
layer neural network with 256 units
et al. 2021
breast cancer
wine
MNIST
boston
10a
1.0
1.0
10a
RMSE
0.8
0.8
10
ADVI
1MC
0.6
10a iterations
0.6
10
10a
0.1
10a
10
103
iterations
iterations
iterations
Figure 4: Bayesian logistic regression and BNN regression approximated by mixtures of Gaussians.
we present the results of our algorithms IBW and MD for N = 5, on these Bayesian inference tasks, plotting the accuracy or Root Mean Squared Error (RMSE) score on the test set over iterations. The evolution of log-likelihood and unnormalized KL (ELBO) are deferred to E.6] As additional baselines, we compared our results with HMC and ADVI methods (except on MNIST -where we did not manage find a working set of hyperparam in such high dimensions). and only plotted the final samples results provided by stan. Other experiments, such comparing the
In Figure 4

---

performance of our methods with unimodal variational approximations, e.g. using N = within our framework or popular Laplace approximations (e.g. Diagonal and K-FAC [Ritter al., 2018) on MNIST, are also deferred to F.66 they show the relevance of using several components to capture different modes, for such multimodal posteriors. Regarding MNIST, for scalability issues, analogy with Blundell et al. [2015), we coded mean-field version of our algorithm (see also Appendix F.6 for details) where each weight marginal fitted by a Gaussian mixture with our Algorithm [ We observed that we achieved the same order of performance for IBW and MD. We also performed the optimization with NGD updates but found out that if the means were not initialized within a very small ball (increasing the chances of missing some modes in the target), the variance estimated by NGD could become negative, which made the optimization difficult practice. In contrast, our scheme guarantees the positivity of by construction.
Conclusion
Mixtures of isotropic Gaussians provide simplified yet powerful tool in variational inference, balanc- ing expressivity for multimodal target distributions with computational and memory efficiency. We presented two optimization schemes, that implement joint optimization descent, and on the variances through adapted geometries for the space of variance matrices, such the Bures Von Neumann entropy ones, guaranteeing that they remain positive. Our numerical experiments validate their relevance for different types of target posterior distributions. Future work include establishing more theoretical guarantees regarding our schemes and mixtures of isotropic Gaussians. For instance, comparing the approximation error of full covariance mixtures versus isotropic ones, would be helpful understand why observe empirically great computational cost gain for a very modest increase of the KL loss (c.g., IBW vs BW). Then, studying optimization guarantees for these schemes is of interest.
the means through gradient
Acknowledgements P-T and AK thank Apple for their academic support in the form of research funding. ML acknowledges support from the French Defence procurement agency (DGA). This project was also funded by the European Union (ERC, Optinfinite, 101201229). The views and opinions expressed are, however, of the author(s) only and do not necessarily reflect those of the European Union the European Research Council Executive Agency. Neither the European Union nor the granting authority can be held responsible for them.
References
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2008
Oleg Arenz, Gerhard Neumann, and Mingjun Zhong. Efficient gradient-free variational inference using policy search. In International conference on machine learning, pages 234-24--20 334.243. PMLR, 2018.
Oleg Arenz, Philipp Dahlinger, Zihan Ye, Michael Volpp, and Gerhard Neumann. unified per- spective natural gradient variational inference with Gaussian mixture models. Transactions of Machine Learning Research, 2023.
Manuel Arnese and Daniel Lacker. Convergence of coordinate ascent variational inference for log-concave measures via optimal transport. arXiv preprint arXiv:2404.08792, 2024.
Pierre-Cyril Aubin- Frankosrankoskki, Anna Korba, and Flavien Léger. Mirror descent with relative smoothness in measure spaces, with application to Sinkhorn and EM. Advances in Neural Information Processing Systems, 35:17263-17277 2022.
David Barber and Christopher Bishop. Ensemble learning for multi-layer networks. In M. Jordan, Kearns, and Solla, editors, Advances Neural Information Processing Systems, volume 10. MIT Press, 1997.
Amir Beck and Mare Teboulle. Mirror descent and nonlinear projected subgradient methods for convex optimization. Operations Research Letters, 3163):167-175, 2003.
10

---

Rajendra Bhatia, Tanvi Jain, and Yongdo Lim. On the Bures- Saesss-asssetteee distance between positive definite matrices. Expositiones Mathematicae, 37(2):165-191, 2019.
Christopher M. Bishop. Pattern recognition and machine learning. Information Science and Statistics. Springer, New York, 2006.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association, 1125118)839-877,, 2017.
Denis Blessing, Xiaogang Jia, Johannes Esslinger, Francisco Vargas, and Gerhard Neumann. Be- yond ELBOs: A large-scale evaluation of variational methods for sampling. arXiv preprint arXiv:2406.0744 07442,, 2024.
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International conference on machine learning, pages 1613-1622. PMLR, 2015.
Silvère Bonnabel, Mare Lambert, and Francis Bach. Low-rank plus diagonal approximations for Riccati-like matrix differential equations. SIAM Journal on Matrix Analysis and Applications, pages 1669-1688, 2024.
Lev M Bregman. The relaxation method of finding the common point of convex sets and its application to the solution of problems in convex programming. USSR computational mathematics and mathematical physics, 7(3):200-217 1967.
Diana Cai, Chirag Modi, Charles Margossian, Robert Gower, David Blei, and Lawrence Saul. Eigenvi: score-based variational inference with orthogonal function expansions. Advances in Neural Information Processing Systems, 37:132611-132221, 2024a.
Diana Cai, Chirag Modi, Loucas Pilliud-Vivien, Charles C Margossian, Robert Gower, David M Blei, and Lawrence K Saul. Batch and match: black-box variational inference with score-based divergence. International Conference on Muchine Learning, 2024b.
Edward Challis and David Barber. Gaussian Kullback-Leible approximate inference. Journal of Machine Learning Research, 14(68):2239-2288 2013.
Yongxin Chen, Tryphon Georgiou, and Allen Tannenbaum. Optimal transport for Gaussian mixture models. IEEE Access, 7:6269-622 6278, 2019.
Sinho Chewi, Jonathan Weed, and Philippe Rigollet. Statistical optimal transport. arXiv preprint arXiv:2407.18163, 2024.
Julie Delon and Agnès Desolneux. A Wasserstein-type distance in the space of Gaussian mixture models. SIAM Journal on Imaging Sciences, 13(2):96-977, 2020.
Michael Ziyang Diao, Krishna Balasubramaniann Sinho Chewi, and Adil Salim. Gaussian variational inference via JKO in the Bures- Banes-Wassee Wasserstein space. In International Conference on Machine Learning, pages 7960-7991. PMLR, 2023.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. Interna- tional Conference on Learning Representationss 2017.
Justin Domke, Guillaume Garrigos, and Robert Gower. Provable convergence guarantees for black- box variational inference. Advances in neural information processing systems, 2023.
Samuel Gershman, Matt Hoffman, and David Blei. Nonparametric variational inference. International Conference of Machine Learning, 2012.
Tom Huix, Anna Korba, Alain Durmus, and Eric Moulines. Theoretical guarantees for variational inference with fixed-variance mixture of Gaussians. International Conference of Machine Learning. 2024.
Pavel Izmailov, Sharad Vikram, Matthew D Hoffman, and Andrew Gordon Gordon Wilson. What are Bayesian neural network posteriors really like? In International conference on machine learning, pages 4629. PMLR, 2021.

---

Yiheng Jiang, Sinho Chewi, and Aram-Alexandie Pooladian. Algorithms for mean-field variational inference via polyhedral optimization the Wasserstein space. The Thirty Seventh Annual Conference on Learning Theory, pages 2720-2721. PMLR, 2024.
Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the Beaker-Planck equation. SIAM Journal on Mathematical Analysis, (9(1):1-17, 1998.
Ivan Kobyzev, Simon JD Prince, and Marcus Brubaker. Normalizing flows: An introduction and review of current methods. IEEE transactions on paltern analysis and machine intelligence, 43 (11):3964-3979, 2020.
Alp Kucukelbir, Dustin Tran, Rajesh Ranganath, Andrew Gelman, and David M Blei. Automatic differentiation variational inference. Journal of machine learning research, 18(14): 1-14, 2017.
Daniel Lacker. Independent projections diffusions: Gradient flows for variational inference and optimal mean field approximations. arXiv preprint arXiv:2:39.. 13322, 2023.
Marc Lambert, Sinho Chewi, Francis Bach, Silvère Bonnabel, and Philippe Rigollet. Variational inference via W Wasserstein gradient flows. Advances in Neural Information Processing Systems, 35: 14434-14447, 2022.
Hugo Lavenant and Giacomo Zanella. Convergence rate of random scan coordinate ascent variational inference under log-concavity. SIAM Journal on Optimization, 34(44:3750-3761, 2024.
Wu Lin, Mohammad Emtiyaz Khan, and Mark Schmidt. Fast and simple natural- natural-gradien gradient variational inference with mixture of exponential-family approximations. International Conference on Machine Learning, pages 3992-4009 PMLR, 2019.
Chirag Modi, Diana Cai, and Lawrence Saul. Batch, match and patch: low-rank approximations for score-based variational inference. Artificial intelligence and statistics, 2025.
Radford M Neal. Slice sampling. The annals of statistics, 31(3):705-767, 2003.
Arkadij Semenovič Nemirovskij and David Borisovich Yudin. Problem complexity and method efficiency in optimization. 1983.
Manfred Opper and Cédric Archambeau. The variational Gaussian approximation revisited. Neural Comput., 21(3):786-792) March 2009. ISSN 0899-7667.
Felix Otto. The geometry of dissipative evolution equations: the porous medium equation. Communi- cations in Partial Differential Equations, 26(1-2):101-174 2001.
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research, 22(57):1-64, 2021.
Arthur Pewsey. Sinh-aresinh distributions. Biometrika, 96, 11 2009. doi: 10.1093/biomet/as0053.
Garvesh Raskutti and Sayan Mukherjee. The information geometry of mirror descent. IEEE Transactions on Information Theory, 61(3):1451-1457, 2015.
Hippolyt Ritter, Aleksandar Botev, and David Barber. scalable laplace approximation for neural
networks. In 6th international conference on learning representations, ICLR 2018-conference track proceedings, volume International Conference on Representation Learning, 2018
Gareth O Roberts and Jeffrey S Rosenthal. General state space Markov chains and MCMC algorithms. Probability Surveys, 1, 2004.
Adil Salim, Anna Korba, and Giulia Luise. The Wasserstein proximal gradient algorithm. Advances in Neural Information Processing Systems, 22312-123666,, 2020.
Roman Soletskyi, Marylou Gabrié, and Bruno Loureiro. theoretical perspective on mode collapse in variational inference. arXiv preprint arXhv:2410.13300, 2024.
12

---

Charles M. Stein. Estimation of the mean of a multivariate normal distribution. The Annals of Statistics, 9(6):1 1151, 1981, doi: 00.214/aos/1166445533
Esteban Tabak and Eric Vanden-Eijnden. Density estimation by dual ascent of the log-likelihood. 2010.
Cédric Villani. Optimal transport: old and new, volume 338. Springer, 2009.
Rentian Yao and Yun Yang. Mean field variational inference via Wasserstein gradient flow. arXiv preprint arXiv:2207.08074 2022.
Rentian Yao, Xiaohui Chen, and Yun Yang. Wasserstein proximal coordinate gradient algorithms. arXiv preprint arXiv:2405.0460 2024.
13

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No]. or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors their evaluation. While "[Yes] is generally preferable to "[No]" it is perfectly acceptable to answer "[No] provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer is often more nuanced, please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: We propose memory-efficient variational inference (VI) algorithm for mix- tures of Gaussians in the abstract and the introduction. This claim is supported by constructive approach under the assumption of isotropic covariances, along with the use of Hessian-free update scheme. Furthermore, we demonstrate experimentally that the simplitied model does sacrifice approximation accuracy.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
14

---

Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We provide convergence guarantees only for the single Gaussian case (N=1) and under the assumption that the posterior is strongly log-concave. In other cases, the problem becomes nonconvex, and the most we can hope for are descent-type results, namely, that the KL objective decreases along IBW iterations. This guarantee holds for the shared, fixed covariance (Huix et al.) 2024, Prop. 4]. Extending this result to variable covariances, as well as characterizing why the approximation performs so well empirically, is the subject of future work, as mentioned in the conclusion.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results are violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
The constructive proofs of the IBW and MD discrete updates for both the single isotropic Gaussian and the mixture case are provided in the Appendix. Note that for IBW, we present two distinct proofs: one based on the derivation of the necessary optimality condition, and the other based on the projection The projected gradient interpretation is given for both the single Gaussian and mixture of isotropic Gaussian case.
Justification:
Wasserstein gradient flow.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch provide intuition.
15

---

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: For synthetic datasets, we provide detailed information in the Appendix on how the data are generated. For real data, we use only well-known standard benchmarks. The algorithm settings are also described in detail in the Appendix, and the code will be made publicly available to enable full reproducibility of the results.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to model. In general. releasing code and data often one good way to accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
also be provided via detailed
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
The Python scripts will be made publicly available on GitHub.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (kttps://nips.cd cc/
for more details.
16

---

While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https guides Swto//eobeissiiieeeeeae swwpiiiissooooooicco.co for more details.
7 //nips
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
All these details about the experiments are provided in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Error bars are not reported systematically, but only when they are relevant to the experiment.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
17

---

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
The assumptions regarding the computing resources are provided in the Appendix.
Justification:
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics / http://wwwrpp...ccocoa/cctttttta
Answer: [Yes]
Justification: our work is conformed with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations in their jurisdiction..
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification; the scientific impact is addressed in the Conclusion. The subject is too abstract to directly address socictal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinfornation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
18

---

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Not applicable here.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
The authors of the codes used are properly credited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, haperswwthcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
Not applicable here.
19

---

Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Not applicable here.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation,
or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: Not applicable here.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality the research, declaration is not required.
Answer: [No]
Justification: LLM only used for grammar correction.
Guidelines:
20

---

The answer NA means that the core method development in this research does not involve LLMs as any important. original, non non-standard components.
Please refer to our LLM policy (https://neurip:: /ttp://pp.ccc//offfeeeee//////// for what should or should not be described.
21

---

A Appendix
A.1 Geometric structure of Isotropic Gaussians space
We prove here the stability of the isotropic Gaussian model along the Bures- Wasserstein geodesics. We recall that the Bures- wrressasssrrttnin space is the space of non-degenerate Gaussian distributions equipped with the Wasserstein-2 metric.
Proposition A.1. The space of isotropic Gaussians equipped with the geodesically convex subset of the Bures- Wasserstein space, which is itself geodesically convex subspace of the Wasserstein space (P2(R2), W2).
distance is
Proof. Let p = (m EpTa) and q = -Vmmp Egla) be isotropic Gaussian distributions. Since p absolutely continuous, the Wasserstein-2 geodesic, (which is also a Bures- Wasserstein geodesic) between p and g is given by the pushforward measure:
Me ((1 t)Id (T) E [0,
where Id is the identity map and T is the optimal transport map between p and q. Denoting Ex cp and Eg = Egg Td, this optimal transport map T is the affine map:
T(x) my + A(r mp), where A Ep
1/2
A is called the Bures map and satisfies E, = AE,A, which can be easily verified from the definition above. Since the map is linear, it preserves densities, ensuring that the transported measure Pt remains Gaussian. The mean and covariance parameters evolve along the Bures- Wasserstein geodesic between and according to the equations:
m, = (1 - m) + tm
Σ, ((1 1)Ta + (A)E, ()Ta + LA).
Since both and are isotropic, the transport map is:
1/2
cg Ep
T(x) my + a(x mp), where
Then, the covariance Σ, evolves according to:
Σ = ((1 - La + taTa)c,((1 - t)Ia tala) = ((1 1) + 2 CcpTa.
which is clearly isotropic. Hence, the interpolated distribution fe remains isotropic Gaussian for all / E [0, 1]. Thus, the space of isotropic Gaussian distributions geodesically convex in the Bures- Wasserstein space. Since the Mures-Wassersee space is itself a geodesically convex subspace of the Wasserstein space, we complete the proof.
A.2 Proof of Theorem 2.1
A.2.1 Background on Wasserstein gradient flows
Let (2((RR R U {+00} be functional. We say that F is a-convex along Wasserstein-2 geodesics if for any two probability measures μo: Mi PR(Rd) and any geodesic {Me in the Wasserstein-2 space connecting No and μl, we have
F(w)) s (1- ()vo))( + (F(i) 1 () 13 VI E [0, 1).
A Wasserstein gradient flow of F is solution (U))EE(O,TT T > 0, of the continuity equation
0hu Øt
V (mmn) =
that holds in the distributional sense, where v is a subgradient of F at s4: Ambrosio al. 2008
Definition 10.1.1]. Among the possible processes (v:),, one has minimal T.P(p)) norm and is called the velocity field of (ue).. In a Riemannian interpretation of the Wasserstein space Otto 2001
22

---

this minimality condition can be characterized by o, belonging to the tangent space to (2(RR) at pa denoted Tue (R())) which is a subset of T, (/4)), the Hilbert space of square integrable functions with respect to ft, whose inner product is denoted "h The Wasserstein gradient defined this unique element, and is denoted Vw,)F(H)). In particular, if μ P2 (2())) is absolutely continuous with respect to the Lebesgue measure, with density in C* (RR) and such that F(u) oc,
V Ww, (x) (F) for p-a.e. I E IRe [Ambrosio et al. 2008 Lemma 10.4.1], where (1) denotes the first variation of F evaluated at jr, i.c. (if it exists) the unique function F(p): R + R s.t.
lim (F(u+ hg) F())) ((()(()))(())
h-
for all { 5= -リ- nt, where P2(RR).
A.2.2 Proof
The proof relies on tools on the Wasserstein geometry and calculus introduced above, and is a direct application of the one of [Lambert al. 2022 Corollary 31, yet state for completeness. When V²V 2 ald. FF f + KL(μ{)) is a-convex on (R (Villani 2009 Theorem 17.15]. Let us denote p* the minimum of this function and recall we denote BW the Bures- Bures-Wasserstein metric on the manifold (IG, BW). We consider the following gradient flow:
0p. = div (p: ww, with the initial condition Po = Po- 0t
We first want to show that the solution of this problem is unique. Let (pr)e and (qt) be two solutions of the above gradient flow. Then, using differential calculus in the Wasserstein space and the chain rule, we have
O₂ BW2 ạt) = 2 (logp, (4:), Fw,FF + (logg. (p:), F(a))an
where V F denotes the Wasserstein gradient at p (2 and log,. (g:) T - Id E (p:)). where T is the optimal transport map from pa to Moreover since F is a convex, we can write
dp. q E (2
F(g) 2 F(p) + (V. log,(a))
a (p)
2
Thus can write
0, BW2 (pc. 9:) S -2a BW.
Hence, by Grönwall's 's lemma, we obtain
BW a.) S exp(-2ot) BW (po, 9o)
Since both (pt); and (g:). are solution of the gradient flow, Po = go and implies that Vt [0, 1]. = dr- This proves the uniqueness of the solution.
Moreover, if a >> O, we can set q: = p" for all 2 to deduce exponential contraction the gradient flow to the minimizer p". Observe that by definition of the gradient flow, we have on the one hand that
a.F(p.) (V WWr, - Ww,F(p)),
[FF()))
(15)
On the other hand, if a > the convexity inequality and Young's inequality respectively, yield
0 0= F=F(p")) 2 F(p) + (V w. F(P) log, (p"))p + BB a V F(y)l
2 F())
1
2 [1og, (v")
||
a
BW
2a
2
(Bo"))
and hence || [(D]1 z 2a F(p). Substituting this into Eq
15 and applying Grönwall's
inequality again, deduce
F(p:) S exp(-2oc F(po).
23

---

A.3 JKO scheme for isotropic Gaussians
In this section, we derive the continuous equations for the isotropic Bures- Wasserstein gradient flow. Starting from the JKO scheme of Eq (3), we constrain the solution to lie in the space of isotropic
Gaussians IG. We follow the same method than [Lambert et al. 2022 Appendix A].
Let us recall the JKO scheme Wordan et al.
1998 Starting from N(mo, coTa) at time & = 0
we look for the solution = - (m, ela) of:
1 Bw
Pk+1 arg min KL(pTT) + 27
(16)
PEIG
The gradient of the KL with respect to the variance parameters is given by Eq 21
d 1
d 1 Tr V-V..
(L
Tr
VV...
2e
2E
2
Then, for isotropic Gaussian distributions p N(m,ell) and pr N( ExLa) have
BW3(p. = ||m -m&li + d (c + cx 22/ck))
and
V. Bp. (1-V ck
Hence, the first order condition of [T] yield:
(1
Ck
dy
Tr V + E = c (1- + Y Tr {22)))²
d
C (1 (1- ( de
1 Ep [[ m) vV') (17)
Developing equation Eq at first order in 7 we obtain:
= 27 1 J Ep ((- + 0(72))
EN
2
Bp |(- vv'
d
Taking the limit y - yields the differential equation:
2
= 2 Ep mv)
For the first order condition on the mean parameter m of 16 using the gradient of the KL w.r.t. the mean given Eq 20 we obtain:
[2)] + - -m m =
Ep
V log
(18)
Since Ep V log pl = 0 obtain, at the limit + 0:
= Ep V log Ep [V]]
Inspecting Eq [77 we observe that solving the JKO scheme yields an implicit discrete-time time update, where the expectations evaluated under the unknown distribution p. We now derive the explicit form of this update, starting from the formulation in Eq
Recall that for general Gaussians =N(m, 2), Nmm,, Ex). have BW p&) I|m
mill + Tr Σ + Σ. 2 /55'/2,, x//2)
24

---

A.4 Forward Euler scheme for isotropic Gaussians
Derivations of the updates. We now derive the updates given by the scheme Eq The first order condition on is given by:
Ck
2% -V. ex) === ek (1 v. V.F(mk,
d
Using the Taylor expansion for small step y given by (1 (1+z)-1 = x + o(z) we obtain at first order the explicit update:
= (1 (1- -1-2 V. Px))
27
d
The first-order condition m gives the explicit update:
V m F.mmFlmss ((m,,
Riemannian interpretation. The latter scheme can identified to Riemannian gradient descent, on the isotropic Bures-Wart Wasserstein space, i.c. the space IG of isotropic Gaussians equipped with the Wasseesesssrtteiin metric, that we will denote iBW = (IG. BW). To achieve this, we first identify the local tangent space of the isotropic Bures- Wasserstein space. The direction of the tangent vector computed by projecting the Wasserstein-2 gradient of the KL objective onto this tangent space (see Appendix A.2. for the definition). We then follow this projected gradient with step size 7, We can then retract back to the isotropic Bures- Wasserstein manifold using an exponential map. We now detail this approach.
Let F P2( + R a functional. The isotropic Bures-Wassen Wasserstein Gradient of F at E IG, denoted F(p), is the projection of its Wasserstein gradient onto the tangent space to iBW p. p = (mp, epla), this tangent space writes:
TRIB(( fo i a s(x mp)la E R S
which can be identified with the pair (a, 8) E Rd x R. Thus,
F(p) = (RW Vw, F(p) = argmin ||w Vw, F(p)l WET,iBW(R9)
The first conditions in Rd and s R of this problem yield:
1
Ep F(p) and == dep
Ep [[- Vvs"
Indeed,
V. / |la + s(x - mp) VwsF(p)(()]]dpppe))
=0
/ 2(a s(x mp) Vw. Fw:FF))())dd((r) 0
a= Ep Vw F(p)),
and
V. / 1 s(x mp) - (in)-Vw,F(()ee (w,p)(//(ec))))ppp(>>>
/ 2(20 mmp)" (a + s(x mp) V F(.s())))))p((())
send - - (e- (oy = 0
Ep [(x V
8 = dep
25

---

Back to our problem with F(p) (p(p)T), using Eq 20 and (11 and that F(p = V log (#)
(see Appendix A.2.II we have that
v. KL(p)r) Fm Eb) and V. kL(p/=) 2
F.F Ex).
We can follow the direction of the gradient from the tangent space back to the isotropic Bures- Wasserstein manifold using an exponential map which is available in closed form
expy(a, 8) ~N(mp + a, (1 + ()pg))))))))
(19)
where we have adapted the exponential map formula for the Bures- Wrrrs-WWssersteinn space Lamberr et al. 2022, Appendix B.3]. written as expp (a, S) = (mm + a, (S + I) Σp (S+ to the case of Isotropic Gaussians. We can then construct discrete update based follow: we compute the iBW gradient of F(p) = (p(p))) multiplied step size and map it back onto the manifold using this exponential map. Starting from p Am EgLa), this discrete time update is:
27 d
Pk++ expp, c IBW expp.
V. F.F(mm 4k),
which gives the update on the mean and variance parameters:
m6+1 m. V P (m EN)
27 v FmF(m Ex))
6641 = (1
A.5 Background on Mirror Descent
Mirror descent is an optimization algorithm that introduced solve constrained convex prob- lems [Nemirovskij and Yudin 1983), and that uses the optimization updates a cost (or "geometry") that is Bregman divergence [Bregman 1967] whose definition is given below.
Definition A.2. Let X' + R strictly convex and differentiable functional referred to as a Bregman potential. The Øi o--rrgmman divergence is defined for any x, E X' by:
convex set X,
Bø(v(=) = d(y) ((x) - (Vo(*), x).
Let y fixed step-size and G objective function X. Mirror descent with øi-Bregman divergence writes at each step l 2 0 as:
argmin (VG(xx), xx) - B2(x((a))
Th+
ZEX
Writing the first order conditions of the problem above, mirror descent writes
V (Vo(az) vVG(xx)),
Tk+1
where o (2)) supyex (v, .)) ((x) ( is the convex conjugate of d. and V Note that, if X' is subset of dRd and that d(x) || that B2 y) = 1 || 1/1 and mirror descent coincides with gradient descent. Yet, this scheme is more general and is useful for constrained optimization, as if one chooses wisely, the inverse V the so-called "mirror map" V maps the iterates into the domain of é.
A.6 Entropic mirror descent updates
Derivations of the updates (Gaussian case). We now derive the updates given by the scheme Eq (2.3) The first order condition on e is given by:
d log é 2
= (-2 F.F.mmm
(emmaee (x) F
ck
The first order conditions the means gives the same explicit update Eq
my F Fm eb)
26

---

Remark A.3. Note that Eq 2.3 involves factor 1/2 in front of the KL penalization term. Our motivation that in Eq ( the Bures distance BW cæId)
(vé v is a distance between the square root of the matrices and thus its derivative w.r.t leads to a term implying only the square root of the matrices, or in our isotropic Gaussian setting the square root of the scale isotropic
value: V. B.B e,Id) d (1 + V Iu
We find out was judicious have a derivative for the
KT. implying also the square roots of the isotropic coefficient, i.e., V.KL(eld)e.ld)) log (
d dloog( V
Mixture Case. In the mixture of Gaussians case, we derive the following updates from Eq first order condition on e¹ for = 1,.., .N gives:
4.3 The
(- 2N7 (-2N7
2N d log
Fim ( ) + + exp
d Ve
F(
and for the means:
m mi -N-V F((m) =0)
A.7 Proof of Theorem 3.1
In this section, aim to compute the gradients of the (L((Ix) objective function defined Eq 7 To achieve this, a fundamental tool are Stein's ' identities (Stein [981 which relate the derivatives with respect to the parameters the derivatives of the integrand.
A.7.1 Stein's identities for isotropic Gaussians
Lemma A.4. Let p be an isotropic Gaussian N (m, ela) and p(x) its density for Rd.. Assume that π C (Ra) and lim sll-+oo p(z) log (x)) O. We have
V m KL(p)T) = Ep
V Vlog
(20)
V. (T(p)=) = 2c
E - m) V log ())-
(21)
Note that the latter equation does not require to compute the Hessian (of dimension d x d). which can be computationally more efficient than an alternative formula given in Eq (23).
Proof. Recall that p(:; 0) = (2xc) exp (- [|==mmll and KT(ppx) Ep [log (2)] First, note that for 0 = (m, e):
Ve / log O)(=;; / logp(z;; 0)dz,
which comes from the fact that the expectation of the score function is null, namely V log pax; (x) 0)dx = 0. Hence,
V Bp log [ / log
P (x) V vmp(=)dx= = / V log (2(x)) p(z)dx,
where we used V (mp(s) V (x((x) and an integration by parts. We now compute the gradient of p with respect to e:
1 p((x) (2 ml| 1m1 4)
V.p(x)
2 Tr p(z) (- m)(x -- -111
Tr V.p(z).
(22)
Then, we have, using an integration by parts:
V V.E. Ep
log
7 =
P)
2 Tr log P (x) vopx(dd
= - - Tr Vap(a)V log
(2(x))
T
dx
1
EpI(- -[---m"' log
p
2e
27

---

Note that applying twice integration by parts would yield another formula:
V.E. log Er
! π
42
log
1
P
Tr Ep
[
log
(23)
2
2e
2
A.7.2 Stein's identities for mixtures of isotropic Gaussians
The application of the previous results for mixture of isotropic Gaussians is straightforward, Let 1 == ΣN Smi, and recall we denote the associated isotropic Gaussian mixture as kg @ μ (2()) is the
1 Σ ker 6m
Ei=NNNnn e'Ta) where k<(x) = (2mr) 4/2
Gaussian kernel.
Using the fact that the expectation of the score function is null and that V Wicke 0 the gradient of the KL with respect to the mean parameter m' is then given by:
N V mi * &
he 2 x V KL
he
1 N
he
V BR In
BK
V
In
With the same arguments and using Eq 22 the gradient with respect to the variance parameter e is:
li & u1 mT) BR In
kg @ # {)
2N ER 1 [c- (--m')"Vln
1
ke @ 1
V KL
Equivalently, note that for the latter gradient, using twice integration by parts would yicld the formula
V. KL k. 0 = 2N Tr +6
12 ki @ 14 x2)) )
B Bures- Wasserstein gradient flow for mixtures
In this section consider the Bures- Wasserstein gradient flow for mixtures proposed in Lambert et al. 2022 in the particular case of isotropic covariance matrices.
B.1 Additional discussion
Starting from mixture of isotropic Gaussians vo E CN at h: 0, the JKO scheme where constrain the distribution to be such a mixture writes recursively at subsequent step k: + 1:
D arg min |K(()=) + 2 wz(,,
(24)
VECN
Unfortunately, a closed-form solution for this scheme not available, as the Wasserstein distance between mixtures is not tractable. Regarding the geometry of the space of mixtures of Gaussians, Lambert et al. (2022) considered mixing measure with infinitely many components μ E P(O), which can be identified with Gaussian mixture of the form N dp(0), where Ng is Gaussian distribution with parameters E O. Transport maps can be defined for this model, and gradient flows can subsequently be computed [Lambert et al., 2022 Theorem 5]. However, this model is highly overparameterized. For instance, as illustrated in Chewi et al. [2024]. single standard Gaussian in IR can be represented by infinitely many mixing measures of the form (0, Id) = J N( rId) dp(x), where p (0, (1 7) for any [0,
An alternative solution is to constrain the mixing model by considering fixed and finite number of components with uniform weights as we do this work. This resolves the identifiability issue: in particular, co-localized components with weight cannot be confused with single component of weight 2w, since all weights are equal. See [Delon and Desolneux] 2020 Proposition 2] for further details on identifiability for the mixture model. Moreover discrete measures with the same number of atoms are stable along Wasserstein-2 geodesics. This property still holds if we consider discrete
measures the isotropic Bures- Wasessstsntttinn space p 1/ Σ (m)..))) We can therefore consider
28

---

discrete mixing measures with uniform weights as stable and identifiable representatives of Gaussian mixtures.
Following Section 4.2 we can reconsider the JKO scheme 24 the Whe mixing measures:
where we replace the W2 distance with
1 min (L(U)T)+ + W (R.
2
which, at the limit + is equivalent
N 2N Σ
min KL(U(T)+ +
(25)
The problem in Eq. 25
corresponds to the discrete scheme introduced for the full covariance case E' e in [Lambert et al., 2022 Appendix BJ, without further justification. This scheme can therefore be reinterpreted as gradient flow on the space of discrete mixing measures.
B.2 Flow of mixtures of isotropic Gaussian
We now solve the problem 25 when v is mixture of isotropic Gaussians. We recall our notation for an isotropic Gaussian mixture: K2 Sp E N(m², eetaa) where kes + W=W"m,, (mmm,eeLa)) as well as our loss function:
N
F(mm*, e el = KL
Σ
N(m La)) ア
KL(Ki @ M(T).
The derivative of this loss w.r.t. the parameters of the Gaussian mixture are given in Appendix and we report them here:
A.7.2
hig 8 1 ]
F
Ek,
Vin
(26)
N
1
[ (--m')Te m ke & 1
F
Ex Abms
(27)
e
Using the derivative of the Bures-Waad Wasserstein distance computed in Appendix order condition w.r.t. for the JKO-like scheme (25):
A.3
obtain the first
d d(1
d1-
c c = 2N
2
F
verk'
and following Appendix
A.3 as + O we obtain the flow:
2 E+ E (-- (--m²) m-" In
ke & 11
On the other hand, first order condition on the mean gives the implicit update:
m = m NY m.
and at the limit -> we obtain the flow:
hi 314
m
Ex,,
Vnn
π
B.3 Diserete update and Gaussian particles
Derivations of the updates. For the scheme given in Eq (II. we obtain the following first-order conditions on the parameters for 1...., N:
ck
V V. [[mm e (() ss E = (1+
2N7
v. F[mm <(1
1
2N
d
29

---

and using Taylor expansion for small step size (1 x)= = -1-2 + o(z) we obtain:
é' = ch (1 (1-2N7 Ve F F(mm <()3=1) 2
2N
(28)
The first-order condition on m² gives the explicit update:
m m N/V F ((m)) )
(29)
Riemannian interpretation. We now show that this update has also Riemannian interpretation, extending our geometric analysis from Appendix t.44 to the case of mixtures. We can compute the isotropic Bures- Wasserstein gradient for each Gaussian component, by projecting the Wasserstein-2 gradient of the KL objective to the IBW-tangent space of each component. Each Gaussian particle (component) follows its own trajectory ruled by V. Adopting this point of view, we obtain the following system of updates for 10r1=11--- 1,...,N. N:
P exppi, (BB))))
where exppi is the exponential map from the iBW tangent space at pt defined in Appendix A.4 then get:
We
F(v) projr (B4)
(BW(R) V W. F(v)
argmin wET, (BB((4))
||W ww,
with w = (a, 8) E Rdd x R. Together with 26 27 it gives:
a = Ep V Ww, F(u) NV m F F(mm² (
8 de' Ep [[- m V Fv,F()) = 2N
V F([m², c 11)))))
Using Eq (19 we obtain the discrete updates Eq
28
B.4
Background Wasserstein distances for Gaussian mixtures (Delon and Desolneux 2020
In Delon and Desolneux 2020), the authors introduce the MW, distance as a Wasserstein distance between Gaussian mixtures, where the transport plan is itself constrained be Gaussian mixture (with any number of components). We denote by GMM M the latter space.
Namely, let Po- P1 being two general Gaussians mixtures po = =aa, TPPó and p1 = =EK, ripl. the MW2 distance is defined as:
(P P1)
min
||=|
x2ll x2112d7(11100 .2) 2 W (po. P1),
which is upper bound to the true Wasserstein distance since the transport plan is constrained.
It has been shown in DDelnn and Desolneux
2020 Proposition 4] that this distance is also equal to:
MW Pı)
min WE
Why
N,), Eg)),
where S(TO, TI) set of coupling matrices between the vector of weights mo and कi of the two mixtures defined by S(mo, m1) = (W E MK. (R=x,, VVi, E, Wij = Té: Vj. Σ, Wij = m/}, where we note Maxp (R+) the set of matrices of size with positive values. Moreover, Delon and Desolneux 2020) showed that the optimal transport plan takes the form:
Ko Ki
(x.
W Wi,p6(()6 (z)
where W is the optimal coupling matrix and T.BW is the BW transport map from Gaussian component Gaussian component j. The transport plan 7 a GMM with at most K K2 Gaussian components which are degenerated.
30

---

where Wew is the Wasserstein distance between mixing measures defined in Section 4.2 The lower and upper bounds for MW2 given in [Delon and Desolneux 2020 Proposition 6] then transfer to Mixture model with fixed number of components. exactly N Gaussians with equal, fixed weights:
We now consider the case of mixture of over [I....,N.. we have:
p N Σ
M(m
Σ
Σ P
N
where p¹ the 8th Gaussian component of the mixture.
The 2W2 distance takes an even simpler expression when considering a distance between two mix- tures po- Pi of exactly N Gaussians with equal, fixed weights. If we note GN the set of permutations
N Σ
MP. PI)
min O N
BW (m) (a())) = W (Po. P1),
Who he and give:
N Tr
2 N
W2(Po. PI) bmn (p, Px) S W2(po.P1) +
TrE
k=1
The last term simplifies for isotropic Gaussians and we finally obtain:
0 < (p. - w vx) S 2V2
(30)
where e is the maximal variance of the mixtures When e - 0, such that the Gaussian mixture degenerates into an empirical measure, the two distance matches.
Geodesics on mixtures Finally, when considering mixtures with N equal weights, the transport plan has exactly N components and can be written:
7x,
Σ (6(x)) (m)()))
N
such that the mixture model with exactly N components is stable along the geodesics transported by this plan. Indeed, the intermediate measure between two GMM pa and pi is given by the formula for t [0, 1]:
Fr Patty where Pi(x) (1 - t)x + ty.
Applying this to specific case, obtain:
N Pt Σ ((1 t)Td + IT (6) #rb
(5)
where pt has exactly N components, proving that our GMM structure with N components is stable along the geodesics.
We may wonder if transport map exists in our simpler framework of mixtures with a fixed number of components and equal weights. Unfortunately, this is not the case, as illustrated in Figure [] where we observe that the map between the two mixtures is not bijective and cannot be represented by function T(x).
C Natural gradient descent updates
In this section, we give more details on Natural gradient descent on IG, which corresponds to algorithm proposed by Lin et al. 2019). NGD adapts standard gradient descent to the geometry of the parameter space, by preconditioning the Fuclidean gradient with the inverse Fisher information matrix.
31

---

OT matrix
Tmap
40
100
20
40
60-
80
Figure 5: Optimal transport plan between two mixtures of two Gaussians with equal weights 1 On the left is the Wasserstein distance W2 case, where the optimal transport plan is not constrained. On the right is the MW2 case, where the optimal transport plan constrained to be a mixture of Gaussians. In the W2 case, there exists bijective map. We don't have such a bijective map for the MW2 case. Indeed, in the right figure, some points have two images. These figures are generated using the Python Optimal Transport library (https://pythonn gi tttb.iibo////
Exponential Family An isotropic Gaussian p Nmm,, cla) belongs to the exponential family. Its density, in the canonical form, writes:
p(2; v) ) exp (11, A(1))
with the natural parameter 7 (m 72) (m - the sufficient statisties S(x)
(ar [|-||12) T and the log partition function A(r)
Inll" á log(-2n2). It follows that 242
V AA))) = Ep(-In) [S(x)] == 3 and A²A(r) = I(n) with 3 and ((v) being respectively the mean parameter and the Fisher information matrix. Also VA-1= = VA where (y) = sup,, A(n) the Legendre transform of A.
NGD as Mirror Descent (Raskutti and Mukherjee 2015 Let + (() be the optimization objective in the natural- aatural-arramtter space, and define the corresponding mean-space objective ,3 + f(3) A"
The natural-gradient step on f of size 7 is
M = π - - = - 7 --((AA((m))) 'v fm))
(31)
On the other hand, mirror descent on with the Bregman potential x(3) (B) updates writes
VA*(3e+1) = (3.) vf
(32)
By the chain rule,
V vff( V2 A Vf(m)) (-2 V(()),
substituting this into the mirror descent update recovers exactly the NGD step above. The correspond- ing optimization scheme is:
1
argmin V (f())), Ba) + BA (31
3
7
32

---

where the geometry is induced by the Bregman divergence B generated by the Legendre transform A" of the partition function. This divergence is equal to the KL divergence between two isotropic Gaussians:
1 (a.
|lm mull d dlog Ek
KL(N(m, cld)
(33)
EK
Ck
Explicit Updates in (m,e) Define P(m, e) = f(m, [Imm de), and let vf(3k) = (91 92)
with V mmF' 91 + 2mg2, V.F = dgz. Equivalently 31 and (32
give
m8+1
ma
m&+1
ma Ek+1 791Ek+1
? vF((3))
1 c++T
1
2
20p+1
Z6k
Ex
The variance update writes
27 +
v. F.FLm
(34)
cx+1
d
and
m
791
ck+1
Ek
m&
Ek
791
6k
1
264-192
mk
? [m F(MRR <k) 2ma222)
Ek 1 2c=
m&
1
+ 27/2
~ mmmk(mm::
C₁ 1 Y99
E/
m7
V m F F(m Ek)
mg
7ck+1 V F ((mx; ck).
(35)
Mixture case
We now extend NGD to mixture of N isotropic Gaussians with equal weights 1/N:
1 N
p(z;n)
(2m) (2m) 4/2 , S(z)) A(n'))
with r/' = ( n T exponential-famiy components, we apply NGD to each component.
m
for = 1,..., N, which is a convex combination of
Writing natural gradients in (m², e²) -space gives
m
m ck
9
c+1
Ny
04
2 2
cg.
where for each j
gf
(13' N
Bd
1=1
and V mi F= of 2m's 9/. F = d o Equivalently, the updates are
1
2N
+ V.F
7
k41
m ma+1 = ml N N
F.
Remark C.1. Note that while KL(-I-) is known to be Bregman divergence the space of probability
distributions over R& AuubinFFrankowsii et al. 2022). it is not a Bregman divergence on Rd x R+*
Indeed, note that Eq (33 does not decouple the mean and variance terms, resulting in coupled updates
in Eq 33 35)
33

---

D Approximation error for mixtures of isotropic Gaussians
In this section, we investigate the approximation error that can be achieved within the variational family we consider in this work. Previously, [Huix et al., 2024 Theorem 7] established that the approximation error of VI within the family of mixtures of Gaussian distributions with equal weights and constant isotropic covariance in the (reverse) Kullback-Leileer tends to 0 N tends to infinity, under the assumption that the target distribution writes as an infinite mixture of these isotropic Gaussian components with same covariance. Below, we derive similar result for our richer family, i.e., mixtures of (isotropic) Gaussian distributions with equal weights (and possibly different covariances). Note that and computations would the same for Gaussians with full covariance matrix.
deriving the results for mixtures of isotropic Gaussians, but the result
Assumption D.1. There exists p" IRd IR+** such that the target writes
π = 1 k.th dy (m,
(36)
where kz"(()) :- (x 1::--10 m) for any Rdd.
Recall that we use the notation PN ki 2p kT dp(m, c) with μ Ei= 6m for an isotropic Gaussian mixture with N components (where p t Eÿ=1 (im) , and in the notation kgg μr, is identified to the vector (cI...., We now state and prove generalization of [Huix et al. 2024 Theorem 7].
Theorem D.2. Let CN = { N(me, mn.el3 (R" x Rt))N }. Suppose that Assumption D.1 holds, then
log( N) +1
(m) (m) dr.
min KL(p)T) & C PECN
where C² =
N
Proof. We denote
DN min KL(pTT), nECN
PN argmin KL(plT). PECN
For any m E Rd, we consider PNin E CN+1 defined as
PN (1- (1-a) + akm
-. with a = N+T. By definition of DN, we have that, DN+1 S (L(pW+flm) Denoting f(x) x log v, we have KL ( f ( where we define:
m. PN+1
L. PN
km
FN+1
(1
Po
π
Defining B(x) zlogs-rt1 for x E R1 V11. By Lemma D.3 this function decreasing; and [1-112
since (x) ro(x) Va. we have B(rN+1(x)) < B(ro(x)) Ie follows that
f(rN+1) = FN++ log(rN+1) S TN+1 + ((ro))
km
1 B(ro) + +a mm _ 11²
=o +
π
I.
2 km
2a k. π
+ tro toolog(r))
B(ro)
B(ra)(ro 1). (37)
34

---

Moreover, we have:
DN+1
Dx+1dp' (m,
- / ( (m, e)
JJ fr (dx+d))dddp" (m,c)
$ f km dndp" (m, c) + / log(ro)dzdp (m, c)
π
+ a
m
+
B B(ro)dxdp"
+
+
20
B(ro)(ra 1)drdy (m,, from
π
ro(z) (og(((())())))))d) + a² ʃ
1.mh (x)
a+
B(ro()))de dp" E)
<(x)
/1 / / [\""Brr((t)(ro(() 1)dr (m),
+ 2a
/ (0(0)))g((7(()))d(()
a
(a)
//
(x) ((ro(=)))da dp" (m,
+ a²
(b)
π(z)
/
+ 20
B(ro(z))((o(s)) ())(x)
(c)
We first observe that we can write
a in function of Dv- Indeed, ro(x) = (1 - 2 so
ro(x) log(ro(z))dT(())
/ (1 a) PN
log a1-a)_0)
PN di
π
π
(1 - aa) / pN((x) (z) log
(N x(x))
dx + (1 a) log(1 a))
= (1 - (1-a) DN + (1 ax)
-x)))log log(1 - a).
For the second term ( we have that lim Bxx) and since B decrease, B(x) S 1, thus B(ro(z)) I 1, this implies:
b a² J
kT (x)²
B(ro(x))dr dp" (m, s a² //
k.hh (x)²
dr dp (m,e)
TΦx)
π(x)
=²
And for the third term (c). we have that -2))(x VJ see Lemma D.4 Thus,
cc 2a: B(ro(=))(ro(z) 1)x(z)dr S $2a 2a / (Vro(z) 1) a (x)
= =2a 2a / 1 (1- (1- TT()) (x))dx 20x
PN (x)
=2aVI-o - / VPN(z) <(z)dx 2a
2 - a (1- (1-H2(pN:*))) x)) 22x
2
- - 2a.
35

---

where we have used the definition of the squared Hellinger distance (f. - =1-1 (f(z)g(z)dr and the property stating that for any densities of probability f.g. 0 - (ffgg)) S Finally, we have
DN+1 S DN -o) + +2a (2(2a111 a
<1-a) + a²c²
(38)
using Lemma
D.5 stating that a + (1 - a) log a) + 2a (vI- - 1) 0
The previous inequality (38) is true for any n 2 and recalling that a x+1
we have,
Dat1 S (1 - Dn + a a²
(n+ 1)Dn+1 -nD, S
C
n1 n+
nD, Σ (n + -DD, S C
NDN S (log( NN +
DN S C2 (og(N) + 1
where the Harmonic number En= x+1 has been bounded by log N) +
In the proof above we used the following lemmas from (Huix al. 2024] (Lemma 8 to 10 therein).
We provide their proofs for completeness.
Lemma D.3. The function ((x)
log
E B [1}, decreasing.
aa 1y
Proof. For all x E [+1 the gradient B writes:
VB(r)= (x - log x 2(x log +
(x 1)3
For E (0, 1), the denominator is strictly negative and the numerator strictly positive, thus VB(x)= S 0.
For æ E (1, oo), the denominator is stricly positive and the numerator is strietly negative, thus (B(r) O.
So B is decreasing on both intervals, and lim = and lim = by Hospital's rule.
Lemma D.4. The function B satisfies: (xx)(x 1) S - 1 (
Proof. Let C(x)== == (xx) - 1) x log x-2+1
For x E (0, 1), log 2 implies
log --1
1
x zlogg --1
< x Vi,
V2
V
For E 00), log < implies
log --1
<
glogz
2
Va.
v
z-1
vr
and C(x) - V -1
v < 0. s-1
log
Lemma D.5. Let consider
Vn E N, then (1-a)))(1a))))) logg) a) (VI 1) < 0.
36

---

Proof. We have:
a (1 1- - (r) + 2a (VI-a-1) - aa a --0011--)00000000000 20r 0x
Ia aa - 2 2+
2a-2+2v1-2 00
using log(1 a) < aa and the fact that 1/ 1) < 1 for the first and second inequality. For
the last inequality we have used the fact that the before last expression decreasing and is equal to when a goes to 0.
E
Additional experiments and details
E.1 Updates for the full-covariance matrices scheme of Lambert et al. [2022) Section 5.2]
In this section detail the updates for the full-covariance scheme of Lambert al. 2022 Section 5.2]. The parameter space is θ = IRd x $4 (the space of means and covariance matrices). Consider initializing this evolution finitely supported distribution po:
po
N
(m/ 2
It has been checked in Lambert et al. finite mixture distribution:
[022
that the system of ODEs thus initialized maintains
N
1 Pt
N
0
(m x()))=
where the parameters system, for i E [N]
8.0) (m,
(i) (1))) evolve according to the following interacting particle
m (9) = EV In Vt (y(()))
1 2 In 1/ (r,
(8) 2:(1)
11
EV In V (x)
where y:(0)
Po(0) and f (o). In their experiments, the ODEs were solved using fourth-order Runge-Kutte Kutta scheme. our experiments we applied the BW SGD described in Lambert
et al. 2022 Algorithm 1].
E.2
Experimental details
As mentioned in Section 6 we detail here our experimental setup and hyperparameters whose values are provided in Table
Initialization of the variational mixture: For N N given number of components, initialize the variational mixture by sampling the means in a ball of size [-8 where E IR and setting each covariance matrix r La, where t For the GD algorithm (mean optimization only, following Huix et al. [2024)) the variances are initialized in the same way but kept fixed during optimization.
Optimization hyperparameters: We set the step-size y, the number of iteration Nierr the number of Monte Carlo samples to Bynd 10 for gradient estimation, and to BKL 1000 for the KL objective estimation.
Normalizing Flows: For the NF baseline, we used simplified RealNVP architecture Dinh al. 2017], based on the code available /tttps://githii http://rrlouugabbbii with coupling layers and hidden
tutorial- samplinnsmmppngg enhanced- geee
37

---

mog Table 1: Hyperparameters
MOG Init.
Optim
Targets
Exp.
1ittr
Figure
mog
15
2
10-1
103
{Aig: Fig. Nigh (8,5,5)
Figure
mog
10ad
10
10
-0-2d--
10s
( Tigs Muj (102d-1,,,55
Figure
mog
30
100
10-2
10
(sig: Fig: Ni)) 110,1,10
Figure
breast_cancer
20
10
10-2
10
prror 100
wine
20
10
10-2
102
o2 prior = 100
boston
10
10
10
6
10a
{o poprior h3 = (10,50}
Figure 6
funnel (a)
5
0.5
10-2
10a
o² 1.2
sinh-arcsinh (b-1)
10
2
10-3
10a
skw (-0.2, -.2))
sinh-arcsinh (b-2)
10
2
10
10
10a
skw = (-0.2, -0.5)
mog (c-1)
10
5
10-2
104
{pt, rig, Not (3,2,4]
mog (c-2)
10
5
10
10a
{pt, Thg Nigt 14,1,4]
mog (C-3)
10
5
10
10
{pt, Tig. Nig) = [3,2,4]
Figure 10
30
100
10-3
10a
{shg Tig: No¹ 520,5,5)
Figure 11
mog
30
100
10-3
10a
{sig: Fig. Nig) {10, 10, 10)
dimension 124, which yields a Neural Network (NN) with 4976 parameters dimension Our isotropic mixture model has (Nd + 1) 3N parameters, even for large N, the NF model remains more complex and costly to optimize. Therefore, for each target distribution, we tuned the learning rate and number of iterations for the NF method separately, rather than using the same settings as for the VI mixture methods, since their optimization dynamics differ significantly.
To generate target MoG distributions, as in the initalization of variational MOG, fix 8ig and Ngg Each component covariance matrix is constructed by sampling a random symmetric positive-defintte matrix (full diagonal or isotropic) and scaling it by rig. We draw raw weights uniformly in f,....,2.0 2Ng and normalize them to one. In Figure c) all weights are equal except in case (c-3), where one component has weight 0.1 and the remaining ones share the remaining mass cqually and the component means are placed at tipt,=pt)
MOG targets:
Datasets: We have used popular datasets from the UCI repository, as well as MNIST. The training ratio has been set to 0.5 for UCI datasets and 0.8 for MNIST.
Computational resources: All experiments (except MNIST) were conducted on a MacBook Air (M3, 2024) with Apple M3 processor and 16 GB of RAM. The MNIST experiments were run on an NVIDIA 50-90 GPU. Experiment runtimes ranged from a few seconds to up to two hours.
E.3 Additional 2-D examples
We present more experiments on 2D synthetic target distributions on Figure These target distribu- tions defined below.
38

---

100
10
15
101
10
10
10
10
10
10
10
10
Banalatt
Aauilat
Aatubitt
(4)
(b-1)
10-0
10-1
10
107
12
12a
10
12
10
20
10'
10
12'
30
10'
(b-2)
(e-1)
10
10
10
10
10:
12
101
10:4
10
12
iivutions
(o-2)
(c-3)
Figure 6: Evolution of the KL divergence over iterations together with the optimized variational mixture density on different type of target distribution: (a) Funnel, (b) sinh-arcsinh normal distribution and (c) Gaussian mixture. Optimization performed with different methods (IBW, MD, BW, NGD, GD and NF) and varying N values.
Funnel distribution: The funnel distribution
Neal 2003 in dimension d = 2 has density
p(01,02) (x1;0,,0))
(x2;; e,e").
for x (x1, E R22 We follow the setting of Cai et al. [2024a] by fixing 0² 1.2. Although unimodal, this "funnel" shape difficult to capture with isotropic Gaussians. We experimented with N = 5, 20, 40 components, but even for large N, our isotropic mixtures struggled, and the BW and methods still outperformed them.
Sinh-aresinh normal distribution:
This distribution Pewsey 2009 applies sinh sinh-arcsinh transfor- mation multivariate Gaussian to control the skewness skw and tail weight Let
Zo (m,, E),
Z = sinh(r r sinh -Zo) skw).
In our experiments, we use = (0.8,0.8), m = (0, Σ (11 0.4 ) and vary the skew
0.4
parameter skw specified in Table
In Figure 7 visualize the optimized Gaussian components of the target density, highlighting the advantages of allowing each component to have its own value.
E.4 Bias induced by uniform weights mixture
We investigate the bias induced by our choice of the variational family, i.e. mixtures (of isotropic Gaussians) with uniform weights. Our goal this section is to analyze how such a constraint affects the approximation of a bimodal target distribution with highly unbalanced mode weights.
39

---

IBW
MD
IBW
IBW-s
MD-s
IBW
MD
Figure 7: Optimized isotropic Gaussians (represented as circles) over the target distributions for the proposed methods and their shared-variance variants. (a) corresponds to the Funnel distribution Figure 6 a); (b) corresponds to the target used in Figure
We consider the target a two-dimensional Gaussian mixture with components and weights (p, 1 p) p) where 0.1. We seek the optimal approximation within our variational family of N isotropic Gaussian components with uniform weights. To mitigate optimization issues and being trapped in a local minima, we initialize the means of our variational approximation by sampling the initial components means in two small balls centered around each target mode (where the number of initial means sampled in ball is proportional to Np and N(1 - p) respectively).
As expected, when N 2, the weights constraint induces strong bias: the approximation either ignores the low-weight mode or assigns equal mass to both modes. However, as N increases, this bias progressively vanishes. For sufficiently large N, the model reallocates more components the dominant mode and fewer to the lighter one, effectively recovering the correct proportions despite the uniform-weige weight constraint.
To quantify how well the variational mixture captures both the mode locations and their effective weights, we drew B = 10, 000 samples from both the variational distribution and the target distri- bution and applied K-means clustering with Relusters = 2, then analyzed the cluster assignments and proportions. The reported emrors are: E₁ the mean Euclidean distance between the target and variational modes, and Eg the absolute difference between their corresponding cluster weights. The first metric captures whether the variational modes well located, while the second whether their relative proportions are well captured. Each experiment repeated 10 times; we report the average result and error bars. ing that uniform-weight Gaussian mixtures can approximate highly unbalanced multimodal targets remarkably well when sufficiently overparameterized. We also display the estimated variational density and Gaussians components (circles) along with the target density.
measures
shows that both errors rapidly decrease with N, demonstrat-
E.5 Additional high-dimensional mixtures
We first provide the full marginals of the experiment described in Figure 2 in Figure performed experiments in 10 Figure and d = 50 Figure II
We also
deasity
-10
-00
10
-10
10
-
10
-10
10
-10
10
-10 0
10
-10.0
10
-10
10
10 0
10
Figure 9: Marginals for MD, IBW, BW and NGD (d -d=20 20).
40

---

0.4
IBW
MD
0.3
0.22
N values Nm5
N valuuss
N=
N=40
N=100
IWW
target
Figure 8: Errors, visualization and estimation of the bias induced by the uniform weights constraint function of N for two dimensional unbalanced mixture of Gaussians target.
10a
A
10a
density
IBW
10-1
MD
NCD
-25
0
25
-
25
25
25
25
25
0
25
25
0
25
10"
10a
iterations
Figure 10:
Marginals (left) and KL objective (right) for MD, IBW, BW and NGD (d 10).
40
102
M
10
JBW
M
M1
10
BW
NGD
10
10a
iterations
Figure 11: Marginals (left) and KL objective (right) for MD, IBW, BW and NGD (d 50).
E.6
More details on the Bayesian inference examples
In this section we provide some background on the Bayesian inference examples of Section as well additional experiments.
F.
Definition of the target distributions
Let D (/,,,)))< be labeled dataset, where v, Rdd and y is the associated label.
Binary logistic regression: We model the probability of a binary label E {0, 1} given x; and parameter 2 E 2d by
Thi 4; 2) (x =) (1 (x 2))²
where = =
+ ) ) is the logistic function. The likelihood is
C(D 2) (3/ xi, z),
41

---

and the log-likeell likelihood is
((D 2) Σ
log (vs 22,22
Σ [u (x) 2) log(1 + c)
im1
With Gaussian prior x(2) (0, o prior Ta), the posterior is
<(2 D) & L(D 2) (())
(39)
and its gradient is
V. Tolog((2) D) = Σ (i o 2)) vi 2 1=1
Multi class logistic regression: For I. classes, let = = (21....,
..h....,,20 21.
with each 2ạ E Rd. Then
exp(x, 2t)
(3/: = 1 24, El= exp(x (0 21)
1= 1=1....
Linear regression: The classical linear model is
= z I4 + Eis
& ~ ~0(,00)
so that
y/ ~ ((o,
and the ordinary least squares estimator
= arg min ZER
Σ(u: (v: 2 x;,²
1=1
for which we are able to find a close form when X = (x:) invertible.
In our Bayesian setting we aim at finding distribution on z, put prior on it, and approximate the resulting posterior =(2 (D) via variational inference.
Bayesian Neural Network: In the Bayesian neural network (BNN) setting, the linear predictor 2 =Ta; replaced with neural network output f (a: 2) and model
y: ([ =), o²).
In our experiments we use a single hidden layer with hidden units, ReLU activation function and output dimension Thus, the dimension of parameters and thus of the problem is
(daata + 1) + ch + 1).
For T-class classification task, 7. and the BNN output class probabilities T = =1 Iii 2)
[a; =)t-
Once the variational approximation to the posterior is optimized, we can make predictions by Bayesian model averaging:
ply 12) x)) = ( x., 2) xpost dz, or i = f(x 2) T (2) de.
When d is large (e.g. MNIST, where d ~ 10*), sampling or expectation under full d-dimensional mixture becomes too expensive. To address this, we adopt a mean-field-style approximation: we model the posterior product of identical univariate Gaussian-mixture marginals,
N Σ
1
j j= 1.....dd
N
so that all dimensions share the same N-component mixture. This reduces both memory and computational cost while retaining multimodality in each coordinate. We updated m', using the
42

---

presented algorithms. In this setting, we follow classical deep-learnigg framework. We use single- layer neural network with 256 hidden units and ReLU activation. The means of the variational mixture of Gaussians are initialized by sampling from normal distribution, and a Gaussian prior is placed on the model parameters.
Laplace Approximation: In the Bayesian setting, a well-known approach to approximate the posterior distribution is to fit it with a Gaussian. Given the posterior r==[D) defined in Equation 39 Laplace approximation methods consider a Gaussian ,N=",>) where
argmin (z), ((2) = - log (2] = ((D 2) - log (22),
2ER
is the Maximum Posteriori (MAP), and Σ* is the inverse of the hessian of U evaluated at z". Indeed,
performing second-order Taylor expansion of U around x* yields:
U(2) U(()) (z- =²)² VU(=)) +/(2 (12==2)'(( - =1)=((==2)
H ==
which corresponds to the negative log of gaussian with mean =* and covariance Σ* H-1 The MAP estimate can be obtained with standard Euclidean optimization methods (e.g. gradient descent). However, in high-dimensional settings, the Hessian H is typically too large compute invert directly, so it must be approximated. In our experiments, we focus on two practical approximations: the Diagonal approximation and the Kronecker-Factorey (K-FAC) approximation [Ritter et al. 2018
Diagonal approximation:
(2) 027
H 1iag((h......................... ha), 14
18
K-FAC approximation:
Rs
Ht),
H₂ Rd G. An
:= ED [a-- G Ep 9
where O--1 are layer inputs and ge the backpropagated pre-activation gradients; expectations empirical over D and all quantities evaluated at z" arg
arg min, U(2).
E.6.2 Real data experiments
Comparison with Laplace Approximation. We compare our method using Gaussian components with Laplace approximations MNIST. On the training set, all methods achieve similar accuracy and negative log-likelihood (NLL). while on the test set our approach yields much better performance, indicating it generalizes better with multiple modes. To illustrate the effect of using multiple Gaussian components, we perform PCA on the means obtained with MD and IBW for N and N 5, along with the MAP estimate used in the Laplace approximation. As shown in Figure 13 the means for N = 5 do not collapse, highlighting the richer posterior structure captured by multiple components.
and N
43

---

breast
breast cancer
breast cancer
breast cancer
breast cancer
10
18W
1.0
Jog Uheiihood
$1,50
5ID
10
10
NGP HMC
10-
8.
Zo>
ADV]
10
10 nerctions
10T
Ierations
Herraions
Berations
Iterations
wine
wine
wine
Log baellbood
1.0
10'
11.00
1.0%
TD
10
0.6
10
10
107
10
0.6
10
10
10" Berations
Herntion
Berotions
iterations
Steratinss
MNIST
MNIST
2.00
2.00
MNIST
Iug 1h Bhound
1.75
3.75
$3,00
1.25
5000 Herabions
Heralions
5000
1,00
5000 Sterations
Deerstoos
1.00
Recolions
10
boston
10
boston
beston
10"
bosion
Iug- wwwlloll
$1.00
10
ROSSE
107
5000 Iberatioos
5000 sterations
5000 Iterations
10'
5000 Iterations
Iterations
Train
Test
Figure 12: Evolution of the ELBO, accuracy (or RMSE), and log-lik log-likelihood over iterations on the train set (left) and test set (Tight) for the breast_cancer (upper row), wine (middle row), MNIST and boston (bottom row) datasets for N 5.
44

---

PCA
N=1
1.0
400
200
MAP
Acc.
0.8
200
-400
-400
-200
0
200
400
0.6
X;
PCA
400
0
200
400
N =
1.0
-200
400
MD
X1
Acc.
0.8
Lap. KFAC
PCA
Lap. Diag.
400
200
0.6
x
200
400
-400
Test set
400
200
0
200
400
Figure 13: Comparison of Laplace and our methods IBW and MD on MNIST, Left: train and test accuracies for N 1 and N = 5. Right: Principal Component Analysis (PCA) projections on the three principal axis of Gaussian component means obtained with MD, IBW, and Laplace approximations.
45


Reviews:
[{"review_id": "xuL0op0X6f", "reviewer": "Reviewer_M74f", "summary": "The paper studies variational inference when the variational family consists of mixtures of isotropic Gaussians where each component has the same weight. The authors propose optimization algorithms and perform numerical experiments that show comparable performance to existing methods with lower memory cost.", "questions": "For the results shown in Figure 2, Was the target a Gaussian mixture with 10 components and d=20? Why not show the comparison for multiple values of N, not just 15?", "limitations": "The paper would have been stronger with some of the items listed as future works.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n\n- The topic is interesting, mixtures of isotropic Gaussians are complex enough to solve interesting problems, but simple enough to be relatively efficient.\n\n- The definitions and assumptions are stated clearly.\n\n- The authors propose two optimization algorithms and compare their results to exiting works.\n\n\n**Weaknesses**\n\nA weakness of this paper is the lack of theoretical guarantees on the approximation error of their scheme vs a full covariance matrix, which the authors mention as future work. \n\nThe numerical experiments show performance similar to BW, which optimizes full covariance matrices. Except for the gain in memory, I believe the only reported gain in computational cost is the one shown in Figure 2.", "was_revised": true, "final_justification": "I have read the authors' rebuttal and will keep my score.", "rebuttal": "We thank the reviewer for his positive comments on our paper. We address his main concerns below.\n\n**Weakness 1: lack of theoretical guarantees on the approximation error of our schemes vs full covariance matrix**\n\nWith all due respect, we believe there may have been a misunderstanding in the reviewer’s comment regarding “the theoretical guarantees on the approximation error of their scheme vs a full covariance matrix.”\nSpecifically, the term *approximation error* refers to the minimal KL divergence achievable within a given variational family, which is independent of the optimization scheme. In contrast, *optimization error*—or equivalently, convergence guarantees—is indeed dependent on the chosen optimization method. We will clarify this distinction more explicitly in the revised manuscript.\n\n- Regarding establishing **approximation guarantees** for VI within a family of mixtures of Gaussians, this is a fundamentally hard problem. Recent guarantees have been obtained in Huix et al 2024 for MOG with constant weights and identical (shared) isotropic covariance matrix, and we have been extending this result (inheriting the same upper bound in logN/N, where N is the number of components) with different covariance matrices (either **isotropic or full**), see end of Section 5 l262 which refers to our result in Appendix D. Yet, this common bound (valid for both cases - isotropic or full) does not enable to capture the small increase in KL objective we notice in our experiments when comparing isotropic to full covariance matrices; which is why we mention in future work that studying further the approximation error is worth pursuing, in particular establishing whether this rate is tight (by proving a lower bound) is an open research question to the best of our knowledge. \n\n- Regarding establishing **convergence guarantees** for our IBW (isotropic Bures-Wasserstein) and MD scheme, we acknowledge that we lack such guarantees in the general case (N>1 and general target distribution). We would like to emphasize that the original paper Lambert et al 2022, which introduced the method employing full covariance matrices for mixtures of Gaussians (namely BW in our paper), which is the full covariance matrix extension of our scheme IBW, also did not provide theoretical guarantees for N>1. This reflects a broader challenge in the field of VI with mixtures of Gaussians: outside the convex setting—such as a single Gaussian approximating a log-concave target—obtaining rigorous approximation error bounds becomes significantly more difficult since we move into the non-convex regime. This is a well-known limitation across related methods, not unique to our approach,  and remains an open problem in general. \n\nWe provided global convergence guarantees (in continuous time) only for N=1 and when the posterior is strongly log-concave (Prop 2.1). Indeed, the KL is convex in the Wasserstein geometry, and isotropic Gaussians are preserved along Wasserstein geodesics. However, when the posterior is not log-concave and N>1, this structure is lost, and we face a non-convex problem. In such settings, the most one can hope for are descent-type results: that the KL objective decreases along IBW iterations; see for instance Prop 4, Huix et al. 2024 for mixtures of isotropic Gaussians with shared covariance, where the VI objective is shown to be smooth w.r.t. the Wasserstein geometry. For our MD scheme, we would need to establish similar descent results, but this time relying on relative smoothness of the objective with respect to another geometry induced by the KL as a Bregman divergence, see for instance (see [1]). However, establishing such results for both our schemes require a strong technical study that we believe can be deferred to future work. \n\n[1] Aubin-Frankowski, P. C., Korba, A., & Léger, F. (2022). Mirror descent with relative smoothness in measure spaces, with application to sinkhorn and em. Neurips.\n\n**Weakness 2: lack of illustration of the gain in memory for our isotropic variance schemes vs full covariance case**\n\nWe acknowledge that Figure 2 is the sole empirical demonstration of our computational advantage; however, we believe it effectively captures the core trade-off. It illustrates the scaling behavior of both the KL objective and running time across varying dimensions, revealing a clear transition in computational cost **from linear (IBW) to quadratic (BW)** with respect to dimensionality. This scaling behavior is theoretically predictable, as emphasized throughout the paper.\nThe observed difference in computational complexity directly stems from the reduction in parameter space—from O(N(d + d²)) to O(N(d + 1))—when using the isotropic family. As dimensionality increases, the computational benefits of this reduction become more pronounced, while the gap in KL divergence remains small. This is precisely what makes our method attractive for high-dimensional settings.\nIn such regimes, full covariance methods often become computationally infeasible on standard hardware due to their memory and runtime demands—a limitation that is theoretically inevitable given their quadratic scaling. Rather than serving as standalone evidence, Figure 2 validates our theoretical predictions. While we are happy to provide further experiments, we believe this figure already conveys the central insight: our method offers comparable approximation quality at a fraction of the computational cost, enabling scalable inference in high dimensions.\n\n**Question about Figure 2**\n\nWe realised that our formatting was confusing. We apologize for this and will update the organization of the figures for greater clarity. The issue stems from our layout: we discuss Figure 3 first, followed by Figure 2, but the description of Figure 3 appears to the left of Figure 2. This spatial arrangement made it unclear which figure we were describing, and we believe the reviewer was referring to Figure 3, not Figure 2. \n\nIn Figure 3, the goal of the experiment is to demonstrate that our methods (IBW, MD) approximates the target mixture as well as  the full covariance version (BW) of the algorithm. We also wanted to highlight that NGD suffers from mode collapse. \nWhile we could have conducted this analysis for different N values, that was not the goal of these particular experiments, which was to show that our schemes were working in high dimensions for a relatively large N (20). Also, we illustrated the effect of N in Figure 1 (and 6 in the Appendix), so for readability reasons we chose not to show the comparison with multiple N. But we agree with the reviewer that investigating the role of this hyperparameter is interesting, and we will add more plots on the same experiments with different N values in the appendix."}, {"review_id": "5ngZ4qhLD7", "reviewer": "Reviewer_cDhA", "summary": "The authors present an algorithm for optimizing $KL(q\\mid\\pi)$ where $q$ denotes a mixture of $N$ isotropic Gaussians, and $\\pi$ denotes a target distribution. The algorithm uses a gradient descent step for updating the means and either a Bures or entropic mirror descent update for the variances. The proposed approach, in both flavors, is compared to existing methods for VI with mixtures of Gaussians on UCI regression tasks (Bayesian logistic regression) and MNIST (Bayesian neural network). Additionally, 2d and 20d mixtures of Gaussians are used as a target for comparison of the methods.", "questions": "Below are my suggestions that I think would improve the paper:\n\n1. Where are the HMC results in figure 1? They appear in the legend, but I cannot see the line in the figure. I also do not see details on HMC for this experiment in the text. Also the color scheme seems to have additional errors, where is ADVI and what method does the purple line represent? Can the authors please correct me if I'm missing something or comment on the issues in the figure and what it looks like after being fixed?\n\n2. In the \"Bayesian posteriors\" experiments, I think comparisons with standard Bayesian neural network tools would be appropriate, for example a Laplace approximation (KFAC [1], diagonal, low-rank, etc.) or SWAG [2]. As it stands, we see some evidence that the proposed algorithm (both flavors) does well compared to the included baselines but naturally the question is, \"will the standard unimodal approaches do just as well?\". Have the authors done this comparison? I understand the BNNs example may not be the focus, but it has been included so the standard baselines should likely be included too. \n\n3. The authors emphasize the value of having a unique variance for each member of the mixture, as opposed to Huix et al. 2024, but I don't see shared variances as a baseline in the experimental results. Can the authors please correct me if I'm wrong or explain why this baseline is missing and if they plan to include it?\n\n4. Why propose two variants of algorithm 1? I'm sure I'm missing something, but I don't see the benefit of having the algorithm contain two options when they perform similarly on all experiments. \n\nRight now my score is \"reject\" as I think 1-3 suggest the proposed method hasn't been explored enough yet. The writing is quite strong and I enjoy the ideas so I'm happy to increase my score if the above items are addressed. For 2 I think it would be sufficient to compare with a diagonal and a KFAC Laplace approximation since they are good unimodal and expressive unimodal baselines in my mind. \n\n**References**\n\n[1] Ritter, Hippolyt, Aleksandar Botev, and David Barber. \"A scalable laplace approximation for neural networks.\" 6th international conference on learning representations, ICLR 2018-conference track proceedings. Vol. 6. International Conference on Representation Learning, 2018.\n\n[2] Maddox, Wesley J., et al. \"A simple baseline for bayesian uncertainty in deep learning.\" Advances in neural information processing systems 32 (2019).", "limitations": "NA", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- The paper is exceptionally clear and well crafted. \n- The background for the separate variants of algorithm 1 provides good context and motivates\n- I like the framing of the problem and the proposed solution. \n\nWeaknesses:\n- Incomplete numerical results and missing baselines, see 2 and 3 in questions. \n- Incomplete / confusing figures, see 1 in questions. \n- Unclear design choices, see 4 in questions.", "was_revised": true, "final_justification": "I feel I still have remaining concerns which make me less confident in the paper. I will increase to borderline reject to reflect the hard effort by the authors and their willingness to engage in the response period. My remaining issues are:\n- my initial issue with figure 1 was that the baselines seemed mislabeled but now it seems they were simply not included in the first place so I am less confident in the results presented being representative of what people will see if they use this approach. \n- the authors made a comment in the responses about using larger networks than used in their experiments and upon a second look at the paper after this comment I noticed the experiments were on the small side and this comment was not justified in my mind.\n\nI believe these issues are sufficient reason to reject a paper from neurips but I am not going to protest if this paper is accepted.", "rebuttal": "We thank the author for his careful review and address his main concerns below.\n\n**Q1:  legend of Fig 1**\n\nThe reviewer is right, the legend of Fig 1 contains a typo  - ADVI and HMC were baselines included in a later experiment (Figure 4 and 10 on breast_cancer, wine and boston (RMSE) datasets) - we corrected it in the revised version. \n\n**Q2: additional unimodal (N=1) baselines for BNN**\n\nWe acknowledge that in the case of BNN, we did not compare with the monoGaussian case (N=1, where N is the number of components in the mixture of Gaussians -MOG- variational posterior). We originally thought that our previous experiments for simple MOG in 2d were already highlighting the critical failure of using one mode for multimodal targets. Note that BNN posteriors are likely to be multimodal (see Figure 3 in [1] for instance). Yet, we acknowledge that including the monoGaussian case as a naive baseline for BNN would have been interesting. We have run these experiments and provide our results below. \n\nWe observe that a MOG with N>1 performs better in terms of accuracy and log-likelihood than a mono-Gaussian. Note that we only report the results on one dataset of our algorithms and the requested baselines for readability; we will add complete convergence curves and benchmarks in the revised paper.\n\nOn MNIST: Accuracy (higher the better) and negative log-likelihood (lower the better) after convergence on the test set:\n-  IBW\n    - N=1: Acc. 0.872, Neg. LL. 1.594\n    - N=5: Acc. 0.962, Neg. LL. 1.508\n- MD\n    - N=1: Acc. 0.871, Neg. LL. 1.596\n    - N=5: Acc. 0.961, Neg. LL 1.509\n- Laplace Approximation\n    - Diag: Acc. 0.929, Neg. LL . 1.779\n    - KFAC: Acc. 0.948, Neg. LL 1.683\n\nThese results show that there is a clear improvement from going to the unimodal (N=1) to multimodal (N>1) VI posterior. We believe this may be due to the fact that multimodal variational posterior enable to explore other basins/modes of good weights for the neural network (see upcoming paragraph) which enhance good generalization, as illustrated in Fig 3 in [2].\nWe thank the reviewer for the suggestion and will include these experiments in the paper. \n\nWe have also investigated numerically whether our algorithms find different modes.  For this analysis, we used two metrics to assess the diversity of the optimized mixture components. We applied PCA to all collected means from different methods (IBW, MD, NGD, and MAP) and visualized them in a scatter plot (which will be included in the revised paper). Here, MAP refers to the Maximum A Posteriori estimate, which corresponds to the mean of the variational Gaussian approximation used with the Laplace method. Note that the MAP estimate is identical for both diagonal and KFAC variants of the Laplace Approximation, as these methods differ only in their covariance matrix approximations. Since we are not able to include the plot here, we provide the average Euclidean distance between distinct mixture component means in the PCA-projected space. We also provide average cosine similarity between component vectors in the original high-dimensional parameter space (excluding self-similarity). We find:\n- IBW with N = 5: pairwise Euclidean average 455, cosine-similarity: 0.034\n- MD with N = 5: pairwise Euclidean average 455, cosine-similarity: 0.034\n\nThe identical metrics between IBW and MD reflect the fact that these methods converge to very similar solutions (see answer to Q4). The low cosine similarity (0.034) indicates that components are nearly orthogonal in the high-dimensional parameter space, confirming the effectiveness of our multimodal approach for capturing complex posterior distributions in BNNs. We will include this experiment. \nOverall, these results demonstrate that mixture components maintain substantial diversity and do not collapse to identical solutions, and enable to get better results than a unimodal variational posterior. \n\n**Q3: about shared-variance (within MOG components) baseline**\n\nHuix et al. 2024 proposed an algorithm that updates a uniform sum of Diracs convolved with a Gaussian kernel of variance epsilon . In their algorithm, the **epsilon parameter is fixed**, and they update only the Dirac positions. Typically, we found numerically at the very beginning of our study that only optimizing the means could lead to situations where the algorithm was trapped in local minimas, and mean particles would stay stuck far from the target. This was a big motivation to optimize the variances as we saw it fixed that problem. Yet, indeed you are right that it is possible to consider as a baseline a version of our algorithms, IBW and MD, where this epsilon parameter would be identical between components (while our schemes, in their full generality, allow different isotropic variances). There are several reasons why we did not present the case of optimizing this shared epsilon:\n- Low computational cost: since epsilon is a 1D parameter, updating it scales with N rather than with d, so updating it is not computationally costly. As an example, for BNNs with millions of parameters, it represents only 5 parameters when N = 5, which is negligible.\n- Readability reason: at the beginning of our numerical study, we conducted experiments with this approach (shared epsilon within the MOG) and found it performed better than not updating epsilon (i.e. the scheme of Huix et al 2024), but worse than having a different epsilon for each Gaussian component in the mixture. Since we needed to make choices about benchmarking—we are comparing with Huix et al.'s method, the full covariance approach of Lambert et al., the NGD method from Lin et al., and two versions of our own method (plus a Normalizing Flow as a general non-specific benchmark)—we chose the clearer, more readable option. Yet, we acknowledge that it could be another useful baseline. We have run such baseline and obtained more results that we present below. We report the percentage of improvement achieved by using different epsilon values for each component compared to sharing a single epsilon across the variational mixture. The percentages below are averaged over 5 experiments.\n    - Figure 1 Toy 2D MOG (IBW / MD):\n        - N=5: 28.9% / 42.3%\n        - N=10: 64.3% / 83.3%\n        - N=20: 66.7% / 100.0%\n    - Figure 6 (a) Toy 2D Funnel (IBW / MD):\n        - N=5: 34.6% / 35.2%\n        - N=20: 48.0% / 62.3%\n        - N=40: 50.0% / 56.0%\n    - Toy Funnel 10D (IBW / MD)\n        - N=5: 27.8% / 18.7%\n        - N=20: 28.5% / 21.9%\n        - N=40: 31.6% / 28.2%\nWe will add examples of optimized Gaussians in 2D in the revised paper to demonstrate the impact and importance of having a different variance for each Gaussian component. Indeed for 2D toy examples we are able to visualise optimized Gaussians (represented by a circle) of very different size, and we see it is really beneficial for distribution as the Funnel’s one or skewed one (a, b-1, b-2 in Fig 6). \n\n**Q4: about why we proposed and included two different algorithms, IBW vs MD**\n\nWe acknowledge that the two algorithms obtain similar final performances, and that our study does not designate (yet?) a clear winner. We included both schemes because both are novel - (while one is an extended version of Lambert et al 2022) but they significantly differ from the competitors (eg NGD). Note also that the two schemes have the same computational cost O(N(d+1)). Both algorithms minimize the same objective but with respect to different geometries. In standard optimization, it is well known that choosing an appropriate geometry can enhance the speed of convergence (for instance, if the objective is smooth and convex in the chosen geometry, see [3]). This is also known on the space of measures ([4], [5]). In our case, it is not clear in advance whether the VI objective is more suited to the Wasserstein or entropic mirror descent geometries. Such investigation is among what we mentioned in the future work in the conclusion, but we believe this requires  a significant theoretical study. Note also that this VI objective (and its smoothness, convexity) is dependent both on the VI family (mog with isotropic Gaussians) but also the log-posterior target (which in practice in bayesian inference applications, depends on the loss, the data, and the model). Recall that for instance, KL(.|pi) is convex in the Wasserstein-2 geometry when pi is log concave on the family of mono Gaussians. Hence, our goal was to investigate empirically the performance of these two schemes for a wide range of posteriors and hyperparameters within the VI family (e.g. number of components). Note that while the algorithms obtain the same final performance (i.e. they find approximately the same solution), their trajectories are different (because they use  different geometries). This, we have noticed in our experiments early in our numerical study, by tracking the evolutions of the variances for each scheme, and these trajectories were generally different. We will include such plots illustrating the (different) behavior of our schemes in the revised version for more clarity. \n\n We hope our responses have addressed the reviewer’s concerns and that they will consider revising their score.\n\n[1] Izmailov, P., Vikram, S., Hoffman, M. D., & Wilson, A. G. G. (2021). What are Bayesian neural network posteriors really like?. ICML.\n\n[2] Wilson, A. G., & Izmailov, P. (2020). Bayesian deep learning and a probabilistic perspective of generalization. Neurips.\n \n[3] Lu, H., Freund, R. M., & Nesterov, Y. (2018). Relatively smooth convex optimization by first-order methods, and applications. SIAM Journal on Optimization.\n\n[4] Aubin-Frankowski, P. C., Korba, A., & Léger, F. (2022). Mirror descent with relative smoothness in measure spaces, with application to sinkhorn and EM. Neurips.\n\n[5] C Bonet, T Uscidda, A David, PC Aubin-Frankowski, A Korba. (2024) Mirror and preconditioned gradient descent in wasserstein space. Neurips"}, {"review_id": "lk7B43x3qS", "reviewer": "Reviewer_pR5Q", "summary": "The paper presents a memory-linear variational-inference scheme that uses an equal-weight mixture of isotropic Gaussians. Two geometry-aware O(d) variance updates—Isotropic Bures-Wasserstein and Entropic Mirror-Descent—combined with gradient steps on the means yield a scalable alternative to full-covariance mixtures, normalising flows, and HMC.", "questions": "1. Can you provide guidance or an adaptive strategy for choosing the number of components N?\n\n2. On a synthetic with a 90/10 probability split between two modes, how large is the bias induced by equal weights? Please quantify.\n\n3. More broadly, how does your method mitigate the bias caused by the equal-weight constraint?\n\n4. For the two update methods such as IBW and MD updates, how practitioners choose between them in practice?", "limitations": "Authors mention the equal-weight constraint and the memory-accuracy trade-off as limitations. They could further discuss how memory scales if N must grow with dimension and potential biases on highly skewed posteriors.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Quality: The submission is technically sound and claims are well supported by theoretical analysis and experimental results.\n\nClarity: The submission is clearly written.\n\nSignificance: This technique has the potential to improve the SOTA of variational inference.\n\nOrginalty: The new formuation of using an equal-weight mixture of isotropic Gaussians is a balance between computation efficiency and representation power.\n\nStrenghs:\n\n1.\tThe algorithm has a transparent O(d) complexity per component.\n\n2.\tThe writing is clear and the theoretical derivations are both rigorous and easy to follow.\n\n3.\tExperiments are comprehensive, extending up to a large-scale BNN.\n\nWeakness: \n\n1.\tNo guidance is given for choosing the number of components N. A sensitivity study is also missing.\n\n2.\tThe equal-weight assumption may introduce significant bias for skewed posteriors, yet no benchmark quantifies this effect.", "was_revised": true, "final_justification": "This is a decent and well-written paper. \nI have no concern of the paper.", "rebuttal": "We thank the reviewer for his careful review and positive comments on our paper. We answer his questions below.\n\n**Q1: guidance or adaptive strategy for choosing N (number of components in the variational MOG approximation)**\n\nIn general, we found that the larger the N, the better the result. For instance in Figure 1, while the target has 5 components, a larger N (e.g. 10, 20) enables to get a lower error than N=5. Intuitively, even with a large N, optimizing the variances enables to capture quite general shapes (see also Appendix E for more experiments, e.g. E.3 for 2D visual examples). \nThis observation is in accordance with our theoretical results (see end of Section 5 and Appendix D) that states that as N goes to infinity, the approximation error within a MOG family with N components goes to zero (assuming the algorithm achieves this error, i.e. the latter finds the minimizer within the family of N-MOG, this would explain why “the larger N the better”). This reflects the fact that for a given number of components, this family of MOG is rich enough (in particular shaping the variances is crucial) to fit a wide variety of targets.  \nIn practical applications, this trade-off must be carefully balanced by the user.\nMoreover, the number of modes in the target distribution is typically unknown in advance. Yet, for BNN for instance, we noticed that even with a low number of modes (N=5), the variational approximation of such complex posteriors already leads to very satisfying accuracy results.  Yet (see answer to Q2 of Reviewer cDhA for more results on BNN and the effect of choosing different number of modes N), for these typically multimodal posterior, we saw that increasing N enables to better cover the target density and we are less likely to miss a mode. \n\n**Q 2-3 : about the bias when using constant weights in MOG, and its mitigation**\n\nWe conducted experiments on a target mixture of Gaussians in two dimensions with two components and weights [0.9, 0.1] as requested by the reviewer. The means were sampled from a ball of radius 5, and covariances were scaled by a factor of 2. We benchmarked against the algorithms presented in the paper and the baselines used throughout (BW, IBW, MD, NGD) for N = 2, 3, 5, 10 components. As expected with N = 2, since the weights are not optimizable, we observe significant bias induced by the uniform weights constraint. The approximation either misses the low-density mode or assigns equal probability to both modes. With N = 3 already, the bias begins to diminish as the method allocates 2 Gaussians near the heavy mode and 1 near the light mode (with N=5, respectively 4 and 1). When N = 10, we achieve probabilities close to the target distribution, with approximately 9 Gaussians clustering around the heavy mode and 1 around the light mode. \nMethodology: We drew B = 10,000 samples from both the variational distribution and the target distribution, then applied K-means clustering with n_clusters = 2. We analyzed the cluster assignments and proportions. The reported errors are: (1) the average Euclidean distance between the target modes and the variational modes, and (2) the average absolute difference between the target cluster weights and variational cluster weights. The first error captures whether the variational approximation has found the right modes, while the second one measures whether it has found them in the right proportions. \nFor readability in our table we round coordinates to one decimal place and probabilities to two decimal places. Our numerical results are the following.\n\nTarget: 0.10 at (-4.6, 0.6); 0.90 at (4.7, -2.2)  \n- IBW \n    - N = 2: 0.50 at (3.6, -2.2); 0.50 at (5.8, -2.2), Errors: 4.8, 0.4\n    - N = 3: 0.33 at (-4.6, 0.5); 0.67 at (4.7, -2.2), errors: 0.02, 0.23\n    - N = 5: 0.20 at (-4.6,0.5); 0.80 at (4.7, -2.2), errors: 0.05, 0.10\n    - N = 10: 0.10 at (-4.6, 0.6); 0.90 at (4.7, -2.2), errors: 0.02, 0.00\n- MD \n    - N = 2: 0.50 at (3.6, -2.2); 0.50 at (5.8, -2.2), errors: 4.9, 0.4\n    - N = 3: 0.33 at (-4.6, 0.5); 0.67 at (4.7, -2.2), errors: 0.33, 0.30\n    - N = 5: 0.20 at (-4.6,0.5); 0.81 at (4.7, -2.2), errors: 0.05, 0.10\n    - N = 10: 0.10 at (-4.6, 0.6); 0.90 at (4.7, -2.2), errors: 0.05, 0.01\nThe detailed table will be included in the revised paper along with additional benchmarks and visualizations. Note that a similar experiment was already conducted in Appendix E, Figure 6, c-3 with 4 modes and different weights, but without quantifying the bias induced by the equal weights restriction.\nTo conclude, indeed for a fixed N there might a bias induced by the fixed weights, but the latter bias is easily mitigated by increasing N.  \n\n**Q4: choice of IBW vs MD**\n\nSee answer to Q4 of Reviewer cDha."}, {"review_id": "jEBPXum7wt", "reviewer": "Reviewer_Fvva", "summary": "The authors present a general variational framework based on mixtures of isotropic Gaussians. Despite being less flexible than the counterparts with full covariance matrices, the proposal is also capable of approximating multimodal posteriors, while being much more efficient in terms of computational and memory costs. In this context, the paper derives two variational inference algorithms, which are evaluated in common approximate inference tasks. The experimental comparison with other strategies from the literature show promising results.", "questions": "- At which scenarios the introduced inference algorithms (IBW and MD) are more indicated? Guidelines for choosing between the two schemes would be a welcomed inclusion to the work.\n\n- Line 252: The decoupling of the updates on the means and variances is an advantage or a disadvantage here?\n\n- Some references are incomplete or show the arXiv versions, instead of published versions, e.g., lines 408, 437, 442, 447.", "limitations": "The authors state the limitations of their work.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "The paper is well organized and written. Although it demands a substantial background from the reader, the authors provide useful references across all the sections. The related work is up to date and highlight the introduced contributions. The experimental section covers useful tasks (low-mid dimension inference and approximate posterior for standard machine learning models). I also praise the authors for the provided code and comprehensive supplementary material.\n\nI did not find great weaknesses in the paper, but I did miss a comparison with more recent approximate inference methods based on deep generative models. The authors do use a simple Normalizing Flow as a baseline, but strategies based on GFlowNets [1,2], Diffusion models [3,4], or Flow Matching [5] could also be considered or at least included in the discussion.\n\nReferences\n\n[1] Nikolay Malkin, Salem Lahlou, Tristan Deleu, Xu Ji, Edward J. Hu, Katie Everett, Dinghuai Zhang, Yoshua Bengio: GFlowNets and variational inference. ICLR 2023\n\n[2] Salem Lahlou, Tristan Deleu, Pablo Lemos, Dinghuai Zhang, Alexandra Volokhova, Alex Hernández-García, Léna Néhale Ezzine, Yoshua Bengio, Nikolay Malkin: A theory of continuous generative flow networks. ICML 2023: 18269-18300\n\n[3] Francisco Vargas, Will Sussman Grathwohl, Arnaud Doucet: Denoising Diffusion Samplers. ICLR 2023\n\n[4] Wasu Top Piriyakulkij, Yingheng Wang, Volodymyr Kuleshov: Denoising Diffusion Variational Inference: Diffusion Models as Expressive Variational Posteriors. AAAI 2025: 19921-19930\n\n[5] Jonas Wildberger, Maximilian Dax, Simon Buchholz, Stephen R. Green, Jakob H. Macke, Bernhard Schölkopf: Flow Matching for Scalable Simulation-Based Inference. NeurIPS 2023", "was_revised": true, "final_justification": "The authors addressed my points and I maintain my favorable score.", "rebuttal": "We thank the reviewer for his very positive comments about our paper. His questions and pointed weaknesses are commented below. \n\n**Weakness - about additional neural-network based baselines**\n\nWe thank the reviewer for the references, most of them we were not familiar with. We originally included baselines that were relatively standard in recent VI papers, and our point when including NF was to highlight that the use of Neural Networks induced a significant computational overhead compared to our scheme, while yielding unsatisfying results. Yet we acknowledge that the alternative neural-network based approaches you mention could perform very differently and would be interesting to compare.  We will mention them in the revised version in the discussion. \n\n**Q1: choice of IBW vs MD**\n\nSee answer to Q4 of Reviewer cDha. \n\n**Q2: about coupled updates between mean and variances in NGD**\n\nWe found that this coupled update was a disadvantage, at least numerically. Indeed, the update of the means (the gradient size) is multiplied by the current variance (see Equation 17). Hence, large variances lead to large moves for the mean, and we found this could even lead to mode missing. For instance in the experiment in Figure 3, we needed to initialize our variational posterior with rather large variances (100) to enforce space exploration, and even with reasonable learning rates (0.01) we noticed that behavior. More generally, in our experiments we found that this scheme was more sensitive to the choice of the step size and hence more numerically unstable. \n\n**Q3 : error in the format of references**\n\nThank you for spotting this - we cleaned it in our revised version."}]

Author rebuttal:

Meta-review: This paper addresses a restricted variational family, i.e., mixture of isotropic Gaussians with uniformly fixed weights and demonstrates its validity and usefulness in the context of expressivity vs memory-efficiency. Two geometry-aware variance update rules, which are decoupled from the mean update, are presented in the context of mixtures of isotropic Gaussians. Most of reviewers feel that the paper is well written and has interesting contributions. However, there are some concerns in approximation/convergence guarantees and empirical comparisons. During the rebuttal period, some concerns were clarified but the concern on experiments with large scale models was not fully resolved. Nevertheless, thanks to extensive communications with reviewers, we have reached consensus on this paper.
Decision: Accept (poster)