On the Convergence of Black-Box Variational Inference
Authors: Kyurae Kim
Jisu Oh
Kaiwen Wu
Yian Ma
Jacob R. Gardner
Venue: NeurIPS 2023 poster (2023)
Keywords: black-box variational inference
stochastic gradient descent
Bayesian inference
variational inference
probabilistic machine learning
Bayesian machine learning
variational Bayes
TL;DR: This paper provides the first convergence proof for black-box variational inference as-is, from which practical insights are obtained

Abstract: We provide the first convergence guarantee for black-box variational inference (BBVI) with the reparameterization gradient.
  While preliminary investigations worked on simplified versions of BBVI (e.g., bounded domain, bounded support, only optimizing for the scale, and such), our setup does not need any such algorithmic modifications.
  Our results hold for log-smooth posterior densities with and without strong log-concavity and the location-scale variational family.
  Notably, our analysis reveals that certain algorithm design choices commonly employed in practice, such as nonlinear parameterizations of the scale matrix, can result in suboptimal convergence rates.
  Fortunately, running BBVI with proximal stochastic gradient descent fixes these limitations and thus achieves the strongest known convergence guarantees.
  We evaluate this theoretical insight by comparing proximal SGD against other standard implementations of BBVI on large-scale Bayesian inference problems.

On the Convergence of Black-Box Variational Inference
Kyurae Kim University of Pennsylvania kyrkim@seas upenn. edu
Jisu Oh North Carolina State University
Kaiwen Wu University of Pennsylvania kaivenwudseas upenn edu
joh266ncsu. edu
Yi-An Ma
Jacob Gardner University of Pennsylvania jacobrgāseas. upenn edu
University of California, San Diego y yianmaQucsd.eduu
Abstract
We provide the first convergence guarantee for black-box variational inference (BBVI) with the reparameterization gradient. While preliminary investigations worked on simplified versions BBVI (e.g., bounded domain, bounded support, only optimizing for the scale, and such), our setup does not need any such algorith- mic modifications. Our results hold for log-smooth posterior densities with and without strong log-concavity and the location-scale variational family. Notably, our analysis reveals that certain algorithm design choices commonly employed in practice, such nonlinear parameterizations of the scale matrix, can result in suboptimal convergence rates. Fortunately, running BBVI with proximal stochas- tic gradient descent fixes these limitations and thus achieves the strongest known convergence guarantees. We evaluate this theoretical insight by comparing proxi- mal SGD against other standard implementations of BBVI on large-scale Bayesian inference problems.
Introduction
Despite the practical success of black-box variational inference (BBVI; Kucukelbir et al., 2017; Ran- ganath el al., 2014; Titsias & Greeila,,,,,,a,, 2014), also known as stochastic gradient variational Bayes and Monte Carlo variational inference, whether it converges under appropriate assumptions on the target problem have been an open problem for a decade. While our understanding of BBVI has been advancing (Bhatia et al., 2022; Challis & Barber, 2013; Domke, 2019, 2020; Hoffmann & Ma, 2020), a full convergence guarantee that extends to the practical implementations as used in probabilistic programming languages (PPL. such as Stan (Carpenter al., 2017), Turing (Ge et al., 2018), Tensorflow Probability (Dillon et al., 2017), Pyro (Bingham et al., 2019), and PyMC (Patil al., 2010) has yet to demonstrated.
Due to our lack understanding a consensus on how we should implement our BBVI algorithms has yet to be achieved. For example, when the variational family is chosen to be the location-scall family, the "scale" matrix can be parameterized linearly or nonlinearly, and both parameterizations are used by default in popular software packages. (See Table in Kim et al. 2023.) Surprisingly, as will show, seemingly innocuous design choices like these can substantially impact the convergence of BBVI. This is critical BBVI has been shown to be less robust (e.g., sensitive initial points, stepsizes, and such) than competing inference methods such as Markov chain Monte Carlo (MCMC). (See Dhaka et al., 2020; Domke, 2020; Welandawe evaluation of BBVI algorithms has been relying on expensive empirical evaluations (Agrawal et al., 2020; Dhaka et al,, 2021; Giordano el al., 2018; Yao et al., 2018).
al., 2022; Yao et al., 2018.) Instead, the
37th Conference on Ncural Information Processing Systems (NeurIPS 2023).

---

To rigorously analyze the design of BBVI algorithms, we establish the first convergence guarantee for the implementations precisely as used in practice. We provide results for BBVI with the repa- rameterization gradient (RP; Kingma & Welling, 2014; Titsias & Lázaro-Grediila, 2014) and the docation-scale variational family, arguably the most widely used combination in practice. Our re- sults apply to log-smooth posteriors, which is a routine assumption for analyzing the convergence of stochastic optimization (Garrigos & Gower, 2023) and sampling algorithms (Dwivedi et al., 2019, $2.3). The key is to show that evidence lower bound (ELBO; Jordan et al., 1999) satisfies regular- ity conditions required by convergence proof's of stochastic gradient descent (SGD; Bottou, 1999; Nemirovski et al., 2009; Robbins& & Monro, 1951), the workhorse underlying BBVI.
Our analysis reveals that nonlinear scale matrix parameterizations used in practice are suboptimal: they provably break strong convexity and sometimes even convexity. Even if the posterior is strongly log-concave, the ELBO is not strongly convex anymore. This contrasts with linear parameterizations, which guarantee the ELBO to be strongly convex if the posterior is strongly log-concave (Domke, 2020). Under linear parameterizations, however, the ELBO is no longer smooth, making optimiza- tion challenging. Because of this, Domke (2020) proposed to use proximal SGD, which Agrawal & Domke (2021, Appendix A) report to have better performance than vanilla SGD with nonlinear parameterizations. Indeed, we show that BBVI with proximal SGD achieves the fastes/ known con- verges rates of SGD, unlike vanilla BBVI Thus, we provide a concrete reason for employing proxi- mal SGD. We evaluate this insight on large-scale Bayesian inference problems by implementing an Adam-like (Kingma & Ba, 2015) variant of proximal SGD proposed by Yun et al. (2021).
Concurrently to this work, convergence guarantees on BBVI with the RP and the sticking-the-landing estimator (STL; Roeder er al., 2017) under the linear parameterization were published by Domke et al. (2023). To achieve this, they show that a quadratic bound on the gradient variance is sufficient to guarantee the convergence of projected and proximal SGD. In contrast, we focus on analyzing the ELBO under nonlinear parameterizations and connect it to existing analysis strategies. A more in-depth comparison of the two works is provided in Appendix E.
1 Convergence Guarantee for BBVI: Theorem 3 establishes a convergence guarantee for BBVI with assumptions matching the implementations used in practice. That is, without algorithmic simplifications and unrealistic assumptions such as bounded domain or bounded support.
Theorem 2 shows that, for location variational families, nonlinear scale parameterizations prevent the ELBO from being strongly-convex even when the target posterior is strongly log-concave.
Optimality of Linear Parameterizations:
Convergence Guarantee for Proximal BBVI: Theorem guarantees that, proximal SGD is used, BBVI on fl-strongly log-concave posteriors can obtain solution c-close to the global optimum with O (1/E) iterations.
Evaluation of Proximal BBVI in Practice: In Section 5, we evaluate the utility of proximal SGD on large-scale Bayesian inference problems.
2 Background
Notation Random variables are denoted in serif (e.g., x, x), vectors in bold (e.g., x. x), and matrices are in bold capitals (e.g. A). For vector Rd. we denote the inner product as x x and xx, the f--norm as ||x|| Vx For a matrix A, ||AI|, vir(A denotes the Frobenius norm. $ is the set of positive definite matrices. For some function f. Dif denotes the ith coordinate of Vf, and C* (XC, 3) is the set of k-time differentiable continuous functions mapping from X' to y.
2.1 Black-Box Variational Inference
Variational inference (VI, Blei et al., 2017; Jordan el al., 1999; Zhang et al., 2019) aims to minimize the exclusive (or backward/reverse) Kullback-Leibler (KL) divergence as:
minimize DRL (92.7) $ FFz-qq log 7 (z) H (q2). ZEA
where DKL (q2, is the KL divergence,
H the differential entropy, is the (target) posterior distribution, and aa is the variational distribution,
While alternative approaches to VI (Dieng et al., 2017; Hernandez-Lobato el al., 2016; Kim et al., 2022; Naesseth el al., 2020) exist, so far, exclusive KL. minimization has been the most successful. We thus use "exclusive KL minimization" as synonym for VI, following convention.
2

---

Equivalently, one minimizes the negative evidence lower bound (ELBO, Jordan et al., 1999) F:
minimize F (2) - Ex-q1 - log x) H (q2),
AEΛ
where log p (z, x) is the joint likelihood, which proportional the posterior as 7 (z) or & (Z.
[xx z) (z), where (x z) is the likelihood and (z) is the prior.
2.2 Variational Family
In this work, focus the following variational family, ( is equivalence in distribution.)
Definition 1 (Reparameterized Family). Let p be some d-variate distribution. Then, qa that can be equivalently represented as
Definition (Location-Scale Reparameteri-
zation Function). [ + p defined as (u) _ Cu+m
~~ qa
z z= Jg (u): u~p,
with A containing the parameters for forming the location m E me and scale C = C() E Raxd is called the location-scale reparameteri- zation function.
is said to be part of reparameterized family generated by the base distribution o and the reparameterization function JΛ
The location-scale family enables detailed theoretical analysis, as demonstrated by (Domke, 2019, 2020; Fujisawa & Sato, 2021; Kim et al., 2023), and includes the most widely used variational families such as the Student-t, elliptical, and Gaussian families (Titsias & Lázzrr-Grredilla 2014).
Handling Constrained Support For common choices of the base distribution Φ. the support of qu is the whole Rd Therefore, special treatment is needed when the support of π is constrained. Kucukelbir er al. (2017) proposed to handle this by applying diffeomorphic transformation denoted with x often called bjectors (Dillon et al., 2017; Fjelde et al., 2020; Leger, 2023), to qa such that
- (z);
~~ qua
q2.
such that the support of q4.2 matches that of 7. For example, when the support of is R4, one can choose 1 exp. This approach, known as automatic differentiation VI (ADVT), is now standard in most modern PPLs.
Why focus on posteriors with unconstrained supports? When bijectors are used, the entropy of 92. H (q2). needs be adjusted by the Jacobian of &p (Kucukelbir er al., 2017). Jæ-1. However, applying the transformation to π instead of 92 is mathematically equivalent and more convenient. In fact, bijectors can be automatically incorporated into our notation by implicitly setting
p (x Z) ( | ()-1(z) (z)) and (z) (x (z)) -y-1 (z)).
such that ((5) (xx 5)) ((), where से is the constrained posterior that we are actually interested in. Therefore, our setup in Section 2.1, where the domain of is taken to be the unconstrained R4, already encompasses constrained posteriors through ADVI.
Lastly, we impose light assumptions on the base distribution Φ, which are already satisfied by most variational families used in practice. (i.i.d:: independently and identically distributed.)
Assumption (Base Distribution). @ is d-variate distribution such that & and u (u,, ........) Ug) with i..dd. components. Furthermore, is (i) symmetric and standardized such that Iuu; 0, Fuf 1, Fuf 0, and (ii) has finite kurtosis Fuf kp co.
The assumptions on the variational family we will use throughout this work are collectively summa- rized in the following assumption:
Assumption 2. The variational family is the location-scale family formed by Definitions I and 2 with the base distribution & satisfying Assumption I.
2.3 Scale Parameterizations
For the "scale" matrix C (2) in the location-scale family, any parameterization that results in positive-definite covariance CCT E s is valid. However, for the ELBO to ever be convex, the entropy HH (q2) must be convex, which requires the mapping 2 + CCT to be convex. To ensure this, we restrict C to (lower) triangular matrices with strictly positive eigenvalues, essentially, Cholesky factors. This leaves two of the most common parameterizations:
3

---

Definition 3 (Mean-Field Family.).
Definition 4 (Full-Rank Cholesky Family).
C= D (s)
C= Dg (s) + L,
where the d elements of s forms the diagonal and A E A such that
where the d elements of's forms the diagonal, L is d- by-d strictly lower triangular, and À E A such that
A = {m. s) m E R R E 8).
A = {m,, L)
m E R s E 8, vec (L) E R
Here, S is discussed in the next paragraph, D& (s) E Rdxd is diagonal matrix such that D& (s) _ diag (& (s)) = diag ( (s.)..... (s.)), and & is a function we call a diagonal conditioner.
Linear v.s. Nonlinear Parameterizations When the diagonal conditioner is a linear function (xx) = x, we say that the covariance parameterization is linear. In this case, to ensure that C is Cholesky factor, the domain of s is set as 8 R4. On the other hand, by choosing nonlinear conditioner Φ R R4, we can make the domain of to be the unconstrained 8 Rd. Because of this, nonlinear conditioners such as the softplus (x) - log (1 + exp (x)) (Dugas et al., 2000) are frequently used practice, especially for mean-field. (See Table by Kim el al., 2023).
2.4 Problem Structure of Black-Box Variational Inference
Exclusive KL minimization VI is fundamentally a composite (regularized) optimization problem
(()
(ELBO)
where f(2) _ Eg-qz e (z) is the energy term, {(z) - log p x)) is the negative joint log- higlliiklihhod, and h (2) _ -HI (qp.x) the entropic regularizer. From here, BBVI introduces more structure.
An illustration of the taxonomy shown in Figure 1. In particular, BBVI has an infinite sum structure (IS). That is, it cannot be represented as a sum of finite subcomponents as in ERM. Furthermore,
IS
CP
BBVI
F (A) = Eu f (A; h (A)
ICP
VI
= Eu e (Ti (u)) + h (2),
(SP RP)
where (2: _ e (Jz (u))
RP
Theoretical Challenges The structure of BBVI has multiple challenges that have hindered its theoretical anal- ysis: (i) the stochasticity of the Jacobian of and (ii) The infinite sum structure.
ERM
FS
CP
Composite
IS
Infinite Sum
For Item (i), we can see that in
RP
Reparameterized
FS
Finite Sum
0J2 (u)
3J (u) di g ud;
(T (u))
(Tz
ve (I2 (u))
ERM
Empirical Risk Minimization
32
where g (2; u) (Ve o J2) (u), both the Jacobian of Jj and the gradient of the log-likelihood, g, depend on the randomness u. Effectively decoupling the two is a major challenge to analyzing the properties of the ELBO and its gradient estimators (Domke, 2019, 2020).
Figure 1: Taxonomy of variational in- ference. Within BBVI, this work only considers the reparameterization gradi- ent (BBVI n RP, shown in dark red). This leaves out BBVI with the score gradient (BBVI \ RP, shown in light red). The set VI FS includes sparse variational Gaussian processes (Titsias, 2009), while the remaining VI (FS U IS U RP) includes coordinate as- cent VI (Blei et al., 2017).
For Item (ii), the problem is that recent analyses of SGD (Garrigos & Gower, 2023; Gower et al., 2019; Nguyen er al., 2018; Vaswani al., 2019) have increas- ingly been relying on the assumption that f (2; u) is smooth for all such that
Vf (A; - v (d': S LIIA 2'll
for some I. < co. This is sensible the support of u is bounded, which is true for the ERM setting but not for the class of infinite sum (IS) problems. Previous works cireumvented this issue by assuming (i) that the support of u is bounded (Fujisawa & Sato, 2021) which implicitly changes the variational family, or (ii) that the gradient Vf is bounded by a constant (Buchholz et al., 2018: Liu & Owen, 2021) 1) which contradicts strong convexity (Nguyen et al., 2018).

---

The Evidence Lower Bound Under Nonlinear Scale Parameterizations
Under the linear parameterization (f(x) X), the properties of the ELBO, such smoothness and convexity, have been previously analyzed by Challis & Barber (2013); Domke (2020); Titsias & Låzaro-Gredilla (2014). We generalize these results to nonlinear conditioners.
3.1 Technical Assumptions
Letg (): be the ith coordinate ofg (2: and recall that denote the ith element of Establishing convexity and smoothness of the ELBO under nonlinear parameterizations depends on pair of necessary and sufficient assumptions. To establish smoothness:
Assumption The gradient of e under reparameterization, g. satisfies
|Eg (ă; u) u,4" (s;)) - L,g
for every coordinate = 1,.. d. any A E A, and some 0 < L, < co.
Here, &" is the second derivative of $. The next one is required to establish convexity:
Assumption 4. The gradient of l under reparameterization, g, satisfies
Egi (A; ui 20
for every coordinate = 1,...
Intuitively, these assumption control how much ve and T2 rotate the randomness u. (Notice that the assumptions are closely related the matrix Cov (g (A; u), the covariance between g and u.) However, the peculiar aspect of these assumptions is that they are not implied by the convexity and smoothness of é. Especially, Assumption strongly depends on the internals of Vé.
3.2 Smoothness of the Entropy
Under the linear paramoterization, Domke (2020) has previously shown that the entropic regularizer term h is not smooth. This fact immediately implies the ELBO is not smooth. However, certain nonlinear conditioners do result a smooth regularizer.
Lemma 1. If the diagonal conditioner & is -oglos-smooth, then the entropic regularizer (2) is Ly-smooth.
Proof. See the full proof in page 24.
Example The following diagonal conditioners result smooth entropic regularizer:
Let & (x) = softplus (x). Then, Lx -smooth with L, ~ 0.167096.
2. Let $ (x) = exp (x). Then, h is -mmoth for arbitrarily small Ln-
This might initially suggest that diagonal conditioners are a promising way of making the ELBO globally smooth. Unfortunately, the properties the energy, f, change unfavorably.
3.3 Smoothness of the Energy
Inapplicability of Existing Proof Strategy Previously, Domke (2020, Theorem 1) have proven that the energy is smooth when $ is linear. The key step was to use Bessel's 's inequality based on the observation that the partial derivatives of the reparameterization function form unit bases in expectation. That is,
T.A (u) aT
E
= linj.
04,
02,
where را= is an indicator function that is 1 only when and otherwise.
Unfortunately, when & is nonlinear, the partial derivatives 052(/// for = 1,.... P no longer form unit bases: while they still orthogonal in expectation, the lengths change nonlinearly depending on 2. This leaves Bessel's inequality inapplicable. To circumvent this challenge, we establish a replacement for Bessel's inequality:
Lemma Let be n X symmetric random matrix, where is bounded ||H||2 almost surely. Also, let be an m X random matrix such that [[EJ Jll2 < co. Then,
|E [EJTHJ]|, [ Jll2-
Proof. See the full proof in page 24.
Remark 1. By assuming that the joint log-likelihood é is smooth and twice-difierentiable, we re- trieve Theorem 1 of Domke (2020) by setting J to the Jacobian of and H to the Hessian of l under reparameterizatton.

---

Remark 2. While our reparameterization function's 's partial derivatives still form orthogonal bases, they need not be; unlike Bessel's inequality, Lemma 2 does not require this. This implies that Lemma 2 is strategy more general than Bessel's inequality.
Equipped with Lemma 2, we present our main result on smoothness:
Theorem 1. Let l be Lm-stooth and twice differentiable. Then, the following results hold:
(i) If & is linear, the energy f is Lm-oohhth.
(ii) If Φ is 1-Lipschitz, the energy e is (Lg + Lg)-smooth if and only if Assumption 3 holds.
Proof. See the full proof in page 27.
Combined with Lemma I, this directly implies that the overall ELBO is smooth.
Corollary I (Smoothness of the ELBO). Let e be L-gsmooth and Assumption 3 hold. Further- more, let the diagonal conditioner be I-Lipschitz continuous, and ty-log-smooth. Then, the ELBO is (Le + L33 + Lgg -smooth.
The increase of the smoothness constant implies that we need to use a smaller stepsize to guarantee convergence when using a nonlinear $ Furthermore, even on simple L-smooth examples Assump- tion 3 may not hold:
Example 2. Let e (z) = (1/2 z Az and the diagonal conditioner be & (x) = softplus (x). Then,
(i) if is dense and the variational family is the mean-field family or (ii) A is diagonal and the variational family is the Cholesky family,
Assumption holds with Lg 0.26034 (max Au).
(iii) A is dense but the Cholesky family is used, Assumption does not hold.
Proof. See the full proof in page 29.
ELBO with ((x) = softplus(x)
ELBO with ((x) = x
Example 2 illustrates that establishing the smooth- ness of the energy becomes non-trivial under nonlin- parameterizations. Even when smoothness does hold, the increased smoothness constant implies that BBVI will be less robust to initialization and step- sizes. Furthermore, in the next section, we will show much more grave problem: nonlinear parameteri- zations may affect the convergence rate.
Lower-Boune Bower-Bounding Quadratic
Tangent Line S -
I
10
5
0
3.4 Convexity of the Energy
5
5
The convexity of the ELBO under linear parameteri- zations has first been established by Titsias & Låzaro- Gredilla (2014, Proposition 1) and Domke (2020, Theorem 9). In particular, Domke (2020) show that, when & is linear, if é is u-strongly convex, the energy also j-strongly convex. However, when using a nonlinear & with co-domain of R4, which is the whole point of using a nonlinear conditioner, strong convexity of l never transfers to f.
S1
Figure Optimization landseape result- ing from different & on a strongly-convex e. is the counter-example of Proposition Item (ii). ((x) = x preserves strong convex- ity as shown by the lower-bounding quadratic (red dotted line ....). & softplus violates the first-order condition of convexity (black dotted line ......
Theorem 2. Let é be u-strongly convex. Then, we have the following:
(i) !f & is linear, the energy f is u-strongly convex.
(ii) If & is convex, the energy convex if and only if Assumptiio 4 holds.
(iii) If is such that & C (R, R4), the energy f is not strongly convex.
Proof. See the full proof in page 33.
The following proposition provides some conditions for Assumption 4 to hold not hold.
Proposition 1. We have the following:
(i) If é is convex. then for the mean-field family, Assumption holds. (ii) For the Cholesky family, there exists a convex é where Assumption 4 does not hold.
Proof. See the full proof in page 31.

---

For any continuous, differentiable nonlinear conditioner that maps only to non-negative reals, the strong convexity of l does lead to a strongly-convex ELBO. This phenomenon visualized in Fig- ure 2. The loss surface becomes flat near the optimal scale parameter. This problem becomes more noticeable as the optimal scale becomes smaller.
Nonlinear conditioners are suboptimal. As the dataset grows, Bayesian posteriors are known to "contract" characterized by the Bernstein-von Mises theorem (van der Vaart, 1998). That is, the posterior variance becomes close to 0. This behavior also applies to misspecified variational posteri- ors as shown by Wang & Blei (2019). Thus, for large datasets, nonlinear conditioners mostly operate in the regime where they are suboptimal (locally less strongly convex). But linear conditioners re- sult in a non-smooth entropy (Domke, 2020). This dilemma originally motivated Domke to consider proximal SGD, which we analyze in Section 4.2.
4 Convergence Analysis of Black-Box Variational Inference
4.1 Black-Box Variational Inference
BBVI with SGD repeats the steps:
1 2,11 2, --------- Y/ (2,) ( where VJ
Eveee ve (JA (um))
(1)
M m=1
with um ~ Φ is the M-sapple reparamcterization gradient estimator and Y, is the stepsize. (Sec Ku- cukelbir et al., 2017 for algorithmic details.)
With our results Section 3 and the results of Khaled & Richtárik (2023); Kim et al. (2023), we obtain convergence guarantee. To apply the result of Kim et al. (2023), which bounds the gradient variance, we require an additional assumption.
Assumption 5. The negative log- lellikikihhodd (z) - llg (x x) is u-quarratically growing for all z E [d such that
- Zuko like (z) like
where Zike the projection of to the set of minimizers of iike and l infzerd like (z).
This assumption is weaker than assuming that the likelihood satisfies the Polyak-Lojjssiwweez in- equality (Karimi et al., 2016).
Theorem 3. Let Assumption 2 hold, the likelihood satisfy Assumption 5, and the assumptions of Corollary 1 hold such that the ELBO F is Lp-smooth with Lp Le + Lg + Ly. Then, the iterates generated by BBVI through Equation (1) and the M M-sample reparameterization gradient include an e-stationary point such that minosisT-1 EIIVF (2,))12 S f.ror any E > 0 if
(F (F FT2² 1 (d, kd)
T20
UME+
for some fixed stepsize Y. where Cd. = d + Kop for the Cholesky family and C Cd, s) 2kp Vd 1 for the mean-field family.
Proof. See the full proof in page 35.
Remark 3. Finding an e-stationary point of the ELBO has an iteration complexity of (dI3xM-1E-* for the Cholesky family and o (V for the mean-field family.
4.2 Black-Box Variational Inference with Proximal SGD
Proximal SGD For a composite objective F f + h, proximal SGD repeats the steps:
[ (VF (2,) (2) A) + h (2) l 211/
44+1 prox hr.h (A. y: Vf arg min AeΛ
(2)
2y:
where prox is known the proximal operator and 71. Vr is stepsize schedule.
In the context of VI, proximal SGD has previously been considered by Altosaar et al. (2018); Diao el al. (2023); Khan el al. (2016, 2015). Their overall focus has been developing alternative algo- rithms by generalizing J|2 - 1 to other metries. contrast, Domke (2020) considered proximal SGD with the regular Euclidean metrie 1a - 2 ||2 for overcoming the non-smoothness of h under

---

linear parameterizations. Here, we prove the convergence of this scheme and show that it retrieves the fastest known convergence rates in stochastic first-order optimization.
Proximal Operator for BBVI In context, h is the entropy of qa in the location-scale family. For this, Domke (2020) show that the the proximal update for SI,..., Sd,
y,o (5)) + 2 5 + 4%, -
For other parameters, the proximal operator is the regular gradient descent update in Equation (1)-
Gradient Variance Bound We first establish a bound on the gradient variance. In ERM, contem- porary strategies do this by exploiting the finite sum structure of the objective (Section 2.4). Here, we establish variance bound for RP estimator that does not rely the finite sum assumption.
Lemma 3 (Convex Expected Smoothness). Let e be Le-smooth and p-strongly convex with the variational family satisfying Assumption with the linear parameterization, Then,
V/ (2:u) - V (2'; u)ll 2 2L C (d, k) By (2,2)
holds, where B (2,2') -2) (2)) 2') is the Bregman divergence, Le/H
is the condition number, (d, + Kp for the Cholesky family, and C (d. 2kp V
for the mean-field family. Proof. See the full proof in page 36.
Furthermore, the gradient variance at the optimum must be bounded:
Lemma 4 (Domke, 2019; Kim al., 2023). Let e be Lg-smooth with the variational family satisfying Assumption 2 and a I-Lipschitz diagonal conditioner &. Then, the gradient variance at the optimum 1 arg min tEA IF (2) is bounded
o² S Cd. .) L3 (lL - fatmmmlip m
1[C* m
M
where 2 is a stationary point of e, m and C are the location and scale formed by 2", the constants are Cdd. d + kg for the Cholesky family and C (d, () 2kg Va 1 for the mean-field family, ky is the kurtosis of q as defined in Assumption 1. Proof. The full-rank case is proven by Domke (2019, Theorem 3), while the mean-field case is basic corollary of the result by Kim et al. (2023, Lemma 2).
The dimensional dependence in the complexity of BBVI is transferred from the variance bound in Lemma 4. Unfortunately, for the Cholesky family, this dimensional dependence in the variance bound is tight (Domke, 2019).
Remark 4.
Main Result With the gradient variance bounds, we now present our complexity result. The proof is identical to Theorem 3.2 by Gower et al. (2019), where they use 2-stage decreasing stepsize schedule: the stepsize initially held constant and then reduced in a 1/t rate.
Theorem 4. Let l be smoshh and y-strongly convex. Then, for any E O, BBVI with proximal SGD in Equation (2), the M-sample reparameterization gradient estimator, variational family satisfying Assumption 2 with the linear parameterization guarantees
El AT 201 I if
for t 5 5T4
802
4T 1l. 2 1112
Y: = 2LxxC(d.p) 2++1
Z max
+
4T₂
for t 4Tx,
cvVE
(1+1)²
where o² is defined in Lemma 4, Tx = [x2C (d. M)
Lellu is the condition number; F (2). C(d, p) kp for the Cholesky family, and
Euler's constant, 2" arg min FF
C(d, Φ)
2k,Vd 2000) + for the mean-field family.
Proof. See the full proof in page 38.
Remark 5. BBVI with proximal SGD on f-stroogly convex and Le-smooth & has a complexity O (x2dM-1<-1) for the Cholesky family and O (x22 /d2 for the mean-field family.
Remark 6. We also provide a similar result with a fixed stepsize in Theorem of Ap- pendix F.3.2. In this case, the complexity is O (x2 loge for the Cholesky family and 0x22VM--e- loge=1) for the mean-field family.

---

Proximal SGD ((x)
SGD {x) =x
SGD ((x) softplus(x)
T for e-accuracy
10a
10a
10
[T
rrmm
mrmm
10
10
10T
10
10
10-2
10
10
10
10-²
10-1
Stepsize (y)
Figure 3: Stepsize versus the number ifterations for vanilla SGD and proximal SGD to achieve DKL (qi(q2.7) πt S E = under different initializations for Gaussian posteriors. The initializations (Ag)) are I, 10 10- from left to right, respectively. The average suboptimality iteration was estimated from 10 independent runs. For each run, the target posterior was a 10-dimensional Gaussian with covariance with condition number 10 and smoothness of L 100.
ProxGen-Adam 6(x)
Adam ((x)
Adam 6(x) softplus(x)
7000
1x100
2x10'
LME-clection
LME--adon
BT-tennis
LR- LR-keggu
ELBO
-8000
-11x10"
-5110'
1x
-0000
-1.2x10*
4x10
-1x10'
-1x10*
-1.3x10*
6x100
2.5k 5k
7.5k 10k
1 5k 10k 15k 20k
10k
20k
30k
10k
20k
30k
Iteration
LR-song
LR-buzz
2x10'
LR-eleetric
-1x10"
AR-ceg
-1x10*
-2x10'
ELBO
1x10"
1x10
-2x10'
-3x10"
-2110*
4x10"
-3x10"
4x"00
-4x10"
-3x10"
-7x10"
-5x10"
10k
20k
30k
25k
50k
10k
20k
5k
10k
15k
20k
Iteration
Figure 4: Comparison of BBVI convergence speed (ELBO v.s. Iteration) of different optimiza- tion algorithms. The error bands are the 80% quantiles estimated from 20 (10 for AR-eeg) indepen- dent replications. The results shown used base stepsize of = 10-3 while the initial point was m. 0, C. I. Details on the setup can be found in the of Section 5.2 and Appendix G.
Experiments
5.1 Synthetic Problem
Setup We first compare proximal SGD against vanilla SGD with linear and nonlinear parameter- izations on a synthetic problem, which is log-smooth, strongly log- log-concave, and the exact solution known. While similar experiment was already conducted by Domke (2020), here we include nonlinear parameterizations, which were not originally considered. We run all algorithms with a fixed stepsize infer multivariate Gaussian with full-rank covariance matrix. The variational approximation is a full-rank Gaussian formed by q = N0, 1) and the Cholesky parameterization.
The results are shown in Figure 3. Proximal SGD clearly the most robust against ini- tialization. Also, SGD with the nonlinear parameterization $(x) = softplus(x) is much slower to converge under all initializations. This confirms that linear parameterizations are indeed superior for both robustness against initializations and convergence speed.
Results
5.2 Realistic Problems
Setup We now evaluate proximal SGD on realistic problems. In practice, Adam (Kingma & Ba, 2015) is observed be robust against stepsize choices (Zhang el al., 2019). The reason why Adam performs well on non-smooth, non-convex problems is still under investigation (Kunstner el al., 2023; Reddi el al., 2023; Zhang el al., 2022). Nonetheless, to compare fairly against Adam, we implement recently proposed variant of proximal SGD called ProxGen (Yun et al., 2021), which

---

includes an Adam-like update rule. The probabilitic models and datasets are fully described in Ap- pendix G. We implement these models and BBVI on top of the Turing (Ge et al., 2018) probabilistic programming framework. Due to the size of these datasets, we implement doubly stochastic subsam- pling (Titsias
laaaal,,,,,,lll, 2014) with batch size of B = 100 (B = 500 for BT-tennis) with M = 10 Monte Carlo samples. For batch subsampling, we implement random-reshufling which is faster than independent subsampling both empirically (Bottou, 2009) and theoretically (Ahn et al., 2020; Haochen & Sra, 2019; Mishchenko et al., 2020; Nagaraj et al., 2019). We also observe that doubly stochastic BBVI benefits from reshuffling, but leave a detailed investigation to future works.
Results Representative results are shown in Figure 4, with additional results in Appendix H. Both ProxGen-Adam and Adam with linear parameterizations converge faster than Adam with nonlin- ear parameterization. Furthermore, for the case of election and buzz, Adam with the nonlinear pa- rameterization converges much slower than the alternatives. When using linear parameterizations, ProxGen- Adam appears to be generally faster than Adam. We note, however, that due to the dif- ference in the update rule between ProxGen-Adam and Adam, proximal operators alone might not fully explain the performance difference. Nevertheless, the results of our experiment do conclusively suggest that linear parameterizations are superior.
6 Discussions
Conclusions In this work, we have proven the convergence of BBVI. Our assumptions encompass implementations that are actually used in practice, and our theoretical analysis revealed limitations in some of the popular design choices (mainly the use of nonlinear conditioners). To resolve this issue, we re-evaluated the utility of proximal SGD both theoretically and practically, where it achieved the strongest theoretical guarantees in stochastic first-order optimization.
Related Works To prove the convergence of BBVI, early works have a-priori "assumed" the reg- ularity of the ELBO and the gradient estimator (Alquier & Ridgway, 2020; Buchholz el al., 2018; Khan et al., 2016, 2015; Liu & Owen, 2021; Regier et al., 2017). Towards a more rigorous under- standing, Domke (2019); Fan al. (2015); Kim al. (2023); Xu el al. (2019) studied the reparam- eterization gradient, Xu & Campbell (2022) studied the asymptotics of the ELBO, Challis & Barber (2013): Domke (2020); Titsias& & Lázaro- Lázaro-Gredilla (2014) established convexity, and Domke (2020) established smoothness. On the other hand, Bhatia et al. (2022); Hoffiman & Ma (2020) established rigorous convergence guarantees by considering simplified variant of BBVI where only the scale is optimized, and Fujisawa& & Sato (2021) assumed that the support of & is bounded almost surely. Meanwhile, under similar assumptions to ours, Diao et al. (2023); Lambert er al. (2022) recently es- tablished convergence guarantees for proximal SGD BBVI with a wasserrassr metric. Their computational properties differ from BBVI as they require Hessian evaluations. Also, understanding BBVI, which is VI with a Euclidean metric, is an important problem due to its practical relevance.
Limitations Our work has multiple limitations: (i) Our results are restricted to the location-scale family, (ii) the reparameterization gradient, and (iil) smooth joint log-likelihoods. However, the location-scale family with the reparameterization gradient is the most widely used combination in practice, and replacing the smoothness assumption is an active area of research in stochastic optimiza- tion. For our results on proximal SGD, we further assume that the joint log- likelihood is j-strongly convex (cquivalently strongly log-concave posteriors). It is unclear how to extend the guarantees to only smooth but non-log-coneave joint liceliikiioods.
Open Problems Although we have proven that the mean-field dimensional family has a dimension dependence of (Va). empirical results suggest room for improvement (Kim et al., 2023). Therefore, we pose the following conjecture:
Conjecture 1. Under mild assumptions, BBVI for the mean-field variational family converges with only logarithmic dimensional dependence or no explicit dimensional dependence at all.
This would put mean-field BBVI in a regime clearly faster than approximate MCMC (Freund et al., 2022). Also, is unknown whether the O (x2) condition number dependence dependence is tight. In fact, for proximal SGD BBVI
Sarres-Wsseesstenn space, Diao el al. (2023) report a dependence of O (x). Lastly, it would be interesting to see whether natural gradient VI (NGVI; Amari, 1998; Khan & Lin, 2017) can achieve similar convergence guarantees. While is empirically known that NGVI often converges faster (Lin et al., 2019), theoretical evidence has yet to follow.
10

---

Acknowledgments and Disclosure of Funding
The authors would like thank Justin Domke for discussions the concurrent results, Javier Bur- roni for pointing out a mistake in the earlier version of this work, and the anonymous reviewers for their constructive comments.
K. Kim and J. R. Gardner were funded by the National Science Foundation Award [IIS-214564]] while Y.-A. Ma was funded by the National Science Foundation Grants INSF-SCALE MoDL- 2134209] and INSF-CCF-2112665 (TILOS)], the U.S. Department Of Energy, Office of Science, and the Facebook Research award.
References
Agrawal, Abhinav, & Domke, Justin. 2021. Amortized Variational Inference for Simple Hierarchical Models. Pages 21388 21399 of: Advances in Neural Information Processing Systems, vol. 34.
Curran Associates, Inc. (page 2)
Agrawal, Abhinav, Sheldon, Daniel R, Domke, Justin. 2020. Advances in Black-Box VI: Nor- malizing Flows, Importance Weighting, and Optimization. Pages 17358-177 17369 of: Advances in Neural Information Processing Systems, vol. 33. Curran Associates, Inc. (page 1)
Ahn, Kwangjun, Yun, Chulhee, & Sra, Suvrit. 2020. SGD with Shuffling: Optimal Rates without Component Convexity and Large Epoch Requirements. Pages 77526-17531 of: Advances in
Neural Information Processing Systems, vol. 33. Curran Associates, Inc. (page 10)
Alquier, Pierre, & Ridgway, James. 2020, Concentration of Tempered Posteriors and of Their Vari- ational Approximations. The Annals of Statistics, 48(3), 1475-1497. (page 10)
Altosaar, Jaan, Ranganath, Rajesh, & Blei, David. 2018. Proximity Variational Inference. Pages
1961 1969 of: Proceedings of the International Conference on Artificial Intelligence and Statis- tics. PMLR, vol. 84. JMLR. (page 7)
Amari, Shun-ichi. 1998. Natural Gradient Works Efficiently in Learning. Neural Computation, 10(2), 251 276. (page 10)
Bertin-Mahieux, Thierry, Ellis, Daniel P.W., Whitman, Brian, & Lamere, Paul. 2011. The Million Song Dataset. In: Proceedings of the International Conference on Music Information. (page 40) Bhatia, Kush, Kuang, Nikki Lijing, Ma, Yi-An, & Wang, Yixin. 2022 (July). Statistical and Compu- tational Trade-Offs in Variational Inference: A Case Study in Inferential Model Selection. arXiv Preprint arXiv:2207.1 1208.. arXiv. (pages 10)
Bingham, Eli, Chen, Jonathan P., Jankowiak, Martin, Obermeyer, Fritz, Pradhan, Neeraj, Karalet- sos, Theofanis, Singh, Rohit, Szerlip, Paul, Horsfall, Paul, & Goodman, Noah D. 2019. Pyro: Deep Universal Probabilistic Programming. Journal of Machine Learning Research, 20(28), 6.
(page 1)
Blei, David M., Kucukelbir, Alp, & McAuliffe, Jon D. 2017. Variational Inference: A Review for Statisticians. Journal of the American Statistical Association, 112(518), 859 877. (pages 2, 4)
Bottou, Léon. 1999. On-Line Learning and Stochastic Approximations. Pages 9-42 of: On-Line Learning in Neural Networks, first edn. Cambridge University Press. (page 2)
Botton, Léon. 2009. Curiously Fast Convergence of Some Stochastic Gradient Descent Algorithms.
(page 10)
Buchholz, Alexander, Wenzel, Florian, & Mandt, Stephan. 2018 Quasi-Monte Carlo Variational Inference. Pages 668-677 of: Proceedings of the International Conference on Machine Learning. PMLR, vol. 80. JMLR. (pages 4, 10)
Carpenter, Bob, Gelman, Andrew, Hoffmann, Matthew D., Lee, Daniel, Goodrich, Ben, Betancourt, Michacl, Brubaker, Marcus, Guo, Jiqiang, Li, Peter, & Riddell, Allen. 2017. Stan: A Probabilistic
Programming Language. Journal of Statistical Software, 76(1). (page 1)
Carvalho, Carlos M., Polson, Nicholas G., & Scott, James G. 2009, Handling Sparsity via the Horse- shoe. Pages 73-80 of: Proceedings of the International Conference on Artificial Intelligence and Statistics. PMLR, vol. 5. JMLR. (page 41)
Carvalho, Carlos M., Polson, Nicholas G., & Scott, James G. 2010. The Horseshoe Estimator for Sparse Signals. Biometrika, 97(2), 480. (page 41)
11

---

Challis, Edward, & Barber, David. 2013. Gaussian Kullback-Leibler Approximate Inference. Jour nal of Machine Learning Research, 14(68), 2239 2286. (pages 1, S, 10)
Christmas, Jacqueline, & Everson, Richard. 2011. Robust Autoregression: Student-T Innovations
Using Variational Bayes. IEEE Transactions on Signal Processing, 59(1), 48. (page 41)
Dhaka, Akash Kumar, Catalina, Alejandro, Andersen, Michael R, ns Magnusson, Må, Huggins, Jonathan, & Vehtari, Aki. 2020. Robust, Accurate Stochastic Optimization for Variational In- ference. Pages 10961-10973 of: Advances in Neural Information Processing Systems, vol. 33.
Curran Associates, Inc. (page I)
Dhaka, Akash Kumar, Catalina, Alejandro, Welandawe, Manushi, Andersen, Michael R., Huggins, Jonathan, & Vehtari, Aki. 2021. Challenges and Opportunities in High Dimensional Variational Inference. Pages 7787-7798 of: Advances in Neural Information Processing Systems, vol. 34,
Curran Associates, Inc. (page 1)
Diao, Michael Ziyang, Balasubramanian, Krishna, Chewi, Sinho, Salim, Adil. 2023. Forward- Backward Gaussian Variational Inference via JKO in the Bures- Wassersesnn Space. Pages 7960- 7991 of: Proceedings of the International Conference on Machine Learning. PMLR, vol. 202.
JMLR. (pages 7, 10)
Dieng, Adji Bousso, Tran, Dustin, Ranganath, Rajesh, Paisley, John, & Blei, David. 2017. Vari- ational Inference via x Upper Bound Minimization. Pages 2729-2738 of: Advances in Neural Information Processing Systems, vol. 30. Curran Associates, Inc. (page
Dillon, Joshua V., Langmore, Ian, Tran, Dustin, Brevdo, Eugene, Vasudevan, Srinivas, Moore, Dave, Patton, Brian, Alemi, Alex, Hoffman, Matt, & Saurous, Rif A. 2017 (Nov.). TensorFlow Distri- butions. arXiv Preprint arXiv: 1ee.1111..100000 10604. arXiv. (pages 3)
Domke, Justin. 2019. Provable Gradient Variance Guarantees for Black-Box Variational Inference. Pages 329-338 of: Advances in Neural Information Processing Systems, vol. 32. Curran Asso- ciates, Inc. (pages 3, 4, 8, 11, 22,
Domke, Justin. 2020. Provable Smoothness Guarantees for Black-Box Variational Inference. Pages 2587-2596 of: Proceedings of the International Conference Machine Learning. PMLR, vol. 119. JMLR. (pages I, 2, 3, 5, 6, 7, 8, 9, 10, 20, 25, 27, 33)
Domke, Justin, Garrigos, Guillaume, & Gower, Robert. 2023. Provable Convergence Guarantees for Black-Box Variational Inference. In: Advances in Neural Information Processing Systems (to Appear). New Orleans, LA, USA: arXiv. (pages 2, 11)
Dua, Dheeru, & Graff, Casey. 2017. UCI Machine Learning Repository. (page 40)
Duchi, John, Hazan, Elad, & Singer, Yoram. 2011. Adaptive Subgradient Methods for Online Learn- ing and Stochastic Optimization. Journal of Machine Learning Research, 12(Jul), 2121-2159.
(page 20)
Dugas, Charles, Bengio, Yoshua, Bélisle, François, Nadeau, Claude, & Garcia, René. 2000. Incor- porating Second-Ordes Functional Knowledge for Better Option Pricing. In: Advances Neural Information Processing Systems, vol. 13. MIT Press. (page 4)
Dwivedi, Raaz, Chen, Yuansi, Wainwright, Martin J., & Yu, Bin. 2019. Log-Concave Sampling: Metropoiis-Hastings Algorithms Are Fast. Journal of Machine Learning Research, 20(183), 1-
42. (p 2)
Fan, Kai, Wang, Ziteng, Beck, Jeff,, Kwok, James, & Heller, Katherine A. 2015. Fast Second Order
Stochastic Backpropagation for Variational Inference. Pages 1387 1395 of: Advances in Neural Information Processing Systems, vol. 28. Curran Associates, Inc. (page 10)
Fjelde, Tor Erlend, Xu, Kai, Tarek, Mohamed, Yalburgi, Sharan, & Ge, Hong. 2020. Bijectors.jl: Flexible Transformations for Probability Distributions. Pages 1-17 of: Proceedings of The Sym- posium on Advances in Approximate Bayesian Inference. PMLR, vol. 118. JMLR. (page 3)
Freund, Yoav, Ma, Yi-An, & Zhang, Tong. 2022. When Is the Convergence Time of Langevin Algorithms Dimension Independent? A Composite Optimization Viewpoint. Journal of Machine Learning Research, 23(214), 1-32. (page 10)
Fujisawa, Masahiro, & Sato, Issei. 2021. Multilevel Monte Carlo Variational Inference. Journal of Machine Learning Research, 22(278), 1-44. (pages 3, 4, 10)
12

---

Garrigos, Guillaume, & Gower, Robert M. 2023 (Feb.). Handbook of Convergence Theorems for (Stochasticc) Gradient Methods. arXiv Preprint arXiv:2301.1123: arXiv. (pages 2, 4, 4,21, 38)
Ge, Hong, Xu, Kai, & Ghahramani, Zoubin. 2018. Turing: A Language for Flexible Probabilistic In- ference. Pages (682-1690) of: Proceedings of the International Conference on Machine Learning.
PMLR, vol. 84. JMLR. (pages 10)
Gelman, Andrew, & Hill, Jennifer. 2007.
Data Analysis Using Regression and Multi- level/ffierarhhial Models. Analytical Methods for Social Research. Cambridge; New York: Cambridge University Press. (page 40)
Giordano, Ryan, Broderick, Tamara, & Jordan, Michael I. 2018. Covariances, Robustness, and Variational Bayes. Journal of Machine Learning Research, 19(51), 1-49. (page 1)
Giordano, Ryan, Ingram, Martin, & Broderick, Tamara. 2023 (Apr.). Black Box Variational Inference with a Deterministic Objective: Faster, More Accurate, and Even More Black Box. arXiv Preprint arXiv:2304.05527. arXiv. (page 41)
Goldberger, Ary L.,, Amaral, Luis A. N., Glass, Leon, Hausdorff, Jeffrey M., Ivanov, Plamen Ch., Mark, Roger G., Mietus, Joseph E., Moody, George B., Peng, Chung-Kang, & Stanley, H. Eugene. 2000. PhysioBank, PhysioToolkis and PhysioNet: Components of a New Research Resource for Complex Physiologic Signals. Circulation, 101(23). (page 41)
Gorbunov, Eduard, Hanzely, Filip, & Richtarik, Peter. 2020. A Unified Theory of SGD: Variance Reduction, Sampling, Quantization and Coordinate Descent. Pages 680-690 of: Proceedings of the International Conference on Artificial Intelligence and Statistics. PMLR, vol. 108. JMLR. (pages 27)
Gower, Robert Mansel, Loizou, Nicolas, Qian, Xun, Sailanbayev, Alibek, Shulgin, Egor, & Richtárik, Peter. 2019. SGD: General Analysis and Improved Rates. Pages 5200-520 5209 of: Pro- ceedings of the International Conference on Machine Learning. PMLR, vol. 97. JMLR. (pages 4, 3,)1,38)
Haochen, Jeff, & Sra, Suvrit, 2019. Random Shuffling Beats SGD after Finite Epochs. Pages 2624- 2633 of: Proceedings of the International Conference on Machine Learning. PMLR, vol. 97. JMLR. (page 10)
Hemandez-Looato, Jose, Li, Yingzhen, Rowland, Mark, Bui, Thang, Hernandez-Lobato, Daniel, & Turner, Richard. 2016. Black-Box Alpha Divergence Minimization. Pages 1511 1520 of: Pro- ceedings of the International Conference on Machine Learning. PMLR, vol. 48. JMLR. (page 2)
Hoffman, Matthew, & Ma, Yian. 2020. Black-Box Variational Inference Parametric Approxima- tion to Langevin Dynamics. Pages 4324-4341 of: Proceedings of the International Conference on Machine Learning. PMLR, vol. 119. JMLR. (pages 1, 10)
Jager, F., Taddei, A., Moody, G. B., Emdin, M., Antolič, G., Dorn, R., Smrdel, A., Marchesi, C., & Mark, R. G. 2003. Long-Term ST Database: A Reference for the Development and Evaluation of Automated Ischaemia Detectors and for the Study of the Dynamics of Myocardial Ischaemia. Medical and Biological Engineering and Computing, 41(2), 172-182. (pages 40, 41)
Jordan, Michael I., Ghahramani, Zoubin, Jaakkola, Tommi S. & Saul, Lawrence K. 1999. An Introduction to Variational Methods for Graphical Models. Machine Learning, 37(2), 183-233
(pages 3)
Karimi, Hamed, Nutini, Julie, & Schmidt, Mark. 2016. Linear Convergence of Gradient and Proximal-Gradient Methods under the chine Learning and Knowledge Discovery in Databases. Lecture Notes in Computer Science. Cham: Springer International Publishing. (page 7)
Condition. Pages 795-81 of: Ma-
Kawala, François, Douzal-
Ahlame, Gaussier, Eric, & Diemert, Eustache. 2013. Prédic- tions d' d'activité Dans Les Réseaux Sociaux En Ligne. Puge 16 of: Actes de La Conférence Sur Les Modèles et I! Analyse Des Réseaux Approches Mathématiques et Informatique. (page 40)
Khaled, Ahmed, & Richtárik, Peter. 2023. Better Theory for SGD in the Nonconvex World. Trans- actions of Machine Learning Research. (pages 7, 21,3,, 35)
Khan, Mohammad, Lin, Wu. 2017. Sunjugate-Computatie
Variational Inference: Convert- ing Variational Inference in Non-Conjugate Models to Inferences in Conjugate Models. Pages 878-887 of: Proceedings of the International Conference on Artificial Intelligence and Statistics. PMLR, vol. 54. JMLR. (page 10)
13

---

Khan, Mohammad Emtiyaz, Babanezhad, Reza, Lin, Wu, Schmidt, Mark, & Sugiyama, Masashi. 2016. Faster Stochastic Variational Inference Using Proxinal-Gradient Methods with General Divergence Functions. Pages 319-328 of: Proceedings of the Conference on Uncertainty in Arti-
ficial Intelligence. 1AI'16. Arlington, Virginia, USA: AUAT Press. (pages 7, 10)
Khan, Mohammad Emtiyaz E, Baque, Pierre, Fleuret, François, & Fua, Pascal. 2015. Kullback-
Leibler Proximal Variational Inference. In: Advances in Neural Information Processing Systems,
vol. 28. Curran Associates, Inc. (pages 7. 10)
Kim, Kyurae, Oh, Jisu, Gardner, Jacob, Dieng, Adji Bousso, & Kim, Hongseok. 2022. Markov Chain Score Ascent: A Unifying Framework of Variational Inference with Markovian Gradients. Pages 34802-348I6 of: Advances in Neural Information Processing Systems, vol. 35. Curran Associates, Inc. (page 2)
Kim, Kyurae, Wu, Kaiwen, Oh, Jisu, & Gardner, Jacob R. 2023. Practical and Matching Gradient Variance Bounds for Black-Box Variational Bayesian Inference. Pages 16853-16876 of: Proceed- ings of the International Conference on Machine Learning, PMLR, vol. 202. Honolulu, HI, USA: JMLR. (pages 1, 3, 4, 7. 8, 10, 21, 22, 34, 35, 36)
Kingma, Diederik P., & Ba, Jimmy. 2015. Adam: A Method for Stochastic Optimization. In: Pro-
ceedings of the International Conference on Learning Representations. (pages 2, 20)
Kingma, Diederik P., & Welling, Max. 2014 (Apr.). Auto-Encoding Variational Bayes. In: Proceed-
ings of the International Conference on Learning Represenlations, (page 2)
Kucukelbir, Alp. Tran, Dustin, Ranganath, Rajesh, Gelman, Andrew, & Blei, David M. 2017. Auto- matic Differentiation Variational Inference. Journal of Machine Learning Research, 18(14), 1-45. (pages 3,
Kunstner, Frederik, Chen, Jacques, Lavington, Jonathan Wilder, & Schmidt, Mark. 2023 (Feb.). Noise Is Not the Main Factor behind the Gap between Sgd and Adam on Transformers, but Sign Descent Might Be. In: Proceedings of the International Conference on Learning Representations. (page 9)
Lambert, Mare, Chewi, Sinho, Bach, Francis, Bonnabel, Silvére, & Rigollet, Philippe. 2022. Vari- ational Inference via Wasserstein Gradient Flows. Pages 14434-14447 of: Advances in Neural Information Processing Systems, vol. 35. Curran Associates, Inc. (page 10)
Leger, Jean-Benoist. 2023 (Jan.). Parametrization Cookbook: A Set of Bijective Parametrizations for Using Machine Learning Methods in Statistical Inference. arXiv Preprint arXiv:2301.088:
arXiv. (page
Lin, Wu, Khan, Mohammad Emtiyaz, & Schmidt, Mark. 2019. Fast and Simple Grauial- Variational Inference with Mixture of Exponential-Famie Family Approximations. Pages 3992 4002 of: Proceedings of the International Conference on Machine Learning. PMLR, vol. 97. JMLR.
(page 10)
Liu, Sifan, & Owen, Art B. 2021. Quasi-Monte Carlo Quasi-Newton Variational Bayes. Journal of Machine Learning Research, 22(243), 1-23. (pages 4, 10)
Magnusson, Måns, Bürkner, Paul, & Vehtari, Aki. 2022 (Nov.). Posteriordb: A Set of Posteriors for Bayesian Inference and Probabilistic Programming. (page 40)
Mishchenko, Konstantin, Khaled, Ahmed, Richtarik, Peter. 2020. Random Reshuffling: Sim- ple Analysis with Vast Improvements. Pages 17309-177 17320 of: Advances in Neural Information Processing Systems, vol. 33. Curran Associates, Inc. (page 10)
Naesseth, Christian, Lindsten, Fredrik, & Blei, David. 2020. Markovian Score Climbing: Variational Inference with (L(p))). Pages 15499-15510 of: Advances in Neural Information Processing Systems, vol. 33. Curran Associates, Inc (page
Nagaraj, Dheeraj, Jain, Prateek, & Netrapalli, Praneeth. 2019. SGD without Replacement: Sharper Rates for General Smooth Convex Functions. Pages 4703-4711 of: Proceedings of the Interna- tional Conference on Machine Learning. PMLR, vol. 97. JMLR. (page 10)
Nemirovski, A., Juditsky, A., Lan, G., & Shapiro, A. 2009. Robust Stochastic Approximation Ap- proach to Stochastic Programming. SIAM Journal on Optimization, 19(4), 1574-1609. (page 2)
Nguyen, Lam, Nguyen, Phuong Ha, van Dijk, Marten, Richtarik, Peter, Scheinberg, Katya, & Takac, Martin. 2018. SGD and Hogwild! Convergence without the Bounded Gradients Assumption.
14

---

Pages 3750-3758 of: Proceedings of the International Conference on Machine Learning. PMLR, vol. 80. JMLR. (page4) 4)
Patil, Anand, Huard, David, & Fonnesbeck, Christopher. 2010. PyMC: Bayesian Stochastic Mod-
elling in Python. Journal of Statistical Software, 35(4). (page 1)
Ranganath, Rajesh, Gerrish, Sean, Blei, David. 2014. Black Box Variational Inference. Pages
314-822 of: Proceedings of the International Conference on Artificial Intelligence and Statistics. PMLR, vol. 33. JMLR. (page 1)
Reddi, Sashank J., Kale, Satyen, & Kumar, Sanjiv. 2023 (May). On the Convergence of Adam and Beyond. In: Proceedings of the International Conference on Learning Representations. (page 9)
Regier, Jeffrey, Jordan, Michael I, & McAuliffe, Jon. 2017. Fast Black-Box Variational Inference through Stochastic Trust-Region Optimization. Pages 2399-2408 of: Advances in Neural Infor- mation Processing Systems, vol. 30. Curran Associates, Inc. (page 10)
Robbins, Herbert, & Monro, Sutton. 1951. A Stochastic Approximation Method. The Annals of
Mathematical Statistics, 22(3), 400-407. (page 2)
Roeder, Geoffrey, Wu, Yuhuai, & Duvenaud, David K. 2017. Sticking the Landing: Simple, Lower- Variance Gradient Estimators for Variational Inference. Pages 6928-6937 of: Advances in Neural Information Processing Systems, vol. 30. Curran Associates, Inc. (pages 2, 21)
Shannon, Paul, Markiel, Andrew, Ozier, Owen, Baliga, Nitin S., Wang, Jonathan T., Ramage, Daniel, Amin, Nada, Schwikowski, Benno, & Ideker, Trey. 2003. Cytoscape: A Software Environment for Integrated Models of Biomolecular Interaction Networks. Genome Research, 13(11), 25044
(page 40)
Titsias, Michalis. 2009, Variational Learning of Inducing Variables in Sparse Gaussian Processes. Pages 567-574 of: Proceedings of the International Conference Artificial Intelligence and Statistics. PMLR, vol. 5. JMLR. (page 4)
Titsias, Michalis, & Lázaro-Gredilla, Miguel. 2014. Doubly Stochastic Variational Bayes for Non- Conjugate Inference. Pages 1971-1979 of: Proceedings the International Conference on Ma- chine Learning. PMLR, vol. 32. JMLR. (pages 1, 2, 3, 5, 6, 21)
van der Vaart, A. W. 1998. Asymptotic Statistics. First edn. Cambridge University Press. (page 7)
Vaswani, Sharan, Bach, Francis, & Schmidt, Mark. 2019. Fast and Faster Convergence of SGD for Parameaameter Models and an Accelerated Perceptron. Pages (195-1204 of: Proceedings of the International Conference on Artificial Intelligence and Statistics. PMLR, vol. 89. JMLR. (page 4)
Wang, Yixin, & Blei, David. 2019. Variational Bayes under Model Misspeciffcation. In: Advances in Neural Information Processing Systems, vol. 32. Curran Associates, Inc. (page 7)
Welandawe, Manushi, Andersen, Michael Riis, Vehtari, Aki, & Huggins, Jonathan H. 2022 (Mar.). Robust, Automated, and Accurate Black- Box Variational Inference. arXiv Preprint arXiv:2203. 11.22033.1944. 15945. arXiv. (page !)
Wright, Stephen J., & Recht, Benjamin. 2021. Optimization for Data Analysis. New York: Cam-
bridge University Press. (page 21)
Xu, Ming, Quiroz, Matias, Kohn, Robert, & Sisson, Scott A. 2019. Variance Reduction Properties of the Reparameteriization Trick. Pages 2711-2720 of: Proceedings of the International Conference on Artificial Intelligence and Statisties. PMLR, vol. 89. JMLR. (page 10)
Xu, Zuheng, & Campbell, Trevor. 2022. The Computational Asymptotics of Gaussian Variational Inference and the Laplace Approximation. Statistics and Computing, 32(4), 63. (page 10)
Yao, Yuling, Vehtari, Aki, Simpson, Daniel, Gelman, Andrew. 2018. Yes, but Did It Work?: Eval- uating Variational Inference. Pages 5581-55999 of: Proceedings of the International Conference
on Machine Learning. PMLR, vol. 80. JMLR. (page 1)
Yun, Jihun, Lozano, Aurelie C, & Yang, Eunho. 2021. Adaptive Proximal Gradient Methods for
Structured Neural Networks. Pages 24365-24378 of: Advances in Neural Information Processing Systems, vol. 34. Curran Associates, Inc. (pages 2, 9, 20)
Zhang, Cheng, Butepage, Judith, Kjellstrom, Hedvig, & Mandt, Stephan. 2019. Advances in Varia- tional Inference. IEFE Transactions on Pattern Analysis and Machine Intelligence, 41(8), 2008- 2026. (pages 2, 9)
15

---

Zhang, Yushun, Chen, Congliang, Shi, Naichen, Sun, Ruoyu, & Luo, Zhi- QhaiQQan. 2022. Adam Can Converge without Any Modification on Update Rules. Pages ral Information Processing Systems, vol. 35. Curran Associates, Inc. (page 9)
of: Advances in Neu-
16

---

On the Convergence of Black-Box Variational Inference Appendix
Table of Contents
Introduction
Background
2
2.1
Black-Box Variational Inference
2
2.2
Variational Family..
3
2.3
Scale Parameterizations
3
2.4
Problem Structure of Black-Box Variational Inference
4
The Evidence Lower Bound Under Nonlinear Scale Parameterizations
5
3.1
Technical Assumptions
5
3.2
Smoothness of the Entropy
3.3
Smoothness of the Energy
5
3.4
Convexity of the Energy
4 Convergence Analysis of Black-Box Variational Inference
7
4.1 Black-Box Variational Inference
7
4.2
Black-Box Variational Inference with Proximal SGD
7
5
Experiments
9
5.1
Synthetic Problem
9
5.2
Realistic Problems
9
6 Discussions
10
Computational Resourees
B Nomenclature
18
C Definitions
19
D ProxGen Adam for Black-Box Variational Inference
20
E Detailed Comparison Against Domke et (2023)
21
F
Proofs
22
F.1
Auxiliary Lemmas
22
F.2
Properties of the Evidence Lower Bound
24
F.2.1
Smoothness
24
F.2.2
Convexity
30
F.3
Convergence of Black-Box Variational Inference
34
F.3.1
Vanilla Black-Box Variational Inference
34
F.3.2
Proximal Black-Box Variational Inference
36
Details of Experimental Setup
40
Additional Experimental Results
42
17

---

A Computational Resources
Table I: Computational Resources
Type
Model and Specifications
System Topology
2 nodes with 2 sockets each with 24 logical threads (total 48 threads)
Processor
1 Intel Xeon Silver 43 10, 2.1 GHz (maximum 3.3 GHz) per socket
Cache
1.1 MiB LI, MiB L2, and 36 MiB L3
Memory
250 GiB RAM
Accelerator
1 NVIDIA RTX A5000 per node, 2 GHZ, 24GB RAM
Running the experiments took approximately a week.
B
Nomenclature
Symbol
Definition
Description
Section
A
Variational parameters Parameters the target model æ location-scale reparameterization function
2.1
z
2.1
T
- c(2) m
2.2
Random vector before reparameterization Base distribution of u Location parameter (part of 2)
2.2
q
2.2
2.2
Scale parameter (part of 2)
2.2 and 2.3
Diagonal conditioner
2.3
k
Kurtosis (non-central 4th moment) of Diagonal of C using the diagonal conditioner &
2.3
D& (s)
2.3
Strietly lower triangular part of C Elements forming the diagonal of
2.3
2.3
(z)
A log p (z,
Negative joint likelihood
2.4
f (2)
A Ex-ga l (z)
Energy
2.4
h (2)
A -HI
Negative entropy
2.4
F F(A)
f(2) + h (A)
Negative ELBO
2.1 and 2.4
f (();
l (I (u))
Negative Log-Likelihooe under reparameterization
2.4
(2:4)
Ve (J (2))
Gradient of the Log-likelihood under reparamcterization
2.4
M
Number of Monte Carlo samples
2.1
Stepsize of (proximal) SGD at iteration / Reparameterization gradient estimator of the energy
4.1 and 4.2
Vf
4.1
18

---

C Definitions
For completeness, we provide formal definitions for some of the we used throughout the paper.
Definition (Smoothness). A function f 2 H is said be L-smooth the inequality
I f (z') | z'Il
holds for all z' E 2.
This assumption, also occasionally called Lipschitz smoothness, restricts the amount the gradient can change for a given distance. When f twice differentiable, an equivalent condition the Hessian to be bounded:
Definition (Smoothness). A twice differentiable function f R. said be L-smooth if the inequality
(z)
holds for all z E 2.
Remark Assuming a function f is smooth is equivalent to assuming that can be upper bounded by quadratic function everywhere.
Remark 8. When the log-density log π of a probability measure II L-smooth, log π can be upper bounded everywhere by the log-density of'a Gaussian.
Definition (Strong Convexity). twice differentiable function f ba is said to be j-strongly convex if the inequality
l - z'' (V)
holds for all z, z' E Rd and some μ
Remark 9. If Definition holds only for μ = 0, f is said to be (non-strongly) convex.
Remark 10. Assuming function f is strongly convex is equivalent to assuming that can be lower bounded by a quadratic.
Definition (Strongly Log- Log-Concave Measures). For probability measure II in Euclidean measurable space (Rd (Rd). , P), where B (Re) is the o-algebra of Borel-mmasurable subsets of Rd, P is the Lebesgue measure, we say II is u-strongly log-concave if its log-density log π (z) R is f-stroogly convex for some μ > 0.
Rdd +
Remark 11. Definitio 8 holds only for µ 0, II is said to be (non-strongly) log-concave.
Remark 12. When II is u-strongly log-concave, log π can lower bounded everywhere by the log-density of a Gaussian.
19

---

D ProxGen Adam for Black-Box Variational Inference
Algorithm 1: ProxGen-Adam for Black-Box Variational Inference
Input:
Initial variational parameters 2y, base stepsize a, second moment stepsize 82. momentum stepsize (f1,111=1 small positve constant
for / I,...TT do
estimate gradient of energy Vf
g: Vf + Vh (2)
2 مr,r + (1 $1,1) d
11 $zU; + 82) gi
Fi+1 diag (a/ (v /D(+1 + E))
2 2, - 1++14++1
51+1 - getscale (d++1)
8,+1 - S1+1 + 2
V sit1 + 47 S/+1
1, setscale S(+1)
end
(By convention, all vector operations are elementwise.)
Adaptive and matrix-valued stepsize-variants of SGD such as Adam (Kingma & Ba, 2015), Ada- Grad (Duchi al., 2011) are widely used. The matrix stepsize of Adam at iteration t given as
r++ = diag (a/ (ut+ + +)),
where D, the exponential moving average of the second moment, a the "base stepsize." Further- more, the matrix stepsize is applied to the moving average of the gradients, a scheme often called the (heavy-ball) momentum, denoted here dz.
Recently, Yun et al. (2021) have proven the convergence for these adaptive, momentum, and matrix- valued stepsize-based SGD methods with proximal steps. Then, the proximal operator is applied as
prox (2, (4,rr;ā,) = arg min { h ha) aa= 2.) (2 (2-22)
For Adam, the matrix-valued stepsize a diagonal matrix. Thus, the proximal operator of Domke (2020) for each s; forms independent 1-dimensional quadratic problems. Thus, the proximal step is given in the closed-form
prox (5)) = (V
4Ys, -t).
where, dropping the index & for clarity, 3; is the element of 2, corresponding to Sj, Ys, denotes the stepsize of s; (a diagonal element of of Combined with the Adam-like stepsize rule, the algorithm is shown in Algorithm 1.
see the differences with vanilla Adam. Notably, ProxGenAdam does not perform bias correction the estimated moments. Furthermore, while some implementations of Adam decay P1 keep it constant. It is possible that these differences could result in a different behavior from vanilla Adam. However, in this work, we follow the original implementation by Yun al. (2021) closely possible and leave the comparison with vanilla Adam to future works.
Difference with Adam In Algorithm
20

---

E Detailed Comparison Against Domke et al. (2023)
In this section, we contrast our results against those of Domke et al. (2023). First, the main challenge to establishing a convergence guarantee for BBVI has been on bounding the gradient variance. In particular, Domke (2019) proved that the variance of the reparameterization gradient for the energy, VJ, is bounded as
ElIVJ (1)1 allA 21|;
for some finite positive constants a, & depending on the problem constants d, L, kg. Domke et al. (2023) call gradient estimator satisfying this bound be "quadratic variance" estimator. Fur- thermore, they prove that the closed-form entropy (CFE; Kucukelbir et al., 2017; Titsias & Lázaro- Gredilla, 2014) estimator:
CFF (2) A (T)
(A) + Vh
and the STL estimator by Roeder al. (2017):
VF STL (2)
Σ
V2l (I1 (um) + 2 log Qu (I (um))
y=d
m=1
where q, also qualify quadratic variance estimators.
Unfortunately, it has been unknown whether SGD is guaranteed to converge with quadratic vari- ance estimator except for strongly convex objectives (Wright & Recht, 2021, 85). Domke et al. (2023) expand the boundaries of SGD and prove that projected and proximal SGD with a quadratic variance estimator converges for both convex and strongly convex objectives. In particular, for the location-scale variational family, the linear parameterization, and log-concave objectives, they prove complexity of O (1/e²), and for strongly log-concave objectives, they prove complexity of (1/E).
On the other hand, Kim et al. (2023) and Lemma 12 developed the bound in Equation (3) to be of the form of
VP (2)11, & (A(F(2) IVF (2)1 +C
(4)
for some positive finite constants A, B, C, for which the convergence of SGD for convex, strongly convex (Garrigos & Gower, 2023; Gorbunov al., 2020), and non-convex objectives (Khaled & Richtárik, 2023) have already been proven. Applying these results to log-smooth and log- quadratically growing objectives, we prove a complexity O (1/c*), while for strong log-concave objectives, we also prove complexity of O (1/c).
Overall, both approaches can be summarized follows: focused establishing gradient vari- ance bounds of known convergence proofs, while Domke et al. (2023) aimed to prove that the bound by Domke (2019) is sufficient guarantee convergence. Note that, for strongly log-concave objec- tives, Equation (3) immediately implies Equation (4). Therefore, both approaches intersect in the case of strongly log-concave objectives. Indeed, Theorem 8 and the analogous result of Domke et al. (2023) are both based on the same proof strategy by Gower et al. (2019).
21

---

Proofs
F.1 Auxiliary Lemmas
Lemma Let (x)==
= x. Then, the parameterization linear in the sense that Ja bilinear
function such that
(2) (4)
for any 2, E A.
Proof.
J2- (u) (C u (m- m')
+ (L L')) u (m m'),
using the fact that & is the identity function,
(Dz - D, (s') + (L -L-L") L')) u (m m')
= (C C(2)) u (m + m')
(C m) - (C (2')u+m))
= (4)- F2 (u).
The linearity with respect is obvious.
Lemma Let the linear parameterization be used. Then, for any 2, E the imner product of the Jacobian of the reparameterization function satisfies the following equalities for any u E R.d.
ii) For the Cholesky family (Domke, 2019, Lemma 8),
@0%%
(u)
3JA (u)
= (1+ (1+1u1]))
82
(ii) For the mean-field family (Kim et al., 2023, Lemma 1),
T (u)
T. (u)
(1+ ||
aA
aa
where U diag (u...... Ma))
Lemma Let the linear parameterization be used. Then, for any A and any z E Rd the
following relationships hold.
(i) For the Cholesky family (Domke, 2019, Lemma 2),
E(1 + [[ll()) |J (u) z||2 (d 1) |m =|I + (d + I|?
(ii) For the mean-field family (Kim et al., 2023, Lemma 2),
E(1 + (1++10P1)))11 (u) 21i- (Vdk kp Va + 1) Im =2ll² + (2kp Va 1)|C
Corollary 2. Let the linear parameterization be used and 2, E be any pair of variational parameters.
(i) For the Cholesky family,
E [| 152 (4) - (w) < (kp + - 1'12
13
(ii) For the mean-field fimm-fiiimmmmmii
(1 (0111() (u) Ja (v) (2k V2 1)] 2'12
Proof. The results are direct consequence of Lemma 7 and Lemma 5.
22

---

Proof of (i) We start from Lemma 7 as
[1+1|u[ß) J- (u) z||² = (d (m m²) z1ī + kd -))
setting 0,
= (d + 1) m (d (g) -
and since k 2 3 by the property of the kurtosis,
$ (d kpp (Ilm m
|Ic C (2')11²)
(d kpp 2
Proof of (ii) Similarly, for the mean-field family, we can apply Lemma 7 as
E(1 + ||0 - (w) =1l² 2 (Vkk, dkp + kp Vd [ m²) 21ī + (2kp Va + ||c ll
setting 0,
= =VVkk dkp + kpvd+ I|m m' (2kg - C/I
and since kp 2 3 by the property of the kurtosis,
s (2kp V 1) (11m-m' mmm -1C(
= 2k V + 1)112 2'112.
Lemma 8. For the linear parumeterization,
EIJA (u) - Ty. (u)li 11a
2'112 12-2'112
for any 2, E A.
Proof. First notice that, for linear parameterizations, we have
FllTa (w) E E||Cuu + mll
Eu CT Cu + ||m][ + 2m T CEu
tr T Cu) + |m 1| 2m T CEu,
rotating the elements of the trace,
tr (arCTCEuu CEuu ||ml[2 + 2m CEu,
applying Assumption
tr(cTTc) |m|
[[C + ||m||
Combined with Lemma 5, we have
E (u) - Tz. (u)liz = El|T2>2/ (4)li, = ||2 -2'1l"
23

---

F.2
Properties of the Evidence Lower Bound
F.2.1
Smoothness
Lemma 1. If the diagonal conditioner & is Liploog-smooth, then the entropic regularizer (2) is Lp-smooth.
Proof.
The entropic regularizer is
(2)) -
Σ log & (s)),
and depends only on the diagonal elements $1, Sq of C. The Hessian of is then diagonal matrix, where only the entries that correspond to Si...., Sp are non-zero. The Lipschitz smoothness constant then the constant Ln co that satisfies
hh (2)
d² log Φ
Ln
ds
dsf
for all = 1, l,....dd d, which is the smoothness constant of s log & (s²).
Lemma Let H be symmetric random matrix, where is bounded ||H[]] almost surely. Also, let J be an m X random matrix such that [|EJ Jll2 co. Then,
|ESTHJ]|z L||EJT
Proof. By the property the Rayleigh quotients, for symmetric matrix 4, its maximum eigenvalue given the variational form
Hx max (H) S O max (H)² =|H|
sup |
where max (A) is the maximal eigenvalue of A. Notice the relationship with the 2 norm. The inequality is strict only if all eigenvalues are negative.
From the property above,
||U HJ] sup x (FJTHJ)x
By reparameterizing as = Jx,
sup Ey T Hy, |
and the property of the e operatoo norm,
sup E [||H|| 11112
sup IF] [||HH||]|Jxx]] 1xlg
From our assumption about the maximal eigenvalue of H,
< L sup E||Jx|| |x,,1
denoting the & vector norm as a quadratic form as,
L sup x (FJT
again, by the property of the e--operator norm,
< L 2ll sup |x||
|x|251
24

---

Lemma 9. For 1-L.ipshiiz diagonal conditioner 4, the Jacobian of the location-scale reparam- eterization function J2 satisfies
0J2 (u)
8. (u)
E
1
02
32
Proof. For notational clarity, we will occasionally represent Ja
Jz (u) = J u2:
such that Fi uA: denotes the ith component of
From the definition of J2 it is straightforward to notice that Jacobian the concatenation block matrices
aJa (2)
8J% (u) J,
3T% (u)
Jm
and Jz òvec (L)
õm
Ōs
The block form deterministic identity matrix
J 3J% (
I,
which is shown by (Domke, 2020, Lemma 4).
The proof strategy is as follows: we will directly compute the squared Jacobian through block matrix multiplication. The key is that, after expectation, the resulting matrix becomes diagonal. Then, the l2 operator norm, or maximal eigenvalue, follows trivially as the maximal diagonal element.
First,
aJ (2:u)
T
a. (2:u)
E
E
J
J,
Jz]
02
a2
J
JH m
J
J m
J
J
EJ
E
E
FJ
EJ
For J,, the entries are
cJ (u) dsy
4 (s))
which is diagonal matrix. Thus, by Assumption 1,
EJ = O,
EJ J3 = diag ($' (s))
For Jg. the entries are
a. (2;u) uxli=j-
aLjk
To gather some intuition, the case of d = looks like the following:
U2
U3
U4
U1
U3
Ua
J =
U1
U2
U4
u
U2
It is crucial to notice that the ith row does not include uj. This means that, the matrix J. J has entries that are either 0, or X' us;) for # J. which EØ' (s;) u,u, = 0 by Assumption 1. Therefore,
EJ.J = O.
25

---

Finally, the elements of J JL are
u2:
8J (A;
Σ
E Σ
Un
li = j= (Eunum) jmm
JLjk jk
iLim
1=0
1=0
where the last equality follows from Assumption 1, which forms an identity matrix
EJ
Therefore, the expected-squared Jacobian is now
EJ, E
EJ J J
TTA (u)
8 (u)
EJ
02
aA
FJ
FJ
I
diag ($ (s))
which, conveniently, is a diagonal matrix. The maximal singular value of a block-diagome matrix the maximal singular value of each block. And since each block diagonal with only positive entries, the largest clement forms the maximal singular value. As we assume that $ is I-Lipchitz, the element of all blocks lower-bounded by 0 and upper-bounded by 1. Therefore, the maximal singular value of the expected-squared Jacobian is bounded by 1.
26

---

Theorem 1. Let é be Lg-Smooth and twice differentiable. Then, the following results hold: (i) If Φ is linear, the energy f is L.gssmotth
(ii) If & is 1-Lipschitz, the energy e is (Le + Ly>-mmothh if and only if Assumptiio holds.
Proof. For notational clarity, will occasionally represent Jj
JA (w) = - u2:
such that Fi u); denotes the ith component of J. By the Leibniz and chain rule, the Hessian the energy follows as
f2 EV3 (Jz (4))
E 8Ji (u)
8JA (4)
d
(2;((1:4)
42l (I2 (u))
+
E
e
(4))
02
0A
222
t=1
Tim
T
When J' is linear with respect to A, it is clear that we have
Tr I (A;
0.
(5)
222
Then, Trom is zero. In contrast, Tin appears for both the linear and nonlinear cases. Therefore, T non fully characterizes the effect of nonlinearity in the reparameterization function.
Now, the triangle inequality yields
IV (2)ll, = Tín + Toolllz S [ + noTlll
where equality is achieved when either term is O. On the contrary, the reverse triangle incquality states that
Ti 1T | (2)12-
This implics that, either Tim or T pon is unbounded, the Hessian is not bounded. Thus, ensuring that Tio and T Too bounded is sufficient and necessary to establish that f is smooth.
Proof (i) The bound on the linear part, Timn follows from Lemma 2 as
| 12 =
E
OJA (u)
v2 (TA (u))
aJA (u) 02
02
TJ (u)
aJ (u)
$ Lg
22
ad
and from the I-Lipschitzness of $ Lemma 9 yields
< Le.
When & is linear, it immediately follows from Equation (5) that
| (2)l22 = l|Tuml S Le,
which is tight as shown by Domke (2020, Theorem 6).
Proof of (ii) For the nonlinear part T Tmon- the fact that Fj (2; is given
Fi u): m; (s)) u, Ln
J
27

---

The second derivative of Jj is clearly non-zero only for the nonlinear part involving $n Sg. Thus, Tuon follows
d
8T (A:u)
non = E
Σ
Dee (Js (4)
822
d
(2;u)
E Σ
8i uA;
O22
i=1
E
E i=1 Si
(4:4)
ds²
Furthermore, the second-order derivatives with respect to $p.... ....... Sd given
0J (2;
1==00" (5j).
jsj
Considering this, the only non-zero block of T non forms diagonal matrix as
Egi (2; (F
(2:u)
daf
Σ 1=1
g: (2:
ōs
Ega (d; gFga(dzu) ds²
Egi (2; &" (s1)
Ega (2; &" (sa) Val
This implies that the only non-zero entries of non lie its diagonal. Since the e₂ norm of diagonal matrix is the value of the maximal diagonal element,
|T 1onlll2
< max Eg (A:u) 4" (s:) U: S Ly.
i=l,...,d
where Lg is finite constant if Assumption 3 holds. On the contrary, if finite 4g does not exist, [| ll cannot be bounded. Therefore, the energy smooth if and only if Assumption 3 holds. When does, the energy f is Ly + L, smooth.
28

---

Example 2. Let 8 (z) = (1/2) Z' Az and the diagonal conditioner be $ (x) = softplus (x). Then,
if A is dense and the variational family is the mean-field family or
(ii) if A is diagonal and the variational family is the Cholesky family,
Assumption holds with Ly 0.26034 (m.. Aq) (iii) If A is dense but the Cholesky family is used, Assumption does not hold.
Proof. Since the gradient is combined with reparameterizzation, we have
v (z) Az,
g (A:u) = A (Cu m)
Then, for each coordinate i = 1,,.... d, we have
Egi (A:u) ui (st) =
Σ Σ
CjUK + LAymy (1)"
Σ Σ (n(CREENNU(""())) + Σ AmmmEEurd""( (si).
ksl
and from Assumption
(') Σ AicCi
J ksj
*" (s) Cn.
Furthermore, the diagonal of C involves & such that
Eg: (2; u) uid"(s;) A (s:) 6" (s:)+
(s))ss
Tdog
j<i
Twn
For the softplus function, we have for any finite s, and we have
$" (s) < 1
sup Φ (s) $" (s) ~ 0.26034,
where the supremum was numerically approximated. Then, it is clear that Taiug is finite as long as the diagonals A are finite. Furthermore, we have the following:
(i) If A is diagonal, then Tofr is 0.
(ii) If is dense but C' is diagonal due to the use of the mean-field family, Ton is again 0.
(iii) However, when both and are not diagonal, Ton can be made arbitrarily large.
29

---

F.2.2 Convexity
Lemma 10. Let l be convex. Then, for convex nonlinear Φ. the inequality
(Vaf (2) - 2') < E(Vg (), (4) T2 (u))
holds for all à E A if and only if Assumption holds. For the linear parameterization, the inequality becomes equality.
Proof. First, notice that the left-hand side is
1²= - (V)E ( 2') =E =
>
t=1 dā, )8 ((2:u), 2.)
2.
[ (w)
<V
i=1
By restricting us to the location-scale family, we then get
TT (u)
a.A (u)
=E
Σ
g (2: (m; m!) +
Σ
(2: u) (Lij -iij)
3m:
رياة
iJ
convexity with respect
convexity with respect
OJA (w) as₁
Σ
g (2s
convexity with respect to
and plugging the derivatives of the reparameterization function,
=E Es (2: (m, mm;)
ΣΣ (ă:u) (Li Lig) (0) (s
j
On the other hand, the right-hand side follows as E (Vg((u;)) IA (u) J2 (u)
=E (2: m') + (g (2: (L- L') (g(d:)), (0(s) (0())
=E Σ si u): (m; - m3) ΣΣ U281 ug (2:
Es (2: u, (s)) (s1)
The convexity with respect the m and L is clear from the first terms; they are equal. The statement is now to the last term. That is, the statement holds if
IF Σ gi (A:u) u; & (s) sf) s E Σ gi (2:u) ui ( (s)) & (5)))
(6)
For this, we will show that Assumption both necessary and sufficient.
Proof of sufficiency Equation (6) holds if
Eg, (2: u; 0
for all 1, d, which is non other than Assumption 4. Proof of necessity Suppose that the inequality
(2) 2 2') < ((g Fx. (u)
holds for all Å E A, implying
Eug (A:))
s/) < Σ Euggi (i; ((p(s)) - #ss)).
For any 2, we are free set any 2' E A and check whether we can retrieve Assumption for this specific A. Now, for each axis i, set sj sj for all j # then
E Eug gi (d:u) &' (s) 2i < E un uig. (+(s:) $(s))).
30

---

Since Φ is assumed to be convex such that
$ - ( (s,) $ (5)).
it follows that
E Eu gi (2; 2 0.
Therefore, for any à A it must be that Assumption 4 holds.
Proposition We have the following:
(i) If l is convex, then for the mean-field family, Assumption 4 holds.
(ii) For the Cholesky family, there exists a convex & where Assumption 4 does not hold.
Proof. For (i), the key property is the monotonicity of the gradient.
Proof of (i) For the nean-field family, recall that
C = (s:).
Also, observe that
Cu + m (C1141 + mi. ... Caaua + ma). By the property of convex functions, Ve is monotone such that
(Ve (z) V (z)) z') 0.
Now, by setting z = Cu + m and z' = Cu + m Cjune, we obtain (ve (Cu + m) ve (Cu m Cju,e,), Cjuee;) 2
for every = l.....d.
For the mean-field family, Cu + Cjule is now independent of U;- Thus, C,u,D, (Cu + m) 2 E Cuu, D (Cu + m - e,) =a (EE)) (Cu+ c,C)))
=0,
where D,f denotes the ith axis of Vf. Since Cj 0 by design,
D (Cu m)
Eu,Dif (Cu+ m) > 0,
which is Assumption
Proof of (ii)
We provide an example that proves the statement. Let ((z) x Az. Then, g (A: e (Iz (u)) (Cu + m) ACu Am.
Suppose that we choose 2 such that
C 1 1
and m = 0. Also, setting
-2
A
-2
we get strongly convex function 6. Then,
1
-2
0
[-1
-2
41
- 2u2
(2: = ACu
-2
5
1
U2
=
U2
3uj + 5u2
Finally, we have
Egi (A:4) = E(-u, - 22z) = -1 <
which violates Assumption
31

---

Lemma 11. For any function f E (R,,R,, there is no constant 0 < L < co such that
If(x) - -()) 2 x|- y/.
Proof. Suppose for the sake of contradiction that such I. > 0 exists. Letting y - x gives If) 2 for all E R. For each x, either fx) -L f (x) L. holds. We discuss cases based on the value f0))..
If (') 2 L, we claim that f'(x 2 L for all E R. Otherwise, f'(x) < L for some x implies ('(x V -L. By the intermediate value theorem (f" is continuous), there exists a point y between and that attains the value (y) O, which is contradiction.
Now that f'(x) 2 L > 0 for all x, f an increasing function. For any 0, we have
f(x) f(x) f(0) f(0)
= -If(x) - f(0)+ + f(0)
I -L|x f(0).
Here, we can plug x' f(0)
(1x)=-1-11 f(0)
f(0) -I((0) + f(0) 10,
which implies that f(x) $ R4, which is a contradiction.
Now discuss the second case f'( v -L. By similar argument, f'(x -L for all x R. Thus, f is a decreasing function. For any x > 0, we have
f(x) f(x) - f())) f(0)
S-Lx+f(0)
Picking x' f(0) results in f(x) E 14, which a contradiction.
32

---

Theorem 2. Let é be j-strongly convex. Then, have the following:
(i) If Φ is linear, the energy f is u-strongly convex.
(ii) If & is convex, the energy f convex if and only if Assumption 4 holds.
(iii) If Φ is such that & C² (R, R4), the energy f is not strongly convex.
Proof. The special case (i) is proven by Domke (2020, Theorem 9). We focus on the general state- ment (ii).
is j-strongly convex, the inequality
(z) (z))
z'112
(8)
holds, where the general convex case is obtained special case with 0. The goal is to relate this to the (u-strong- Jconvexity of the energy with respect to the variational parameters given by
(()=
1
Proof of (ii) tion, we have
Plugging the reparameterized latent variables to Equation (8) and taking the expecta-
Ee (4)) (Th (4) E(Ve (0)) Jy. (4)) ト E (u) JTa (4)l
2
E(Ve (0) J1 (u) To (4)) /
EI (u) J3 (4))
(2)- > E(Vg (4' T2 (4) EIIT (u) T2. (u)
Thus, the energy is convex if and only if (g (A: u) J (u) Jx (u)) 2 2 2') holds. This is established by Lemma 10.
Proof of (iii) We now prove that, under the nonlinear parameterization, the energy cannot be
strongly convex. When the energy is convex, it is also strongly convex if and only if
2
E (4) - T2 (u)1 2
" 1| 2'11
From the proof of Domke (2020, Lemma 5), follows that
E (u) Tz. (v)I = ||C c'I + ||m m' n2'12.
Furthermore, under nonlinear parameterizations, ||C - C' -f' + ||m m'1
=|( (s) D+ (s') (L- (L- + -|mm mm'1
expanding the quadratic,
= =|p (s) D3 (s' + I|L (D
L'); + ||m m²'l
and since D& (s) and L reside different sub-spaces, they are orthogonal. Thus,
= Dg (s) - D& |1 + IL L'll + ||m m
= 11 (s) - & (s)) IL 2 + ||m m'll
For the energy term to be strongly convex, Equation (9) must be bounded below by I| 2'112 Evi- dently, this implies that a necessary and sufficient condition is that
$$ (sH) & (5/a 2 L |S - 2/11
by some constant Lipschitz condition. Unfortunately, there is no such continuous and differentiable function & R - R4, as established by Lemma 11. Thus, for any diagonal conditioner & C² R4), the energy cannot be strongly convex.
co. Notice that the direction of the inequality is reversed from the
33

---

F.3 Convergence of Black-Box Variational Inference
F.3.1 Vanilla Black-Box Variational Inference
Theorem 5. Let the variational family satisfy Assumption 2, the likelihood satisfy Assumption 5, and the assumptions of Corollary 1 hold such that the ELBO, F, is Lp-smooth with Lp = Le L + Lg. Then, if the stepsize satisfy y < 1/Lp, the iterates of BBVI with SGD and the M-sample reparameterization gradient estimator satisfy
2 C (d. Q)
min EIIVF( S 0515T-1
(11 Ziikell + --((F=ffi)
M
yT + y2 4lppeex c(d.))) (F - FF).
where
joint = proje (z)
is the projection of z onto set of minimizers of is the projection of z onto set of minimizers of l like. is the condition mumber.
Like = proj Luke
(z)
X = Le/u
Fx inf F (2),
DEN
like = inf fike (z),
AERd
C(d, q) = d KØ
for the Cholesky nonlinear,
C(d, ②) = 2k vd+1
for the mean-field nonlinear,
M
is the number of Monte Carlo samples.
Proof. Khaled & Richtárik (2023, Theorem 2) show that, if the objective function F is Lp-stooth and the stochastic gradients satisfy the ABC given as
EIJVF (a)I S A (2) F²) BIIVFI,
for some <A,B, B, c0, SGD guarantees
211 + LF (F (2n) - F*).
min 0S1ST-1
yT
Under the conditions of Corollary 1, is Lm-sth with Lp Le + Ly + Lgg Furthermore, under Assumption 5, Kim et al. (2023) show that the Monte Carlo gradient estimates satisfy
E| V (2)11, < 4L3C (d,
(2 BI
μM
2L3C (d. μM
4L3C (d. (F* lftee)
||Zjoint - Zike +
μM
This means that the ABC condition satisfied with constants
4L3
2L (d, Φ) Ejoint Enkell + μM
4L3
Cd. p),
B 1,
C= P.M
(d, p) (F* -itee)
J.M
Plugging these constants in, obtain
2L 2L (d,
min
EIVS(a))) <
(11z)ainc - Eineell 11, + 2 (F*5 (fine))
0S1ST-1
UM
T
2 yT
2 c(d.)) (F F+).
4L2
uM
Substituting the condition number yields the stated result.
34

---

Theorem 3.
Let Assumption 2 hold, the likelihood satisfy Assumption 5, and the assumptions of
Corollary 1 hold such that the ELBO F is Lp-smooth with Lp = Le L $ Ly. Then, the iterates generated by BBVI through Equation (1) and the M- M-sample reparameterization gradient include an E-stationary point such that minogisT-1 EIIVE (2.)112 S f.for any > 0 if
(F (F)) Lp (C,.C.dd,, (d,
T20
HMM
for some fixed stepsize y. where C (d. q) d + kp for the Cholesky family and C C Φ) 2kg vd 1 for the mean-field family.
Proof. As corollary to Theorem Khaled & Richtárik (2023, Corollary 1) show that, for Lp- smooth objective function F, a gradient estimator satisfying the ABC condition, an c-stationary point can be encountered if
6 LFAT L 2LpC)
12 (Ap)
12 (20)
2C
= min
T M
max B,
22
22
Under Assumption 5, Kim al. (2023) show that the Monte Carlo gradient estimates satisfy
413CC (d, +))
E| E| (2)11, 2
(F (2) - F*)) BI
μM
(d.
AL3C(d, (F* lfkee)
I - Ziikells +
μM
μM
This means that the ABC condition is satisfied with constants
4L2 WM (d.
2L C Cd. p) μM
A=
B=1 1,
C=
Enieellz +2 -ft)).
where
joint = proje (z)
is the projection of z onto set of minimizers of l the projection Σ onto set of minimizers of like
Like
proj (z))
Fa inf F (2),
ZEA
like
inf llike (z),
ZERd
C(d, p) d + kp
for the Cholesky family,
C(d, p) 2kp Vd+1
for the mean-field family,
M
the number of Monte Carlo samples.
Plugging these constants in, we obtain
48 (Ao) - (d, Zikko 12 (F+ e(ine)) MME2
12 (2())))FFF))
T>
1,
g2
MME2
- (F)
JMM
where we omitted the dependence kp and the minimizers of & and like.
35

---

F.3.2 Proximal Black-Box Variational Inference
Lemma 3 (Convex Expected Smoothness). Let e be Le-smooth and u-strongly convex with the variational family satisfying Assumption with the linear parameterization. Then,
af (2:u) - Vx. (2': < 2L C (d. By (2.2)
holds, where By (2, _ ( (2') 2') is the Bregman divergence, x L.e/Kt
is the condition number, C (d. + kp for the Cholesky family, and C (d. 2k pkpVd.
for the mean-field family.
Proof. First, have
Vf (2:u) 22 Vx (2:=0)) E (4)
(T (V))I
= =|| a (u)
aty (u) a2' g (2', u)
ad 8 (a. u)
For the linear parameterization, the Jacobian of J3 does not depend on 2. Therefore,
aT (u)
E
-8 - (g
02
and Lemma yields
Jy (u) Elg (2, 8 (A', 4)22
where
Jy (u) = + |u]
for the Cholesky family and
b (u) = + |1 for the mean-field family.
From now on, we apply the strategy of Domke (2019, Theorem 3) for resolving the randomness That is,
Ely - (V) 1e ( (u)) from the ---smoothness of f.
s L3 (4) (u) Jz. (4))
and applying Corollary 2,
s <L I - 2'1
The last step follows the approach of Kim et al. (2023), where we convert the quadratic bound into bound involving the energy. Recall that the u-strongly convexity of implies
2 k
(10)
From Lemma 8, we have
12 C (d, Φ) L - 2'll c(d,) JTz (1)
and by μ-strongly convexity,
213
(())))(((((()))))(())0))))))))-- (Ve (1) (v) - Tx (u))
1
213
f - (2::0:0(0000 (4))
213
#
Finally, by applying the equality in Lemma 10,
213 # ((1)))((0-))))--7741111)11)))
36

---

Lemma 12 (Variance Transfer). Let l be Le-smooth and u-strongly convex with the variational family satisfying Assumption 2 with the linear parameterization. Also, let VJ be M M-sample gradient estimator of the energy. Then,
tV (A) S 4L,x (,,xC(d)9)
2tr By (u.A))2uVV])2))
M
Lel is the condition number, By is the Bregman divergence defined in Lemma 3, Cd, q) d for the Cholesky fumily, and C. (d,p) Φ) = 2k Vd for the mean-field family.
Proof. First, the M-sample gradient estimator defined as
M Σ m) m=1
VJ (2)
M
where Um ~ q. Since u1 "i....ee Um are independent and identically distributed, we have
tv (2) M
Vr (2:u).
From here, given Lemma 3, the proof is identical with that of Garrigos & Gower (2023, Lemma 8,20), except for the constants.
Theorem 6. Let be -mooohhhh and j-strongly convex. Then, BBVI with proximal SGD in Equa- tion (2), M-Monte Carlo samples, variational family satisfying Assumption 2, the linear param- the iterates satisfy
eterization, and fixed stepsize 0 y 2L x(dxC(d
2yo²
FllÂT 2 (1 yH) I1A - 22112
where = Lg/p is the condition number, o² is defined in Lemma 4, 2 = arg F (2). C(d, p) = d + kpp for the Cholesky family, and C(d, p) 2k V d + for the mean-field family.
Proof. Provided that
(A.6.1) the energy f' is u-strongly convex,
(A.6.2) the energy f is -mmoohh,,
(A.6.3) the regularizer is convex,
(A.6.4) the regularizer is lower semi-continuous,
(A.6.5) the convex expected smoothness condition holds,
(A.6.6) the variance transfer condition holds, and
(A.6.7) the gradient variance o² at the optimum is finite such that σ² <
the pr identical to that of Garrigos & Gower (2023, Theorem 11.9), which is based on the results of Gorbunov et al. (2020, Corollary A.2).
In our setting,
(A.6.1) is established by Theorem 2,
(A.6.2) is established by Theorem 1,
(A.6.3) is trivially satisfied since h is the negative entropy,
(A.6.4) is trivially satisfied since h is continuous,
(A.6.5) is established Lemma 3,
(A.6.6) is established Lemma 12,
(A.6.7) is established in Lemma 4.
The only difference is that, we replace the constant Lmax the proof of Garrigos & Gower to L C (d, p) This stems from the different constants in the variance transfer condition.
37

---

Theorem 7. Let é be Lggsmmoth and u-strongly convex. Then, for any O, BBVI with proximal SGD in Equation (2), M-Monte Carlo samples, variational family satisfying Assumption 2, and
the linear parameterization guarantees EllAT
10 I
min ( μ 20 2L xc,xC(,
M
1 4o²
2x
C (d.
(2 - -2-1
T max
log
E
4²²
M
6
where Le/js o² defined in Lemma 4, 2 arg arg 2EA F (A). C(d. p) d + kp for the Cholesky family, and C(d. p) 2kp Vd 1 for the mean-field family.
Proof. This is corollary the fixed stepsize convergence guarantee in Theorem as shown by Gar- rigos & Gower (2023, Corollary 11.10). They guarantee an E-accurate solution as long as
6 2 1
2L )
T max 1 40 4²²
ma μ
log [211Ao - -
y min
2 2
In our notation, of = 0² and Lmxx = LgxC (d, P)/M.
Theorem 8. Let l be l.g-soothh and j-strongly convex. Then, BBVI with proximal SGD in Equa- tion (2), the M- M-sample reparameterization gradient estimator, variational family satisfying Assumption 2, the linear parameterization, T 4Tx, and a stepsize schedule of
M
for t 1 4Tx
2LxxC(d.0) 21+1
=
for 1 1> 4Tg-
(1+1)'H
where Tx (x x Lelu u is the condition number, C(d, = d + Kg for the Cholesky family, and C(d, p) 2kp Vd + for the mean-field family, then the iterates satisfy
El 2 201 16 12 - 2² 20
802 μ2T
e2T2
where o² defined in Lemma 4, e is Euler's 's constant, and 1 = arg min F (2).
Proof. Under our assumptions, Theorem holds, of which the proof is essentially obtaining the recursion
ElA,+1 1-1 (1- - 2- + 27?0²
Instead of fixed stepsize, we can apply the decreasing stepsize rule in the proof statement, then which the proof becomes identical to that of Gower al. (2019, Theorem 3.2). We only need to replace L with Imax in the proof of Garrigos & Gower (2023, Theorem 11.9). This, in our notation, Lamux L2xC (d, f)/M.
Theorem 4. Let e be sm-sohhh and u-strongly convex. Then, for any с 0, BBVI with proximal SGD in Equation (2), the M-sample reparameterization gradient estimator, variational family
satisfying Assumption with the linear parameterization guarantees
El Ar 201 2 < if
for t 5 4Tx
80²
4T 4T A - 20
7:
2L. xL
21+1
for 1 4Txx
max
42E
+
eVE
4Tg
(1+1)"
M- (d. where o² is defined in Lemma 4, Tx = (d.
x Lelp is the condition number; is Fuler's 's constant, X" arg argminzeN F (2), C(d, p) kp for the Cholesky family, and
C(d, 2k, + for the mean-field family.
Proof. The computational complexity follows from the smallest number of iterations T such that
16T
2
802
E||AT - 2012
+
<
c2T2
42T
38

---

By multiplying both sides with T as
80² 422
16T, [2, - 12
T
T
2 0,
(11)
e²
we can see that we are looking for the smallest positive integer that is larger than the solution of quadratic equation with respect to T. This given
8a² u²a
802
T
1
+
+
64c
6²
T 2
2c
Applying the inequality V a + S Va + Vb,
R
8d2
Tk 110-2
T 1Ã0-2
n2
μ
+ 64c
+
64c
2c
20
16o
+ Ve
8Tx
|A0-2"
4²
2E
8G2
20
2 1l2
μ22 E
evee
Thus, EE|2r 2²'1/ S e can be satisfied with number of iterations at least
802
4T₂llAo
22112
T Z max
+
4Tx
1226
ove
39

---

G
Details of Experimental Setup
Table 2: Summary of Datasets and Problems
Abbrev.
Model
Dataset
N
LME-election
1988 U.S. presidential election (Gelman Hill, 2007)
90
11,566
Linear Mixed Effects
LME-radon
U.S. household radon levels (Gelman Illll, 2007)
391
12,573
BT-tennis
Bradicy-Terry
ATP World Tour tennis
6030
172,199
LR- RR-keggu
KEGG-undireeted (Shannon er al., 2003)
31
63,608
LR-song
million songs (Bertin- Mertiin et al., 2011 buzz in social media (Kawala et al., 2013)
94
515,345
Linear Regression
LR- LR-buzz
81
583,250
LR-electric
household electric
15
2,049,280
AR-ecg
Sparse Autoregression Long-term ST ECG (Jager el al., 2003)
63
20,642,000
Linear Regression (LR-*) We consider a basic Bayesian hierarchical linear regression model
T Nr (0, 10"),
OB- N (0, a d~N; (0,
B ~ N (0. a a~ a~N(0.00))
> I + a, d²),
where a weakly informative half-normal hyperprior N4, normal distribution with the support re- stricted to 4 is assigned the hyperparameters. For the datasets, we consider large-scale regres- sion problems obtained from the UCI repository (Dua & Graff, 2017), shown in Table For all datasets, we standardize the regressors X; and the outcomes yi.
Radon Levels (MLE-radon) MLE-radon is radon level regression problem by Gelman & Hill (2007). It fits a hierarchical mixed-effects model for estimating household radon levels across dif- ferent counties while considering the floor elevation of each site. The model is described as
N4 (0, 12), da ~ N4 (o, 12), Ha ~ NV. 10").
N (0.1001))
81 N (o, 10a).
B2 . 10²)
Ma J&E
H; [<county,] | & log (uppm,) floor, $2
log radon; ~ .V (H;
which uses variable slopes and intercepts with non-centered parameterization. The dataset was ob- tained from PosteriorDB (Magnusson et al., 2022). Also, for the radon regression problem, the Minnesota subset is often used due computational reasons. Here, we use the full national dataset.
Presidential Election (MLE-election) MLE-election is a model for studying the effects of socio- logical factors on the 1988 United States presidential election (Gelman Hill, 2007). The model is described as
age N' (0, 100?),
G Gedu N (0, 100").
agexcedu ~ N (0, 100?)
N' (o, 100?),
region ~ NV' 100?),
state
N0, 02 1),
edu .r ~N, 22uu b bagexedu - (0,0 o ozgoxedu 1)
age
Nv (0. 2 I), Bregion N (0, o
state
B (0, 10021)
P: 81 + $2 black; A3 female; B4 Uprov,i + As female; black;
bagel [age]] [edu])du,] + bagexedul aage edu;] b [state,] + b [ [egionn;]
y; bernoulli (p;).
The dataset obtained from PosteriorDB (Magnusson al., 2022).
40

---

Bradley-Terry (BT- Tennis) BT-Tennis is Bradley-Terry model for estimating the skill of pro- fessional tennis players used by Giordano et al. (2023). The model described
N (0,1)
0 N'' (0,
Pi 0[win;] - 0[1ss;
yi ~ bernoulli (p),
where win, los, are the indices of the winning and losing players for the ith game, respectively, While we subsample over the games i = ,,,,,NN each player's involvement is sparse in that each player plays only handful of games. Consequently, the subsampling noise is substantial. Therefore,
we use a larger batch size of 500. Similarly to Giordano et al. (2023), we use the ATP World Tour data publically available online
Autoregression (AR-ecg) AR-ecg is a linear autoregressive model. Here, we use a Student-t like- lihood as originally proposed by Christmas Everson (2011). While they originally imposed an automatic relevance detection prior the autoregressive coefficients, shrinkage prior (Carvalho el al., 2009, 2010). Since the horseshoe is known to result in complex posterior geometry, this should make the problem more challenging. The model is described
instead set horseshoe
ad = 10-?
=== = 10
ad a Ba = 107
d gamma (ca.Pa)
σ inverse-gamma (a,, fa),
t ~ cauchy, (0, 1),
À ~ cauchy, (0, 1).
0 ~ N' diag (2))
y[n] stduent-t 0ny[n 1] 02y(n - 02y[n-2] 2]
where d is the degrees--offfreedom for the Student-t likelihood, cauchy, is half-Cauchy prior.
For the dataset, use the long-term electrocardiogram measurements of Jager al. (2003) obtained from Physionet (Goldberger et al., 2000). The data instance we used has duration of 23 hours sampled at 250 Hz with 12-bit resolution over range of +10 millivolts. During the experiments, we observed that the hyperparameters suggested by Christmas & Everson are sensitive to the signal amplitude. Therefore, scaled the signal amplitude to be +10.
tennis-aan
41

---

H
Additional Experimental Results
6000
-6000
T = 5,000
10
10
0= 10
ELBO
-8000
ELBO
-8000
-1x10*
-1x10*
-1,2x10*
-12x10
10 10 0.01 0.1
25k
50k
25k
50k
1
25k
50k
Base Stepsize (u)
lteration (T)
(a)
LME-election
-1.7x10'
-1.7x10*
T- 5,000
10
10
a 10
ELBO
-1.75x10*
-1.75x101
-1.8<10*
-1.8x10*
10 10"
0.01
0.1
25k
50k 1
25k
50k
25k
50k
Base Stepsize (u)
Iteration (T)
(b) LME-radon
-1x10'
-1x10*
T 5,000
10
10
10
-1.1x10*
-1.1x10"
ELLB
ELBO
-1.2x10*
-1.2x10
-1.3x10
1.3x10
10
50k
25k Iteration
25k
50k
Base Stepsize (a)
(c) BT-tennis
10y
10y
T = 5,000
a 10
a = 10
a= 102
04
ELBO
ELBO
-1x10'
-1x10'
-2x10
-2x10*
-3x10"
-3x10"
10
10'3
0.01
0.1
25k
50k
25k
50k
25k
50k
Base Stepsize (c)
Iteration (T)
(d) LR-keggu
Figure 5: BBVI convergence speed (ELBO v.s. Iteration) and robustness against stepsize (ELBO at T = 50, 000 v.s. Base stepsize). The error bands are the 80% quantiles estimated from 20 (10 for eR-eeg) independent replications. The initial point was m. = 0, C. = I.
42

---

0
T- 5,000
a- 10
10
Q
10
ELBO
-1x100
ELBO
=x100
-2x10
-2x10
-3x100
-3x10
10" 10"
0.01
0.1
25k
50k
25k
50k
25k
50k
Base Stepsize (o))
Iteration (T)
(a) LR-song
T- 5,000
10
10
10
-1x100
-1x10
ELBO
ELBO
-2x10
-2x10
-3x10"
-3x10"
10 10
0.01 0.1
25k
50k
25k
50k
25k
50k
Base Stepsize (o))
Iteration (T)
(b)
LR-buzz
0
0
T 5,000
a
10
02
10
-1x10"
-1x
ELBO
-2x10"
-2x10
3x100
-3x10"
10
10
0.01 0.1
25k
50k
25k
50k
25k
50k
Base Stepsize (ct)
Iteration (T)
(c) LR-electric
T- 50,000
10"
10"
10
-1.9x10*
-1.9x10*
ELBO
ELBO
-1.95x10*
-1.95x10"
-2x10
-2x10
10
1
0.01 0.1
25k
50k
25k
50k
25k
50k
Base Stepsize (()
Iteration (T)
(d) AR-ecg
Figure BBVI convergence speed (ELBO v.S. Iteration) and robustness against stepsize (ELBO at T = 50, 000 v.s. Base stepsize). The error bands are the 80% quantiles estimated from 20 (10 for AR-eeg) independent replications. The initial point was mo 0, Co I.
43


Reviews:
[{"review_id": "dgNZrcrqiJ", "reviewer": "Reviewer_gnfU", "summary": "The paper establishes the first complete convergence result of BBVI as it is used in practice. It is shown that a linear parametrization of the covariance leads to better convergence rates, which is confirmed in some numerical experiments. Finally, a proximal version for BBVI due to Domke et al 2020 is analyzed and shown to perform favorably in experiments to Adam.", "questions": "1) When using the linear parametrization, could there be an issue that the scale parameter becomes negative? Or is this somehow handled by the proximal operator of the entropy?\n\n2) Do the proposed convergence results also carry over to ProxGen? Perhaps it is out of scope for this paper, but maybe a comment in the experiments section could be helpful.\n\n3) How well does the method work for modern neural networks  (ResNets, transformers)? There is many claims in the community that variational inference doesn't work well in these settings (e.g. worse performance than MAP inference). It would be highly interesting to see whether the proximal scheme solves these problems.", "limitations": "All limitations are addressed.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "- The paper is well-written and clearly a lot of time was spent by the authors to make the results look clean and polished. \n- The proposed proximal scheme for BBVI is implemented in the probabilistic programming language Turing, which may be of interest to the community once the code is released.\n", "weaknesses": "1) The results and experiments are incremental (but nontrivial) extensions of what is presented in (Domke 2020).\n\n2) The paper missed to cite the recent line of works\n* https://arxiv.org/abs/2205.15902\n* https://proceedings.mlr.press/v202/diao23a.html\n* https://proceedings.mlr.press/v202/diao23a.html\nwhich give very strong convergence results for a BBVI-like algorithm by interpreting it as a Wasserstein gradient flow.  A theoretical (and maybe even practical) comparison to these recent works would make the paper stronger. \n\n3) It is known that BBVI is inferior to natural-gradient algorithms which exploit the KL / Fisher-geometry (https://arxiv.org/abs/2107.04562).  This line of works could be mentioned in the introduction.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review.\n\n> The paper missed to cite the recent line of works\n\nThank you for the suggestions. We will add them to the next version.\n\n> It is known that BBVI is inferior to natural-gradient algorithms which exploit the KL / Fisher-geometry (https://arxiv.org/abs/2107.04562).  This line of works could be mentioned in the introduction.\n\nGiven the importance of NGVI, we agree to include it in the introduction and discussion sections. However, we would note that BBVI is a broader algorithm in the sense that one can use variational distributions outside of the exponential family, which our theory does include, and amortized variational families, which we plan on working on next. \n\n> When using the linear parametrization, could there be an issue that the scale parameter becomes negative? Or is this somehow handled by the proximal operator of the entropy?\n\nYes, the proximal operator ensures that the scale is never negative. Furthermore, using projection operators as suggested by Domke (2020) is also an effective way of ensuring this. In practice, however, choosing a stable initial point (large scale) with a small enough step size appears to be enough to ensure convergence, though there is still room for investigation on whether it is a robust enough choice.\n\n> Do the proposed convergence results also carry over to ProxGen? Perhaps it is out of scope for this paper, but maybe a comment in the experiments section could be helpful.\n\nUnfortunately, ProxGen assumes the problematic “bounded gradient variance” assumption, thus our results cannot be immediately applied. It might be possible to extend their proof to include our setting. As suggested, we will add a discussion about this.\n\n> How well does the method work for modern neural networks (ResNets, transformers)? There is many claims in the community that variational inference doesn't work well in these settings (e.g. worse performance than MAP inference). It would be highly interesting to see whether the proximal scheme solves these problems.\n\nGiven the theoretical and empirical results on “pruning” (Trippe et al., 2018; Coker et al., 2022; Huix et al., 2022), we conjecture that the proximal scheme will not fix the problem. Exclusive KL may simply be inappropriate for deep models, although more investigation is certainly needed.\n\n### References\nTrippe, Brian, & Turner, Richard. 2017. Overpruning in Variational Bayesian Neural Networks. Tech. rept. arXiv:1801.06230. ArXiv.\n\nCoker, Beau, Bruinsma, Wessel P., Burt, David R., Pan, Weiwei, & Doshi-Velez, Finale. 2022. Wide Mean-Field Bayesian Neural Networks Ignore the Data. Pages 5276–5333 of: Proceedings of the International Conference on Artificial Intelligence and Statistics. PMLR.\n\nHuix, Tom, Majewski, Szymon, Durmus, Alain, Moulines, Eric, & Korba, Anna. 2022 (July). Variational Inference of Overparameterized Bayesian Neural Networks: A Theoretical and Empirical Study. Tech. rept. arXiv:2207.03859. ArXiv\n"}, {"review_id": "y02tpNL6Gv", "reviewer": "Reviewer_C3v1", "summary": "This paper tries to prove convergence of BBVI algorithm in a more general setting compared to previous work: the main being that the target needing to be only log-smooth and not needing to be log-concave, which means that the objective can be non-convex. The paper claims that non-linear parameterizations of the variational parameters (the mean, the cholesky factor L for full rank matrix and scale for mean field), often used in practice can break strong convexity, even if the target is log-concave and thus plain SGD with non-linear conditioning is sub-optimal while proximal SGD with linear parameterization can give the fastest convergence for BBVI.", "questions": "1. If I understood correctly, the object of treatment in this paper is the objective function and not the gradient, both Theorems 1 and 2 are for the objective rather, I wanted to know why did the authors concentrate on objective and not the gradient itself.\n2. Is it expected that the gradients will follow the same behaviour as objective which is the main object of treatment in this work, especially when there is a non-linear transformation involved in the parameterization, i.e. $ \\nabla_{\\lambda}f(\\mathbf{t}_{\\lambda}(\\mathbf{u})) $ is different from $\\nabla f$ ?  \n3. Does the work cover the case when gradients are computed with mini-batching which brings another source of stochasticity in addition to randomness of base distribution ? \n4. Maybe the authors can also address, how their work is different/similar to the one in 'Provable convergence guarantees for black-box\nvariational inference' on arxiv, how their claims match to it ? \n5. Why was the exp conditioner/transformation not used in expts. for Figure 4, when it is the most popular one in my opinion ? \n6. I did not like this statement so much: 'This would put mean-field BBVI in a regime comparable to the best-known convergence rates in\napproximate MCMC' since MCMC is aymptotically exact while MFVI is not and can be arbitrarily different from target.", "limitations": "\nMinor\n1. Typo in line 115:'inifinite'\n2. For the uninitiated, maybe introduce $L_{h}$ as Lipschitz smooth, if it does .. ?\n3. Maybe explain what do all the constants $L_{f}, L_{h}, L_{s}$ mean at the beginning of Sec 3. I guess $L_{f}, L_{h}$ mean smothness constants for likelihood/energy part and entropy part in the objective respectively,\n\nI am overall inclined in favour of accepting the paper given some of the questions and concerns are addressed by the authors. I need also a bit more time in analysing this work and recent similar work, where some of the claims are similar but some are different.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is solid and covers a lot of theory behind VI, Figure 1 is quite informative.\n2. The figures are great and descriptive and support the theory and text.\n3. I am not good with theory, but I found the theorems and inequalities fine \n4. Covers relevant and contemporary literature quite well, although a paper with similar ideas and claims recently appeared on arxiv. \n", "weaknesses": "1. I think, the paper is trying to unify many different parameterizations, and so at times the paper has become a bit harder to read, and in some places, authors can motivate more why and how certain theorems can have pratical impact, where things are generally followed and discussion is only theoretical and where theorems have a practical impact. \n2. Then my main concern with this work, is that its claims can be easily misinterpreted, when readers see: 'BBVI converges', especially when there is no general consensus of what black box means, I think BBVI was first used as an acronym by the Ranganath 2013 paper which actually used score gradients. This nomenclature of algorithms is unfortunate because many new readers also think of Neural Networks powered inference algorithms, when they hear Black-box VI and assume Normalizing flows and Neural Network based Variational inference also as BBVI. With the recent explosion of methods, where a user can practically choose any divergence measure(any member of the f-divergence family) and any approximating family and use MC samples and auto diff to estimate gradients for inference, while this paper is narrow in the scope that it only considers exclusive KL, well behaved log-concave Gaussian (loc-scale) variational family and RP gradients, it will be unfair to say that 'BBVI converges', and therefore I strongly suggest authors to change the title of the paper. \n3. In results section, it seems a bit odd and unsettling to me that the linear parameterization is claimed to be better than non-linear in all cases. Another submission in this conference, claims that non-linear(log) parameterization is theoretically more stable (also in Domke 2019), and shows some practical experiments confirming the same. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review.\n\n> I think, the paper is trying to unify many different parameterizations, and so at times the paper has become a bit harder to read\n\nWe agree that the paper is quite dry in the submitted form. Due to the page limit, we had to remove many of our results' commentaries. We will add more practical discussions of our results in the next version.\n\n> Then my main concern with this work, is that its claims can be easily misinterpreted, when readers see: 'BBVI converges', especially when there is no general consensus of what black box means, I think BBVI ... actually used score gradients. \n\nIn our defense, much recent work on BBVI with Monte Carlo gradients does seem to use “BBVI” synonymously. Recent papers such as the work of Domke (2019,2020), Hoffman and Ma (2020), Regier et al. (2017), Manushi et al. (2022) all use the term BBVI despite none of these works considering the score gradient. We are happy to use more general terms such as Monte Carlo VI and clarify this in the abstract. However, we’d also like to find some middle ground here and not deviate substantially from conventional technical vocabulary norms set in the recent literature,\n\n>  I strongly suggest authors to change the title of the paper.\n\nWe agree that the title might be perceived overly broadly. We propose to change the title to: “On the Convergence and Scale Parameterization of Black-Box Variational Inference.” Please let us know if the reviewer has better suggestions.\n\n> In results section, it seems a bit odd and unsettling to me that the linear parameterization is claimed to be better than non-linear in all cases. Another submission in this conference, claims that non-linear(log) parameterization is theoretically more stable (also in Domke 2019) \n\nSome of the authors have also seen that paper during bidding, although we couldn’t read it (as we didn’t bid for it due to the apparent conflict of interest). Thus, unfortunately, it would be impossible and inappropriate for us to comment on their results. However, Thm 2 is quite clear; you lose strong convexity with non-linear parameterizations, and thus one loses the guarantee to converge fast. To us, it was thus unsurprising that the linear param. did well in the experiments. Furthermore, we are unaware that Domke (2019) discusses non-linear params. His recent paper ($\\S$ 5; Domke et al., 2023) agrees with our theoretical analysis: “*An alternate parameterization would also have implications for the (strong) convexity of the ELBO.*”\n\nIn practice, the performance strongly depends on the choice of initialization. So this might explain the different conclusions. For instance, when initialized in the non-smooth regime (small initial scale), the linear parameterization starts by diverging (visualized in Domke 2020 Fig 1), resulting in slow convergence. However, as shown in Fig 3 in our paper and Domke (2020), proximal SGD fixes this sensitivity.\n\n> If I understood correctly, the object of treatment in this paper is the objective function and not the gradient ... I wanted to know why did the authors concentrate on objective and not the gradient itself.\n\nFor studying the properties of SGD, one needs to analyze both the objective and gradient estimator. In particular, the landscape of the objective (smoothness, convexity) often determines the convergence rate. We thus focus on the properties of the ELBO.\n\n> Is it expected that the gradients will follow the same behaviour as objective which is the main object of treatment in this work?\n\nWe kindly request clarification for this question.\n\n> Does the work cover the case when gradients are computed with mini-batching ... ?\n\nYes, this should be fairly trivial to do whenever the linearity of the variance is used. We will add a discussion on how to do this in the future version. We also found that more insight could be drawn in the doubly stochastic setting, which we are attempting to study in a follow-up paper.\n\n> Maybe the authors can also address, how their work is different/similar to the one in 'Provable convergence guarantees for black-box variational inference' on arxiv, how their claims match to it ?\n\nWe indeed saw that paper. Before we comment on the paper at all, note that the paper is concurrent work (potentially a concurrent NeurIPS submission), which we need to tread carefully about considering this paper here, as the NeurIPS FAQ suggests we should generally not do (see the “policy on comparisons to recent work”).\n\nWhile we emphasize the policy above, we are willing to contrast the two. Our paper was more interested in the effect of parameterizations, while that paper was more focused on proving the convergence of SGD with quadratic variance estimators, which were studied in Domke (2019). The only overlap is for the convergence of proximal SGD with a decreasing stepsize for strongly convex ELBO, for which the complexity matches perfectly: $\\mathcal{O}(\\kappa^2 \\frac{1}{\\epsilon})$.\n\n> Why was the exp conditioner/transformation not used ... ?\n\nTo our knowledge, the softplus conditioner is also popularly used (see Tab 1 of Kim et al. 2023.) Furthermore, from the conclusions of Thm 2, we expected the performance of the softplus to be representative of all nonlinear conditioners.\n\n> I did not like this statement so much: 'This would put mean-field BBVI in a regime comparable to ... approximate MCMC' since MCMC is aymptotically exact while MFVI is not and can be arbitrarily different from target.\n\nHere, we’d like to emphasize the use of the word approximate, and will make this clearer in the text. The fastest convergence rates in MCMC are achieved with “approximate” MCMC methods, which are asymptotically biased, such as unadjusted Langevin. The cited paper indeed discusses the unadjusted Langevin algorithm, which is an approximate MCMC method.\n\n> For the uninitiated, maybe introduce Lipschitz smooth ...\n\nThanks for the suggestion. We will add more discussion on these assumptions in the future version.\n"}, {"review_id": "eB76gfl8ma", "reviewer": "Reviewer_MsbX", "summary": "The authors analyze the smoothness and convexity of the ELBO under different parameterizations (linear vs nonlinear) of the scale for location-scale families, building on the work of Domke (2020). This enables convergence analysis for BBVI with standard and proximal stochastic gradient descent. Their main findings are that 1) the energy term of the ELBO is smooth under certain conditions (Theorem 1, ii) on the diagonal conditioner $\\phi$ and reparameterization gradient, and 2) nonlinear diagonal conditioners break (strong) convexity of the energy, affecting SGD convergence rates, and 3) establishing rates for obtaining $\\epsilon$-optimal solutions in BBVI for both convex and nonconvex energy terms.\n", "questions": "* The summary of contributions (lines 51-58) seems overly informal to me. What is a “full” guarantee? What converges? “Precisely as used in practice” may mean different things to different people. “Suboptimal” in what sense (item #2)? Maybe more precision would be better here. \n* The formatting of lines 67-70 is odd, and perhaps these quantities should just be defined in-line.\n* In line 102, I believe an additional “diag” is missing: should it be $diag(\\phi(s)) = diag(\\phi(s_1),\\dots, \\phi(s_d))$?\n* Line 115, “inifinite” is a typo. This taxonomy discussion I find detrimental to the paper, as it creates more confusion than clarity; Figure 1 and its caption could be removed altogether.\n* Is the distinction between the LHS and RHS of the equation after line 127 that of a “total derivative” vs a partial derivative with respect to $\\lambda$?\n* Line 130-132 has a typo in phrasing “have been”, “alienating”\n* The proof of Theorem 1 seems to imply that $f$, in addition to being smooth, is twice-differentiable to compute the Hessian. If this isn’t implied by some other conditions, it should be stated explicitly perhaps for completeness.\n* The statement of Theorem 2 differs in the main paper and the supplement, this should be corrected.\n* Line 174, “becomes” is a typo\n", "limitations": "The theoretical contributions make clear the assumptions that must be met for the results to hold. As stated, the 1-Lipschitz condition for $\\phi$ could be discussed more.\n", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "* The work extends previous results meaningfully: Domke (2020) considered the case of a linear diagonal conditioner, while this work focuses on nonlinear ones while including the linear case for completeness. \n* Beyond showing convergence merely occurs (per the title), rates are established both for standard SGD iterates (Theorem 3) and proximal SGD in the case of a convex $f$.\n* The variational family considered (location-scale) encompasses a wide range of distributions usually used for VI; one additional step of  generalizability is the use of the bijector $\\psi$ to make the work relevant to ADVI in general.\n* The work has theoretical and practical merit, advising users on the tradeoffs between linear and nonlinear diagonal conditioners along with providing clearly established lemmas, assumptions and proofs that further work might build upon. \n", "weaknesses": "* The story is a bit disjointed. Assumption 4 and Theorem 3 might be better placed directly after Example 4 to complete a story about convergence of BBVI for nonlinear $\\phi$ and general nonconvex $f$ that satisfies the smoothness conditions (among others). This might have even been a logical stopping point for a paper with the given title. Theorem 2 seems better placed nearby the section on Proximal SGD, as it is precisely the convexity of $f$ that Section 3.4 requires.\n* The exponential diagonal conditioner (line 108-109), commonly used in practice,  does not satisfy the 1-Lipschitz assumption (Theorem 1, (ii)). I put this as a weakness, but maybe this point should be emphasized instead and used to evangelize the merits of softplus, which evidently does satisfy this assumption. \n* The case where $f$, which is usually $-\\log p(x,z)$ in this work when $\\psi$ is the identity, is convex in $z$ is uncommon in many practical problems of significance for VI when the density $p(x,z)$ is multimodal. This diminishes the significance of Theorem 2 (breaking convexity) and Section 3.4.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review.\n\n> The story is a bit disjointed. Assumption 4 and Theorem 3 might be better placed directly after Example 4 to complete a story about convergence of BBVI for nonlinear \n and general nonconvex $f$ that satisfies the smoothness conditions (among others). This might have even been a logical stopping point for a paper with the given title. Theorem 2 seems better placed nearby the section on Proximal SGD, as it is precisely the convexity of \n that Section 3.4 requires.\n\nThank you for the suggestions. We will reorganize our paper in the next version accordingly.\n\n> The exponential diagonal conditioner (line 108-109), commonly used in practice, does not satisfy the 1-Lipschitz assumption (Theorem 1, (ii)). I put this as a weakness, but maybe this point should be emphasized instead and used to evangelize the merits of softplus, which evidently does satisfy this assumption.\n\nWe point out that Theorem 2 does say something about non-Lipschitz conditioners; the exp conditioner will also result in a slower convergence rate since it breaks strong convexity. Given the breadth of this result, we thought it didn’t matter much which conditioner is used; any non-linear conditioner is bad. Nevertheless, we will add the suggested comment in the limitations.\n\n> The case where $f$, which is usually $- \\log p(x, z)$ in this work when $\\psi$ is the identity, is convex in $z$ is uncommon in many practical problems of significance for VI when the density $p(x, z)$ is multimodal. This diminishes the significance of Theorem 2 (breaking convexity) and Section 3.4.\n\nWe believe that the significance of our theory should be taken qualitatively rather than quantitatively. Our results say that when the posterior is “easy” such that the landscape is Gaussian, nonlinear conditioners will fail to take advantage of this, unlike linear conditioners. If we extrapolate this intuition to posteriors that are non-Gaussian but not too much, it is still conceivable that linear conditioners will result in faster convergence. It is also known that, for large datasets, many Bayesian posteriors become Gaussian as a consequence of the Bernstein von Mises theorem. Thus we believe our results do explain the performance difference of our experiments, where none of the problems are provably Gaussian (strongly log-concave.)\n\n> Is the distinction between the LHS and RHS of the equation after line 127 that of a “total derivative” vs a partial derivative with respect to $\\lambda$\n\nSorry for the confusion. We should have added a written description. The RHS is derivative with respect to $\\lambda$, while the LHS is the composition of $t_{\\lambda}(u)$ and the gradient of $f$ denoted as $\\nabla f$. That is, $\\nabla f \\circ t_{\\lambda}(u)$.\n\n> The summary of contributions (lines 51-58) seems overly informal to me. What is a “full” guarantee? What converges? “Precisely as used in practice” may mean different things to different people. “Suboptimal” in what sense (item #2)? Maybe more precision would be better here.\n\nThank you for the comments. We will be more precise in the contribution summary and state that we do not assume any unrealistic modifications to the algorithm, such as bounded domain or bounded gradient. For the “suboptimal,” we will state that the nonlinear conditioners break strong convexity.\n\n> The proof of Theorem 1 seems to imply that $f$, in addition to being smooth, is twice-differentiable to compute the Hessian. If this isn’t implied by some other conditions, it should be stated explicitly perhaps for completeness.\n\nThank you very much for catching this! Indeed we missed that assumption. We will add it in the next version.\n\n> The statement of Theorem 2 differs in the main paper and the supplement, this should be corrected.\n\nSorry for the confusion. We will fix this in the next version."}, {"review_id": "U4wSGW9BSh", "reviewer": "Reviewer_TwpX", "summary": "The paper proves convergence results for black-box variational inference (BBVI) with ordinary stochastic gradient descent (SGD) and proximal SGD under several assumptions: it considers only the reparameterization gradient setting with the location-scale variational family (in particular, mean-field and Cholesky parameterization) and a symmetric base distribution (with other mild requirements); it also requires an assumption involving the reparameterized gradient and the diagonal conditioner of the scale matrix and another one about the growth of the bijector. As an improvement over many previous works, it does not need to assume bounded domain, bounded support, or regularity conditions about the evidence lower bound (ELBO) directly.\nThe paper notes that the theoretical results show that for the location-scale family, nonlinear scale parameterizations are suboptimal (but widely used in practice).\nThis claim is tested experimentally on a synthetic and realistic problems. The empirical results confirm this claim and suggest that proximal methods converge faster.", "questions": "- What is meant by the \"infinite sum structure\" in line 115?\n- Theorem 3 claims that the iterates generated by BBVI include an $\\epsilon$-stationary point. I would have expected a statement of the form \"the iterates will be $\\epsilon$-stationary points from some point onwards\". Does that mean that BBVI will usually move away from the stationary point again?\n- In Theorem 3, $\\kappa$ is defined but not used, $\\gamma$ and $M$ are used but not defined (at least I failed to find their definition) and $C$ takes sometimes one and sometimes two arguments. Could you clarify the statement of the theorem?\n- Could you define the \"proximal operator\" (line 209)? It is defined for the specific setting, but not what it means in general. Some intuition would also be helpful.\n\nTypos:\n- line 3: comma should go after the parentheses\n- line 14: \"box\" missing\n- line 60: did you mean \"sans-serif\"?\n- line 212: remove \"this\"\n- line 230: Theorem 12 is not in the main part of the paper.", "limitations": "Limitations (in the form of assumptions) are mentioned in the paper, but the list is incomplete (see \"Weaknesses\"). It would be good if all limitations were collected in one place in the paper.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper improves on previous convergence results on BBVI by lifting assumptions needed in previous work. It even goes further by empirically testing the theoretical prediction of linear scale parameterization being superior. Therefore it is interesting and relevant to the variational inference community.\n\nThe presentation of the paper is mostly clear and clean despite the technical content. In particular, highlighting the assumptions in the text is very helpful for the reader.\n\nThere is a lot of information about the experimental setup in the appendix and the code is available, which is good for reproducibility.", "weaknesses": "The main weakness of the paper is how it deals with the limitations: while it is good that the assumptions are highlighted in the text, there are more limitations that should be pointed out. For example, the paper only deals with reparameterization gradient BBVI, which is not mentioned in the abstract. In a similar vein, the title oversells the paper by a lot, making it seem as though it solved a much more general problem. (If the next paper generalizes these results or lifts some of the assumptions, should they give it the same title?) The paper also lacks a limitations section.\n\nAnother weakness is that the experimental evaluation (for the most part) uses the Adam optimizer (and a proximal version) even though the theoretical results concern standard and proximal SGD. Furthermore, it explores the combinations Adam&linear, Adam&nonlinear, ProxAdam&linear, but not ProxAdam&nonlinear. I think it would be good to report the results of this combination too, in order to get a more complete picture.\n\nFinally, the results are quite technical, with a lot of variables involved in the bounds and there seem to be some errors in Theorem 3 (see my question to the authors below). It would be good the collect all the variable names and their meaning (or the position of their first occurrence in the text) in one table, so the reader can look up the meaning in one place. As a minor point, it would also help if expected values could include brackets around their arguments, to clarify their scope.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for review.\n\n> The main weakness of the paper is how it deals with the limitations: while it is good that the assumptions are highlighted in the text, there are more limitations that should be pointed out. For example, the paper only deals with reparameterization gradient BBVI, which is not mentioned in the abstract.  In a similar vein, the title oversells the paper by a lot, making it seem as though it solved a much more general problem. (If the next paper generalizes these results or lifts some of the assumptions, should they give it the same title?) \n\nThank you for the suggestion. We will add this point in the abstract and introduction. We have used of BBVI to mean “BBVI with reparameterization gradients” primarily because we are following recent work on the topic e.g., Domke (2019,2020), Kim et al (2023), Hoffman and Ma (2020), Regier et al. (2017), none of which use score gradients. Nevertheless, we agree that reducing the scope of the title would be more appropriate. We thus propose to change the title to: “On the Convergence and Scale Parameterization of Black-Box Variational Inference.” Please let us know if the reviewer has better suggestions.\n\n> The paper also lacks a limitations section.\n\nWe originally had a separate limitation section, but had to remove it due to space issues. We will add it back in the next version.\n\n> Another weakness is that the experimental evaluation (for the most part) uses the Adam optimizer (and a proximal version) even though the theoretical results concern standard and proximal SGD. Furthermore, it explores the combinations Adam&linear, Adam&nonlinear, ProxAdam&linear, but not ProxAdam&nonlinear. I think it would be good to report the results of this combination too, in order to get a more complete picture.\n\nWe didn’t include ProxAdam&nonlinear because proximal SGD removes the need to use nonlinear conditioners. We could have included it, but didn’t see the point.\n\n> It would be good the collect all the variable names and their meaning (or the position of their first occurrence in the text) in one table, so the reader can look up the meaning in one place. As a minor point, it would also help if expected values could include brackets around their arguments, to clarify their scope.\n\nThanks for the suggestion. We will add a separate nomenclature section in the future version.\n\n> What is meant by the \"infinite sum structure\" in line 115?\n\nWe agree that we should have explained this. In the SGD literature, “finite sum” describes an objective function that can be represented as an infinite sum of functions, e.g. $\\frac{1}{N} \\sum_{n=1}^N f_n(x)$. In contrast, the VI objective cannot be represented as such an objective, as it is an expectation over an infinite number of points. To contrast with the finite sum setting, this setup has been called the “infinite sum.” We will add an explanation and citations to this in the next version.\n\n> Theorem 3 claims that the iterates generated by BBVI include an \n-stationary point. I would have expected a statement of the form \"the iterates will be \n-stationary points from some point onwards\". Does that mean that BBVI will usually move away from the stationary point again?\n\nThis is a limitation of our current understanding of SGD in the nonconvex smooth setting in general, not a problem specific to BBVI. Most of the known analyses prove that the average gradient norm over the iterates is bounded. This in turn implies that $ \\min_t || \\nabla F (x_t) || < \\mathbb{E}_t || \\nabla F (x_t) || $. Thus one can only say that the trajectory includes a stationary point (the point achieving the minimum norm.)\n\n> In Theorem 3, gamma is defined but not used, kappa and M are used but not defined (at least I failed to find their definition) and C takes sometimes one and sometimes two arguments. Could you clarify the statement of the theorem?\n\nThanks for pointing out that $\\kappa$ is not used. Here $gamma$ is the step size, $M$ is the number of samples, and $C$ is used as a general placeholder for arbitrary constants. As suggested before, we will add a separate nomenclature table to resolve the confusion.\n\n> Could you define the \"proximal operator\" (line 209)? It is defined for the specific setting, but not what it means in general. Some intuition would also be helpful.\n\nThank you for the suggestion. Yes, we should have included an explanation.\n\nIn general, a proximal operator (or proximity operator) is a tool commonly used in optimization. For a convex function $h$ it is generally described as $\\mathrm{prox}_{\\gamma, h}(x) = \\arg\\min_z \\frac12 \\| x - z \\|^2 + \\gamma h(z)$, where $\\gamma$ is the stepsize. We want to find a point close to $x$ and simultaneously minimize the function $h$. While it might not be obvious, this operator is a generalization of the gradient descent update $x_t - \\gamma g_t$, which can be obtained by setting $h$ to be the first-order linearization of the objective function. \n\nA typical use-case of proximal operators is when $h$ is non-smooth or even non-differentiable. In those cases, the proximal operator may have a closed-form expression and favorable convergence properties than, say, subgradient descent. In our case, the proximal operator is used to circumvent the fact that the entropy term is non-smooth.\n\nWe will add this explanation in the next version.\n\n\n"}]

Author rebuttal:

Meta-review: All reviewers agreed this paper would make a solid contribution to NeurIPS. In the camera ready, please carefully go through the reviews, your rebuttal, and the discussion, and update the paper accordingly. Some particular notes:

- include noted papers in the literature review
- edit the title to avoid misleading claims
- there was some concurrent work brought up during the review process (I noted that one reviewer brought up a paper with seemingly contradictory claims to the present paper, and another that may have overlap); please include a discussion of these works if reasonable to do so, with a mention that they are concurrent. Note that the neurips policy encourages comparison to concurrent work in the camera ready, and I think it would be beneficial for readers in this instance.
- a comprehensive limitations section
- some suggested reorganizations of the text/results
Decision: Accept (poster)