Title: Natural Gradient VI: Guarantees for Non-Conjugate Models
Venue: NeurIPS 2025
Accepted: yes
Keywords: Variational inference

natural gradient

mirror descent

non-convex 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.

3

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 nonconjugate 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.

1 Introduction

Variational Inference (VI) is a powerful framework for approximating Bayesian posteriors by casting inference as an optimization problem [JGJS99, BKM17]. Unlike sampling-based approaches such as Markov chain Monte Carlo (MCMC; [MRR + 53]) VI enables scalable posterior approximation through stochastic optimization, often with significant computational advantages. In practice, however, the exact gradient of the variational objective-such as the evidence lower bound (ELBO)-is 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 BlackBox VI [WW13, TLG14, RGB14], this form of gradient-based 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 of Natural Gradient Variational Inference (NGVI) [HTRK07], which replaces standard gradients with natural gradients [Ama98] that account for the underlying information geometry of the variational family. NGVI follows the steepest descent direction (on average) with respect to the Kullback-Leibler (KL) divergence, rather than the Euclidean norm. Empirical evidence

0

10

20

30

40

50

10

-

3

10

-

1

10

1

10

3

10

5

Function Value Gap

step size=0.5

step size=0.3

step size=0.1

0

10

20

30

40

50

10

-

3

10

-

1

10

1

10

3

10

5

Function Value Gap

step size=0.5

step size=0.3

step size=0.1

Iteration

Iteration

suggests that this geometry-aware approach leads to faster or 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 ), a basic variant of NGVI, under conjugate likelihoods, by framing SNGD as a 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 or Poisson regression, where the objective becomes non-convex and the geometry is not well-behaved. For instance, Figure 1 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 in non-conjugate models.

Contributions. This work develops a theoretical foundation for stochastic natural gradient variational inference (NGVI) in 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 a compact set of parameters. Our results apply to a broad class of non-conjugate models in mean-field parameterization, including logistic regression, and provide explicit smoothness constants (polynomial in the problem dimension) over arbitrary compact subsets of the parameter space.

Projected Stochastic NGVI Algorithm. Based on the relative smoothness analysis, we 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 algorithm 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 non-asymptotic convergence to a stationary point at rate O (1 / √ T ) , and further show that under concavity of the log-likelihood, a hidden convexity structure [FHH25] allows us to prove global convergence at rate O (1 /T ) .

By extending convergence guarantees to the non-conjugate regime, this work addresses a fundamental challenge in the theory of NGVI and provides a 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, Dom20], authors of [DGG23] propose two gradient descent-based algorithms that achieve convergence rates of O (1 / √ T ) using convexity, and O (1 /T ) using strong convexity. Concurrently, Kim et al. [KOW + 23] analyze the convergence behavior of gradient-based VI under various parameterization 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] establish the first O (1 /T ) convergence rate of NGVI assuming the model is 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, HBWP13, LT21], particularly for non-conjugate models [KBL + 16, KL17, SEH18, TR19]. Notably, [LKS19] and [ADY + 23] extend natural gradient methods to mixtures of exponential-family variational distributions, while [CAAK19] and [JCM24] focus on their application in online learning scenarios. In reinforcement learning, a series of works has focused on natural policy gradient [Kak01, AKLM21, Xia22, SLZ + 23]. Natural gradient variational inference (NGVI) has also been applied across various domains, including the training of variational Bayesian neural networks [ZSDG18, OSK + 19], Kalman filtering [LZM + 25], and multimodal optimization [MAM + 25].

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 in analysis of SMD [HR21, VYB + 22]. In the non-convex setting, [GLZ16] derive an O (1 / √ T ) convergence rate to a stationary point, albeit under very strong assumptions and with large mini-batches. 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 generating 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 typically highly non-convex but admits certain structures analogous to our problem, see, e.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, b ∈ R d , let a ⊙ b denote their entrywise (Hadamard) product. Let S d + denote the cone of d × d symmetric positive definite matrices. For symmetric matrices A,B ∈ R d × d , we write A ⪰ B if A -B ∈ S d + . For distributions p and q , we write D KL ( q ∥ p ) = E q ( z ) log( q ( z ) /p ( z )) for the Kullback-Leibler divergence from p to q . When needed, we write q ( z ; θ ) to emphasize the parameterization of the distribution.

2.1 Variational Inference and Exponential Families

Let z ∈ R d be a latent variable and D = { ( x i , y i ) } n i =1 be the observed dataset. Given a prior p ( z ) and likelihood model p ( D| z ) , the goal of Bayesian inference is to compute the posterior p ( z | D ) ∝ p ( z ) p ( D| z ) . Since this posterior is generally intractable, variational inference (VI) approximates it by minimizing the KL divergence from p ( z | D ) to a tractable distribution q ( z ) within a family Q :

This is equivalent to minimizing the negative evidence lower bound (negative ELBO) [BN06]:

In this work, we assume Q is a minimal exponential family [WJ + 08], i.e., each q ∈ Q has the form

where η is the natural parameter, ϕ ( z ) is the sufficient statistic, and A ( η ) is the log-partition function. The dual expectation parameter is defined as ω := E q ( z ; η ) [ ϕ ( z )] .

Let D A := { η ∈ R d : A ( η ) < ∞} be the natural parameter domain, and let Ω be the set of corresponding expectation parameters. The gradient map ∇ A : D A → Ω is bijective, with inverse ∇ A ∗ , where A ∗ ( ω ) := sup η ∈D A ( ⟨ η, ω ⟩ -A ( η )) is the convex conjugate of the log-partition function. This duality ensures that for any distribution q ∈ Q , parameterized by either η or ω , we have

For a comprehensive overview of exponential families and duality, we refer readers to [WJ + 08].

Remark. With a slight abuse of notation, we denote the negative ELBO as ℓ ( η ) := ℓ ( q ( z ; η )) or ℓ ( ω ) := ℓ ( q ( z ; η )) when working in either parameterization.

Example: Gaussian Family. For the multivariate Gaussian family q = N ( µ, Σ) , the natural parameter is η = ( λ, Λ) with λ = Σ -1 µ and Λ = -1 2 Σ -1 , and the expectation parameter is ω = ( ξ, Ξ) with ξ = µ and Ξ = Σ+ µµ ⊤ . The standard parameterization ( µ, Σ) and Cholesky form ( µ, 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 non-Euclidean geometries by replacing the squared Euclidean distance with a Bregman divergence.

Definition 2.1 (Bregman Divergence) . Given a strictly convex and differentiable function Φ : U → R , the Bregman divergence between u, v ∈ U is D Φ ( u, v ) := Φ( u ) -Φ( v ) -⟨∇ Φ( v ) , u -v ⟩ .

The update rule for SMD on a differentiable objective ℓ : U → R is in

where ˆ ∇ ℓ ( u t ) is a stochastic gradient and γ t is a step size.

In variational inference with exponential families, the natural parameter η admits a geometry governed by the Fisher information matrix I ( η ) = ∇ 2 A ( η ) , where A ( η ) is the log-partition function. Stochastic NGVI in this space preconditions the gradient by the inverse Fisher matrix [RM15]. Applied to the variational objective ℓ , the update becomes:

where ω t = ∇ A ( η t ) is the corresponding expectation parameter, see Section C.2 for details. Using the duality between η and ω , and the identity η = ∇ A ∗ ( ω ) , we can re-write the update in the expectation parameter space as

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 ω , with geometry induced by KL divergence. Indeed, for exponential families, the Bregman divergence associated with A ∗ coincides with KL divergence: D A ∗ ( ω, ω ′ ) = D KL ( q ( z ; ω ) ∥ q ( z ; ω ′ )) , see e.g., [WG24].

2.3 Mean-field Parameterization

In the following, we assume that the prior is standard Gaussian, p ( z ) = N (0 , I ) , and the variational family Q is the mean-field Gaussian family, with mean µ = ( µ i ) 1 ≤ i ≤ d and diagonal covariance matrix Σ = diag(( σ i ) 1 ≤ i ≤ d ) . The corresponding expectation parameter ω = ( ξ, Ξ) is given by ξ = µ and Ξ = Σ + diag( µ ⊙ µ ) , defined over the domain

This parameterization is widely adopted due to its tractability and interpretability, allowing efficient coordinate-wise updates while still capturing key aspects of the posterior distribution. Moreover, it often provides a favorable trade-off between computational complexity and approximation quality in high-dimensional settings. It 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 ℓ ( ω ) with a particular focus on the properties useful for establishing non-asymptotic convergence of NGVI. First, we note that ℓ ( ω ) is coercive in the expectation parameter, i.e., it grows to infinity ℓ ( ω ) → ∞ when the parameters approach the boundary of the set ω → ∂ Ω , see Section D for more details. This is a useful property which guarantees the existence of the minima of ℓ ( ω ) . Second, it is known that ℓ ( ω ) is non-convex w.r.t. ω in the general non-conjugate likelihood setting, i.e., when likelihood does not belong to the variational family Q , see Section E for details. This non-convexity is the main challenge for establishing non-asymptotic convergence of NGVI. In what follows we aim to use modern tools from non-convex optimization, which will allow us to provide a better understanding of NGVI landscape and characterize its convergence. The rest of the section is organized as follows. In Section 3.1, we will prove that, despite non-convexity, both the lower and upper curvature of ℓ ( ω ) 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 Polyak-Lojasiewicz (PL) condition [Pol63, Loj63], which allows us to show fast convergence of NGVI despite non-convexity.

3.1 Relative Smoothness of Variational Objective

We start with the definition of α -β relative smoothness [BBT17, LFN18], which is a generalization of smoothness and weak convexity to Bregman geometry,

Definition 3.1. Let Φ : U → R be a differentiable and strictly convex function on a convex set U . A differentiable function ℓ is said to be α -β smooth relative to Φ for some α ∈ R , β > 0 if

If α ≥ 0 , ℓ is also called α -strongly convex relative to Φ .

In our problem we will mostly deal with the case of negative curvature α < 0 . Then if L = β = -α , we refer to ℓ ( · ) as being L -smooth relative to Φ , or relative smooth with parameter L .

According to Proposition 1.1 of [LFN18], α -β relative smoothness is equivalent to

In NGVI setting, we hope to find conditions under which negative ELBO objective (1)

is L -relative smooth on Ω w.r.t. A ∗ .

Relative Smoothness of KL Divergence Term Since A ∗ ( ξ, Ξ) = -1 2 log det(Ξ -diag( ξ ⊙ ξ )) (see Section C.1), and the KL divergence between Gaussian distributions admits a closed-form

the Hessians of the two functions coincide. Thus KL divergence term is 1 -1 smooth relative to A ∗ .

Relative Smoothness of Log-Likelihood Term Now we consider the log-likelihood term in (6). We explain the intuition for the derivation in the univariate case and state the general result in high dimensional case, deferring the formal proof to Theorem 3.2. 2

In order to compute the Hessian matrix, ∇ 2 E q ( z ; ω ) [ f ( z )] , for some four times differentiable function f , we need the following useful lemma [Pri58, Bon64, OA09].

Lemma 3.1 (Bonnet's and Price's Gradients) . Let q ( z ) be the probability density function (PDF) of the multivariate Gaussian, N ( µ, Σ) , and assume f is twice continuously differentiable, then

In the univariate case with standard parameters, µ = ξ , σ 2 = Ξ -ξ 2 , we apply Theorem 3.1 and chain rule to obtain

Using Theorem 3.1 and chain rule again and we get the following result.

Proposition 1. When d = 1 , let f ( z ) := -log p ( D| z ) , then the Hessian of the log-likelihood term equals

Moreover, it is straightforward to see that

Therefore, proving α -β relative smoothness is equivalent to finding α , β such that for all ω ,

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 , D > 1 , define the bounded sets in the spaces of standard and expectation parameters:

Then we present our main result on relative smoothness on such bounded sets for any U and D . 3 Theorem 3.2 (Sufficient Conditions for Relative Smoothness) . Let f ( z ) = -log p ( D| z ) be a four times continuously differentiable function in z on R d . Assume sup z ∈ R d sup i =1 , ··· ,d |∇ i f ( z ) | ≤ L 1 , sup z ∈ R d sup i =1 , ··· ,d sup j =1 , ··· ,d |∇ 2 ij f ( z ) | ≤ L 2 . Then the log-likelihood term is L -smooth relative to A ∗ on ˜ Ω with

Proof sketch. First, we use the approach mentioned above to compute the Hessian matrices ∇ 2 E q ( z ; ξ, Ξ) [ f ( z )] and ∇ 2 A ∗ ( ω ) . To handle high-order derivatives appearing in ∇ 2 E q ( z ; ξ, Ξ) [ f ( z )] , we apply Stein's Lemma (Theorem F.3) to bound them by lower-order derivatives. Finally, we find α and β as in (9) by proving the positive semidefiniteness of the corresponding matrices.

Example 1. For logistic regression, -log p ( y | x, z ) = log(1 + e -yx ⊤ z ) , where x ∈ R d is a data point and y ∈ {-1 , 1 } is the label. Since -∇ i log p ( y | x, z ) = -σ ( -yx ⊤ z ) yx i and -∇ 2 ij log p ( y | x, z ) = σ ( -yx ⊤ z )(1 -σ ( -yx ⊤ z )) x i x j , by writing ∥ x ∥ ∞ := max i | x i | , we have

Therefore, ℓ ( ω ) is smooth relative to A ∗ with parameter O ( dD 2 ∥ x ∥ ∞ ( U + D )(1 + U ∥ x ∥ ∞ )) .

Comparison to conjugate case. For conjugate models, ℓ ( ω ) is 1 -smooth and 1 -strongly convex relative to A ∗ [WG24], which makes it 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 ℓ ( ω ) 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 ℓ ( ω ) exhibits hidden convexity when the log-likelihood is concave in z , as in logistic and Poisson regression. Formal proofs of statements in this section are 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 ℓ ( ω ) = H ( c ( ω )) for an invertible map c : Ω → Θ . Then ℓ is hidden convex with modulus µ C > 0 , µ H ≥ 0 if:

The set Θ is convex and H : Θ → R is µ H -strongly convex w.r.t. Euclidean geometry.

The map c is invertible and ∃ µ C > 0 : ∥ c ( ω 1 ) -c ( ω 2 ) ∥ 2 ≥ µ C ∥ ω 1 -ω 2 ∥ 2 , ∀ ω 1 , ω 2 ∈ Ω .

The result below follows from the fact that the strong convexity of f ( u ) transfers to the expected value E q ( u ; µ,C ) [ f ( u )] under the Cholesky parameterization ( µ, C ) , where C is lower triangular with positive diagonals and CC ⊤ = Σ [Dom20, Theorem 9].

Proposition 2. If log p ( D| z ) is concave in z , then the restriction of ℓ ( ω ) : Ω → R on ˜ Ω is hidden convex with modulus µ C = (4 U 2 +4 D +1) -1 / 2 and µ H = 1 .

Given hidden convexity, we establish the PL inequality based on the analysis in [FHH25].

Proposition 3. If log p ( D| z ) is concave and a stationary point ω ∗ lies in ˜ Ω , then ω ∗ is a global minimum. Furthermore, ℓ ( ω ) satisfies the PL inequality in the relative interior ri( ˜ Ω) :

The assumption that ω ∗ ∈ ˜ Ω is mild. Under weak conditions, ℓ ( ω ) is coercive: ℓ ( ω ) → ∞ as ∥ ( µ, Σ) ∥ → ∞ or det(Σ) → 0 (see Section D), thus for large enough U and D , ˜ Ω contains a stationary point. Note that although the PL condition holds, it does not immediately imply the function value convergence of NGVI, since for 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-SNGD ) algorithm and prove an O (1 / √ T ) convergence rate under the relative smoothness condition in Section 4.2. Finally, if the log-likelihood is concave, we show in Section 4.3 that Proj-SNGD achieves O (1 /T ) convergence to the global minimum.

4.1 Projected Stochastic Natural Gradient Descent

Given an initial point ω 0 ∈ Ω , a number of iterations T and step sizes { γ t } 0 ≤ t ≤ T -1 , we introduce the update rule of our projected variant of SNGD (with ω 0 ∈ ˜ Ω )

Here Proj ˜ Ω ( ω ) denotes the non-Euclidean projection of ω onto ˜ Ω induced by geometry of A ∗ . This involves 1) a transformation from the expectation parameter ( ξ, Ξ) to the standard parameter ( µ, Σ) , 2) a Euclidean projection of ( µ, Σ) onto ˜ P , and 3) a reverse transformation back to the expectation space. Note that in the mean field case, the projection in the second step can be performed efficiently by a simple entry-wise clipping. Denote the clipping function clip [ a,b ] ( x ) = min { max { a, x } , b } , then the second formula of (11) is equivalent to

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 ˜ Ω , 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

Equipped with the relative smoothness of ℓ ( · ) , we can prove that Proj-SNGD converges to a stationary point with an O (1 / √ T ) rate. We use ˆ ∇ ℓ ( ω ) to denote the stochastic gradient and assume E [ ˆ ∇ ℓ ( ω ) | ω ] = ∇ ℓ ( ω ) . 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 ˆ ∇ ℓ ( ω t ) .

Assumption 1. There exists V ≥ 0 such that where ω + t +1 := argmin ω ∈ ˜ Ω γ t ⟨∇ ℓ ( ω t ) , ω ⟩ + D A ∗ ( ω, ω t ) is the output of a Proj-NGD step with exact gradient.

Remark. Assumption 1, first proposed by [HR21], reduces to E ∥ ˆ ∇ ℓ ( ω t ) -∇ ℓ ( ω t ) ∥ 2 in standard gradient descent case, however, it does not depend on any particular norm in general. This assumption was justified and used in [WG24] to show convergence of SNGD in conjugate setting. In Section I, we prove Assumption 1 is satisfied for our (non-conjugate) logistic regression models.

For a general setting, we will use the Bregman Forward-Backward Envelope (BFBE; [ATP21, FH24]) as the convergence criterion to a stationary point.

Definition 4.1 (Bregman Forward-Backward Envelope (BFBE)) . For some ρ > 0 , the BFBE at ω ∈ ˜ Ω is defined as E ρ ( ω ) := -2 ρ min ω ′ ∈ ˜ Ω [ ⟨∇ ℓ ( ω ) , ω ′ -ω ⟩ + ρD A ∗ ( ω ′ , ω )] .

In the Euclidean case with unconstrained domain, BFBE becomes the squared norm of the gradient E ρ ( ω ) = ∥∇ ℓ ( ω ) ∥ 2 . In non-Euclidean case it is a natural generalization of stationarity criteria to the geometry induced by A ∗ , and if ω ∈ ri( ˜ Ω) , then E ρ ( ω ) = 0 if and only if ∥∇ ℓ ( ω ) ∥ = 0 . It is shown in [FH24] that BFBE is the strongest criterion available for non-convex SMD . Next, we establish non-asymptotic convergence of Proj-SNGD using BFBE criterion.

Theorem 4.1 (Convergence of Proj-SNGD ) . Assume ℓ ( ω ) has a bounded domain ˜ Ω ⊆ Ω and ri( ˜ Ω) contains at least one stationary point ω ∗ . Suppose ℓ is smooth with respect to A ∗ with parameter L . Then under Assumption 1, for constant step size γ t = γ = min { 1 2 L , √ λ 0 V 2 LT , } , Proj-SNGD satisfies

where λ 0 := ℓ ( ω 0 ) -ℓ ∗ and ¯ ω T is sampled from { ω 0 , . . . , ω T -1 } with probabilities γ t / ∑ T -1 i =0 γ i .

The proof is in Section H. Theorem 4.1 implies that Proj-SNGD converges at a rate of O (1 / √ T ) in the number of iterations with the presence of randomness in gradients. In noiseless setting ( V = 0 ), we obtain a faster O (1 /T ) 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 (1 /T ) rate if the negative log-likelihood is convex. This fast rate relies on the PL inequality established in Proposition 3, but requires the following additional technical assumption.

Assumption 2. For any ω ∂ , let n be an outward normal direction at ω . 4 Then it holds that

∈ ˜ Ω ω .

In the proof of fast convergence of Proj-SNGD (see Theorem 4.2 below), we require that for every ω ∈ ˜ Ω ,

where ρ > 0 is a constant that can be chosen arbitrarily large. When ω lies in the interior of ˜ Ω , the condition (13) is satisfied for sufficiently large ρ . When ω lies on the boundary, Assumption 2 ensures that the gradient ∇ ℓ ( ω ) points towards the interior of ˜ Ω under the geometry induced by A ∗ . It further guarantees that, for sufficiently large ρ , ω ∗ defined in (13) lies in ˜ Ω .

The example below shows that Assumption 2 can be satisfied for a univariate Bayesian linear regression model.

Example 2. We consider a univariate Bayesian linear regression model with a single data point ( x, y ) , where x, y, z ∈ R , p ( z ) = N (0 , 1) and p ( y | x, z ) = N ( xz, 1) . It can be shown that (see Section H.4) Assumption 2 is satisfied if the following set of inequalities holds:

Therefore, by carefully choosing U and D , Assumption 2 can be satisfied for any data point ( x, y ) .

Now we are ready to present the O (1 /T ) global convergence guarantee of Proj-SNGD algorithm.

Theorem 4.2 (Fast Convergence of Proj-SNGD ) . Suppose ri( ˜ Ω) contains a stationary point ω ∗ and ℓ is smooth with respect to A ∗ with parameter L . Suppose log p ( D| z ) is concave in z . Let µ B := µ 2 C 9 U 2 D 2 . Under Assumptions 1 and 2, for the step size scheme

the iterates of Proj-SNGD satisfy

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 (1 /T ) rate. Notably, this rate matches the rate in [WG24] under conjugate assumption, and the rates in [DGG23, KOW + 23] under strong convexity conditions. However, we work in much more general setting of non-conjugate models and in non-Euclidean geometry.

One limitation is our technical Assumption 2 that can be difficult to verify for a 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 2 may fail, the performance of Proj-SNGD is no worse (and sometimes better) than SNGD .

5 Experiments

Numerical Stability of SNGD and Proj-SNGD . In this experiment, we again focus on Poisson regression as discussed in Figure 1. We consider data ( x, y ) = (0 . 9 , 24) where y is sampled from Poisson( e 4 x ) , fixed initialization of variance parameter σ 2 0 = 2 , and stochastic initialization of mean parameter µ 0 ∼ Unif([ -3 , 0]) . We aim to demonstrate numerical stability of Proj-SNGD and SNGD . 5 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, we observe an increase in ℓ 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 Figure 2. Moreover, Proj-SNGD exhibits faster convergence and, despite the constraint an additional ˜ Ω , it converges to the same optimal solution as SNGD .

0

10

20

30

40

50

10

-

3

10

-

1

10

1

10

3

10

5

Function Value Gap

step size=0.5

step size=0.3

step size=0.1

0

10

20

30

40

50

10

-

3

10

-

1

10

1

10

3

10

5

Function Value Gap

step size=0.5

step size=0.3

step size=0.1

Iteration

Iteration

0

200

400

600

800

1000

Epoch

600

800

1000

1200

1400

Negative ELBO

Prox-SGD

Proj-SGD

SNGD

Proj-SNGD

10

-

2

10

-

1

γ

0

1000

2000

3000

4000

5000

6000

Number of Iterations

Prox-SGD

Proj-SGD

SNGD

Proj-SNGD

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 consider logistic regression on a subset of MNIST dataset [LeC98] with labels 6 or 8 ( n = 11769 , d = 784 ). We use mini-batches of size 2000 and set the step size γ t = γ 0 / √ t , where γ 0 is a hyperparameter to be finetuned. We run the algorithm for 1000 epochs (6000 iterations). The results averaged over 5 independent runs are shown in Figure 3. In the left panel of Figure 3, we observe that non-Euclidean algorithms converge faster than Euclidean ones with fine-tuned step-size. The right panel illustrates robustness to the initial step size γ 0 , with lower iteration counts indicating faster convergence. Non-Euclidean algorithms reach the threshold in around 1000-2000 iterations for a wide range 0 . 05 ≤ γ 0 ≤ 0 . 2 , while Euclidean algorithms only do so at γ 0 = 0 . 02 . We also observe that Proj-SNGD slightly outperforms SNGD . Therefore, non-Euclidean algorithms, especially Proj-SNGD , are more robust to 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 a 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 O (1 /T ) convergence rate, we invoke an assumption on 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.

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. On 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):251-276, 1998.

[ATP21] Masoud Ahookhosh, Andreas Themelis, and Panagiotis Patrinos. A bregman forwardbackward linesearch algorithm for nonconvex composite optimization: superlinear convergence to nonisolated local minima. SIAM Journal on Optimization , 31(1):653685, 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¨ ur Informatik, 2017.

[BBT17] Heinz H Bauschke, J´ erˆ ome Bolte, and Marc Teboulle. A descent lemma beyond lipschitz gradient continuity: first-order methods revisited and applications. Mathematics of Operations Research , 42(2):330-348, 2017.

[BKM17] David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association , 112(518):859-877, 2017.

[BN06] Christopher M Bishop and Nasser M Nasrabadi. Pattern recognition and machine learning , volume 4. Springer, 2006.

[Bon64] Georges Bonnet. Transformations des signaux al´ eatoires a travers les systemes non lin´ eaires sans m´ emoire. In Annales des T´ el´ ecommunications , volume 19, pages 203220. Springer, 1964.

[CAAK19] Badr-Eddine Ch´ erief-Abdellatif, Pierre 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¨ el Castillo and Paul Egels. Posterior and variational inference for deep neural networks with heavy-tailed weights. arXiv preprint arXiv:2406.03369 , 2024.

[CFLM23] Michael Celentano, Zhou Fan, Licong Lin, and Song Mei. Mean-field variational inference with the tap free energy: Geometric and statistical properties in linear models. arXiv preprint arXiv:2311.08442 , 2023.

[DDM18] Damek Davis, Dmitriy Drusvyatskiy, and Kellie J MacPhee. Stochastic model-based minimization under high-order growth. arXiv preprint arXiv:1807.00255 , 2018.

[DGG23] Justin Domke, Robert Gower, and Guillaume Garrigos. Provable convergence guarantees for black-box variational inference. Advances in Neural Information Processing Systems , 36:66289-66327, 2023.

[Dom20] Justin Domke. Provable smoothness guarantees for black-box variational inference. In International Conference on Machine Learning , pages 2587-2596. PMLR, 2020.

[DWZZ24] 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 general bregman divergence. In International Conference on Artificial Intelligence and Statistics , pages 3493-3501. PMLR, 2024.

[FHH25] Ilyas Fatkhullin, Niao He, and Yifan Hu. Stochastic optimization under hidden convexity. SIAM Journal on Optimization , 2025.

[GLZ16] Saeed Ghadimi, Guanghui Lan, and Hongchao Zhang. Mini-batch stochastic approximation methods for nonconvex stochastic composite optimization. Mathematical Programming , 155(1):267-305, 2016.

[HBWP13] Matthew D Hoffman, David M Blei, Chong Wang, and John Paisley. Stochastic variational inference. Journal of Machine Learning Research , 14(1):1303-1347, 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´ arik. 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 in the conjugate exponential family. Advances in Neural Information Processing Systems , 25:2888-2897, 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 , 37:131104-131153, 2024.

[JGJS99] Michael I Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. An introduction to variational methods for graphical models. Machine learning , 37:183233, 1999.

[Kak01] Sham M Kakade. A 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 proximal-gradient methods with general divergence functions. In Proceedings of the Thirty-Second Conference on Uncertainty in Artificial Intelligence , page 319-328, 2016.

[KL17] Mohammad Khan and Wu Lin. Conjugate-computation variational inference: Converting 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 variational inference: Should we stick the landing? In International Conference on Artificial Intelligence and Statistics , pages 235-243. PMLR, 2024.

[KOW + 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 , 36:44615-44657, 2023.

[Lan12] Guanghui Lan. An optimal method for stochastic composite optimization. Mathematical Programming , 133(1):365-397, 2012.

[Lan23] Guanghui Lan. Policy mirror descent for reinforcement learning: Linear convergence, new sampling complexity, and generalized problem classes. Mathematical programming , 198(1):1059-1106, 2023.

[LeC98] Yann LeCun. The mnist database of handwritten digits. http://yann. lecun. com/exdb/mnist/ , 1998.

[LFN18] Haihao Lu, Robert M Freund, and Yurii Nesterov. Relatively smooth convex optimization by first-order methods, and applications. SIAM Journal on Optimization , 28(1):333-354, 2018.

[LKS19] Wu Lin, Mohammad Emtiyaz Khan, and Mark Schmidt. Fast and simple naturalgradient variational inference with mixture of exponential-family approximations. In International Conference on Machine Learning , pages 3992-4002. PMLR, 2019.

[Loj63] Stanislaw Lojasiewicz. Une propri´ et´ e topologique des sous-ensembles analytiques r´ eels. Les ´ equations aux d´ eriv´ ees 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ˆ am Le Minh, Julyan Arbel, Thomas M¨ ollenhoff, Mohammad Emtiyaz Khan, and Florence Forbes. Natural variational annealing for multimodal optimization. arXiv preprint arXiv:2501.04667 , 2025.

[Mar20] James Martens. New insights and perspectives on 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 calculations 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 optimization , 19(4):1574-1609, 2009.

[NY83] Arkadij Semenoviˇ c Nemirovskij and David Borisovich Yudin. Problem complexity and method efficiency in optimization . Wiley, 1983.

[OA09] Manfred Opper and C´ edric Archambeau. The variational gaussian approximation revisited. Neural computation , 21(3):786-792, 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 in neural information processing systems , 32:4264-4277, 2019.

[Pol63] Boris T Polyak. Gradient methods for the minimisation of functionals. USSR Computational Mathematics and Mathematical Physics , 3(4):864-878, 1963.

[Pri58] Robert Price. A useful theorem for nonlinear devices having gaussian inputs. IRE Transactions on Information Theory , 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 814-822. PMLR, 2014.

[RM15] Garvesh Raskutti and Sayan Mukherjee. The information geometry of mirror descent. IEEE Transactions on Information Theory , 61(3):1451-1457, 2015.

[SEH18] Hugh Salimbeni, Stefanos Eleftheriadis, and James Hensman. Natural gradients in practice: Non-conjugate variational inference in gaussian process models. In International Conference on Artificial Intelligence and Statistics , pages 689-697. PMLR, 2018.

[SLZ + 23] 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-43971, 2023.

[Ste81] Charles M Stein. Estimation of the mean of a multivariate normal distribution. The annals of Statistics , pages 1135-1151, 1981.

[TLG14] Michalis Titsias and Miguel L´ azaro-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 International Conference on Machine Learning , pages 6145-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 Variational Inference. In International Conference on Machine Learning , pages 5339853421. PMLR, 2024.

[WJ + 08] Martin J Wainwright, Michael I Jordan, et al. Graphical models, exponential families, and variational inference. Foundations and Trends® in Machine Learning , 1(1-2):1305, 2008.

[WW13] David Wingate and Theophane Weber. Automated variational inference in probabilistic programming. arXiv preprint arXiv:1301.1299 , 2013.

[Xia22] Lin Xiao. On the convergence rates of policy gradient methods. Journal of Machine Learning Research , 23(282):1-36, 2022.

[ZH18] Siqi Zhang and Niao He. On the convergence rate of stochastic mirror descent for nonsmooth nonconvex optimization. arXiv preprint arXiv:1806.04781 , 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.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

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 in 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 be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as 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]

Justification: The limitations of the paper are summarized in the last section (Limitations and Future Work).

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 a 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 authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or 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 it 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 address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory assumptions and proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should 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 conclusions of the paper (regardless of whether the code and data are provided or 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, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is 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, or 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 to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in

some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The data is publicly available. The code will be released to enable reproduction 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.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a 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.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on 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. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental setting/details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

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 be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment statistical significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: The medians and interquartile ranges are reported in the plots for each experiment involving randomness.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer 'Yes' if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out 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

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

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 or failed experiments that didn't make it into the paper).

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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 a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: This is a theoretical paper on optimization, and it has no foreseeable positive or negative societal impacts.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is 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, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation 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

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

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 no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: This paper uses the MNIST dataset, which is provided through PyTorch's built-in torchvision library. The original dataset is properly cited in our work.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

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 package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 not available online, the authors are encouraged to reach out to the asset's creators.

13. New assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: This paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of 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 is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and research with human subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well 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 not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional review board (IRB) approvals or equivalent for research with human subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or 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.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

16. Declaration of LLM usage

Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods 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, or originality of the research, declaration is not required.

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://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

Contents

Appendix

A Summary of Assumptions in this Paper

Below is a summary of the assumptions used in this paper for establishing the convergence of ProjSNGD :

The objective ℓ 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.

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 Section I).

The relative interior ri( ˜ Ω) contains a stationary point ω ∗ . This condition is required for both Theorem 4.1 and Theorem 4.2. It is a mild assumption, since ℓ is coercive under very weak conditions (see Section D).

The descent direction points inward on the boundary of ˜ Ω . This assumption is only required in Theorem 4.2, and it is satisfied, for example, in a univariate Bayesian linear regression model (see Section H.4).

The log-likelihood log p ( D| z ) is concave in z . This is only needed in Theorem 4.2, and holds for many models, including logistic regression and Poisson regression.

B Descent Lemma for Mirror Descent

In this section, we prove a descent lemma for mirror descent.

Lemma B.1. Let Φ be a strictly convex distance generating function on a closed domain Ω ⊆ R d with induced Bregman divergence D Φ . Assume that ℓ : U → R is L -smooth relative to Φ . Consider (deterministic) mirror descent update with step size γ

then with step size γ = 1 /L , the objective decreases after each iteration, i.e.,

Theorem B.1 indicates that if a mirror descent update with step size γ causes an increase of the objective, the relative smoothness parameter must be greater than 1 /γ . 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 α -β relative smoothness in univariate case (Theorem F.1) implies

thus the smoothness parameter should grow exponentially with respect to σ 2 0 . Therefore, Figure 1 is consistent with our theoretical findings.

Proof. By the second part of Lemma F.1 in [FH24] and the definition of u t +1 , we have that for all v ∈ U ,

Then, we use the relative smoothness of ℓ to get that for all v ∈ U ,

Finally, we choose v = u t and the result follows.

C Background on Exponential Family

In this section, we will provide background information about the exponential family. In Section C.1, we will derive the natural and expectation parameters of a 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 if it takes the following form:

where η is the natural parameter, ϕ ( z ) is the sufficient statistic, and A ( η ) is the log-partition function. The expectation parameter is defined as ω = E q ( z ; η ) [ ϕ ( z )] . Now we let q ( z ) be the PDF of Gaussian random vector N ( µ, Σ) , then

Therefore, the sufficient statistics are z and zz ⊤ , and we have the following fact.

Fact 1. For multivariate Gaussian distribution q , the expectation parameters of q are given by ξ := E [ z ] = µ and Ξ := E [ zz ⊤ ] = Σ + µµ ⊤ . The natural parameters of q are given by λ := Σ -1 µ and Λ := -1 2 Σ -1 .

Moreover, by plugging the definition of λ and Λ , one have

Next, we compute A ∗ via the definition of convex conjugate:

It's easy to see that the optimal solution of the previous problem is given by

Moreover, we have the following result.

Fact 2. The convex conjugate A ∗ equals

Based on Fact 2, we can check that

This also indicates that ∇ A and ∇ A ∗ are inverse operators of each other.

Moreover, when d = 1 , it's straightforward to check that

This result will be useful in the proof of relative smoothness, e.g., in Section 3.1 (8). Next, we consider the mean field family,

Fact 3. For multivariate Gaussian distribution q under mean-field parameterization, the expectation parameters of q are given by ξ i := E [ z i ] = µ i and Ξ ii := E [ z 2 i ] = Σ ii + µ 2 i entry-wise. Equivalently, we have

Similarly, we also have λ := Σ -1 µ and diagonal matrix Λ = -1 2 Σ -1 .

In addition, Fact 3 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 ProjSNGD in (11) (given stochastic gradient ˆ ∇ ℓ ( ω t ) ) can be performed in O ( d ) time.

C.2 Derivation of SNGD

In this section, we show that formula (4) holds, i.e.,

where η t and ω t are the natural and expectation parameter of the same distribution. The proof is adapted from [RM15] and [WG24].

Proof of Equation (4) . We first prove that I ( η ) = ∇ 2 A ( η ) . Since

we have

By definition of Fisher information and duality between η and ω ,

Moreover, we have

Then, we use the chain rule to get

Multiplying I ( η ) -1 on both sides gives the desired result.

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 ℓ ( ω ) 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 ( ω ) → + ∞ as u ∈ U , u → ∂ U or ∥ u ∥ → ∞ , where ∂ U denotes the boundary of U .

Notice that coerciveness may potentially depend on specific parameterization. We consider 4 different parameterizations of Gaussian variational family in our work:

Standard:

(

µ,

Σ)

Expectation:

ω

= (

µ,

Σ+

µµ

⊤

)

Cholesky:

c

(

ω

) = (

µ, C

)

Natural:

In Cholesky parameterization, C is defined as the Cholesky factor of Σ . Since our focus is the convergence of NGVI, we are mainly interested in the landscape properties of ℓ ( ω ) in expectation parameterization. We first show objective ℓ ( µ, Σ) is coercive in standard parameterization ( µ, Σ) ∈ R d ×S d + .

Theorem D.1 (Coerciveness of the Variational Objective in Standard Parameterization) . Let q ( z ) = N ( z ; µ, Σ) be a Gaussian variational distribution with mean parameter µ ∈ R d and covariance matrix Σ ∈ S d + . Let the prior be the standard normal distribution p ( z ) = N ( z ; 0 , I ) . Suppose the log-likelihood function satisfies a sub-quadratic growth condition: there exist constants c 0 , c 1 ≥ 0 and an exponent r ∈ [1 , 2) such that for any observed data D and all z ∈ R d ,

Then the variational objective ℓ : R d ×S d + → R defined in (6) is coercive, in the sense that

The sub-quadratic growth assumption (16) is a very mild condition which is satisfied by most problems in practice. For example, the log-likelihood function log p ( D| z ) = log p ( y | x, z ) for dataset D = { ( x i , y i ) } n i =1 diverges to -∞ when ∥ z ∥ → ∞ in Bayesian regression and logistic regression, and has at most linear growth in Poisson regression for bounded { x i } n i =1 .

With this result, we can prove that the objective is coercive in all parameterizations as listed below.

Expectation parameterization. The domain of expectation parameter is Ω = { ( ξ, Ξ) : ξ ∈ R d , Ξ -ξξ ⊤ ∈ S d + } . For ξ , ∥ ξ ∥ → ∞ if and only if ∥ µ ∥ → ∞ . For unboundedness of Ξ , since ∥ Ξ ∥ ≤ ∥ Σ ∥ + ∥ µ ∥ 2 , ∥ Ξ ∥ → ∞ implies at least one of ∥ Σ ∥ → ∞ and ∥ µ ∥ → ∞ holds. Moreover, if ( ξ, Ξ) → ∂ Ω , i.e., det(Ξ -ξξ ⊤ ) → 0 , this is equivalent to det(Σ) → 0 .

Natural parameterization. The domain of natural parameter is D A = { ( λ, Λ) : λ ∈ R d , -Λ ∈ S d + } . For λ , ∥ λ ∥ ≤ ∥ Σ -1 ∥∥ µ ∥ = ( λ min (Σ)) -1 ∥ µ ∥ , thus ∥ λ ∥ → 0 implies either det(Σ) → 0 or ∥ µ ∥ → 0 . The proof of Λ is obvious.

Cholesky parameterization. Since C is defined as a lower-triangle matrix with positive diagonal entries, ∥ C ∥ → ∞ if and only if ∥ Σ ∥ → ∞ , and det( C ) → 0 if and only if det(Σ) → 0 .

Proof of Theorem D.1. It's clear that the boundary of R d × S d + is R d × { Σ : Σ is symmetric and det(Σ) = 0 } . To show coerciveness, we need to prove that under the following 3 cases, the objective ℓ ( µ, Σ) will diverge to infinity:

Recall from (6) that the variational objective is given by

For the first term, we can find the closed-form solution of KL divergence between to Gaussian distributions:

For the second term, the sub-quadratic growth condition (16) implies that

It remains to upper bound E [ ∥ z ∥ r ] . Let x ∼ N (0 , I ) be a standard Gaussian random vector, then by Minkowski inequality and Jensen's inequality we have

Therefore, since 1 ≤ r < 2 , it can be shown that for some constant C > 0 , it holds that

Finally, we have

If det(Σ) → 0 , the first term will diverge to + ∞ . If ∥ µ ∥ → ∞ or ∥ Σ ∥ → ∞ , the terms in the square brackets dominate as r < 2 , hence ℓ ( µ, Σ) will diverge to + ∞ .

E Examples of Non-convex Objectives in Non-conjugate Models

In this section, we present two non-conjugate models where the objective is 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 = { ( x i , y i ) : x i ∈ [ -1 , 1] , y i ∈ {-1 , 1 }} n i =1 and latent variable z = ( w,b ) ∈ R 2 . The prior p ( z ) is standard Gaussian. Then we have

In the following we will write s i := σ ( -y i ( wx i + b )) for 1 ≤ i ≤ n , where σ ( · ) is the sigmoid function. On the convex subset

we can compute the second derivative with respect to σ 2 2 using Theorem 3.1 as

A necessary condition of convexity of ℓ ( ω ) is that the second derivative w.r.t. σ 2 2 is non-negative.

Note that as σ 2 1 , σ 2 2 → 0 , w,b → 0 and s i → 1 / 2 in probability for all 1 ≤ i ≤ n , thus

Then there exists δ > 0 such that when σ 2 1 = σ 2 2 = δ ,

holds for all 1 ≤ i ≤ n . Then we can choose n > 32 /δ 2 to show that ℓ ( ω ) is non-convex when σ 2 1 = σ 2 2 = δ , i.e.,

Poisson Regression. In this example, we are given dataset D = { ( x i , y i ) : x i ∈ R d , y i ∈ N } n i =1 with latent variable z ∈ R d . We assume that y | x ∼ Poisson( λ ) , where

The expected log-likelihood admits the following closed-form solution:

Then we can compute the Hessian w.r.t. ξ :

We consider the following convex subset where the covariance matrix equals identity, i.e.,

We can find x i and ξ such that 0 < x ⊤ i ξ < 2 for all 1 ≤ i ≤ n . Then we have exp ( x ⊤ i ξ + 1 2 x ⊤ i (Ξ -ξξ ⊤ ) x i ) > 1 , and thus

where we used the fact that D KL ( q ( z ; ω ) ∥ q ( z ; ω ′ )) = D A ∗ ( ω, ω ′ ) . Moreover,

Therefore, for some c > 0 , if rescale the dataset D ′ = { ( cx i , y i ) : x i ∈ R d , y i ∈ N } n i =1 and evaluate ℓ ( ω ) at ξ ′ = ξ/c , we have

if c is sufficiently large.

F Tighter Relative Smoothness in Univariate Case ( µ, σ ∈ R )

In this section, we will provide tighter relative smoothness guarantees in univariate case, i.e., µ, σ ∈ R , using the equivalent conditions for relative smoothness in (9).

F.1 Necessary and Sufficient Conditions for Relative Smoothness

Let f ( z ) := -log p ( D| z ) . Inequalities (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.,

Let B = E q [ ∇ 3 f ( z )] , C = E q [ ∇ 4 f ( z )] . Using the explicit form of the Hessians (7) and (8), the relative smoothness condition (17) is equivalent to

Theorem F.1 (Necessary and Sufficient Conditions for Relative Smooth Condition (17), Univariate) . Let f be a four times continuously differentiable function on R . Then for Gaussian variational family Q = { q ( z ; ω ) : ω = ( ξ, Ξ) ∈ Ω } and some positive constant β , condition (17) holds if and only if for all µ ∈ R , σ 2 > 0 ,

Remark. Theorem F.1 is the immediate result of Sylvester's criterion of positive-semidefiniteness of matrix M ( β ) .

The first condition corresponds to the non-negativeness of the bottom right entry.

The second condition corresponds to the non-negativeness of the top left entry. Note that this is not a quadratic function of µ since B and C implicitly depend on µ . However, if we impose a uniform bound | C | ≤ β 4 and | B | ≤ β 3 for some β 4 , β 3 ≥ 0 for all µ ∈ R , σ 2 > 0 , the following inequality will be a sufficient condition for it to hold:

(20) holds for all µ ∈ R if and only if: either the coefficients of quadratic and linear terms are both zero ( ∇ 3 f ( z ) = 0 ), or the discriminant is non-positive, which gives

The last condition corresponds to the non-negativeness of the determinant. A direct computation of the determinant gives

Note that the coefficients of the 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,C with the upper bounds of | B | , | C | to obtain (21).

We also need to find conditions under which the inequality on the left side of (9) holds, i.e.,

which is equivalent to

We apply Sylvester's criterion as 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 = { q ( z ; ω ) : ω = ( ξ, Ξ) ∈ Ω } and some constant α , condition (22) holds if and only if for all µ ∈ R , σ 2 > 0 ,

Remark. The first condition implies an important necessary condition for relative convexity (i.e., α ≥ 0 ): ∇ 4 f ( z ) ≥ 0 for all z ∈ R . Indeed, if ∇ 4 f ( z 0 ) < 0 for some z 0 and ∇ 4 f ( z ) is continuous, we can always pick some ( µ, σ 2 ) such that E q [ ∇ 4 f ( z )] < 0 (e.g. µ = z 0 and σ 2 is sufficiently small), then we cannot find any α ≥ 0 such that the first condition holds.

F.2 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]) . Let Z ∼ N ( µ, Σ) , then for any differentiable function g , we have

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 in multivariate case)

Combining the parts above, we can get the following sufficient conditions.

Corollary F.3.1 (Sufficient Conditions for Relative Smoothness, Univariate) . Under the assumptions of Theorem F.1, if we further assume that sup z |∇ f ( z ) | ≤ L 1 and sup z |∇ 2 f ( z ) | < L 2 for some constant L 1 , L 2 ≥ 0 , then on the restriction of the parameter set { ( µ, σ 2 ) : 0 < σ 2 ≤ D } for some D > 0 , condition (17) holds with

Moreover, relative weak convexity condition (22) holds with parameter

Proof. For the choice of β , we first check three conditions in Theorem F.1 one by one:

Since σ 2 ≤ D and sup z |∇ 2 f ( z ) | ≤ L 2 , use the upper bound on | E q [ ∇ 4 f ( z )] | in (24) and we have

Thanks to our global bound on | E [ ∇ 4 f ( z )] | and | E [ ∇ 3 f ( z )] | in (24), we have

holds for any µ ∈ R .

In order to show the RHS is non-negative for any µ ∈ R , we first assume µ ≥ 0 (similar results follow if µ ≤ 0 ). In this case, we only need to prove

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-positive, that is,

This is again a quadratic function of β , and the inequality holds if we pick a sufficiently large β . More specifically, the larger root of (25) is

Therefore, β satisfies the inequality (25).

We again use the upper bound of | E q [ ∇ 4 f ( z )] | and | E q [ ∇ 3 f ( z )] | in (24),

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 a relative strong convexity guarantee of the objective ℓ ( ω ) , since the KL divergence term in (6) is 1-relatively convex, we need to guarantee that -D 2 L 2 -√ 2 D 2 L 1 > -1 .

As a result, under the assumptions of Corollary F.3.1, we conclude that E q [ f ( z )] is relatively smooth with respect to A ∗ with parameter

Example 3. Below are some examples where the likelihood p ( D| z ) = p ( y | x, z ) for D = { ( x, y ) } is specified:

Linear Bayesian Regression. In linear Bayesian regression where p ( y | x, z ) = N ( xz, σ 2 ) , it's easy to see that -log p ( y | x, z ) is a quadratic function in z , hence the necessary and sufficient conditions in Theorem F.1 and Theorem F.2 are satisfied with α = β = 0 . Therefore, the negative ELBOis 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 ( z ) = -log p ( y | x, z ) = log(1 + e -xyz ) , then it can be shown that |∇ f ( z ) | ≤ | x | , |∇ 2 f ( z ) | ≤ x 2 4 for all z ∈ R . Therefore, we can set σ 2 ≤ D and E q [ f ( z )] is relatively smooth with parameter L = x 2 D 8 + √ 2 D | x | 2 according to Corollary F.3.1.

Poisson Regression In Poisson regression, we write f ( z ) = -log p ( y | x, z ) = e zx -xyz + c for y ∈ N + and some constant c . Then E q [ ∇ 4 f ( z )] = x 4 e µx + σ 2 x 2 / 2 . In order to satisfy the first necessary and sufficient condition σ 4 E q [ ∇ 4 f ( z )] ≤ 2 β in the necessary and sufficient conditions Theorem F.1, we have to upper-bound | x | , | µ | and σ 2 . Therefore, it is in general difficult to prove relative smoothness for Poisson regression without further assumptions (e.g., bounded domain).

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 ∇ 2 E q ( z ; ω ) [ f ( z )] and ∇ 2 A ∗ ( ω ) . Lemma G.1. Let f ∈ C 4 ( R d ) , then for ω = ( ξ 1 , Ξ 11 , · · · , ξ d , Ξ dd ) ∈ R 2 d , we have for p, q ∈ { 1 , · · · , d } ,

̸

where H ij = E q [ ∇ 2 ij f ( z )] , B ij = E q [ ∇ 3 ijj f ( z )] and C ij = E q [ ∇ 4 iijj f ( z )] .

Proof. We apply Bonnet's and Price's gradients (Theorem 3.1), we can easily get the following partial derivatives

Next, we use Theorem 3.1 again to get for i = j ,

And for i = j ,

Moreover, for either i = j or i = j , we have

̸

Lemma G.2. For mean field family q ω ( z ) , we have

where each ∇ 2 A ∗ i ( ω i ) is a 2 × 2 matrix corresponding to univariate case, i.e.,

Proof. From Section C.1 we know that

̸

̸

̸

̸

Since the coordinates are independent, ∇ 2 ω i ω j A ∗ ( ω ) = 0 for all i = j . Therefore, ∇ 2 A ∗ ( ω ) is a block diagonal matrix, and it's easy to show that each block

In order to prove the theorem, we still need a lemma to transform high-order derivative conditions to lower-order ones with Stein's Lemma F.3. The proof of this lemma is deferred to Section G.2.

Lemma G.3. For f ∈ C 2 ( R d ) , if Z ∼ N ( µ, Σ) where Σ is a diagonal matrix, we have that for any i, j ∈ { 1 , · · · , d } ,

With the three lemmas above, we are ready to prove Theorem 3.2.

Proof of Theorem 3.2. First, we will show that with β = O ( dD 2 U ( L 1 + L 2 U )+ dD 3 ( L 1 + L 2 U )) , ∇ 2 E q ( z ; ω ) [ f ( z )] ⪯ β ∇ 2 A ∗ ( ω ) holds, i.e., M ( β ) := β ∇ 2 A ∗ ( ω ) - ∇ 2 E q ( z ; ω ) [ f ( z )] is positive semidefinite. According to Theorems G.1 and G.2, for p, q ∈ { 1 , · · · , d } ,

̸

̸

̸

Then, we will prove that with β = O ( dD 2 U ( L 1 + L 2 U ) + dD 3 ( L 1 + L 2 U )) , M is positive semidefinite. In the following, we will drop the dependence of β for convenience.

Note that the assumptions in the theorem and Theorem G.3 imply that | B pq | ≤ σ -2 q sup z |∇ p f ( z ) | ≤ DL 1 and | C pq | ≤ σ -2 q sup z |∇ 2 pp f ( z ) | ≤ DL 2 . It is obvious that | H pq | ≤ L 2 .

Define matrix ˜ M = M -2 3 diag( M ) , where diag( M ) is the diagonal matrix by setting all nondiagonal entries of M to 0. Then for any v ∈ R 2 d ,

Here B is the set of sets of cardinality d , which contain exactly one element of { 1 , 2 } , one element of { 3 , 4 } , ..., one element of { 2 d -1 , 2 d } . The set B has cardinality 2 d . Equation (27) can be proved by comparing the coefficient of each term v i v j M ij on both sides. Therefore, to prove v ⊤ Mv ≥ 0 for any v ∈ R 2 d , we only need to show each term on the RHS is non-negative. In other words, we only need to show

For any B ∈ B , the matrix ˜ M B is PSD, where ˜ M B ∈ R d × d is the submatrix of ˜ M by choosing the i -th rows and columns where i ∈ B .

For any p ∈ { 1 , · · · , d } , the submatrices of M

are PSD.

For the first set of conditions, we aim to find β such that the matrices ˜ M B are diagonally dominant and the diagonal entries are positive. For positiveness, if β ≥ O ( dD 2 U ( L 1 + L 2 U ) + dD 3 ( L 1 + L 2 U )) , then

For diagonal dominance, we fix i = 2 p -1 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 B ii is dominating in this matrix. Indeed, since the 2 p -th row does not appear in ˜ M B , we only need to guarantee that ˜ M 2 p -1 , 2 p -1 ≥ ∑ 2 d i =1 ,i / ∈{ 2 p -1 , 2 p } | ˜ M 2 p -1 ,i | :

̸

Now, applying the same reasoning for i = 2 p , we obtain

̸

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 ( p ) is non-negative. For every fixed p ∈ { 1 , · · · , d } the determinant is given by

Let the determinant be non-negative and we get

Since the larger root of β 2 = aβ + b ( a, b ≥ 0 ) is given by β = -a + √ a 2 +4 b 2 , the condition β ≥ a + √ b = a 2 + a + √ 4 b 2 ≥ a + √ a 2 +4 b 2 guarantees that β 2 ≥ aβ + b always holds. Therefore, the determinant is non-negative if β ≥ O ( D 2 U ( L 1 + L 2 U ) + √ D 3 U ( L 1 + L 2 U )) = O ( D 2 U ( L 1 + L 2 U )) .

Clearly, β = O ( dD 2 U ( L 1 + L 2 U ) + dD 3 ( L 1 + L 2 U )) satisfies all the conditions.

Therefore, using the same approach, one can show that -M ( α ) ⪰ 0 with -α taking the same value as β . As a result, -E q [log p ( D| z )] = E q [ f ( z )] is L -smooth relative to A ∗ on ˜ Ω , where L = O ( dD 2 U ( L 1 + L 2 U ) + dD 3 ( L 1 + L 2 U )) .

Finally, we need to prove that α = -O ( dD 2 U ( L 1 + L 2 U ) + dD 3 ( L 1 + L 2 U )) satisfies ∇ 2 E q ( z ; ω ) [ f ( z )] ⪰ α ∇ 2 A ∗ ( ω ) , i.e., ∇ 2 E q ( z ; ω ) [ f ( z )] -α ∇ 2 A ∗ ( ω ) = -M ( α ) ⪰ 0 . We can again choose α < 0 and | α | so large that the matrix -M ( α ) is diagonally dominant with positive diagonal 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 ( α ) .

G.2 Proof of Theorem G.3

We first restate Theorem G.3 below.

Lemma. For f ∈ C 2 ( R d ) , if Z ∼ N ( µ, Σ) where Σ is a diagonal matrix, we have that for any i, j ∈ { 1 , · · · , d } ,

Proof. For i = j , we have

where g ( Z i ) is a univariate function and

Since Z i ∼ N ( µ i , σ 2 i ) , we apply Stein's lemma (Theorem F.3) twice and we have

where X i := ( Z i -µ i ) /σ i ∼ N (0 , 1) . Using Lemma G.4 and we get the following bound on E [ ∇ 2 ii f ( Z )] :

For i = j , we define g ( Z ) = ( Z j -µ j ) f ( Z ) and use Theorem F.3 to get

̸

Reading the i -th coordinate gives

Then we define h ( Z ) = ∇ i f ( Z ) and Theorem F.3 gives

Reading the j -th coordinate gives

Combining (28) and (29) and we have the following bound of E [ ∇ ij f ( Z )] :

where X i := ( Z i -µ i ) /σ i and same for X j . Note that X i , X j are independent because we assume Σ is diagonal. Finally, we use the standard result that E | X | ≤ 1 for standard Gaussian variable X to get that

E [ ∇ ij f ( Z )] = σ -1 i σ -1 j E [ f ( Z ) X i X j ] ≤ σ -1 i σ -1 j sup z | f ( z ) | E | X i | E | X j | ≤ σ -1 i σ -1 j sup z | f ( z ) | . By symmetry, -E [ ∇ 2 ij f ( Z )] ≤ σ -1 i σ -1 j sup z | f ( z ) | holds as well.

Lemma G.4. Let X ∼ N (0 , 1) , then

Proof. First, we have

Note that

then we can compute the expectation in closed form

G.3 Proof of Proposition 2: Hidden Convexity of ℓ ( ω )

We define Θ := { ( µ, C ) : C ∈ S d + and is diagonal } . This set corresponds to the Cholesky parameterization of the expectation parameters contained in Ω . Furthermore, we define c : Ω → Θ to be the one-to-one map from expectation parameter to Cholesky parameter of the same distribution.

ω = ( µ, Σ+diag( µ ⊙ µ )) c ( ω ) = ( µ, C ) , where CC ⊤ = Σ Expectation parameters, Ω Cholesky parameters, Θ mapping c ( ω )

For bounded expectation parameter domain ˜ Ω , we also define its counterpart in Cholesky parameterization ˜ Θ := { ( µ, C ) : C is diagonal and D -1 ≤ C 2 ii ≤ D, ∀ 1 ≤ i ≤ d } . With a slight abuse of notation, we also use c to denote the one-to-one map from ˜ Ω to ˜ Θ .

In the following, we will verify that the negative ELBO ℓ ( ω ) is hidden convex as defined in Definition 3.2 on the bounded domain ˜ Ω .

Proof of Proposition 2. Define L : ˜ Θ → R ,

where q ( z ; θ ) = q ( z ; µ, C ) = N ( µ, CC ⊤ ) is the Gaussian vector with Cholesky parameterization.

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:

Since each -log C ii is convex in C ii , the entropy is convex in C . Then we can conclude that the second term E q θ ( z ) [log q θ ( z )] is convex in θ .

For the first term of (30), we apply Theorem 9 of [Dom20], which states that the µ H -strong convexity of f ( z ) for some function f implies the µ H -strong convexity of E q ( z ; θ ) [ f ( z )] w.r.t. the Cholesky parameter θ . Therefore, we only need to prove the strong convexity of -log p ( y | z ) -log p ( z ) in z . Indeed, -log p ( y | z ) -log p ( z ) is 1-strongly convex in z because log p ( y | z ) is concave in z and

is 1-strongly convex in z .

Therefore, L ( θ ) , as the sum of a 1-strongly convex function (first term) and a convex function (second term), is 1-strongly convex. Then we conclude that the first property in Definition 3.2 is satisfied with µ H = 1 .

Proof of the Second Property in Definition 3.2. Now we compute µ c by computing the Lipschitz coefficient of the inverse map c -1 on ˜ Θ . We reparameterize ω and θ as vectors: ω = ( ξ 1 , Ξ 11 , · · · , ξ d , Ξ dd ) ∈ R 2 d and θ = ( µ 1 , C 11 , · · · , µ d , C dd ) ∈ R 2 d , respectively. Note that ω = c -1 ( θ ) = ( µ 1 , µ 2 1 + C 2 11 , · · · , µ d , µ 2 d + C 2 dd ) , and ∇ c -1 ( θ ) = diag( G 1 , · · · , G d ) ∈ R 2 d × 2 d , where

For each block G i , the largest singular value of G i is bounded by

The last inequality above holds since

and Tr( G i G ⊤ i ) = 1 + 4 µ 2 i +4 C 2 ii ≤ 4 U 2 +4 D +1 using the fact that ( µ, C ) ∈ ˜ Θ .

Hence we conclude that ∥∇ c -1 ( θ ) ∥ ≤ √ 4 U 2 +4 D +1 , and c -1 ( · ) is √ 4 U 2 +4 D +1 -Lipschitz continuous. Then the second property is satisfied with µ C = (4 U 2 +4 D +1) -1 / 2 .

H Missing Proofs in Section 4

In this section, we will prove the convergence guarantees of Proj-SNGD in Section 4. 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, respectively.

H.1 Smoothness and Strong Convexity of A ∗

In this subsection, we will establish the smoothness and strong convexity properties of A ∗ in Euclidean geometry in the bounded set

It is obvious that A ∗ is globally 1-strongly convex in the non-Euclidean geometry induced by A ∗ itself. However, it is only globally strictly convex in the Euclidean geometry, and strongly convexity only holds in a bounded domain of Ω . In order to transform the PL inequality in Euclidean geometry (10) to the PL inequality in non-Euclidean geometry (see Section H.3), such smooth and strong convexity results are necessary.

Lemma H.1. A ∗ is strongly convex with parameter C S := ( D (4 U 2 +2 D +1)) -1 with respect to Euclidean norm. Moreover, it is smooth with parameter C L := 9 U 2 D 2 2 .

For each block ∇ 2 A ∗ i ( ω i ) , its smaller eigenvalue satisfies

Since ∇ 2 A ∗ ( ω ) is a block diagonal matrix, we conclude that λ min ( ∇ 2 A ∗ ( ω )) ≥ 1 D (4 U 2 +2 D +1) , hence

For the smoothness part, we solve for the larger eigenvalue directly to get

where we used √ a + b ≤ √ a + √ b in the first inequality, and 2 µ 2 i + σ 2 i > 0 in the second inequality. Then we conclude that A ∗ is smooth with parameter C L := 9 2 U 2 D 2 .

H.2 Proof of Theorem 4.1: Convergence of Proj-SNGD

We first introduce the standard assumption on gradient variance, which is widely used in optimization.

Assumption 3. For all ω t ∈ ˜ Ω , there exists V ≥ 0 such that

We also define the Bregman Moreau envelope of function ℓ as

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 ∗ is 1-strongly convex and ℓ is smooth with respect to A ∗ with parameter L . Let { γ t } 0 ≤ t ≤ T -1 be non-increasing with γ 0 ≤ 1 / (2 L ) . For { ω i } T -1 i =0 generated by (32) , let ¯ ω T be randomly chosen from ω 0 , · · · , ω T -1 with probability p t = ˜ γ t / ∑ n -1 i =0 ˜ γ i , then under Assumption 3, we have

where λ 0 = 2( ℓ ( ω 0 ) -ℓ ∗ ) and ℓ ∗ := inf ω ∈ ˜ Ω ( ω ) . If we use constant step size γ t = min { 1 2 L , √ λ 0 V 2 LT } then

Note that ℓ 1 /ρ ( ω ) ≤ ℓ ( ω ) for all ω ∈ Ω . Thus ˜ λ 0 ≤ 2( ℓ ( ω 0 ) -ℓ ∗ ) . Theorem 4.1 is essentially the direct result of Theorem H.2 except two minor differences:

Instead of the update rules of Proj-SNGD as defined in (11), Theorem H.2 assumes a more direct update rule

We will prove the equivalence of two expression in Theorem H.3.

Theorem H.2 considers a different assumption on gradient variance, and Theorem H.2 assumes 1-strong convexity of the distance generating function A ∗ . In Proposition 4 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 ˜ ω t +1 be the output of Proj-SNGD as in (11) , and let ω t +1 be the direct update rule as in (32) . Then ˜ ω t +1 = ω t +1 . Moreover, (11) can be performed entry-wise in O ( d ) time.

Proof. Recall that in the update rule of Proj-SNGD as in (11), we define ω t +1 , ∗ ∈ Ω such that

By the duality of SMD , ω t +1 , ∗ is also the optimal solution without constraints, i.e.,

Use the definition of ω t +1 and (33), we have

Therefore, to find ω t +1 , we only need to compute ω t +1 , ∗ and then project it onto ˜ Ω under the geometry induced by A ∗ . Next we show that ˜ ω t +1 solves the optimization problem above. Use the fact that for mean-field parameterization (see Section C.1),

in the standard parameter space we have

The last equality holds due to the mean field assumption. Therefore, we can solve ω t +1 by optimizing over ( µ i , Σ ii ) independently.

For each entry i , (34) is a quadratic function in µ i with positive quadratic term and the (unconstrained) minimum is attained at µ i = ( µ t +1 , ∗ ) i . If ( µ t +1 , ∗ ) i < -U , (34) is an increasing function on [ -U, U ] , and if ( µ t +1 , ∗ ) i > U , the function is decreasing on [ -U, U ] . Therefore, the minimum on [ -U, U ] is always attained at ( µ t +1 ) i = clip [ -U,U ] (( µ t +1 , ∗ ) i ) .

Similarly, (34) is a convex function in Σ ii and the (unconstrained) minimum is attained at Σ ii = (Σ t +1 , ∗ ) ii . It's also straightforward to check that the minimizer is given by (Σ t +1 ) ii = clip [ D -1 ,D ] ((Σ t +1 , ∗ ) ii ) . The closed-form solution of ( µ t +1 , Σ t +1 ) coincides with the standard projection of ( µ t +1 , ∗ , Σ t +1 , ∗ ) onto ˜ P under Euclidean geometry. By transforming ( µ, Σ) back to expectation parameter space, we conclude that ˜ ω t +1 = ω t +1 .

Since the computation of ∇ A ∗ ( · ) and its inverse ( ∇ A ∗ ( · )) -1 involves only entry-wise calculation (see Section C.1), one can also compute

in O ( d ) time. Finally, the transformation between ( µ, Σ) and ( ξ, Ξ) , and the projection both require O ( d ) time. Therefore, ω t +1 can be calculated in O ( 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 t ≥ 0 , define

In step II (one step progress on the Lyapunov function), formula (15) of the proof in [FH24], we have

here we define ˆ ω := prox l/ρ ( ω ) .

Note that E [ ⟨ ˆ ∇ ℓ ( ω t ) - ∇ ℓ ( ω t ) , ˆ ω t -ω + t +1 ⟩ | ω t ] = 0 because the ˆ ω t -ω + t +1 has no randomness given ω t , and the gradient estimator is unbiased. Moreover, by bounded variance assumption (12),

E [ ⟨ ˆ ∇ ℓ ( ω t ) -∇ ℓ ( ω t ) , ω + t +1 -ω t +1 ⟩ | ω t ] ≤ γ t V 2 . Therefore, we use the same assumptions as in the proof that ρ = 2 L and γ t ≤ 1 / (2 L ) and we have

This is the same formula as formula (17) in [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 of Theorem 4.2: Fast Convergence of Proj-SNGD

To establish Theorem 4.2, we rely on 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 to PL inequality with a difference that the gradient norm is replaced with BFBE, a non-Euclidean first-order stationarity measure. Recall that BFBE is defined in Definition 4.1 as

Assumption 4 (Bregman Prox-PL condition) . There exists some constant ρ > 3 L and µ B > 0 such that for all ω ∈ ˜ Ω ,

Theorem H.4. [Theorem 4.7 in [FH24]] Suppose A ∗ is 1-strongly convex and ℓ is smooth with respect to A ∗ with parameter L . Let Assumptions 3 and 4 hold. For step size scheme

we have

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 1 and 2, whereas Theorem H.4 relies on Assumptions 3 and 4 together with the 1-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 a 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 1-strongly convex.

Second, we will prove that Assumption 4 is implied by the PL inequality (10) in ˜ Ω under Assumption 2, by proving that BFBE (non-Euclidean measure of stationarity) is greater than ∥∇ ℓ ( ω ) ∥ 2 (Euclidean measure of stationarity) up to a constant in the following proposition.

Proposition 5 (PL implies Bregman Prox-PL under Assumption 2) . Under Assumption 2, for sufficiently large ρ which may depend on ω , it holds that

Together with (10) , we conclude that Assumption 4 holds with parameter

Remark. In Theorem 4.2, we require that Proposition 5 holds for all iterates { ω t : 0 ≤ t ≤ T -1 } . Since this set is finite, we can take the maximum value of ρ corresponding to all ω t , ensuring that Proposition 5 holds uniformly along the entire trajectory.

The proof of Proposition 5 relies on another metric of first-order condition, which is closely related to BFBE:

Definition H.1 (Bregman Gradient Mapping (BGM)) . For some ρ > 0 , the BGM at ω ∈ Ω is defined as

where D sym A ∗ ( ω, ω + ) := D A ∗ ( ω, ω + ) + D A ∗ ( ω + , ω ) is the symmetrized Bregman divergence, and

The following lemma from [FH24] reveals the connection between the two measures:

Lemma H.5. [Lemma 4.2 in [FH24]] For any ω ∈ Ω and any ρ > 0 , we have

Proof of Proposition 5. Recall in Theorem H.1, we have shown that A ∗ is C L -smooth and C S -strongly convex with respect to the Euclidean norm in the bounded set ˜ Ω . By Lemma H.5 and smoothness of A ∗ , we obtain the following lower bound of BFBE

where ω + = argmin ω ′ ∈ ˜ Ω ⟨∇ ℓ ( ω ) , ω ′ ⟩ +2 ρD A ∗ ( ω ′ , ω ) , following the definition in Definition H.1. Then, we can show that without domain constraint, ω ∗ := argmin ω ′ ∈ Ω ⟨∇ ℓ ( ω ) , ω ′ ⟩ +2 ρD A ∗ ( ω ′ , ω ) satisfies

By strong convexity of A ∗ , we have

If ω + = ω ∗ holds, we can plug (39) into the lower bound of BFBE (38) and we conclude that

This completes the proof of Proposition 5. Therefore, it remains to prove ω + = ω ∗ .

For ω lying in the interior of ˜ Ω , the condition holds for sufficient large ρ as ∥∇ ℓ ( ω ) ∥ is bounded on ˜ Ω and 2 ρD A ∗ ( ω ′ , ω ) dominates in the optimization problem. The strong convexity of D A ∗ then ensures that ∥ ω ∗ -ω ∥ is very small.

For ω ∈ ∂ ˜ Ω , we will need Assumption 2. Since ∇ A ( η ) and ∇ A ∗ ( ω ) are inverse operators of one another (see Section C.1), (39) implies

as ρ →∞ .

Then under Assumption 2, for any ω ∈ ∂ ˜ Ω and any outward normal direction n ω , there exists some ε > 0 such that

Then we have

for some sufficiently large ρ . Hence ω ∗ ∈ ˜ Ω and ω + = ω ∗ .

Step 2. Now we aim to strengthen the bound in Theorem H.4 to a last-iterate bound. In the proof of Theorem H.4 (which can be found in Appendix D in [FH24]), we can derive a recursion of form

Therefore, we do not need to assume that Λ τ ≥ ε for all τ = 0 , · · · , t to obtain the same recursion. As a 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 be satisfied for a univariate Bayesian linear regression model, if the parameters U and D satisfy (14).

Proof. Using the explicit form of ∇ 2 A ∗ ( ω ) in Section C.1, we can compute

Since

Then we have and

Next, we consider the 4 constraints which define the boundary of ˜ Ω .

ξ = U . If this constraint is active, the unique outward normal direction is n ω = (1 , 0) , and Assumption 2 requires U ( x 2 +1) > xy .

2. Ξ -ξ 2 = D . If this constraint is active, the unique unnormalized outward normal direction is n ω = ( -2 ξ, 1) , and Assumption 2 requires D ( x 2 +1) -1 > 0 .

4. Ξ -ξ 2 = D -1 . If this constraint is active, the unique unnormalized outward normal direction is n ω = (2 ξ, -1) , and Assumption 2 requires -D -1 ( x 2 +1) + 1 > 0 .

3. ξ = -U . If this constraint is active, the unique outward normal direction is n ω = ( -1 , 0) , hence Assumption 2 requires -U ( x 2 +1) < xy .

For ω 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.

I Variance of the Gradient Estimator in Logistic Regression

In this section, we aim to show that Assumption 1 is satisfied for logistic regression in mean-field setting. We assume that the dataset D = { ( x i , y i ) : x i ∈ R , y i ∈ {-1 , 1 }} n i =1 . We begin with Theorem I.1 to bound the gradient of the KL divergence term. In Theorem I.2 we will bound the gradient associated with the log-likelihood term.

Proof. Using the closed-form solution of KL divergence between two multivariate Gaussian distributions and the chain rule, we have

Therefore, we have

where we write ∇ ξ D KL ( q ( z ; ω ) ∥ p ( z )) as a vector in R d and use the fact that U, D ≥ 1 .

Lemma I.2. In D , denote s 1 = max i ∥ x i ∥ and s 2 = max i ∥ x i ⊙ x i ∥ . Then for all 1 ≤ i ≤ n and z ∈ R d we have

Here ∇ 2 z log p ( y i | x i , z ) is defined as a d -dimensional vector where j -th entry equals ∇ 2 z j z j log p ( y i | x i , z ) .

Proof. For the j -th entry,

where σ ( · ) is the sigmoid function. Then we have

With the same approach as in Lemma I.2, we have the following bounds on the log-likelihood term and the gradient ∇ ℓ ( ω ) .

Next, we define our gradient estimator. Consider mini-batch gradient estimator

where each i k 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-likelihood, thus we approximate the expectation with N samples, i.e.,

where z 1 , · · · , z N are iid samples of the variational distribution q . It is obvious that E [ ˆ ∇ MC ℓ ( ω )] = ˆ ∇ ℓ ( ω ) , so the gradient estimator is unbiased.

Similar to (41), we can obtain the following bound on the stochastic gradient ˆ ∇ ℓ ( ω ) .

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.,

Proof.

Finally, we are ready to prove that Assumption 1 holds for logistic regression.

Theorem I.4. Consider stochastic gradient estimator (42) . Then with s 1 = max i ∥ x i ∥ and s 2 = max i ∥ x i ⊙ x i ∥ , Assumption 1 is satisfied with some V 2 > 0 , which is at most polynomial in n, d, U, D, s 1 and s 2 .

Proof. We first summarize our notations.

Here ω + t +1 , ∗ (and ω t +1 , ∗ ) are obtained by doing an (S)MD update from ω t with step size γ t , ω t +1 = argmin ω ∈ ˜ Ω D A ∗ ( ω, ω t +1 , ∗ ) and ω + t +1 = argmin ω ∈ ˜ Ω D A ∗ ( ω, ω + t +1 , ∗ ) . We decompose (12) into the sum of two terms.

Next, we aim to bound ∥ ω t +1 -ω t ∥ . By C S -strong convexity of A ∗ and definition of ω ∗ t +1 , we have

Moreover, by strong convexity we have

Then by triangle inequality we get

Then we use the definition of SMD step in (11) to get

Similarly, we have

Plug (45) and (46) into (44) and we get

J Experiment Details and Additional Results

J.1 Pseudocode of the Algorithms

We first present the pseudocode of Proj-SNGD algorithm proposed in Section 4.1.

Algorithm 1 Proj-SNGD

Require: Initialization ω 0 ∈ ˜ Ω , number of iterations T , step sizes { γ t } 0 ≤ t ≤ T -1

2: Compute stochastic gradient ˆ ∇ ℓ ( ω t )

1: for t = 0 , 1 , . . . , T -1 do

3: Compute ω t +1 , ∗ ∈ Ω such that

6: Sample ¯ ω T from { ω t } 0 ≤ t ≤ T -1 with probability p t = γ t / ∑ T -1 i =0 γ i

4: Set

5: end for

7: Return ¯ ω T

Next, we provide the pseudocode of Prox-SGD and Proj-SGD from [Dom20, DGG23]. We write

Therefore, using the definition of L ( θ ) in (30), we have

Algorithm 2 Prox-SGD

Require: Initialization µ 0 ∈ R d , C 0 ∈ S d + and diagonal, iterations T , step sizes { γ t } 0 ≤ t ≤ T -1 ,

2: Compute stochastic gradient ˆ ∇ µ ˜ L ( µ t , C t ) , ˆ ∇ C ˜ L ( µ t , C t )

1: for t = 0 , 1 , . . . , T -1 do

3: Set µ t +1 = µ t γ t ˆ µ ˜ L ( µ t , C t ) , C t +1 , ∗ = C t γ t ˆ C ˜ L ( µ t , C t

4: For each 1 ≤ i ≤ d , set

-∇ -∇ )

5: end for

6: Return ( µ T , C T )

In Algorithm 3, M is defined as the smoothness parameter of the joint log-likelihood log p ( z, D ) with respect to z . Recall that the clipping function in (50) is defined as clip [ a,b ] ( x ) = min { max { a, x } , b } .

Algorithm 3 Proj-SGD

Require: Initialization µ 0 ∈ R d , C 0 ∈ S d + and diagonal, iterations T , step sizes { γ t } 0 ≤ t ≤ T -1 , smoothness parameter M

1: for t = 0 , 1 , . . . , T -1 do

3: Set µ t +1 = µ t -γ t ˆ ∇ µ L ( µ t , C t ) , C t +1 , ∗ = C t -γ t ˆ ∇ C L ( µ t , C t )

2: Compute stochastic gradient ˆ ∇ µ L ( µ t , C t ) , ˆ ∇ C L ( µ t , C t )

4: For each 1 ≤ i ≤ d , set

5: end for

6: Return ( µ T , C T )

0

100

200

300

400

500

Epoch

4000

6000

8000

10000

Negative ELBO

Prox-SGD

Proj-SGD

SNGD

Proj-SNGD

10

-

3

10

-

2

10

-

1

γ

500

1000

1500

2000

Number of Iterations

Prox-SGD

Proj-SGD

SNGD

Proj-SNGD

0

0

500

1000

1500

2000

2500

3000

Epoch

20000

40000

60000

80000

100000

120000

140000

Negative ELBO

Prox-SGD

Proj-SGD

SNGD

Proj-SNGD

10

-

4

10

-

3

γ

0

8000

10000

12000

14000

Number of Iterations

Prox-SGD

Proj-SGD

SNGD

Proj-SNGD

J.2 Experimental Results on Additional Datasets

In this section, we 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 implementation can be found in Section J.4.

We consider logistic regression on Madelon dataset ( n = 2600 , d = 500 ). We use mini-batches of size 2000 and set the step size γ t = γ 0 / √ t , where γ 0 is a hyperparameter to be tuned. We run the algorithm for 1000 epochs (2000 iterations). The results of 5 independent runs are shown in Figure 4. Similar to the results of the MNIST experiment, we observe that Proj-SNGD slightly outperforms SNGD , and both non-Euclidean algorithms achieve faster convergence than Euclidean counterparts. Moreover, non-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 a subset of CIFAR-10 dataset with pictures of cats and dogs ( n = 10000 , d = 3072 ). Weuse mini-batches of size 2000 and set the step size γ t = γ 0 / √ t , where γ 0 is a hyperparameter to 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. In the right panel of Figure 5. Notably, non-Euclidean algorithms can accommodate larger step sizes and are the most robust to step size tuning.

0

200

400

600

800

1000

Epoch

600

800

1000

1200

1400

Negative ELBO

U=

∞

, D=

∞

U=4, D=20

U=3, D=15

U=2, D=10

γ

0

0

100

200

300

400

500

4000

6000

8000

10000

Negative ELBO

U=

∞

, D=

∞

U=3, D=400

U=2, D=300

U=1.5, D=200

10

-

3

10

-

2

10

-

1

500

1000

1500

2000

Number of Iterations

U=

∞

, D=

∞

U=3, D=400

U=2, D=300

U=1.5, D=200

Epoch

γ

0

J.3 Choice of U and D in 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 = ∞ , D = ∞ ), projection with U = 4 , D = 20 , U = 3 , D = 15 and U = 2 , D = 10 . Smaller values of U and D yield a smaller smoothness coefficient and a larger hidden convexity coefficient, leading to stronger theoretical guarantees.

Madelon Dataset. For this dataset, we consider the following four different settings: no projection, 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 as those used for the MNIST dataset: no projection, projection with U = 4 , D = 20 , U = 3 , D = 15 and U = 2 , D = 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 a sub-optimal solution. The right panel of Figure 6 confirms that smaller U and D lead to faster convergence. Therefore, a 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 projection 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 7 and 8), and Proj-SNGD will converge to a suboptimal solution. If they are too large, the algorithm will not benefit from projection (see SNGD without projection in Figures 6 to 8).

The slow convergence of SNGD without projection is due to the poor tolerance for large step sizes, which causes divergence in 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 is consistently observed across 3 different datasets. Moreover, this effect is consistent with our theoretical upper bounds for the relative smoothness parameter. This is because our upper bound

Number of Iterations

6000

5000

4000

3000

2000

1000

10

-

2

10

-

U=

, D=

U=4, D=20

∞

∞

U=3, D=15

U=2, D=10

1

0

500

1000

1500

2000

2500

3000

Epoch

20000

40000

60000

80000

100000

120000

140000

Negative ELBO

U=

∞

, D=

∞

U=4, D=20

U=3, D=15

U=2, D=10

10

-

4

10

-

3

γ

9000

10000

11000

12000

13000

14000

15000

Number of Iterations

U=

∞

, D=

∞

U=4, D=20

U=3, D=15

U=2, D=10

0

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-SNGD . We use N = 2000 samples from variational distribution q to estimate the expected log-likelihood (see the definition of stochastic gradient estimator (42)).

In the left panel of Figure 3 (MNIST dataset), we set γ 0 = 0 . 05 for Proj-SNGD and SNGD , and γ 0 = 0 . 01 for Prox-SGD and Proj-SGD . In the left panel of Figure 4 (Madelon dataset), we set γ 0 = 0 . 5 for non-Euclidean algorithms, and γ 0 = 0 . 01 for Euclidean algorithms. In the left panel of Figure 5 (CIFAR-10 dataset), we set γ 0 = 10 -3 for non-Euclidean algorithms, and γ 0 = 2 × 10 -4 for Euclidean algorithms.

All experiments, except for CIFAR-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 it takes about 1 minute. The CIFAR-10 experiment is run on an NVIDIA GeForce RTX 3090 GPU, requiring roughly 8 minutes for 3000 epochs.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n### 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.\n\nQUESTIONS:\nn/a\n\nLIMITATIONS:\nThe limitation section in the paper is appropriate.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-18"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrenghs\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.\n\nQUESTIONS:\nWhy 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.\n\nLIMITATIONS:\nYes, 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nNot Applicable.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\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* 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?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-01"}