Optimization Guarantees for Square-Root Natural-Gradient Variational Inference
Authors: Navish Kumar
Thomas Möllenhoff
Mohammad Emtiyaz Khan
Aurelien Lucchi
Venue: Accepted by TMLR ()

Abstract: Variational inference with natural-gradient descent often shows fast convergence in practice, but its theoretical convergence guarantees have been challenging to establish. This is true even for the simplest cases that involve concave log-likelihoods and use a Gaussian approximation. We show that the challenge can be circumvented for such cases using a square-root parameterization for the Gaussian covariance. This approach establishes novel convergence guarantees for natural-gradient variational-Gaussian inference and its continuous-time gradient flow. Our experiments demonstrate the effectiveness of natural gradient methods and highlight their advantages over algorithms that use Euclidean or Wasserstein geometries.

Published in Transactions on Machine Learning Research (05/2025)
Optimization Guarantees for Square-Root Natural-Gradient Variational Inference
Navish Kumar
nawish.kumarSuumbbe.u.
University of Basel, Basel, Switzerland
Deportment of Mathematics and Computer Science
Thomas Möllenhoff
RIKEN Center for AI Project, Tokyo, Japan
Mohammad Emtiyaz Khan
emtiyaz.khon Driken..
RIKEN Center for Al Project, Tokyo, Japan
Aurelien Lucchi
aurelien, aurelien. www.euilccchiiuiiass....
University of Basel, Basel, Switzerland Department of Mathematics and Computer Science
Reviewed on OpenReview: https: // openreview. net/ for//f id- 0ue77
Abstract
Variational inference with natural-gradient descent often shows fast convergence in practice, but its theoretical convergence guarantees have been challenging to establish. This is true even for the simplest cases that involve concave log-likelihoods and use a Gaussian approximation. We show that the challenge can be circumvented for such cases using a square-root parameterization for the Gaussian covariance. This approach establishes novel convergence guarantees for natural-gradient variational-Gaussiann inference and its continuousttine gradient flow. Our experiments demonstrate the effectiveness of natural gradient methods and highlight their advantages over algorithms that use Euclidean or Wasserstein geometries.
1 Introduction
Variational inference (VT) is widely used many areas of machine learning and can provide a approxima-
tion the posterior distribution (Jaakkola Jordan 1997 Wainwright al.) 2008 Graves 2011
Kingma &
Welling 2013). VI works by formulating Bayesian inference
optimization problem
restricted class
of distributions which is then solved, for example, by using gradient descent (GD) Grraes 2011
Ranganath
et al. 2014 Blundell et al.) 2015). A faster alternative is to use natural-gradient descent (NGD) Ammri
1998 which exploits the information geometry of the posterior approximation and often converges much
faster than GD, for example, see Figure []), Khan & Nielsen (2018 Fig. 1b), or Lin et al. 2019 Fig. 3)
The updates can often implemented using message passing (Knowles Minka 2011 Hoffman
2013
Khan & Lin 2017) and sometimes even as Newton's method or dcep-learning optimizers (Khan et al.
2017
2018
Salimans & Knowles
2013). Due to this, there has been a lot of interest. in NGD based VI (Honkela
et al.
2008 Hensman et al. 2012 Salimbeni et al. 2018 Osawa et al.) 2019 Adam et al. 2021 Khan & Rue
2023
Despite their fast convergence, little has been done to understand the theoretical convergence properties of NGD, and its continuous counterpart, which we call natural- -radieent (NG) flow. Instead, most works focus on GD (Alquier et al., 2016 Domke 2019 2020 Domke al.] 2024 Kim et al., (2024). For NGD, Khan et al. (2016) derived convergence results for general stochastic, non-convex settings. Still, their convergence rates slow and there is plenty of room for improvement, especially for concave log-likelihoods. For such

---

Published in Transactions on Machine Learning Research (05/2025)
likelihoods, recently, Chen et al. 2023 studied the contimuuus-tmme flow. There also some studies on stochastic NG variational inference; for example, Wu & Gardner (2024) prove non-asymplotic convergence rates for conjugate likelihoods. However, the conjugacy assumption limits their applicability to more general settings. Both in discrete and continuous time, however, there is a gap between theory and practice for convergence guarantees of NGD.
A major difficulty in obtaining stronger convergence guarantees for NGD and NG flow lies in a subtle technical issue: commonly used parameterizations can often destroy the concavity property of the log-likelihood (Céélief-bbdellatiff et al.
2019 Sec. 4.3). It is well-known that the expectation of concave log-likeli for Gaussian approxination is not concave with respect to either the covariance or precision matrix, but only with respect to the Cholesky factor of the covariance (Challis & Barber 2013). Therefore, commonly used parameterizations, such as natural or expectation parameterization of the exponential family, destroy the concavity property and prohibit the application of tools from convex optimization to analyze the properties of NGD. Our goal in this paper to address this issue.
We circumvent this problem by considering variational Gaussian inference that is defined by the square root of the Gaussian covariance matrix. We show that discretizing the NG flow with this square-root parameterization results in an NGD update based on the Cholesky factor. This idea was proposed by Tan (2025). though it was also previously explored in a broader context by Lin et al. (2023 2024). Our method extends this idea to include any square-root parameterization. This way of parameterizing helps maintain the concavity of log- liellioe functions, which in turn facilitates convergence guarantees for both NG flow and NGD. The key contributions of our work are outlined as follows:
We prove that the NG flow exhibits an exponential convergence rate for a strongly concave log- likelihood, as stated in Theorem | When discretized using the squarc-root paramcterization, this aligns with the update provided by Tan 20225 Theorem 1). The proof of convergence hinges on establishing a Riemannian Polyak-F.ojasewiicz (PT.) inequality, which is shown to be valid when the lowest eigenvalue of the metric tensor bounded.
We prove an exponential rate of convergence (Theorem 2 for NGD using Cholesky and under strongly-concave log-likc;;hood; the proof extends to any square-root parameterization.
We present empirical results showcasing the fast convergence of NGD, attributed to its Newton-like update. These results are illustrated in Figure and Figure 3 is important highlight that our experiments employ the piecewise bounds described in MMarlin et al. 2011) for computing expectations. This raises the question of how the resulting bias affects NGD convergence, To address this, we include discrete-time time analysis with bias in our supplementary section, see Theorem 3 in Appendix D The reason for our preference for the piecewise over the stochastic implementation is justified in Appendix
Our proof addresses the additional challenges posed by the non-smoothness of entropy and does not require any proximal projection operators such as those used by Domkel (201)); Domke et al. (002)). Our results can be useful for further research on proving such results for generic NGD updates. The use of square-roots could lead to suboptimal convergence, for instance, it may not converge in one step for quadratic problems (as shown in Figure []), but this only implies that even better results are possible for the general case. We provide an insight into why this happens in Section 3.5 noting that square-root-basee updates essentially compute an approximate estimate of the inverse of the preconditioning matrix that is used in Newton's method.
2 Variational Inference
The goal of variational inference is to approximate the posterior distribution by a simpler distribution, such as the Gaussian distribution. Given a Bayesian model with n likelihoods, denoted by p(D,(0) for the i'th data example with parameter 0, and a prior p(0) over the parameters, Bayesian inference aims to find the posterior p(0) Di, D2..... . D.). In variational inference, we aim to find approximation gr (0) to the posterior by
2

---

Published in Transactions on Machine Learning Research (05/2025)
Natural-Gradient
VN SR-VN
VN
Euclidean
SR-VN
Wasserstein
(a)
(b)
(c)
Figure 1: (a) Natural-gradient descent (circle) converges much faster (in just step) than algorithms that use Euclidean (cross) or Wasserstein (star) geometries which take 10 and 470 iterations respectively. The illustration is a 2-D Bayesian linear regression (quadratic loss in the background for all three figures). The 3 methods are taken from Khan & Rue (2023); Ranganath al. (2014); Lambert et al. (2022b) (b) For smaller step-sizes (of the order = 10 or less), SR-VN and VN perform almost identically, (c) For larger step-sizes they start to differ, step-size for VN here is 1.
solving the following problem over the space E Ω of the valid varialional porameters,
min rest
+ KL le- (0)l
(1)
where denote ((0) = - log.-oogpi PDD 0). [ to be the KL divergence, and y 0 is scalar parameter commonly referred to the lemperature. When 7 1, the objective is referred to the negalive evidence lower bound (ELBO) (Jaakkola & Jordan 1997 Blei et al.) 2017), which is also the focus of our analysis. The value of plays a crucial role in shaping the postcrior variance. If y is set too high, the KL term dominates and forces the variational posterior (0) to collapse to the prior p(0), eliminating any learned posterior structure. Conversely, if y is set too low, the KT penalty becomes negligible, which can lead to high-variance gradient estimates and unstable training (Shen al. 2024b).
In many problems, the negative log- likeliool (NLL) ((0) is convex; it could also be arbitrary loss function, giving rise to generalized posterior (Zhang) 1999 Catoni 2007). The regularizer, denoted by R(0) = log p( (0), too is chosen to be convex. Popular examples include Bayesian linear and logistic regressions, as well as Gaussian process models. In such cases, a more convenient form is used to separately define two terms as the expected log-loss and negative-entrops respectively, that is f(T) = Eg-(o) yR(0)), H(T) = E (0) [log (- Using these, we can rewrite Equation (1) by expanding the KL divergence as
min ((T) where L(T) == f(T) + *H(T), TEs
(2)
and where the first term involves the expectation of convex loss function while the second term is the entropy. This reformulation allows us to connect VT to wide variety of problems in supervised, unsupervised, and reinforcement learning which can be formulated as a minimization of the form
min ((0) where {()) = (()) +R(0).
(3)
DRRP
VI approaches can be seen as a variant of such problems where instead of finding a single minimizer 0", we seek parametric distribution (-- (0) by solving Equation 2 This closely related to techniques robust or global optimization that use (0) to smooth the loss and mitigate the influence of local minima (Mobahi & Fisher III 2015 Leordeanu & Hebert 2008 Hazan et al.) 2016), often with taking value
of O. This approach is also referred to as Variational Optimization (VO) (Staines & Barber 2012). VO can Icad us to the minimum 0' because, for y 0, ((T) is an upper bound on the minimum value of f, that is, ming E(0)= < minren L(T). Therefore minimizing L(T) minimizes ε(0), and when the distribution q puts all its mass on 0", we recover the minimum value. Problem 2 with ? = 0 is also encountered across various other domains, including random scarch (Babal 1981), stochastic optimization (Spall 2005), and evolutionary strategies (Beyer (001). In this context, 9, (0) serves as the 'search' distribution employed locate the global minimum of black-box function, {(0). Notably, in the realm of reinforcement learning, this
3

---

Published in Transactions on Machine Learning Research (05/2025)
distribution can take on the form of policy distribution designed to minimize the expected value-function {(0) (Sutton et al. 1998). sometimes incorporating entropy regularization techniques (Williams & Peng 1991 Teboulle 1992; Mnih et al. 2016).
3
Background on NGD for VI
We will review VI using GD and NGD methods and discuss the NG flow and NGD with the Cholesky factor, and for both, we will provide convergence guaranlees in the later sections.
3.1 Variational Gaussian Inference with GD
Throughout, we will focus on VI with a Gaussian distribution 9-(0) := .N(0 m, V) where m E Rd is the mean, V E Rdxd the covariance. The variational parameter T (m, V) is only defined within the set of valid parameters Q the set of where V is positive definite, that is S2 {m, V) m E Rd V > E [ddx]. J. 2 can then be rewritten as, m C(m, V). The function L is differentiable under mild The problem
conditions on the interior of Q' even when & is not differentiable, This makes it possible to apply gradient-based optimization methods to optimize it. A straightforward approach to minimize L is to use Gradient Descent (GD) as shown below,
discussed by Staines & Barber (2012)
GD: V4+1 V, pVyL(m, Vi), M4+1 m, pV m V.),
(4)
where p > 0 is a step size. This approach is often referred to as the black- VI (BBVI) in the literature Ranganath al.) (2044), and offers simple and convenient updates implement using modern automatic- differentiatton methods and the reparameterization trick (Graves) 2011 Blundell et al.) 2015 Titsias & Unfortumately, such direct gradient-descent methods can be slow in practice.
2014).
3.2
Variational Gaussian Inference with NGD
An alternative approach is natural-gradient descent which exploits the information geometry (Amari 2016) of q to speed up convergence. Given the Fisher information matrix (FIM) (Amari 2008) of 4r (0), that is F = (o) log o (O) 4 (.)) which is positive-defiitte for all 7 E Q in our setting, the NGD for VT in the natural-parameter space (with t as natural parametrization) is given as follows (Khan & Rue
2023
NGD: Tet = - V,
55)
The preconditioning of the gradient by the FIM leads proper scaling in each dimension which often leads faster convergence. A notable connection between NGD for / and Newton's method for 3 was established by Khan al. 2017 2018 in the case of Gaussian distribution with mean and the precision S, which is the inverse of the covariance S V-1. Denoting the expected gradient and the expected Hessian, respectively, as:
= Bo.N
H, = Eo v.) V. (V)(())
(6)
The resulting updates,
VN: S = + pH,, m = m. 1(S -gg,
(7)
referred to as the variational Newton (VN) update, optimize instead of 3 The standard Newton's update
for problem X is recovered by approximating the expectations at the mean m and using step-size p = when y = 1; see Khan & Rue (2023). Since the precision lies in a positive-definite matrix space, the update 7 may violate this positivity constraint of the parameter space ? (Khan & Nielsen 2018 For example, this happens when the loss E(0) is non-convex.
3.3 Natural Gradient Flow
The NG flow minimizes the KL divergence functional on a parameterized manifold of Gaussian densities the objective stated in Equation 1 That is, NG flow is continuous time update of the NGD (Equation 5
4

---

Published in Transactions on Machine Learning Research (05/2025)
Algorithm Square-Root Variational Newton (SR-VN)
Parameter: Step-size p 0 is chosen using Equation tril is defined Equation 10
Initialize (mo: Co) satisfy Assumption ! for t = ...... do Update 8: and H, using Equation 6
5: C++1 <- C. pC, tri H,C. y
6: m4+1 f m. C, gr
7: end for
for y= that is,
dT: - ---,'VV((T) V+[(T,)
-F KL (+(()) p(0 D)
(8)
dt
The updates for the NG flow in T = (m, V) parameterization were provided in Chen et al. 2023 Equation (4.18)) as:
dV. VVVVV di - dt
dm:
(9)
Upon discretization with respect to the natural parameters, this results in the VN update with y = 1. Refer to A Appendix|] B for a comparison the updates, as obtained from VN. However, due to the non-convex nature of KL even with concave log- ligeeiklihodd for (m, V) parameterization (see further discussion in Section 3.6 instead focus
analyzing the NG flow in the squarc-root parameterization, that is 7 (m, C) E 22 where Ω = {(m, C) É Rd, C > 0 E Rdxdj under the assumption that C is square and lower- triangular with positive real diagonal entries, ensuring its invertibility. Let us define function that takes the swerrtriangular part of a matrix and halves its diagonal, that is
Aij
i>
tril [ri
=
2 A
= J,
(10)
0
<.j.
Then, derive flow equations for the Cholesky factor, we first note from Murray 20166 Equation (7)) that the derivative of the Cholesky factor C, can be obtained from the derivative of the covariance V, CCC
dCr. = C. Cri
dV
cV:C,TI cT
(11)
dt
dt
which after using Equation
9 leads to the following flow equations:
C dC. = C, - H.C.), H.C.,, dt -CC. g:-
dm,
(12)
dt
In Theorem we prove the convergence of the flow dynamics in Equation 12 under the concave log- likelihood. Importantly, the NG flow remains the same regardless of the parameterization, and therefore the flow dynamics in Equation (9) and Equation (1) trace out the same trajectory on the Gaussian manifold.
Hence, we focus on square-root parametrization solely for its theoretical convenience in proving convergence.
3.4 Square-Rood Root Variational Newton (SR-VN)
In this section, we introduce a different NGD update called square-root variational Newton (SR-VN), outlined Algorithm [ which utilizes the Cholesky factor C of the covariance matrix V. SR-VN is derived through forward Euler discretization of the NG flow defined in square-root parameterization (Equation (12). The resultant update aligns with independently proposed updates by Tan (2025), who defined VN-like update using the Cholesky factor, given
C++1 C, pC, Cr VoL(m,, C.)], m m. PC,.O Vmm(mm, C.).

---

Published in Transactions on Machine Learning Research (05/2025)
Now, calculated in Equation 34 in Appendix A.2
we have
Vel(m,, Ct) = (H, (C,,~<(C,C))'))CC
(13)
Finally, by seeing that V VmC(m,, Ct) = Bt, we can rewrite the updates from Tan 2025
Theorem 1) as
SR-VN: C++1 C. pC, trillC, H,C, -11, M m. ,C,,C,
- g:
(14)
allowing us to immediately see that these updates indeed come from direct forward Euler discretization of
Equation 12). In Theorem 2 we establish convergence guarantees for SR-VN. We also note here that other
updates based on the Cholesky factor have been considered in the literature, for example, see (Sun et al. 2009
Salimbeni et al. 2018 Glasmachers et al. 2010
3.5 SR-VN as an Approximation of VN
It is interesting to see that SR-VN performs a Newton-like update in m (Equation a matrix inversion the preconditioner (such as the inversion of the Hessian in VN, Equation (7)). This makes it casier to implement and analyze. However, we show in Appendix B that, reality, SR-VN computes an approximation of the inverse and is therefore only suboptimal compared to VN. Looking at Figure D bb, can see that when the step-size is very small, SR-VE VN and VN converge to the same dynamics because they both approximate the NG flow. However, the quadratic offset introduced by the inverse approximation used in SR- VN leads it to perform suboptimally compared to VN. VN, on the other hand, computes the full inverse of the FIM defined using natural parameters (Equation (5 and it shows its superiority by achieving one-step convergence in Figure [c) when the step-size is larger, specifically when it is set to 1.
14 ). but does not require
3.6 Variational Parameters and Convexity
NGD updates are often writlen in the natural-parameter space 5 but this can destroy the convexity of the problem. Specifically, given a convex function E(O), it is well-known that Eo. v)() is jointly convex w.r.t. m and the Cholesky factor of V, but not w.r.t V or its inverse (Challis & Barber 2013
Domke 2019p. That is, even for the simplest convex cases, such as logistic regression, the convexity is lost the natural and expectation parameter spaces. For example, given n training data points {x,, y:) } with X; E Rd and Yi [-1,+11, in Bayesian logistic regression minimize the following loss that contains the Ez-regularization with regularization constant 3 > 0:
(0)
Σ log(1 e x:)]]| (xx)))
B |101²
(15)
The Hossian takes the form Volin (0) XDXT + SI, where X _x....., x, IT Rexd is the data matrix, D axd is a diagonal matrix such that DR (0 x1)(1 o (0 (0xx))) with the sigmoid function 1 +e Given this form of the Hessian, it is clear that the Hessian is positive-defnite (and hence regularized logistic loss is B-strongly convex), however, the expected loss (0) ((,) is non-convex in natural and expectation parameterizations, or even in the (m, V) parameterization of the Gaussian.
This issue makes harder to analyze convergence guarantees NGD in Equation ( in general, and therefore for the algorithm given in Equation (7). The same true for the NG flow in Equation (9). because KL is non-convex in (m, V). The issue is specifically noted by AlérieefAAbdllatte et al. 2019 for NGD. They analyze the convergence of online VI by using tools from convex optimization, but their proof techniques do not extend to NGD because of the loss of convexity in the natural or expectation paramcters. Our goal in this paper is to address this challenge. Our key idea is to simply switch to the square-root parameterization where the convexity properties are preserved. We consider specific NGD variants that use such square-root parameterization and derive strong convergence guarantees for them. Additionally, by leveraging the strong convexity of the ligelikelihood we show that KL functional satisfies a Riemannian PL inequality in square-root parameterization, which uvlly leads us to prove convergence for NG flow.

---

Published in Transactions on Machine Learning Research (05/2025)
4 Assumptions
We will use the following assumptions for our analysis. Assumption 1 (Initialization). We initialize Algorithm with To = (mo: Co) = (mo, Cot) for constant Co > 0 and mo E Rd.
Assumption 2 (Strong Convexity). 0, that is, (2) > 81. Then, from Domke 2019 Theorem 9) [(m, C) is 8-strongly-come convee in square-root pormmetrization T (m, C), that is
We assume that the negative log- log-likelihood {(O) is 8-strongly-come conwer in
[(T') 2 [(T) (V+ - -1 -11
(16)
Assumption 3 (Lipschitz Gradient). We assume that the negative log- log-likeldhood {(O) is M-Lapschitzzsmoooh in 0, that is, () & MI. Then, from (Domke, 2020, Theorem 1) C(m, C) is M M-Lipschittz-smottt in square-rool puromelrization T (m, C), that is
M || -11
C(T') < C(T) (V+L()), - T)
(17)
Assumption 4 (Bounded Iterates). Let the iterates generated at time E No from Algorithm 1 be C,. Then the following two statements are true.
1. There erist constants 0 < S < 0o such that & S |Cell I Eu for all t 2 0.
2. There crists a constant Amiu > 0 such that V, & AminT I for all t 0.
Remark 1. The largest singular value of V, (that is, (|Vel can be upper bounded using Assumption 4 follows: ||Vel| = | ||C,Il < C, 1 S
Assumptions and hold for commonly encountered loss functions, e.g. the l--regularized logistic loss f
Assumption ensures that the updates remain bounded throughout optimization, and enforces the positive-defniteness of V.
15
The assumptions hold for Logistic Regression. We saw in Section 3.6 that the Hessian of the logistic loss takes the form XDX B1, where the diagonal matrix D has entries bounded between 0 and 1/4. Hence, the Hessian of the regularized logistic loss is bounded as long as the data matrix X is bounded (which is often the case), i.c., 3T & H < (T where Ç = Manxx(DX in part of Assumption 4 holds with an appropriate selection of step-size in Equation (14.. Additionally, Assumption 2 and part Assumption f are met because regularized logistic regression is 3-strongly convex.
+ B. Thus, the assumption of a bounded square root
Therefore, Assumptions g apply to the original VN's V update when working with the regularized logistic loss. As confirmed empirically, the same conclusion extends to the SR-VN's C update.
5 Analysis
In this section, will analyze the convergence guarantees for the NG flow (Equation 12 and SR-VN (Algorithm D. We will do this by applying tools from convex optimization, considering set of assumpt related to boundedness and convexity. To prove convergence, we need Assumption 2 and part 2 of Assumption 1 for the NG flow, and Assumptions [ for SR-VN.
5.1 Flow Convergence with Riemannian PL
We first show that, under the assumption of strongly convex NLL (Assumption 20, the KL functional satisfies a local Riemannian PL inequality in the square-root parameterization (refer to Lemma This result, paves the way for proving the convergence of the NG flow, as outlined in Theorem
We first state the following lemma that ensures that the FIM bounded, and thus well-behaved, which is crucial for intire analysis. The detailed proof for all the stated results be found in Appendix C
7

---

Published in Transactions on Machine Learning Research (05/2025)
Lemma 1 (Bounded FIM eigenvalues). For (m, C), given Assumptions 4
we conclude that
Laxed
(Lexd
>
I xx2
F
Xnax
Laee Lxex
where
-min
Amin:
min
and Xaax = E2
(18)
min
2
We know that any minimizer dr. = -N(0mm, V.) satisfics Khan & Rue 2023
Lambert et al.
2022a
Eg Eg (0) (0) Vol(O)] 0, and V E@ (0) V (O))
(19)
Existence of limit points. The limit points m", C* exist because of the strong convexity assumption (Assumption 2) which implies there exists T' (m" C*) that uniquely minimizes L. The existence of this T' can be argued using the fact that C is continuous and coercive and Ω (the set of r's) can be extended to closed set (but the minimizer can be shown to lie inside Ω due to Equation 19 exists due to the Weierstrass extreme value theorem.
Then, global minimum
As a consequence of Lemma | one can prove the following lemma, given that NLL is strongly-convex, which establishes a PT. inequality where we measure the magnitude of the gradient the (natural) Riemannian geometry. This will be needed prove the convergence guarantocs for the Riemannian gradient flow and descent.
Lemma (Riemannian PI, Inequality). The KL. functional salisfies a local Tolyak-Eojasieen (PI.) inequality with PL constant μ 8x that is,
V KT. (0) (( D) 2 2p
KT. (r) 11 p(0 DD)
KI. 9r. (0) 1( 1)) (1x(011))]
(20)
FF
Remark 2. Lemma is a 'local' PL inequality because it only holds for the chosen parametrization. This is because of the use of strong-convexity parametrization. Since one does not know convexity will hold all parametrizations there is no guarantee that the PL inequality will hold globally.
ELBO for its proof (sce Appendix C.2 which only holds for (m, C)
Once the Riemannian PL condition is established, we can derive a convergence rate using a standard Lyapunov function analysis. As mentioned earlier, this analysis can be modified to incorporate a bias term as shown in Appendix D
Thcorem 1. Under Assumptions and the NG flow dynamics defined in Equation
12
satisfy
KL qT. (0) |1 p(0 D)
KL 9r. (0) 1| D)
I E 220
200 KL 4T
(0) || p(0 D) KKL 2r. (0) 1l D)
5.2 Convergence of SR-VN
2 give what we believe the first rigorous convergence guarantees for optimizing using Algorithm Specifically, we show an exponential rate of convergence for Algorithm I as stated in Theorem 2 with à detailed proof provided in Appendix C.4 The literalure has suggested various updates, as seen in references Lin et al. (2019 2021b a): Khan & Rue 2023). It is also worth highlighting that our proof extends to generic NGD algorithms that use any square-root paramotrization. The proof technique used derive Theorem can easily be adapted to prove convergence for these other variants and is omitted here
Under Assumption
avoid redundancy.
8

---

Published in Transactions on Machine Learning Research (05/2025)
Diabetes-scale
Mushrooms
Covtype-scale
10a
VN
Training ELBO
10a
SR-VN
Training ELBO
BW-GD
10*
Test NLL
10a
10:
10
10a
10a
2000 Iteration
4000
0
2000
4000 Iteration
6000
8000
0.0
0.5
1.0 seconds
1.5
Diabetes-scale
Mushroams
Leukemia
10a
102
Test NLL
101
Test NLL
10a-
Test NLL
10
10a
109
4000
2000
4000 6000 Iteration
8000
0.0
0.1
0.2
Iteration
seconds
For small-scale LIBSVM datasels, show the training ELBO and test NLL w.r.l number of iterations (left and middle panel) and time (right panel). The min and max of the plotted values are displayed around their averages (taken over five random initializations). We also subtracted the min values achieved (for both training and testing) over all iterations from all three methods and only plotted the resulting values. We can see that, in most cases, SR- VN and VN perform similarly while BW-GD tends to be slower. However, sometimes SR-VN could also slower than the two, for example, see Appendix F.3
Figure 2:
Theorem 2 (Global Convergence). Given Assumptions
1 and considering the limit points m, and C, in
Equation 19 the updates provided by Algorithm 1 satisfy
f(m Ct-1) L(m.; C.) S 2n15 (
(C(mo. Co) - L(m., C.));
(21)
where
MP251
5 PEP
5p
" min Aminp
(22)
2
For convergence, we require that 0 < 270) <1 This constraints the slep-size p >> 0 satisfy cerlain conditions. The permissible step-sizes satisfying the above constraint are discussed in Appendia C.5.
along
with the time complerity of SR-VA in Equation 55
6 Experiments
In this section, we present results on the Bayesian logistic regression problem introduced in 15 to optimize 2 where we compare SR-VN to two other methods, namely 1) BW-GD Lambert et al. (2022b Algorithm and 2) Variational Newton in Equation (7). The BW-GD updates are given as:
BW-GD: m m, a (x(mm, (Ve(O))
M, = I & (6 S.), V M.V,
where a > is the step-size. Due to space constraints, details of all the datasets used and the setting of various algorithmic parameters for these methods are given in Appendix F The additional results and experiments are provided in Appendix F.2
6.1 Bayesian logistic regression
We run our experiments on both small-scale and large-scale datasets compare the aforementioned methods. The way measure the performance of methods is via two metrics: the negative ELBO on the training

---

Published in Transactions on Machine Learning Research (05/2025)
Phishing
Covtype-scale
MNIST
10a
105
Training ELBO
Training ELBO
10
Training ELSO
104
10
10a
10
10a
10a
10a
0
250
500
750
1000
2000
4000
6000
1000
2000 Iteration
3000
Iteration
Iteration
Phishing
Covtype-scale
MNIST
10
10f
10
Test NLL
10a
Test NLL
10a
Test
10a
102
10a
10a
10a
10a
x
250
500
750
1000
2000
4000
6000
1000
2000
3000
Iteration
Iteration
Iteration
Figure 3: Comparison on Large-Scale Datasets. The same trends hold at scale: both SR-VN and VN exhibit comparable performance, whereas BW-GD tends to be slower.
dataset and the NLL on the test dataset. We report the average of the NLL over all test points, i.e. we compare the NLL on the test set computed as follows: -Eer log(1 + exp( - (0T //)eeeesst where Ô is parameter estimate and Peest is the number of examples in the test set.
Datasets. We consider eight different LIBSVM dataset H Chang & Lin 2011 consisting of five small and three large-scale datasets. The description of these datasets is provided in Table 2 of Appendix F Here, we
show results for two small-scale datasets (see Figure 2), namely Diabetes-scale (n = 768, = 8, Atrain = 614) and Mushrooms (n = 8124, d = 112, Merain = 64, 99). For large-scale datasets (sce Figure 3 we show MNIST
(n = 70, 000, d = 784, Perain = 60, 000), Covtype-scale (n = 581, 012, 54, Perain = 500, 000), and Phishing (n 11, 055, d 68, Atrain
844) datasets. All these datasets fall in the case where case, choose the Leukemia dataset (n- 38, d = 129, Merain 34) and show time plots to be inclusive of all possible different settings. Note that different datasets use binary labels other than {-1,+1}, we mapped them all to the same labels so that loss evaluation is uniform across all datasets. All experiments are performed on NVIDIA GeForce RTX 3090 GPUs.
d. For the d > n
We can observe that in the majority of instances, SR-VN performs similarly to VN and both of them are always better than BW-GD. Most often, VN is slightly better than SR-VN, indicating its superiority, and motivating its analysis as future work. As side note, one can observe that the wall-clock time shoots up for SR-VN in the rightmost panel of Figure 2 The reason for this lies in the difference between the computational complexity of algorithmss presented, this is discussed in more detail in Appendix F.T
7
Related Works
NGVI NGVI approximates intractable posterior distributions by employing natural gradients, which consider the geometry of the parameter space, instcad of standard gradients to update the variational parameters Hollman et al. 2013). This approach is believed to lead to more robust and efficient optimization Amari 1998), especially when dealing with complex models and high-dimensional data (Shen et al.) 2024a). Deterministic NGVI (Honkela et al. 2008 2010 Godichon-Baggioon et al. 2024). in contrast to its stochastic counterpart, computes the exact natural gradient at each iteration. This can be computationally more demanding but can also lead to more accurate and stable updates. When feasible, the benefit of this approach is that one can use a constant step-size.
Available https: //wwu. cww.c.iee...... .ntu. http:///hia//tttttto//aaaaaaaaea
10

---

Published in Transactions on Machine Learning Research (05/2025)
Convergence Analysis. While NGVI has shown promising empirical performance- for example, in latent
Dirichlet allocation topic models (Hoffman et al. 2013). Bayesian neural networks (Khan & Nielsen
2018
Khan et al. 2018; Osawa et al. 2019 Shen et al. 2024a), probabilistic graphical models (Johnson et al,
2016), and large-scale Gaussian processes (Hensman et al. 2013 2015 Salimbeni et al. 2018 -te thcoretical
analysis of its convergence rales is still an active area of research Theis & Holfiman 2015). Existing works
often focus on specific model classes or makes simplifying assumptions to derive convergence guarantees. For instance, (Wu & Gardner 2024) showed that stochastic NGVI, exhibits non-asymptotic convergence rate of O(1/T) for conjugate models, where T represents the number of iterations. This rate is comparable to SGD, but stochastic NGVT often has has better constants in the convergence rate, leading to faster convergence in practice (Khan & Nielsen 2018). However, for non-conjugate models, stochastic NGVI with canonical parameterization operates on a non-convex objective, making it challenging to establish a global convergence rate. In the context of Gaussian process models, (Tang & Ranganath 2019) demonstrated the effectiveness of natural gradicnts in both conjugate and non-conjugate scenarios, showing potential for faster convergence compared to traditional gradient- based mechods.
Despite these advancements, comprehensive understanding of the convergence behavior of NGVI in general scttings remains an open question (Regier ct al. 2017). This is partly due to the difficulty in analyzing NGD for commonly-used parameterizations that can destroy the concavity of the log-likelihood (Chérief-Abdellttit 2020). Our work addressed this issue by considering a square-root parameterization for variational Gaussian inference, which helps maintain the concavity of the log- liglliiel and facilitates the convergence analysis. We proved an exponential convergence rate for both NG flow and NGD under this parameterization, bridging the gap between theory and practice for NGD convergence guarantees. Similar efforts to analyze NGD convergence have been made by Khan ct al. (2016) for general stochastic, non-convex scttings and (Chen et al. 2023) for continuous-time flow with concave log- log-likeihooods However, these studies either have slow convergence rates or focus on restrictive settings, leaving room for improvement in more general cases.
Parameterization Choices. The choice of paramcterization plays a crucial role in the effectiveness of NGVT (Lin et al.) 2019). Different parameterizations can lead to different geometries in the parameter space, affecting the trajectory of the optimization process. Studies have explored various parameterizations for NGV1, including natural and expectation parameters (Honkela et al. 2010 Khan & Rue 2023 Moddichnn-Baggion et al. 2024]. Ko et al." (2024) proposed an alternative approach based on score- score-based divergence that can be optimized by a closed-form proximal update for Gaussian variational families with full covariance matrices. They proved that in the limit of infinite batch size, the variational parameter updates converge exponentially quickly to the target mean and covariance. Our work focused on square-root parameterizations for variational Gaussian inference, extending the idea introduced by Tan (2025) who utilized the Cholesky factor. This paramotorization helps maintain the concavity of log-likelihood functions, which in turn facilitates convergence guarantees for both NG flow and NGD. While this approach may lead to suboptimal convergence in some cases, it provides a valuable framework for analyzing and understanding the convergence behavior of deterministic NGVI.
8 Discussion and Conclusion
In this paper, we studied the performance of NG flow and NGD for optimizing the ELBO, a popular objective function in VI. Our findings suggest that the effectiveness of these methods can be affected by the choice of parameterization and the underlying geometry of the optimization landscape; however, a definitive link has yet to be established. Motivated by the need for a principled convergence analysis, we focused on a square-root parameterization for variational Gaussian inference. This approach, building upon the work of Domke (2019), has the promise of preserving the concavity of the log- likeliooo,,,,, a property often lost in traditional parameterizations By preserving concavity, we were able to establish theoretical guarantees for the convergence of both NG flow and NGD under this parameterization, thus bridging a gap between thelretical analysis and practical implementation of NGD.
Our empirical evaluations further validated the benefits of the squarc-root parameterization. We observed that it generally performs comparably to the standard natural parameterization, exhibiting fast convergence due to its Newton-like update. Moreover, our experiments, particularly the toy (ill-conditioned) problem in
11

---

Published in Transactions on Machine Learning Research (05/2025)
Figure highlighted the role of the square-root parameterization in maintaining stable updates, especially in scenarios where traditional parameterizations might be sensitive to parameter choices errors in the Fisher information matrix estimation.
Despite the advancements, challenges and open questions remain. Our findings pave the way for future research in several key directions, this includes extending our analysis stochastic NGVi-common real-world settings-to study how noise influences optimization and whether the square-root parameterization remains effective, along with the validity of key assumptions (e.g., Assumptions 4 for more complex applications like deep learning.
Broader Impact Statement
Our work makes several contributions: i) we propose theoretical analysis using a square-root parameterization for the Gaussian covariance, which we expect can encourage further research into establishing convergence guarantees for natural-ge natural-gradient variational inference, particularly for cases involving concave log-likelihoods. ii) Our experimental results reveal that natural gradient methods outperform algorithms relying on Euclidean Wasserstein geometries, showcasing their practical advantages in terms of convergence speed and effectiveness, which could impact the development of more robust inference algorithms in probabilistic modeling.
Acknowledgments
Navish Kumar and Aurelien Lucchi acknowledge the financial support of the Swiss National Foundation, SNF grant No 207392.
References
Vincent Adam, Paul Chang, Mohammad Emtiyaz E. Khan, and Arno Solin. Dual parameterization of sparse variational gaussian processes. Advances in Neural Information Processing Systems, 34:11474 11486, 2021.
Pierre Alquier, James Ridgway, and Nicolas Chopin. On the properties of variational approximations of Gibbs posteriors. J. Mach. Learn. Res. (JMLR), 17:239:1-239:41, 2016.
Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):251-276, 1998.
Shun-Ichi Amari. Information geometry and its applications: Convex function and dually flat manifold. In LIX Fall Colloquium on Emerging Trends in Visual Compuling, pp. 75 102. Springer, 2008.
Shun-ichi Amari. Information geomelry and ils applications, volume 194. Springer, 2016.
N Baba. Convergence of a random optimization method for constrained optimization problems. Journal of Optimization Theory and Applications, 33:451-461, 1981.
Hans-Georg Beyer. The theory of evolution strategies. Springer Science Business Media, 2001.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American slatistical Associalion, 112(518):859 877, 2017.
Charles Blundell, Julien Cornebise, Koray Kavukeuoglu, and Daan Wierstra. Weight uncertainty in neural networks. In International Conference on Machine Learning (ICML), 2015.
O Catoni. Pac-bayesian supervised classification: The thermodynamies of statistical learning. institute of mathematical statistics lecture notes maccsraph series 56. IMS, Beachwood, OH. MR2483528, 5544465, 2007.
Edward Challis and David Barber. Gaussian kullback-leibler approximate inference. Journal of Machine Learning Research, 14(8), 2013.
Chih-Chung Chang and Chih-Jen Lin. LIBSVM: a library for support vector machines. ACM transactions intelligent systems and technology (TIST), 2(3):1-27, 2011.
12

---

Published in Transactions on Machine Learning Research (05/2025)
Yifan Chen, Daniel Zhengyu Huang, Jiaoyang Huang, Sebastian Reich, and Andrew M Stuart. Gradient flows for sampling: Mean-field models, gaussian approxinations and affine invariance. arv preprinl arXiw:2302.11024 2023.
Badr-Edding Chérief- Abdellatf...
Convergence rates of variational inference in sparse deep learning. In International Conference on Machine Learning, pp- 1831-1842. PMLR, 2020.
Badr-Eddine Alérref-Abdellatti Pierre Alquier and Mohammad Emtiyaz Khan. A generalization bound for online variational inference. In Asian Conference on Machine Learning (ACML), 2019.
Jon Dattorro. Convex optimization and euclidean distance geometry. Book, 2015. attps://ccrra..saanfff http:/aattarrr0077766/000000000. pd.
Michael Ziyang Diao, Krishna Balasubramanian, Sinho Chewi, and Adil Salim. forward-backwart gaussian variational inference via jko in the busss-wwssrrstein space. In International Conference on Machine Learning, pp. 7960-7991. PMLR, 2023.
Justin Domke. Provable gradient variance guarantees for black-box variational inference. Advances in Neural Information Processing Systems, 32, 2019.
Justin Domke. Provable smoothness guaranlees for black-box variational inference. In International Conference on Machine Learning (ICML), 2020.
Justin Domke, Robert Gower, and Guillaume Garrigos. Provable convergence guarantees for black-box variational inference. Advances in neural information processing systems, 36, 2024.
Tobias Glasmachers, Tom Schaul, Sun Yi, Daan Wierstra, and Jürgen Schmidhuber. Exponential natural evolution strategies. In Proceedings of the 12th annual conference on Genetic and evolutionary computation, pp. 393-400, 2010.
A Godichon-Baggioni D Nguyen, and M-N Tran. Natural gradient variational bayes without fisher matrix analytic calculation and its inversion. Journal of the American Statistical Association, pp. 1 12, 2024.
Alex Graves. Practical variational inference for neural networks. In Advances in Newral Information Processing
Systems (NeurIPS), 2011.
Elad Hazan, Kfir Yehuda Levy, and Shai Shalev-Shwartz. On graduated optimization for stochastic non-couvex problems. In International conference on machine learning, pp. 1833-1841. PMLR, 2016.
James Hensman, Magnus Rattray, and Ncil Lawrence. Fast variational inference in the conjugate exponential family. Advunces in neural informalion processing syslems, 25, 2012.
James Hensman, Nicolo Fusi, and Neil D Lawrence. Gaussian processes for big data. arXio preprinl 2013.
James Hensman, Alexander Matthews, and Zoubin Ghahramani. Scalable variational gaussian process classification. In Artificial Intelligence and Statistics, pp. 351-360. PMLR, 2015.
Matthew D Hoffman, David M Blei, Chong Wang, and John Paisley. Stochastic variational inference. Journal of Machine Learning Research, 2013.
Antti Honkela, Matti Tornio, Tapani Raiko, and Juha Karhunen. Natural conjugate gradient in variational inference. In Neural Information Processing: 14th International Conference, ICONIP 2007, Kitakyusha, Japan, November 13-16, 2007, Revised Selected Papers, Part II 14, pp. 305 314. Springer, 2008
Antti Honkela, Tapani Raiko, Mikael Kuusela, Matti Tornio, and Juha Karhunen. Approximate riemannian conjugate gradient learning for fixed-form variational bayes. The Journal of Machine Learning Research, 11:3235-322 32288, 2010.
Tommi S Jaakkola and Michael I Jordan. variational approach to bayesian logistic regression models and their extensions. In International Conference on Artificial Intelligence and Statistics (ATSTA TS), 1997.
13

---

Published in Transactions on Machine Learning Research (05/2025)
Matthew J Johnson, David K Duvenaud, Alex Wiltschko, Ryan P Adams, and Sandeep R Datta. Composing graphical models with neural networks for structured representations and fast inference. Advances in neural information processing systems, 29, 2016.
Mohammad Khan and Wu Lin. Conjugate-computttion variational inference: Converting variational inference non-conjugale models to inferences in conjugate models. In Artificial Intelligence and Statistics, pp- 878 887.887.20 PMLR, 2017.
Mohammad Khan, Didrik Nielsen, Voot Tangkaratt, Wu Lin, Yarin Gal, and Akash Srivastava. Fast and scalable bayesian deep learning by pertur-pprturbation in adam. In International conference machine learning, pp. 2611 2620. PMLR, 2018.
Mohammad Emtiyaz Khan and Didrik Nielsen. Fast yet simple natural-gradient descent for variational inference in complex models. In 2018 International Symposium on Information Theory and Its Applications (ISITA), pp. 31 35. IEEE, 2018.
Mohammad Emtiyaz Khan and Håvard Rue. The bayesian learning rulc. Journal of Machine Learning Research, 24(281):1-46 2023.
Mohammad Emtiyaz Khan, Reza Babanezhad, Wu Lin, Mark Schmidt, and Masashi Sugiyama. Faster stochastic variational inference using proximal-gradient methods with general divergence functions. In UAI'16, UAI'16, pp. 319 328, Arlington, Virginia, USA, 2016. AUAI Press.
Mohammad Emtiyaz Khan, Wu Lin, Voot Tangkaratt, Zuozhu Liu, and Didrik Nielsen. Variational adaptive- newton method for explorative learning. arXiv preprint or Xin: 71 1... 05560, 2017.
Kyurae Kim, Jisu Oh, Kaiwen Wu, Yian Ma, and Jacob Gardner. On the convergence of black-bar variational inference. Advances in Neural Information Processing Systems, 36, 2024.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint xrx:133312.6114 2013.
David Knowles and Tom Minka. Non-conjugate variational message passing for multinomial and binary regression. Advances in Neural Information Processing Systems, 24, 2011.
Joohwan Ko, Kyurae Kim, Woo Chang Kim, and Jacob R. Gardner. Provably scalable black-box variational inference with structured variational familics. In Proceedings of the 41st International Conference on Machine Learning, ICML/24. JMLR.org, 2024.
Mare Lambert, Silvere Bonnabel, and Francis Bach. The recursive variational gaussian approximation (r-vga). Statistics and Computing, 32(1) (10, 2022a.
Mare Lambert, Sinho Chewi, Francis Bach, Silvère Bonnabel, and Philippe Rigollet. Variational inference via wasserstein gradient flows. Advances in Newral Information Processing Systems, 35:14484-14447, 2022b.
Marius Leordeanu and Martial Hebert. Smoothing-based optimization. In 2008 IEEE Conference on Compuler Vision and Pottern Recognition, pp. 1-8. IEEE, 2008
Wu Lin, Mohammad Emtiyaz Khan, and Mark Schmidt. Fast and simple natural-gradient variational inference with mixture of exponential-family approxinations. In International Conference Machine Learning, pp- 3992-4002 PMLR, 2019.
Wu Lin, Frank Nielsen, Mohammad Fatiyaz Khan, and Mark Schmidt. Structured second-order methods via natural gradient descent. arXiv preprint 1018:22107,0888, 2021a.
Wu Lin, Frank Niclsen, Mohammad Emtiyaz Khan, and Mark Schmidt. Tractable structured natural-gradient descent using local parameterizations. In International Conference on Machine Learning, pp. 5680-6691. PMLR, 2021b.
14

---

Published in Transactions on Machine Learning Research (05/2025)
Wu Lin, Valentin Duruisseaux, Melvin Leok, Frank Nielsen, Mohammad Emttyaz Khan, and Mark Schmidt. Simplifying momentum-baseed positive-definite submanifold optimization with applications to deep learning. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of
Proceedings of Machine Learning Research, pp. 21026 21050. PMLR, 23 29 Jul 2023. URL https: //proceedings... mlr. ppwess/v02//iii/ html
Wu Lin, Felix Dangel, Runa Eschenhagen, Kirill Neklyudov, Agustinus Kristiadi, Richard E. Turner, and Alireza Makhzani. Structured inverse-free natural gradient descent: Memory-cfficient and mumerically-stable KFAC. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 29974 29991. PMLR, 21 27 Jul 2024. URL pttps://proceedeee
mlr. wwr..presssssee 11/233/11n2245 html
Tianle Liu, Promit Ghosal, Krishnakumar Balasubramanian, and Natesh Pillai. Towards undorstanding the dynamics of gaussian-slein variational gradient descent. Advances in Newral Information Processing Systems, 36, 2024.
Luigi Malagò and Giovanni Pistone. Information geometry of the Gaussian distribution in view of stochastic optimization. In Proceedings of the 2015 ACM Conference on Foundations of Genetic Algorithms, pp- 150 162, 2015.
Benjamin M Marlin, Mohammad Emtiyaz Khan, and Kevin P Murphy. Piccewisc bounds for estimating bernoulli-ogiistic latent gaussian models. In ICML, pp. 633 640, 2011.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukeuoglu. Asynchronous methods for deep reinforcement learning. In International conference on machine learning, pp. 1937. PMLR, 2016.
Hossein Mobahi and John Fisher ITI. A theoretical analysis of optimization by gaussian continuation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 29, 2015.
Tain Murray. Differentiation of the cholesky docomposition. arXiv preprint Xix::1600.0556 2016.
Kazuki Osawa, Siddharth Swaroop, Mohammad Emtiyaz F. Khan, Anirudh Jain, Runa Eschenhagen, Richard E Turner, and Rio Yokota. Practical deep learning with bayesian principles. Advances in neural information processing systems, 32, 2019.
K. B. Petersen and M. S. Podersen. The matrix cookbook, nov 2012. URL http::/www2. compute dtu. dk/ html Version 20121115.
Rajesh Ranganath, Sean Gerrish, and David Blei. Black box variational inference. In Artificial intelligence and statistics, pp. 822. PMIRR, 2014.
Jeffrey Regier, Michael Jordan, and Jon McAuliffe. Fast black-box variational inference through stochastic trust-region optimization. Advances in Neural Information Processing Systems, 30, 2017.
Tim Salimans and David A. Knowles. Fixed-form variational posterior approximation through stochastic lincar regression. Bayesian Analysis, 8(4), December 2013. ISSN 1936-0975. doi: 10.1214/13 ba858. URL
http://dx..o dor.ddddoo.ooo/110/1 wtt.:/orrrr.....1////--aa-
1214/13-BA85
Hugh Salimboni, Stefanos Eleftheriadis, and James Hensman. Natural gradients in practice: Non-conjugate variational inference in gaussian process models. In International Conference on Artificial Intelligence and Statistics, pp. 687. PMI.R, 2018.
Yuesong Shen, Nico Daheim, Bai Cong, Peter Nickl, Gian Maria Marconi, Clement Bazan, Rio Yokota, Iryna Gurevych, Daniel Cremers, Mohammad Emtiyaz Khan, and Thomas Mollenhoff. Variational learning is effective for large deep networks. In Proceedings of the 4ist International Conference on Machine Learning, 1CML.24. JMr.R.rrg 2024a.
15

---

Published in Transactions on Machine Learning Research (05/2025)
Yuesong Shen, Nico Daheim, Bai Cong, Peter Nickl, Gian Maria Marconi, Clement Bazan, Rio Yokota, Iryna Gurevych, Daniel Cremers, Mohammad Emtiyaz Khan, et al. Variational learning is effective for large deep networks. arXiv preprint orXic:2402.176.. 2024b.
James C Spall. Introduction to stochastic search and optimization: estimation, simulation, and control. John Wiley & Sons, 2005.
Joe Staines and David Barber. Variational optimization. arXiv preprint orXiv:1212.4507, 2012.
Yi Sun, Daan Wierstra, Tom Schaul, and Jürgen Schmidhuber. Efficient natural evolution strategies. In Proceedings of the 1 11th Annual conference on Genetic and evolutionary computation, pp. 539 546, 2009.
Richard S Sutton, Andrew G Barto, et al. Introduction reinforcement learning. vol. 135, 1998.
Linda SL Tan. Analytic natural gradient updates for cholesky factor in gaussian variational approximation. Journal of the Royal Statistical Society Series B: Statistical Methodology, pp. qkaf001, 2025.
Da Tang and Rajesh Ranganath. The variational predictive natural gradient. In International Conference Machine Learning, pp. 6145-6154. PMLR, 2019.
Mare Teboulle. Eutropic proximal mappings with applications to nonlinear programming. Mathematics of Operations Research, 17(3):670-699, 1992.
Lucas Theis and Matt Hoffman. A trust-region method for stochastic variational inference with applications streaming data. In International conference on machine learning, pp. 2503 2511. PMT.R, 2015.
Michalis Titsias and Miguel Sezaro-Gredilla. Doubly stochastic variational bayes for non-conjugate inference.
In International conference on machine learning, pp. 1979. PMLR, 2014.
Martin J Wainwright, Michael I Jordan, et al. Graphical models, exponential families, and variational inference. Foundations and Trendsø in Mochine Learning, 1(1 2):1 2):1-305, 2008.
Ronald J Williams and Jing Peng. Function optimization using connectionist reinforcement learning algorithms. Connection Science, 241-2288 1991.
Kaiwen Wu and Jacob R Gardner. Understanding stochastic natural gradient variational inference. arXiv preprint arXiv:2406.018 2024.
Tong Zhang. Thcoretical analysis of class of randomized regularization methods. In Proceedings of the twelfth annual conference Compulational learning theory, pp. 156 156-163, 1999.
16

---

Published in Transactions on Machine Learning Research (05/2025)
Table of Contents
Gradient Calculations
18
A.1 Gradients in (m, V) Parametrization
A.2 Gradients in (m, C) Parametrization
19
B SR-VN as an Approximation of VN (Detailed)
20
C Analysis Proofs
21
Proof of Lemma
22
Proof of Lemma
23
C3 Proof of Theorem
23
Proof of Theorem
23
C.5 Permissible Step Sizes for SR- VN
25
D Biased SR-V Convergence
26
E Piecewise versus stochastic implementation
27
F Experiments Detail
28
1 Runtime comparison
29
2 Algorithmic Details and Additional Results
30
3
Iteration Plots with Accuracy
31
Time Plots with Accuracy
33
Comparing Geometries: SR- VN vs. Alternatives
35
Non-convex Model
36
17

---

Published in Transactions on Machine Learning Research (05/2025)
A Gradient Calculations
This section focuses on showing calculations that lead the gradients used to prove results in the main paper. In Appendix D in Khan & Rue (023)), the gradients used to define VN were calculated using the moment parametrization (which were calculated under this paramcterization. The goal of this section is to use these gradients to derive the gradients w.r.t T = (m, V) and t = (m, C) parameterizations.
will define in bit) μ (Mr, (2), i.e. gradient of the objective L(u)
Before we begin, it useful to learn about the two important parameterizations for the Exponential family. The first one is called moments or mean/ mean/erpeer exppecateitn parameters, and are given as:
M1 E Rd μ2 = + mm S
(23)
where 84 = (X Raxd - x 0} is the set of symmetric positive-defiiite matrices. The second parametrization given via the natural paramelers as
1 S sd
Sm E Rdd X2
(24)
The way these two parameterizations are interrelated is via the following relation
# = (AH = H(H())= = A,
(25)
where (X) is a convex function, and its conver conjugate H(U) returns the entropy of the Gaussian distribution given its moments μ.. These details are explained in Chapter 3 in Wainwright et al. (0008)
see
also the paper by Malagò & Pistone 2015
A.1 Gradients in (m, V) Parametrization
Now let us look at the derivatives of f w.r.t and f12, but first we note that:
25 24
H(H) A
Sm,
(26)
Let us define
g(T) = Eø (o)())), H(T) = Ba [z((0),,
where T = (m, V) when 97 = N(m, V), and when this is the case, we omit writing g(m, V), H(m, V) in favour of only writing g. H. Then, Equation 10 Khan & Rue (2023) is given that:
Fp f14 = Em~N(m (v (v m g Hm.
(27)
Combining the above two equations, we get:
V ((H) Vm Eo_N(m,v)((00) +
V μi Eo [((0) + V H(H)
g Hm Sm,
(28)
and from Equation 11 in Khan & Rue 2023
we have
V us C(H) = (H-78)
(29)
We take M(D) the negative entropy this paper.
18

---

Published in Transactions on Machine Learning Research (05/2025)
This further leads to
0c(m)
OC
0p1 dm,
dc
T 0p2
(VmC(H));
+ tr
Om,
ppm
Op12
am,
ac
ac
0pp
Σ
0pp
0p12
0ms
1
Oc
aC
0C
OC
0pm
Σ
0pz
m& +
Σ
my
+
0pp2
m
ki
il
Oc
Oc
OC
+ Σ
my + 2
174
0pm
0pz
0uz
lepfi
+ (V Lm); (M L)
b-th rOW
vmL= Vu,
C Lr,C
2V M2 Lm
28 29
Hm + +mm, +2 + 7S) T m
g.
(30)
Similarly,
(VvL(u)) 0C(K)
OC
T 1p
0C 8pp
ir
0V:
tl
iV
0112
OVi
O0
0p1
Σ
0pz
as
OC
aC
0pl
0pp ij
VvL
uu2C
@L(1))
0p1
V (()))
0p
O(S)
H - -S))
(31)
A.2 Gradients in (m, C) Parametrization
Let us represent (()( as L(p1,12(C) the chain rule for composite matrix functions given in Equation 1884 in Dattorro (015) combined with Equation 137 in Petersen & Pedersen 2012), we arrive at
u2 (C)) where μn = m (independent of C) and #2 (C) = V = CCT, Using
VcL(u(C)) ((V (V C(H)) (VCH((C)) + tr ((p))) VeH2(C)) tr ((H)) VCH2(C)).
which leads to
ac(M(C))
0c 0pz
tr
00ij
Oju 0C.;
19

---

Published in Transactions on Machine Learning Research (05/2025)
Note that 0pa is fourth-order tensor whose entries are 8C
Cj
i, + i
a(cc)) )m
2Cij
i, i
(32)
aCig
Ckj
o =
0
else.
where we used the expression (CC )at = En Ckn Cill= = =En Ckn Cta. Then
@C(p(C))
Oc
0hz
tr
8Cig
8uz
aCis
OC
0hz
Σ
0pl2
8Cij
k2l
Oc T
0p2
dc T
0pt
ac
0pp2
Σ
Σ
+ Σ
0pz
00i,
0pm
aCis
aCis
k=
k=i,=i
M
Ik
k
Uk
OL
OL
0C
Cij +2
Cij
+
Ckj
8p4
0pp
0p12
ii
ki
7C
ac
Σ
Ctj + Σ
Ckj
0pp
0p12
ki
OL
OC C
C
+
0p2
01z
i9
Therefore, we conclude that
0L(2z(C)) (V (2 C+ (V =
(33)
8C
(l(H²(C))
29
HC - -C
(34)
ac
as Wu VvL= VVC=(VvC)+
B SR-VN an Approximation of VN (Detailed)
For 1, the update of S, from VN in Equation can be rewritten in terms of V as,
V=l1 (1 V + pH,
(1 C- C pH.
C. (1- p)T pC, H,C, Cj
Ci
C 1+ (1-
p c. H.C. (.CCC
VAI
C.
(1- p)
Inverse
This above inverse can be approximated using the following truncated Neumann series for any symmetric
matrix A,
(I
to give
C
-- P
C. H.H.C.+ (1- ) C.H.C.CCHH..CC cT.
V
(1 (1-p)
20

---

Published in Transactions on Machine Learning Research (05/2025)
Let us define, W. H,H,C, W. tril W.]. and make use of the following identity: w = W, + W! to arrive at
C₁
I p
W + Wo + P
(1 (1-p) p)) )
2
W. + w. c
V41
(1 --p) p)
(1 (1-p)
C2
||- P
2(1- We 2 3 p
W C
(1 1- p)
(1-p)
oo C. (1 p] pw, C. 0(p2))
(35)
where (a) comes from also expanding p) p+ p² + Op2)), at all occurrences.
Now, let us write down the update for V441 using the updates from SR-VN 14] (writing a for the step-size for SR- VN here). We have
V4+1 C
=C, - aC, trilCC H.C. (aC,-I)((C, - ac, tril(C H.C, -I) I))
a 2 ((1+
T 1 - -ailC ((1+ I 1- [nilIC 7 c.
C,
((1+ a) I - ww. (1++g)
2) g)
a - T-aw, T C.
C.
[-2| + W WT]+ W. w..w. C
C.
I+a(I - [ + w.) +
4
= =C, (1+ oW, c. + +0(0)
(36)
where a is chosen to be p. Comparing Equation 35 with Equation (36) shows that the two updates for V only match up to the first order. The presence of quadratic difference in Equation (36) ariscs due to an alternative approximation of the inverse, highlighting the distinctions in the update methods and exposing the less-than-optimal nature of SR-VN.
C Analysis Proofs
Table 1: Vectorization Notation for a square matrix A.
(def)
A
lower triangular matrix derived from A by replacing all supra- diagonal elements by zero
(def)
diag(A)
diagonal matrix derived from A by replacing all non-diagonal elements by zero
Ä
(def)
Ā A diag(A)/2
(def)
vec(A)
vector obtained by stacking the columns of A in order from left to right
(def)
(e)
vector obtained from vec(A) by omitting supra-diagonal elements.
(def)
K
commutation matrix such that K vec(A) vec(A) A T)
(def)
elimination matrix such that L vve(AA = vech(A)
21

---

Published in Transactions on Machine Learning Research (05/2025)
C.1 Proof of Lemma 1
Proof. In Tan 2025
Lemma 1), the inverse FIM in (m, C) parametrisation has the following form:
Fm
V
F7'(m, C)
Fc =
0
F22
(37)
where, given N L2)/2
F L(a @ - CT
(38)
Let work first to prove the LHS of the claim. From part of Assumption For F we proceed follows:
4
we first see that Fmml
> Amin
|F''l12 2 FFc
a77
2025
Lemma 1()))
[2L[T C C-TINICC-
Now, notice that N - [1 + ||L (2) < since | S 1, while being a permutation matrix. Since
L is also form permutation matrix (with additional zeros appropriate places), we have that | S Combining these facts with the norm inequality, we obtain
F 112 12L(1C P (AN(&CCI)LL 2 II C- 1 C-
())
21|C-11
2
:
|
2
2
Amin
where (a) comes from using |A & Bl, S
A
B Bl
AA, B E Raxd and |B |B
Now let us prove the RHS of the claim. From Remark inequality as follows:
we have F = V c For Fc we use the norm
Fc
11,
[L
||
(LNL > IL @ & c
2
1
S
[LNL >
||, ||C
2
Assumption
4
c2
[[NL
>-1112
(39)
2
Let us see how the operator LNL T [d(d+1)/2 + Iddd+1)/2 acts on vech(A) E 1Rd(d11/22 By definition, we have
K L LT vcch(A)
(LNLT) (A)
L
2
(LLL ) vech(A) +
(LL vech(A))
La(4+1//2
(LK vec(A) A) + voch(A))
(:. vech(A) = vec(A))
2
(L vec(A )) + vech(A)).
(: K vee(A) vec(A T))
(vech(A) (A) )) (ech(A))),
(: L L (eo(AT) vech((T))
(40)
22

---

Published in Transactions on Machine Learning Research (05/2025)
where Equation (40)
implies that the operator LNL is equivalent to halving the sub-diagonal entries while keeping the diagonal of the matrix A intact. Therefore, LNL has eigenvalues ween 1/2 and 1, which implies that its inverse has eigenvalues bounded between and 2, leading to [LNLL S 2. Thus, after plugging in Equation 399 this bound, we have
F 12 - Eu.
(41)
allowing us to complete the proof of the main claim.
C.2 Proof of Lemma
Proof. Notice that Equation 22 is equivalent
|++((T]] 22
2p - [(T.)),
(42)
because KL
9r (0) || p(0)
can be written
KL 9- (0)
p(0) loggee V) + Em Eeg-N(m.v v)[[(0]] + const.
[(T) + const.
Then simply note that by Lemma we obtain that
| F. 2 Aginn
|||
[(m, C) - L(mm,, C,) 20
where now μ Ammin is indeed the desired the PI, constant.
C.3 Proof of Theorem
Proof. For conciseness, we will refer KL 9r, (0) |l p(0 D) Lyapunov function E(t) = 22tt KL.), then
= KLt throughout the proof. Let us define the
Leumna
= (KL; KL.) But v KL2 +) The energy functional {() therefore monotonically decreasing, and hence
2 (2n(KL; KL.) V KLLllee
0.
E(t) 2ut (KL, KL.) s E()) KLg KL.
(43)
C.4 Proof of Theorem 2
Proof. From Assumption we start by writing the following Taylor expansion for L(m, C) when both and C change.
L(m++1; C(+1) < L(m,, C.) + V C(m,, C.) C m.) ir((VoL(m. (C,1 C.))
M
2 m
m +
2 1 - C.Il..
(44)
By plugging in values from Equation
14 in the joint Taylor expansion above, we arrive at:
C(m++1, C++1) L(m,, C.)
< - S-(VmL(mi, C.) C,CT( V Cm,, C.)) + Mp² |.C.C, VmL(me, C.)l
Mp²
-prr (VoLmm, C.)) C,H, +
|CC.H
(45)
2
8
23

---

Published in Transactions on Machine Learning Research (05/2025)
where H. tril[C, VeL(m, Ct)l. Using the following facts
_m, C,C, mm C₁) 1 Amin V C,)ll
,,,,, JC,C//VC(mm vm cJll < |c,C, [[m, C.)Il Remark
2 | [m,, C.)Il
and applying the lower for bound Amin assumed in Assumption
to term A, we arrive at
x Amin || mmm,, C+) + Mp
Mm,,m, C
2
which then gives us
A & Wm Vm,(m; C+)
(46)
where
Mp²
Aminp.
Now for term B, let us first look at the following term by repeatedly applying the matrix norm inequality to obtain
[| I || [B||))
(a)
(b) I c.,12
JC,HH P < | H
|
2
2
+
||iag H|
(2 2
||diag (H)
2
2
5
|c
2 H
H
2
||N
H
F
2
P
5
|c,l
C
K
s
5
|all
K
(e)
5
|
2
5 |
VeC C.I
2
where (a) comes from using the fact that VA, B E RDDD we have | - BI S || 2 and choosing A H and B = diag(H)/2. In (b) and (c) we have used ||||| S ||A Fs as A is a lower triangular
form obtained by sending all supra-diagonal elements of A to zero. For the first term B, we have
Ir (VVcl(m, C.)) C.H.) |Vo[(m, Co)ll |C,.Hll
(r < |||| |B| r)
|all VeL(m, C.ll
2
This leads to finally obtain
B - T call (Vo C 5Mp²
|Ca,1 [VoC(m, C
Now using the bounds in Assumption
we arrive at
B < we I C
(47)
where
5p
V 2
WC
4
Now plugging Equation
46
and Equation (47 in Equation
45
to get
A
A < wm
V [(me, C+) + we [eLLm, C.)l
A
A <
V C(m, 1
(48)
24

---

Published in Transactions on Machine Learning Research (05/2025)
where m w and we require that nj 0. Now from Assumption by joint d-strongly convexity of C in (m, C), we have
1
C(m;, C2) L(m", C*)
(m
28
V (m, C(m,, C
26
(A - A
2m
(1 2
(49)
from which conclude that indeed obtain descent for in (m, C) if we require
0<1 - 2mj6 + = 1
(50)
26
This would further lead the global convergence
A s (1 (1-2m78) 2
D =..0
(51)
0
C.5 Permissible Step Sizes for SR-VN
Theorem 2 requires the condition 0 < 28 hold. This might prompt the reader to inquire whether specific step size 1) can be determined to ensure the fulfillment of this condition. We answer this question below in the context of a constant step size that we also employ in our experimental work.
Constant step size. Let Amax = E be the upper bound for the largest eigenvalue of Vt, as derived in
Remark | then Equation 22 becomes
Mp2cd 2
5 pe
5
7 min
Amiup
Note that both arguments inside the min above are quadratics in p. Let us choose p such that
S (axx
(52)
where
Amin
2 &
and P2 V
P1
X M
X2
MC
are the extremal points of the quadratic equations the first and second terms inside the min, respectively. Choosing p = Pi for the first argument and p P2 for the second argument yields
& 2M
7 S min
Ainn
(53)
2X2 D
M
In order to ensure that the rate of convergence is contractive, we require (1 - 2n²) < 1, that is 0 < 2nj < 1.
With the choice above, 7) is obviously positive, so let us locus on 2716 < = 1) < 23 From the first quantity in the min appearing in Equation (53 we need
X
min
1
min
<
(54)
212
M
28
M
max
max
The quantity Ania a lower bound on the inverse condition number of the preconditioning matrix V. and is Amax thus upper bounded by Since M the smoothness constant of L(m, C), then the quantity upper bounded by the inverse condition number of the Hessian of L(m, C) which is itself upper bounded by 1.
25

---

Published in Transactions on Machine Learning Research (05/2025)
This means that Equation [54] is always satisfied with the proposed choice of step size. The same argument applies to the second quantity in the min appearing in Equation 53
Intuition: We here give an intuitive explanation for the choice of the step size p. By Theorem 2 we see that the speed of convergence exponential with respect to time t. The contraction factor C₁ := (1 +
directly depends on the choice of step size through 7 defined in Equation (22 From the latter equation, one can observe that larger step sizes can be used if the objective function is smoother (which is controlled by M). In turn, larger step sizes allow for faster convergence as the algorithm takes larger steps towards the minimum, which means that (1 2nd) is smaller and therefore C. contracts at faster rate.
Complexity Calculation. To compute the complexity, we require [(,,, C₁) - -m,,, C.) S e, meaning from Equation 21 we must have
2))) Co) - [(m., C.)) Se = iS log(e/(C(mo)> Co) - L(m,, C.))) log(1 211)
Now using Equation we arrive
logle/(C(moo Co) - L(m.; C.))) log(e/(C((mn. Co) - L(m., C.))) 5 Xgnin
2 min
(55)
log(1
log(1
E
M max
where the switching of inequality happens because the log(-) term in the denominator is negative for both terms inside min.
D Biased SR-VN Convergence
Definition 1 (Biased Gradient Oracle). map g IRd+dxd x D - Rd 8.1.
)(T) g + (T)
(56)
for a bias by Rd+dxd + Rd Definition 2 (Biased Hossian Oracle). A map H dd+d x D - Rdddd s.t.
H(r) H -T)
(57)
for a bias by Iddddxd + Baxd
We assume that the bias is bounded: Assumption 5. (Bounded biases) There erists constants 0 S mg my < 1, and Ch, Gn. 2 0 such that
[T]( my C(m;, C.)l + Cg VT E Bdddddd
||b(T]) < mn VeLmm., Co)ll Cie VT Rd+dxd
With the prerequisites established, SR-VN update. This bias may arise from estimating the expectation using piecewise bounds presented in Marlin et al. (2011).
present theorem demonstrating the convergence of the biased
Theorem 3 (With Bias Convergence). Given Assumplions and considering the limit points m. and C.
in Equation 19 the updates provided by Algorithm 1
satisfy
L( C++1) L(m,, C.) S (1 2n5) 21 (C Co) L(m,, C.)) + Wm +wo-oo wo
Proof. Following similar steps, can arrive at the following versions of Equation 46 and Equation
47
A Wm
H
B we (,)
(58)
26

---

Published in Transactions on Machine Learning Research (05/2025)
where Vel(me, Ci) (H; ((,,C,,) Now using Assumption 5 one can get
A s Wm
H
Wm
H + by(TT)] S wm
[He 2 [Ib)( s wm (1) mg) Hel (3) (59)
Similarly,
B we VcL(m,, C.) wc
H. + +x,+t))))
F
H. (H.- C, 1
+wo (bbr(t))C.11
wc
wc (VcC(m,,C))) + [Tb
Swo (1+ (eee))ll/C(mmmm call +
(60)
Now plugging Equation 59 and Equation
60 in Equation
45 to get
A A++ A
< Wi (1 + mg) Cmm,, C.) wo(l + mn) [VC[m,, C.) + wmL wolH
(61)
A H+1
< - (m) C.) + w wc
(62)
where 7 min Wm (1 + mg), wc + mH}} and require that 7j > O. Now from Assumption
by joint
d-strongly convexity of L in (m. C), we have
C(m;, C.) - L(m", CC) S 1 (m,
25
A
V Vmm, Ci)l
28
62
(A A++1) + 2nd (wm 22 + wo(h)
2nd
[+1
(1 2n10 + Wm 2 + wohh
(63)
Remark 3. is worth noting that the piccewise mothod for computing expectations, given in Marlin et al. (2011), involves computing a 1-D integral which one can compute below numerical float precision. In other words, the bias with piccewise bounds can be made arbitrarily small by increasing the number of pieces. The bias is therefore very small, and hence the terms Cy and G) are negligible which shows that SR-VN converges long as this holds.
E Piecewise versus stochastic implementation
In this section, we make comparison between the piecewise and the stochastic implementation of the expectation in ELBO, see Figure g This plot highlights why the piecewise implementation is more accurate and laster representation of the deterministic algorithm we analyzed, and thus the preferred choice for the experiments carried out in the paper.
27

---

Published in Transactions on Machine Learning Research (05/2025)
Piecewise and Stochastic (minibatch = 32, MC samples= 1) implementations for SR-VN
(n, (800, 2)
(n, d) (200, 2)
1.0
10
Piecewise
10
Training ELBO
Stochastic
Test Accuracy
0.8
10a
NLL
10a
10
0.6
101
10
0.4
10
10
0.2
0
200
400
600
800
0 200
400 600
800
0
200
400
600
800
Iteration
Iteration
Iteration
Figure 4: The provided comparison highlights the differences between the piecewise implementation and the (minibatch + MC sampling) stochastic implementation for SR-VN, The results clearly demonstrate that the stochastic approach is much slower compared the piecewise method. Additionally, while the stochastic implementation converges to slightly higher ELBO, the piecewise implementation reaches a more accurate solution more rapidly. These findings are based iterates averaged over different initializations. The dataset laken is the displayed in Figure taken with different n's.
F Experiments Detail
For all experiments, we first grid search to tune model hyper-parameters, where the search is performed in specific range of values. The resultant values were then fixed during experiments. The statistics of the datasets and the model hyper-paramneters used are given in Table 2 In practice, implementing deterministic NGD is often feasible by employing piecewise analytical linear and quadratic bounds to compute expected values. For example, please refer to (Marlin et al. 2011]. Our experiments adopt this technique and replicate the setup outlined in (Khan Lin 2017).
Table 2: Dataset Statistics and Model Hyperparameters
Step Sizes
Dataset
3
N
d
Neroin
VN
SQ-VN VN
BW-GD
Australian-scale
10-5
690
14
552
5 X 10
5 5 10-3
4.4 x 10
Diabetos-scale
10-2
768
614
5 X 10
5 x 10-3
9 10
breast-cancer
10-1
683
10
546
9 x 10-3
6 10-3
6.3 10-3
Mushrooms
10-2
8124
112
6499
2.5 x 10
2.5 x 10-4
8.5 x 10-5
Phishing
10-2
11055
68
8844
4 X 10
x 10
8 x 10
MNIST
10-1
70000
784
60000
5 x 10-6
5 x 10-6
10-6
Covtype-scale
10
581012
54
500000
10-5
10-5
10 10-6
Leukemia
2 x 10
38
7129
34
X 10
5 x 10-6
1.5 10
28

---

Published in Transactions on Machine Learning Research (05/2025)
F.1 Runtime comparison
It is worth noting that each algorithm differs in computational complexity not because of the training data whhhh is the same across all-but due to its specific update operations after computing the expected gradient g and Hessian H. Since these quantitics cost the same for cach method, we focus on the subsequent matrix/vector steps.
Runtime comparison synthetic binary classification fixed sample-size (n=2000) and increasing dimensions (d) Theoretical Runtime: olfid)
Example: n, (2000,22 2)
Training Runtime
VN
SR-VN
10th
BW-GD
log (Runtime (sec))
logifld)
Feature
10:2
Class
Feature
1000 1500
2000
1000 1500
2000
(dimensions)
(dimensions)
Figure 5: We compared the run times of all three algorithms across increasing dimensions while keeping the number of data points (n) fixed. The run times were averaged over 1000 iterations per algorithm. The theoretical run times, derived in rebuttal answer, functions of and are plotted reflect the floating point operations. The results indicate that the run times can vary significantly among the algorithms, depending on the specific data and dimensionality.
1. Complexity for VN:
S update: scalar multiplication + matrix addition + O(2d).
update: matrix inversion Odd)), plus matrix-vector product Ocd)), plus matrix addition O(d)), for a total O(2d2 + dd).
Overall cost: Odd + 4d2).
2. Complexity for SR-VN:
C update: includes C HCC (O(d)) plus some matrix additions//ubbtractions (O(d)) and multiplications (O)))).
update: CCT multiplication (O(d)) plus matrix products (O(d))).
Overall cost: d3 + d + g)
3. Complexity for BW-GD:
m update: scalar multiplication + matrix subtraction - O(37).
M update: two matrix additions + scalar operations - O(2ff)
V update: matrix inversion + (dd)) plus two matrix multiplications + O(2d)), totaling (3d3))
Overall cost: (3d 5d
These matrix/vecton operations dominate the wall-clock time for each algorithm. To illustrate how runtime scales in practice, we compare VN, SR-VN, and BW-GD in Figure 5 as the dimension grows, highlighting their differing execution costs.
29

---

Published in Transactions on Machine Learning Research (05/2025)
F.2 Algorithmic Details and Additional Results
In Table provide descriptive statistics for different metrics achieved by each algorithm when it is for every dataset considered.
Table 3: Performance Comparison on different Datasets
Dataset
Algorithm
Train ELBO
Test NLL
Test Accuracy
Time (sec)
VN
196.82 82
54.61
0,877
0.53
Australian-scale
BW-GD
196.82
54.61
0.877
0.53
SQ-VN
196.82
54.61
0.877
0.53
VN
301.18
79.72
0.74
0.7
Diabetes-scale
BW-GD
301.18
79.72
0.74
0.49
SQ-VN
301.88
79.72
0.74
0.53
VN
52.86
13.62
0.956
0.61
Breast-cancer
BW-GD
52.86
13.62
0.956
0.61
SQ-VN
52.86
13.62
0.956
0.61
VN
57.51
6.69
1.0
1.04
mushrooms
BW-GD
78.41
7.02
0.999
1.04
SQ-VN
53.56
6.43
0.999
1.05
VN
339.41
727.96
0.706
0.36
Leukeumia
BW-GD
564.65
796.76
0.676
0.38
SQ-VN
339.38
727.99
0.706
0.37
VN
1267.77
334.19
0.938
0.99
Phishing
BW-GD
1377.645
349.10
0.933
0.97
SQ-VN
1268.06
333.37
0.939
0.98
VN
6317.52
1017.64
0.989
0.22
MNIST
BW-GD
6367. .323
1102.34
0.988
0.23
SQ-VN
6301.87
1012.37
0.989
0.2
VN
27066.64
10741.37
0.745
0.95
Covtype-scale
BW-GD
28036.23
11259.20
0.714
0.98
SQ-VN
27065.02
10740.73
0.745
0.91
30

---

Published in Transactions on Machine Learning Research (05/2025)
F.3
Iteration Plots with Accuracy
Breast-cancer-scale
1.0
VN
Training ELBO
10a
SR-VN
10a
Test Accuracy
0.8
Zoomed
BW-GD
Test NLL
10a
0.95
10-1
0.6
10-1
0.90
0.4
10-1
10-3
0.85
0.2
1045
0
2000
4000
0
2000 Iteration
4000
2000
4000
Iteration
lteration
Australian-scale
10a
10a
0.85
Training ELBO
Test Accuracy
Zoomed
10a
Test NLL
109
0.80
0.89
10
0.75
0.86
10
0.84
10-3
0.70
10-4
20
875
1730
2585
0.65
1000
2000
3000
1000
2000
3000
1000
2000
3000
Iteration
Iteration
Iteration
Diabetes-scale
0.7
Training ELBO
101
Test Accuracy
Test NLL
0.6
10a
0.76
10-1
0.5
0.74
0.72
10-3.
0.4
0.70
100
0.68
0.3
0
2000
4000
2000 Iteration
4000
2000
4000
Iteration
Iteration
Mushrooms
10a
1.0
105
0.9
Training ELBO
Test Accuracy
Zoomed
104
Test NLL
103
103
0.7
0.8
10a
10a
0.6
0.6
10a
0.5
2470
4920
7370
2000
4000 6000
8000
2000 4000 6000
8000
0
2000
4000
6000
8000
Iteration
Iteration
Iteration
Figure 6: Plots with respect to the number of iterations for small-scale datasets.
31

---

Published in Transactions on Machine Learning Research (05/2025)
MNIST
1.0
10
10
Training ELBO
Test Accuracy
Zoomed
104
Test NLL
10a
10a
0.94
10a
0.4
0.92
102
0.2
0.90
10
20
510
755
500 Iteration
1000
500 Iteration
1000
0
5.00
1000
Iteration
Covtype-scale
104
0.75
Training ELBO
104
Test Accuracy
0.70
Zodmed
Test NLL
0.65
10
0.70
0.60
103
0.55
0.65
102
0.50
875
1730
2585
1000
2000
3000
0
1000
2000
3000
1000
2000
3000
Iteration
Iteration
Iteration
Phishing
10a
104
0.9
Training ELBO
104
Test Accuracy
Zoomned
Test TestNL NLL
103
0.8
103
102
0.7
102
10a
0.6
0.6
20
1980
3940
5900
2000
4000
6000
2000
4000
6000
2000
4000
6000
Iteration
Iteration
Iteration
Leukemia
10a
0.7
Training ELBO
102
Test Accuracy
Zoomed
Test NLL
0.6
10a
0.68
10a
0.5
0.64
10
0.62
0.4
636
1252
1868
1000
2000
0
1000
2000
0
1000
2000
Iteration
Iteration
Iteration
Figure 7: Plots with respect to the number of iterations for large-scale datasets.
32

---

Published in Transactions on Machine Learning Research (05/2025)
F.4
Time Plots with Accuracy
Australian-scale
10a
10a
0.85
Training ELBO
Test Accuracy
10a
Test NLL
10
0.80
10-1
0.75
10 2.
10 3.
0.70
0.65
0
5
10
15
0.0
0.1
0.2 seconds
0.3
0.0
0.1
0.2
0.3
seconds
seconds
Breast-cancer-scale
1.0
103
10a
Training ELBO
Test Accuracy
0..8
Test NLL
10a
10-1
0.6
10-1
0.4
10-4
10-3
0.2
5
10
15
0.0
0.1
0.2
0.3
0.4
0.0
0.1
0.2
0.3
0.4
seconds
seconds
seconds
Diabetes-scale
0.7
Training ELBO
10a
Test Accuracy
Test NLL
0.6
10a
10-1
0.5
10
0.4
100
0.3
0
20
40
2
2
seconds
seconds
seconds
Mushrooms
10
1.0
105
0.9
Training ELBO
Test Accuracy
104
Test NLL
10a
0.8
103
0.7
10a
0.6
0.5
10
100
200
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
seconds
seconds
seconds
Figure 8: Plots with respect to time for small-scale datasets.
33

---

Published in Transactions on Machine Learning Research (05/2025)
MNIST
1.0
10a
10
Training ELBO
Test Accuracy
0.8
104
Test NLL
10a
0.6
10a
102
0.4
102
0.2
10
200
400
600
0.0
0.1
0.2
0.0
0.1
0.2
seconds
seconds
seconds
Covtype-scale
104.
0.75
Training ELBO
104
Test Accuracy
0.70
0.65
10a
Test
0.60
103
0.55
102
0.50
500 seconds
1000
0.00
0.25
0.50
0.75
0.00
0.25
0.50
0.75
seconds
seconds
Phishing
10a
104
0.9
Training ELBO
Test Accuracy
104
Test NLL
103
0.8
103
10
0.7
102
10
0.6
100 seconds
200
300
0.0
0.2
0.4
0.6
0.8
0.0
0.2
0.4 seconds
0.6
0.8
seconds
Leukemia
1
0.7
Training ELBO
102
Test Accuracy
Test NLL
0.6
10a
10a
0.5
10
0.4
500
1000
0.0
0.1
0.2
0.3
0.0
0.1
0.2
0.3
seconds
seconds
seconds
Figure 9: Plots with respect to time for large-scale datasets.
34

---

Published in Transactions on Machine Learning Research (05/2025)
G Comparing Geometries: SR-VN vs. Alternatives
In this section, we discuss two complementary comparisons for SR-VN: (1) a theoretical analysis of its convergence rate relative to BW-GD, and (2) an empirical evaluation against an enhanced method tailorod for the BW geometry.
Thcoretical rates comparison, SR-VN BW-GD: To estimate the convergence rate for BW-GD, refer to (Lamberl et al. 2022b Theorem 4), which under the assumption of strong log-concavity provides exponential rate in the Wasserstein distance once we ignore the variance term because of its deterministic nature. Consequently, in terms of asymptotic guarantocs, BW-GD and SR-VN both exhibit exponential convergence rates. However, the specific details namely, which geometry (Fisher Rao vs. Wasserstein) is deployed, and which measure (KI. parameter space vs. Wasserstein distance probability space) is analyzed separate these methods apart. As result, establishing a theoretical comparison challenging, leading to focus experimental evidence in the next section.
Additional experimental comparisons. Although BW-GD is a natural baseline for SR-VN in the BW geometry, more recent studies have proposed improved and more stable stochastic BW methods (c.g., Diao
g0 Liu et al. (2024)). We focus here on Liu el al. (2024), which uses SVGD and has demonstrates better performance than Diao et al. (2023). Specifically, we implemented a deterministic version of (Liu ett al.] 2024 Algorithm 1) i.e., replacing sample estimates with the full expectation and refer to this
et al.
method SVGD-density. It is worth mentioning that Liu et al. (2024) demonstrated improvements in the stochastic setting compared to BW-GD. However, since we focus on the deterministic implementation here, it
remains unclear whether the same advantages apply this is precisely what we aim to address next.
We evaluated SVGD-density (using various step sizes) on the same 2D Bayesian linear regression problem from Figure 1 as shown Figure 10 (upper). Additionally, we tested it on the Bayesian Logistic Regression scenario in Figure as shown Figure (lower). both plots, SVGD-density shows the slowest convergence among the compared algorithms. Although it can approach the performance of BW-GD, given finely tuned step size e.g. in Figure 10 (upper), its performance is generally worse compared to SR-VN or VN.
SVGD-censity (eps1-1c-07, cps2-1c.07)
105
SVGD-Gensity (eps) 04e074 opc2-40-07)
DiGD-density (eps]=6e-077 eps2=6e-07)
NLL (log-scale)
SRAVN
104
103
10a
2000
4000
6000
8000
10000
Iteration
Breast-cancer-saale
1.0
10a
SR-VN
10a
Training ELBO
BW-GD
Test Accuracy
0.8
SVGD-density
Test NLL
102
10
0.6
10
0.4
109
10-3
10-1
0.2
1000
2000
3000
0
1000
2000
3000
0
1000
2000
3000
Iteration
Iteration
Iteration
Figure 10: Comparison of SVGD-density: Bayesian linear regression (top) and Bayesian logistic regression (bottom).
35

---

Published in Transactions on Machine Learning Research (05/2025)
H Non-convex Model
In this section, we evaluate all algorithms on the same Bayesian logistic regression loss in Equation 15 with the addition of non-convex regularizer Σ Ei=m mē/(1 + mz).
The results in Figure 1 indicate that all compared methods still converge despite the non-convexity introduced in the mean parameter, suggesting promising directions for further investigation into non-convex scenarios.
Breast-cancerscaae
103
10
102
SR-VN
0.9
Training ELBO
BW GD
10
Test Accuracy
0.8
10
Test NLL
100
0.7
10
10-1
0.6
10
10
0.5
10-2
0.4
0
1000
2000
3000
0
1000
2000
3000
1000
2000 Iteration
3000
Iteration
Iteration
Figure 11: To evaluate performance under non-convex conditions, augment the Bayesian logistic regression problem Equation (15 with non-convex regularization term: Σi= mf/(1 + mf).
36


Reviews:
[{"review_id": "yJ62hxd9zI", "reviewer": "Reviewer_TQto", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* First rigorous convergence guarantee for discrete-time variational NGD\n\nWeaknesses: \n* The main new update, eq 14, is already published by Tan (2025)\n* The assumption of convex loglikelihood limits applicability, excluding most ML settings of interest\n* The story is incomplete: SR-VN makes the problem convex but VN still performs about as well. Why?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies NGD for gaussian variational inference. It uses the previously establisehd fact that for a concave loglikelihood, the variational objective is convex wrt square root parameterization of the covariance, but not wrt to the covariance itself or the precision. The main contributions are to apply convex optimization tools to the square-root parameterization to prove convergence for NG flow and for NGD discretized wrt the Cholesky factor.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I think the paper would benefit from an experiment on a non-convex model. There’s value in papers that prove theoretical guarantees in restricted settings and test empirically whether performance tends to extend beyond those assumptions. On the other hand a paper that only addresses convex models is rather limited.\n\nParameterization does not matter for NG flow. The introduction is careful on this point, only stating the algorithm depends on the parameterization in the discretized case, and the end of sec 3.3 explains the square root parameterization is used in the flow case only as a proof technique since it allows application of convex analysis. Still it could improve clarity to be more explicit about this in the introduction.\n\nAfter eq 1, the description of $\\gamma$ seems to be reversed. Setting $\\gamma$ too high should cause the posterior to collapse to the prior. Setting $\\gamma$ too low should produce no regularization and hence excess sampling variance.\n\nThe remark before eq 5 that $\\boldsymbol{\\tau}$ is the natural parameter is unnecessary. The equation holds regardless. On the other hand the choice of natural parameters is essential for eq 7.\n\nSec 3.2: It may help to highlight that the Fisher of $q_{\\boldsymbol{\\tau}}(\\boldsymbol{\\theta})$ used here is different from the Fisher of $p(\\boldsymbol{y}|\\boldsymbol{\\theta})$ used in non-variational NGD (especially since terms related to the latter appear in eq 6).\n\nThe statement that eq 8 is a continuous-time version of eq 5 holds only when $\\gamma=1$.\n\nThe statement after eq 9 could be made more precise: “Upon discretization with respect to the natural parameters, this results in the VN update with $\\gamma=1$.”\n\nSec 3.5: It may help to point out that, in addition to its connection to Newton, eq 7 with $\\gamma=\\rho=1$ is the exact Bayes update in the conjugate case, while discretizations based on $\\boldsymbol{V}$ or $\\boldsymbol{C}$ are only approximations.\n\nEq 15: missing $-$ and $)$\n\nAfter eq 15: The Hessian is $\\boldsymbol{X}^\\top \\boldsymbol{D} \\boldsymbol{X}$ with $\\boldsymbol{D}\\in\\mathbb{R}^{n\\times n}$.\n\nsame paragraph: the loss is $\\beta$-strongly convex, not $\\delta$-strongly\n\nAssumption 3: The citation should be to Domke 2020 and the premise should be that $\\bar{\\ell}$ is $M$-Lipschitz-smooth\n\nAssumption 4 for logistic regression: It’s not obvious to me that the bounds on $\\boldsymbol{H}$ imply bounds on $\\boldsymbol{C}$. The reasoning \"This validates the assumption of a bounded square root” is too casual since $\\boldsymbol{C}$ is not the square root of $\\boldsymbol{H}$ but instead is the result of the update in Eq 14.\n\nSec 6.1 NLL definition: sum should run to $n_{\\rm test}$ and missing $-$ sign\n\nCovtype-scale: extra commas in $n$ and $n_{\\rm train}$\n\n\"Our findings indicate that the success of these methods is heavily influenced by the choice of parameterization”: I don’t see this in the results. The parameterization matters for proof techniques but does not necessarily matter for performance. The experiments mainly show the natural and square-root parameterizations perform comparably, except in the conjugate case where natural parameterization gives the exact Bayes update (fig 1c).", "broader_impact_concerns": "none", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "fMcuDSK9zL", "reviewer": "Reviewer_Cj54", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strenghts:\n- Background on Variational Inference and NG methods is well explained\n- In general, explanations of the paper are clear and intuitive\n- Square root Gaussian parametrization is an elegant solution to address the lack of concavity issue\n- Theoretical analysis is rigorous and provides interesting bounds on SR-VN\n- Paper is clearly organized - e.g. background, assumptions, main results, experiments - which makes the whole proof and reasoning easy to follow\n\nWeaknesses:\n- No proper literature review - section 2 cites briefly some previous work done in the general area but I would have wanted a much clearer review of what has been done in this specific topic\n- Fairly limited discussion section - I would have expected much more discussion on the experiment result as well as on the strenghts of the assumptions made in section 4", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper provides a theoretical analysis of the convergence of Natural Gradient methods for Variational Inference, in the context of a Gaussian square root parametrization, i.e. Variational Inference with a Gaussian distribution and using a Cholesky decomposition of the covariance matrix. This setting provides significant advantages, e.g. maintaining concavity of the objective which allows to use tools from convex optimization for the analysis.\nAfter introducing some background on Variational Inference and Natural Gradient methods (Natural Gradient Descent and its continuous form Natural Gradient Flow), the paper presents its main algorithm, SR-VN (Square Root Variational Newton), using this Cholesky decomposition to derive a new update for a Newton method. It also  provides theoretical guarantees of its exponential convergence in the case of a strongly concave log-likelihood. It finally presents some empirical results on Bayesian Logistic Regression problems.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See weaknesses section.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "XVoXoigauq", "reviewer": "Reviewer_442L", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n\nThe paper studies theoretically the convergence of Natural Gradient Descent Variational Inference algorithm. It provides new insights on the convexity of the problem, shows the convergence of the Natural Gradient Flow, and provides a new discretization of the flow with a parametrization using the mean and the square root of the covariance, which is shown to converge. This new findings are very interesting, and provide a nice justification of the convergence of Variational Inference algorithm with Natural Gradient Descent. Moreover, the paper is very clear and pleasant to read.\n\n\n**Weaknesses:**\n\nA little weakness might be the choice of baselines in the Bayesian Logistic Regression problem. While it makes senses to compare the algorithm with the Bures-Wasserstein gradient descent, there are some more recent works which focused also on this problem. Also, I would be curious to see a comparison of the theoretical guarantees for the Bures-Wasserstein gradient descent with those derived in this work.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Variational Inference is a common technique to approximate posterior distribution. However, it is not theoretically well understood. This paper focuses on Variational Inference with the approximation belonging to the Gaussian family, and studies a Natural-Gradient algorithm to minimize the Kullback-Leibler divergence over this family. It is shown that while the minimization problem is not convex in the mean and covariance even when the target distribution is log-concave, it is with respect to the mean and the square root of the covariance. Leveraging this observation, the authors show that the Natural Gradient Flow converges exponentially fast, showing that the Kullback-Leibler divergence satisfies a Riemannian PL inequality. Moreover, they propose a discretization of the Natural Gradient Flow on the mean and the square root of the covariance. Then, they show that this algorithm converges towards the minimum under reasonable assumptions (strong convexity and smoothness of the loss). Finally, they apply it to a Bayesian Logistic Regression problem, and show that the proposed algorithm has the same performances as the Natural Gradient Descent Variational Inference algorithm, and better performances than the Bures-Wasserstein Gradient Descent.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "These are not change which would be crucial for securing my recommendation for acceptance, but it might strengthen the work.\n\n- Since you compare with BW-GD, it would be a nice addition to mention the convergence rates known for this method, and compare it with the convergence rate derived in Theorem 1 and 2.\n- On the applications, it could be nice to have some comparison on a simple Gaussian target (in the same spirit as Figure 1, but e.g. also plotting the convergence), additionally to the Bayesian Logistic Regression.\n- In relation with the weakness I mentioned, it could be nice to compare with or at least mention [1], which proposes a scheme to solve this problem which is more stable than BW-GD. I do not expect it to have better results though. Another work, which could be pertinent to mention and/or compare with, is [2], which uses SVGD. They also applied their scheme to Bayesian Logistic Regression, and demonstrated better results than [1].\n- Less related, but note that there are some works trying to improve the convergence of Wasserstein gradient descent by changing the geometry of the space, see e.g. [3].\n\n\n[1] Diao, M. Z., Balasubramanian, K., Chewi, S., & Salim, A. (2023). Forward-backward Gaussian variational inference via JKO in the Bures-Wasserstein space. In International Conference on Machine Learning (pp. 7960-7991). PMLR.\n\n[2] Liu, T., Ghosal, P., Balasubramanian, K., & Pillai, N. (2024). Towards understanding the dynamics of gaussian-stein variational gradient descent. Advances in Neural Information Processing Systems, 36.\n\n[3] Cheng, Z., Zhang, S., Yu, L., & Zhang, C. (2024). Particle-based variational inference with generalized wasserstein gradient flow. Advances in Neural Information Processing Systems, 36.", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Reviewers unanimously recommend acceptance. Please see above for motivation of request for minor revision.
Decision: Accept (poster)