Natural Gradient VI: Guarantees for Non-Conjugate Models
Authors: Fangyuan Sun
Ilyas Fatkhullin
Niao He
Venue: NeurIPS 2025 poster ()
Keywords: Variational inference
natural gradient
mirror descent
non-convex optimization
Primary area: optimization

Abstract: Stochastic Natural Gradient Variational Inference (NGVI) is a widely used method for approximating posterior distribution in probabilistic models. Despite its empirical success and foundational role in variational inference, its theoretical underpinnings remain limited, particularly in the case of non-conjugate likelihoods. While NGVI has been shown to be a special instance of Stochastic Mirror Descent, and recent work has provided convergence guarantees using relative smoothness and strong convexity for conjugate models, these results do not extend to the non-conjugate setting, where the variational loss becomes non-convex and harder to analyze. In this work, we focus on mean-field parameterization and advance the theoretical understanding of NGVI in three key directions. First, we derive sufficient conditions under which the variational loss satisfies relative smoothness with respect to a suitable mirror map. Second, leveraging this structure, we propose a modified NGVI algorithm incorporating non-Euclidean projections and prove its global non-asymptotic convergence to a stationary point. Finally, under additional structural assumptions about the likelihood, we uncover hidden convexity properties of the variational loss and establish fast global convergence of NGVI to a global optimum. These results provide new insights into the geometry and convergence behavior of NGVI in challenging inference settings.

Natural Gradient VI: Guarantees for Non-Conjugate Models
Fangyuan Sun
Iyas Fatkhullin 1,23+
Niao He
*Department of Computer Science, ETH Zurich
2ETH Al Center
3 Industrial and Systems Engineering, Georgia Institute of Technology ch, niao. heβinf .ww.f.eehecccmmm ch
fansunōstudent. ethz. ch, ilyas. fat.hu.lthhlllll
Abstract
Stochastic Natural Gradient Variational Inference (NGVI) a widely used method for approximating posterior distribution in probabilistic models. Despite its em- pirical success and foundational role variational inference, its theoretical un- derpinnings remain limited, particularly in the case of non-conjugate likelihoods. While NGVI has been shown to be a special instance of Stochastic Mirror Descent, and recent work has provided convergence guarantees using relative smoothness and strong convexity for conjugate models, these results do not extend to the non- conjugate setting, where the variational loss becomes non-convex and harder analyze. In this work, we focus on mean-field parameterization and advance the theoretical understanding of NGVI in three key directions. First, derive suffi- cient conditions under which the variational loss satisfies relative smoothness with respect suitable mirror map. Second, leveraging this structure, we propose a modified NGVI algorithm incorporating mon-Euclidean projections and prove its global non-asymptotic convergence stationary point. Finally, under additional structural assumptions about the likelihood, we uncover hidden convexity prop- erties of the variational loss and establish fast global convergence of NGVI to a global optimum. These results provide new insights into the geometry and con- vergence behavior of NGVI in challenging inference settings.
Introduction
Variational Inference (VI) is powerful framework for approximating Bayesian posteriors by casting inference as an optimization problem [JGJS99, BKM17]. Unlike sampling-bais based approaches such Markov chain Monte Carlo (MCMC; [MRR + 53]) VI enables scalable posterior approximation through stochastic optimization, often with significant computational advantages. In practice, how- ever, the exact gradient of the variational objective- such the evidence lower bound (ELBO)- rarely available in closed form, especially for complex or non-conjugate models. Consequently, gradients must be estimated from data using Monte Carlo sampling. Since the advent of Black- Box VI [WW13, 13, TLG14, RGB 44, this form of gradient-basee stochastic optimization in Euclidean parameter spaces has become the de facto standard in practical applications.
While gradient descent in the Euclidean space of parameters has become the standard tool in modern VI, it ignores the intrinsic geometry of the space of probability distributions. This has motivated the development Natural Gradient Variational Inference (NGVI) [HTRK07], which replaces stan- dard gradients with natural gradients [Ama98] that account for the underlying information geom- etry of the variational family. NGVI follows the steepest descent direction (on average) with re- spect to the Kullback-Leible (KL) divergence, rather than the Euclidean norm. Empirical evidence
*.SSS is supported by ETH AI Center, LE, is funded by ETH AI Center Doctoral Fellowship.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

step s8E size-05
step siso-0,5
Fraeiie Value Ceo
10"
Penetion Value Gag
-0P
stor-
step sizp=0.1
slpp vinomO.
00-1
1011
10ml
10-3
40
Eeeration
Iteration
Figure 1: Convergence of SNGD on Poisson regression with different initialization. The function value gap refers to the difference between the objective value at iteration and the optimal value; performance averaged over 10 runs for each algorithm. In the left panel, SNGD is initialized with of = 2. We observe instability at the initial iteration and slow convergence with step sizes 0.5 and 0.3. In contrast, with a different initialization of 0.4 shown in the right panel, we observe stable behavior and fast convergence for larger step sizes.
suggests that this geometry-aware approach leads to faster more stable convergence in various applications, including Bayesian neural networks [ZSDG18, OSK 19] and probabilistic filtering (SEH18, LZM 25]. However, theoretical understanding of NGVI remains limited particularly in the case of non-conjugate models, where the posterior is not contained in the variational family and the optimization landscape becomes non-convex.
Recent work [WG24] established the first non-asymptotic convergence guarantees for stochastic natural gradient descent (SNGD), basic variant of NGVI, under conjugate likelihoods, by framing SNGD as special case of stochastic mirror descent (SMD; [NY83]). Their analysis hinges on the strong convexity and smoothness of the variational objective in the Bregman geometry induced by KL divergence. Unfortunately, these guarantees fail to extend to non-conjugate models such as logistic Poisson regression, where the objective becomes non-convex and the geometry is not well-behaved. For instance, Figure indicates that 1-smoothness of the objective as in conjugate models does not hold globally (see Section B for more detailed explanation). This observation serves as a motivation for our subsequent analysis of relative smoothness non-conjo non-conjugate models.
Contributions. This work develops theoretical foundation for stochastic natural gradient varia- tional inference (NGVI) the non-conjugate setting. Our main contributions are as follows:
Relative Smoothness of the Variational Objective. We derive sufficient conditions under which the negative ELBO is smooth relative to the Bregman geometry induced by the KL divergence on compact set of parameters. Our results apply to broad class of non-conjugate models in mean-field field parameterization, including logistic regression, and provide explicit smoothness constants (polyno- mial in the problem dimension) over arbitrary compact subsets of the parameter space.
Projected Stochastic NGVI Algorithm. Based the relative smoothness analysis, introduce a projected variant of NGVI called Proj-SNGD that enforces updates within a compact parameter domain using non-Euclidean projections. This approach leverages the dual structure of exponential families and admits efficient implementation in the mean-field Gaussian case. The resulting algo- rithm improves numerical stability while preserving the geometric fidelity of the natural gradient updates.
Convergence Guarantees via Mirror Descent. We analyze Proj-SNGD through the lens of stochastic mirror descent (SMD), extending recent non-convex convergence theory [FH24]. We establish asn-asymptotic convergence to stationary point at rate O (1/VTT), and further show that under concavity of the log-like likeliooooo hidden convexity structure [FHH25] allows us to prove global convergence at rate (1/T).
By extending convergence guarantees to the non-conjugate regime, this work addresses a fundamen- tal challenge in the theory of NGVI and provides principled framework for variational inference in complex, non-linear Bayesian models.

---

1.1 Related Work
In the Euclidean setup, A large body of prior work provides convergence guarantees of standard VI algorithms. Building on smoothness and convexity properties established in [TLG14, Dom20J, authors of [DGG23] propose two gradient descent- based algorithms that achieve convergence rates of (11/VT) using convexity, and (1/TT) using strong convexity. Concurrently, Kim et al. [KOW 23] analyze the convergence behavior of gradient-based VI under various parameteri- zation and prove similar convergence bounds. Furthermore, a linear rate in the case of conjugate likelihoods is established in [KMG24]. In the non-Euclidean setup, Wu and Gardner [WG24] es- tablish the first ((1/T) convergence rate of NGVI assuming the model conjugate. For general overview on natural gradient descent, we refer to the survey by Martens [Mar20].
Since its introduction in [Ama98], natural gradient descent has been extensively studied in the context of variational inference [HRL12. HBWPI3, LT21], particularly for non-conjugate mod- els [KBL+ 16. KL17, SEH18, TR19]. Notably, [LKS19] and [ADY+ 23] extend natural gradient methods to mixtures of exponential- -ammly variational distributions, while [CAAK19] and [JCM24] online learning scenarios. In reinforcement learning, a series of works has focused on natural policy gradient [Kak01, AKLM21, Xia22, SLZ*23J. Natural gradient vari- ational inference (NGVI) has also been applied across various domains, including the training of variational Bayesian neural networks IZSDGI8, OSK*19), Kalman filtering [LZM 255, and multi- modal optimization [MAM+25] 25].
focus on their application
There exists an extensive body of work on the convergence of stochastic mirror descent (SMD) under convexity assumptions [NJLS09]. Later, [Lan12, AZO17] show that mirror descent can be accelerated, similar to Nesterov's method. Relative smoothness, which was introduced concurrently in [LFN18, BBT17], has become an important tool analysis of SMD [HR21, VYB+22) In the non-convex setting, [GLZ16] derive an (1/VTT) convergence rate to a stationary point, albeit un- der very strong assumptions and with large
This requirement is relaxed by [ZH18] and (DDM18), who establish similar convergence guarantees without relying on mini-batching. More recently. Fatkhullin and He [FH24] further relax the smoothness assumptions on the distance gen- erating function using a distinct proof technique and improved convergence criterion. Non-convex SMD has also been studied in the context of reinforcement learning, where the value function is typ- ically highly non-convex but admits certain structures analogous to our problem, see, c.g., [Lan23]. Additionally, alternative variance assumptions for SMD are explored by [Hen24].
2 Background
Notation. We write || for the Euclidean norm of a vector or the operator norm of a matrix, and ( to denote the standard inner product in Euclidean space. For vectors a, E Rd, let @ b denote their entrywise (Hadamard) product. Let St denote the cone of d symmetric positive definite matrices. For symmetric matrices A, B E Raxd we write A > B if A B E S For distributions p and g. we write (x ([p) = (gtz) log(g(=)/p(a) for the Kullback-Leibler divergence from p to q. When needed, we write (2; 0) to emphasize the parameterization of the distribution.
2.1 Variational Inference and Exponential Families
Let Rdd be latent variable and D {I 34) y be the observed dataset. Given a prior p(2) and likelihood model p(D z), the goal of Bayesian inference is to compute the posterior p(= (D) xx p(z)p(D) [2). Since this posterior is generally intractable, variational inference (VT) approximates by minimizing the KL divergence from p(= (D) to a tractable distribution g(z) within a family Q:
min DKL (g(=) ([ q
This is equivalent to minimizing the negative evidence lower bound (negative ELBO) [BN06]:
{()) = - (og =)] + DKI(((()) ll((=))
In this work, we assume Q is minimal exponential family [W]+ i.e., each q E Q has the form n(=;n) = 11()) exp ((0(=). A(r)),
3

---

where is the natural parameter, ((2) is the sufficient statistic, and A(7) is the log- pege-pattiion function. The dual expectation parameter is defined as := Ex(z))) (=)=..
Let DA = {np E R A(n) < oo) be the natural parameter domain, and let 52 be the of corre- sponding expectation parameters. The gradient map VA DA - f2 is bijective, with inverse VA", where A A"(w) := De ((7, w) - A(n)) is the convex conjugate of the log-partition function. This duality ensures that for any distribution y E Q, parameterized by either vi w, we have
V =
VA((w)= 17.
For a comprehensive overview of exponential families and duality, we refer readers to [W]+08].
Remark. With a slight abuse of notation, we denote the negative ELBO as E(n) == {(((2;))) or ((w) = E(g(2; m)) when working in either parameterization.
Example: Gaussian Family. For the multivariate Gaussian family 4 = Mp, E), the natural parameter is 1) (A, A) with x Σー'm p and A - - and the expectation parameter is (E, E) with μ and Σ + unr". The standard parameterization (μ, Σ) and Cholesky form (n, C)) (C' is the Cholesky factor of )) are often used in practice.
2.2 NGVI as Stochastic Mirror Descent
Stochastic mirror descent (SMD) generalizes standard SGD to Euclideld geometries by replac- ing the squared Euclidean distance with a Bregman divergence.
Definition 2.1 (Bregman Divergence). Given a strictly convex and differentiable function + u - IR, the Bregman divergence between u, U is (4( == 4(u)) Φ(v) (Vo(v), v).
The update rule for SMD on differentiable objective l:ul + Risi in
primal form:
Ue+11 argmin 7: (VE(ue), u) + ue) wEUl
equivalently in
(2)
(1(41) = (4(u) - V((c),, uet1 = argmin Dz (u, D(41), dual form:
(3)
wEll
where Vl(u.) is a stochastic gradient and Tt is a step size.
In variational inference with exponential families, the natural parameter admits geometry gov- erned by the Fisher information matrix I(n) = v A²)(,,, where A(r) is the log-partition function.
Stochastic NGVI in this space preconditions the gradient by the inverse Fisher matrix [RM15]. Ap- plied to the variational objective f, the update becomes:
M 1h Imn) v((n) = 7t V((we),
(4)
where Wi AA(T)) is the corresponding expectation parameter, see Section C.2 for details. Using the duality between and w, and the identity y (A*(w), can re-write the update in the expectation parameter space as
SNGD: VA' V (w:) (w:)
74 Ve(w)).
This matches the update rule of stochastic mirror descent with mirror map A". Hence, NGVI can be viewed as mirror descent over the expectation parameter w, with geometry induced by KL diver- gence. Indeed, for exponential families, the Bregman divergence associated with A" coincides with KL divergence: D (w, w²) = DKL (g(zw( (((2="))) see e.g., [WG24].
2.3 Mean-field Parameterization
In the following, assume that the prior is standard Gaussian, p(2) (0, I). and the variational family @ is the mean-field Gaussian family, with mean μ = (44)ςısa and diagonal covariance matrix Σ diag( (o) sissd) The corresponding expectation parameter ç = μ and = Σ + di @ μ)) defined over the domain
(E. E) is given by
Ω (E. E) & Rd, E - diag(s @ S and is diagonal).

---

This parameterization is widely adopted due to its tractability and interpretability. allowing efficient coordinate-vee updates while still capturing key aspects of the posterior distribution. Moreover, often provides a favorable trade-off between computational complexity and approximation quality in high-dimensional settings. has been extensively studied in recent theoretical works, including applications to Bayesian deep neural networks [CE24], high-dimensional Bayesian linear models [CFLM23), and particle-based variational inference [DWZZ24).
3 Landscape Properties of Non-Convex NGVI
In this section, we investigate the landscape properties of variational objective ((w) with a particular focus on the properties useful for establishing non-asymptotic convergence of NGVI. First, we note that f(w) is coercive in the expectation parameter, i.c., it grows to infinity {(w) + oc when the parameters approach the boundary of the set w + 82, see Section D for more details. This is useful property which guarantees the existence of the minima of f(w). Second, it is known that f(w) non-convex w.r.t. the general non-conjugate likelihood setting. ie.. when likelihood does not belong to the variational family Q, see Section E for details. This non-convexity is the main challenge for establishing -onassyptooicc convergence of NGVI. In what follows we aim use modern tools from non-convex optimization, which will allow us to provide better understanding of NGVI landscape and characterize its convergence. The rest of the section organized as follows. In Section 3. we will prove that, despite non-convexity, both the lower and upper curvature of f(w) are bounded in the non-Euclidean geometry following the formalism of relative weak convexity and relative smoothness [LFN18]. In Section 3.2, we uncover the hidden convexity properties of the objective and connect them with us show fast convergence of NGVI despite non-convexity.
(PL) condition [Pol63, Loj63], which allows
3.1 Relative Smoothness of Variational Objective
We start with the definition of a-8 relative smoothness [BBT17, LFN18), which is generalization of smoothness and weak convexity to Bregman geometry,
Definition 3.1. Let & : U + R be differentiable and strietly convex function
conver set A differentiable function is said to be a-B smooth relative to 4 for some a E R, > 0 Of
aD4(U, u) S {(v) E(u)= (VE(u), v- u) S (D,(u.)). Vu,eeu.
a 0, is also called a-strongly convex relative 4.
In our problem we will mostly deal with the case of negative curvature < 0. Then if 8 -a, we refer to {(-) as being L-smooth relative to b or relative smooth with parameter L.
According to Proposition 1.1 of [LFN18], a-B relative smoothness is equivalent to
(v) = (() = (V),
for any u U.
(5)
In NGVI setting, we hope to find conditions under which negative ELBO objective (1)
((w) E. (og 2)]
+
DKI((((=) P()))
(6)
log-likelihood term
KL divergence
L-relative smooth on w.r.t. A".
Relative Smoothness of KI. Divergence Term Since A² -
- log det diag(E @ {))
(see Section C.I), and the KI. divergence between Gaussian distributions admits closed-form
Dx1 2; g. [. (l)) = log det(E diag(E @ {)) + Tr(3),
the Hessians of the two functions coincide. Thus KL divergence term 1-1 smooth relative to A",
Relative Smoothness of Log-Likelihood Term Now we consider the log ligelikelihood term in (6). We explain the intuition for the derivation in the univariate case and state the general result high dimensional case, deferring the formal proof to Theorem 3.2.²
For simplicity, we assume that 2 = {(x, y)} only contains a single data point. there are n data points and each - E, [log p(æ,, 31/2)] is a.-33v-reattve smooth, then [og ((D)((())) z)] is En. ag 2 2 B, relative smooth, since -E, (log pDD]z)] >" E. flog p(2;, p()),))))))
5

---

In order to compute the Hessian matrix, ( If (2)] ) for some four times differentiable function f. we need the following useful lemma [Pri58, Bon64, OA09].
Lemma 3.1 (Bonnet's s and Price's Gradients). Let q(2) be the probability density function (PDF) of the multivariate Gaussian, N (1, 2), and assume f is twice continuously differentiable, then
V V Egg [Vf()),
In the univariate with standard parameters, μ 0² = apply Theorem 3.1 and chain rule to obtain
0p
d0²
[v))
マ E
( f())
[()) v²f(z)
2
00
O0²
V=
+
B.y [()=
0E = [V))
Using Theorem 3.1 and chain rule again and we get the following result. Proposition When d ler f(z) = - log 2). then the Hessian of the aigliikllhood term equals
V (EE [f(2)] fe) (() 1A
(
f(2) - (7)
( [(s)
[V'((
Moreover, is straightforward to see that
2p 0²
V'Ar(w) = σ -4
(8)
Therefore, proving a-3 relative smoothness is equivalent to finding 3 such that for all
2 (w) ( log (D (D| 3 A" (w:)
Using the same approach, we can compute the Hessian matrices for any d > 1 (see Theorems G.1 and G.2).
For any U > 1, define the bounded sets in the spaces of standard and expectation parameters:
Standard:
P = (/1 Σ) lad < U, Σ is diagonal, D- < Σ. < D, 1 1 d}, Ω ((S, (E. diag(E @ {)) E P) C Ω.
Expectation:
Then present our main result relative smoothness such bounded sets for any U and D Theorem 3.2 (Sufficient Conditions for Relative Smoothness). Let f(z) log a)) be four times continuously differentiable function in 2 on Rd. Assume sup=eRa sup V. [=)] S L,, supzeRs sup,=1,... supj=1,...d ( < L2. Then the log-likelihood term L-smooth relative A on Ω with
L= O(DDUUL + L2U) + dD L2U)).
Proof sketch. First, we use the approach mentioned above to compute the Hessian matrices V=Egte:g and (w)" To handle high-order derivatives appearing in E we apply Stein's ' Lemma (Theorem F.3) to bound them by lower-order derivatives. Finally, we find and 3 as in (9) by proving the positive semidefiniteness of the corresponding matrices.
Example 1. For logistic regression, - -log 2) = log(1 + e --vpr' ²). where E Ed is a data point and E (-1,1} the label. Since , ply) xr x) = -o(( ((-y""
-o(-yx 22yx;; and
Vi, log x, = o(-yr''"ppp 2) or (-yx" 2) ,, by writing || = max, |zg), we have
sup sup V. logp(y x 2) < || L1,
2 i=ن...,
sup sup sup
log 2) < || = Lg-
* 1=, dj=
Therefore, {(w) is smooth relative A with parameter OLDD 1| (U + D) + U(|=+)))))
A stronger statement with tight characterization for the 1 can be found in Section F.

---

Comparison to conjugate case. For conjugate models, f(w) is 1-smooth and 1-strongly convex relative A" [WG24], which makes very well-conditioned strictly (but not strongly) convex objective. For non-conjugate models including logistic regression and Poisson regression, however, is typically negative, indicating that ((w) is a non-convex objective. The relative smoothness constant in this case scales polynomially with the size of the set P and dimension d.
3.2 Hidden Convexity of Variational Objective under Log-concave Likelihood
We show that f(w) exhibits hidden convexity when the log-likelihood is concave in 2, as in logistic and Poisson regression. Formal proofs of statements in this section relegated to Section G.3.
A function has hidden convexity if it becomes convex under a reparameterization c(-). Formally: Definition 3.2 (Hidden Convexity; [FHH25]). Let & : Ω + R satisfy f(w) H(c(w)) for an invertible map e: Q + O Then is hidden convex with modulus PCC > pa 2 0 if:
1. The set θ convex and H : e + R is str-snyyggly convex w.r.t Euclidean geometry. 2 Ncllwi-w22l
2. The map is invertible and 3uc > 0 lel:>>>:::/aaaaaa////
Ywi, E Q.
The result below follows from the fact that the strong convexity of [(u) transfers to the expected value (yyie C) [f)) under the Cholesky parameterization (N, C),, where C is lower triangular with positive diagonals and CCT
[Dom20, Theorem 9].
Proposition 2. If log pDD z) is concave 2, then the restriction of ((W) Ω R on 5 is hidden convex with modulus pC (4U 4D + 1)-1/2 and рa 1.
Given hidden convexity, we establish the PL inequality based on the analysis in [FHH255]. Proposition If log 22) concave and a stationary point w" lies in Ω then w" a global minimum. Furthermore, f(w) satisfies the PI. inequality in the relative interior ri())):
[ 2 (P(ww) -8), Ve E ri({)),
with PC from Proposition 2.
(10)
The assumption that ~' E Ω is mild. Under weak conditions, E(w) is coercive: {(w) + oc as I1 | + d or de((E) + 0 (see Section D), thus for large enough U and D, Ω contains a sta- tionary point. Note that although the PL condition holds, it does not immediately imply the function value convergence of NGVI, since for non non-Euclidean algorithms the gradient norm may converge arbitrarily slow. In the next section, we introduce an additional mild assumption (Assumption 2), sufficient for the function value convergence.
4 Convergence of Non-Convex NGVI
This section analyzes convergence properties of NGVI. In Section 4.1, we introduce our projected stochastic natural gradient descent (Proj-SNGDD algorithm and prove an O(1/VT) convergence rate under the relative smoothness condition in Section 4.2. Finally, if the log- lig-liieldhood is concave, we show in Section 4.3 that Proj-SNGD achieves O(/TT) convergence to the global minimum.
4.1 Projected Stochastic Natural Gradient Descent
Given an initial point wo E Ω2 number of iterations T and step sizes (1losees-1s we introduce the update rule of our projected variant of SNGD (with wo E ②
Proj-SNGD:
(A"(w))1,,))) V.A 7/, V((w),
We+1 Projø (w)41,))
(11)
Here Projg(w> denotes the non-Fuclidean projection of w onto Ω induced by geometry of A". This involves 1) a transformation from the expectation parameter (E) the standard parameter (/1, E),
2) Euclidean projection of (H, Σ) onto P. and 3) a reverse transformation back the expectation space. Note that in the mean field case, the projection in the second step can be performed efficiently by simple entry-wise clipping. Denote the clipping function clip/a,)) min(max (a, x) b). then the second formula of (11) is equivalent to
(M++1): clip; clip (1441,)))
(E++1) cliptyp. PIliid d.
7

---

Importance of projection. The projection onto the bounded set Ω is necessary for two reasons. Theoretically, as shown in Section 3. both relative smoothness and hidden convexity hold on the set Q, and SNGD can potentially escape this set (as shown in Figure 1), invalidating its convergence guarantees. Empirically, we will show in Section 5 that projection improves the numerical stability.
4.2 Convergence using Relative Smoothness
prove that Proj-SNGD converges to a sta- tionary point with an (LIVVT) rate. We use Vi(w) denote the stochastic gradient and assume ((V((W) w] = Vl(w). Randomness may stem from using mini-batches rather than the entire dataset when computing the gradient. Next, we impose the following assumption on the variance of stochastic gradients V((w))
Equipped with the relative smoothness of E(-),
Assumption 1. There exists V 2 such that
7 (((((((( - V((we), wTt - w(+1) ) wel S V2 for all wa E Ω
(12)
where w = argmin wet 7/ (VE(w)), w) + DA- (w, wi) is the outpur of Proj-NGD step with exact gradient.
Remark. Assumption 1, first proposed by [HR21], reduces to (l)))((e)) - V((w))) )|12 in standard gradient descent case, however, does not depend on any particular norm in general. This as- sumption was justified and used in [WG24] to show convergence of SNGD in conjugate setting. In Section I. we prove Assumption is satisfied for our mnoncconjgaate) logistic regression models.
For a general setting, we will use the Bregman FH24]) as the convergence criterion to stationary point.
-Bakkdrdd Envelope (BFBE; [ATP21,
Definition 4.1 (Bregman Forward-Backward Envelope (BFBE)). For some p > 0, the BFBE at w Ô is defined as (w(w) == -2 / w/ w) + pD (w/, ..).
In the Euclidean case with unconstrained domain, BFBE becomes the squared norm of the gradient (p(w) [[V 1| In non-Euclidean case it is natural generalization of stationarity criteria the geometry induced by A" and w E ri()2), then Ep(w) = and only if [[[[)]] = O. is shown in [FH24] that BFBE is the strongest criterion available for non-convex SMD. Next, we establish non-asymptotic convergence Proj using BFBE criterion.
Theorem 4.1 (Convergence of Proj-SNGD). Assume E(w) has bounded domain Ô C S and ri( contains at least one stationary point w". Suppose is smooth with respect to A with parameter
L. Then under Assumption 1, for constant step size %o- = min satisfies
Ao VLLT:
Proj-SNGD
Lλo +
LV
E[EsL (wT )]
18
T
T
where do = E(wo) e and WT is sampled from {wo..... w1 with probabilities 7t/ ETT %
The proof is in Section H. Theorem 4. implies that Proj-SNGD converges at a rate of O (L/VTT) in the number of iterations with the presence of randomness in gradients. In noiseless setting (V = 0), we obtain a faster (1)TT) convergence rate. With Theorem 4.1, one may plug in the relative smoothness parameter derived from Theorem 3.2 to obtain the corresponding convergence rates.
4.3 Fast Convergence under Log-concave Likelihood
In this section, we prove that Proj-SNGD attains the global minimum at a fast O(LTT) rate if the negative log-lik logelikel is convex. This fast rate relies on the PL inequality established in Proposi- tion but requires the following additional technical assumption.
Assumption 2. For any aΩ, les n, be outward normal direction at 4 Then holds that
(- (-(V - V(((), < 0.
*For compact
(w Rd (س) 0,1 < < k), the set of outward normal directions E OP is defined as the normalized cone of gradients of the active constraints, i.e., cone (Ve(w) Re, ||v]| = 1).
(w)) 0,
8

---

In the proof of fast convergence of Proj-SNGD (see Theorem 4.2 below), we require that for every Ω
W. == argmin ((Ve(()))))))))))) 2pD (w', ,w) E 52
(13)
where > 0 is a constant that can be chosen arbitrarily large. When w lies in the interior of Q the condition (13) is satisfied for sufficiently large p. When w lies on the boundary, Assumption ensures that the gradient V((w) points towards the interior of f under the geometry induced by A". further guarantees that, for sufficiently large p, w, defined in (13) lies in 12
The example below shows that Assumption can be satisfied for a univariate Bayesian linear re- gression model.
Example 2. We consider univariate Bayesian linear regression model with single data point (r, y). where 2, 2. E R. p(2) = (o, 1) and p(y x; 2) = N (æ2,, 1). It can shown that (see Section H.4) Assumption 2 is satisfied if the following set of inequalities holds:
- -i + 1) < xy < U( + 1), D < x² + +1 D.
(14)
Therefore, by carefully choosing U and D, Assumption 2 can be satisfied for any data point (a, v).
Now we are ready to present the (I/TT) global convergence guarantee of Proj-SNGD algorithm. Theorem 4.2 (Fast Convergence of Proj-SNGD). Suppose ri({t) contains a stationary point w' and & is smooth with respect A" with parameter Suppose log p(D =) is concave z. Let
2² μB Dr Under Assumptions 1 and for the step size scheme
1
ift S T/2 and T S 6L otherwise,
7t
26
(15)
Ma(t
(2/2))+12T.
the iterates of Proj satisfy
192LAo
MBT
648LV2
B(((wT,+) - -80
exp
+
PB
12T.
where E- = f(w*),, Wr,+ = argmin (Vl(w))
(Vl w²) + LD x (wl, we) and No == ((wb) l"
The proof of Theorem 4.2 can be found in Section H.3. In the noiseless setting, V = 0. Theorem 4.2 implies linear convergence and in the stochastic case we have the O(1TT) rate. Notably, this rate matches the rate in [WG24] under conjugate assumption, and the rates in [DGG23, 23] under strong convexity conditions. However, we work in much more general setting non-conjugate models and in non-Fuclidean geometry.
One limitation is our technical Assumption 2 that can be difficult to verify for more complex model than a Bayesian linear regression. However, empirical results (see Section 5) indicate that with moderate values of U and D, even if Assumption may fail, the performance of PNGj is no worse (and sometimes better) than SNGD.
5 Experiments
Numerical Stability of SNGD and PNoj- In this experiment, we again focus on Pois- regression discussed in Figure 1. We consider data (z. 3) = (0.9, 24) where y is sampled from (e*s fixed initialization of variance parameter o = 2, and stochastic initialization of mean parameter μ4 [ 0]). We aim to demonstrate numerical stability of Proj-SNGD and SNGD, The results of 10 independent runs are shown in Figure 2. The solid line represents the median and the shaded area indicates the first and third quartiles. In the left panel of Figure 2, ob- serve an increase in f at the initial iteration when using SNGD with larger step sizes. However, with the same initialization, the increase is absent for Proj-SNGD, as shown in the right panel of Fig- ure 2. Moreover, PNGj- exhibits faster convergence and, despite the constraint an additional û, it converges to the same optimal solution as SNGD.
*Poisson regression admits closed-fom form gradient, and thus no randomness is involved in the training process. The only source of randomness is the initialization of noo

---

step s8E size-05 step sizp=0.1
stpp siso-0,5
Fraciie Value Ceo
10"
Penetion Value Gag
-0P
stor- slpp vinomO.
00-1
10"1
10ml
10-3
40
Eeeration
Iteration
Figure 2: SNGD (left panel) and Proj-SNGD (right panel) applied Poisson regression problem with different step-sizes and initialized with of = 2. Proj-SNGD is used with U = 4 and D = 25. Non-Euclidean projection improves convergence and sensitivity of SNGD. The left panel here is the same as the left panel in Figure 1. The right panel shows that projection fixes SNGD even when starting with the same (large) initial point of
Negative ELBO
1400
SNGD
ProjSGD
Number of Iorrationss
2000
Pnoj Smi-S0
SNGD
SG1
1200
FreySWGD
Pnjj SNDD
3000
2000
600
Egoch
Figure 3: Euclidean and non- aun-Euciddenn algorithms on MNIST dataset. Left: Objective during optimization with tuned step size. Right: Number of iterations before the objective falls below f(w) < 700 for different initial step-sizes "/o- Non-Fuclidean algorithms show consistently better performance, tolerate larger step-sizes and more robust to step-size tuning.
Experiment on MNIST Dataset. In this experiment. we compare non-Euclidean algorithms (Proj-SNGD and SNGD) with Euclidean algorithms (Prox-SGD and Proj-SGD) proposed in [Dom20]. Details of implementation and additional results can be found in Section J. We con- sider logistic regression
subset of MNIST dataset [LeC98] with labels or 8 (n 11769, d 784). We use mini-batches of size 2000 and set the step size "/- - 70/ /l/ where yo is hy- perparameter to be finetuned. We run the algorithm for 1000 epochs (6000 iterations). The results observe
averaged over 5 independent runs are shown in Figure 3. In the left panel of Figure that non-Euclidcan algorithms converge faster than Euclidean ones with tne-tuned step-size. The right panel illustrates robustness the initial step size Yo. with lower iteration counts indicating faster convergence. Non-Fuclidean algorithms reach the threshold in around 1000-2000 iterations for wide range 0.05 < "to < 0.2, while Euclidean algorithms only do at to 0.02. We also observe that Proj-SNGD slightly outperforms SNGD. Therefore, non-Euclidean algorithms, espe- cially Proj-SNGD, are more robust step-size, potentially reducing the tuning burden in practice.
Additional experimental results on other datasets can be found in Section J.2.
6 Limitations and Future Work
While our work makes significant progress in understanding NGVI in non-conjugate models, there are certain limitation we want to discuss. First, our study of landscape properties is restricted to a compact domain. Whether global relative smoothness and/or PL inequality holds remains an open question. Second, our analysis is limited to the mean-field variational family, extending the analysis to full-covariance Gaussian family is an important direction for future work. Third, to obtain (1/T) convergence rate, we invoke an assumption the behavior of objective on the boundary of the domain, which might be challenging to verify, and it would be important to remove this assumption.
10

---

References
[ADY+23] Oleg Arenz, Philipp Dahlinger, Zihan Ye, Michael Volpp, and Gerhard Neumann. A unified perspective on natural gradient variational inference with gaussian mixture models. Transactions on Machine Learning Research, 2023.
[AKLM21] Alekh Agarwal, Sham M Kakade, Jason D Lee, and Gaurav Mahajan. the theory of policy gradient methods: Optimality, approximation, and distribution shift. Journal of Machine Learning Research, 22(98):1-76, 2021.
[Ama98] Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):51-277, 1998.
[ATP21] Masoud Ahookhosh, Andreas Themelis, and Panagiotis Patrinos. A bregman forward- backward linesearch algorithm for nonconvex composite optimization: superlinear convergence to nonisolated local minima. SIAM Journal on Optimization, 31(1):653- 685, 2021.
[AZO17] Zeyuan Allen-Zhu and Lorenzo Orecchia. Linear Coupling: An Ultimate Unification of Gradient and Mirror Descent, In 8th Innovations in Theoretical Computer Science Conference, volume 67, pages 3:1-3:22. Schloss Dagstuhl Leibniz-Zentrum für In- formatik, 2017.
[BBT17] Heinz H Bauschke, Jérôme Bolte, and Mare Teboulle. A descent lemma beyond lips- chitz gradient continuity: first-order methods revisited and applications. Mathematics of Operations Research, 42(2):330-3) -38, 2017.
[BKM17] David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association, 1126518):55--77,, 2017.
[BN06] Christopher M Bishop and Nasser M Nasrabadi. Pattern recognition and machine learning. volume Springer, 2006.
[Bon64] Georges Bonnet.
Transformations des signaux aléatoires travers les systemes non linéaires sans mémoire. In Annales des Télécommunications, volume 19, pages 203- 220. Springer, 1964.
[CAAK19] Badr-Eddine Chérief-Abdellatif, Pierr Alquier, and Mohammad Emtiyaz Khan. A generalization bound for online variational inference. In Asian conference on machine learning, pages 662-677. PMLR, 2019.
[CE24] Ismaël Castillo and Paul Egels. Posterior and variational inference for deep neural networks with heavy-tailed weights. arXiv preprint (rXiv:2406.03369 2024.
[CFLM23] Michael Celentano, Zhou Fan, Licong Lin, and Song Mei. Mean-field variational in- ference with the tap free energy: Geometric and statistical properties in linear models. arXiv preprint 1xXiv:2311.08444 2023.
[DDM18] Damek Davis, Dmitriy Drusvyatskiy, and Kellie MacPhee. Stochastic model-based minimization under high-order growth. arXiv preprint arXiv: 2018.
[DGG23] Justin Domke, Robert Gower, and Guillaume Garrigos. Provable convergence guaran- tees for black-box variational inference. Advances in Neural Information Processing Systems, 011666899663277 2023.
[Dom20] Justin Domke. Provable smoothness guarantees for black-box variational inference. In International Conference on Machine Learning, pages 2587-2596. PMLR, 2020.
[DZZ224] Qiang Du, Kaizheng Wang, Edith Zhang, and Chenyang Zhong, A particle algorithm for mean-field variational inference. arXiv preprint arXiv:2412.20385, 2024.
[FH24] Ilyas Fatkhullin and Niao He. Taming nonconvex stochastic mirror descent with gen- eral bregman divergence. In International Conference on Artificial Intelligence and Statistics, pages 3493-3501. PMLR, 2024.
11

---

[FHH25] Ilyas Fatkhullin, Niao He, and Yifan Hu. Stochastic optimization under hidden con- vexity. SIAM Journal on Optimization. 2025.
[GLZ16] Saeed Ghadimi, Guanghui Lan, and Hongchao Zhang. Mini-batch stochastic approxi- mation methods for nonconvex stochastic composite optimization. Mathematical Pro- gramming, 155(1):267-305, 2016.
[HBWP13) Matthew D Hoffman, David M Blei, Chong Wang, and John Paisley. Stochastic varia- tional inference. Journal of Machine Learning Research, 14()):1003-347, 2013.
[Hen24] Hadrien Hendrikx. Investigating variance definitions for mirror descent with relative smoothness. arXiv preprint arXiv:2404. 12213, 2024.
[HR21] Filip Hanzely and Peter Richtárik. Fastest rates for stochastic mirror descent methods. Computational Optimization and Applications, 79:717-766, 2021.
[HRL12] James Hensman, Magnus Rattray, and Neil Lawrence. Fast variational inference the conjugate exponential family. Advances in Neural Information Processing Systems, 15:2888-28977 2012.
[HTRK07] Antti Honkela, Matti Tornio, Tapani Raiko, and Juha Karhunen. Natural conjugate gradient in variational inference. In International Conference on Neural Information Processing, pages 305-314. Springer, 2007.
[JCM24] Matt Jones, Peter Chang, and Kevin P Murphy. Bayesian online natural gradient (bong). Advances in Neural Information Processing Systems, 17:33104--31153, 2024.
[JGIS99] Michael Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. An introduction to variational methods for graphical models. Machine learning, 37:183- 233, 1999.
[Kak01] Sham M Kakade. natural policy gradient. Advances in neural information processing systems, 14:1531-1538, 2001.
[KBL+ 16] Mohammad Emtiyaz Khan, Reza Babanezhad, Wu Lin, Mark Schmidt, and Masashi Sugiyama. Faster stochastic variational inference using with general divergence functions. In Proceedings of the Thirty-Second Conference on Uncertainty in Artificial Intelligence, page 319-328, 2016.
methods
[KL17] Mohammad Khan and Wu Lin. Conjugate-computatio variational inference: Convert- ing variational inference in non-conjugate models to inferences in conjugate models. In International Conference on Artificial Intelligence and Statistics, pages 878-887.
PMLR, 2017.
[KMG24] Kyurae Kim, Yian Ma. and Jacob Gardner. Linear convergence of black-box varia- tional inference: Should we stick the landing? In International Conference on Artifi- cial Intelligence and Statistics, pages 235-245--002 245.243.222 PMLR, 2024.
[KOW*23] 23] Kyurae Kim, Jisu Oh, Kaiwen Wu, Yian Ma, and Jacob Gardner. On the convergence of black-box variational inference. Advances in Neural Information Processing Systems, 04446155446577 2023.
[Lan12] Guanghui Lan. An optimal method for stochastic composite optimization. Mathemat- ical Programmingg 133(1):365-397, 2012.
[Lan23] Guanghui Lan. Policy mirror descent for reinforcement learning: Linear convergence, new sampling complexity, and generalized problem classes. Mathematical program- ming, 198(1): 1059-11109116 2023.
[LeC98] Yann LeCun. The mnist database of handwritten digits. http://yann lecun. com/exdb/nnist// 1998.
12

---

[LFN18] Haihao Lu, Robert M Freund, and Yurii Nesterov. Relatively smooth convex opti- mization by first-order methods, and applications. SIAM Journal on Optimization, 28(1) 28(1)3333-354 2018.
[LKS19] Wu Lin, Mohammad Emtiyaz Khan, and Mark Schmidt. Fast and simple natural- gradient variational inference with mixture of exponential- approximations. In International Conference on Machine Learning, pages 3992-4002, PMLR, 2019.
[Loj63] Stanislaw Lojasiewicz. Une propriété topologique des sous-ensembles analytiques réels. Les équations aux dérivées partielles, 117(87-89), 1963.
[LT21] Yueming Lyu and Ivor W Tsang. Black-box optimizer with stochastic implicit natural gradient. In Joint European conference on machine learning and knowledge discovery in databases, pages 217-232. Springer, 2021.
[LZM 25] Hua Lan, Shijie Zhao, Yuxiang Mao, Zengfu Wang, Qiang Cheng, and Zhunga Liu. Noise adaptive kalman filtering with stochastic natural gradient variational inference. IEEE Transactions on Aerospace and Electronic Systems, 2025.
[MAM*25] Tåm Le Minh, Julyan Arbel, Thomas Möllenhoff, Mohammad Emtiyaz Khan, and Florence Forbes. Natural variational annealing for multimodal optimization. arXiv preprint arXiv:2501.04600 2025.
[Mar20] James Martens. New insights and perspectives the natural gradient method. Journal of Machine Learning Research, 21(146):1-76, 2020.
[MRR 53] Nicholas Metropolis, Arianna W Rosenbluth, Marshall N Rosenbluth, Augusta H Teller, and Edward Teller. Equation of state caleulations by fast computing machines. The journal of chemical physics, 21(6):1087-1092, 1953.
[NJLS09] Arkadi Nemirovski, Anatoli Juditsky, Guanghui Lan, and Alexander Shapiro. Robust stochastic approximation approach to stochastic programming. SIAM Journal on opti- mization, 1744)99,,,,, 2009.
[NY83] Arkadij Semenović Nemirovskij and David Borisovich Yudin. Problem complexity and method efficiency in optimization. Wiley, 1983.
[OA09] Manfred Opper and Cédric Archambeau. The variational gaussian approximation re- visited. Neural computation, 2163):786-79,, 2009.
[OSK 19] Kazuki Osawa, Siddharth Swaroop, Mohammad Emtiyaz E Khan, Anirudh Jain, Runa
Eschenhagen, Richard E Turner, and Rio Yokota. Practical deep learning with bayesian principles. Advances neural information processing systems, 14:226444277 2019.
[Po163] Boris T Polyak. Gradient methods for the minimisation of functionals. USSR Compu- tational Mathematics and Mathematical Physics, 344)8844-878, 1963.
[Pri58] Robert Price. useful theorem for nonlinear devices having gaussian inputs. IRE Transactions Information Theeory 4(2):69-72, 1958.
[RGB14] Rajesh Ranganath, Sean Gerrish, and David Blei. Black box variational inference. In International Conference on Artificial Intelligence and Statistics, pages 812-822. PMLR, 2014.
[RM15] Garvesh Raskutti and Sayan Mukherjee. The information geometry of mirror descent. IEEE Transactions on Information Theory, 1511:11441 2015.
[SEH18] Hugh Salimbeni, Stefanos Fleftheriadis, and James Hensman. Natural gradients in practice: Non-conjugate variational inference in gaussian process models. In Inter- national Conference on Artificial Intelligence and Statistics, pages 689-697. PMLR, 2018.
[3L]++33] Youbang Sun, Tao Liu, Ruida Zhou, PR Kumar, and Shahin Shahrampour. Provably fast convergence of independent natural policy gradient for markov potential games. Advances in Neural Information Processing Systems, 36:43951-43999 2023.
13

---

[Ste81] Charles M Stein. Estimation of the mean of a multivariate normal distribution. The annals of Statistics, pages 1135-1151, 1981.
[TLGI4] Michalis Titsias and Miguel Lázaro-Gredilla. Doubly stochastic variational bayes for non-conjugate inference. In International conference on machine learning, pages 1971-1979. PMLR, 2014.
[TR19) Da Tang and Rajesh Ranganath. The variational predictive natural gradient. In Inter- national Conference Machine Learming, pages 3145-6154, PMLR, 2019.
[VYB+22] Nuri Mert Vural, Lu Yu, Krishna Balasubramanian. Stanislav Volgushev, and Murat A Erdogdu. Mirror descent strikes again: Optimal stochastic convex optimization under infinite noise variance. In Conference on Learning Theory, pages 65-102. PMLR, 2022.
[WG24] Kaiwen Wu and Jacob R. Gardner. Understanding Stochastic Natural Gradient Vari- ational Inference. In International Conference on Machine Learning, pages 53398- 53421. PMLR, 2024.
[W]+ Martin J Wainwright, Michael I Jordan, et al. Graphical models, exponential families, and variational inference. Foundations and Trends@ in Machine Learning. 1(1-2):1- 305, 2008.
[WW13] David Wingate and Theophane Weber. Automated variational inference in probabilistic programming. arXiv preprint arXiv: 1299, 2013.
[Xia22] Lin Xiao. On the convergence rates of policy gradient methods. Journal of Machine Learning Research, 23(282):1-366 2022.
[ZHI8] Siqi Zhang and Niao He. On the convergence rate of stochastic mirror descent for nonsmooth nonconvex optimization. arXiv preprint arXiv:1806.0444 0477880000048,70 2018
[ZSDG18] Guodong Zhang, Shengyang Sun. David Duvenaud, and Roger Grosse. Noisy natural gradient as variational inference. In International Conference on Machine Learning, pages 5852-5861. PMLR, 2018.
14

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification:
The paper's contributions and scope are clearly summarized in the abstract and introduction, and are well supported by the assumptions, theorems, and propositions presented throughout the work.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
fine to include aspirational goals motivation as long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
The limitations of the paper are summarized in the last section (Limitations and Future Work).
Justification:
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The au- thors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the ap- proach. For example, a facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't it acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory assumptions and prool's
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
15

---

Answer: [Yes]
Justification: All the assumptions are clearly stated, and complete proofs of all theorems and propositions are provided in the supplemental material.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theo- rems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data provided not)?
Answer: [Yes]
Justification: The pseudocode for all algorithms used in this paper can be found in the supplementary material. Model parameters (including step size, batch size, etc.) are also provided in Experiments section and the supplementary material. The code will also be released.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model ge.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate the research performed.
While NeurlPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case thors are welcome to describe the particular way they provide for reproducibiliity In the case of closed-source models, it may be that access to the model is limited in
16

---

some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
The data publicly available. The code will released enable reproduc-
tion of the experimental results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cr cc/
public/guides httee//giidss/oddeuuiiiiise "CodeSummissionP crniiissooooicyyy for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //nips wwwpccccp.c/cccccccccccccccllllt wwtpcc//ggg/eeieee Sttp://eeeeSeeeeeeeeeeeeee seiionPolcyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: Necessary implementation details are provided in the Experiments section and the supplementary material.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should presented in the core the paper to level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: The medians and interquartile ranges are reported in the plots for each exper- iment involving randomness.
Guidelines:
The answer NA means that the paper does not include experiments.
17

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state The authors should prefer- ably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory. time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Information on computer resources, including type of compute worker and time of execution, is included in the supplementary material.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't t make it into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: : http://eerrpppeeeee cc httbbiiic//iicccccciieieeeeeee
Answer:
[Yes]
Justification:
The research conforms with the NeurIPS Code of Ethics in every respect.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiciio))
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This is a theoretical paper on optimization, and it has no foreseeable positive negative socictal impacts.
18

---

Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This work is theoretical and it does not involve releasing any high-risk data or models.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by re- quiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: This paper uses the MNIST dataset, which is provided through Py PoToch's built-in torchvision library. The original dataset is properly cited in our work.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URI.
19

---

The name of the license (e.g.. CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the pack- age should be provided. For popular datasets, paperswith wwwoeheeeee...o www.ooo///cc//o datasttss has curated licenses for some datasets. Their licensing guide can help determine the li- cense of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors are encouraged reach out to the asset's ss creators.
13. New assets
Question: Are new assets introduced in the paper well documented and the documenta- tion provided alongside the assets?
Answer: [NA]
Justification:
This paper does not release new assets.
Guidelines:
The answer means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the pa- per include the full text instructions given participants and screenshots, applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, cura- tion, or other labor should be paid at least the minimum wage the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
20

---

Depending on the country in which research is conducted, IRB approval (or equiva- lent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
The LLM is only used for writing and visualizing purpose.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: ptttps://neuripee cc//oonferrncces//22////L for what should or should not be described.
21

---

Contents
Introduction
1.1
Related Work
2
Background
2.1
Variational Inference and Exponential Families
2.2
NGVI as Stochastic Mirror Descent
2.3
Mean-field Parameterization
3
Landscape Properties of Non-Convex NGVI
3.1
Relative Smoothness of Variational Objective
3.2 Hidden Convexity of Variational Objective under Log- concave Likelihood
4
Convergence of Non-Convex NGVI
4.1
Projected Stochastic Natural Gradient Descent
4.2
Convergence using Relative Smoothness
4.3 Fast Convergence under Log-concave Likelihood
8
5
Experiments
Limitations and Future Work
10
Appendix
24
A Summary of Assumptions in this Paper
24
Descent Lemma for Mirror Descent
25
Background on Exponential Family
26
C.I Facts about Gaussian Variational Family
26
C.2 Derivation of SNGD
27
D Coerciveness of the Variational Objective for Gaussian Prior and Gaussian Variational Family
28
E Examples of Non-convex Objectives in Non-conjugate Models
30
F Tighter Relative Smoothness in Univariate Case (n4, E RR)
32
F.1 Necessary and Sufficient Conditions for Relative Smoothness
32
F2 Simple Sufficient Conditions for Relative Smoothness
33
Missing Proofs in Section
35
G.1 Proof of Theorem 3.2: Sufficient Conditions for Relative Smoothness
35
G.2 Proof of Theorem G.3
38
22

---

G.3 Proof of Proposition 2: Hidden Convexity of {(W)
39
H Missing Proofs in Section
H.1 Smoothness and Strong Convexity of A
H.2 Proof of Theorem 4.1: Convergence of Proj-SNGD
41
H.3 Proof of Theorem 4.2: Fast Convergence of Proj-SNGD
................................
44
H.4 Proof of Example 2
46
I Variance of the Gradient Estimator in Logistic Regression
47
J Experiment Details and Additional Results
52
J.1 Pseudocode of the Algorithms
52
J.2 Experimental Results Additional Datasets
53
J.3 Choice of U and D Proj-SNGD
54
J.4 Implementation Details
55
23

---

Appendix
A Summary of Assumptions in this Paper
Below is a summary of the assumptions used in this paper for establishing the convergence of Proj- SNGD:
1. The objective e is smooth relative to A This assumption is required for both Theorem 4.1 and Theorem 4.2, and can be verified using Theorem 3.2.
2. The stochastic gradient is unbiased with bounded variance (Assumption 1). This is needed for both Theorem 4.1 and Theorem 4.2, and holds for logistic regression models (see Sec- tion I).
3. The relative interior ri({)) contains a stationary point w" This condition required for both Theorem 4.1 and Theorem 4.2. It is a mild assumption, since E is coercive under very weak conditions (see Section D).
4. The descent direction points inward on the boundary of Ω. This assumption only required in Theorem 4.2, and it is satisfied, for example, in univariate Bayesian linear regression model (see Section H.4).
5. The log-likelihoop ( is concave in This is only needed in Theorem 4.2, and holds for many models, including logistic regression and Poisson regression.
24

---

B Descent Lemma for Mirror Descent
In this section, we prove descent lemma for mirror descent.
Lemma B.1. Let & be strictly convex distance generating function closed domain 12 C TR. with induced Bregman divergence De- Assume that E:U - I-smooth relative to t. Consider (deterministic) mirror descent update with step size "y
argmin >VV((u)) u) + Do(u, uz),
1e++1
sEu
then with step size L,L, the objective decreases after each iteration, i.e.,
((u+1)) & {((u).
Theorem B. indicates that mirror descent update with step size y causes an increase of objective, the relative smoothness parameter must be greater than 1/7. Thus in the left panel of Figure 1, the smoothness parameter is greater than 1/0.3, hence 1-relative smoothness fails in the non-conjugate model.
In addition, the necessary condition of a-3 relative smoothness in univariate case (Theorem F.1) implies
B 2
2
thus the smoothness parameter should grow exponentially with respect to of- Therefore, Figure is consistent with our theoretical findings.
Proof. By the second part of Lemma F.1 in [FH24] and the definition of u++1, have that for all E U
(Vi(u)),) u) + - Do, 2 (Vl(us), ue-1) - Da - Doo uet1).
Then, we use the relative smoothness of to get that for all v E U,
((++1) S ((w)) + (Ve(u)), ue+1 uz) + (u
1 =((u)) (Vl(u)) uz) + (Ve(ue); m(41) Do (o(uut,)u)) 1 <((u:) (V((u)), u.) (Vl(u)), v) D Da(o, u141).
Finally, we choose u = 4 and the result follows.
25

---

C Background on Exponential Family
In this section, will provide background information about the exponential family. In Section C.1, we will derive the natural and expectation parameters of Gaussian distribution. We will prove the simple form of NGD update (formula (4)) in Section C.2.
C.1 Facts about Gaussian Variational Family
Recall that q belongs to exponential family it takes the following form:
((2;)) = h(2) exp ((o(=), n) - A(1)),
where v) is the natural parameter, ((2) is the sufficient statistic, and A(r) is the log-partition function.
The expectation parameter is defined as = (((m)) [+)]] Now we let g(z) be the PDF of Gaussian random vector (H, E), then
(z)) x exp
(z ( H)) 14) - (z n)
log det(2))
2
- -
x exp
(zz
+ (z, (,,--1, u) - 11,20-113 log det
Therefore, the sufficient statistics are z and zz and we have the following fact.
Fact 1. For multivariate Gaussian distribution 9. the expectation parameters of q are given by E == [[2] = (t and E = E[zz]]= = = pxp" The natural parameters of are given by X == Σ-1
and A = --
Moreover, by plugging the definition of À and A, one have
A(A, A) = (1, E-'m) + 2 log det (E.) = 1 ATA
log (et(-2A).
Next, we compute A via the definition of convex conjugate:
A (E, E) max A.A-0
(E. + (E, + ATTA A 1 2
log det( det(--A) 221
It's easy to see that the optimal solution of the previous problem is given by
X = -2A A (EE -E)
Moreover, have the following result.
Fact 2. The convex conjugate A equals
(E = 1 log (E ((TT--)T)= )++ + d + $ log2, i EET > 0, otherwise.
L00, +00,
Based on Fact 2, we can check that
VxA(A,A) = 1 - VAA(A,A) 1 _ - -
-A='=E, =E,
2 A-
VeA' -E--ET)--E-A
-
A.
This also indicates that V and V A are inverse operators of each other.
Moreover, when d = 1, it's straightforward to check that
22p o2
<2Ar(w) σ
-p1
This result will be useful the proof relative smoothness, e.g., in Section 3.1 (8).
Next, we consider the mean field family.
d
<(2) x exp
Σ
+
Pi
2i +
log Eis
25,
Ejj
26

---

Fact 3. For multivariate Gaussian distribution under mean- parameterization, the expectation parameters of q are given by & := [[2,] μi and E;; == 1| Ei + pu entry-wise. Equivalently, we have
5 = p. E == + diag(a 1t).
Similarly, also have X Eー and diagonal matrix A -
In addition, Fact implies that the transformation between any two of the natural, expectation and standard parameters can be performed efficiently in O(d) time. Therefore, every update of Proj- SNGD in (11) (given stochastic gradient V((we)) can be performed in (dd time.
C.2 Derivation of SNGD
In this section, we show that formula (4) holds, i.e.,
[((m)) = Vl(w)),
where тI and W, are the natural and expectation parameter of the same distribution. The proof adapted from [RM15] and [WG24].
Proof of Equation (4).
We first prove that I(n) v²A(m) Since
(( h(2) exp( m) A(7)).
we have
V. Volog((2;)) x(z) V <(=) [[]())] By definition of Fisher information and duality between 7) and w.
I(n) =E,V, -og, log 1 Cov(o(2)).
Moreover, have
-V / (()))(2)d
Vx(g(2))o(2)d=> de
/
=(((()(((()) B[o(=))) Cov((((=)
Then, we use the chain rule to get
(n),(,) V,fl = v;A(T) = I(n)
Multiplying IIm) on both sides gives the desired result.
27

---

D Coerciveness of the Variational Objective for Gaussian Prior and Gaussian Variational Family
In this subsection, we aim to investigate the coerciveness of the variational loss f(w) with Gaussian variational family.
Definition D.1 (Coerciveness on Bounded Domain). Let f u - R be a real-valued function defined on an open set U. We say f is coercive if f(w) too as E U. - ou or ||a|| + oc, where öu denotes the boundary of
Notice that coerciveness may potentially depend on specific parameterization. We consider 4 differ- ent parameterizations of Gaussian variational family in our work:
Standard:
(4: E)
Expectation:
w = (1, E + pp )
Cholesky:
c(w) = C)
1 17 Σ 11, - 5ー
Natural:
In Cholesky parameterization, C is defined the Cholesky factor of E. Since focus is the convergence of NGVI, we are mainly interested in the landscape properties of f(w) in expectation parameterization. We first show objective ((U, Σ) is coercive in standard parameterization (n, Σ) R² x S4.
Theorem D.1 (Coerciveness of the Variational Objective in Standard Parameterization, Lef g(z) N p. )) be a Gaussian variational distribution with mean parameter E IR and covariance matrix E E St. Let the prior be the standard normal distribution p(2) = 0,; Suppose the
log-likelihood fanction satisfies a sub-quadratic growth condition: there exist constants Co, Ci 2 and an exponent r E [1,2) such that for any observed data 'D and all z E IRd
log (D 2) S co + 1=1""
(16)
Then the variational objective E: IRd x S + RR defined in (6) is coercive, in the sense that
11,,, Σ)] + 0 or del(E) +
E(p, Σ) + co.
The sub-quadratic growth assumption (16) is a very mild condition which is satisfied by most prob- lems in practice. For example, the ligelioeld function log p(D [) log p(y xr 2) for dataset D = { y4) y)) diverges to -oo when ||=|| + ac in Bayesian regression and logistic regression, and has at most linear growth in Poisson regression for bounded {x:),-
With this result, we can prove that the objective coercive in all parameterizations as listed below.
1. Expectation parameterization. The domain of expectation parameter is Ω {(E, & E IR - <E" E } For &. || + oo if and only if || + co. For unboundedness of since | < |Σ + 2 || + oc implies at least one of |E|| - oo and |||||| + co holds. Moreover, if E) + 0n, i.e., de < ) + 0, this is equivalent to det(EE) - 0.
2. Cholesky parameterization. Since C defined as a lowerltriangle matrix with positive diagonal entries, ||C|| oo if and only if ||=||| co, and det(C) - 0 if and only if (E) + O.
3. Natural parameterization. The domain of natural parameter is DA = ((A,A) : e Rdd - $4.. For X, ||||| <
= (Amin - thus ||| + 0 implics either de) + 0 or || + O. The proof of A is obvious.
Proof of Theorem D.I. It's clear that the boundary of Rd x S is Rdd { : Σ is symmetric and det(E) = O}. To show coerciveness, we need to prove that under the following cases, the objective E(1, Σ) will diverge infinity:
1. ||||| + 20.
28

---

2. det(E) + 0,
3. |-||| + oC.
Recall from (6) that the variational objective is given by
E(p, Σ = (x - 1,, [z]].
For the first term, we can find the closed-form solution of KL. divergence between to Gaussian distributions:
(RL(a(=) 2
log (2) + Tr(E) + ||/] - d]
For the second term, the sub-quadratic growth condition (16) implies that
-E, 1o x)] 2 co oi =|[[==11"]
remains upper bound [|"]]||]] Let (0, I) be a standard Gaussian random vector, then by Minkowski inequality and Jensen's 's inequality we have
[E] =1 (Elll + 11)2=/"
S ||l/| + (E[(II (=1()/))//201333//1)))
S |ell +
|all Tr(E)//2.
Therefore, since 1 < < 2, it can be shown that for some constant C > 0, it holds that
||||-||] < C(lldI + Tr(2))/))
Finally, we have
1 - log det(2)) + Tr(2) + 1|al12 d) co C(C(ll + Tr())/)).
euu, 2 2 2
If det(E)) 0, the first term will diverge +0c. If ||l|| + co || square brackets dominate as r < 2, hence E(p, Σ) will diverge to +00.
oc, the terms in the
29

---

E Examples of Non-convex Objectives in Non-conjugate Models
In this section, we present two non-conjugate models where the objective non-convex. These examples are also discussed in Section 5 of [WG24], but we include them here for completeness.
Logistic Regression. We consider 1-dimensional logistic regression model with data D {(z,, y4) [-1, 1]. E [-1,1}) and latent variable = (w,b) R The prior p(2)
standard Gaussian. Then we have
((w) Σ
E (( [log(1 + exp(-V((zz + b))) DxL (q(w, b) ((w,
In the following will write 8; o0-y:(ur; + b)) for 1 S i n, where o(-) is the sigmoid function. On the convex subset
Ω fw E): : E s and E diag(o? 02)) S 2
can compute the second derivative with respect to o3 using Theorem 3.1
V (
Σ Eg (6s - 6s, + 1)] + 20
A necessary condition of convexity of f(w) is that the second derivative w.r.t. of is non-negative.
Note that of, 0b 0, w. 0 and - 1/2 in probability for all 1 S n, thus
lim () [s:(1 =;s - 6s; + =)]=-
32 < 0.
o..0
4
Then there exists 8 > 0 such that when of = oī. = 8,
8;(1 - (6s 68; + 1) S 64
holds for all 1 < i < Then we can choose > 32/0² to show that ((w) is non-convex when of = o2 i.c.
Vag Vagl(w) (i0gE-)
0.
Poisson Regression. In this example, we are given dataset D = {(x, U:) :2 E Rd, y: E N)
with latent variable 2 E Rd. We assume that y x ~ Poisson(A), A), where
A =xpxp x).
The expected log-likelihood admits the following closed-form solution:
1 (E - -55 )ar:) )
- Ex(()) log(ulær 2)
Σ
ya E + exp x + 2
Then we can compute the Hessian w.r.t. g:
Vellw) = Σ
exp
x T 1 (E - E5 )zs
x. [(x7 & 2)0,0, + Vi
2
i=1
We consider the following convex subset where the covariance matrix equals identity, i.e.,
S2 = tw (E. E) E E Rd, =- = EE + 213.
We can find xy and & such that < xTT6 < 2 for all S i
S n. Then we have
exp (x. E + J (E S5 )ax) > 1, and thus
Villw) 3 Σ x E(x g 22) 2)xa + VEA(w),
30

---

where we used the fact that DKL || (((;w")) = DA- (w),") Moreover,
= + Σ- E-1
Therefore, for some c > 0, rescale the dataset D' {(rr; Ui) 2z E R ; E N and evaluate ((w) at E' - E/c, we have
Vil(w) 6/=
cx 5(2 2),,2, +
+ 25 x 0,
mm
if è is sufficiently large.
31

---

Tighter Relative Smoothness in Univariate Case (1, a E R)
In this section, we will provide tighter relative smoothness guarantees in univariate case, i.c. μe, a
R. using the equivalent conditions for relative smoothness in (9).
E1 Necessary and Sufficient Conditions for Relative Smoothness
Let f(z) == log p(D z). Incqualities (9) give the necessary and sufficient conditions under which
relative smoothness holds. We first find such that the inequality on the right hand side of (9) holds, i.e.,
Aw) - V B [/(=)] ž 0.
(17)
Let = = relative smoothness condition (17) is equivalent to
f())), C v f(=)]. Using the explicit form of the Hessians (7) and (8), the
A
2n
2 B 12C]
1 - WC) 2n
M(3)
= 0.
(18)
B [C]
C
Theorem F.1 (Necessary and Sufficient Conditions for Relative Smooth Condition (17), Univariate). Let f be a four times continuously funtiionn
function on IR. Then for Gaussian variational family and some positive constant 3, condition (17) holds if and only
Q = (g(z;w) : w = (E, E [t}
for all R. o² > 0,
Bn V f(z)] 23,
2
x² 0f 2
((e)
(19)
8
E
(E,lV* [{)]) 2 0.
296
4g²
Remark. Theorem F.1 is the immediate result of Sylvester's criterion of positive-semideefiiteenss of matrix MB).
The first condition corresponds to the non-negativeness of the bottom right entry.
The second condition corresponds the non-negativeness of the top left entry, Note that this is a quadratic function of since B and C' implicitly depend on However, if we impose uniform
bound [C| < B1 and |B] s B3 for some 31, B3 2 0 for all μ R, o2 > 0, the following inequality will be sufficient condition for it to hold:
28 64
3a
11 233
2 0.
(20)
(20) holds for all y E and only if: either the coefficients of quadratic and linear terms are both zero (V f(z) 0), or the discriminant is non-positive, which gives
3
82
(21)
o2
20
4
The last condition corresponds to the non-negativeness of the determinant. direct computation of the determinant gives
(22 - C)
(28
8
-
C
(8-2)
del(M(())
=1²
20
3
C 4)
8 04
2 2
3
8
C 4
B²
+ μ
2B
B
2
4
2
Note that the coefficients of quadratic and linear terms are both 0, therefore we only need to guarantee that the constant term is non-negative. Interestingly, this condition is very similar to the non-positive discriminant condition (21): ) one substitutes B, with the upper bounds of [B], |C| to obtain (21).
32

---

We also need to find conditions under which the inequality on the left side of (9) holds, i.e.
(v) 0, V V By(csw)
(22)
which is equivalent to
2 + 12C =C1 + o 2 [ [C]+ -((a)= = -
2
=
1 2
We apply Sylvester's s criterion before to obtain the following result.
Theorem F.2 (Necessary and Sufficient Conditions for Relative Weak Convexity (22), Univariate). Let f be a (a.e.) fourth differentiable function on R. Then for Gaussian variational family Q (g(z;w)
E) E Q2 and some constant a, condition (22) holds if and only if for all E R, o² > 0,
('E,IV'((=)) 20x a
24²a
+ " V'
22 [()))
0²
(23)
a²
Eg
(E [( 2 0.
206
40²
Remark. The first condition implies an important necessary condition for relative convexity (i.e., 2 0): V*f(2) > 0 for all 2 E R. Indeed, if V (2o) < 0 for some 2n and V f²((2) is continuous, we can always pick some (j,, o²) such that E,IV* f(z)] 0 (e.g. 1 = 2g and o² is sufficiently small), then we cannot find any 2 0 such that the first condition holds.
F22 Simple Sufficient Conditions for Relative Smoothness
In the following, we aim to provide sufficient conditions for (17) and (22) with lower-order derivative conditions by applying Stein's lemma.
Lemma F.3 (Stein's Lemma [Ste81]). Ler Z (p, Σ), then any differentiable function g. we have
(Z [4]] = [
Theorem F.3 can be used to reduce the order of derivatives. More specifically, we can show that (see Theorem G.3 and Section G.2 for a more general proof multivariate case)
| f(=)] - 2 sup V f2)),
|V, f())] o0 sup Vf(()).
(24)
Combining the parts above, we can get the following sufficient conditions.
Corollary F.3.1 (Sufficient Conditions for Relative Smoothness, Univariate). Under the assump-
tions of Theorem F.I. we further assume that sup, f < L1 and sup, 2 f() L2 for
some constant L1, L2 2 0, then on the restriction of the parameter set {(u, o²) : o² < D) for some D O, condition (17) holds with
D
2DD 2
B 12+
L2
Li
Moreover, relative weak convexity condition (22) holds with parameter
D
V2D
- L2
L1.
2
2
Proof. For the choice of B, we first check three conditions Theorem F.1 one by one:
1. Since x² < D sup, [f2 < T.2, use the upper bound on Eg V V+S(=)) we have
in (24) and
E, [()] < o² sup [V2 < DL2 < 2
33

---

2. Thanks to our global bound on f())] and f()]] in (24), we have
23 (ox -EI n² 2μE,IV²f(z)) +
3 z o² ot o²
23
L2
2L1 1²
B o²
|
r²
holds for any E R. In order show the RHS is non-negative for any μ we first assume μ 2 (similar results follow if μ S 0). In this case, we only need to prove
(23 120²) 2L
30 20 0.
This is a quadratic function of μ with positive quadratic term coefficient (this is the first condition), hence a sufficient condition is that the discriminant is non- non-positive, that is,
AL10² - 4,80 - L2a²) < 0.
(25)
This is again quadratic function of 5 and the inequality holds if we pick a sufficiently large B. More specifically, the larger root of (25) is
3 L20² + VL 88,002 L20² + L202
+ 2/2L102
D
V2D
L2+
L1
-3.
4
4
2
2
Therefore, B satisfies the inequality (25).
3. We again use the upper bound of EE, V$ (elV((()) f(=)) and E,, V [())] in (24).
B E, f(2]B
(E, (((
8²
L
L3
2
402
4
206
4o
4
In fact, RHS is non-negative if and only if (25) holds.
The proof of relative weak convexity is similar.
Remark. In order to obtain relative strong convexity guarantee of the objective [(w), since the KL divergence term in (6) is I-relatively convex, need to guarantee that
D
V2D
L2
2 L1 -1.
2
As result, under the assumptions of Corollary F.3.1, we conclude that [[(]( is relatively smooth with respect to A" with parameter
D
22D
L 2
12
LT.
2
Example 3. Below some examples where the likelihood p(D [)) plyl x for D {(x, y)) specified:
Linear Bayesian Regression. In linear Bayesian regression where p(y xr, 2) = N(xz, o²), it's easy to see that - logp(y [x, z) is a quadratic function z. hence the necessary and sufficient conditions in Theorem F.1 and Theorem F.2 are satisfied with = B 0. Therefore, the negative ELBO is relatively smooth with respect to A" with parameter 1, and is also relatively strongly convex with parameter 1. This coincides with the result in [WG24] when the model is conjugate.
Logistic Regression. For logistic regression, let f()) can be shown that /V](=) S |x), |V²](=)| < 4 for all E R. Therefore, we can set 0² < D and [x)] is relatively smooth with parameter L = z'D x
log x, 2) log(1 <=**)), then
+ /2D|=l according to Corollary F.
Poisson Regression In Poisson regression, we write f(z)
log log x. x. e zzyz
order to satisfy the first necessary and sufficient condition o (V < 2, the necessary and sufficient conditions Theorem F.I. have upper-bound |o|| |/c| and の² Therefore, it is in general difficult to prove relative smoothness for Poisson regression without further assumptions (c.g., bounded domain).
for y E N4 and some constant c. Then E, [V)(()) = x - =2=/2
34

---

G Missing Proofs in Section 3
G.1 Proof of Theorem 3.2: Sufficient Conditions for Relative Smoothness
We begin with two lemmas which compute the Hessian matrices V [f(z) and (4))
Lemma G.I. Let f E CA then for = Eu, Sd. Edd) E Read, we have for q (1,.., d}.
Bmmlp Cpplip:
j 2p 1 i= j = 2p,
4 Cppn
1 3Bpp Cpplip):
(i.j) = (2p 1,
or (i,j) = (2p, 2p 1);
H pm Baaty Baply + CaHlpHa
(i.j) (2p 1,29
1), + 9.
2 BBps Cpalt)),
(1.j) = (2p 1,24)
or (i.j) = (2p. 29
1), p) + q.
Cpg.
(i.j) = (2p, 2q). .P # q.
where Hig E Bij = E,y ( (3 [=)] and Cij = (V, [., (=)]
Proof. We apply Bonnet's and Price's gradients (Theorem 3.1), we can easily get the following partial derivatives
V.. -.BB
[[[=]] - -
Next, we use Theorem 3.1 again to get for +
Ve,c B[f())] Eyl (z)())
= =E,[VVz](>>> (=)) - Viy + (2y)))
= -Hi- - Biglt- - BigH4: + CigHiitt.
And i
V2 VV, - Vi.f()E
-E::V Via Vin + Vhnr(=)) vaf(2))
= = --2Ba44 + CRug.
Moreover, for either i j or # j. we have
1 v=)()) = 2 1 -11gn(()))) = =BBs,b - -CgH)).
RE,,, [[f[)) = VE,,
V VG, = _ Ve,, VV - +,,
= Cis
1
Lemma G.2. For mean field family qu (2), we have
(w) diag(V 2 (w)), V2 Al2))) ... V2 (wa))
where each (wi) is a x 2 matrix corresponding to univariate case, i.e.
(2 of
u4
1
44
Proof. From Section C.1 we know that
Ve (E EETT ) E-1 m
(w) = - (3 EE ) - x-
2
2
35

---

Since the coordinates are independent, V block diagonal matrix, and it's easy to show that each block
w A (w) = 0 for all i o Therefore, (2A7(w) a
V Wiw! (w) V2 (7w)) of
2
+ o
-14
14 2
In order to prove the theorem, we still need a lemma to transform high-order derivative conditions lower-order ones with Stein's 's Lemma F.3. The proof of this lemma is deferred Section G.2.
Lemma G.3. For f E (IR)), Z ~ ((,, )) where E is diagonal matrix, we have that for any i (1,..., d}.
E y [()) S of o, sup [f())
With the three lemmas above, are ready to prove Theorem 3.2.
Proof of Theorem 3.2. First, we will show that with (dd L2U)) (d) + l22U),.
( [[()] 3 3 (w)) holds, i.e., M(3) == 8V2 (w) V (=))))(((() is positive
semidefinite. According to Theorems G.I and G.2, for p, & E {1, 11.... .d.
( +z)) + 2B Bppppp Cpply
i= j 2p 1,
Cpp
i j 2p,
20
3 {Bmm Cppll),
(i, (2p 1, 2p) (i,j) (2p, 2p 1),
M(3)
Hpg Braln + Bap Caplipan
(i, -(2p 1, 2q
1), + q.
(Bpg Cpaltp)
(1,5) 1,2-1,2g)
or (i.j) (2p, 24 1), + 4.
Ppp:
(i,j) (2p, 2q), #p*#
(26)
Then, we will prove that with 3 = OLD + L2U) + dD + L2U)), M is positive semidefinite. In the following, we will drop the dependence of /3 for convenience.
Note that the assumptions in the theorem and Theorem G.3 imply that BByl T.2 sup, pvf(=)] < DL, and |Cpgl I o, sup, V (
DL2 It is obvious that
|Hpyl I L2-
Define matrix M = M- ddiag(M) M). where diag(M) is the diagonal matrix by setting all non- diagonal entries of M to 0. Then for any E R2ad
2d
T MM
U.
22-
ΣΣΣ
Miy
1=1
BEBiEB jEB
(27)
d
2
U2p U2p 2+j M2p 2+i
km
Here B is the set of sets of cardinality d, which contain exactly element of {1,2), one element of (3, 4). ... one element of 1, The set B has cardinality 2d. Equation (27) can be proved by comparing the coefficient of each term M. on both sides. Therefore, to prove v Mr 2 0 for any E 2d we only need to show each term on the RHS is non-negative. In other words, we only need to show
For any B E B, the matrix M² is PSD, where MB e [dd is the submatrix of M by choosing the i-th rows and columns where i B.
For any p E {1.-. d}. the submatrices of M
[M (M2p-1,2p-
M2p- 1,2p Map.2p
(p)
M2p.2p-1
are PSD.
36

---

For the first set of conditions, we aim to find 3 such that the matrices M are diagonally dominant and the diagonal entries are positive, For positiveness, if 8 2 OLDD L2U) (L) + T), then
M2p 1,2p-1
07 (2 + o²) + 2Bppllp Cppll
"
20p
23 2 1ppl Appllppl
D 2DL1U -D2UU
20.
1
1 D-23 2 2
M2
Cpp
DL2 20. 4
20
4
For diagonal dominance, we fix 2p for some p 1,... d. and show that for every M B containing the i-th row and i-th column, the diagonal element M is dominating in this matrix. Indeed, since the 2p-th row does not appear in MB, we only need to guarantee that M2p-1,20 1,2p 2
2d
2d
M2p-12 1,2p-1 2
MM2p-1.8 1,8l
3
(2p o
Bypllp [ Cppllp
Σ
-
Hpi Boips BipHy Critpptss
30
p
i#
+ 2
Bp - -Cpilpl
B
2 d(L2 2DL, + +DZU²² ? (DL, DL2U)
3 p
= = OldD'U + L2U)).
Now, applying the same reasoning for 2p, we obtain
2d
M2p.2p
2
Σ
M2p.il
i= 1,2p}
8
Bip Cipliel 1 Cpxl
1
Cpp 2
60
12
2
ipp
8
d
(DL, + DLzU) + dDL.2
6 p
2
= 2 O + L2U)).
The second set of conditions can be easily verified. Since we have proved the non-negativeness of diagonal entries, to prove PSD we only need to show that the determinant of M(r) is non-negative. For every fixed E (1, the determinant is given by
B (2p2 20 (2p of) + 2Bppl' Cpnti)
(
3
Cop)
Mp oh 2
1 (Brr CroHte))
20
Let the determinant be non-negative and we get
8² 22D'UL + L2U))S+ O + L2U)).
Since the larger root of 82 aB + b
-at Vo the condition is given by 8 82 a vb g + a+ atVo=41b guarantees that 82 > a3 b always holds. Therefore, the determinant is non-negative 3 (T + L2U)) + / [L, + L2UT) (T + LzU)).
2
37

---

Clearly, 3 = L2U)+ + dD' + L2U)) satisfies all the conditions.
-O(dDU - (D (1 I2U)) satisfies ž 0. We can again choose a < 0 and |a] so large that the matrix (M is diagonally dominant with positive diago- nal entries. Since the previous proof provides upper bounds on the absolute value of non-diagonal entries of M, the same bound also applies to the absolute value of non-diagonal entries of -M (a)..
Finally, we need prove that VERRgeew) č AV2A*( i.e., (()]
= -
Therefore, using the same approach, one can show that (M > 0 with -a taking the same value as B. As result, By log p(D) 2²] [x]f(2)] is L-smooth relative to A on Q where L = OLL + L2U) ( + L2U)).
G.2 Proof of Theorem G.3
We first restate Theorem G.3 below.
Lemma. For ʃ (TR)), if Z M(,, Σ) where E is a diagonal matrix, we have that for any i,j (1....,) ,d),
E { o, sup [f()).
Proof. For i j. we have
E [[V)((Z)]
((zi)
Rd-1 ((=1) (z)==
/ d(2i)
30-1
V2) (or ((2--)) al2-)) 21) de-1) )))
dzi
(2,))
where ((2)) is a univariate function and
((7:) [ al= ((==))((==e) Zi) dz)
[g(Z.)) S sup [/())
zERt
Since 74- Nn4, Of), we apply Stein's lemma (Theorem F.3) twice and we have Elg" (Z:)] = "El (Z.)(Z. - p:)]
-,,"EEe (Z,) - -4) + g(2)) o, "
= =o, "E(g(()) - "
= =a, ?Elg (x) -))]
where X; = (Z, - 1i)/o ~ NNO0 1). Using Lemma G.4 and we get the following bound on (Z))
[[V))(()) = (z)) < 0, " - 1)] 2 sup [(=)
1| s o sup [/(=))-
For * j. we define g(Z) = (Z; - (z) and use Theorem F.3 to get E[f(Z)(Z, -2)(2 p)) Elg(Z) ز)] yE()
Reading the i-th coordinate gives - (2
- [4]] = (2) - -
(28)
Then we define l(Z)= v.f(Z) and Theorem F.3 gives
(l -14] = EE)))(
Reading the j-th coordinate gives
B[(Z, ,, V.f()) E[h()) ==))= - (29)
38

---

Combining (28) and (29) and we have the following bound of B by f(Z)]
where X; (Z, pa)/a; and same for X,. Note that X,, X, are independent because is diagonal. Finally, we use the standard result that E X S for standard Gaussian variable X to get that
= - S o, oj sup xx of oj sup [(()]-
By symmetry, Symmeury.-EiVi..... -E[V < प. o sup= [=)) holds well.
Lemma G.4. Let X ~ NO, 1), then
I 1. E| x = 2
Proof. First, have
E||x² - = =[X² -]
px|s1 -1) = 41 (1
Note that
d [I : = (1 2²)e
dr
then we can compute the expectation in closed form
[11xejo.1(1- - x2) = V2m
(1 =212-
dr
1
IC
222m
0
V2TT
Hence X||X² - 11] 4E [E X²) 2 2 < 1.
G.3 Proof of Proposition 2: Hidden Convexity of ((w)
We define + ((, C e S and is diagonal). This set corresponds to the Cholesky param- eterization of the expectation parameters contained in Q. Furthermore, we define e: Ω + θ to be the one-to-one map from expectation parameter to Cholesky parameter of the same distribution.
Expectation parameters, Ω
Cholesky parameters,
mapping c(w)
c(w) = (H, C): where CC
w (j,, Σ + diag(u O ni)
For bounded expectation parameter domain Ô. we also define its counterpart Cholesky parame- terization é = (14 C) C is diagonal and D-1 < C2 S D, V1 sisid With a slight abuse of notation, we also use e to denote the one-to-one map from Ô to é.
In the following, we will verify that the negative ELBO E(w) is hidden convex defined in Defini- tion 3.2 on the bounded domain Ω
Proof of Proposition 2. Define L R,
((0) = ((c(w)) = Ey(x:0)-- log =) logp(2)) + Ey ((2; 0)},
(30)
where q(2;0) = q(2; C) (CT) is the Gaussian vector with Cholesky parameterization.
39

---

Proof of the First Property in Definition 3.2. We first verify the strong convexity of L on Ô. It is obvious that the domain é is convex. For the second term in (30), we note that it is simply the negative entropy of Gaussian distribution:
(2:00) [log ((2;)) = - d 2
log(( log det(CCT) = Const
log Cis.
Since each - log Cs convex in Cu. the entropy is convex in C. Then we can conclude that the second term [g)(=) [=(2] is convex in 0.
For the first term of (30), we apply Theorem of [Dom20], which states that the jur-strong convexity of f(2) for some function implies the -tr-stoong convexity of (z)) [z() w.r.t. the Cholesky
parameter 0. Therefore, we only need to prove the strong convexity of log p(y =) - p(z) in z.
Indeed, loy z) log p(=) is 1-strongly convex in because log p(y 2) is concave in z and
d log (2x)
1
- log
2
2
is I-strongly convex in z.
Therefore, L()), as the sum of I-strongly convex function (first term) and convex function (second term). is I-strongly convex. Then we conclude that the first property in Definition 3.2 is satisfied with μ
Proof of the Second Property in Definition 3.2. Now we compute He by computing the Lip- schitz coefficient of the inverse map c-1 on Q. We reparameterize w and as vectors: w = (f1, Ei. EEd Eda) E 2ddd and 1 (M1, CI fdd Cda) R2dd, respectively. Note that (0) (HI, Cf Md, m Cld), and VC) = diag(G1, ,Ga) R2dx2d where
0
G = 2/1 2Cs²)
For each block Gj. the largest singular value of G, is bounded by
max (G.) /mmax ( < [TGGGTT 4U 4
The last inequality above holds since
1
2/1s 4C
G.G.T
2/14
and TrG,, = 1 4/4 + 40 < 2 4D using the fact that (μ, (),C) Ô.
Hence conclude that (V)] 44 4D + I, and c-) is /40= 4D + 1-Lipschitz continuous. Then the second property is satisfied with pc (4U2 4 + 1)-1/2,
40

---

H Missing Proofs in Section 4
In this section, we will prove the convergence guarantees of Proj-SNGD in Section We first prove an auxiliary result of smoothness and strong convexity of A in the Euclidean geometry in Section H.1. Next, we prove Theorem 4.1 and Theorem 4.2 in Section H.2 and Section H.3, respec- tively.
H.1 Smoothness and Strong Convexity of A
In this subsection, we will establish the smoothness and strong convexity properties of A in Eu- clidean geometry in the bounded set
á = ((C |&al U, -iaig((E is diagonal, D= S (E-diag(EoE)) D, LA sis d].
It is obvious that A is globally 1-strongly convex in the -ucllieid geometry induced by A' itself. However, is only globally strictly convex in the Euclidean geometry, and strongly convexity only holds in bounded domain of Ω2. In order to transform the PL inequality in Euclidean geometry (10) to the PL. inequality in non-Fuclidean geometry (see Section H.3), such smooth and strong convexity results are necessary.
Lemma H.1. A is strongly convex with parameter Cs == (D(4U² + + 1)) with respect Euclidean norm. Moreover, smooth with parameter Cr. 90'p"
Proof. As shown in Lemma G.2
= diag( V2 (2AT(w1),.". ..,, V²
V² 1 (2/ of
M4 for sisd.
of
Hi
For each block (w,) its smaller eigenvalue satisfies
Amin (wx) > det (v (TV2A7(())) (4 + 20? +1) 2 D(AU2 +D +1)
Since 2²A*(w) is block diagonal matrix, we conclude that Amin (V2A' 2 D(aU=42D+1) hence
V2Ar(w) D(AU2+2D+) 1. For the smoothness part, we solve for the larger eigenvalue directly to get
24 of + 0.5 + (2p + of 0.5)² 4u
Amaxe ('A(ww)))
20
2n of +.5 + |2 + of - 0.5) + 2 2/mml
20
2nº + of + 0.5 + 2n + o + 0.5 + 2/44
20
22 + D2 2
2 U2D2
where we used Va bb b S Va+ vb in the first inequality, and 2 + > in the second inequality. Then we conclude that A is smooth with parameter C₁ == 3 U D2
H.2 Proof of Theorem 4.1: Convergence of PNGD
We first introduce the standard assumption on gradient variance, which is widely used in optimiza- tion.
41

---

Assumption 3. For all wi E Ω there exists V 2 0 such that
B B(ll)((w)) V(( [2 < V2.
(31)
We also define the Bregman Moreau envelope of function é as
(1/ = min w'ESt
[((w)) + pD (w',
In order to prove Theorem 4.1, we will rely on the following theorem on the convergence of SMD under relative smoothness condition. Note that this theorem assumes Assumption 3.
Theorem H.2. [Theorem 4.3 in [FH24]] Suppose A 1-strongly convex and is smooth with respect to A with parameter L. Let (70)0StgT-1 be non-increasing with To & 1/(2T). For (wiHi generated by (32), let WT be randomly chosen from wo: W1 with probability pt. = T// Ei=o %/o then under Assumption 3, we have
3 + 6LV El= 7
[(T)) 5
Σi=
where Ao = 2(((wo) - e*) and l = inf infwef(w).
If we use constant step size Th.
min
do V VLLT 2L
1
then
LAo
LV2Xo
ElEsl (wz) ) 18
+ +9/
T
T
Note that (1/p(w)= < ((w) for all w E Q. Thus Àn < ((((w)) - ft). Theorem 4.1 is essentially the direct result of Theorem H.2 except two minor differences:
1. Instead of the update rules of Proj-SNGD as defined in (11), Theorem H.2 assumes more direct update rule
Wh+1 = argmin y (Vf(we)) + DA- ww)
(32)
wen
We will prove the equivalence of expression Theorem H.3.
2. Theorem H.2 considers different assumption gradient variance, and Theorem H.2 sumes 1-strong convexity of the distance generating function A*. In Proposition we will show that the same result holds also under Assumption 1, and that 1-strong convexity assumption can be removed with this new assumption.
Lemma H.3. Let Wi+1 be the output of Proj-SNGD as in (11), and let wi+1 be the direct update rule as in (32). Then Wt+1 Wi+1. Moreover, (11) can be performed entry-wise in ((d) time.
Proof. Recall that in the update rule of Proj-SNGD as in (11), we define W1+1, E Ω such that A VA* (w:) % (((w))
(33)
By the duality of SMD, 11+ is also the optimal solution without constraints, i.e.,
Wh+ argmin 7/ (VE(w)), + DA-(w,,w))
wen
Use the definition of W441 and (33), we have
We+1 argmin
7/ ( (V((w, + D (w, wi)
weì
argmin(V (w:) V (w)+,,,))) w) (w Ar(w)) (VA'(w)) w/)
wen
argmin D (w, w,), + Const
weì
argmmn D (w, Const. weñ
42

---

Therefore, to find W1+1, we only need to compute w1+1,+ and then project it onto O under the geometry induced by A". Next we show that W/+1 solves the optimization problem above. Use the fact that for mean-field parameterization (see Section C.I),
(E, (2 diag(; @ E)) VEA (5, = - diag(E @ E))
in the standard parameter space we have
1 + 2 log del - (2tt,,
(14+1, Ex+1) = argmin (n.EEP
log logdet(E))
2
1 1+ E diag(a + n))
2
d
argmin (e.E)eP
log Σn - (Ee+1, (1441,)) (E + Z (z + m4) 2
(34)
The last equality holds due to the mean field assumption. Therefore, we can solve wh+1 by optimiz- ing over (14, Ex)) independently.
For each entry i, (34) quadratic function in μr with positive quadratic term and the (uncon- strained) minimum is attained at () If (Me+1,=); < -U, (34) is an increasing function on [-U, U). and if (pe+1.=); > U, the function decreasing on [-U, U7. Therefore, the minimum on UU is always attained at (14+1); clipp-ewr ((#e11,))))
Similarly, (34) is a convex function in Eg and the (unconstrained) minimum is attained at En (E+11)) It's also straightforward check that the minimizer is given by (E++1) clipip (EE+1,, )a). The closed-form solution of (J4+1, E coincides with the standard pro- jection of (M(+1. E++1,) onto P under Euclidean geometry. By transforming (m, Σ) back to pectation parameter space, we conclude that W Wt-1-
Since the computation of (-) and its inverse (V.* involves only entry-wise calculation (see Section C.I), one can also compute
W1+1,+ (VA*)- (w)) 5/ V((we)
in (d)) time. Finally, the transformation between (μ. Σ) and (E, E) and the projection both require (d) time. Therefore, Wi+1 can be calculated in ((d) time. This completes the proof.
Proposition 4. The results of Theorem H.2 also hold if Assumption 3 is replaced by Assumption 1, without requiring A to be 1-strongly convex.
Proof. For 1 2 O, define
N, == [(j(wur - E + +h-Ip((((u)) -E").
In step II (one step progress the Lyapunov function), formula (15) of the proof in [FH24], have
Net1 - (help-L)Dae L)
2(p + L) Ep+ (wr) pri (elwe))) - Ve(ws), - wi)
NOP
+ + V((wi),
-w(+1) p(1 - L) DA-
= =À--mpl L. (6), 2(p + L)
7pp E Eptl(w)) + Pr - Ve(we), Ci -with)
+ + V((we), wett - w(+1) ) - p( -
here we define a = (w))
Note that BK(V((w))) - Ve(ur), w | = because the W; w has no randomness given W/, and the gradient estimator is unbiased. Morcover, by bounded variance assumption (12),
43

---

B(((((r)) - Ve(we), wi+1 - -w--1 wi] I nV2. Therefore, we use the same assumptions as in the proof that - 21. and "/e _ 1/(2T.) and we have
E wil S²x - upll L) (W),
Yp 2(0 L) L) (0t + p?2v²
- p(1 - 2-
<A Tt /n E3r. Esr (wx) + 21,22
This the same formula as formula (17) [FH24], where the last inequality holds due to the non-negativeness of Bregman divergence. Note that in the original proof of Theorem H.2. 1-strong convexity assumption is used only when deriving (17) from (15). However, in our new proof, we can arrive at (17) without invoking the strong convexity condition. Therefore, the same results hold following the proof of the theorem.
H.3 Proof Theorem 4.2: Fast Convergence of Proj-SNGD
To establish Theorem 4.2, we rely Theorem 4.7 from [FH24] (restated as Theorem H.4 below). Before doing so, we introduce the Bregman Prox-PL condition, which serves as an assumption in Theorem H.4. It is similar PL inequality with difference that the gradient norm is replaced with BFBE, a non-Euclidean first-order stationarity measure. Recall that BFBE is defined Defi nition 4.1 as
Ep(w) = 2p min ((V((w),'' - + pD (w),w),
(35)
w'
Assumption 4 (Bregman Prox-PL condition). There exists some constant 3L and μæ > 0 such that for all w E Ω
Em(w) 2 (p
Theorem H.4. [Theorem 4.7 in [FH24]] Suppose A I-strongly convex and is smooth with respect to A with parameter L. Let Assumptions 3 and 4 hold. For step size scheme
ift S T/2 and S ps²
6L
2L
(36)
6
otherwise,
#B
(1/21)+121
we have
192LAo
HBT
648LV2
min E[l(wr,)) -8
exp
pB
12L
We observe that there are slight differences between Theorem 4.2 and Theorem H.4 in assumptions and upper bounds. Specifically. Theorem 4.2 requires Assumptions and whereas Theorem H.4 relies on Assumptions and 4 together with the I-strong convexity of A In addition, Theorem 4.2 gives the last-iterate bound, while Theorem H.4 does not. In the first step, we will prove that the assumptions in Theorem H.4 can be derived from those required in Theorem 4.2. In the second step, we will prove stronger last-iterate bound.
Step 1. First, following the same approach as in Proposition 4, we can show that Assumption 3 can be replaced by Assumption 1, without requiring A" to be I-strongly convex.
Second, will prove that Assumption 4 is implied by the PL inequality (10) in á under Assump- tion 2, by proving that BFBE (non-Euclidean measure of stationarity) is greater than [[((()) (Euclidean measure of stationarity) up to constant in the following proposition.
Proposition 5 (PL implies Bregman Prox-PL under Assumption 2). Under Assumption 2, for suffi- ciently large p which may depend on w, it holds that
E,((W: Z
2C, 20, (IV(()) Ve á.
Together with (10), we conclude that Assumption 4 holds with parameter
14
PB
2C1
44

---

Remark. In Theorem 4.2, we require that Proposition 5 holds for all iterates fwr - Since this set is finite, we can take the maximum value of p corresponding to all W/e ensuring that Proposition 5 holds uniformly along the entire trajectory.
The proof Proposition 5 relies on another metric first-order condition, which is closely related to BFBE:
Definition H.1 (Bregman Gradient Mapping (BGM)). For some p > O, the BGM at w E Q is defined as
A = p² D w
(37)
where D (w, (e), == DA w) + DA- (w- w) is the symmetrized Bregman divergence, and w+ = argmin V ((W), + pDa-(w',u)
EeER
The following lemma from [FH24] reveals the connection between the two measures: Lemma H.5. [Lemma 4.2 in [FH24]] For any E Q and any >> 0, we have
(w) (w).
Proof of Proposition 5. Recall in Theorem H.1, we have shown that A is Cr.-smooth and Cs- strongly convex with respect to the Euclidean norm in the bounded set f. By Lemma H.5 and smoothness of A", we obtain the following lower bound of BFBE
2 (w) Ep(w) 2
2p D (w, wt)
(38)
= 2p² (VA (W (ww wt
2p²
2 [VA (w
C2
where w argmin_/ed (V((w), w/) + 2pD,, 1' w)) following the definition in Definition H.1. Then, can show that without domain constraint, W. argmin (V")) (w', w) satisfies
VA*(w.) V 2p
V((w).
(39)
By strong convexity of A*, we have
I|.. w|ll -
-VA
|e(W)]l-
2pCs
w = w. holds, we can plug (39) into the lower bound of BFBE (38) and we conclude that
Ep(w) 2 2p² [VA" VA*
2C, V((W)]
C2
This completes the proof of Proposition 5. Therefore, it remains to prove wt W..
For w lying in the interior of Ω the condition holds for sufficient large p as ||V[())] is bounded on Ω and 2pD x- (w',u ,w) dominates in the optimization problem. The strong convexity of D then ensures that ||w. w|l is very small.
For 0Ω, will need Assumption 2. Since VA(n) and VA*(w) are inverse operators of one another (see Section C.I), (39) implies
w. = -VA-)- VAT(W)
2p V((W))
v[A( o(1/p)
w 2p
1
+(1/p)
2p
45

---

as p 00.
Then under Assumption 2, for any w E 02 and any outward normal direction ny,, there exists some > 0 such that
VV(ww), no) < -c.
Then we have
e 2 ((1pp) <0
(nw; W. - w) <
2p
for some sufficiently large p. Hence w. E S2 and w = wa.
Step Now we aim to strengthen the bound in Theorem H.4 last-iterate bound. the proof of Theorem H.4 (which can be found in Appendix D in [FH24]), we can derive recursion of form
A4t1 S Ax
7
A2/2
+
3
(40)
11
7//
A, 2T
3
Therefore, we do not need to assume that A, 2 for all T = 0, t to obtain the same recursion. As result, the same approach directly gives the upper bound of the last iterate.
H.4 Proof of Example 2
In this section, we prove that Assumption 2 can satisfied for univariate Bayesian linear regres- sion model, if the parameters U and D satisfy (14).
Proof. Using the explicit form of (2)
(w) in Section C.I, we can compute
0²
2po²
4n²00² +20
2400
Since
E()
Eo,logg(e x, 2)] - DKL (g ll
1
Bgl(y (v xz) x2) 99+ +
(-log 0² + + + 11a 1)
2
1( xp1²
+ +z²0² - loo 0² + 0² + p²= -1] 1]
[( x²)² + +xx + 1)( - 2 log(3 - {²) + 5² 1].
Then we have
V(($,3) = -ry + ミー
- -zy 4
2+11 2(--
22+1 1
20
and
n1 + 1)
(VA =o²
(2 + 0² 0²) + +1) - 2pay
{xx² + 1)
=o² (E + E (2² + +1) - 25xy)
Next, we consider the 4 constraints which define the boundary of ŋ.
1. U. If this constraint is active, the unique outward normal direction is n (1,0), and Assumption 2 requires U(x2 +1) > xy.
2. E - D. If this constraint is active, the unique unnormalized outward normal direction is n, = (-26, 1), and Assumption 2 requires D(22 + 1) - 1 0.
3. -U. If this constraint is active, the unique outward normal direction is n. (-1, 0), hence Assumption requires -U(22 + < zy.
4. = 2 = D- If this constraint is active, the unique unnormalized outward normal direction is n (26, -1). and Assumption 2 requires D- + + 0.
For w at corners (multiple constraints are active), n. can be any normalized linear combination of the outward normal directions shown above. Then (14) remains sufficient to ensure Assumption 2.
46

---

I Variance of the Gradient Estimator in Logistic Regression
In this section, aim show that Assumption is satisfied for logistic regression in mean-field setting. We assume that the dataset = {( vi) : 24 E R, yi E {- We begin with Theorem I.1 bound the gradient of the KL divergence term. In Theorem I.2 will bound the gradient associated with the log-likelihood term.
Lemma I.1. For any w E S2
3v
Proof. Using the closed-form solution of KL divergence between two multivariate Gaussian distri- butions and the chain rule, we have
&
V DKL((((;w)) |1
"ii - -
1 1 Eis E
Va, DkL(0(2w))) ll((z)
2
Therefore, we have
(V || [[))]
S VeDDL [I) (v ||
Svd
sup 1sisd
| DKL (X p(=))\|
sup Isisd
VV Duu( I(()))I)
<Vd
UD + 2
max - D). 11 1-
D
I Vd
UD+
2
dV
where we write V. DeD (g:l((((;w)) p2)) as a vector in Rd and use the fact that U, 2
Lemma 1.2. In D, denote $1 max; ||>||| and s2 max; || xill- Then for all 1 S and E Rd we have
V, log x1, V log (i, + Ell S Us2
lol I z)1l <
$2
4
Here
v log p(y 2112)) is defined as d-dimensional vector where j-th entry equals log p(y; Dig 2).
42
.yz,
Proof. For the j-th entry,
V log logp(ui( =) + V2 logp(wi mis =)]
v., log(1 e - E, v 2,, log(1 e - m)1
= - a-v11 =)yeri - E 2) - o( o1-yior 2) 2!)
S af (-ya =)yaaii) (6 ( 2) oo- (-1102 2) 2))ar!
U
|xyl| + Kijs
2 Eyzy
logp(ui 2)1 = syz log(1 (1 + e -Mi 2)]
= =lo(- -yo =)(1 a1-y1o) 2)) 2))a²))
1if
47

---

where o--) is the sigmoid function. Then we have
|| - V: log p(y: m4, 2) + v2 log p(ys o 2) 0
= |10 - V 2 log Is 2) + Eg E3 log p(gr(u) Iolllaa 24, 2)sssll
[1()z)]] + الر)
$ ||ت|||| + 4 ||l o x:
88 + Us2
4
v log 2, 2]ll
||(- [|--
log (2,,)))[ll
1|l
With the same approach as in Lemma 1.2, we have the following bounds the log- oog-likelhhoo term and the gradient V((w).
Us2
V [ p(u:
Ii, =)]1| || v= log | =) + V2 log I =) ② Elll S +
4
V E log p(ys Mi, =)]1 ||[[- p(us I.. =)ll| S
$2
8
Σ V E
E
logp(y x4 =)I1| V -
log 24, =)]ll]
DKL I(=)))
(41)
Us2
82
3vdUD
In n
$1
+
+
4
2
Next, we define our gradient estimator. Consider mini-batch gradient estimator
(.MBE((w)
logp(Ui. Iia. =)] DeDkL(q(=)))
m
7
V logp(yn (x)2,2) + V? log Fie 2) O E
+
Ve Kl. (g(2)g()))))) llp(2)).
Ve,MBl(w)
V= log ((yia Tia 2)] + V=DKI(((>) (p(=))
m
E 72 log p(y,, Iai 2)] + V=DKu(((2) (r(=))
2m
where each is is sampled uniformly from {1, n). and m is the batch size. There is usually no closed-form solution of the expectation of the log- lieliiol, thus we approximate the expectation with samples, i.e.,
Vel(w)=
logp(u,, Jea, 21) V2 log log / 21) E
mN
l=1
V DeL (1(=)
(42)
m
V=((w)= 2mN k=1
v logp(y I 21) + D=DKL(((>))) (lp22)
48

---

where 21, are iid samples of the variational distribution is obvious that E|V MC [cf(wW)] = Vi(w), so the gradient estimator is unbiased. Similar to (41), we can obtain the following bound on the stochastic gradient Vl(w).
(e(())l
n
vg log 1 21)1 +
|| - Vz log I =1)
MN
t=1
+ || V Dxr. (1(z;)))
(43)
Us2
82
dvDUD
81
2
Next we prove that this estimator satisfies bounded variance assumption (Assumption 3), which will also be useful for the proof of Assumption 1.
Lemma I.3. The Monte Carlo gradient estimator satisfies Assumption 3, i.e,
B[] - [[(w)]]3] S 16n²
((s1+
(011+4.
Us2
2
+
64
Proof.
|8 [[] V((W)]13
21 V((W) [MB(WW)) + 2E MBl(W) [
2E
log |.i,, 2)) + E log log Figr 21]
2
71
E
V. log p(yi. | 2) + CV2 p(Ui Die 2]]
m
m
2E
log
|zig 21]
n
E
V log plun, [xر
2m
2m
1=1
1=1
m
12
21E
VE log p(Viz D44 2)]
VEI- 7B[-logp(114 log Dies 2)]
m
4 max
72
log 22)22) + V p(y 22),2) o
k=
+ 4n² max
V E togr(n Ia )]
2
m
max 2,2
V logp( Iie 2)]
4n max V [x,,
2m
+ 8n² (max
V log p(us 2., 2]]1" max | log 1(y Iz, =)]
V? log p(u: [x,,2)) 2
< 2n
max ||
V_ log x 2) log p(u: |x,, 2) @ E1 + max
+ 12n² (max
V | log |x,, ]) + max | log p(u |x,, =)]1
(01
Us2 2
<16n²
Finally, we ready prove that Assumption holds for logistic regression.
49

---

Theorem L4. Consider stochastic gradient estimator (42). Then with max, | ll and $2 max; |2 @ xell, Assumption 1 is satisfied with some V > 0, which is at most polynomial in n.d, d, U, D, 81 and 82.
Proof. We first summarize our notations.
Proj.
SMD with (e(w)) (stochastic)
W1+1,-
MD with Vi(az) (exact)
Proj.
W
Wi+1
Here wit (and +x+ are obtained by doing an (S)MD update from wa with step size Yts Wh+1 argminwed D (w, 1,+) and wi+1 = argmin A (w, w(t1,,))
We decompose (12) into the sum of two terms.
[[(VE(w))) Velwr), we41) []]
7
((((E(W)) Ve(wo), wtt - -w)) wi] + 1 - Ve(we), wh+1) |
< 7t
ve(wo))lll will [|l + "
Ell ve(We)llke4t well wi).
(44)
Next, we aim bound ||4441 - well. By Cs-strong convexity of A and definition of wf41, we have
||++1,* we+1ll S
V Cs
DA- (A-
D (w
V
Cs
<
V
(Da-(wr (wr, + D wx,+)
Os
V
(VA41)))11)) VA* wi+1,+ - wi)
Cs
2
V
(w (w:)l.
Cs
Moreover, by strong convexity we have
1
||++1,1 well <
IVA' (Vr1(1))1.)) -AA
Cs
Then by triangle inequality we get
V2+1
|| will || I ||w++1 W1+1,-l + || wills
| VA (wr+1,)) VA* (wx)).
Cs
Then use the definition of SMD step in (11) to get
V2+ +
V2+1
||++1 -wll
- vA*(wi)) <
(45)
Cs
Cs
Similarly, we have
V2+1
1 will <
Yell [[([((()))]
(46)
Cs
50

---

Plug (45) and (46) into (44) and we get
B((((w)) Ve(we)) - we+1)) wir
71
v2+1
El| V((wu) ve(wo)Il] [ + v2 +
E1
E B[] Ve(walll] [
Cs
Cs
v2+1
V VEllÍ(((.)) - ve(w))]1
sup (l]Ve( + ||([(w))]])
Cs
V2+ +1
4n/ (s U/32
2n
91
+
Us2
+
82
+ +3VOUD)
Cs
64
4
8
22 1
U82
2
$2
[88]] + 222s2 + 22)
12V Dnl
(11
Cs
4
64
51

---

Experiment Details and Additional Results
J.1 Pseudocode of the Algorithms
We first present the pseudocode Proj-SNGD algorithm proposed Section 4.1.
Algorithm Proj-SNGD
Require: Initialization wy E f2 number of iterations T. step sizes {1eJucis-11 1: for t = .,....... T 1 do
2: Compute stochastic gradient V((w))
3: Compute Wh+1,- E f2 such that
V (u1+11),,)) = (A4(w++1) 74 V((w))
(47)
4: Set
wi+1 = Proja (w)+1.o)
(48)
5: end for
6:
Sample wy from fwi JwiSoTT--- with probability 7 Ei= %
7:
Return WT
Next, we provide the pseudocode of Prox-SGD and Proj-SGD from [Dom20, DGG23J. We write
[()) L(,, C) E- log p(y)
z) - log p())].
Therefore, using the definition of L(0) in (30), we have
L(0) = L()) + [t: Dog (2;
Algorithm ProX-SGD
Require: Initialization Nu Rd,, Co S4 and diagonal, iterations T, step sizes fntoseer-1.
1: for =0 101...., - do
2: Compute stochastic gradient VnL(He, C.). VoI(H, C.)
3: Set Me+1 = V Ci), C++1, = C2 - vo Ct)
4: For each S < d, set
(C+t)a== (Ct)a v (C.) 47h
(49)
5: end for
6: Return (HT. Cr)
In Algorithm 3, M is defined as the smoothness parameter the joint likeeliodoood log p(z, D) with respect to Recall that the clipping function in (50) is defined clip (I) min/max(a, 6).
Algorithm 3 Proj-SGD
Require: Initialization 140 RRd Co E Sd and diagonal, iterations T, step sizes (r smoothness parameter M
for = 1,1,..,,'-11 do
2: Compute stochastic gradient V Liee Ci), VeL(M,.C) C.)
3: Set Mt+1 = #4 % VuL(ur, Co), Ct+1. = C. - nt VoL Ct)
4: For each i < d, set
(C++1)ss = clip (asPp//WMMtoo(CCHHHL..a ((C(11x)s)
(50)
5: end for
6: Return (pT: Cr)
52

---

10000
2100
Fies-SDD ProjSGD
Nummber of (temttions
Pri-SGD
Negatve ELIDO
5000
SNCE
ProjSVGD
SNGD
6000
1000
4000
210 3001
2001
Epoch
Figure 4: Euclidean and -on-Euclidean algorithms on Madelon dataset. Left: Objective during optimization with tuned step size. Right: Number of iterations before the objective falls below ((w) v 3000 for different initial step sizes "oo- non-Buclidean algorithms show consistently better performance, tolerate larger step sizes and are more robust to step size tuning.
Pro
Proj-SDD
Nurbber of Ioenatioss
14000
PGi
Negotive ELBO
120000
SNGE
100000
Froj-SNGI
12000
Poij-SGD
80020
40000
2000
1000
1500
2000
2500
2000
10
Epoch
Figure 5: Euclidean and non- Eunlidann algorithms CIFAR-10 10 dataset. Left: Objective during optimization with tuned step size. Right: Number of iterations before the objective falls below ((w) 35000 for different initial step sizes "0- Euclidean algorithms are more robust step size tuning and achieve faster convergence in the initial phase of optimization, but non-Euclidean algorithms exhibit faster convergence after 3000 epochs.
J.2 Experimental Results on Additional Datasets
In this section, compare the performance of Euclidean and non-Euclidean algorithms on Madelon and CIFAR-10 dataset.
Experiment on Madelon Dataset. In this experiment, we compare non-Euclidean algorithms (Proj-SNGD and SNGD) with Euclidean algorithms on Madelon dataset. Details of implemen- lation can be found in Section J.4.
We consider logistic regression Madelon dataset (n 2600, d = 500). We mini-batches of size 2000 and set the step size N 7 Vt, where 7o is hyperparameter to be tuned. We run the algorithm for 1000 epochs (2000 iterations). The results of 5 independent runs are shown in observe that Proj-SNGD slightly
Figure 4. Similar to the results of the MNIST experiment, outperforms SNGD, and both an-Euulidean algorithms achieve faster convergence than Euclidean counterparts. Moreover, Euclidean algorithms, especially Proj-SNGD, are more robust to step size and admit larger step sizes.
Experiment on CIFAR-10 Dataset. In this experiment, we consider logistic regression on sub- set of CIFAR-10 dataset with pictures of cats and dogs (n 10000, 3072). We use mini-batches of size 2000 and set the step size %/- = 7 /t,, where %o a hyperparameter be tuned. We run the algorithm for 3000 epochs (15000 iterations). The results of 5 independent runs are shown in Figure Figure 5. In the left panel of Figure 5, we observe that non-Euclidean algorithms converge faster during the initial phase (before 1500 epochs). However, Euclidean algorithms reach the optimum faster overall. the right panel of Figure Notably, Fon-Fuclidean algorithms can accommodate larger step sizes and are the most robust to step size tuning.
53

---

1100
Umi, D=20
Usio,
Ntumber of (bemtionn
0000
Umoo, Deog
D=
Negolive ELBO
Usd, D-15
Umed D=
D-10
4000
De
1000
3000
800
2000
1000
500
1000
Epoch
Figure 6: Convergence of Euclidacd algorithms on MNIST dataset. Left: Objective during optimization with "to 0.05. Right: Number of iterations before the objective falls below {(w) s 700 for different initial step sizes "/o-
10000
Umao,
2000
Usio, Doroo
Und Do-000
Nemther of Iterationns
T-3, D-
Nogative ELDO
6000
Uw2 D=
1000
D-200
T-1.5, D-210
1000
4000
100)
2141
300
Egoch
Figure 7: Convergence of non-Eud non-Euclidean algorithms on Madelon dataset. Left: Objective during optimization with To 0.5. Right: Number of iterations before the objective falls below ((w) S 3000 for different initial step sizes "/u-
J.3 Choice of U and D Proj-SNGD
In this section, we examine the effect of projection in the Proj-SNGD algorithm.
MNIST Dataset. We compare four different settings: no projection (U = oc, D oc), projection with U = 4, D = 20, U = 3, D = 15 and 2, D 10. Smaller values of U and D yield smaller smoothness coefficient and larger hidden convexity coefficient, leading to stronger theoretical guarantees.
Madelon Dataset. For this dataset, consider the following four different settings: projec- tion, projection with U = 3, D = 400, U = 2, D = 300 and U = 1.5, D = 200.
CIFAR-10 Dataset. For CIFAR-10 dataset, we adopt the same settings those used for the MNIST dataset: no projection, projection with U = 4, = 20, U = 3, D = 15 and U = 2, 10.
As shown in the left panel of Figure 6, we note that projection can accelerate convergence. However, when U and D are too small (e.g., U = 2 and D 10), Proj-SNGD may converge to sub-optimal solution. The right panel Figure 6 confirms that smaller U and D lead to faster convergence. Therefore, moderate choice of U and D provides both strong theoretical guarantees and favorable empirical performance. Similar phenomena can also be observed in Figures 7 and 8, where projec- tion with suitable choices of U and D) can accelerate and stabilize the training process. If U and D are too small, the optimal solution may fall outside the search space (see right panel of Figures and 8), and SNG will converge to a suboptimal solution. If they are too large, the algorithm will not benefit from projection (see SNGD without projection Figures 6 to 8).
The slow convergence of SNGD without projection is due to the poor tolerance for large step sizes, which causes divergence the initial phase. In contrast, projection stabilizes this in the initial phase and accelerates convergence (see the Poisson regression example in Figure 2). This effect consistently observed across different datasets. Moreover, this effect is consistent with our theoretical upper bounds for the relative smoothness parameter. This is because our upper bound
54

---

140000
Uneo0
Negalive ELDO
120000
Usen Dwild
Nutbber of (temttoon
1.000
U-d) D-10
D=
1000000
Des0
10000
D-10
S00000
12000
000000
11000
10000
10000
20000
9000
1000
2000
3000
Eposh
Figure 8: Convergence of non-Euclidean algorithms on CIFAR-10 10 dataset. Left: Objective during
optimization with "o 10-3 3 Right: Number of iterations before the objective falls below [(w) s 35000 for different initial step sizes YYo-
on depends on the diameter of the compact set, indicating that relative smoothness may not hold globally, which leads to divergence behavior in the initial phase.
J.4 Implementation Details
In all experiments involving Proj-SGD, we set M D (see Algorithm 3 for the definition of M) for a fair comparison with Proj- We use N = 2000 samples from variational distribution to estimate the expected log-likelihood (see the definition of stochastic gradient estimator (42)).
In the left panel Figure 3 (MNIST dataset), we set Tu 0.05 for Proj-SNGD and SNGD, and To = 0.01 for Prox-SGD and Proj-SGD. In the left panel of Figure 4 (Madelon dataset), we set To 0.5 for non-Euclidean algorithms. and "o 0.01 for Euclidean algorithms. In the left panel of Figure (CIFAR-10 dataset), we set To = 10 3 for non-Euclidean algorithms, and To = 2 x 10"H for Euclidean algorithms.
All experiments, except for 10, are conducted on an Apple M3 Pro CPU. On the MNIST dataset, training for 1000 epochs takes approximately 5 minutes, and on the Madelon dataset takes about I minute. The CIFAR-10 experiment is run on an NVIDIA GeForce RTX 3090 GPU, requiring roughly 8 minutes for 3000 epochs.
55


Reviews:
[{"review_id": "kfUmWXFKc1", "reviewer": "Reviewer_icec", "summary": "The paper proposes a theoretical study of Stochastic Natural Gradient Descent in the particular setting of Variational Inference, assuming a Gaussian prior and mean-field Gaussian variational family. The study begins by establishing some properties of the loss landscape (coerciveness, relative smoothness and hidden convexity). Convergence guarantees are then provided under fairly strong assumptions.", "questions": "* Lines 284-286: “empirical results indicate that […] even if Assumption 2 does not hold, the performance of Proj-SNGD is no worse and sometimes better than SGD.” Can you expand on this?\n* Is the big O characterization of  $L$ given by Theorem 3.2 practically relevant?\n* I don’t really understand how $(35)$ is obtained. What are the steps that lead to this equation?\n* In the Poisson regression experiment, what is the function value gap?\n* What is the difference between Figures 1 and 2? Does projecting achieve the same effect as changing the initialization in this example?", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n\nThe results overall are quite insightful, and the authors use a wide range of concepts and techniques/previously established results to support their analysis and prove their claims. The loss landscape analysis is of particular interest.\n\n**Weaknesses**\n* The main limitation, as emphasized by the authors, is Assumption 2. Unlike Assumptions 1 and 3, it cannot be checked directly and is susceptible to not being satisfied even in the restricted setting considered in the paper.\n* While the paper is well-structured, its clarity could be improved. This is partly because each result relies on its own set of technical assumptions. A dedicated section outlining all the assumptions would significantly improve readability and be a valuable addition. It would be helpful to have more than just the two (three, counting the appendices) assumptions explicitly stated in the paper.\n* The results are limited to specific problems in a purely Gaussian setting. They do not seem to be easily generalizable to other settings, as a significant part of the arguments rely on the properties of Gaussian distributions.\n* There is no explicit construction of the step-size sequence.\n\n**Additional remarks**\n* Given how much effort went into writing the appendices, a summary of the results already established and used in the proofs (when the proof consists in checking the application conditions of some theorem) would have been a nice touch.\n* In Section 2.2, $I$ denotes both the Fisher Information Matrix and the identity matrix.\n* Line 864: “1-relatively smoothness” -> 1-relative smoothness\n* Line 960: “the objective is coerciveness […] listed above” -> the objective is coercive […] listed below.", "was_revised": true, "final_justification": "This paper develops theoretical results on NGVI in non-conjugate settings, assuming a Gaussian prior and variational family. While some theoretical arguments are interesting and insightful, strong assumptions besides Gaussianity restrict the applicability of the guarantees obtained. The experiments conducted by the authors during the rebuttal period convincingly justify their method. However, this is mainly a theoretical paper, which is why I decided to keep my original score of 4 - Borderline accept.", "rebuttal": "We thank the reviewer for the careful reading and valuable feedback. Below, we address each of the weaknesses and questions in detail.\n\n### Weaknesses\n\n> The main limitation, as emphasized by the authors, is Assumption 2. Unlike Assumptions 1 and 3, it cannot be checked directly and is susceptible to not being satisfied even in the restricted setting considered in the paper.\n\nWe agree that Assumption 2 is primarily a technical condition and can be challenging to verify beforehand. The key condition we need in the proof of Theorem 4.2 is $\\omega^+=\\omega_*$ (see Appendix H.3), where $$\\omega^+\\coloneqq\\arg\\min_{\\omega' \\in\\tilde{\\Omega}}\\\\{\\langle\\nabla\\ell(\\omega),\\omega'\\rangle+2\\rho D_{A^\\*}(\\omega',\\omega)\\\\},\\quad \\omega_*\\coloneqq \\arg\\min_{\\omega'\\in\\Omega}\\\\{\\langle\\nabla\\ell(\\omega),\\omega'\\rangle+2\\rho D_{A^*}(\\omega',\\omega)\\\\}.$$When $\\omega\\in\\partial \\tilde\\Omega$, this condition does not necessarily hold, so we impose Assumption 2 to avoid this scenario.\n\nHowever, it is possible to replace Assumption 2 with another algorithm-independent assumption. The intuition is that if the gradients $\\nabla \\ell(\\omega)$ point towards the interior of $\\tilde \\Omega$ when $\\omega\\in\\partial \\tilde\\Omega$, then for sufficiently large $\\rho$, we have $\\omega_*\\in\\tilde\\Omega$, and hence $\\omega^+=\\omega_*$.\n\nSince $\\nabla A(\\eta)$ and $\\nabla A^*(\\omega)$ are inverse operators of one another (see Appendix C.1), (35) implies $$\\omega_\\*=(\\nabla A^\\*)^{-1}\\left(\\nabla A^\\*(\\omega)-\\frac{1}{2\\rho}\\nabla\\ell(\\omega)\\right)=\\omega-\\frac{1}{2\\rho}\\nabla^2 A(\\eta)\\nabla \\ell(\\omega)+o(1/\\rho)=\\omega-\\frac{1}{2\\rho}(\\nabla^2 A^\\*(\\omega))^{-1}\\nabla\\ell(\\omega)+o(1/\\rho)$$as $\\rho\\to\\infty$. Then we impose the following assumption.\n\n**Assumption 2'.** *For any $\\omega\\in\\partial \\tilde\\Omega$, let $\\mathbf{n}_\\omega$ be the outward normal direction at $\\omega$. Assume that$$\\left<-(\\nabla^2 A^\\*(\\omega))^{-1}\\nabla\\ell(\\omega),\\mathbf{n}_\\omega\\right><0.$$*\n\nLet $W\\coloneqq\\\\{\\omega_t:0\\leq t\\leq T\\\\}\\cap \\partial \\tilde\\Omega$ denote the set of iterates that lie on the boundary of $\\tilde \\Omega$. Note that $W$ is a finite set. If Assumption 2' is satisfied, there exists some $\\varepsilon>0$ such that $$\\left<-(\\nabla^2 A^\\*(\\omega))^{-1}\\nabla\\ell(\\omega),\\mathbf{n}_\\omega\\right><-\\varepsilon$$for all $\\omega\\in W$. Then we have $$\\left<\\mathbf{n}_\\omega, \\omega_*-\\omega\\right><-\\frac{\\varepsilon}{2\\rho}+o(1/\\rho)<0$$\n\nfor some sufficiently large $\\rho$. Hence $\\omega_*\\in\\tilde\\Omega$.\n\nIn the following, we present an example where Assumption 2' holds. We consider a univariate Bayesian linear regression model with a single data point $(x, y)$, where $x,y,z\\in \\mathbb{R},p(z)=\\mathcal{N}(0, 1)$ and $p(y|x,z)=\\mathcal{N}(xz, 1)$. It can be shown that Assumption 2' is satisfied if the following set of inequalities holds:$$-U(x^2+1)<xy<U(x^2+1),\\text{ and }D^{-1}<x^2+1<D.$$Therefore, by carefully choosing $U$ and $D$, Assumption 2' can be satisfied. We will include the derivation of this fact in the revision.\n\n> While the paper is well-structured, its clarity could be improved. This is partly because each result relies on its own set of technical assumptions. A dedicated section outlining all the assumptions would significantly improve readability and be a valuable addition. It would be helpful to have more than just the two (three, counting the appendices) assumptions explicitly stated in the paper.\n\nThank you for your suggestion. We will add a list of assumptions and a summary of key external results used in the proofs, especially when the proof mainly involves verifying these conditions. This will make the paper easier to follow and more readable.\n\n> The results are limited to specific problems in a purely Gaussian setting. They do not seem to be easily generalizable to other settings, as a significant part of the arguments rely on the properties of Gaussian distributions.\n\nWe acknowledge that our analysis relies on properties of the Gaussian family. We choose this setting due to its analytical simplicity and its widespread use in practice. This restriction is standard in the literature (e.g., [1,2]), as relaxing it would require fundamentally different techniques. We view extending the results beyond the Gaussian setting as an important direction for future work.\n\n> There is no explicit construction of the step-size sequence.\n\nThank you for the comment. Our theorems use specific step size schemes, which are standard in non-convex optimization literature. For Theorem 4.1, the step-size is constant and is provided in line 1251, Appendix H.2: $$\\gamma_t=\\gamma=\\min\\left\\\\{\\frac{1}{2L},\\sqrt{\\frac{\\lambda_0}{V^2LT}}\\right\\\\}.$$ For Theorem 4.2, we follow the same step size scheme as Theorem 4.7 in [3], which we did not explicitly mention in the paper:$$\\gamma_t=\\begin{cases}\n    \\frac{1}{2L},&\\text{if }t\\leq T/2\\text{ and }T\\leq\\frac{6L}{\\mu_B},\\\\\\\\\n    \\frac{6}{\\mu_B(t-\\lceil T/2\\rceil)+12L}, &\\text{otherwise}.\n\\end{cases}$$ In particular, it is important in both cases that $\\gamma_t \\leq 1/(2L)$ to avoid divergence. We acknowledge that this could cause confusion and will revise the main text to clearly state the step size schemes. \n\n### Questions\n\n> Lines 284-286: “empirical results indicate that […] even if Assumption 2 does not hold, the performance of Proj-SNGD is no worse and sometimes better than SGD.” Can you expand on this?\n\nIn our experiments (Figures 2 and 3), Assumption 2 is not satisfied. This can be inferred from the different loss behaviors of SNGD and Proj-SNGD: such differences indicate that projection occurred during Proj-SNGD, and when projection occurs, the iterate must lie on the boundary of $\\tilde\\Omega$. We will make this explicit in the revision to avoid confusion. Despite this, Proj-SNGD converges faster and demonstrates greater robustness to step size tuning. These observations suggest that Proj-SNGD can outperform SGD algorithms even when Assumption 2 does not hold.\n\n> Is the big O characterization of $L$ given by Theorem 3.2 practically relevant?\n\nThe Big-O characterization of $L$ in Theorem 3.2 is primarily of theoretical interest, as it provides a conservative estimate of the real smoothness parameter. It offers insight into how $L$ scales with problem parameters, including $d,U,D$ and the properties of the log likelihood function, which can guide step size selection in practice. In particular, the dependence on dimension $d$ is at most linear, and the dependence on $U$ and $D$ is polynomial of degree at most 3. This characterization is practically relevant because $L$ directly influences the convergence rate of Proj-SNGD, and implies the polynomial iteration complexity in these parameters.\n\n> I don’t really understand how (35) is obtained. What are the steps that lead to this equation?\n\n(35) is obtained by explicitly solving the optimization problem in line 1323. By definition, the Bregman divergence equals $D_{A^\\*}(\\omega',\\omega)=A^\\*(\\omega')-A^\\*(\\omega)-\\left<\\nabla A^\\*(\\omega),\\omega'-\\omega\\right>.$ Substituting it into the problem in line 1323 and applying the first-order optimality condition yields $\\nabla\\ell(\\omega)+2\\rho(\\nabla A^\\*(\\omega')-\\nabla A^\\*(\\omega))=0$, which coincides with (35). The second derivative $2\\rho \\nabla^2 A^\\*(\\omega')\\succ 0$ because $A^\\*$ is strictly convex.\n\n> In the Poisson regression experiment, what is the function value gap?\n\nFunction value gap at iteration $t$ is the difference between the function value at iteration $t$ and the optimal value, i.e., $\\ell(\\omega_t)-\\ell^*$. \n\n> What is the difference between Figures 1 and 2? Does projecting achieve the same effect as changing the initialization in this example?\n\nFigures 1 and 2 correspond to the same experimental setting. The left panel of Figure 1 is identical to the left panel of Figure 2. The right panel of Figure 1 illustrates the performance of SNGD under a different initialization, and the right panel of Figure 2 shows the performance of our proposed Proj-SNGD with the original initialization. Both projection and changing the initialization in this example can stabilize and accelerate the training process, but the training trajectories of the two are not identical. The main insight from this experiment is that the objective is less well-behaved in the nonconvex setting (Figure 1), and that our Proj-SNGD converges even with initializations for which original SNGD initially diverges (Figure 2).\n\n---\n\nWe also thank the reviewer for the helpful suggestions on clarity and presentation in the additional remarks, and we will incorporate these suggestions in the revision to improve readability. We believe that the reviewer’s feedback will help us improve the quality of the manuscript.\n\n### References\n\n[1] Kaiwen Wu and Jacob R. Gardner. Understanding Stochastic Natural Gradient Variational Inference. In International Conference on Machine Learning, pages 53398-53421. PMLR, 2024.\n\n[2] Navish Kumar, Thomas Möllenhoff, Mohammad Emtiyaz Khan, and Aurelien Lucchi. Optimization guarantees for square-root natural-gradient variational inference. Transactions on Machine Learning Research, 2025.\n\n[3] Ilyas Fatkhullin and Niao He. Taming nonconvex stochastic mirror descent with general bregman divergence. In International Conference on Artificial Intelligence and Statistics, pages 3493–3501. PMLR, 2024."}, {"review_id": "KWA7jdFGsW", "reviewer": "Reviewer_Kcs6", "summary": "The paper studies on theoretical aspects of stochastic natural gradient variational inference. It connects NGVI with the stochastic mirror descent and advance theoretical understanding in the case of non-conjugate likelihood settings and provide non-asymptotic error bounds.", "questions": "1. While the finite error bounds in Theorem 4.1 and Theorem 4.2 are established, they are conditional upon the choice of the step size. The current theorem only demonstrates the error bound for some step size (not necessarily any step size schedule). Can you further elaborate on influence of the step size choice in error bounds and their rate? Will it hold regardless of how large the initial step size is chosen?\n\n2. In the projection step introduced in the Section 4.1., one needs to calibrate the projection radius. How should practitioners decide projection radiuses? Can you perhaps elaborate more on the influence of the choice of projection radiuses? I understand that the authors have stated moderate values of U and D work empirically, but it's not clear what they mean by moderate values. Perhaps more extensive numerical experiments would be beneficial to justify the projection step.", "limitations": "Not Applicable.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths: The paper provides a solid theoretical understanding of NGVI, expanding the existing works to the setting where the variational loss can be non-convex. Characterizing a hidden convexity properties of the variational loss, it extends theoretical insights on the convergence of NGVI. The organization of the paper is well-structured with clear exposition.\n\nWeaknesses: No major significances, while I found some incomplete sentences to be edited, e.g., page 2, line 72, \"In the Euclidean set up, ...\" and page 5 line 167, \"the Hessian of the two functions coincide\".\n\nAlso, I suggest providing more references behind characterizations they are using. For instance, the duality of the natural parameter in the exponential family and the dual expectation parameter are stated without providing further references.", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful and helpful feedback. We will address each of the weaknesses and questions in detail.\n\n### Weaknesses\n\n> I found some incomplete sentences to be edited, e.g., page 2, line 72, \"In the Euclidean set up, ...\" and page 5 line 167, \"the Hessian of the two functions coincide\".\n\nThank you for pointing this out. We will revise these sentences in the revision.\n\n> Also, I suggest providing more references behind characterizations they are using. For instance, the duality of the natural parameter in the exponential family and the dual expectation parameter are stated without providing further references.\n\nThank you for your suggestion. Most of the content in Section 2.1 is based on [1], which provides detailed discussion of the duality between the natural and expectation parameters in exponential families. We will revise the paper to make this clearer.\n\n### Questions\n\n> While the finite error bounds in Theorem 4.1 and Theorem 4.2 are established, they are conditional upon the choice of the step size. The current theorem only demonstrates the error bound for some step size (not necessarily any step size schedule). Can you further elaborate on influence of the step size choice in error bounds and their rate? Will it hold regardless of how large the initial step size is chosen?\n\nOur theorems use specific step size schemes, which are standard in non-convex optimization literature. For Theorem 4.1, the step-size is constant and is provided in line 1251, Appendix H.2: $$\\gamma_t=\\gamma=\\min\\left\\\\{\\frac{1}{2L},\\sqrt{\\frac{\\lambda_0}{V^2LT}}\\right\\\\}.$$ For Theorem 4.2, we follow the same step size scheme as Theorem 4.7 in [2], which we did not explicitly mention in the paper:$$\\gamma_t=\\begin{cases}\n    \\frac{1}{2L},&\\text{if }t\\leq T/2\\text{ and }T\\leq\\frac{6L}{\\mu_B},\\\\\\\\\n    \\frac{6}{\\mu_B(t-\\lceil T/2\\rceil)+12L}, &\\text{otherwise}.\n\\end{cases}$$ In particular, it is important in both cases that $\\gamma_t \\leq 1/(2L)$ to avoid divergence. We acknowledge that this could cause confusion and will revise the main text to clearly state the step size schemes. \n\n> In the projection step introduced in the Section 4.1., one needs to calibrate the projection radius. How should practitioners decide projection radiuses? Can you perhaps elaborate more on the influence of the choice of projection radiuses? I understand that the authors have stated moderate values of $U$ and $D$ work empirically, but it's not clear what they mean by moderate values. \n\nBy \"moderate values\", we mean that:\n\n1. $U$ and $D$ should not be too small, otherwise the optimal solution may fall outside $\\tilde \\Omega$. For instance, in the left panel of Figure 4 on page 52, Proj-SNGD with $U=2,D=10$ converges to a suboptimal solution.\n\n2. $U$ and $D$ should not be too large, otherwise the algorithm will not benefit from projection. \n\nIn our experiments (see Appendix J.2), we run a small number of iterations with several candidate values of $U$ and $D$, and select the sets of parameters that lead to faster convergence.\n\n> Perhaps more extensive numerical experiments would be beneficial to justify the projection step.\n\nWe have conducted additional experiments on new datasets to further justify the projection step.\n\n1. Madelon dataset [3] ($n=2600,d=500$). For SNGD and Proj-SNGD algorithms with different choices of $U$ and $D$, we report the median number of iterations required for the loss to fall below 3000 (denoted as “NA” if not achieved in 2000 iterations); values in parentheses indicate the first and third quartiles (similar to Figure 4).\n\n|                      | $\\gamma_0=0.1$           | $\\gamma_0=0.2$           | $\\gamma_0=0.5$           |\n|:--------------------:|:-------------------------:|:-------------------------:|:-------------------------:|\n| $U=\\infty,D=\\infty$ | 1314 (1224, 1358)        | 530 (528, 546)           | 250 (228, 264)           |\n| $U=3,D=400$         | 914 (852, 956)           | 320 (318, 328)           | 108 (108, 108)           |\n| $U=2,D=300$         | 766 (756, 812)           | 270 (262, 270)           | 100 (98, 106)            |\n| $U=1.5,D=200$       | NA                       | NA                        | NA                        |\n\n2. A subset of the CIFAR-10 dataset [4] containing images of cats and dogs ($n=10000,d=3072$). Similarly, we report the median number of iterations required for the loss to fall below 35000 (denoted as “NA” if not achieved in 15000 iterations); values in parentheses indicate the first and third quartiles.\n\n|                      | $\\gamma_0 = 5 \\times 10^{-4}$    | $\\gamma_0 = 10^{-3}$          |\n|:--------------------:|:--------------------------------:|:-----------------------------:|\n| $U = \\infty, D = \\infty$ | NA                             | NA                            |\n| $U = 4, D = 20$          | 12206 (12138, 12528)            | 9127 (9073, 9177)             |\n| $U = 3, D = 15$          | 11703 (11522, 12007)            | 9253 (9039, 9387)             |\n| $U = 2, D = 10$          | NA                               | NA                             |\n\nThe results above suggest that projection with appropriately chosen hyperparameters can accelerate the convergence of NGVI.\n\n---\n\nThank you again for the insightful questions and helpful suggestions. We hope our responses have addressed your concerns, and we will revise the manuscript accordingly to improve the clarity and completeness of our manuscript.\n\n### References\n\n[1] Martin J Wainwright, Michael I Jordan, et al. Graphical models, exponential families, and variational inference. Foundations and Trends in Machine Learning, 1(1–2):1-305, 2008.\n\n[2] Ilyas Fatkhullin and Niao He. Taming nonconvex stochastic mirror descent with general bregman divergence. In International Conference on Artificial Intelligence and Statistics, pages 3493–3501. PMLR, 2024.\n\n[3] Isabelle Guyon. Madelon. UCI Machine Learning Repository, 2004. DOI:https://doi.org/10.24432/C5602H.\n\n[4] Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009."}, {"review_id": "M91C1S5V91", "reviewer": "Reviewer_nVxV", "summary": "The paper investigates the convergence properties of non-conjugate stochastic gradient descent on the variational inference ELBO leveraging recent work on the convergence of stochastic natural gradient descent (SNGD) for conjugate models but presently advanced to the non-conjugate setting. Notably, the proposed procedure includes a projection step (akin to how similarly a projection step of Proj-SGD is explored in [DGG23]) forming the proj-SNGD relying on theoretically derived results wrt. the smoothness properties of the associated ELBO and explores recent results from stochastic mirror descent of [FH24] to establish a convergence theory of the proposed proj-SNGD of O(1/sqrt(T)) with enhanced convergence of O(1/T) for concave log-likelihood functions. Notably, the manuscript also establishes non-trivial sufficient conditions for the relative smoothness of the ELBO function and properties wrt. hidden convexity (leveraging [FHH23]) in the present context for log-concave likelihood functions. On one synthetic and real dataset proj-SNGD is explored in practice and found to perform well when compared to standard SGD and SNGD notably with enhanced convergence and robustness to initial step size for the proj-SNGD procedure.", "questions": "Why are there no error bars on Figure 2 – the results appear somewhat anecdotal in the current presentation? -  it would be important here to see how the convergence behavior is across many initializations.\n\nThe experimentation is rather limited considering one synthetic and one real dataset. The paper would improve by establishing the convergence properties on more problems to highlight that the observed benefits especially of proj-SNGD when compared to SNGD can be observed in general. The gains are interesting but with the limited experimentation it could be established more convincingly by considering more problems than presently given in the main paper and small additional experimentation on MNIST given in the supplementary material. \n\nIt would further be interesting to increase the problem domains systematic from very smooth to less smooth and probe the strengths and limitations of the derived theoretical results relying on smoothness and projections as well as provide a general guidance on tuning U and D in the projection procedure as briefly investigated in the supplementary.", "limitations": "Yes, the line of research is foundational for inference in the non-conjugate VI setting and positive and negative societal impacts are not so relevant to discuss in this work.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strenghs\n* The paper is well-positioned in the literature\n* The approach is novel and with a strong technical depth.\n*  The problem considered is interesting and important providing theoretical results in the typical non-conjugate setting of VI which is a strongpoint.\n* The theoretical contribution is non-trivial and explores recent results in the literature and combines several interesting insights wrt. the relative smoothness of the ELBO and convergence of the proposed proj-SNGD procedure.\n* The limited experimentation supports the theoretical results and practical use of the proposed procedure.\n* The established theory is in depth described and derived with an extensive supplementary material providing necessary details. \n\nWeaknesses\n* The main weakness of the paper is the experimentation which is limited and could be more convincing (see questions).\n\n\nQuality: \nThe manuscript is well written, the results derived non-trivial, but the experimentation is limited and could be strengthened by considering more problems and datasets to highlight the general strength of the proposed framework which will also increase the impact and make the results more convincing.\n\n\nClarity: \nThe paper is generally well written and clear although the material covered extensive and with a high technical depth. There are a few minor typos such as \n\nThe realative smoothness -> the relative smoothness\n\nI further found the example below lemma 3.1. unclear. Can you please clarify the steps in the equalities between lines 176 and 177. I do not understand how the right hand side is obtained from the expressions in the middle. \n\nOn line 126 please establish the PL inequality when first used and describe what PL abbreviates to make the manuscript more self-contained.\n\n\nSignificance:\nThe proposed framework establishes convergence properties of non-conjugate VI inference based on stochastic natural gradient descent which is an important and non-trivial contribution to the field. The experimentation is rather limited and the gains of the projection step wrt. convergence properties when comparing the proposed proj-SNGD to the SNGD appears to improve upon convergence and robustness but the gains are somewhat limited and it would be interesting to see if there are cases where substantial gains can be achieved by the proposed projection step and how general these gains can be observed across models and datasets.\n\n\nOriginality:\nThe manuscript provides new and original results and non-trivial theorems establishing smoothness and convergence properties of SNGD in the non-conjugate setting which can warrant publication.", "was_revised": true, "final_justification": "I find that the authors' responses in their rebuttal sufficiently clarifies and address my concerns. I appreciate the additional experimentation provided including the investigation of influence of initialization and inclusion of experimentation on additional data sets as requested. In light of this and the authors' rebuttals and answers to the concerns of the other reviewers I maintain my positive score.", "rebuttal": "We sincerely appreciate the reviewer’s careful evaluation and insightful feedback. Below, we provide detailed responses to the identified weaknesses and questions.\n\n### Weaknesses\n\n> There are a few minor typos such as The realative smoothness -> the relative smoothness\n\nThank you for pointing this out. We will correct these typos.\n\n> Can you please clarify the steps in the equalities between lines 176 and 177. I do not understand how the right hand side is obtained from the expressions in the middle.\n\nSince $\\mu=\\xi,\\sigma^2=\\Xi-\\xi^2$, the chain rule yields $\\frac{\\partial \\mu}{\\partial \\xi}=1,\\frac{\\partial \\sigma^2}{\\partial \\xi}=-2\\xi, \\frac{\\partial \\mu}{\\partial \\Xi}=0$ and $\\frac{\\partial \\sigma^2}{\\partial \\Xi}=1$. Applying Lemma 3.1 together with the results above gives the formulas shown between lines 176 and 177.\n\n> On line 126 please establish the PL inequality when first used and describe what PL abbreviates to make the manuscript more self-contained.\n\nThank you for the suggestion. The abbreviation \"PL\" first appeared in line 216. We introduced the full term “Polyak–Lojasiewicz” in line 152, and in the revision we will make it clear at this point that we will use the abbreviation “PL.”\n\n### Questions\n\n> Why are there no error bars on Figure 2 – the results appear somewhat anecdotal in the current presentation? - it would be important here to see how the convergence behavior is across many initializations.\n\nIn Figure 2, Poisson regression admits closed-form gradients (see line 1005 in Appendix E). Therefore, for a fixed initialization, the setup is deterministic, and error bars are not necessary.\n\nTo examine the convergence behavior under different initializations, we conduct an additional experiment, where $\\mu_0\\sim \\mathrm{Unif}([-3, 0])$ is sampled 10 times, and we compare the performance of SNGD and Proj-SNGD with different step sizes. We report the median number of iterations required for the function value gap to fall below $2\\times 10^{-4}$; values in parentheses indicate the first and third quartiles. The results indicate that projection still accelerates convergence under different initializations. We will include these new results in the revised version of the paper.\n\n|            | Step size = 0.5      | Step size = 0.3      | Step size = 0.1       |\n|:----------:|:----------------------:|:----------------------:|:----------------------:|\n| **SNGD**   | 27 (20.5, 30)         | 33 (26.5, 35.5)       | 47 (47, 47.75)        |\n| **Proj-SNGD** | 9 (9, 9)            | 11 (11, 11)           | 47 (47, 47.75)        |\n\nFor $\\mu_0>0$, this effect is less pronounced, as SNGD exhibits more stable behavior in this regime.\n\n> The experimentation is rather limited considering one synthetic and one real dataset. The paper would improve by establishing the convergence properties on more problems to highlight that the observed benefits especially of proj-SNGD when compared to SNGD can be observed in general. The gains are interesting but with the limited experimentation it could be established more convincingly by considering more problems than presently given in the main paper and small additional experimentation on MNIST given in the supplementary material.\n\nWe have conducted additional experiments on new datasets to further highlight the advantages of Proj-SNGD over SNGD.\n\n1. Madelon dataset [1] ($n=2600,d=500$). For SNGD and Proj-SNGD algorithms with different choices of $U$ and $D$, we report the median number of iterations required for the loss to fall below 3000 (denoted as “NA” if not achieved in 2000 iterations); values in parentheses indicate the first and third quartiles (similar to Figure 4).\n\n|                      | $\\gamma_0=0.1$           | $\\gamma_0=0.2$           | $\\gamma_0=0.5$           |\n|:--------------------:|:-------------------------:|:-------------------------:|:-------------------------:|\n| $U=\\infty,D=\\infty$ | 1314 (1224, 1358)        | 530 (528, 546)           | 250 (228, 264)           |\n| $U=3,D=400$         | 914 (852, 956)           | 320 (318, 328)           | 108 (108, 108)           |\n| $U=2,D=300$         | 766 (756, 812)           | 270 (262, 270)           | 100 (98, 106)            |\n| $U=1.5,D=200$       | NA                       | NA                        | NA                        |\n\n2. A subset of the CIFAR-10 dataset [2] containing images of cats and dogs ($n=10000,d=3072$). Similarly, we report the median number of iterations required for the loss to fall below 35000 (denoted as “NA” if not achieved in 15000 iterations); values in parentheses indicate the first and third quartiles.\n\n|                      | $\\gamma_0 = 5 \\times 10^{-4}$    | $\\gamma_0 = 10^{-3}$          |\n|:--------------------:|:--------------------------------:|:-----------------------------:|\n| $U = \\infty, D = \\infty$ | NA                             | NA                            |\n| $U = 4, D = 20$          | 12206 (12138, 12528)            | 9127 (9073, 9177)             |\n| $U = 3, D = 15$          | 11703 (11522, 12007)            | 9253 (9039, 9387)             |\n| $U = 2, D = 10$          | NA                               | NA                             |\n\nThe results above suggest that, with properly chosen hyperparameters $U$ and $D$, Proj-NGVI can outperform NGVI.\n\n> It would further be interesting to increase the problem domains systematic from very smooth to less smooth and probe the strengths and limitations of the derived theoretical results relying on smoothness and projections as well as provide a general guidance on tuning U and D in the projection procedure as briefly investigated in the supplementary.\n\nThank you for raising this practical point. An illustrative experiment on the effect of the domain for the MNIST dataset is presented in Figure 4, Appendix J.2. Based on this, a broad guideline for the choice of $U$ and $D$—intended as intuition rather than a strict prescription—is as follows:\n\n(1) $U$ and $D$ should not be too small, otherwise the optimal solution may fall outside $\\tilde \\Omega$. For instance, in the left panel of Figure 4, Proj-SNGD with $U=2,D=10$ converges to a suboptimal solution.\n\n(2) $U$ and $D$ should not be too large, otherwise the algorithm will not benefit from projection. \n\nIn our experiments, we run a small number of iterations with several candidate values of $U$ and $D$, and select the sets of parameters that lead to faster convergence.\n\n---\n\nWe appreciate the reviewer for thoughtful remarks and suggestions, which we believe will improve the overall quality of the manuscript.\n\n### References\n\n[1] Isabelle Guyon. Madelon. UCI Machine Learning Repository, 2004. DOI:https://doi.org/10.24432/C5602H.\n\n[2] Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009."}, {"review_id": "CFtz8S5tXI", "reviewer": "Reviewer_tPUM", "summary": "The paper analyzes the convergence of natural gradient variational inference (NGVI) with stochastic gradients applied to non-conjugate models. Obtaining convergence guarantees for stochastic NGVI under realizable/verifiable assumptions has been an open problem for decades. The main challenge has been the fact that, under typical smoothness and log-concavity assumptions on the target, the ELBO objective is no longer convex and smooth with respect to the canonical variational parametrizations of NGVI. This paper addresses this open issue by constraining the variational parameters to a compact set. This enables the ELBO to be relatively smooth and even Polyak-Lojasiewicz (PL). In light of this, for log-smooth and log-Hessian-smooth likelihoods, the paper reports a $O(1/\\sqrt{T})$ convergence rate for reducing the squared gradient norm of the ELBO. If the likelihood is additionally strongly log-concave, the paper reports a $O(1/T)$ rate for reducing ELBO. Some experimental results in favor of confining the parameters to a compact set are demonstrated.", "questions": "n/a", "limitations": "The limitation section in the paper is appropriate.", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "### Strengths\n* The paper demonstrates that the ELBO can be relatively smooth if the variational parameters are confined to a compact set. As I haven't seen this result before, this is new and potentially interesting.\n* The fact that the ELBO can satisfy PL due to hidden convexity is interesting.\n\n### Weaknesses\n* The fact that the ELBO is relatively smooth under compactness is unsurprising. Compactness is very strong, and most regularity assumptions start to hold under this assumption due to continuity. However, in addition to this, the relative smoothness result requires both the log-likelihood and its gradient to be Lipschitz. The need for both at the same time is quite strong.\n* It is unclear to me if Assumption 1 can be simply assumed away. In [1], Assumption 1 was explicitly established by studying special cases. However, that justification does not apply here since [1] did not study NGVI with Monte Carlo gradients. Given that the assumptions used in the paper are already strong enough, I would expect Assumption 1 to be verified explicitly.\n* The experiments in Section 5 have a correctness issue. In particular, the implementation of Prox-SGD seems to be incorrect, which is why it appears that it is not converging in Figure 3. I checked the code, and the error appears to be that the stochastic gradient used for Prox-SGD is differentiating the whole ELBO $\\nabla \\ell(z)$ instead of just the energy term $\\nabla_{w} \\mathbb{E}_{z \\sim q(\\cdot; w)} -\\log p(D|z)$. In Prox-SGD, the proximal operator handles the entropy term. Therefore, the gradient should only be taken on the energy term. Otherwise, the algorithm ends up minimizing the entropy term twice, which is what seems to be happening.\n* It is unclear why Assumption 2 is needed and what it implies. Why has this to be assumed when projections to a compact set are already being used?\n\nThe problem of proving the convergence of stochastic NGVI on non-conjugate models is a genuinely difficult problem where very little progress has been made. Therefore, showing any useful theoretical results is a significant contribution. However, the set of assumptions are really strong to shed much light into the behavior of NGVI. Still, I think any theoretical result on this problem could serve as an initial step for future works. As such, I am in favor of accepting the submission. The reason for the rejection-leaning score is due to the correctness concern in the experiments and the unclear meaning of Assumption 2. If these concerns are addressed, I will happily raise my score.\n\n## Major Comments\n* The title is too non-descriptive. It is not obvious that this paper is about a theoretical analysis of the convergence of NGVI.\n* Discuss the assumptions in the abstract, especially since they are strong. The fact that the domain is enforced to be bounded and both the log-likelihood and its gradient need to be Lipschitz continuous should be stated.  Citing the parametrization considered in the paper would also be useful.\n* What is also not clear to me is why the paper focuses on joint likelihoods with an explicit prior-likelihood pair. Since the paper considers the non-conjugate setting, it should suffice to just consider a generic unnormalized target, which should simplify the setup and the notation. Similarly, the regularity results like Theorem 3.2  and Proposition 2,3 assume regularity of the likelihood $\\log p(\\mathcal{D}\\mid z)$. This doesn't seem to be necessary. Why not just assume the regularity of the joint? (and equivalently, the \"target unnormalized density\")\n* Line 211: Isn't this about \"strong\" convexity, not just convexity? If it is the case, this should be made more precise.\n* Proposition 2: I think the constant $\\mu_C$ can easily be improved. If the diagonal of a triangular matrix is strictly positive, then the diagonal elements correspond to its eigenvalues. Also, the eigenvalues coincide with the singular values. Then I think, for all $i$, \\sigma_{\\mathrm{max}}(G_i) \\leq \\max\\left(1, 2C_{11}, \\ldots, 2C_{dd}\\right) \\leq 1 + 2 \\max_{j} C_{jj} \\leq 1 + 2 \\sqrt{D}$ holds?\n\n## Minor Comments\n* Line 26: BBVI, essentially non-conjugate ELBO maximization with SGD, was concurrently developed by [1,2] as well.\n* Line 29: While I know that the term \"Euclidean gradient descent\" has been used in some works, I personally feel that this is not descriptive enough. (Which Euclidean space?) Gradient descent in the space of parameters is probably a better term in my opinion.\n* Line 31: Cite something for NVGI. My understanding is that Amari's paper doesn't specifically describe VI with natural gradients. Therefore, NGVI should be attributed to some other paper.\n* Line 38 \"where the posterior is not contained in the variational family\": This statement doesn't seem necessary.\n* Line 42: Cite something for mirror descent.\n* Line 72: Convexity was established in [2]. Domke 2020 generalized their result to strongly convex targets.\n* Line 184: Why not change $U$, $D$ to $M$, $S$ so that $M$ corresponds to the bound on $\\mu$ and $S$ corresponds to the bound on $\\Sigma$.\n* Line 297: The use of projected and proximal SGD was first proposed in [3].\n\n1. Wingate, D., & Weber, T. (2013). Automated variational inference in probabilistic programming. arXiv preprint arXiv:1301.1299.\n2. Titsias, M., & Lázaro-Gredilla, M. (2014, June). Doubly stochastic variational Bayes for non-conjugate inference. In International conference on machine learning (pp. 1971-1979). PMLR.\n3. Domke, J. (2020, November). Provable smoothness guarantees for black-box variational inference. In International Conference on Machine Learning (pp. 2587-2596). PMLR.", "was_revised": true, "final_justification": "The authors resolved my original major concerns on the correctness of the experimental results and the peculiar Assumption 2. Therefore, I now lean towards accepting the paper.", "rebuttal": "Thank you for taking the time to thoroughly review our manuscript and share constructive suggestions. We address each of your weaknesses and major comments in detail below.\n\n### Weaknesses \n\n> The fact that the ELBO is relatively smooth under compactness is unsurprising. Compactness is very strong, and most regularity assumptions start to hold under this assumption due to continuity. However, in addition to this, the relative smoothness result requires both the log-likelihood and its gradient to be Lipschitz. The need for both at the same time is quite strong.\n\nWe agree that compactness is a strong assumption, and that requiring both the log-likelihood and its gradient to be Lipschitz is restrictive. These conditions partly reflect the inherent complexity of NGVI under the relative smoothness framework. As illustrated in Appendix F, for univariate case, while it is possible to relax the compactness assumption on $\\mu$, the necessary and sufficient conditions for relative smoothness involve higher-order (third and fourth order) derivatives of the log-probability, highlighting the technical difficulty of characterizing smoothness in this setting. We will explicitly address the limitations of these assumptions in the revision and emphasize that developing approaches to relax them is an important direction for future work.\n\n> It is unclear to me if Assumption 1 can be simply assumed away. In [1], Assumption 1 was explicitly established by studying special cases. However, that justification does not apply here since [1] did not study NGVI with Monte Carlo gradients. Given that the assumptions used in the paper are already strong enough, I would expect Assumption 1 to be verified explicitly.\n\nThank you for raising this point, and we agree that Assumption 1 requires careful verification. In Appendix I, we have explicitly checked that Assumption 1 holds for logistic regression model. We will make this clarification more explicit in the revision. \n\nWe did not find Assumption 1 in [1] and assume you meant to refer to some other paper. Could you please share the correct reference?\n\n> The experiments in Section 5 have a correctness issue.\n\nThank you very much for catching this issue. You are correct that there was an error in the gradient computation for Prox-SGD. We have corrected this mistake and rerun all the experiments. The updated results on MNIST dataset (see below) show that Prox-SGD converges as expected and achieves performance comparable to Proj-SGD. We will include these corrected results in the revised version. Thank you again for pointing this out.\n\nWe report the median negative ELBO after 1000 epochs (6000 iterations) across 5 runs with tuned step size; values in parentheses indicate the first and third quartiles. These results confirm that Prox-SGD converges with the corrected implementation and achieves performance comparable to Proj-SGD.\n\n|        | Prox-SGD                      | Proj-SGD                      | SNGD                          | Proj-SNGD                      |\n|:------:|:-----------------------------:|:-----------------------------:|:-----------------------------:|:------------------------------:|\n| $-\\text{ELBO}$ | 691.07 (688.26, 691.75)        | 691.99 (689.44, 692.32)        | 696.27 (695.00, 696.91)        | 691.19 (690.95, 691.85)         |\n\nWe also report the median number of iterations required for the loss to fall below 700 with tuned step size; values in parentheses indicate the first and third quartiles. The results below also show that Prox-SGD converges at a rate comparable to Proj-SGD. However, SNGD algorithms converge much faster than SGD algorithms.\n\n|               | Prox-SGD               | Proj-SGD               | SNGD                    | Proj-SNGD              |\n|:--------------|:----------------------:|:----------------------:|:-----------------------:|:-----------------------:|\n| **# Iterations** | 1905 (1843, 2050)       | 2121 (2109, 2169)       | 1087 (910, 1166)         | 764 (681, 783)          |\n\n\n> It is unclear why Assumption 2 is needed and what it implies. Why has this to be assumed when projections to a compact set are already being used?\n\nThe key condition we need in the proof of Theorem 4.2 is $\\omega^+=\\omega_*$ (see Appendix H.3), where $$\\omega^+\\coloneqq\\arg\\min_{\\omega' \\in\\tilde{\\Omega}}\\\\{\\langle\\nabla\\ell(\\omega),\\omega'\\rangle+2\\rho D_{A^\\*}(\\omega',\\omega)\\\\},\\quad \\omega_*\\coloneqq \\arg\\min_{\\omega'\\in\\Omega}\\\\{\\langle\\nabla\\ell(\\omega),\\omega'\\rangle+2\\rho D_{A^*}(\\omega',\\omega)\\\\}.$$When $\\omega\\in\\partial \\tilde\\Omega$, this condition does not necessarily hold, so we impose Assumption 2 to avoid this scenario.\n\nHowever, it is possible to replace Assumption 2 with another algorithm-independent assumption. The intuition is that if the gradients $\\nabla \\ell(\\omega)$ point towards the interior of $\\tilde \\Omega$ when $\\omega\\in\\partial \\tilde\\Omega$, then for sufficiently large $\\rho$, we have $\\omega_*\\in\\tilde\\Omega$, and hence $\\omega^+=\\omega_*$.\n\nSince $\\nabla A(\\eta)$ and $\\nabla A^*(\\omega)$ are inverse operators of one another (see Appendix C.1), (35) implies $$\\omega_\\*=(\\nabla A^\\*)^{-1}\\left(\\nabla A^\\*(\\omega)-\\frac{1}{2\\rho}\\nabla\\ell(\\omega)\\right)=\\omega-\\frac{1}{2\\rho}\\nabla^2 A(\\eta)\\nabla \\ell(\\omega)+o(1/\\rho)=\\omega-\\frac{1}{2\\rho}(\\nabla^2 A^\\*(\\omega))^{-1}\\nabla\\ell(\\omega)+o(1/\\rho)$$as $\\rho\\to\\infty$. Then we impose the following assumption.\n\n**Assumption 2'.** *For any $\\omega\\in\\partial \\tilde\\Omega$, let $\\mathbf{n}_\\omega$ be the outward normal direction at $\\omega$. Then it holds that$$\\left<-(\\nabla^2 A^\\*(\\omega))^{-1}\\nabla\\ell(\\omega),\\mathbf{n}_\\omega\\right><0.$$*\n\nLet $W\\coloneqq\\\\{\\omega_t:0\\leq t\\leq T\\\\}\\cap \\partial \\tilde\\Omega$ denote the set of iterates that lie on the boundary of $\\tilde \\Omega$. Note that $W$ is a finite set. If Assumption 2' is satisfied, there exists some $\\varepsilon>0$ such that $$\\left<-(\\nabla^2 A^\\*(\\omega))^{-1}\\nabla\\ell(\\omega),\\mathbf{n}_\\omega\\right><-\\varepsilon$$for all $\\omega\\in W$. Then we have $$\\left<\\mathbf{n}_\\omega, \\omega_*-\\omega\\right><-\\frac{\\varepsilon}{2\\rho}+o(1/\\rho)<0$$\n\nfor some sufficiently large $\\rho$. Hence $\\omega_*\\in\\tilde\\Omega$.\n\nIn the following, we present an example where Assumption 2' holds. We consider a univariate Bayesian linear regression model with a single data point $(x, y)$, where $x,y,z\\in \\mathbb{R},p(z)=\\mathcal{N}(0, 1)$ and $p(y|x,z)=\\mathcal{N}(xz, 1)$. It can be shown that Assumption 2' is satisfied if the following set of inequalities holds:$$-U(x^2+1)<xy<U(x^2+1),\\text{ and }D^{-1}<x^2+1<D.$$Therefore, by carefully choosing $U$ and $D$, Assumption 2' can be satisfied.\n\n### Major Comments\n\n> The title is too non-descriptive. It is not obvious that this paper is about a theoretical analysis of the convergence of NGVI.\n\nWe will revise our title to a more descriptive one reflecting the main theoretical contribution of the paper, e.g., ``Natural Gradient VI: Guarantees for Non-Conjugate Models''. \n\n> Discuss the assumptions in the abstract, especially since they are strong. The fact that the domain is enforced to be bounded and both the log-likelihood and its gradient need to be Lipschitz continuous should be stated. Citing the parametrization considered in the paper would also be useful.\n\nThank you for this suggestion. We will revisit the abstract mentioning the key assumptions and the parameterization considered in the paper.\n\n> What is also not clear to me is why the paper focuses on joint likelihoods with an explicit prior-likelihood pair. Since the paper considers the non-conjugate setting, it should suffice to just consider a generic unnormalized target, which should simplify the setup and the notation. Similarly, the regularity results like Theorem 3.2 and Proposition 2,3 assume regularity of the likelihood $\\log p(\\mathcal{D}|z)$. This doesn't seem to be necessary. Why not just assume the regularity of the joint? (and equivalently, the \"target unnormalized density\")\n\nThank you for this observation. We adopt the explicit prior–likelihood decomposition in order to leverage the structure of the KL term: by isolating $D_{\\mathrm{KL}}(q(z)\\|p(z))$ we can directly verify that this term is 1-1 relative smooth with respect to $A^*$ (see line 166). The remaining term is the expected negative log-likelihood $-\\mathbb{E}_q[\\log p(\\mathcal{D}|z)]$, for which we establish separate regularity assumptions.\n\n> Line 211: Isn't this about \"strong\" convexity, not just convexity? If it is the case, this should be made more precise.\n\nYes, you are correct. In this case, $\\mu_H>0$, which corresponds to strong convexity rather than mere convexity.\n\n> Proposition 2: I think the constant $\\mu_C$ can easily be improved. If the diagonal of a triangular matrix is strictly positive, then the diagonal elements correspond to its eigenvalues. Also, the eigenvalues coincide with the singular values. Then I think, for all $i$, $\\sigma_{\\mathrm{max}}(G_i)\\leq\\max\\left(1,2C_{11},\\ldots,2C_{dd}\\right)\\leq1+2\\max_{j}C_{jj}\\leq 1+2\\sqrt{D}$ holds?\n\n\nThank you for your thoughtful attempt to improve the bound. However, the argument does not apply in this case because $G_i$ is not symmetric, so its eigenvalues do not necessarily coincide with its singular values. The singular values must therefore be computed directly from the definition.\n\n---\n\nWe also thank the reviewer for the detailed suggestions on citations, terminology and presentation in minor comments. We have already included additional related work in Appendix B, where more papers on NGVI and mirror descent are cited. \n\nWe hope our responses have addressed your concerns, and we will incorporate your constructive and thoughtful remarks in the revision, which we believe will greatly improve the overall quality of our work.\n\n### References\n\n[1] David Wingate and Theophane Weber. Automated variational inference in probabilistic programming. arXiv preprint arXiv:1301.1299, 2013."}]

Author rebuttal:

Meta-review: This paper analyzes stochastic nature gradient descent in the context of variational inference. In particular, it considers the setting of non-conjugate models, where the loss is nonconvex. In this process, a modified NGVI algorithm is proposed that includes non-Euclidean projections, and the convergence of this method is established. The paper also provides an empirical study on regression problems.

Reviewers initially had mixed opinions on the paper. The main reason for this was because of the limited setting of the theoretical results: reviewers pointed out many assumptions were quite strong, e.g., relying on Gaussianity with additional restrictions, such as compactness and bounded derivatives. However, after the discussion phase with the authors, reviewers were more positive about accepting the paper, as the authors cleared up some questions about the assumptions and experiments; the authors also provided additional experiments. 

During the reviewer discussion phase, the reviewers agreed that despite the limited applicability of the results, the paper provides valuable insights into stochastic NGDVI. Thus, I recommend accepting this paper. We encourage the authors to revise the paper according to the discussion with the reviewers.
Decision: Accept (poster)