Nearly Dimension-Independent Convergence of Mean-Field Black-Box Variational Inference
Authors: Kyurae Kim
Yian Ma
Trevor Campbell
Jacob R. Gardner
Venue: NeurIPS 2025 poster ()
Keywords: variational inference
variational Bayes
mean-field approximation
stochastic optimization
SGD
Bayesian inference
Primary area: probabilistic_methods
TL;DR: We prove that black-box variational inference with the mean-field Gaussian variational family converges in a rate with an explicit dimension dependence of only $O(\log d)$.

Abstract: We prove that, given a mean-field location-scale variational family, black-box variational inference (BBVI) with the reparametrization gradient converges at a rate that is nearly independent of explicit dimension dependence. Specifically, for a $d$-dimensional strongly log-concave and log-smooth target, the number of iterations for BBVI with a sub-Gaussian family to obtain a solution $\epsilon$-close to the global optimum has a dimension dependence of $\mathrm{O}(\log d)$. This is a significant improvement over the $\mathrm{O}(d)$ dependence of full-rank location-scale families. For heavy-tailed families, we prove a weaker $\mathrm{O}(d^{2/k})$ dependence, where $k$ is the number of finite moments of the family. Additionally, if the Hessian of the target log-density is constant, the complexity is free of any explicit dimension dependence. We also prove that our bound on the gradient variance, which is key to our result, cannot be improved using only spectral bounds on the Hessian of the target log-density.

Nearly Dimeensin-Inndeeedeett Convergence of Mean-Field Black-Box Variational Inference
Kyurae Kim University of Pennsylvania kyrkim@seas upenn.
Yi-An Ma University of California San Diego yianmaQucsd. edu
Trevor Campbell
Jacob R. Gardner University of Pennsylvania jacobrgšseas. upenn edu
University of British Columbia trevor@stat.o ca
Abstract
We prove that, given mean-field location-scaa variational family, black-box vari- ational inference (BBVI) with the reparametrization gradient converges at rate that is nearly independent of any explicit dimension dependence. Specifically, for a d-dimensional strongly log-concave and log-smooth target, the number of iter- ations for BBVI with a sub-Gaussian family to obtain a solution e-close to the global optimum has an explicit dimension dependence no larger than O(log d). This a significant improvement over the (dd) dependence of full-rank location- scale families. For heavy-tailed families, we prove a weaker O(d//k) dependence, where & is the number of finite moments of the family. Additionally, if the Hessian the target log-density is constant, the complexity is free of any explicit dimen- sion dependence. We also prove that our bound on the gradient variance, which is key to our result, cannot be improved using only spectral bounds on the Hessian of the target log- density..
Introduction
Variational inference (VI; Blei et al., 2017; Hinton and Camp, 1993; Jordan et al., 1999; Peterson and Hartman, 1989) is an effective method for approximating intractable high-dimensional distribu- tions and models with tall datasets. Among various VI algorithms, black-box VI (BBVI; Kucukelbir et al., 2017; Ranganath et al., 2014; Titsias and Lázaro-Greddee which minimizes the exclusive KL divergence (Kullback and Leibler, 1951) via stochastic gradient descent (SGD: Bottou et al.. 2018; Robbins and Monro, 1951) in the space of parameters, is widely used due to flexibility apply to wide range of variational families with only minor modifica- tions (Bingham et al., 2019; Carpenter et al., 2017: Fjelde et al., 2025; Ge et al., 2018; Patil et al., 2010). Specifically, location-scale variational families-ae in which a base distribution is mutated by an affine transformationn-remmie
2014; Wingate and Weber, 2013),
popular choice, encompassing those with diagonal scale matri- ces (the meann-felld approximation; Hinton and van Camp, 1993; Peterson and Hartman, 1989), as well as scale matrices with low rank (Ong et al., 2018: Rezende et al., 2014; Tomezak et al., 2020) and full-rank (Kucukelbir al., 2017: Titsias and Lázaro- GGeeila, 2014) factors.
The choice of the variational family is generally known to affect the convergence speed of BBVI. where families that are more "expressive." those that contain more complex distributions, result in slower convergence. For example, in location-scal families, it has been empirically observed that mean-field families often provide faster convergence to an accurate posterior approximation than full-rank families (Agrawal et al., 2020; Giordano et al., 2018, 2024; Ko et al., 2024; Zhang et al., 2022). This because full-rank families often require running SGD with a smaller step size and
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

for longer; even given a large computation budget, BBVI on full-rank family may not converge adequately (Ko al., 2024). Therefore, choosing the expressiveness of the family corresponds to trading statistical accuracy for computational efficiency (Bhatia et al., 2022). In order to control this trade-off for our benefit, a clear theoretical understanding of the relationship between convergence speed and expressiveness is needed.
Formally, consider the setting of approximating a ju-ttrongly log-concave and L-log-smooth target, where K. A L/p j. is the condition number. For BBVI with the reparametrization gradient (Kingma and Welling, 2014; 4; Rezende et al., 2014; Titsias and Lázaro-Grediilla, 2014) on a full-rank location- scale family, an e-close solution to the global optimum in squared distance in parameter space can be obtained after at least O(dk² (0) iterations (Domke, 2019; Kim et al., 2023a). For mcan-field location- scale families, on the other hand, the iteration complexity improves to et al., 2023a). While this clearly better than the (cd) explicit dimension dependence of full-rank families, it has been conjectured that a better dependence is more likely (Kim et al., 2023a).
(Kim
In this work, we positively resolve this conjecture by obtaining stronger convergence guarantees for BBVI on mean-field location-scale families (Section 3). In particular, under the conditions stated above, we prove that BBVI with a mean-field location-scale family with sub-Gaussian tails can obtain an e-accurate solution in squared distance after O d)x²E-) iterations. Hcavier-tailed families achieve weaker O(d2/&x²c=) iteration complexity guarantee, where k is the number of finite moments of the variational family. For the Student-t variational family with high-enough degrees of freedom v, this corresponds to ((d2/(=22) explicit dimension dependence. In addition, if the Hessian of the target log-density is constant, any mean-field location-scale family attains a O(x"E iteration complexity without any explicit dependence on d.
The key element of the proof is careful probabilistic analysis of the variance of the reparametriza- tion gradient (Section 4): In general, the reparametrizztion gradient of the scale parameters contains heavy-tailed components that grow not-so-slowly in d. However, for mean-field families, only a single random coordinate turns out to be heavy-tailed. Through a probabilistic decomposition, the influence of this heavy-tailed component can be averaged out over all d coordinates. Then the lighter-tailed components of the gradient dominate d increases, resulting in benign dimension dependence (Lemma 1.... We also provide lower bound (Proposition 4.2) showing that our analy- sis cannot be improved when using only spectral bounds on the Hessian of the target log-density.
Preliminaries
We denote random variables in sans serif (e.g.. U. U). 5% C Rdxd denotes the set of d x positive definite (PD) matrices, Ddd C Rdxd denotes the set of diagonal matrices, and D C Dd n n is its positive definite subset. (:-)) and |1-112 denote the Euclidean inner product and norm. For matrix A E Reddd ||Al| / /tr(AT A) is the Frobenius norm, ||Al|2 ma) (A) is the lg operator norm, where mmax(.) and Gmin (.) are the largest and smallest singular values.
Notation
2.1 Problem Setup
Our problem of interest is an optimization problem over some space A C IP of the form of
minimize XEA
(F()) e f(2) + (A)) where S(A) Ezzga e (z)
é: R2 + R is measurable function we refer to the "target function", h A + R is potentially non-smooth convex regularizer, and the expectation Ezzga {(Z) is assumed to be intractable.
BBVI is special case of Eq. (1) where f = log is the negative (unnormalized) log-density of some distribution x with respect to the Lebesgue measure and h(X) = -H[[x)] is the negative differential entropy of qx- Then F is the exclusive Kullback-Leibler divergence DKL. (Kullback and
Leibler, 1951) up to an additive constant (Jordan et al., 1999), where Eq. (1) reduces to
minimize xEA
{
(DKL TT x - EEzgg log (Z) - H (qx))
(2)
We assume is supported on Rd, which, unless discrete-valued variables are involved, is often valid after appropriate support transformations (Kim et al., 2023a, $2.2). Such a setup for BBVT has been proposed by Kucukelbir et al. (2017), and now encompasses most practical use of BBVI with the

---

reparametrization gradient as implemented in Stan (Carpenter et al., 2017), PyMC (Patil et al., 2010), Pyro (Bingham et al., 2019), and Turing (Fjelde et al., 2025; Ge al., 2018).
For the purpose of a quantitative theoretical analysis, we will consider the following properties: R, we say & is L-(Lipschitz Jsmooth if there exists some L E (0, +00) such that. for all 2, E R
Definition (Smoothness). For some Φ Rdd
[[Vo(=) Vo(z')ll2 S L||= 2/12-
Definition (Strong Convexity). For some á : R + R, we say o is u-strongly convex if there exists some constant ja (0, L] such that, for all 2, 2' E Rdd
(Vo(=), 2²) 6(2) 6(z')
2 -|
In the context of BBVI, assuming that & log π is both p-strongly convex and L-smooth is equivalent to assuming is p-strongly log-concave and L-log-Lipschitz smooth, respectively, which common in the analysis of MCMC (Chewi, 2024) and (Arnese and Lacker, 2024; Diao et al., 2023; Domke et al., 2023; Kim et al., 2023a; Lambert et al., 2022; I Lavenant and Zanella, 2024).
2.2 Variational Family
We consider the location-scale family (Casella and Berger, 2001, $3.5):
Definition 2.1 (Location-Sae Scale Variational Family). A family of distributions Q is referred to as location-scale variational family if there exists some univariate distribution ip dominated by the Lebesgue measure such that each member of Q indexed by X = (m,C) E Rd x C, where C C Rdxd and qx E Q, satisfies
Z ~ yX
z = Tx (4)
where
Tx (U) & Cu+m,
(UI,..., Ua);
U Lid. 8:
and d is equivalence in distribution. Then Tx is referred to as the "reparametrization functionn" while and C are referred to the location and scale parameters, respectively.
In addition, we impose mild regularity assumptions on the moments of the base distribution:
Assumption 2.2. &p satisfies the following: (i) It is standardized such that IEu, = 0 and Eu? 1, (ii) symmetric such that EU;"" 0, and (ii) its kurtosis is finite such that Eu,
The location-scale family with Assumption 2.2 encompasses many variational families used in prac- tice, such as Gaussians, Student-f with a high-enough degrees freedom v. Laplace, and so on, and enables the of the reparametrization gradient.
While the choice of yp gives control over the tail behavior the family, the choice of the structure of the scale matrix C' gives control over how much correlation between coordinates of l the varia- tional approximation can represent. This ability to represent correlations is often referred to as the "expressiveness" of variational family, where the most expressive choice is the following:
Definition 2.3 (Full-Rank Location-See Scale Family). We say Q is a full-rank location-scaa family if it satisfies Definition 2.1 and, for any C E C. C. is invertible and the squared Cs, CCT, sppn the whole space of dense Rdd posilive definite matrices as {CCT C S4
Typically, full-rank location-scale families are formed by setting C to be the set of invertible trian- gular matrices (the "Cholesky factor parametrization" Kucukelbir et al., 2017; Titsias and Lázaro- Gredilla, 2014) or the set of symmetric square roots (Domke, 2020; Domke et al., 2023). Adding further restrictions focus on the case where C' C is restricted to be diagonal such that C C D which known as the mean-field approximation (Hinton and van Camp, 1993; Peterson and Hartman, 1989):
C forms various subsets of the broader location-scale family. In this work,
Definition 2.4 (Mean-Fieh) Field Location-Seen Scale Family). We say Q is mean-field lecation-ssall if it satisfies Definition 2.1 and all C' C are diagonal such that C C D
2.3 Algorithm Setup
Recall that BBVI is essentially SGD in the space of parameters of the variational distribution. There- fore, we have to define the space of parameters. For this, we use the "linear" parametrization:
A {> (m, C)) m Rd C e [ c R..
(3)
3

---

Under this parametrization, the desirable properties of l easily transfer to f. For instance, if & p-strongly convex and L-smooth, f is also μestrongly convex and L-smooth (Domke, 2020). This contrasts with "non-linear parametrizations" commonly used in practice, such as making the diagonal positive by C exp(Ac.)). )- Such practice rules out transfer of strong convexity and smoothness (Kim et al., 2023a) unless constraints such as C₁ 2 5 for some 8 >> 0 are enforced (Hotti et al., 2024), (Though they can sometimes be beneficial by reducing gradient variance; Hotti et al., 2024; Kim et al., 2023b.) The flip side of using the linear parametrization is that we must now enforce the constraint > O. Furthermore, h then becomes non-smooth with respect to C:
h(A)
(xx(
-lo det Cl d (p))
-E log Cal d (+)
(4)
This corresponds to log-barrier functions (Parikh and Boyd, 2014, $6.7.5), which are non-smooth. Thus, the optimization algorithm must somehow deal with these difficulties (Domke, 2020).
In this work, we will rely on the proximal variant of stochastic gradient descent (SGI; Bottou, 1999; Bottou et al., 2018; Nemirovski et al., 2009; Robbins and Monro, 1951; Shalev-Shwartz et al., 2011), often referred to stochastic proximal gradient descent (SPGD; Nemirovski et al., 2009). Proximal methods are a family of methods that rely on proximal operators (Parikh and Boyd, 2014), which well defined as long as the following hold: Assumption 2.5. h
R U {+oo] convex, bounded below, and lower semi-continuous.
The non-smoothness of h and the domain constraint are handled by the proximal operator
prox.,, - arg min { (n(x)++++)))))) 1 >
XEA
while the intractability of is handled through stochastic estimates of V (Definition 2.6). For step size schedule (7z) f, an unbiased estimator of V.((A)) = EV fAx U). and a sequence of i..dd noise (U:) for each / 2 0, SPGD iterates
A++1 prox (A+ VT( u.)
In the case of BBVI with a mean-field location-scale family, the proximal operator of Eq. (4) is identical that of log-barrier functions (Parikh and Boyd, 2014, $6.7.5):
prox, & (A = (m, C)) (m, C"), where C₁ = (1/2) On + (C
Instead of using SPGD, one can also use projected SGD, where C' is projected to subset where F is smooth (Domke, 2020) and the "closed-form entropy" gradient VF V Vh Vh (Kucukelbir et al., 2017; Titsias and Lázaro-Gredinn indistinguishable (Domke al., 2023), and the need for setting a closed domain of C' is inconvenient. Therefore, we only consider SPGD. But our results can easily be applied to projected SGD.
2014). However, the resulting theoretical guarantees are
For f. we will use the classic reparametrization gradient (Ho and Cao, 1983; Rubinstein, 1992): Definition 2.6 (Reparametrization Gradient). For differentiable function E R R.
v)(A.)) & V (T) (Tx(u) V8 (Tx (u)) where U U~ P
an unbiased estimator of V f such that VAEZ~g ((z) = (A))..
The reparametrization gradient, also known as the push-in gradient or pathwise gradient, was in- troduced to VI by Kingma and Welling (2014); Rezende et al. (2014); Titsias and Lázaro-Gredilla (2014). (See also the reviews by Glasserman 1991; Mohamed et al. 2020; Pflug 1996.) empir- ically observed to outperform alternatives (Kucukelbir et al., 2017; Mohamed al., 2020) such as the score gradient (Glynn, 1990; Williams, 1992) and de facto standard whenever is differentiable. (Though theoretical evidence of this superiority is limited to the quadratic setting; Xu et al., 2019.)
2.4 General Analysis Stochastic Proximal Gradient Descent
Analyzing the convergence of BBVI corresponds to analyzing the convergence of SPGD (or more broadly, of SGD) for the class of problems that corresponds to BBVI. For this, we will first discuss sufficient conditions for the convergence of SPGD and the resulting consequences.
Assumption 2.7 (Lipschitz Gradients in Expectation). There exists some constant L E [0, oo) such that, for all X, X E A,
(I(ī)();)) A](: u) < L X112

---

Assumption 2.8 (Bounded Variance). There exists some constant a [0, 00) such that, for all
E
A_ E arg min F()).
Both assumptions were initially used by Bach and Moulines (2011, Assumptions H2 and H4) to analyze the convergence of SGD. Here, Assumption 2.7 serves as an analog of L-smoothness, and thus determines the largest stepsize we can use. The strategy of combining Assumptions 2.7 and 2.8 is referred to "variance transfer" (Garrigos and Gower, 2023, $4.3.3). Previously, for analyzing BBVI, slightly different assumption called saadrratcally-bounder variance (V)>hich assumes
the existence of a, 3 E [0, +oo) such that, for all À E A, E| [A; U) < a || X. ll + 3 holds- has been commonly used (Domke, 2019; Domke et al., 2023; Kim al., 2024b). While similar, our assumptions result in a constant-factu improvement in the resulting bounds.
For the analysis, will use two-stage step size schedule (Gower et al., 2019, Theorem 3.2):
70
if L S 4. if 2 +1, where 0 70 S 2L
P
2(+T))+ (c+r+1)=
(5)
This operates by first maintaining a fixed step size 7o until some switching time l. E (O.....7
7).
and then switches to the 1/1 schedule of Lacoste-Julien al. (2012) with an offset r 2
Under Assumptions 2.7 and 2.8, we can now provide a complexity guarantee for solving Eq. (1) via SPGD. Since BBVI consists of a subset of Eq. (I), establishing Assumptions 2.7 and 2.8 and invoking the following result will constitute our complexity guarantee for BBVI.
Proposition 2.9. Suppose is u-strongly convex, h satisfies Assumption 2.5, and f satisfies As- sumptions 2.7 and 2.8. Then, for the global optimum X. arg minxes F(A), A _ IlAo Xa 2l, some la, T, and 7/0 (explicit the proof), SPGD with the step size schedule in Eq. (5) guarantees
T T 0 0² oC 1
log 12a
C
A2
+ u²a
L
log A² + +1
=
E|AT 2..
Ve'
Proof. See the full proof in Appendix B.1.1, p. 19.
This result is a slight improvement over past analysis of SPGD with Eq. (5) (Domke al., 2023, Theorem 7). In particular, the dependence on the initialization A has been improved to be loga- rithmic instead of polynomial. Furthermore, it encompasses the case where we have "interpolation" (o2 = O; Kim et al., 2024b; Schmidt and Roux, 2013; Vaswani et al., 2019) automatically resulting O(lg 1/c) complexity. A similar result for vanilla SGD, where the dependence on both and A is optimized simultaneously, was reported by Stich (2019). However, this result required schedule that depends on the maximum number of iterations T. Compared with this, our result provides an any-time guarantee that holds for any number of iterations.
For non-strongly convex f, using the strategy of Domke al. (2023, Theorem 8 and 11) should yield corresponding (1/22) complexity guarantee under the same set of assumptions. However, this requires fixing the horizon T in advance, and it currently unknown how to obtain an anytime O(I/VT) convergence bound for SGD under Assumptions 2.7 and 2.8 or QV. If one moves away from the canonical SGD update by incorporating Halpern iterations (Halpern, 1967), it is possible to obtain any-time convergence under QV-like assumption (Alacaoglu et al., 2025).
3
Main Results
3.1 General Result
For our results, we impose an additional assumption that is generalization of L-smoothness under twice differentiabilityy of f.
Assumption 3.1. f is twice differentiable and, for all z E Rdd, there exist some matrix H E Rddx and constant [0, 20) satisfying
| < x and [| H|I 58.
Notably, if is twice differentiable, j-strongly convex, L-smooth, it already satisfies Assumption 3.1 with H = L and 8 477. If é is only I-smooth, it satisfies it with H = Oaxa and å = L. The key advantage of this assumption, however, that it characterizes Hessians that are not necessarily well-conditioned, but almost constant. This crucially affects the dimension dependence.
5

---

Given our assumptions on the target function É, variational family Q, and our choice of gradient esti- mator, we can guarantee that SPGD applied problem structure corresponding to BBVI (Eq. (1))
achieves a given level of accuracy after (g H, 8, 14, (..,,,,)))),))) number of iterations:
Theorem 3.2. Suppose the following hold:
1. f is p-strongly convex and satisfies Assumption 3.1 and S mmin (H) S mmax (1) S L.
2. h satisfies Assumption 2.5.
3. Q a mean-field location-scale family, where Assumption 2.2 holds.
4. Vf is the reparametrization gradient.
Denote the global optimum X = (m,,C.) = arg min F(X), the irreducible gradient noise as of - |mm,- =112 + ||C lijs and the stationary point of l as z 2 arg ming f(z). Then, for some
Lx T, and To (explicit in the proof), SPGD with the step size schedule Eq. (5) guarantees
T
0 gld, H. v) (o2c + log (11 - x.13)c-1/2)
E|AAT A. sc,
where
g(d, H. p. .p) 2 21 (T) ( 4(62/ (1/2) + max j=I....d
Ut)
Proof. The full proof can be found in Appendix B.2.1, p. 25.
- Tx (Lemma A.3), which is the squared cost of a coupling between Qyr and qx our guarantee also translates to guarantee in Wasserstein- 2 distance: ET 2.112 < + EW2(gar: )2.)² < In the general case where d > O, the dimension dependence enters through E maxxx=1....d UT. which depends on the order-statistics of the base distribution sp. In case E is quadratic, corresponding distribution in the BBVI context, there exists some such that (=((2)= = H for all x E Rd. Thus, independent convergence rate. We will Assumption 3.1 holds with d 0, implying dimension-independen present additional special cases with more explicit choices of & in the next section.
Due to the identity | X12 Ey-ps ( [TA
æ being Gaussian target
In case we do not want to assume Assumption 3.1 and only assume that is p-strongly convex and L- smooth instead, we can replace them with the generic choices of H L and ó - which hold for all is that are p-strongly convex, I-smooth, and twice differentiable. This then makes the role of the condition number к. & L/p more explicit.
Corollary 3.3. Suppose is is twice differentiable, u-strongly convex, and L-smooth. Then, denot- ing the condition number as & L/ p. Theorem 3.2 holds with L-
g d, I
2 4p (1/2) + ra)(k+ (K (1 + + (r j= maxx U
This makes the O(k²) condition numher dependence explicit, but the downside that we lose di- mension independence in the case of ill-conditioned quadratic !s. This fact suggests that dimension dependence is more fundamentally related to how close the Hessian is to a constant rather than how well-conditioned it is.
3.2 Special Cases with Benign Dimension Dependence
We now present some special cases of Theorem 3.2, which has yet exhibit explicit dependence on dimensionality. As mentioned in the previous section, dimension dependence depends on the order statistics of ip, which is related the tail behavior of yp.
Variational Families with Sub-Gaussian Tails. The most commonly used variational family in practice is the Gaussian variational family. More broadly, for sub-Gaussian variational families, u? is sub-cxponential and therefore admits a moment generating function (MGF) (Wainwright, 2019, Theorem 2.6), which leads to Ollog d) explicit dimension dependence.
Proposition 3.4. Suppose there exists some t 0 such that the MGF of uē satisfies Ma²(t) oo. Then
max Uf S (1/1) (log M (1) + logd). 1=1 $=
For example, if q is a standard Gaussian, then
g(d, H, 6, v) < 8(1 + (52 /12) (22 22 + 16 log
Proof. The full proof can be found in Appendix B.2.2, p. 26.

---

Variational Families with Finite Higher Moments. For families with tails heavier than sub- Gaussian, however, Uf may not have an MGF. While we then lose the O(log d) dependence. we may still obtain a polynomial dependence that can be better than O(vd) /d) obtained in previous works (Kim al., 2023b). In particular, the result that will follow states that the highest order of the available moments determines the order of dimension dependence. For Student-t families, this implies that using a high-enough degree of freedom
make the dimension dependence benign.
Proposition 3.5. Suppose, for k 2 the kih moment of Uī is finite as T2k Euf* c0. Then
E max uf S V2 d 1/kk :=.....d
12k
For example, if is a Student-t with 4 degrees freedom and unit variance, then g(d, H, p1 S 8(1)/1///²) (6²) //²) 16 + /2 v²d )²)=
Proof. See the full proof in Appendix B.2.3, p. 28.
4 Analysis of Gradient Variance
4.1
Overview
The key technical contribution of this work is analyzing the gradient variance and thus establishing the constants L (Assumption 2.7) and o² (Assumption 2.8), which boils down to analyzing
(Tx(U) VE(Tx(U))
(Ux VE(Tx
E
V f(:
vf(:
E
OÅ
Ox
(U)
E
(VE(TA(U) (V(Tx)
where the equality follows from the fact that the Jacobian (T/ does not depend on A. For mean-field location-scale variational families, the squared Jacobian follows as
TTT (U) ax
(T) Ia + U2
OA
where U A diag(U1, , Ua) (Kim al., 2023b). This implies that
E | (A: (T(''
VE - (Tx (0)) (TU
V (Tx (b)))
(6)
Vaste
Our goal is to bound each term by ||> X122
In order to solve the expectations, we need to simplify the Vl terms. For instance, for the gradient of the location Vice, assuming that e is L-smooth allows for quadratic approximation. That is,
Viee Ve (u)) V (Tx (U)) < L (u) Tx (u) -
where the last equality is by Lemma A.3.
Now, it is tempting to use the same quadratic approximation strategy for the gradient of the scale Vicaee Indeed, this strategy was used by Domke (2019) to bound the gradient variance of full-rank location-scall variational families and by Ko al. (2024) for structured location-scale variational families. Unfortunately, this strategy does not immediately apply to mean-field families due to the matrix U. We somehow have to decouple VE(Tx - Ve(Tx and U, but in a way that does not lose the correlation between the two; the correlation leads to cancellations critical to obtaining tight bound. Kim et al. (2023b) used the inequality
Vscale
S
(7)
which resulted in dimension dependence of O(riVd) after solving the expectation. The key question is whether this dimension dependence can be improved. Due to the ordering of norms I1-Il2 & I|., it natural to consider the tighter inequality
Vecale
U (()) - (U))
(This step corresponds to Eq. (8) in the proof sketch of the upcoming result.) The main challenge, however, is solving the resulting expectation in a way that is also tight with respect to d. We will that this requires a careful probabilistic analysis.
7

---

4.2 Upper Bound on Gradient Variance
We now formally state our upper bound on the gradient variance. In the context of proving Theo- rem 3.2, the following lemma implies both Assumption 2.7 and Assumption 2.2. (See the proof of Theorem 3.2.) We provide a corresponding unimprovability result in Section 4.3.
Lemma 4.1. Suppose Assumptions 2.2 and 3.1 hold, Q is mean-field location- and f is the reparametrization gradient. Then, for any X, E R2 x D
E| EI] VX 1)13 < 2(1 + r4) 2 - 462 (12 ( max Ut 01A X112
j=I.....
Proof Sketch. For the proof sketch, we will assume that é is L-smooth instead of taking Assump- tion 3.1. This will vastly simplify the analysis and let us focus on the key elements.
Recall Vicale in Eq. (6). Applying the operator norm and the L-smoothness of yields
Vsuale s ((T
(e(Txx(U))I
S T. (T(()) (u)
(8)
It remains to solve the expectation over Un,..., Ua. Denote (m, C), X (m', C").
m m', and C
recall that Tx(u) = Cu + m (Definition 2.1), and notice that, since Uis diagonal matrix, | ma.,=l,...dd U?. Then we can rewrite Eq. (8) as
Vicale
S
max
03
Ei= (C.:u; + m. C.u, m²)²
I,2
u/)
Ei= (CRU: + +mi)²
max i=l,..,dd
< L'
max i=lgeedd
u/) Σ - (20) uf +2m?)
(Young's inequality)
The problematic term is
Di= = j=l.....d d
C.u = T T(CR, of + + C3..
max
where 1 = arg max; u? is the coordinate of maximum magnitude. Here, uf max;=l....d uf is a heavy-tailed quantity that generally grows fast in d, unlike uf. (e.g. for Gaussian Us uf has MGF but Un does not.) Therefore, a benign dimension dependence might appear futile. Notice, however, that the problematic term only affects a single dimension: the coor- dinate indicated by i. A probabilistic analysis reveals that as d increases, the effect of U becomes averaged out and the effect of the remaining term involving uf dominates. More formally.
Eu E CR
Ei E U 1(.
i + uî uf +i)
where
B 1 i}} E E Bi [}] E = B Blu
Since the maximum of d random variable is always smaller than their sum, the probability of the maximally random event, Píí, = i] = 1/d, kills off the dimensional growth of u³t In fact, using the crude bound Euy < <E>; u = dra, where the last equality is due to Assumption 2.2, is enough to make this term independent of d. The remaining dimension dependence comes from uL:
of E E * {]] E max uf (i. :1] s E max up [[ i
E
max u3
where the last equality follows from Assumption 2.2. Therefore, we finally obtain
Vicale $ 21
Σ
max
43 +r C +r + max I...d=1
j_ u
< 22 (E max U3 + T3 (11m8l2 +1(112
I=......
2L (E max us + r4 || X1I2
The full proof performs an analogous analysis under the more general Assumption 3.1. See the full proof in Appendix B.3.1, 30.
8

---

4.3 Unimprovability
We also demonstrate a lower bound, which implies that Lemma 4.1 cannot be improved when relying on the spectral bounds on V From Eq. (6) and the fundamental theorem of calculus,
E (A 1 z U(VE Ve (
E U (=((z)) 2)dw 12 (9)
where 2 e {z V((=)) = =0} is a stationary point of z a (x)), and z 2 wz + -1- There exists a matrix-valued function with bounded singular values that lower-bounds this quantity:
Proposition 4.2. Suppose Assumption 2.2 holds and Q is a mean-field location-scale family Then, for any > 0, d u. L E (0, +20) satisfying μ S L, there exists a matrix-valued function (z)) Rd + S satisfying pll 3 H z LIa almost surely and a set of parameters X (m, C)
Rd x D such that
E S (Z- =2)d
(L (1-P)²
1
UC joll.s)(E
==.....d1 uf - 1IC0B
where cll, (p) > 0 a constant only dependent on and yp. Proof. The full proof can be found in Appendix B.3.4, p. 36.
For Gaussians, E max; U, is upper bounded O(vd' (Gumbel, 1954, Eq. 1.6), which means the negative term vanishes at a V rate. Furthermore, E max; uf 2 (TE max; u.)² 22oog d) by the well-known lower bound on the expected maximum of i.i.d. Gaussians (Wainwright, 2019, Exercise 2.1 (bb).... Combining these facts with Proposition 4.2 yield a S2(1. log d) bound on Eq. (9).
Remark 4.3. It is not obvious that the rows of our worst-case example H form conservative vector fields. This means that Proposition 4.2 does not assert the existence of a function that satisfies 2 Huurst- However, does suggest that one cannot improve Lemma 4.1 by relying only on spectral bounds on the Hessian.
5 Discussion
5.1 Related Works
Early results analyzing VI had to rely assumptions that either: (i) do not hold on Gaussian targets, (ii) are difficult to verify, or (iii) require bounds on the domain (Alquier and Ridgway, 2020; Buchholz et al.. 2018; Fan et al., 2015; Fujisawa and Sato, 2021; Khan al., 2016; Liu and Owen, 2021: Nguyen et al., 2025; Regier et al., 2017). coordinate-ascen VI (CAVI), in particular, was studied on specific models (Ghorbani et al., 2019; Zhang and Zhou, 2020) only. Under general and verifiable assumptions, Xu and Campbell (2022) obtained asymptotic convergence guarantees, while partial results, such as bounds on the gradient variance (Domke, 2019; Fan et al., 2015; Kim et al., 2023b), or regularity of the ELBO (Challis and Barber, 2013; Domke, 2020; Titsias and Lázaro- Gredilla, 2014), were known.
It was only recently that mon-asympo quantitative convergence under realizable and verifiable as- sumptions was established. For BBVI specifically, Hoffman and Ma (2020) first proved convergence on Gaussian targets (quadratic f), while Domke et al. (2023): Kim et al. (2023a, 2024b) proved the first results on strongly convex and smooth functions with location-scale families. Surendran et al. (2025) extended these results to non-convex smooth functions and more complex variational fam- ily parametrizations, and Cheng et al. (2024) analyzed variant of semi-implicit VI. The results of Domke et al.; Hoffman and Ma; Kim et al., who focused on full-rank families, suggest O(d) dimension dependence in the iteration complexity. On the other hand, Kim et al. (2023a) reported O(vd) dimension dependence for mean-field location-scale families, while conjecturing O(log d) dependence, based on the partial result of Kim et al. (2023b). For targets with diagonal Hessian structure, Ko et al. (2024, Corollary 1) show that mean-fiell families are dimension dimension-indepeeddeer
Apart from BBVI, Wasserstein VI wlhh minimize the KL divergence on the Wasser- stein geometry-provide non-asymptotic convergence guarantees. In particular, the algorithms by Diao et al. (2023): Lambert et al. (2022) optimize over the full-rank Gaussian family, while that of Jiang et al. (2025) optimizes over all mean-field families with bounded second moments. To guarantee 2W2r qA,) )a < e on strongly log log-concave and log-smooth targets, they all report an iteration complexity of O (d -1 log Meanwhile, under the same conditions, Arnese and Lacker

---

(2024); Lavenant and Zanella (2024) analyzed (block) CAVI, and reported an iteration complexity of Od (. Bhattacharya et al. (2025) provides concurrent result on CAVI, but relies on an assumption that departs from log-concavity and smoothness. Finally, Bhatia et al. (2022) analyzes a specialized algorithm optimizing over only the scale of Gaussians, which has gradient query complexity of O(dke- 3), where k is the user-chosen number of rank-1 factors in the scale matrix.
5.2 Conclusion
In this work, we proved that BBVI with mean-field location scale families is able to converge with an iteration complexity with only Ollog d) dimension dependence, as long as the tails of the family are sub-Gaussian. For high-dimensional targets, this suggests substantial speed advantage over BBVI with full-rank families. In practice, the mean-field approximation can be combined with other design elements such as control variates (Boustati et al., 2020; Geffner and Domke, 2018, 2020; Miller et al., 2017; Roeder et al., 2017; Wang et al., 2024) and data-point subsampling (Kucukelbir et al., 2017; Titsias and Lázaro-Gredinn 2014). Our analysis strategy should easily be combined with existing analyses (Kim et al., 2024a,b) for such design elements.
For target distribution with condition number of x: and a target accuracy level e, we now know how to improve the dependence on d and e in the iteration complexity: Using less-expressive families such as mean-field (Theorem 3.2) or structured (Ko et al., 2024) families improves the dependence on d, while applying control variates to gradient estimators (Kim et al., 2024b) improves the depen- dence on e. However, is currently unclear whether the dependence on K tight or improvable. If is tight, it would be worth investigating whether this can be provably improved through algorithmic modifications, for example, via stochastic second-order optimization methods (Byrd et al., 2016; Fan et al., 2015; Liu and Owen, 2021; Meng et al., 2020; Regier et al., 2017).
Another future direction would be develop methods that are able adaptively adjust the compu- tational cost between Oloo d) and (dd) by trading statistical accuracy akin to the method of Bhatia al. (2022). Existing BBVI schemes with "low-rank (-plus-diagona)""" families (Ong et al., 2018; 8; Rezende et al., 2014; Tomezak et al., 2020) result in non-smooth, non-Lipschitz, and non-convex landscape. This not only rules out typical theoretical convergence guarantees but also exhibits unsta- ble and slow convergence in practice (Modi et al., 2025). Furthermore, understanding the statistical side of this trade-off will be an important direction. As of now, our understanding is restricted to either mean-field or full-rank families (Katsevich and Rigollet, 2024; Margossian and Saul, 2023, 2025; Wang and Blei, 2019 2019;,;; Yang et al., 2020; Zhang and Gao, 2020) with little in between except for the work of Bhatia et al. (2022).
10

---

Acknowledgments and Disclosure of Funding
The authors thank Anton Xue for helpful discussions and the reviewers for helpful suggestions.
K. Kim, J. R. Gardner were supported through the NSF award [IS2115644]; Y.-.. Ma was sup- ported by the NSF Award CCF-2112665 (TILOS), the DARPA AIE program, and the CDC-RFA- FT-23-0069; T. Campbell was supported by the NSERC Discovery Grant RGPIN-2025-04208.
References
Abhinav Agrawal, Daniel R Sheldon, and Justin Domke. Advances in black-box VI: Normalizing flows, importance weighting, and optimization. In Advances in Neural Information Processing Systems, volume 33, pages 17358-17369. Curran Associates, Inc., 2020. (page 1)
Ahmet Alacaoglu, Yura Malitsky, and Stephen J. Wright. Towards weaker variance assumptions for stochastic optimization. arXiv Preprint arXiv:2504.09951, 2025. (page 5)
Pierre Alquier and James Ridgway. Concentration of tempered posteriors and of their variational approximations, The Anmals of Statistics, 48(3):1475-1499 1497, 2020. (page 9)
Manuel Arnese and Daniel Lacker. Convergence of coordinate ascent variational inference for log-
concave measures via optimal transport. arXiv Preprint arXiv:2404.00 08792, 2024. (pages 3, 9)
Francis Bach and Eric Moulines.
Non-asymptotic analysis of stochastic approximation algorithms for machine learning. In Advances in Neural Information Processing Systems, volume 24, pages 451-459. Curran Associates, Inc., 2011. (page 5)
Kush Bhatia, Nikki Lijing Kuang, Yi-An Ma. and Yixin Wang. Statistical and computational trade-offs in variational inference: case study inferential model selection. arXiv Preprint arXiv:2207. 11208, 2022. (pages 10)
Anirban Bhattacharya, Debdeep Pati, and Yun Yang. On the convergence of coordinate ascent variational inference. The Annals of Statistics, 53(3):929 962, 2025. (page 10)
Eli Bingham, Jonathan P. Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul Szerlip, Paul Horsfall, and Noah D. Goodman. Pyro: Deep universal probabilistic programming. Journal of Machine Learning Research, 20(28):166, 2019, (pages 1, 3)
David M. Blei, Alp Kucukelbir, and Jon D. McAuliffe. Variational inference: A review for statisti- cians. Journal of the American Statistical Association, 112(518):859-8777 2017. (page 1)
Léon Bottou. On-line learning and stochastic approximations. In On-Line Learning in Neural Networks, pages 9-42. Cambridge University Press, edition, 1999. (page 4)
Léon Bottou, Frank E. Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM Review, 60(2):223-33 -312, 2018. (pages 1,4)
Ayman Boustati, Sattar Vakili, James Hensman, and S. T. John. Amortized variance reduction for doubly stochastic objective. In Proceedings of the Conference on Uncertainty in Artificial Intelligence, volume 124 of PMLR, pages 61-70. JMLR, 2020. (page 10)
Alexander Buchholz, Florian Wenzel, and Stephan Mandt. Quasi-Monte Carlo variational inference. In Proceedings of the International Conference on Machine Learning, volume 80 of PMLR, pages 668-677. JMLR, 2018. (page 9)
R. H. Byrd, S. L. Hansen, Jorge Nocedal, and Y. Singer. A stochastic -ewain method for large-scale optimization. SIAM Journal on Optimization, 26(2): 1031, 2016. (page 10)
Bob Carpenter, Andrew Gelman, Matthew D. Hoffman, Daniel Lee, Ben Goodrich, Michael Betan- court, Marcus Brubaker, Jiqiang Guo, Peter Li, and Allen Riddell. Stan: A probabilistic program- ming language. Journal of Statistical Software, 76(1):1-32 2017. (pages 1,
George Casella and Roger L. Berger. Statistical Inference. Cengage Learning, 2 edition, 2001.
(page 3)
Edward Challis and David Barber. Gaussian Kullback-Leeiblrr approximate inference. Journal of Machine Learning Research, (4(68))2239-22866 2013. (page 9)

---

Ziheng Cheng, Longlin Yu, Tianyu Xie. Shiyue Zhang, and Cheng Zhang. Kernel semi-implicit vari- ational inference. In Proceedings of the International Conference on Machine Learning, volume 235 of PMIR, pages 2248-8269 JMIR, 2024. (page 9)
Sinho Chewi.
Sampling. Unpublished draft, november 3, 2024 edition, 2024. URL httts://chewiiiiee github io/baini./miiiimm pdf. (page:3)
Michael Ziyang Diao, Krishna Balasubramanian, Sinho Chewi, and Adil Salim. forward-backwerr Gaussian variational inference via JKO in the
space. In Proceedings of the International Conference Machine Learning, volume 202 of PMLR, pages 7960-7991. 2023. (pages 3)
Justin Domke. Provable gradient variance guarantees for black-box variational inference. In Ad- vances in Neural Information Processing Systems, volume 32, pages 329-338. Curran Associates, Inc., 2019. (pages 2, 7,9)
Justin Domke. Provable smoothness guarantees for black-box variational inference. In Proceedings of the International Conference on Machine Learning, volume 119 of PMLR, pages 2587-2596.
JMLR, 2020. (pages 3, 4, 9, 25)
Justin Domke, Robert Gower, and Guillaume Garrigos. Provable convergence guarantees for black- box variational inference. In Advances in Neural Information Processing Systems, volume 36, pages 66289-66322 Curran Associates, Inc., 2023. (pages 4, 5, 9, 19)
Kai Fan, Ziteng Wang, Jeff Beck, James Kwok, and Katherine A Heller. Fast second order stochastic backpropagation for variational inference. Advances Neural Information Processing Systems, volume 28, pages 1387-1395, Curran Associates, Inc., 2015. (pages 10)
Tor Erlend Fjelde, Kai Xu, David Widmann, Mohamed Tarek, Cameron Pfiffer, Martin Trapp, Seth D. Axen, Xianda Sun, Markus Hauru, Penelope Yong, Will Tebbutt, Zoubin Ghahramani, and Hong Ge. Turing.jl: A general-purpost probabilistic programming language. ACM Transac- tions on Probabilistic Machine Learning, 1(3): 2025. (pages 3,
Masahiro Fujisawa and Issei Sato. Multilevel Monte Carlo variational inference. Journal of Machine
Learning Research, 22(278):1 44, 2021. (page 9)
Guillaume Garrigos and Robert M. Gower. Handbook of convergence theorems for (stochastic) gradient methods. arXiv Preprint arXiv:2301.1 1235, 2023. (pages 22)
Hong Ge, Kai Xu, and Zoubin Ghahramani. Turing: A language for flexible probabilistic inference. In Proceedings of the International Conference on Machine Learning, volume 84 of PMLR, pages 1692- JMLR, 2018. (pages 3)
Tomas Geffner and Justin Domke. Using large ensembles of control variates for variational inference. In Advances in Neural Information Processing Systems, volume 31, pages 9960-9970. Curran Associates, Inc., 2018. (page 10)
Tomas Geffner and Justin Domke. Approximation based variance reduction for reparameterizatioo gradients. In Advances in Neural Information Processing Systems, volume 33, pages 2397-2407. Curran Associates, Inc., 2020. (page 10)
Behrooz Ghorbani, Hamid Javadi, and Andrea Montanari. An instability in variational inference for topic models. In Proceedings of the International Conference on Machine Learning, volume 97 of PMLR, pages 2221-2231. JMLR, 2019. (page 9)
Ryan Giordano, Tamara Broderick, and Michael I. Jordan. Covariances, robustness, and variational Bayes. Journal of Machine Learning Research, 19(51): 2018. (page 1)
Ryan Giordano, Martin Ingram, and Tamara Broderick. Black box variational inference with a deterministic objective: Faster, more accurate, and even more black box. Journal of Machine Learning Research, 25:1-39, 2024. (page 1)
Paul Glasserman. Gradient Estimation via Perturbation Analysis. Number 116 in The Springer In- ternational Series in Engineering and Computer Science. Springer, New York, NY, 1991. (page 4)
Peter W. Glynn. Likelihood ratio gradient estimation for stochastic systems. Communications of the ACM, (3))))375--84 1990. (page 4)
Eduard Gorbunov, Filip Hanzely, and Peter Richtarik. A unified theory of SGD: Variance reduction, sampling, quantization and coordinate descent. In Proceedings of the International Conference on Artificial Intelligence and Statistics, volume 108 of PMLR, pages 1800690, JMLR, 2020. (page 22)
12

---

Robert Mansel Gower, Nicolas Loizou, Xun Qian, Alibek Sailanbayev, Egor Shulgin, and Peter Richtárik. SGD: General analysis and improved rates. In Proceedings of the International Con- ference on Machine Learning, volume 97 of PMIR, pages 5200-5209. JMLR, 2019. (pages 5, 22)
J. Gumbel. The maxima of the mean largest value and of the range. The Annals of Mathematical Statistics, 25(1):76-844 1954. (page 9)
Benjamin Halpern. Fixed points of nonexpanding maps. Bulletin of the American Mathematical Society, 73(6)9977-91, 1967. (page 5)
Geoffrey E. Hinton and Drew van Camp. Keeping the neural networks simple by minimizing the description length of the weights. In Proceedings of the Annual Conference Computational Learning Theory, pages 5-13. ACM Press, 1993. (pages 1.3)
Y. C. Ho and X. Cao. Perturbation analysis and optimization of queueing networks. Journal of Optimization Theory and Applications, 40(4):559-582, 1983. (page 4)
Matthew Hoffman and Yian Ma. Black-box variational inference as a parametric approximation to Langevin dynamics. In Proceedings of the International Conference on Machine Learning, volume 119 of PMLR, pages 4324-4341. JMLR, 2020. (page 9)
Alexandra Maria Hotti, Lennart Alexander Van der Goten, and Jens Lagergren. Benefits of non- linear scale parameterizations in black box variational inference through smoothness results and gradient variance bounds. In Proceedings of the International Conference on Artificial Intelli- gence and Statistics, volume 238 of PMIR, pages 3538-3546. JMLR, 2024. (page 4)
Yiheng Jiang, Sinho Chewi, and
Pooladian. Algorithms for mean-field variational inference via polyhedral optimization in the Wasserstein space. Foundations of Computational Mathematics, 2025. (page 9)
Norman L. Johnson, Samuel Kotz, and Narayanaswamy Balakrishnan. Continuous univariate distri- butions. volume 1 of Wiley Series in Probability and Mathematical Statistics. Wiley, New York, 2 edition, 1994. (page 26)
Norman L. Johnson, Samuel Kotz, and Narayanaswany Balakrishnan. Continuous univariate distri- butions. volume 2 of Wiley Series in Probubility and Mathematical Statistics. Wiley, New York, 2 edition, 1995. (pages 26, 28, 29)
Michael Jordan, Zoubin Ghahramani, Tommi S. Jaakkola, and Lawrence K. Saul. An introduction to variational methods for graphical models. Machine Leurning, 37(2):183-233, 1999. (pages 1, 2)
Anya Katsevich and Philippe Rigollet. On the approximation accuracy of Gaussian variational inference. The Annals of Statistics, 52(4): 1384)1998,,, 2024. (page 10)
Ahmed Khaled, Othmane Sebbouh, Nicolas Loizou, Robert M. Gower, and Peter Richtárik. Unified analysis of stochastic gradient methods for composite convex and smooth optimization. Journal of Optimization Theory and Applications, 199:499-555 540, 2023. (page 22)
Mohammad Emtiyaz Khan, Reza Babanezhad, Wu Lin, Mark Schmidt, and Masashi Sugiyama. Faster stochastic variational inference using proximal- -raadien methods with general divergence functions. In Proceedings of the Conference Uncertainty in Artificial Intelligence, UAI' 16, pages 319-328, Jersey City, New Jersey, USA, 2016. AUAI Press. (page 9)
Kyurae Kim, Jisu Oh, Kaiwen Wu, Yian Ma, and Jacob R. Gardner. On the convergence of black- box variational inference. In Advances in Neural Information Processing Systems, volume 36, pages 44615-44657. Curran Associates Inc., 2023a. (pages 2. 3. 4, 22)
Kyurae Kim, Kaiwen Wu, Jisu Oh, and Jacob R. Gardner. Practical and matching gradient variance bounds for black-box variational Bayesian inference. In Proceedings of the International Confer- ence on Machine Learning, volume 202 of PMLR, pages 16853-16876. JMLR, 2023b. (pages 4, 7.9)
Kyurae Kim, Joohwan Ko, Yi-An Ma, and Jacob R. Gardner. Demystifying SGD with doubly stochastic gradients. In Proceedings of the International Conference on Machine Learning. vol- ume 235 of PMLR, pages 24210-24247. JMLR, 2024a. (page 10)
Kyurae Kim, Yian Ma, and Jacob R. Gardner. Linear convergence of black-box variational infer- ence: Should we stick the landing? In Proceedings of the International Conference on Artificial Intelligence and Statistics, volume 238 of PMLR, pages 235-243. JMLR, 2024b. (pages 5, 9, 10)
13

---

Diederik P. Kingma and Max Welling, Auto-encoding variational Bayes, In Proceedings of the International Conference on Learning Representations, Banff, AB, Canada, 2014. (pages 2, 4)
Joohwan Ko, Kyurae Kim, Woo Chang Kim, and Jacob R. Gardner. Provably scalable black-box variational inference with structured variational families. In Proceedings of the International Con- ference Machine Learning, volume 235 of PMLR, pages 24896 24931.. JMLR, 2024. (pages 1,
2, 7,9,10 10. 22)
Alp Kucukelbir, Dustin Tran, Rajesh Ranganath, Andrew Gelman, and David M. Blei. Automatic differentiation variational inference. Journal of Machine Learning Research, 18(14):1-44 2017.
(pages 1, 3, 4, 10)
S. Kullback and R. A. Leibler. On information and sufficiency. The Annals of Mathematical Statis- tics, 22(1):29 86, 1951. (pages 2)
Simon Lacoste-Julien, Mark Schmidt, and Francis Bach. A simpler approach to obtaining an O(1(l) convergence rate for the projected stochastic subgradient method. arXiv Preprint arXiv:1212.2002, 2012. (page 5)
Mare Lambert, Sinho Chewi, Francis Bach, Silvère Bonnabel, and Philippe Rigollet. Variational inference via Wasserstein gradient flows. Advances in Neural Information Processing Systems, volume 35, pages 14434-14447 Curran Associates, Inc., 2022. (pages 3,
Hugo Lavenant and Giacomo Zanella. Convergence rate of random scan coordinate ascent vari- ational inference under log-concavity. SIAM Journal on Optimization, 34(4):3750-317 3761, 2024. (pages 3, 10)
Sifan Liu and Art B. Owen. Quasi-Monte Carlo quasi-Newton in variational Bayes. Journal of Machine Learning Research, 22(243):1-23, 2021. (pages 9, 10)
Charles C. Margossian and Lawrence K. Saul. The shrinkage-delinkage trade-off: An analysis of factorized Gaussian approximations for variational inference. In Proceedings of the Conference on Uncertainty in Artificial Intelligence, volume 216 of PMLR, pages 1358-1367. JMLR, 2023.
(page 10)
Charles C. Margossian and Lawrence K. Saul. Variational inference in location-scale families: Exact recovery of the mean and correlation matrix. In Proceedings of the International Conference on Artificial Intelligence and Statistics, volume 258 of PMLR, pages 1466-3474, JMLR, 2025.
(page 10)
Si Yi Meng, Sharan Vaswani, Issam Hadj Laradji, Mark Schmidt, and Simon Lacoste-Julies Fast and furious convergence: Stochastic second order methods under interpolation. In Proceedings of the International Conference on Artificial Intelligence and Statistics, volume 108 of PMLR, pages
1375-1386. JMLR, 2020. (page 10)
Andrew Miller, Nick Foti, Alexander D' Amour, and Ryan P Adams. Reducing reparameterizatioo gradient variance. In Advances in Neural Information Processing Systems, volume 30, pages 7708-3718. Curran Associates, Inc., 2017. (page 10)
Chirag Modi, Diana Cai, and Lawrence K. Saul. Batch, match, and patch: Low-rank approximations for score-based variational inference. Proceedings of the International Conference on Artificial Intelligence and Statistics, volume 258 of PMLR, pages 4510-4518. JMLR, 2025. (page 10)
Shakir Mohamed, Mihacla Rosca, Michael Figurnov, and Andriy Mnih. Monte Carlo gradient estimation in machine learning. Journal of Machine Learning Research, 21(132):1-62, 2020. (page 4)
A. Nemirovski, A. Juditsky, G. Lan, and A. Shapiro. Robust stochastic approximation approach to stochastic programming, SIAM Journal on Optimization, 19(4):1574-169) 2009. (page 4)
Dai Hai Nguyen, Tetsuya Sakurai, and Hiroshi Mamitsuka. Wasserstein gradient flow over varia- tional parameter space for variational inference. In Proceedings of the International Conference on Artificial Intelligence and Statistic, volume 258 of PMLR, pages 1756-1764. JMLR, 2025.
(page 9)
Victor M.-H. Ong, David J. Nott, and Michael S. Smith. Gaussian variational approximation with factor covariance structure. Journal of Computational and Graphical Statistics, 27(3):465-478 2018. (pages 1, 10)
Neal Parikh and Stephen P. Boyd. Proximal Algorithms, volume 1 of Foundations and Trends& in Optimization. Now Publishers, Norwell, MA, 2014. (page 4)
14

---

Anand Patil, David Huard, and Christopher Fonnesbeck. PyMC: Bayesian stochastic modelling in Python. Journal of Statistical Software, 35(4):1-81, 2010. (pages 1.3)
Carsten Peterson and Eric Hartman. Explorations of the mean field theory learning algorithm. Neu- ral Networks, 2(6):475-49)) -494, 1989. (pages 1,3)
Georg Pflug. Optimization of Stochastic Models: The Interface between Simulation and Optimiza- tion. Number v.373 in The Springer International Series in Engineering and Computer Science Ser. Springer, New York, NY, 1996. (page 4)
Rana. Answer to "A bound for the expectation of the maximum independent random variables", 2017. URL https:: //math. stackexchange. com /000//2277221. (page 27)
Rajesh Ranganath, Sean Gerrish, and David Blei. Black box variational inference. In Proceedings of the International Conference on Artificial Intelligence and Statistics, volume 33 of PMLR, pages 814 822. JMLR, 2014. (page 1)
Jeffrey Regier, Michael Jordan, and Jon McAuliffe. Fast black-box variational inference through stochastic trust-region optimization. Advances in Neural Information Processing Systems, vol- ume 30, pages 2399- 2408. Curran Associates, Inc., 2017. (pages 9, 10)
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In Proceedings of the International Conference on Machine Learning, volume 32 of PMLR, pages 2778-128.. JMLR, 2014. (pages 1, 2, 4, 10)
Herbert Robbins and Sutton Monro. A stochastic approximation method. The Annals of Mathemat- ical Statisties, 2203):400-4077 1951. (pages 1,4)
Geoffrey Roeder, Yuhuai Wu, and David K Duvenaud. Sticking the landing: Simple, lower-variance gradient estimators for variational inference. In Advances in Neural Information Processing Sys- tems, volume 30, pages
Curran Associates, Inc., 2017. (page 10)
Reuven Y. Rubinstein. Sensitivity analysis of discrete event systems by the "push out" method. Annals of Operations Research, 39(1):229-250, 1992. (page 4)
Mark Schmidt and Nicolas Le Roux. Fast convergence of stochastic gradient descent under a strong growth condition. arXiv Preprint arXiv: 00..11000037000 2013. (page 5)
Shai Shalev-Shwartz, Yoram Singer, Nathan Srebro, and Andrew Cotter. Pegasos: Primal estimated sub-gradient solver for SVM. Mathematical Programming, 127(1):3-30, 2011. (page 4)
Sebastian U. Stich. Unified optimal analysis of the (stochastic) gradient method. Preprint arXiv:1907.04232. arXiv, 2019. URL http://arxiv, org/ aba/aab/00 04232. (pages 5, 19)
Sobihan Surendran, Antoine Godichon- Bagichon-Bagiioni, and Sylvain Le Corff. Theoretical convergence guarantees for variational autoencoders. In Proceedings of the International Conference on Arti- ficial Intelligence and Statistics, volume 258 of PMLR. JMLR, 2025. (page 9)
Symbol-1. Answer to "Meaningful lower-bound of Va² + b' - when b> 0", 2022. URL
https: //mtth. httpa/////aatsttacceeeeeeee cwwp://hhhhnnecc//////eeeeec/
(page 21)
Michalis Titsias and Miguel Låzaro-Gredilla Doubly stochastic variational Bayes for non-conjugate inference. In Proceedings of the International Conference Machine Learning, volume 32 of PMI.R, pages 1971-1979. JMIR, 2014. (pages 1, 2, 3, 4, 9, 10)
Marcin Tomczak, Siddharth Swaroop, and Richard Turner. Efficient low rank Gaussian variational inference for neural networks. In Advances in Neural Information Processing Systems, volume 33, pages 4610-4622. Curran Associates, Inc., 2020. (pages 1, 10)
Lloyd N. Trefethen and David Bau. Numerical Linear Algebra. Society for Industrial and Applied Mathematics, Philadelphia, 1997. (page 36)
Sharan Vaswani, Francis Bach, and Mark Schmidt. Fast and faster convergence of SGD for over- parameterized models and an accelerated perceptron. In Proceedings of the International Con- ference on Artificial Intelligence and Statistics, volume 89 of PMLR, pages 1195-1204. JMLR, 2019. (page 5)
Martin J. Wainwright. High-Dimenssionll Statisties: A Non-Asymplotic Viewpoint. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, New York, NY, 1st ed edition, 2019. (pages 6, 9)
15

---

Xi Wang, Tomas Geffner, and Justin Domke. Joint control variate for faster black-box variational inference. In Proceedings of the International Conference on Artificial Intelligence and Statistics, volume 238 of PMLR, pages 1639-1647. JMLR, 2024. (page 10)
Yixin Wang and David Blei, Variational bayes under model misspecification. In Advances in Neural Information Processing Systems, volume 32, pages 13357-13367. Curran Associates, Inc., 2019a.
(page 10)
Yixin Wang and David M. Blei. rrequettist of variational Bayes. Journal of the Ameri- can Statistical Association, 1146577:1144-111611 2019b. (page 10)
Ronald J. Williams. Simple statistical gradient-following algorithms for connectionistrriiffrreeee learning. Machine Learning, 8(3):229-256, 1992. (page 4)
David Wingate and Theophane Weber. Automated variational inference in probabilistic program- ming. arXiv Preprint arXiv:1301. 1299. 2013. (page 1)
Ming Xu, Matias Quiroz, Robert Kohn, and Scott A. Sisson. Variance reduction properties of the reparameterizatioo trick. In Proceedings the International Conference on Artificial Intelligence and Statisties, volume 89 of PMIR, pages 271 1-2720. JMLRR 2019. (page 4)
Zuheng Xu and Trevor Campbell. The computational asymptotics of Gaussian variational inference and the Laplace approximation. Statistics and Computing, 32(4), 2022. (page 9)
Yun Yang. Debdeep Pati, and Anirban Bhattacharya, a-variational inference with statistical guaran-
tees. The Annals of Statistics, 48(2):886-905, 2020. (page 10)
Anderson Y. Zhang and Harrison H. Zhou. Thcoretical and computational guarantees of mean field variational inference for community detection. The Annals of Statistics, 48(5):2575-2598, 2020.
(page 9)
Fengshuo Zhang and Chao Gao. Convergence rates of variational posterior distributions. The Annals of Statistics, 48(4):2180-2200) 2020. (page 10)
Lu Zhang, Bob Carpenter, Andrew Gelman, and Aki Vehtari. Pathfinder: Parallel quasi-Newton variational inference. Journal of Machine Learning Research, 23(306):1-49, 2022. (page 1)
16

---

Contents
Introduction
2
Preliminaries
2.1
Problem Setup..
2.2
Variational Family
2.3
Algorithm Setup
2.4
General Analysis of Stochastic Proximal Gradient Descent
3
Main Results
3.1
General Result
3.2
Special Cases with Benign Dimension Dependence
4
Analysis of Gradient Variance
4.1
Overview
4.2
Upper Bound Gradient Variance
8
4.3
Unimprovability.
5
Discussion
5.1
Related Works
5.2
Conclusion
10
Auxiliary Lemmas
18
B Proofs
19
B.1 Proofs of Results in Section 2.
19
B.1.1
Proof of Proposition 2.9
19
B.1.2
Proof of Lemma B.1
22
B.2
Proofs of Results in Section 3
25
B.2.1
Proof of Theorem 3.2
25
B.2.2
Proof of Proposition 3.4
26
B.2.3
Proof of Proposition 3.5
27
B.3
Proofs of Results in Section 4
30
B.3.1
Proof of Lemma 4.1
30
B.3.2
Proof of Lemma B.4
32
B.3.3
Proof of Lemma B.5
33
B.3.4
Proof of Proposition 4.2
35
17

---

A Auxiliary Lemmas
Lemma A.I. Suppose Assumption 2.2 holds. Then T4 iBUr 2
Proof. By Jensen's inequality Eu, 2 (Eu?)². Lastly, Eut 2 (EU?)² 1 by Assumption 2.2.
Lemma A.2. Suppose Assumption 2.2 holds and denote U diag((U (i,......... Ua). Then we have the
following identities: (i) Euu" Lá. (ii) IU ld
Proof. From Assumption 2.2, we know that Euf
1. Then (i) follows from
Euu' 1ig =
Eu.U;
Euf
if i-
i - (o if
Eu
For (ii), we only need to focus on the diagonal since the off-diagonal is already zero.
[EU)
(Ediag(U1,..... ......, (u,,...,,)).,. Ua
Euf 1
ii
Lemma A.3. Suppose Assumption 2.2 holds, Tx is the reparametrization function for location- scale family, and the linear parametrization is used. Then
E| Tx()) - Tx (1)l = y| X1122
Proof.
Denoting À = (m, C*) and X' = (m',
E| TT() Tx [u)
Ell(Cu + m) (C'u - m')l 2
ElI(C C')u+ (m m
EEI(C C')ul² 2/2 2((C CC) m²) El/m m'11
EEI(C C²)ul 2 El m m ml113
(Assumption 2.2) (10)
Lastly,
E|[(C - CT) EC (c-(cc))((c C'uu
(C- c'ju
=tr(C (C-''"""
(cyclic property of trace)
= =tr(C - CT) [(c - I
(Lemma A.2)
-r (C C²) T (C- C*)
||C -
Combining this with Eq. (10) yields the result.
Lemma A4. Suppose is p-strongly convex and Assumption 2.7 holds. Then f is L-Lipschitz smooth, while the constants satisfy the ordering
L S C.
Proof. For all A, X' E A, the unbiasedness of f and Jensen's 's inequality states that
[V)(A) V f(')
Then the j--troong convexity of f and Assumption 2.7 yields the inequality
2 X11 [V)(()) v < E[Vī(A: - - 1 X1I3
from which the statement follows immediately.
18

---

B Proofs
B.1
Proofs of Results Section
B.1.1 Proof of Proposition 2.9
Under the stated assumptions, we first establish convergence bound which bounds BT - A~ 1 after iterations under a given step size schedule. We will invert this convergence bound into complexity guarantee by identifying the conditions on T. l., T, and to that guarantee |E||AT X. < 6 for a given O.
Lemma B.1. Suppose f is j-stroonly convex, h satisfies Assumption 2.5, and f satisfies Assump- tions 2.7 and 2.8. Then, for the global optimum X, = arg minxes F(A), and the step size schedule / in Eq. (5) with any tx 2 O and 2 4C² /2², the contraction coefficient - 1 - pro satisfies
p E (0, 1) and the last iterate of SPGD after T iterations, XT, satisfies
(t. +)²
+ 2 0² (t. + T)²
802 T t
BlJAT 11 |A 1.112p p'
+
(T + +T)²
# (T' TT)
12 (T 7)
Proof. The full proof is deferred to Appendix B.1.2, p. 22.
This is slight generalization of the result by Domke al. (2023, Theorem 7), where the switching time l. was fixed to some l, x C²/n², While the choice of l. x L //²² results in the typical O(1/e) asymptotic complexity, it suffers from a suboptimal polynomial dependence on the initialization error A = a 2.112- Picking an alternative L., which is what we do in the proof, improves the iteration complexity to O(1/c+ + 1/V log A + log( vanilla SGD was demonstrated by Stich (2019), where the dependence on A is also logarithmic. Their step size schedule, however, requires the maximum number of iterations T, which means T must be fixed before running the algorithm. Our step size, on the other hand, does not require T and is therefore an any-time result.
Now, a similar O(1/c) result for
p-strongly convex, h satisfies Assumption 2.5, and VJ satisfies As- sumptions 2.7 and 2.8. Then, for the global optimumm X arg minxea F(A). A & I1 A 2l2. some t., T, and 7/ (explicit in the proof), SPGD with the step size schedule in Eq. (5) guarantees
Proposition 2.9. Suppose
T2 o 0² oC log
C2 A2 o
L log 112
A²
+
E E A. SE.
1 112
Proof. Since f is strongly convex and is convex, F' is also strongly convex. This implies that, by the property of strictly convex functions, F' has unique global optimum, which we denote X.. Then we can invoke Lemma B.I. We will optimize the upper bound over the parameters t., T, 7o. and T that can ensure the c-accuracy guarantee El|AT A. <.
Consider the choice
4L 2
t, = log 1/p
log
1r 2
Io
1.113)
and
Ta
(11)
2C
We will separately analyze the cases of 4. 2 T and I, < T. When 4. 2 T, the second-stage never kicks in. Therefore, apply Eq. (27) with t, = T. Furthermore,
2o00 IA x.113) ) 2T
H
log
log 1/p
is true. An immediate implication is that
log 1/p 2 11
log
1
A. 112 + +12T
log
Aou - 2u113 2 (log1/p)(T) -1)
20002
1A4
1.113 (T-1)
27002
19

---

2 270 d²
11:0 Xo 1 r-1
(12)
14
Considering this fact, applying Eq. (27) with t, T
A 2.113 - P1 1=112p² + 20 o²
(t, = T)
1
- aA 1.112p p" + 11A0 - do 101120
(Eq. (12))
- 2| Ao - Ao 1120 --11
(p< < 1)
The number of required steps for achieving the e-aceuracy requirement follows from
211A0 - A 1120-~1 Se
211au 2.113 s (1/p)=-
21110 do 1111 =
_
log
s <T' - 1)) (1/p)
1
21100 Xv 11 !
log
ST
log 1/p
1
log
211 Ao 1111 :
(log(1/p) 21-1)
p
2C
2 2/10- - 2113-' 1
log
+15T
(1 - p= 7014 = 22/(2(²) (13)
112
For the case t. <T
1
4, =
log
11 - 1.113
(14)
log 1/p
27o02
1
log
M 1100
2-11
log 1/p
2
2y00²
2
1
log
log
1
110
\o 1o1B
This implies
r² W 2000
1
110 2o113
Substituting for this in Lemma B.I.
E| - Ao s 11 - 1=1 (1. +7)
d² (t. +72
802 T-. to
+ +
+
(15)
T2
H
T
12 T
0² (t. +)² 2
02 T - tx
< 400
+
8
/
T
112
T
(t. +7²
0²1 112 T T
S 4%0
+8
4
T
0² 2t, 27
0²1 1
s 4%0 T2
+8
11
127
0 T2
+
'
which is a quadratic function of 1/T with the coefficients
a e 8% o²
b 2 8 0² 122
( (7))
and
20

---

Achieving the e-aceuracy guarantee is equivalent to finding the largest r = 1/T satisfying the inequalities x > 0 and
az² + br c
By the quadratic formula, this is equivalent to finding the largest e satisfying
b + v0
4ac
2a
Therefore, picking any
2a
T2
b 1 + 40c
sufficient to obtain an e-accurate solution. To make the bound more interpretable, after defining 4ac and B = b, we can use the incquality (Symbol-1,
S 3 + BB+a
2
B2
+a
a
Then
2a
2v 62 4ac
V + dac 1
1 +-+2 1
2a
b + V02 + 4ac
40E
Ve"
where we used the incquality va+b S a + Vo. Thus, we have
T F-b-+2 +0-+2
B|AA Ao11 112 Se.
Ve
Substituting 4, and To with the expressions in Eq. (11),
0²1
T
2
(t+² 2 + +2) Ve 870
(16)
1
T 2 8 021 1
1 Ve
4
70 11/2
(t, + +T)
22
T 8 0² 1
4C2
1
+4V2
To
log
2 11:0 -2a11
+1
(Eq. (14))
1 c
12//2
log 1/p
112
Vc
- T
021 1
4/22
2
log
P
2
1100 - 2o
4C2
+
1
(log 1/p 2 - p)
112
2L
/22/2
112
202
11
112
Ve
= 8 22
222
L2
4L
+
log
170
X.
11
+
112 +1
2
L
Ve
Considering this, the sufficient condition for BJ|AT
x,113 <
is now
802
[
C2
16
40 1
T2
log
1oo
- 21
+
(17)
112
112
02
122
[
L
Ve
Combining both cases, that is, Eqs. (13) and (17), we have
T Z max
802 + 40
2C
log =
C
11
11112
d
+
4C
1
2C
log
211X0
+ +1
112
112
112
Σ
Ve"
This implics the stated result.
21

---

B.1.2 Proof of Lemma B.1
The proof closely mirrors the strategy of Garrigos and Gower (2023, Theorem 12.9), which is combination of previous analyses of SPGD (Gorbunov et al., 2020; Khaled et al., 2023) with the analysis of SGD strongly convex objectives with decreasing step size schedule (Gower et al., 2019). The main difference is that Garrigos and Gower utilize a different condition on the gradient variance instead of Assumption 2.7. Specificially, they assume that, for all X, E A, there exists some function of L(u)= supp(u) () + [0, oo) such that, for each u E supp(u), the function f u) A + Rº is Lmu)-smooth with respect to A. This then enables the use of the "convex expected smoothness" (Gorbunov et al., 2020; Khaled et al., 2023) condition, which postulates that, for all X E A, there exists some C < oc such that
V)V - VI( 5 {)),
(18)
where
-x))
(19)
is the Bregman divergence associated with f. Note that Assumption 2.7 and the p-strong convexity of f implies Eq. (18). Therefore. under assumptions, one invoke the results that assume Eq. (18), which was the strategy by some previous analyses of BBVI (Kim et al., 2023a; Ko et al., 2024). Here, we will take a more straightforward approach that uses Assumption 2.7 directly in the convergence proof, but the results are identical to the indirect approach of establishing Eq. (18).
Lemma B.1. Suppose f is u-strongly convex, h satisfies Assumption 2.5, and f satisfies Assump- tions 2.7 and 2.8. Then, for the global optimum X. = arg min F(A), and the step size schedule
/ /2², the contraction coefficient p - 1 - pro satisfies p E (0, 1) and the last iterate of SPGD after T iterations, \T, satisfies
in Eq. (5) with any t- 2 O and T 2 4C²
(1. T) + 270 0² (t. +T2 +
80² T-
BJJAT 1 DA 2.112p- p'
(T T)
# (T+T)
12 (T (T+
Proof. Since f is strongly convex and h is convex, F is also strongly convex. This implies that F has a unique global optimum, which we denote as A.. Furthermore, under the stated assumptions on h, the proximal operator prox,, (-) is non-expansive for any 7 (0, oo) (Garrigos and Gower, 2023, Lemma 8.) and any X, X IR such that
prox,, h (A) prox,1 & (l)1l2 S π|^ - X1I2
(20)
and X, is the fixed-point of the deterministic proximal gradient descent step (Garrigos and Gower, 2023, Lemma 8.18) such that
prox (A. (2. (V.(A) X.
(21)
Using these facts,
|| 1.1 < Ilprox (Ne 22 V u)) -prox (h. (2
(Eq. (21))
< 1i 76 v((;;)) - d+ + 7...())
(Eq. (20))
Expanding the square.
[De41 1.112 < [| 2.11 - - ( - VI - + -V ().)))
Denoting the filtration of the field of the iterates generated up to iteration as Fi,
E [[] - 1.112 [i]
$ || - 1.11 2n (B (B -f(.). - X.) + + [/ - vf(A.11 F.
= 1| 2. 10- VS(A)), X.) + | vf())) Fi.] (22)
where the equality follows from the fact that ff is unbiased conditional on any A, E A.
From the u-strong convexity of J.
(2(())())) v.(Ae), Xx)
22

---

-AA) - =-7//1/>; 2.112 y (Ac:
S -7Hul1:: - 21112 - 2-) de (u-strong convexity of f)
(Fq. (19))
(23)
The gradient variance at A on the other hand, can be compared against the gradient variance at X through the variance transfer strategy as
E
[VI(Ar u) VS(A; U) 1/ Fi
3 [ U) (A.;; U) + VT(L.; - VI(A. 1 Fi]
s 1/ [ (Ae: u) - VI(A.: 41 [i|
- + 22 - V Fi]
(Young's incquality)
< 2 - 11112 1/- + 27/00
(Assumptions 2.7 and 2.8)
= f² 2
(1(N)-((A)) (04)--(A0) ((A.)
Ae - ->a)) + 2720²
(u-stronn convexity of f)
C² 47 ())) + 2710²
(24)
Applying Eqs. (23) and (24) to Eq. (22),
E B[ 1.112 F.] V || do 112 7e (ullle + (21 +27 2 2 (2 + oz)
L² -1) 27200
- - X. + 2
27t
Taking expectation over all randomness, we obtain our general partial contraction bound
L E[D,(Ar,AA)] + 27 (25)
BlA A. 2 < (1 Bll 1. + 2% 27
Due to the form of the step size schedule, SPGD operates in two different regimes: the first stage with fixed step size it = 7Y (1 E (0,...,,.. L. - 1)) and the second stage with decreasing step size
T++1 < % (f E (t......Th T)). In the first stage, 7/ To S 222- Then the Bregman divergence term Eq. (25) is negative such that
Bl - 2.13 (1- (1-TA) A. + +27202
(26)
Unrolling the recursion yields
1-11
Ar,
2.11 < (1 - 201)" ( |AO - 2.11 - 27002
Σ
(1
t=
g²
$ (1 7011)=1111 11o 2.112
20
(geometric series sum formula)
o²
< 0' 11o A- 113 270
(27)
1
From Lemma A.4, we deduce that 7014 μ²² /22C² < 1/2, which implies p E (0, 1).
We now turn to the second stage, where the step size starts decreasing. Notice that, for any > 4 p/, Eq. (5) satisfies
1 2 +)
12
1 2
Tt.
11 (1 (1+++1)
J 44L
<
2
Therefore, 7/ S a for all l 2 0. Again, the Bregman term in Eq. (25) is negative such that _
E E| - A. s a/)) A.1 + 22?00
23

---

Subtituting Yo with the choice in Eq. (5), we obtain
1 2( 2)
BlIA 1-12 +2 0² (2(t + T) +1)
BI - 2.133 -
(1+ 1)²
122 (1+++1)*
(t 7)
Ella 1. +2 0² (2(1 (((+7)++1 +1)
(t (++ 1)²
#2 (++ 1)¹
Multiplying (1 + T + 1)2 to both sides,
0² (2(1+T)+11)
(L (++ do ( +) 2.113
-2 -2
422
(t (++++1)
Let us choose the Lyapunov function Vn _ ( 1)² 11 do 11 of the Lyapunov,
Then the discrete derivative
0² (2(t + T) +1)
o²
V4+1 V/
S 2-
42a (t (++
4²
shows that the energy is increasing only by constant. By integrating the Lyapunov over the time interval t t.....,T-A
T
Vr- Vi, I 8 02 (T t.) 1²²
Vr S Vi, 8 0²
(T tx)
112
0²
(T (+++)BBIAA7 - 2111 - (t, + r) R
x,112
+ 8 IT t.,)
(t. + T)² (T + r)²
0² T tx
B| 2.11 <
Elldc. 21.1
122 (T + TT)
Substuting 1ee. 2x12 1/ with the error in Eq. (27),
2
(1. +
02 T 1
E - 2.133 < ( 2.113 + 270
+ 8
(T +
2²2 (T 7)
(t, + T)
tt +
802 - t
= 1A - 111
(T + T)
3
+ 2no
(T +
+
112 (T +)) 2
(28)
which is our stated result.
24

---

B.2 Proofs of Results in Section 3
B.2.1 Proof of Theorem 3.2
Theorem 3.2. Suppose the following hold:
1. is p-sttrongy convex and satisfies Assumption 3.1 and μ S mmin (H) S max maxx L.
2. h satisfies Assumption 2.5.
3. Q is a mean-field location-scale family, where Assumption 2.2 holds.
4. ff is the reparametrization gradient.
Denote the global optimum X (m,, = arg minXea F()), the irreducible gradient noise as o & ||m, =112 + ||C. l/y- and the stationary point of as f arg min,eza 8(z). Then, for some
l, T, and To (explicit in the proof), SPGD with the step size schedule in Eq. (5) guarantees
T2 T20 g(d. H. p.
o
H. ()):))((0) (o,c-' + to, lo. (1 Ax13))-113))
EJAA X, 11,sc
where
g(d, H, 6, v)) 2 2(1 +() (1)
(1/2) + + max U
j=1.....
Proof. The proof consists of establishing the sufficient conditions of Proposition 2.9 follows:
(i) E is p-strongly convex
f is p-strongly convex.
iii Assumption 3.1
Assumptions 2.7 and 2.8.
Under the linear parametrization, (i) established by Domke (2020, Thm. 9). It remains to establish (ii). Therefore, the proof focuses on analyzing the variance of the gradient estimator V Since Assumption 3.1 holds, Lemma 4.1 states that, for all x, X' E R DA, the inequality
max ut ) X E (): VJ u) 12 < (1) + 482 ( - j.....
holds. Since A C. Rd idd under the linear parametrization, this implies we satisfy Assumption 2.7 with
C = = 2 H 48 (1/2 4/2+ + max Uy) _=,....
(29)
Furthermore, For the specific choice of X, = (m.; C.) = arg minXex F(2) and = (z, Oaxa) (which is not part of A), we have the equality
B|| [IIV)(A,; u) (); - ()V))(
4)
This means Lemma 4.1 also implies Assumption 2.8 with the constant
13 dll
<² El| + ||C.
C².:
(30)
We are now able to invoke Proposition 2.9. Substituting C and σ² in Eq. (28) with the expressions above, we obtain the condition
T max 80?- + 4o,
log
- |Ao
x1133
+ +4+
112 2og
211A0 - 1.133 +1
L
Ve
Finally, substituting for Eq. (29) yields our stated result.
25

---

B.2.2 Proof of Proposition 3,4
The result follows from well-known bound on the expected maximum of sub-exponential random variables. We state the proof for completeness.
Lemma B.2. Let Xi...., Xa be i.i.d. random variables. Suppose there exists some t 0 such that their pooment-generauiee function (MGF) satisfies Mx, (t) < co. Then
1 (log Mx, (t) log d)
E
max
Proof.
B
max X i=.....d
log exp
IF
max
X;
:=l....dd
S log B Eexp
max
Xi
(Jensen's 's inequality)
i=....,
= log E max
exp (tX:)
i=l.r.d
I log E Σ exp (tx;)
i=1
log
M& (t)
(Definition of MGFs)
ldM(,M2 (t))
(X....,, Xy are i..dd)
Dividing both sides by f yields the statement.
Applying Lemma B.2 to uf yields the result.
Proposition 3.4. Suppose there exists some such that the MGF of uf satisfies Moç(1) < co. Then
E max U s (1/t) (log Moy + log d)
iI....d For example, if a standard Gaussian, then g(d,
8(1m1)//²) (82
log
Proof. The first part of the statement is re-statement of Lemma B.2. For the special case of (0, 1). know that uf xi (Johnson et al., 1995, Eq. 29.1), which the x2 distribution with degree freedom. The MGF of x is given as
Myr (t) (1 2t) 2//-
(Johnson al., 1995, Eq. 29.6)
for E (0, 1/2). Then we can invoke Lemma B.2, which suggests
1 - log (1 - 21) + log
max U < min i=.....d tE(0.1/2)
Any fixed choice of / E (0, is valid upper bound. Picking 1= { !)) 2 $ yields
E max Ui 4 + log i=I....
(31)
Furthermore, the kurtosis of the standard Gaussian is 14 Plugging F4 and Eq. (31) into g in Theorem 3.2 yields the statement.
(Johnson al., 1994, Eq. 13.11).
26

---

B.2.3 Proof of Proposition 3.5
The result follows from the following moment-based bound on the expected maximum of random variables, which is a non-asymptotic refinement of the proof by Rana (2017).
Lemma B.3. Let .i...., Xa be i...d. non-negative random variables where, for k 2 2, their kth moment is finite. That is, Ex, = Th < xo. Then
X S d' d' (k/(k- 1)) 10 11/k 1/k
E max
78
Proof. For any > have
/k
000 X; 2 dt +
1 edlli
max X;
=
P
max
max x: 2t dt i=l.....
i= i=l....d
i=.....
+
dF [x; 2 1] dt
(i..d. and P|J S 1)
0
Jedl/k
//
4 lcct-1
k )
k-1
P(x: 2 4 di
edi/k
000
I
//
c+
ktt
P xx; 2 t dt
(ed <)
kck-1
eddi/N
< //
c+ k
kl kP x 2 4 dl
(Decreased lower limit of integral).
Now, from the definition of moments, we know that
/ kt k- P 2 dt
ktt I Ix>e dlppz dt ktth Ix>e dt dP(p;
(Fubini's Theorem)
kth dt dP(a;]
x dP(c:]
Tk
Therefore,
E max I......d
Xs S d/6k
c+
1 Irk
The bound is minimized when setting
k
1/k
Pk
Then
E max
X, < d//
A: 1
rk
1/k
+
1
m
k --1 1
k. Tk
1(k-1)/R
.....d
k
k
//
1
Tk
1/k
+
k
1
k 1-1 1
rk
1/60
///
1 +
1
k
1
1/k
1
k
(k-
1)/k
d//k
T
k k-1
27

---

If the kth moment of U, is finite, this then immediately implies a polynomial O /b) bound on g.
Proposition 3.5. Suppose, for k 2 2, the kth moment of Uī is finite as 12k = Eu. < cc. Then
E max up I v2 d 1/k i=I....d
r2k
For example, if q is a Student-t with v > 4 degrees freedom and unit variance, then g(d, U,,
8(1I11//22) (6²) / 16 v2 v vd
Proof. The first part of the statement directly follows from Lemma B.3, where we simplified (f/=-1)(k-1)/k In particular, for k 2 2, (4/*-1)(6-1)// is monotonically decreasing. Since an order & 2 2 moment exists by the assumption on the degrees of freedom, (+/K-1) (k- 1)/k V
Let's turn to the second part of the statement. We will denote Student-I distribution with u-degrees of freedom as lg. Since l, does not have unit variance (Johnson al., 1995, Eq. 72), we have to set the sampling process from &p to be
U: d V -ー 2 Vs where V3 ~ tv
i.d.
Ug ~~
Now, it is known that v? d W; ~ FDist 1/2) (Johnson et al., 1995, $28.7), where FD 1/2) is Fisher's F-distribution with (V1, 12) degrees of freedom. The kthh raw moment of FDist(V1, 1/2), denoted as my Ew,", exists up to 2k 1/2 V and is given
1/2 T( K) 2222/2 (2 k) ((v2/2)
mg
(Johnson et al., 1995, Eq. 27.43)
V1
/2
This means that we can invoke Lemma B.3 as
2
V V 2) 2
E max Uf = im I...d
E
max W; I....d
< v2
d
1/k
V
with any & < v/2
can the fact that the gamma function satisfies the recursion r(2 + 1) =T(2). which implies r(a/2 + k) = [(a/2) k-1 (a/2 + i) for any a > 0. Therefore,
For mi 1/k
1/k
I [aa2++ 1/k
1)
((a/2)
2
k-1
$
+1)
+1
(AM-GM inequality)
k 1=0
2
a k(k 1)
+
(gcometric series sum formula)
2 k 2
a
Applying this bound to a = 12 = v and a = V1- 1 respectively,
T(1/2+ k) T(v/2+ (w/2++1) 1/4
) T(v/2)
1/k
m&
I(1/2)
k v
2
2
3v
(k < v/2)
4 4
3
4
Also, choosing = 1w/2 - 17, we have dth < 2/((-2)) This yields
2
2
3²
1
E
max
U? <
V2
d
=
<
3
(32)
1=1 Il...
V
4
2 2
28

---

Lastly, the kurtosis of u; = (2 (2-22// 2-22/// V, follows as (Johnson et al., 1995, Eq. 28.5)
T4 V
V (u-2)
("
12- 2 4
3v²
(v 2)
3
Ew?
1
3
1
(v - 2) 4)
2² - 4)
Plugging the bound in Eq. (32) and the value of r4 into g in Theorem 3.2 yields the statement.
29

---

B.3 Proofs of Results in Section 4
B.3.1 Proof Lemma 4.1
Under the assumption that V 3 LIa and twice differentiability, it is well known that 4 LIa = f is L-smooth. We will prove supporting result analogous to this under Assumption 3.1, which will allow us bound the relative growth of Vt.
Lemma B.4. Suppose e: IRd + R satisfies Assumption 3.1. Then, for any W paxd satisfying
||W]|2 < 00,
[((e((=) VE(2))I2 & IWH(2 22)1 + W 1221 2/112
Proof. The full proof is deferred to Appendix B.3.2, p. 32.
Using this, we can now simplify the V terms in Eq. (6). Applying Lemma B.4 to Vis with W la and Young's inequality.
Vie 2 B (Tx (u) Tx (U))ll2 + 8 ITA (u) Tx (u)ll2)
(Lemma B.4)
(Tx (u) Tx (u)) + 25 BI (u) Tv (4)l1
(Young's inequality)
$ { (u) Tx (0)l + 28t BIITA (u) Tx (u)
(Operator norm)
-211 + 5 Ta (u)- - Tx (u)l = =1111 + 0 - -X1
(Lemma A.3)
(33)
Similarly, applying Lemma B.4 to Vicale with W and Young's inequality,
Vicale - (T) (u) Tx (u))ll2 + 011 Tx (u) Tx (u)ll2)² (Lemma B.4)
< 2E 22]|UUF (H) Tx (U)) +20 +20 E UI (T)) (u) (Young's inequality).
V
(34)
Veunsa corresponds to the constant component of the Hessian v whereas V const corresponds to the non-constant residual. Denote the location and scale parameters of X and X
== (m, C)
and X'= 2'mmm', C²).
For Veenst, we can use the following lemma:
Lemma B.5. Suppose Tx is the reparameterization operator of a mean-field location-fanily and Assumption 2.2 holds. Then, for any matrix H E dd and any X, X E R Ddd
(UH(TX(U) - (x) H X122.
See the full proof in Appendix B.3.3, 33.
The remaining part of the proof closely resembles the proof sketch of Lemma 4.1. For convenience, we first restate Lemma 4.1 and then proceed to the full proof.
Lemma 4.1. Suppose Assumptions 2.2 and 3.1 hold, Q mean-field location-family, and f is the reparametrization gradient. Then, for any X, E Rd x Ddd
E - 1 {2(1 + +
46 (122+4++B
U 12 XII
max j=I....d
Proof. Recall Eq. (34). The proof consists of bounding the two terms Veona and V const. First, for Voonst under Assumption 3.1,
Veenst S 14 |1 XI12
(Lemma B.5)
(35)
remains to bound Vnon-const, which our main challenge.
Denote m m m' and Ō A C C" such that
Tx (u) (u) = (Cu + m) - (C' m')
= (C - C')u (m m')
30

---

Cu+ m.
Then
V
Vion-const
E| U (u) - (
= E| U Cu+
< B|| U Cul + 2
(Young's 's inequality)
E
max
uf
Σ 12 +2m²)
d Σ m
2E
max u j=
CR. + 2E
max uf j=
We will focus on the first term. Denoting i. = arg max;=1....d Uf the coordinate of maximum magnitude, we can decompose the expectation by the contribution of the event i, and i, i. That is,
- Euf Ei = Ei=
C E 14 =
VITAS
V V
The expectation of the event = follows as
Vmux
EE [U, (í. i}}
E.. [ [E[ (i. = +H]
(U) 11 i.)
= I. [u/]] [P[. = s]
10 [ d
< E
Σ u
d
max S
j=......d
1
(dra) 14
(Assumption 2.2)
d
(36)
On the other hand, for the event i, i
V Vmon-mux
E U ut 11 (i. + i}}
E
max u u₂ 1 (i. + 17 jt
R
max ui u?]
(1 1))
s#1
< E
max ui jfز
E u
(u, Il Us for all + j)
max
03 [[
(U....., Ua are i.i.d))
j=ر,..d-1
max un
(Assumption 2.2)
Therefore, we finally obtain
Vnon-const 22 > E max Uf + r4 C/ +E max ui j=I.....d
j=
< 2 E max uf + r1) (|| + |
Up S max 03)
max
j=b...d
j=1 j=l....d-1 j=.....
= 2 E max uî + r4 /|
X113
(37)
j=l....d
Combining Eqs. (6), (33) to (35) and (37) yields the statement.
31

---

B.3.2 Proof of Lemma B.4
Lemma B.4. Suppose f: IRd + R satisfies Assumption 3.1. Then, for any W E Raxd satisfying
[[ < oc,
[W(V[(=) Ve(2))
& (WH(2 (IVV()))V[()))l>>>>>> 211122 + 8 W 11212 2/112-
Proof. From twice differentiability of & (Assumption 3.1) and the fundamental theorem of calculus, we know that
[W(((2)) (l(22)) 1002::2m10 w Jw
/.' Vel(tz (1 t)=')(z =') dt
Denoting _ ta + (1 - t)& for clarity.
[W((V((2 V2(22))12z
1// (() =²) dt
$ [[)) 222
(Jensen's inequality)
[I(((2(())) H + H) (z z')112 dt
$ { [IW =1)12- || 11 (((()) - 111l11= z'112) dt
(Triangle inequality)
(IWH(2 =1)112 + M|l ||2| - 2' 11 dt
(Assumption 3.1)
[IWH(2 =1)1 + 1111W/1211 =221/2-
32

---

B.3.3 Proof of Lemma B.5
Lemma B.5. Suppose Tx is the reparameterization operator of mean-field location-family and Assumption 2.2 holds. Then, for any matrix E IR Radd and any X, E R DAA
(U - (x)( < < X112
Proof. For clarity, let us denote C _ C C" and m - m' such that
Tx(u) (x(u) (Cu+ m) (C' u m') =(C (C)- (m mm)
-Cu+m.
Then
[UH((T(4)) (IU m)11
=EIJUČČUIÉ +2 RUHm, HCU: + EE| U
Vica
Vinss
Vie and Venoss are straightforward. Under Assumption 2.2, it immediately follows that
Vio = Um
.m H HU
m T H Hm
(Lemma A.2)
= ||H2
On the other hand,
Venuss -UUHM, UHČU)
m -THHH H (TUFHCCC))
The expectation follows as
[BUHHCU Buf Σ H HuyO
j=1
Hu, + Σ
Hoy HU
j$
0
(Assumption 2.2)
Thus, the cross term Veross vanishes.
Vicale requires careful elementwise inspection in order to apply Assumption 2.2. That is,
Vicate
2
u
Hig HipC
1=1
E
Hig Cuu, + Σ Cijus
+ (HHHCOss
E
u
H.C.P.P. + 2aH
Σ
EHoOsouup
(expand quadratic)
jifa
sife
HACREU 2H C.C.EEEEE
Σ
+ EEuE
(distribute uf²
jzti
33

---

T4 T4ll 05 +
Σ
(Assumption 2.2)
IM
raH + Σ
Hu, EH,CCyEU,HaCAAEEU kdd
(expand quadratic)
T4
T4H Σ H2
(Assumption 2.2)
H (r4 1) Σ H.C.
= ||CÍ| + (r4 - diag 11
Combining everything,
(T) 2)1i Vioc 2Veros Visale
Hm + |HC (ra - 1)]ldiag(HH)II
(38)
From the property of the Frobenius norm, for any matrix A Iddxd, we can decompose
||||
ΣΣ Σ A
A
|diag(A)]] [oo
i=1 iデj
where (A{ a function that zeroes-out the diagonal of A. Then from Eq. (38),
[[(((U)) - 2)
| ||off (HC) + ||diag HC) 1i + (ra- - 1) (HC) ||
|| [Hf ralldiag (
S +
ra/ldiag|( (HC)
(Lemma A.1)
r4 (I) Hm + ||C
S 1 [1
(operator norm)
= rrll A 1-A--11 X1I2.
which is the stated result.
34

---

B.3.4 Proof of Proposition 4.2
For any n,, L E o0) such that js S L, our goal is to obtain a matrix-valued function H worst _d - 8Ão satisfying
ula
Worsst
IA
that, under the choice H = Worsts maximizes the quantity
U ! Hzz 2)dw
(39)
where 2 E {=
V((z) = 0} is any stationary point of f, z -
(U)), and z & wz + (1 w
Given the norm constraint, the worst-case example that maximizes Eq. (39) will be the matrix-valued function that approximately results in
U Hzz 2)dw
LII E
x
for any realization of on Rd. For this, we will establish the relations
|u] (z - =)dw
(I E
LIUI E11
(40)
The first equality in Eq. (40) follows from identifying the conditions where (zz is independent of the value of w. For the specific choice of
Oa
C= diag(6, ..., 6),
any 8 0
(z"") ) is independent of w if it only depends on the quantities
10
arg max |a L...d
and
2
(41)
|| 1
That is, with some abuse of notation, HZz H I.). Lemma B.6. Suppose m = 2 = Og, and for any ó > O, C' = diag(ā, 6). If (( is function of only is and zd then H (z is constant with respect to [0. 1].
Proof. It suffices to show that, under the stated conditions, the values of i, and 2n are invariant to w. For 2 this trivially follows from the assumption that
Wz + (1 w)
wz
2
00
wz
+ w
I 12
2
For L., we use the fact that the diagonal matrix C' is isotropic
arg max _ arg max w Cailup = I....
arg max we [ul i=l....d
arg max |
i=l....d
i=.....d
From H(z²)= = H("",, (2",,,), i.), the integral in Eq. (39) can be solved as
U J. H(z")(z- 2) 2)dw
(UH(Zz)) 2)1
remains to construct H in a way that depends only on 2n and i, such that
(UH 2) = L 2||
Recalling the spectral constraints, this is equivalent to, for all Rd H solving the equation
H(i)z L 2 % subject to nl, - H(z < LI.
(42)
35

---

Notice the equivalence
2
H( 11zz=L12 2
(zz
Lei.
Thus, 2nd and f. contain all the information we need. The following matrix-valued function almost solves Eq. (42):
H Worsst (2) = ald + B + 2 (0,22+22e) +2e
where 2
(43)
2
This function is reminiscent of Householder reflector (Trefethen and Bau, 1997, Eq. 10.4), where some modifications were made to make satisfy the eigenvalue constraint. From the fact that both e,_ and 2 have unit norm, it is apparent that this matrix satisfies Assumption 3.1 with H = ala and 3. Furthermore, by setting the constants as
L - 1 L+ and B
(44)
2
the triangle inequality asserts that the eigenvalue constraint pId H worst Llg is satisfied almost surely.
Given the specific form of worstt restate the proposition for convenience and then proceed to the proof.
now ready to formally prove Proposition 4.2. Let us first
Proposition 4.2. Suppose Assumption 2.2 holds and Q is a mean-field location-scale family. Then, for any t > 0, d O, u, E (0, +o0) satisfying S L, there exists matrix-valued function H (2) Rd + 8 satisfying pId = H z LIg almost surely and a set of parameters X = (m, C) E
Rd x D such that
B USO (Z-
maN U
}oll,s)(
max uf -1 :1....._
where clt, x0) > 0 is a constant only dependent on and sp.
Proof. Recall H Worrst in Eq. (43). By inspection, we know that H Hworat (zm) only depends on the quantities i. and z" Then Lemma 8.6 states that Hworst (z constant function. There- fore,
E U [
Hworst (2" - 2)dw = E WH (2- -2)1
(Lemma B.6)
u aT4 + B
2 [0022 +20,2) )²
(Eq. (43)) (45)
This can decomposed as
3
1
B
U
aT.
(e,
T
2
3 Uei (2Tz) + 2
B
dUz+
U2
(e Z)
2
3 Uei ||/||2 + UŽzZ
3
E
aUz
2
(2 |1/||
aUz +
UUi + (222)
Uz
2
12 |
8
3
2
Uei +
a
Zi,
Uz
2
2
132
22.1 [1]12((TT
||]
Ue;
112 +
a
Zi,
Uzl13 +
a
4
36

---

8 Uf 1
+ E
a
B Zi
||
+B
a + B 22, ||z|2 (P)
(46)
1
1/2
Here, the first term /2/4 2 is the worst-case behavior we expect from solving Eq. (42). The remaining terms V1 and V2 are the error caused by inexactly solving Eq. (42). suffices show that 12//4/, | dominates lower bounds on V1 and 1/2 asymptotically in L and d.
V 2 0 trivially holds and can immediately be lower-bounded. V22 on the other hand, is not neces- sarily non-negative. Therefore, we will use the bound 1/2 2 -|EV//-
3
EV S B
a
E 2
2
U2zl
2
8
S 3
a +
E
Z|
U lz.1
2
S B a
+
3
(EI
1/2
Eu z7,
1/2
(Cauchy-Schwarz)
(47)
2
V/3
For V3, we can use an argument similar to Eq. (36) where we distribute the influence of the maximum coordinate over the d coordinates.
1/3 = Bui = Ef,
d
Σ
B U C2 1(.
C uf E[ iH
(U,_ 1L i.)
C uf] P(i. $
[ut || ā (Bt EE|ZI
d
(48)
The last equality follows by applying Lemma A.2 the identity E| EUT T Cu. By applying Eq. (48) into Eq. (47), we can now notice that V/2 decreases by a factor of Eut /d.
E
EV2 2 -8
a
Eu?
|
E
2
d
8
2 +
E
||
(Assumption 2.2)
d
It is clear that 1/2 vanishes as d + cc. Applying the lower bound V2 into Eqs. (45) and (46), we have
E
H Hworst Hworst(2")) (z =
82
E E
E
[ [z|2]] - 3 a +
[
4
d
g2
g2
E E(u
a3 +
Bu,
||z]|
4
2
d
37

---

LL - 2)
L n?
(L 2=
E 1
+
Euf, |
(Eq. (44))
d
(L 2
L 1²a
L2 + n²
T8[ BR[u
Buf |1211
(Young's inequality)
4
E max U br
(L p)
L
Euf. ||Cull2
(49)
16
remains solve the expectation.
Let us decompose the events where the ith coordinate attains the maximum (4, = i) or not (h. + i) done in Lemma 4.1.
EU 1 = E
Σ
C
li=1
Σ
c
E Uu
Σ
C 17
We are left with the expectation over the event i, * i. For the upper bound in Lemma 4.1, the expectation was solved by noticing that uf and Uí. can be made independent after upper bounding the indicator. For a lower bound, however, breaking up the expectation for u² and už more involved.
E [ uf. 11.= = uf max oy 1i.=d f
}
uf max uy
u²
max uf j#
(50)
jti
By introducing a free variable > O, we can break up the indicator
{ uf < max u/ 121
2 11 U < max UH < max U?
j41
SF
2 1l u² < t, max us > >t} }
1 (uf < t,
max us > >1]
(51)
jti
This then allows the expectation to break up between terms depending on u² and max Uj which is the independence that we were after. That is, applying Eq. (51) to Eq. (50),
E E of 1==] 2 E uf max uf u(u < }11 3f1
max uf ff1
+t)
u3 > $1]
Elu uf < B
max u3 1
max
I.....
j= bill...
=Efuf <t}] (E)
max
ug
-E
u² 1l
uf =t))
j=
I.....
( pl > sjas)
up - -t).
max 1.
Notice that the function (1, + 1 (P dependent on & and the base distribution ip.
> 8) ds strictly positive as long as L 0 and only
38

---

We now obtain our final result by combining the results into Eq. (49). With explicit constants,
EL' max U,
E %/'' Hoorst (2 2)dw
w1=
(L
L
i= I.....
2
4
2 d
U22
P[' > s]ds
E
max
up
-1) I
Substituting c(l, x)) 1o [P[f > ds into this yields the stated result.
39

---

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: Yes, obtain O(log 1) iteration complexity result for BBVI on strong log-concave and log-smooth target distributions. This corresponds to a nearly dimension-indeeendentt iteration complexity.
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 perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings. It is fine to include aspirational goals motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: As stated in Remark 4.3, the main limitation of our work is that the unimprov- ability result Proposition 4.2 does not fully assert that Lemma 4.1 tight for any target function &. It only shows that our specific proof strategy, which uses only spectral bounds on the Hessian of f, is unimprovable. principle, using additional properties of the Hes- sian could result in a tighter bound.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The au- thors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with 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 ap- proach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-tex, system might be used reliably provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
40

---

3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: All assumptions stated in either Sections 2 and 3 or the propositional statements.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theo- rems.
The proof's can either appear in the main paper the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proof's 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 conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [NA]
Justification: The paper does not contain any experiments.
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 make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way accomplish this, but reproducibility 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 model checkpoint, or other means that are appropriate to the research performed.
also be provided via
While NeurlPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to re- produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may tricky some cases, in which case au- thors are welcome to describe the particular way they provide for reproducibility.
41

---

In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [NA]
Justification:
The paper does not contain any experiments.
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/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so No an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for 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. htt/ppblic//guieesete /wt////Suuueiiionnniiieo
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [NA]
Justification: The paper does not contain any experiments.
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 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
Does the paper report error bars suitably and correctly defined other appropri- ate information about the statistical significance of the experiments?
Question:
Answer: [NA]
Justification: The paper does not contain any experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
42

---

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 enror bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should prefer- ably report 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative еπτог rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [NA]
Justification:
The paper does not contain any experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't t make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: pattps://neuripee cc wttpcccccccc/ceecc.cc..ccccccltt httpp///bbl.cccccccceeeeeeeeeeee htt.hhiiiiciceleees.....
Answer: [Yes]
Justification:
Guidelines:
The content of the paper a theoretical study of an inference algorithm and does not involve real data.
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration 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: [No]
Justification: The content of the paper is a theoretical study of an inference algorithm and does not have direct societal consequences.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
43

---

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 spe- cific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that 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 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 mitiga- tion strategies (e-g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper does not involve real data or models.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards allow for controlled use of the model, for example by re- quiring that users adhere to usage guidelines restrictions to access the model or implementing safery 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 challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification:
The paper does not involve real data.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include 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.
44

---

If assets are released, the license, copyright information, and terms of use in the pack- age should be provided. For popular datasets, paperswith wwwoehhehe....ooo condaataesss has curated licenses for some datasets. Their licensing guide can help determine the li- cense of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer: [NA]
Justification: The paper does not involve real data.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset 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. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the pa- per include the full text of instructions given to participants and screenshots, if applicable, well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but the main contri- bution of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection. cura- tion, or other labor should be paid at least the minimum wage the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper deseribe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equiva- lent) may be required for any human subjects research. you obtained IRB approval, you should clearly state this in the paper.
45

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere 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
Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [No]
Justification: Part of the supporting results were obtained after some minor interaction with LLMs. However, all of the proofs were written and proofread by humans. Therefore, LLMs did not play an important, original role nor did they contribute any non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components. wwwpeeew.....oonreeeeeeeeeeeeeee /
Please refer to our LLM policy for what should or should not be described.
46


Reviews:
[{"review_id": "zh3x4qARBl", "reviewer": "Reviewer_S4iz", "summary": "This paper proves that mean field VI converges at an $O(log d)$, where $d$ is the dimension of the parameters, improving on an $O(d)$ rate for full-rank VI. Previously, it was known that the rate could at most be on the order of the square root of d, and faster convergence was speculated but not proven. If the Hessian of the target density is constant, there is no dimension dependence. To obtain this result, there are assumptions about the regularity of the target density (log-concave and log-smooth), the variational family (mean field location-scale family with sub-Gaussian tails), optimization (proximal SGD), and the gradient estimator (reparameterization trick). The assumption of sub-Gaussian tails can be relaxed, leading to, for example in the student-t case, a $O(d^2/(\\nu-2))$ dependence, where $\\nu$ is the degrees of freedom.", "questions": "- Theorem 3.2 states that there exists a switching time and initial step size for which the result holds. Can you comment on the practical significance of having to choose these values? \n\nSmall:\n- Missing citation on line 253\n- Should be $\\lambda’ = (m’, C’)$ in line 290\n- In the equation after line 290, I believe it should be an inequality right after $V_{scale}$", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "This is an exceptionally well-written paper, with a detailed overview of related work and rigorous proofs of all claims. The assumptions are standard and discussed carefully. The contribution is important because it may justify the use of mean field VI for high-dimensional problems. \n\nSection 4, and the proof sketch of lemma 4.1, provides a nice explanation of the challenges in the proof and why past works obtained looser bounds. \n\nI think the lack of experiments is fine for this kind of paper, but it would be nice to understand the practical significance for common problems. \n\nAs the authors remark, it had been speculated that a faster convergence happens in the mean field case, so one weakness is that the result doesn’t seem surprising, but there is value in identifying and proving the rate as this paper does. The proof strategy also seems like a nice contribution.", "was_revised": true, "final_justification": "This is a clearly-written, rigorous theoretical paper, that I think is a clear accept. As some of the reviewers pointed out, the demonstration of the practical impact via experiments would add to the paper, but I also understand the authors response about the challenges in doing so and maintain my original assessment that experiments are not needed for the paper to be accepted.", "rebuttal": "Thank you for your review!\n\n> Theorem 3.2 states that there exists a switching time and initial step size for which the result holds. Can you comment on the practical significance of having to choose these values?\n\nYes. In practice, it is no secret that any application of SGD requires extensive tuning of the step size schedule. Indeed, Proposition 2.11 reflects this fact through the use of the two-step step size schedule in Eq. (4). Thus, in that sense, our theory is qualitatively in agreement with practice. \n\nFurthermore, in previous works [1], the switching time was selected only to depend on the condition number $\\kappa$. This suggests that it is sufficient to adapt the step size schedule (or more precisely, the switching time) according to the target problem only. Proposition 2.11 demonstrates that, in order to further optimize the dependence on the initialization ${\\lVert\\lambda_0 - \\lambda_*\\rVert}_2$, it is necessary to also adapt the switching time according to the properties of the gradient estimator ($\\sigma^2$). That is, our results suggest that whenever the gradient estimator is swapped, the step size schedule should also be adjusted accordingly to obtain faster convergence. We will make sure to mention this fact in the next version.\n\n1. Gower, R. M., Loizou, N., Qian, X., Sailanbayev, A., Shulgin, E., & Richtárik, P. (2019, May). SGD: General analysis and improved rates. In International conference on machine learning (pp. 5200-5209). PMLR."}, {"review_id": "pQZNj8PhW0", "reviewer": "Reviewer_P294", "summary": "In this work the authors investigate the problem of optimizing black box variational inference, in particular regularized problems which use location (a $d$-dimensional vector $m$ designating the \"center\" of the distribution) and scale ( a $d\\times d$-dimensional matrix $C$, which linearly transforms the distribution, may be enforced to be diagonal, low rank plus identity, etc) parameters, which are contained in $\\lambda$. In such problems,  it is theoretically known to that estimate to $\\epsilon$ accuracy is known to take order $d\\kappa^2 \\epsilon^-1$. It has been conjectured that this dependence on $d$ is no so bad. Indeed for the *mean field* setting where the matrix is enforced to be diagonal, this improves to order $\\sqrt{d}\\kappa^2\\epsilon^-1$, yet it is thought that this dependence on $d$ can be further improved.\n\nThis paper presents results showing that estimating $\\lambda$ using a proximal stochastic gradient descent algorithm with correct step sizes can achieve much better dependences on $d$, with $\\log(d)$ dependence in the case of a sub-Gaussian family, and $d^{2/k}$ when $k$ moments exist for the distribution (so $d^{2/(2-\\nu)}$ for Student-t with $\\nu$ degrees of freedom).", "questions": "I don't have any particularly deep questions regarding the paper.\n\n* I wonder if it would be possible to verify the results experimentally?\n* Are there any concrete useable/practical lessons from this paper, other than one need not worry so much about dimensionality when the data is expected to be not heavy tailed.\n* Again the main point keeping me from a strong accept is my lack of certainty regarding whether the impact, originality or usefulness is good enough to warrant the top reviewing score. This may just be because I am not super comfortable with the topic. Anything you could do here to elaborate for a non-expert would be helpful.", "limitations": "I think the limitations are pretty clear from the exposition in the paper that indicate where their results apply.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 2, "strengths_and_weaknesses": "This paper shines light on a very simple yet fundamental technique in a statistical method that is quite popular recently. Understanding how these estimators work in high dimensions is of interest. I found that paper pleasant to read, even though I feel as though I couldn't fully grasp every part, which is probably unavoidable due to the topicality.\n\nI've scanned the proofs and they also look very nice. They seem well organized with clear justifications provided at every step and signposts for the reader to help guide the proof logic.\n\nFor weaknesses, I don't have any significant issues with the work. The paper is purely theoretical result on a phenomenon that was already suspected, on a topic that isn't as mainstream, as some at NeurIPS, with no clear practical impact. Hence I am a bit reluctant to give a strong accept. This line of work decently outside of my field of expertise so its possible that I'm missing something here.\n\nSmall errors:\n\n* line 144: PSGD is never defined, and the sentence \"Instead of using PSGD, one can also use projected SGD\" is a bit confusing since its easy to think that PSGD refers to projected SGD. I think this may actually be an issue in other parts of the paper. Please just make sure the nomenclature is clear and consistent.\n\n* Prop 2.11 (and elsewhere): I don't really like the $T \\ge O(...)$. Maybe $\\Omega$ is better?\n\n* line 193 and 196: \"anytime\" vs \"any-time\"\n\n* line 253: missing citation", "was_revised": true, "final_justification": "I think accept is warrated for the reasons given in my review. I don't think this is quite strong enough for a spotlight paper for the reasons mentioned. A score of 5 is appropriate.", "rebuttal": "Thank you for your review!\n\n> line 144: PSGD is never defined, and the sentence \"Instead of using PSGD, one can also use projected SGD\" is a bit confusing since its easy to think that PSGD refers to projected SGD.\n\nWe apologize for the confusion, this is a typo and should have been SPGD for stochastic proximal gradient descent. \n\n> Prop 2.11 (and elsewhere): I don't really like the $T \\geq O(\\ldots)$. Maybe $T \\geq \\Omega(\\ldots)$ is better?\n\nThank you for the suggestion. Indeed, this is a common point of confusion, but in this case, big-O is correct, whereas big-$\\Omega$ is not quite. This is because $T \\geq O(\\ldots) \\Rightarrow {\\lVert \\lambda_0 - \\lambda_* \\rVert}_2^2 \\leq \\epsilon$ implies that “the smallest $T$ one can choose to guarantee the $\\epsilon$-accuracy condition grows no faster than $\\ldots$.” Therefore, this is an “upper bound” on the sufficient number of iterations one needs to take. On the other hand, if one uses big-$\\Omega$, it would read, “the smallest $T$ one can choose grows faster than $\\ldots$,” which is a statement of necessity rather than sufficiency.\n\n> I wonder if it would be possible to verify the results experimentally?\n\nThank you for the suggestion. We agree that empirical verification of theoretical results is worth doing. Now, the fact that the mean-field approximation results in better scaling has been known empirically for a long time. Therefore, we found that simply showing that the mean-field approximation scales well is not empirically meaningful. Instead, it would be most interesting to verify that the worst-case $O(\\log d)$ dimensional scaling can be observed in practice. However, as noted in Remark 4.3, the theoretical worst-case example that achieves this $O(\\log d)$ scaling may not exist. We tried various informed guesses for $f$ to see if we observe an increasing scaling with respect to $d$. However, all examples either resulted in a constant scaling or even decreasing (!) scaling with respect to $d$. Therefore, we couldn’t design an interesting enough experiment to verify the results in the paper.\n\n> Are there any concrete useable/practical lessons from this paper, other than one need not worry so much about dimensionality when the data is expected to be not heavy tailed.\n\n> Again the main point keeping me from a strong accept is my lack of certainty regarding whether the impact, originality or usefulness is good enough to warrant the top reviewing score. This may just be because I am not super comfortable with the topic. Anything you could do here to elaborate for a non-expert would be helpful.\n\nThis is an important point. Our results say a bit more than one needs to worry about dimensionality. As mentioned in Section 1, we believe VI is the most relevant for obtaining a controllable statistical-computational trade-off. That is, we would like to know how much speed up we can obtain by giving away some amount of statistical accuracy. For this, a quantitative convergence analysis is necessary for variational families of various degrees of expressivity. With this in mind, our quantitative convergence analysis implies that, in the high-dimensional regime, our results suggest moving from a full-rank family to a mean-field family will require running SGD roughly $d$ times longer with a smaller step size. For variational families with an expressivity in between mean-field and full-rank, we can expect to see a dimensional scaling in between $O(\\log d)$ and $O(d)$, where we can use the $O(\\log d)$ result as a baseline for estimating the amount of necessary computation time. (See also the last paragraph of Section 5.2.)"}, {"review_id": "BhaZufcOIt", "reviewer": "Reviewer_FijK", "summary": "In this paper, the authors developed a new theoretical convergence guarantee for black-box variational inference when using mean-field location-scale families. They show that by using stochastic proximal gradient descent, it can achieve a convergence rate that is nearly independent of dimension for target distributions that are log-smooth and strongly log-concave. Moreover, authors have analyzed the variance of the reparameterization gradient and showed that the bound on the gradient variance cannot be improved using only spectral bounds on the Hessian of the target log-density.", "questions": "1. Typo in lines 144 & 185: It looks like the acronym \"PSGD\" is a typo. Should it be SPGD instead?\n2. Should the matrix $C$ in Definition 2.6 be required to be positive definite and diagonal?", "limitations": "All limitations are addressed.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\nThis paper is well-written and clearly structured. The inclusion of proof sketches is helpful for readers. The paper demonstrates that the iteration complexity of mean-field location-scale families, which are commonly used in practice, can be improved to nearly a dimension-independent rate, representing an advancement over previous work. I believe this is a significant contribution to the variational inference community.\n\nWeaknesses:\nThe paper is highly theoretical and lacks numerical experiments or synthetic illustrations. I think including illustrations would help the reader follow the concepts more easily. While the paper focuses on the reparameterization gradient and provides a technical justification for its use, it does not adequately explain why this estimator is prioritized over alternatives like the score function gradient. Providing additional context or justification could strengthen the relevance of the paper's contributions.", "was_revised": true, "final_justification": "This is a technically sound and well-argued paper that addresses an important topic. The authors’ rebuttal addresses my earlier concerns. While the lack of illustrative experiments limits the immediate empirical validation, I understand the challenges involved and find the theoretical contributions compelling. Overall, I remain confident in my original assessment, \"accept\", and believe the paper makes a valuable contribution.", "rebuttal": "Thank you for your review!\n\n> The paper is highly theoretical and lacks numerical experiments or synthetic illustrations. I think including illustrations would help the reader follow the concepts more easily.\n\nThank you for the suggestion. We are open to suggestions for any illustrations or experimental results that would help readers understand the technical content of the paper. Please feel free to suggest anything the reviewer would find useful!\n\n> While the paper focuses on the reparameterization gradient and provides a technical justification for its use, it does not adequately explain why this estimator is prioritized over alternatives like the score function gradient.\n\nAs mentioned in Lines 154-155, we focus on the reparametrization gradient as it is empirically known to perform much better than the score gradient [1,2]. This fact has also been theoretically established in [3] for Gaussian targets. In addition, as mentioned in Lines 76-78, most probabilistic programming frameworks use the reparametrization gradient by default whenever applicable. Therefore, we believe it is fair to say that the reparametrization gradient is the most popular choice in practice. However, we do agree that the score gradient also warrants some attention in terms of a general theoretical analysis. But this would be worth its own paper.\n\n> Providing additional context or justification could strengthen the relevance of the paper's contributions.\n\nThank you for the suggestion. As mentioned in Section 1, a quantitative convergence analysis of mean-field BBVI is practically relevant regarding the statistical-computational trade-off. That is, it is useful to know how long it would take to obtain a full-rank approximation compared to a mean-field approximation. Our paper implies that, in the high-dimensional regime, one should roughly expect a factor of $d$ difference. (See also the reply to Reviewer P294.)\n\n> Typo in lines 144 & 185: It looks like the acronym \"PSGD\" is a typo. Should it be SPGD instead?\n\nYes, thank you for spotting this! It should have been SPGD.\n\n> Should the matrix $C$ in Definition 2.6 be required to be positive definite and diagonal?\n\nThis is somewhat a matter of convention. Mathematically speaking, $C$ only needs to be invertible to result in a well-defined distribution since the resulting covariance $CC^{\\top}$ will be positive definite even if $C$ has negative eigenvalues. However, it is indeed common to define $C$ to have strictly positive eigenvalues, since, in the 1-dimensional, allows to interpret $C$ as a standard deviation. For our purpose, this requirement isn't necessary for defining location-scale families, so we didn't mention it Definition 2.6. Though we should clarify that $C$ at least needs to be invertible. We will fix this in the next version and thank you for pointing this out.\n \n1. 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.\n2. Kucukelbir, A., Tran, D., Ranganath, R., Gelman, A., & Blei, D. M. (2017). Automatic differentiation variational inference. Journal of machine learning research, 18(14), 1-45.\n3. Xu, M., Quiroz, M., Kohn, R., & Sisson, S. A. (2019, April). Variance reduction properties of the reparameterization trick. In The 22nd international conference on artificial intelligence and statistics (pp. 2711-2720). PMLR."}, {"review_id": "apqclVPeQt", "reviewer": "Reviewer_Hyeq", "summary": "This paper analyzes the rate of convergence of a particular optimization algorithm (stochastic proximal gradient descent using reparameterization gradients and a particular step size schedule) applied to mean-field variational inference for problems with strongly log-concave and log-smooth likelihoods using location-scale variational families. The main result is that the rate of convergence only depends weakly on the dimension of the problem: for variational families generated by sub-exponential random variables the dependence is $O(\\log d)$, for variational families with $k$ finite moments it is $O(d^{2/k})$.  The authors then provide some justification that similar proof techniques could not result in a better bound.  The main ingredient in the proof is a careful analysis of the variance of the gradient in this setup, but the authors also use a number of additional clever tricks like a carefully chosen step size schedule.", "questions": "n/a", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n* Overall I thought that the paper was really clear. The exposition and proofs were generally easy to follow (although see below) and well-motivated.\n* The main result is very interesting and suggests that mean-field BBVI should scale to higher dimensional problems more gracefully (at least computationally).\n* The results could be useful in practice (e.g., setting step sizes or choosing a variational family to balance the ability to fit heavier tails with a worse scaling in dimension).\n\nWeaknesses (all minor):\n* There is a lot of emphasis on the scaling being nearly dimension-free, which -- unless I'm missing something --  obscures that some of the assumptions likely do not scale sub-linearly in the dimension.  For example, it would be reasonable to expect the Frobenius norm of the optimal scale matrix, $\\Vert C_* \\Vert_F^2$, to grow linearly in the dimension.  I don't think this is a major issue, but some discussion of how the other parameters in the bound are likely to grow with dimension in practice is warranted.\n* It would be nice if the authors could discuss how relevant (at least in terms of the intuition) the authors think their results are to practical implementations.  For example, not all implementations use proximal SGD, many implementations use parameterizations other than the linear parameterization, the optimal step size and step size schedule depend on generally unknown parameters and so cannot be set like this in practice and so on.  Proofs in these other cases are obviously outside the scope of the paper, but it would be helpful if the authors could discuss which of their assumptions are for technical purposes and which they think are likely actually required of algorithms.\n* I found the proof of Proposition 4.2 hard to follow.  In particular, it was unclear which parts of the proof are just for developing intuition for how they chose $H_{worst}$ and which parts of the proof are actually rigorous.  Relatedly, why is there a $z^w$ in the middle of the display after line 911? (And similarly in the second line of the display after line 921)?\n\nAdditional comments and typos:\n* I think the step size schedule has a nice interpretation -- $t_*$ is chosen so that for the first $t_*$ steps the noise in the gradient is negligible relative to the ``signal''.  As a result, one can use a constant step size as if the gradients were noise free, this also allows one to geometrically converge from the initialization to $\\lambda$ that are close enough to the optimum such that the noise in the gradients begins to dominate.  After this, one is in a regime that is more SGD-like and so a step-size like $1/t$ is required (c.f., Robbins-Monro).\n* I found the phrasing around the formal Assumptions in the main text to be somewhat confusing on my first read.  For example \"Assumptions\" 2.1 and 2.2 read more like definitions, and the actual assumption is that various functions have those properties.  Relatedly, it was unclear throughout when \"Assumptions\" were being stated in the general case vs. being actual assumptions about the problem at hand.  For example, I think Assumption 2.7 should always hold for mean-field VI with location-scale families with linear parameterizations. I understand that it is being used to set up a more general result, but some changes to the wording might make that clearer.\n* It's up to the authors, but Equation (16) appears in, I think, 3 different proofs, and I think it would be cleaner if that were made to be another Lemma in Appendix A instead of needing to refer back to an equation in the middle of a proof.\n* Line 37: “parameterization gradient” —> “reparameterization gradient”\n* Line 95: Shouldn't \"diag\" should be \"vec\"? $z$ is vector valued and $C$ is a matrix, so $Cu$ as defined would also be a $d\\times d$ matrix, I think.\n* Line 109: Latex typo $Q$ --> $\\mathcal{Q}$\n* Line 222: $q_T$ --> $q_{\\lambda_T}$\n* Line 253: missing reference in \"of the result by , states that the highest\"\n* Line 290: second $\\lambda$ should be $\\lambda'$\n* Line 326: \"difficult to verify\" --> \"are difficult to verify\"\n* Line 362: \"should easily be combined\" --> \"should be easy to combine\"\n* Line 669: second $\\lambda$ should be $\\lambda'$\n* Line 735: \"straightly\" --> \"straightforwardly\"\n* Not that it matters, but I believe the first line on p. 21 can be replaced with an equality because it's just algebra\n* Statement of Lemma B.1: assumption should be on $\\widehat{\\nabla f}$ not $\\nabla f$ (appears in both statements of Lemma B.1)\n* \"Equality\" after line 704 should be $\\le$\n* After line 714 is presumably missing a $\\le \\epsilon$\n* I believe the equations after lines 796 and 797 are missing terms that correspond to the function $g$ defined in the Theorem.\n* Second line in the display after line 822: should be $X_i > t$ not $X_i > s$ in the domain of the indicator\n* Norm on right hand side of first line of display after line 891 is missing a subscript 2.\n* In both statements of Proposition 4.2 \"location-family\" --> \"location-scale family\"", "was_revised": true, "final_justification": "Based on my reading, the other reviews, and the authors' rebuttal to all the reviews, I feel that this is an exceptionally clear, well-written, and interesting paper.  As such, I feel that it is a clear \"Accept\".  There are practical limitations to implementing the results in the paper (e.g., the results only apply to very specific algorithmic choices), which makes the paper a bit too purely theoretical in my opinion to be useful to an extremely broad community, preventing me from choosing \"Strong Accept\".  Nevertheless, I enjoyed the paper and think it will make a strong addition to the conference.", "rebuttal": "Thank you for your review!\n\n> There is a lot of emphasis on the scaling being nearly dimension-free, which -- unless I'm missing something -- obscures that some of the assumptions likely do not scale sub-linearly in the dimension. For example, it would be reasonable to expect the Frobenius norm of the optimal scale matrix, ${\\lVert C_* \\rVert}_{\\mathrm{F}}$, to grow linearly in the dimension. I don't think this is a major issue, but some discussion of how the other parameters in the bound are likely to grow with dimension in practice is warranted.\n\nThank you for the suggestion. Indeed, the term dimension-free appears to vary across communities and may not be clear to all readers. In the submission, we tried to avoid the ambiguity by stating “explicit dimension dependence” in both the abstract and the introduction, but not everywhere in the introduction. As such, in the next version, we will spell out “explicit” more often in the introduction and add a mention about implicit dimension dependence in both the introduction and results sections.\n\n> It would be nice if the authors could discuss how relevant (at least in terms of the intuition) the authors think their results are to practical implementations. For example, not all implementations use proximal SGD, many implementations use parameterizations other than the linear parameterization, the optimal step size and step size schedule depend on generally unknown parameters and so cannot be set like this in practice and so on. Proofs in these other cases are obviously outside the scope of the paper, but it would be helpful if the authors could discuss which of their assumptions are for technical purposes and which they think are likely actually required of algorithms.\n\nIndeed, the design space is large, and the algorithm design we chose in the paper does not exactly match practical implementations. In terms of parametrizations, we briefly mention Lines 124-128 that alternative “non-linear parametrizations” are commonly used. However, as mentioned, these usually rule out obtaining convergence guarantees [1] without setting a compact domain [2]. For optimization algorithms outside of proximal SGD, projected SGD is an obvious candidate. However, as mentioned in Lines 145-148, this requires setting an arbitrary compact domain, and both the analysis and convergence guarantees are analogous up to a factor of x2. (See [3].) However, we agree that a more holistic investigation and discussion over all these design choices would be useful for the community. For this, we plan on writing a paper performing an empirical comparison that would hopefully act as a guide to the various pros and cons of these design elements.\n\nFor setting the step size schedule and initial step size, we agree that a discussion about practical considerations would be useful. We will mention the fact that the act of parameter tuning can be seen as identifying the step size (schedules) that depend on the unknown problem constants.\n\n> I found the proof of Proposition 4.2 hard to follow. In particular, it was unclear which parts of the proof are just for developing intuition for how they chose $H_{\\mathrm{worst}}$ and which parts of the proof are actually rigorous. Relatedly, why is there a $z^{w}$ in the middle of the display after line 911? (And similarly in the second line of the display after line 921)?\n\nThank you for pointing this out! We agree that the use of $z^w$ here is extremely confusing. The original intent was to convey that the result of the integral $\\int^1_0 H_{\\mathrm{worst}}(z^w) (z - \\bar{z}) \\mathrm{d}w$ is the same as the pointwise evaluation of $H_{\\mathrm{worst}}(z^w) (z - \\bar{z})$ with respect to any $w \\in [0, 1]$, but the notation is inappropriate. We will fix this in the next version. In terms of the organization of the proof, we will add more signpost text to clearly demarcate where the requirements for the worst-case example are stated and the actual proof begins.\n\n> I found the phrasing around the formal Assumptions in the main text to be somewhat confusing on my first read. For example \"Assumptions\" 2.1 and 2.2 read more like definitions, and the actual assumption is that various functions have those properties.\n\nThank you for the suggestion. We agree that “Definition” would have been better here, and have in fact changed this shortly after our initial submission.\n\n> For example, I think Assumption 2.7 should always hold for mean-field VI with location-scale families with linear parameterizations. I understand that it is being used to set up a more general result, but some changes to the wording might make that clearer.\n\nThank you for the suggestion. We agree that some text to explain the relevance to VI would have been better here. We will improve this in the next version.\n\nAnd thank you for spotting the various typos in the proofs! We will make sure to fix these in the next version.\n\n1. Kim, K., Oh, J., Wu, K., Ma, Y., & Gardner, J. (2023). On the convergence of black-box variational inference. Advances in Neural Information Processing Systems, 36, 44615-44657.\n2. Hotti, A. M., Van der Goten, L. A., & Lagergren, J. (2024, April). Benefits of non-linear scale parameterizations in black box variational inference through smoothness results and gradient variance bounds. In International Conference on Artificial Intelligence and Statistics (pp. 3538-3546). PMLR.\n3. Domke, J., Gower, R., & Garrigos, G. (2023). Provable convergence guarantees for black-box variational inference. Advances in neural information processing systems, 36, 66289-66327."}]

Author rebuttal:

Meta-review: We thank the authors for their submission. 

This work sheds theoretical light on the convergence of a commonly used variational approximation — location scale mean field families. Notably, the authors show that under certain conditions, convergence only weakly depends on the dimension of the parameters.  Reviewers agree that this is a well-written paper with a theoretically sound analysis of a popular algorithm, justifying its use and making more precise under what conditions/likelihoods may enjoy these theoretical guarantees.
Decision: Accept (poster)