Variational inference tries to approximate a complex target distribution with a distribution from a simpler family [1,2,3,4]. If p(z, x) is a target distribution with latent variables z ∈ R d and observed data x ∈ R n , and q w is a variational family with parameters w W ∈ , the goal is to minimize the negative evidence lower bound (ELBO)
.
(
For subsequent discussion, we have decomposed the objective into the free energy l and the negative entropy h . Minimizing f is equivalent to minimizing the KL-divergence between q w to p(•|x), because KL (qw ∥p(•|x)) = f (w) + log p(x) . Recent research has focused on "black box" variational inference, where the target distribution p is sufficiently complex that one can only access it through evaluating probabilities (or gradients) at chosen points [5,6,7,8,9,10,11]. Crucially, one can still get stochastic estimates of the variational objective f and of its gradient, and use these to optimize.
Still, variational inference can sometimes be unreliable [12,13,14,15], and some basic questions remain unanswered. Most notably: does stochastic optimization of f converge to a minimum of the objective? There has been various progress towards answering this question. One line of research seeks to determine if the variational objective f has favorable structural properties, such as smoothness or (strong) convexity [13,16,17] (Sec. 2.1). Another line seeks to control the "noise" (variance or expected squared norm) of different gradient estimators [18,19,20,21] (Sec. 2.2). However, few full convergence guarantees are known. That is, there are few known cases where applying a given stochastic optimization algorithm to a given target distribution is known to converge at a given rate.
We identify two fundamental barriers preventing from analysing this VI problem as a standard stochastic optimization problem. First, the gradient noise depends on the parameters w in a nonstandard way (Sec. 2.3). This adds great technical complexity and renders many traditional stochastic optimization proofs inapplicable. Second, stochastic optimization theory typically requires that the (exact) objective function f is Lipschitz continuous or Lipschitz smooth. But in our VI setting, under some fairly benign assumptions, the ELBO is neither Lipschitz continuous nor Lipschitz smooth.
We obtain non-asymptotic convergence guarantees for this problem, under simple assumptions.
Central contributions (Informal). Suppose that the target modellog p(•, x)is concave and Lipschitz-smooth, and that q w is in a Gaussian variational family parameterized by the mean and a factor of the covariance matrix (Eq. 2). Consider minimizing the negative ELBO f using either one of the two following algorithms:
• a proximal stochastic gradient method, with the proximal step applied to h and the gradient step applied to l , estimating ∇l(w) with a standard reparameterization gradient estimator (Eq. 5), and using a triangular covariance factor;
• a projected stochastic gradient method, with the gradient applied to f = l + h , estimating ∇f(w) using either of two common gradient estimators (Eq. 7 or Eq. 9), with the projection done over symmetric and non-degenerate (Eq. 3) covariance factors
Then, both algorithms converge with a 1/ √ T complexity rate (Cor. 12), or 1/T if we further assume that log p(•, x)is strongly concave (Cor. 13).
We also give a new bound on the noise of the "sticking the landing" gradient estimator, which leads to faster convergence when the target distribution p is closer to Gaussian, up to exponentially fast convergence when it is exactly Gaussian (Cor. 14).
This is achieved through a series of steps, that we summarize below.
1. We analyze the structural properties of the problem. Existing results show that with a Gaussian variational family, if -log p(•, x) is (strongly) convex or Lipschitz smooth, then so is the free energy l . This is for instance known to be the case for some generalized linear models, and we give a new proof of convexity and smoothness for some hierarchical models including hierarchical logistic regression (see Appendix 7.3). The remaining component of the ELBO, the neg-entropy h , is convex when restricted to an appropriate set. It is not smooth, but it was recently proved to be smooth over a certain non-degeneracy set.
2. We study the noise of three common gradient estimators. They do not satisfy usual noise bounds, but we show that they all satisfy a new quadratic bound (Definition 5). For the sticking-thelanding estimator, our bound formalizes the longstanding intuition that it should have lower noise when the variational approximation is strong (Thm. 4).
3. We identify and solve the key optimization challenges posed by the above issues via new convergence results for the proximal and projected stochastic gradient methods, when applied to objectives that are smooth (but not uniformly smooth) and with gradient estimators satisfying our quadratic bound.
Recently, Xu and Campbell [22] analyzed projected-SGD (stochastic gradient descent) for Gaussian VI using the gradient estimator we will later call g ent (7), with a particular rescaling. They show that, asymptotically in the number of observed data, their method converges locally with a rate of1/ √ T , under mild assumptions. Our results are less general in requiring convexity, but are non-asymptotic, apply with other gradient estimators, and give a faster 1/T rate with strong convexity.
Lambert et al. [23] introduce a VI-like SGD algorithm, derived from a discretisation of a Wasserstein gradient flow, and show it converges at a 1/T rate for the 2-Wasserstein distance when the log posterior is smooth and strongly concave. This line was continued by Diao et al. [24], who propose a proximal-SGD method based on the decomposition f = l + h . They obtain a 1/ √ T rate when the log posterior is smooth and convex, and a1/T rate when it is smooth and strongly convex. Unlike typical black-box VI algorithms used in practice, these algorithms require computing the Hessian of the posterior. We analyze more straighforward applications of SGD to the VI problem using standard gradient estimators. Under the same assumptions, our algorithms have the same rates for KL-divergence, which imply the same rates for 2-Wasserstein distance by of Pinsker's inequality and that the total-variation norm upper-bounds Wasserstein distance [25,Remark 8.2].
In concurrent work, Kim et al. [26], consider a proximal-SGD method similar to our approach in Sec. 6. They obtain a 1/T convergence rate, similar to what Cor. 12 gives when the log posterior is strongly concave. They also consider alternative parametrizations of the scale parameters that render the ELBO globally smooth, and obtain a nonconvex result: Under a relaxed version of the Polyak-Lojasiewicz inequality, they can guarantee a1/T 4 rate.
Traditionally, the ELBO was optimized using message passing methods, which essentially assume that p and q are simple enough that block-coordinate updates are possible [2,3,4]. However, in the last decade a series of papers developed algorithms based on a "black-box" model where p is assumed to be complex enough that one can only evaluatelog p(or its gradient) at selected points z . The key observation is that even if p is quite complex, it is still possible to compute stochastic estimates of the gradient of the ELBO, which can be deployed in a stochastic optimization algorithm [ 5,6,8,9,10].
This paper seeks rigorous convergence guarantees for this problem. We study the setting where the variational family is the set of (dense) multivariate Gaussian distributions, parameterized in terms of w = (m, C) , where m is the mean and C is a factor of the covariance, i.e.
We optimize over w W ∈ , where W = {(m, C) : C ≻ 0} and will further require C to be either symmetric or triangular. This does not really change the problem, since for a given covariance matrix Σ there always exists a symmetric (or lower triangular) positive definite factor C such that CC ⊤ = Σ . Now, is it possible to solve this optimization problem? Without further assumptions, it is unlikely any guarantee is possible, because it is easy to encode NP-hard problems into this VI framework [ 27,28]. We discuss below the assumptions we will make to be able to solve the problem.
The properties of the free energy l depend on the properties of the target distribution p . It is necessary to assume that p is somehow "nice" to ensure the problem can be solved. Titsias and Lázaro-Gredilla [17, Proposition 1] showed that if -log p(•, x) is convex, then l is convex too. Challis and Barber [16,Sec 3.2] showed that if the likelihoodp(•|z) is convex and the priorp(z) is Gaussian, then l is strongly-convex. Domke [13,Theorem 9] showed that if -log p(•, x) is µ -strongly convex, then l is µ -strongly convex as well, and that the constant is sharp. Similarly, Domke [13,Theorem 1] showed that if log p(•, x)is M -smooth, then l is also M -smooth, and that the constant is sharp.
In this paper we make two assumptions about the target distribution p : the negative log-probability log p(•, x) must be convex (or strongly convex), and Lipschitz smooth. Section 7.3 (Appendix) gives some example models where these assumptions are satisfied. For example, if the model is Bayesian linear regression, or logistic regression with a Gaussian prior, then -log p(•, x) is smooth and strongly convex. In addition, if the target is a hierarchical logistic regression model, then log p(•, x) is smooth and convex.
Assumptions on p also impact what a minimizer w * = (m * , C * ) of f (w) can look like. Intuitively, if the target log p(•, x)is µ -strongly concave, then we would expect the target distribution to be "peaky". This means that the optimal distribution would be close to a delta function centered at the MAP solution: m * would be close to some maximum of log p(•, x)noted m, and the covariance factor C * would not be too large. This intuition can be formalized: in this context we have ∥C * ∥ 2 + m ∥ * -m∥ 2 2 ≤ d/µ [13,Theorem 10]. Similarly, if log p(•, x)is M -Lipschitz smooth, then we expect that the target is not too concentrated, so we might expect that the optimal covariance cannot not be too small. Formally, it can be shown that the singular values of the covariance factor C * are greater than 1/ √ M [13, Theorem 7].
The properties of the neg-entropy h are inherited from the choice of the variational family and do not depend on p . Since we consider a Gaussian variational family, h is known in closed-form. To avoid some technical issues, we will restrict h so that the covariance factor is positive definite, so h(w) is equal to -ln det C up to a constant if C is positive definite (see Appendix 7.1), and to+∞ otherwise . So h inherits the properties of the negative log determinant, meaning it is a proper closed convex function whenever C is symmetric or triangular (see Appendix 7.2). From an optimization perspective, it is natural to use a gradient-based algorithm. But unfortunately h is not Lipschitz smooth on its domain, because its gradient can change arbitrarily quickly when the singular values of C become small. However, it can be shown [13,Lemma 12] that h is M -smooth over the following set of non-singular parameters (which contains the solution w * , see the previous paragraph) given by
Instead of computing gradients, an optimizer might want to use the proximal operator of h , which can also be computed in closed form (see Sec. 4).
This paper considers gradient estimators based on the "path" method or "reparameterization". These assume some base distribution s is known, together with some deterministic transformation T w , such that the distribution of T w (u) is equal to q w when u s ∼ . Then, we can write
In the case of multivariate Gaussians, the most common choice is to use s = N (0, I) and T w (u) = Cu + m , which we will consider in the rest of the paper.
We will consider three different gradient estimators based on the path-type strategy (Eq. 4). We will provide new noise bounds for each of them (all the proofs are in Appendix 8). Our analysis is mostly based on the following general result [20,Thm. 3].
Theorem 1. Let T w (u) = Cu + m for w = (m, C) . Let ϕ : R d → R be M -smooth, suppose that ϕ is stationary at m, and define w = ( m, 0). Then
Furthermore, the first inequality cannot be improved.
Intuitively, this result says that the noise of a gradient estimator is lower when the scale parameter C is small and the mean parameter m is close to a stationary point.
The first estimator we consider is ∇l(w) only. It is given by taking ϕ = -log p into Eq. 4, i.e.
The next result gives a bound on the noise of this estimator, which is a direct consequence of Theorem 1. The second line uses Young's inequality to bound the noise in terms of (i) the distance of w from w * and (ii) a constant determined by the distance of w * from some fixed parameters w.
Theorem 2. Suppose that log p(•, x)is M -smooth and has a maximum (or stationary point) at m, and define w = ( m, 0). Then, for every w and every solution w * of the VI problem,
Second, we consider an estimator of the gradient of the full objectivel + h. It is obtained by simply taking the above estimator and adding the true (known) gradient of the neg-entropy, i.e.
The noise of g ent can be bounded since it only differs from g energy by the deterministic quantity ∇h(w) , and the fact that-provided w W ∈ L -the singular values of w cannot be too small and so ∇h(w) cannot be too large. This is formalized in the following theorem, where again, the second line relaxes the result into a term based on the distance of w from w * plus constants. (Another type of noise bound [29] for g ent was obtained by Kim et al. [21] for diagonal Gaussians and a variety of parameterizations of the covariance.) Theorem 3. Suppose that log p(•, x)is M -smooth, that it is maximal at m, and define w = ( m, 0). Then, for every L > 0 , for every w W ∈ L and every solution w * of the VI problem,
≤ 4(d + 3)Mfoot_1 ∥w -w * ∥ 2 + 4(d + 3)M 2 ∥w * -w∥ 2 + 2dL.
While g ent used the exact gradient of h , it may be beneficial to use a stochastic estimator h instead.
To derive our third estimator, write the gradient of the ELBO asfoot_0
where the parameters v serve as a way to "hold w constant under differentiation". This leads to our third estimator, called the "sticking the landing" (STL) gradient estimator
Intuitively, we expect thatg STL will tend to have lower noise than g ent when the posterior is wellapproximated by the variational distribution. The reason is that, as observed by Roeder et al. [30], if q w (z) were a perfect approximation of p(z|x) , then the two terms in Eq. 9 would exactly cancel (for every u ) and so the estimator would have zero variance.
Below we formalize this intuition in what we believe is the first noise bound on g STL . Theorem 4. Suppose that log p(•, x) is M -smooth. Consider the residual r(z) := log p(z, x)log qw * (z) for any solution w * of the VI problem, assume that it has a stationary point m, and define ŵ = ( m, 0). Then r is K -smooth for some K ∈ [0, 2M ] , and for all w W ∈ M , E ∥g STL ∥ 2 2 ≤ 8(d + 3)M 2 ∥w -w * ∥ 2 2 + 2(d + 3)K 2 ∥w -ˆw∥ 2 2 (10)
When the target is Gaussian then K = 0 in Eq. 10, meaning that when w = w * , the STL estimator has no variance. This is a distinguishing feature of g STL , as opposed to g energy or g ent .
Three major issues bar applying existing analysis of stochastic optimization methods to our VI setting: 1) non-smooth composite objective 2) lack of uniform smoothness and 3) lack of uniformly bounded noise of the gradient estimator.
The first issue is due to the non-smoothness of neg-entropy function h . This means that under the benign assumption that the target log p is smooth, the full objective l + h cannot be smooth, since a nonsmooth function plus a smooth function is always nonsmooth. This renders stochastic optimization proofs (e.g. those for the "ABC" conditions [29]) that do not use projections or proximal operators inapplicable.
One way to tackle this first issue would be to use a non-smooth proof technique, but these rely on having a uniform gradient noise bound (our third issue). Alternatively, one can overcome the non-smoothness of the neg-entropy function by either using a proximal operator, or projecting onto a set where h is smooth, namely W M . This is the strategy we pursue.
Table 1: Table of known black-box VI properties relevant to optimization. In some cases there could be multiple valid w * or w in which case these results hold for all simultaneously.
Estimator for ∇l g energy = -∇ w log p(Tw (u), x)
Stationary point of l w = (¯z, 0) for any z that is a stationary point of log p(•, x)
g ent , and g STL are quadratically bounded The second issue is that existing analyses for proximal/projected stochastic methods either rely on a uniform noise bound ([31, Cor. 3.6]) or uniform smoothness [32, 33, 34, 35, 36], both of which are not known to be true for VI. By uniform smoothness, we refer to the assumption that log p(Tw (u), x) is M -smooth for every u , with M being independent of u . Instead, we can only guarantee that log p(Tw (u), x) is smooth in expectation, i.e. that l(w) is smooth. Several works [9, Cond. 1][37, Thm. 1][38, Sec. 4][39, Assumption A1][40, Thm. 1][41, Assumption 3.2] assumed that the full objective l + h is smooth, but we are not aware of cases where this holds in practice for VI and in our parameterization (Eq. 2) this cannot be true if log pis smooth.
The third issue is the lack of a uniform noise bound. Most non-smooth convergence guarantees within stochastic optimization [42] assume that the noise of the gradient estimator is uniformly bounded by a constant. But this does not appear to be true even under favorable assumptions-e.g. it is untrue if the target distribution is a standard Gaussian. The best that one can hope is that the gradient noise can be bounded by a quadratic that depends on the current parameters w , and-depending on the estimator-even this may only be true when the parameters are in the set W M . Thus, our main optimization contribution is to provide theoretical guarantees for stochastic algorithms under such a specific noise regime, which we make precise in the next definition.
Definition 5. Let ϕ be a differentiable function. We say that a random vectorg is a quadratically bounded estimator for ∇ϕ at w with parameters (a, b, w * ), if it is unbiased E [g] = ϕ ∇ (w) and if the expected squared norm is bounded by a quadratic function of the distance of parameters w to w * , i.e. E ∥g∥ 2 2 ≤ a w -w ∥ * ∥ 2 2 + b. The noise bounds derived in Section 2.2 for the estimatorsg energy , g ent , and g STL imply that they are uniformly quadratically bounded estimators. See Appendix 8.2 for a table with the corresponding constants a and b .
In this section we give new convergence guarantees for the Prox-SGD algorithm and the Proj-SGD algorithm with quadratically bounded gradident estimators. Because these may be of independent interest, they are presented generically, without any reference to the VI setting. We specialize these results to VI in Section 4. For both algorithms, we present results assuming the problem to be strongly convex or just convex. All proofs for this section can be found in the Appendix (Sec. 9).
Here we want to minimize a function which is the sum of two terms l + h , were both l and h are proper closed convex functions, and l is smooth. For this we will use stochastic proximal gradient.
Definition 6 (Prox-SGD). Let w 0 be a fixed initial parameters and let γ 1 , γ 2 , • • • be a sequence of step sizes. The stochastic proximal gradient (Prox-SGD) method is given by
where g t is a gradient estimator for ∇l(w t ), and the proximal operator is defined as
Theorem 7. Let l be a µ -strongly convex and M -smooth function, and let w = argmin(l) . Let h be a proper closed convex function, and let w * = argmin(l + h) . Let (w t ) t∈N be generated by the Prox-SGD algorithm, with a constant stepsize γ ∈ 0, min{ µ 2a
, 1 µ } . Suppose that g t is a quadratically bounded estimator (Def. 5) for ∇l with parameters (a, b, w * ). Then,
Alternatively, if we use the decaying stepsize
In both cases, T = O(ϵ -1 ) iterations are sufficient to guarantee thatE w ∥ T -w * ∥ 2 ≤ ϵ.
The above theorem gives an anytime 1/T rate of convergence when the stepsizes are chosen based on the strong convexity and gradient noise constants µ and a . Note that we do not need to know precisely those constants: We show in Appendix 9.3 that using any constant step size proportional to T / log T leads to a log T /T rate.
Theorem 8. Let l be a proper convex and M -smooth function. Let h be a proper closed convex function, and let w * ∈ argmin(l + h) . Let (w t ) t∈N be generated by the Prox-SGD algorithm, with a constant stepsize γ ∈ (0, 1 M ]. Suppose that g t is a quadratically bounded estimator (Def. 5) for ∇l with parameters (a, b, w * ). Then,
where θ def = 1 1 + 2aγ 2 and wT def = P T t=1 θ t+1 w t P T t=1 θ t+1
. In particular, if γ = 1 √ aT
, then
Thus, T = O(ϵ -2 ) iterations are sufficient to guarantee thatE f ( wT ) -inf f ≤ ϵ.
Here we want to minimize a function f over a set of constraints W , where W is a nonempty closed convex set and f is a proper closed convex function which is differentiable on W . To solve this problem, we will consider the stochastic projected gradient algorithm.
Definition 9 (Proj-SGD). Let w 0 be some fixed initial parameter, and let γ 1 , γ 2 , • • • be a sequence of step sizes. The projected stochastic gradient (Proj-SGD) method is given by
where g t is a gradient estimator for ∇f(w t ), and the projection operator is defined as
Note that we do not require f to be smooth, meaning that this setting is not a particular case of the one considered in Section 3.1. In fact, the arguments we use in the proofs for Proj-SGD are more closely related to stochastic subgradient methods than stochastic gradient methods. Theorem 10. Let W be a nonempty closed convex set. Let f be a µ -strongly convex function, differentiable on W . Let w * = argmin W (f ) . Let (w t ) t∈N be generated by the Proj-SGD algorithm, with a constant stepsize γ ∈ 0, min{ µ 2a
, 2 µ } . Suppose that g t is a quadratically bounded estimator (Def. 5) for ∇f with parameters (a, b, w * ). Then,
Alternatively, if we use the decaying stepsize
In both cases, T = O(ϵ -1 ) iterations are sufficient to guarantee thatE w
Note that Eq. 13 is a sum of two terms: one that decays exponentially in T and one that decreases only when the stepsize γ is sufficiently small. This has an important consequence: if one uses the gradient estimator g STL and the target distribution is exactly a Gaussian, thenb = 0 (Appendix 8.2), meaning that the algorithm will converge at an exponential rate. This is similar to many results in the stochastic optimization literature showing faster rates hold when interpolation holds [35]. Theorem 11. Let W be a nonempty closed convex set. Let f be a convex function, differentiable on W . Let w * ∈ argmin W (f ) . Let (w t ) t∈N be generated by the Proj-SGD algorithm, with a constant stepsize γ ∈ (0, +∞) . Suppose that g t is a quadratically bounded estimator (Def. 5) for ∇f at w t with constant parameters (a, b, w * ). Then,
where
We now specialize the optimization results of the previous section to our VI setting. We aim to minimize the negative ELBO, f : R
and where V is the vector space of matrices in which the covariance factors C belong. Under the assumption that the log-target log p(•, x)is M -smooth and concave, we propose two strategies to minimize this objective. All the proofs for this section can be found in Appendix 10.
1. Apply the Prox-SGD algorithm to the sum l + h , using a proximal step with respect to h , and a stochastic gradient step with respect to l . The gradient of l will be estimated with g energy (5), which is globally quadratically bounded. The prox of h admits a closed form formula if we choose that V is the space of lower triangular matrices. 2. Apply the Proj-SGD algorithm to minimize f over W M . The gradient of f will be estimated using either g ent (7) or g STL (9), both of which are quadratically bounded on W M . The projection onto W M admits a closed form formula if choose that V is the space of symmetric matrices.
Corollary 12 (Prox-SGD for VI). Consider the VI problem where q w is a multivariate Gaussian distribution (Eq. 2) with parameters w = (m, C) ∈ R d × T d , and assume that this problem admits a solution w * . Suppose that log p(•, x)is M -smooth and concave (resp. µ -strongly concave). Generate a sequence w t by using the Prox-SGD algorithm (Def. 6) applied to l and h (Eq. 16), using g energy (5) as an estimator of ∇l . Let the stepsizes γ t be constant and equal to 1/( p a energy T ) (resp.
be decaying as in Theorem 7 with a energy = 2(d + 3)M 2 ). Then, for a certain average wT of the iterates, we have for
Corollary 13 (Proj-SGD for VI). Consider the VI problem where q w is a multivariate Gaussian distribution (Eq. 2) with parameters w = (m, C) ∈ R d × S d , and assume that this problem admits a solution w * . Suppose that log p(•, x)is M -smooth and concave (resp. µ -strongly concave). Generate a sequence w t by using the Proj-SGD algorithm (Def. 9) applied to the function f = l+h (Eq. 16) and the constraint W M (Eq. 3), using g ent (7) or g STL (9) as an estimator of ∇f . Let the stepsizes γ t be constant and equal to p 2/(aT ) (resp. be decaying as in Theorem 10) with a = a ent = 4(d + 3)M 2
or a = a STL = 24(d + 3)M 2 . Then, for a certain average wT of the iterates, we have for
Corollary 14 (Proj-SGD for VI -Gaussian target). Consider the setting of Corollary 13, in the scenario that log p(•, x)is µ -strongly concave, that we use the g STL estimator, and that we take a constant stepsize γ t ≡ γ ∈ 0, min{
Let us now discuss the practical implementation of these two algorithms (more details and an explicit implementation of the methods are given in Appendix 10.2). These require computing either the proximal operator of the neg-entropy h , or the projection onto the set of non-degenerate covariance factors W M . These can be computed [13] for every w = (m, C) as:
Our theory for the Prox-SGD algorithm formally assumes that the covariance factor C lives in the vector space of lower-triangular matrices. This can be implemented by letting C ∈ R d×d and "clamping" the upper-triangular entries to zero throughout computation. Then the gradient g energy (5) can be computed using automatic differentiation, and the proximal operator can be computed as in the above equation. Or, one may choosew = (m, c) where c ∈ R d×(d+1)/2 is the lower-triangular entries of the matrix so C = tril(c) . Gradients with respect to w = (m, c) can again be estimated using automatic differentiation (including the tril operator), and the proximal operator can be computed by forming C , using the above formula, and then extracting the lower-triangular entries.
Similarly, our theory for the Proj-SGD algorithm formally assumes the covariance factor C lives in the vector space of symmetric matrices. This can be implemented by letting C ∈ R d×d , computing the gradient estimator g ent (7) or g STL (9) over the reals using automatic differentiation, and "symmetrizing"
C ⊤ after applying each gradient update, right before projection. Or, one may choosew = (m, c) where c ∈ R d×(d+1)/2 is the lower-triangular entries of the matrix, so C = symm(c) , where symm is a function similar to tril except creating symmetric matrices. The gradient estimator g ent (7) or g STL (9) can be computed using automatic differentiation (including the symm operator). In this case, to project one would first form C , then use the above formula, and then extract the lower-triangular matrices.
In terms of cost, computing log qw (z) or T w (u) needs Θ(d 2 ) operations. Computing the prox of h requires d operations, but in contrast projecting onto W M requires diagonalizing a symmetric matrix, which takes Θ(d 3 ) operations. We note that this is not necessarily an issue, because 1) eigenvalue decomposition has excellent computational constants in moderate dimensions; 2) computing the target log p(z, x) may itself require Θ(d 3 ) or more operations; 3) it is common to average a "minibatch" of gradient estimates, meaning each eigenvalue decomposition computation is amortized over many ELBO evaluations. Still, in high dimensions, when log p(z, x) is inexpensive, and small minibatches are used, computing such decomposition in each iteration could become a computational bottleneck. In this scenario, the Proj-SGD algorithm is clearly cheaper, with its Θ(d) cost for the proximal step.
While this paper has focused on convergence with dense Gaussian variational distributions, the results in Sec. 3 apply more generally. It is conceivable that the necessary smoothness, convexity, and noise bounds could be established for other variational families, in which case these optimization results would provide "plug in" guarantees. We suspect this might be fairly easy to do with, e.g., elliptical distributions or location-scale families, using similar techniques as done with Gaussians. But it may be difficult to do so for broader variational families.
If -log p is strongly convex and smooth, another proof strategy is possible: Smoothness guarantees that w * ∈ W M and strong convexity guarantees that w * ∈ {w : w -∥ ¯w∥ 2 2 ≤ d/µ} . So one could do projected SGD, projecting onto the intersection of these two sets. The negative ELBO would be strongly convex and smooth over that set, and since ∥w -¯w∥ 2 is bounded, gradient noise is upperbounded by a constant, meaning classical projected SGD convergence rates that assume smoothness, strong convexity, and uniform gradient noise apply. However, projecting onto the intersection of those sets poses a difficulty and this uniform noise bound may be loose in practice.
Several variants of the gradient estimators we consider have been proposed to reduce noise, often based on control variates or sampling from a different base distribution [30,37,43,44,45,46,47,48,49,50]. It would be interesting to establish gradient noise bounds for these estimators.
Various VI variants have been proposed that use natural gradient descent. It would also be interesting to seek convergence guarantees for these algorithms.
VI can be done with "score" or "reinforce" gradient estimators [6,8], which use the identity
in place of how we used Eq. 4 for the "path" estimators we considered. While path estimators often seem to have lower variance, this is not always true: Take an unnormalized log-posteriorϕ(z) where z is scalar. Now, define ϕ ′ (z) = ϕ(z) + √ ϵ sin(z/ϵ), where ϵ is very small. Then, ϕ and ϕ ′ represent almost the same posterior, and score estimators for them will have almost the same variance. Yet, the derivative of the added term is cos(z/ϵ)/ √ ϵ , so a path estimator for ϕ ′ will have much higher variance than one for ϕ . This underlines why smoothness is essential for our guarantees-it excludes posteriors like ϕ ′ . Future work further unravel when score estimators perform better than path estimators.
One attractive feature of VI is that data subsampling can often be used when computing gradients, decreasing the cost of each iteration. While we have not explicitly addressed this, our proofs can be easily generalized, at least for target distributions of the form p(z, x) = p(z) Q n=1 p(xn |z). The only issue is to bound the noise of the subsampled estimator. While our gradient variance guarantees use Theorem 1 from [20], a more general result [20,Theorem 6] considers data subsampling. Very roughly speaking, with uniform data subsampling of 1 datum of a time, the gradient noise bounds in Thms. 2 and 3 would increase by a factor between 1 (no increase) and the number of data, depending on how correlated the data are-less correlation leads to a larger increase. (More precisely, the second term in Thm. 3 would not increase.) These increases would manifest as larger constants a and b in the quadratic bounds, after which exactly the same results hold. However, it is not obvious if the bound for g STL in Thm. 4 can be generalized in this way.
We note M d the space of d×d matrices, S d the subspace of symmetric matrices, and T d the subspace of lower triangular matrices. We use
} to denote the set of parameters with positive definite covariance matrix whose singular values are greater or equal to
variance of vector-valued random variables
This section contains the proofs of all the claims made in Section 2.
We recall and detail here the setup of our problem. Given a target distribution p and observed data x , we want to solve the following Variational Inference problem min
where q w is a multivariate Gaussian variational family, whose density may be written as
We make the assumption (VI) has a non-degenerated solution, i.e. a solution for which the covariance matrix is invertible.
We will impose that our parameters have the form w = (m, C) ∈ R d × V , where V is some vector subspace of M d , the space of d × d matrices. On top of that, we will also want the covariance factor to be positive definite. In other words, we will impose that our parameters are taken from the set
For the sake of simplicity, the reader can assume for most of the paper that V = M d . But when coming to analyze our algorithms, we will see that we should take either V = T d (the subspace of d × d lower triangular matrices) or V = S d (the subspace of d × d symmetric matrices). This choice is dictated by two reasons: it guarantees that our problem remains convex (see Lemma 19), and it lowers the computational cost of our algorithms (see Section 10.2). Fortunately, requiring C to be symmetric (or triangular) and positive definite does not change our problem, as stated in the next Lemma. Lemma 15 (Equivalent problems for triangular/symmetric covariance factors).Suppose that w ∈ R d ×M d is such that q w minimizes KL (qw ∥p(•|x)) over R d ×M d . Then there exists w * = (m * , C * ) such that q w and q w * have the same distribution, with C * being positive definite and symmetric (or lower triangular).
Proof. Write w = ( m, C) , and take m * = ¯m. Let Σ = C C⊤ . Since we assumed that (VI) has a non-degenerated solution, we can assume without loss of generality that Σ is invertible. Therefore we can apply the Cholesky decomposition and write Σ = C * (C * ) ⊤ , where C * is lower triangular and positive definite. We can also define C * = ( Σ) 1/2 to obtain a symmetric and positive definite factor.
Once restricted to W , our problem (VI) becomes equivalent to
where
This equivalence is mostly standard, but we state it formally for the sake of completeness.
Lemma 16 (VI reformulated). Problems (VI) and (18) are equivalent.
Proof. Let w = (m, C) W ∈ , and let us show that KL (qw ∥p(•|x)) = l(w) + h(w) . By definition of the Kullback-Liebler divergence, and using p(z|x) = p(z, x)/p(x) , we can write
If we note H(X) the entropy of a random variable, we can see that E z q ∼ w log qw (z) = -H(z) where z q ∼ w . We can rewrite z via an affine transform as z = Cu + m with u N ∼ (0, I) , so we can use [51,Section 8.6] to write (see also the arguments in [13])
Here H(u) , the entropy of u, is a constant equal to (d/2)(1 + log(2π) . Moreover since w W ∈ , we know that C ≻ 0 , which implies that log | det C| = log det C. Because log p(x) and H(u) are both constants, we deduce that minimizing KL (qw ∥p(•|x)) over w W ∈ is equivalent to minimizing l + h , where h(w) = -log det C if w W ∈ and +∞ otherwise.
We emphasize that working with a triangular or symmetric representation for C is not restrictive, and that doing computations in this setting is pretty straightforward, as we illustrate next.
Lemma 17 (Computations with triangular/symmetric covariance factors). Let V be a vector subspace of 1. This is immediate due to the fact that V is convex, since it is a vector subspace of M d .
) be the canonical injection. Then we clearly have ϕ V = ϕ • ι . Moreover, ι is a linear application whose adjoint is ι * is exactly the orthogonal projection proj R d ×V . So the computation of the gradient follows after applying the chain rule :
. As for the norm, it suffices to observe that the orthogonal projection is a linear operator with norm less or equal than 1.
3. Using that the orthogonal projection is a contractive map gives
4. This is a standard linear algebra result.
Now we state the main smoothness and convex properties for VI problems. We will see in our analysis that the smoothness properties often revolve around the following subset of W
where σ min (C) refers to the smallest singular value of C .
Lemma 18 (Smoothness for VI). Assume that log p(•, x)is M -smooth. Then
Proof.
1. Combine [13, Theorem 1] with Lemma 17.3.
Combine the fact that the gradient of -log det at a matrix X is -X -⊤ with Lemma 17.2.
3. Combine [13, Lemma 12] with Lemma 17.3.
Assume that -log p(•, x) is convex (resp. µ -strongly convex). Then
Proof.
Titsias and Lázaro-Gredilla[17, Proposition 1] (resp. Domke [13, Theorem 9]) with Lemma 17.1. 2. Convexity of h follows from [52, Example 7.13 and Theorem 7.17]. Convexity of W comes from the fact that it is essentially the set of symmetric positive definite matrices (which is equivalent to λ min (C) > 0 ). Let us now prove that W M is convex and closed. First we notice that, because of symmetry, we can rewrite W M as
This set is closed because λ min is a continuous function. To see that W M is convex we are going to use the fact that λ min is a concave function. Take C 1 , C 2 ∈ W M , α ∈ [0, 1] and write
Using Sylvester's criterion, it is easy to see that a lower triangular matrix is positive definite if and only if its diagonal has positive coefficients. In other words, W = {(m, C) ∈ R d × T d : Cii > 0} , which is clearly convex. It remains to verify that h is closed convex even if W is not closed. For all w = (m, C) ∈ R d × T d , we can use the triangular structure of C to write
where we used the notation δ A for the indicator function of a set A , which is equal to zero on A and +∞ outside. Since -log(t) + δ (0,+∞) (t) is a closed convex function, we see that h is a separable sum of closed convex functions, and thus is itself closed and convex.
4. Let d = 2 , let I be the identity matrix, let R = 0 1 -1 0 be a rotation matrix, and define ψ : R → R with ψ(t) = h(I + tR) . Observe that ψ is well defined because I + tR ≻ 0 for all t ∈ R . On the one hand, if h was convex, then ψ would be convex too, as it is a composition of a convex function with the affine function t 7 → I + tR . On the other hand, we can compute explicitly that ψ(t) = -log(1 + t 2 ), which is not convex. Theorem 7].
Suppose now that V = S d , meaning that w * is a minimizer of f over W ⊂ R d × S d . We see that w * must also be a minimizer of f over R d ×M d , because if there was a better solution ŵ = ( m, Ĉ) with Ĉ M ∈ d , we could consider ( Ĉ Ĉ⊤ ) 1/2 ∈ S d which would itself be a better solution than C * , which is a contradiction. So we can apply [13,Theorem 7] Proof. According to the definition of q w (see (17)), the Hessian of log qw is ∇ 2 z (log qw )(z) = -(CC ⊤ ) -1 , and so ∥∇ 2 z (log qw )(z)∥ = σ max ((CC ⊤ ) -1 ) = σ min (C) -2 .
where A is a matrix with a n in the n -th column. Then
Now, take some arbitrary vector v . We have that
Putting both the above results together, we have that
This implies that all eigenvalues of -∇ 2 z log pare in the interval [µ, β]. All the claimed properties follow from this:
• For any eigenvalue λ of ∇ 2 z log p(z, x), we know that λ i (∇ 2 z log p(z, x)) ≤ M = max (|µ| , |β|) .Thus, log p(z, x) is M -smooth over z .
• If µ ≥ 0 , this means all eigenvalues of -∇ 2 z log p(z, x) are positive, which implies that log p(z, x) is convex.
• If µ > 0 , then all eigenvalues of -∇ 2 z log(z, x) are bounded below by α , which impiles that log p(z, x) is µ -strongly convex.
Proof. In this case we have that
Or, more abstractly,
Thus we have that θ min = θ max = 1 σ 2 . It follows from Theorem 24 that -log p(z, x) is µ -strongly convex for µ = λ min Σ -1 + 1 σ2 AA ⊤ and M -smooth for M = λ max Σ -1 + 1 σ2 AA ⊤ . In the particular case where Σ = I , we get that µ = λ min I + 1
Corollary 26. Take a logistic regression model with inputs a n ∈ R d and outputs
Proof. In this case we have that
Or, more abstractly,
The second derivative is non-negative and goes to zero if α → -∞ or α → +∞ . It is maximized at α = 0 in which case ϕ ′′ = 1 4 β 2 . But of course, in this model, the second input is β = x n ∈ {-1, +1} , meaning β 2 = always. Thus we have that θ min = 0 and θ max = 1 . It follows from Theorem 24 that -log p(z, x) is µ -strongly convex for µ = λ min Σ -1 and M -smooth for M = λ max Σ -1 + 1 4 AA ⊤ .
In the particular case where Σ = I , we get that µ = 1 and M = λ max (1 + 1 4 AA ⊤ ) = 1 + 1 4 σ max (A) 2 . Theorem 27. Take a heirarchical model of the form p(θ, z, x) = p(θ) Q i p(zi |θ)p(xi |θ, z i ), where all log-densities are twice-differentiable. Then, log p(θ, z, x)is smooth joinly over (θ, z) if and only if the spectral norm of the second derivative matrices ∇ θθ ⊤ log p(θ, z, x) , ∇ θz ⊤ i log p(θ, z i , x i ), and ∇ z i z ⊤ i log p(θ, z i , x i ) are all bounded uniformly over (θ, z). Proof. The following conditions are all equivalent: 1. log p(θ, z, x)is smooth as a function of (θ, z) 2. ∇ 2 log p(θ, z, x) 2 is bounded above (uniformly over (θ, z), where ∇ 2 denotes the Hessian with respect to (θ, z) and ∥•∥ 2 denotes the spectral norm.) 3. ∇ 2 log p(θ, z, x) is bounded above (uniformly over (θ, z), where ∇ 2 denotes the Hessian with respect to (θ, z) and ∥•∥ denotes a block infinity norm on top of spectral norm inside each block. 4. ∇ 2 θθ ⊤ log p(θ, z, x) 2 , ∇ 2 θz ⊤ i log p(θ, z, x) 2 , and ∇ 2 z i z ⊤ j log p(θ, z, x) 2 are all bounded above (uniformly over (θ, z)) 5. ∇ 2 θθ ⊤ log p(θ, z, x) 2 , ∇ 2 θz ⊤ i log p(θ, z i , x i ) 2 , and ∇ 2 z i z ⊤ j log p(θ, z, x) 2 are all bounded above (uniformly over (θ, z)) 6. ∇ 2 θθ ⊤ log p(θ, z, x) 2 , ∇ 2 θz ⊤ i log p(θ, z i , x i ) 2 , and ∇ 2 z i z ⊤ i log p(θ, z, x) 2 are all bounded above (uniformly over (θ, z)) 7. ∇ 2 θθ ⊤ log p(θ, z, x) 2 , ∇ 2 θz ⊤ i log p(θ, z i , x i ) 2 , and ∇ 2 z i z ⊤ i log p(θ, z i , x i ) 2 are all bounded above (uniformly over (θ, z)) Corollary 28. Take a hierarchical logistic regression model defined by p(θ, z, x) = p(θ)
where
. Then, -log p(θ, z, x)is convex and smooth with respect to (θ, z).
Proof. It's immediate that -log p is convex with respect to(θ, z). We can show that this function is smooth, by verifying the three conditions of Theorem 27.
• p(zi |θ) is a Gaussian with constant covariance, so is smooth with a constant independent of θ . And we can write that log p(xij |z i ) = log 1 + exp -x ij z ⊤ i a ij . This function is also smooth over z i with a constant that doesn't depend on θ . Thus log p(zi |θ) + P j log p(xij |z i ) is smooth over z i with a smoothness constant that is independent of θ .
We remember that unless specified otherwise, we consider that w = (m, C) ∈ R d × V , where V is a vector subspace of M d .
Theorem 1. Let T w (u) = Cu + m for w = (m, C) . Let ϕ : R d → R be M -smooth, suppose that ϕ is stationary at m, and define w = ( m, 0). Then
Furthermore, the first inequality cannot be improved.
Proof. This bound is proven by [20,Thm. 3] in the case that V = M d . We conclude that this bound holds for every subspace V of M d by using the inequality in Lemma 17.2.
Theorem 2. Suppose that log p(•, x)is M -smooth and has a maximum (or stationary point) at m, and define w = ( m, 0). Then, for every w and every solution w * of the VI problem,
Proof. It is a direct consequence of Theorem 1 and Young's inequality
Theorem 3. Suppose that log p(•, x)is M -smooth, that it is maximal at m, and define w = ( m, 0). Then, for every L > 0 , for every w W ∈ L and every solution w * of the VI problem,
Proof. The difference between g ent and g energy is the addition of the constant vector ∇h(w) . Thus, we have that
It remains to bound the final term. We can do this using the closed-form expression for the entropy gradient (see Lemma 18.2), plus the assumption that w W ∈ L ⊂ W , to see that
We conclude with Young's inequality
For the next Theorem, we will need the follwing technical Lemma :
Proof. Develop the squares to write
Since Eu = 0 , we immediately see that E A ⟨ ⊤ b, u⟩ = 0 . Because of symmetry, the third-order moment Eu∥u∥ 2 is also zero, so we obtain E A ⟨ ⊤ b, u∥u∥ 2 ⟩ = 0 . We also have the second-order moment
The conclusion follows after gathering all those identities. Theorem 4. Suppose that log p(•, x) is M -smooth. Consider the residual r(z) := log p(z, x)log qw * (z) for any solution w * of the VI problem, assume that it has a stationary point m, and define ŵ = ( m, 0). Then r is K -smooth for some K ∈ [0, 2M ] , and for all w W ∈ M ,
Proof. Let w W ∈ M be fixed. For the duration of the proof, we take v to be a second copy of w (held constant under differentiation with respect to w ). Now, we can rearrange the estimator in Eq. 9 by making appear q w * (Tw (u)) , namely
p(Tw (u), x) .
Introducing ϕ(z) := log qv (z) -log qw * (z) , and recalling r(z) = log p(z, x) -log q w * (z) , we have
We assume that log p(•, x)is M -smooth, and that V = M d , T d or S d , which implies that w * ∈ W M (see Lemma 20). This in turn guarantees that log qw * is also M -smooth (see Lemma 21). Thus, r is the sum of two M -smooth functions, so it is K -smooth with K ∈ [0, 2M ] . Using Young's inequality, together with Theorem 1 applied to r , we obtain
where ŵ = ( m, 0), with m being a stationary point of r . Now it remains to control the last term of inequality (19). Apply the chain rule onto ϕ(Tw (u)) to write (with the help of Lemma 37)
This gives us directly that
where we used that ∥vu ⊤ ∥ 2 F = v ∥ ∥ 2 ∥u∥ 2 for any vectors v and u . Computing ∇ z ϕ amounts to computing the gradient of the Gaussian density q w . From its definition (see Eq. 17) we have:
Recalling the definition ofϕ = log qv -log qw * , and
In other words, we have
where
We can now combine (20) and ( 21) and use Lemma 29 to write
Let us now introduce the functionκ : R d ×M d → R { ∪ +∞} , defined byκ(w) := E z q ∼ w log qw (z)log qw * (z) if C is invertible, +∞ otherwise. This function is nothing but the Kullback-Liebler divergence between the two Gaussian distributions q w and q w * , and can be computed in closed-form on its domain:
where we note as before w = (m, C) ,
This allows us to compute its gradient at w W ∈ M : 22) and the fact that ∇ w κ(w * ) = 0 (because w * is a minimizer of κ ), we deduce that
Now we remind that w * ∈ W M which implies that log qw * is M -smooth. This means that E z q ∼ w log qw * is M -smooth as well, according to Domke [13,Theorem 1]. On the other hand, we know (see the proof of Lemma 16) thatE z q ∼ w log qw * = h(w) which is M -smooth on W M (see Lemma 18.3). All this implies that κ is 2M -smooth on W M , from which we conclude
The Theorem's main inequality (10) follows after plugging the above inequality into (19). The second inequality follows after using Young's inequality
To conclude the proof, it remains to check that K = 0 whenever p(•|x) is Gaussian. In that case, because our main objective function f (w) is the divergence between q w and p(•|x), it is clear that f is minimized whenever q w = p(•|x) . Therefore, without loss of generality, we can assume that p(•|x) = q w * . This implies that the residual r is constant (r(z) = log p(x) ), and so is 0-smooth.
Proof. Combine the results of Theorems 2, 3 and 4.
This section contains all the proofs for the Theorems stated in Section 3.
Theorem 31 ( Theorem 3.2 in [ 53] ). Suppose we are given a sequence of w t iterates such that for a step size γ t ≤ 1 2L and given constants µ > 0 and σ 2 > 0 we have that
By switching to a decaying stepsize according to 2 for t ≥ t * where t * = 4 L/µ ⌊ ⌋ , we have that
Proof. We divide the proof for steps t that are great or smaller than t * . For t ≤ t * , we have by unrolling (23) starting from t * , . . . , 1 , with γ t ≡ γ gives
Now consider t ≥ t * for which we have that γ t = 1 µ 2t+1 (t+1) 2 which when inserted into (23) gives 4 .
Multiplying through by (t + 1) 2 and re-arranging gives
where we used that 2t+1 t+1 ≤ 2 . Summing up over t = t * , . . . , T and using telescopic cancellation gives
Now using that for t ≤ t * we have that γ = 1 L and ( 25) holds, thus
where we used that (1 -γµ) ≤ 1. Substituting γ = 1 2L , t * = 4 L/µ ⌊ ⌋ and multiplying by (T + 1) 2 gives
which concludes the proof of (24).
Lemma 32. Suppose we are given a sequence wT and constants a, b, B > 0 such that
Thus, T = O(ϵ -2 ) iterations are sufficient to guarantee thatE f ( wT ) -inf f ≤ ϵ.
Proof. First we show that 1 1-θ T ≤ 2 is equivalent to
Indeed this follows since
x for x ≥ 0 . Applying this gives that
it is sufficient to enforce that
Assuming A ≥ p 2/B , this last condition holds if
Since we also impose that BA 2 ≥ 2 we have that 1 BA 2 -1 ≤ 1 thus for (27) to hold it suffices that T ≥ 2. Substituting γ = A √ T and the relaxation 1 1-θ T ≤ 2 into the original bound gives the claimed result.
We start by recalling some standard properties of the proximal operator. Lemma 33. Let γ > 0 and w * ∈ argmin w l(w) + h(w). Assume h(w) is convex and that l(w) is continuously differentiable at w * . Then 1. For all w, w ′ and γ > 0 , prox γh (w)
Theorem 7. Let l be a µ -strongly convex and M -smooth function, and let w = argmin(l) . Let h be a proper closed convex function, and let w * = argmin(l + h) . Let (w t ) t∈N be generated by the Prox-SGD algorithm, with a constant stepsize γ ∈ 0, min{ µ 2a
, 1 µ } . Suppose that g t is a quadratically bounded estimator (Def. 5) for ∇l with parameters (a, b, w * ). Then,
Alternatively, if we use the decaying stepsize γ t = min
In both cases, T = O(ϵ -1 ) iterations are sufficient to guarantee thatE w
Proof. We start by using the non-expansiveness and fixed-point properties of the proximal operator to write
Above the first line uses the fixed point property of the proximal operator that w * = prox γh (w *γ l ∇ (w)) , while the second line uses the fact that the proximal operator is non-expansive. Now, we apply our assumption on the gradient estimator to write
For the last term, we have by strong convexity and the fact that g t is an unbiased estimator that
where the second line follows from the fact that l is µ strongly convex.
Putting the pieces together, we get that
Now the following conditions are equivalent:
Now since γ ≤ 1 µ ⇔ 1 -γµ ≥ 0 we can apply the above recursively, which gives that
But we can bound this geometric sum by
leading to the result that
To prove the anytime result, we can now apply Theorem 31. To do so, we need to map the notation we use here to the notation used in Theorem 31. The mapping we need is L = a µ , and σ 2 = b + M 2 ∥w * -w∥ 2 2 which when inserting into Theorem 31 gives the result.
Corollary 34. Under the conditions of Theorem 7, ifγ = A log T T for 1 µ ≤ A , and T is large enough that
Proof. The previous theorem can only be applied when γ ≤ µ 2a , i.e. when A log T T ≤ µ 2a or 2aA µ ≤ T log T . Supposing that is true, observe that (1 -x) T ≤ exp(-xT ), from which it follows that (1 -γµ) T ≤ exp (-µA log T ) = 1 T µA ≤ 1 T .
Theorem 8. Let l be a proper convex and M -smooth function. Let h be a proper closed convex function, and let w * ∈ argmin(l + h) . Let (w t ) t∈N be generated by the Prox-SGD algorithm, with a constant stepsize γ ∈ (0, 1 M ]. Suppose that g t is a quadratically bounded estimator (Def. 5) for ∇l with parameters (a, b, w * ). Then,
where θ def = 1 1 + 2aγ 2 and wT def = P T t=1 θ t+1 w t P T t=1 θ t+1
. In particular, if γ = 1 √ aT
, then
Thus, T = O(ϵ -2 ) iterations are sufficient to guarantee thatE f ( wT ) -inf f ≤ ϵ.
Proof. Let us start by looking at ∥w t+1 -w * ∥ 2 -∥w t -w * ∥ 2 . Expanding the squares, we have that
Since w t+1 = prox γ t h (w t -γ t g t ), we know from Lemma 33 that
and therefore that w t -w t+1 γ t ∈ g t + ∂h(w t+1 ),
where ∂h(w) is the subdifferential of h . Consequently there exists a subgradient b t+1 ∂ ∈ h(w t+1 ) such that
We decompose the last term of (28) as
(29) For the first term in the above we can use that b t+1 ∂ ∈ h(w t+1 ) is a subgradient together with the defining property of subgradient to write
On the second term we can use the fact that l is M -smooth to write
On the last term we can use the convexity of f to write
By inserting ( 30), ( 31), ( 32) into (29) gives
Now using the above in (28), and our assumption that γ t M ≤ 1 , we obtain
We now have to control the last term of (33) in expectation. To shorten our notation we temporarily introduce the operators
Notice in particular that w t+1 = prox γ t h ( T (w t )) . So the the last term of (33) can be decomposed as
We observe that the last term is, in expectation, is equal to zero. This is becauseprox γ t h (T (w t ))-w * is deterministic when conditioned on w t . Since we will later on take expectations, we drop this term now and keep on going. As for the first term, using the nonexpansiveness of the proximal operator (Lemma 33), we have that
Using the above two bounds in (35) we have proved that (after taking expectation)
Injecting the above inequality into (33) and multiplying through by γ t , we obtain
From now on we assume that the stepsize sequence is constant γ t ≡ γ . Now, applying our assumption on the gradient estimator, we have
Inject this inequality into (36), and reordering the terms gives
where we introduced the constant C def = 1 + 2aγ 2 . We will now do a weighted telescoping by introducing a sequence α t+1 > 0 of weights. Multiplying the above inequality by some α t+1 > 0 , and summing over 0, . . . , T -1gives
To be able to use a telescopic sum in the right-hand term, we need that α t+1 C = α t (an idea we borrowed from [54] ). This holds by choosing α t def = 1 C t , which allows us to write
Let us now introduce wT def =
, using that α 0 = 1 , and by using the convexity of f with Jensen's inequality, we obtain
To finish the proof, it remains to compute the geometric series P T -1 t=0 α t+1 which is given by
Using the above and substituting for C = 1 + 2aγ 2 and θ = 1/C gives Note that with this choice of γ we have that
Thus for T ≥ 2 the result of Lemma 32 gives
5 Projected gradient descent with strong convexity Theorem 10. Let W be a nonempty closed convex set. Let f be a µ -strongly convex function, differentiable on W . Let w * = argmin W (f ) . Let (w t ) t∈N be generated by the Proj-SGD algorithm, with a constant stepsize γ ∈ 0, min{ µ 2a
, 2 µ } . Suppose that g t is a quadratically bounded estimator (Def. 5) for ∇f with parameters (a, b, w * ). Then,
Alternatively, if we use the decaying stepsize
In both cases, T = O(ϵ -1 ) iterations are sufficient to guarantee thatE w ∥ T -w * ∥ 2 ≤ ϵ.
Proof. Plugging in one step of Proj-SGD we have that
where we used that projections are non-expansive. Taking expectation conditioned on w t , and expanding squares we have that
t , where we used E [g t ] = f ∇ (w t ) and definition 5, and that f is µ -strongly convex. Taking full expectation and re-arranging gives
Re-arranging we have that
Using γ t ≡ γ constant, and since γ ≤ 2 µ ⇔ 1 -µγ 2 ≥ 0 we have by unrolling the recurrence that
we have that
which concludes the proof of (13).
To prove (14) we will apply Theorem 31, for which we will switch our notation to match that of Theorem 31. That is, substituting
which now holds for γ t ≤ 1 2L and fits exactly the format of ( 23), excluding the additional hat on µ. Thus we know from following verbatim the proof that continues after (23) that for a stepsize schedule of
for t < t *
1 μ 2t + 1 (t + 1) 2 for t ≥ t * for t * = 4 L/ ⌊ ˆµ⌋ we have that E ∥w T +1 -w * ∥ 2 ≤ 16 L/ ⌊ ˆµ⌋ 2 (T + 1) 2 ∥w 0 -w * ∥ 2 + σ 2 μ2 8 T + 1 . After switching back notation L = a µ , σ 2 = b 2 and μ = µ 2 gives (14). 9.6 Projected gradient descent with convexity Theorem 11. Let W be a nonempty closed convex set. Let f be a convex function, differentiable on W . Let w * ∈ argmin W (f ) . Let (w t ) t∈N be generated by the Proj-SGD algorithm, with a constant stepsize γ ∈ (0, +∞) . Suppose that g t is a quadratically bounded estimator (Def. 5) for ∇f at w t with constant parameters (a, b, w * ). Then,
where
Thus, T = O(ϵ -2 ) iterations are sufficient to guarantee thatE f ( wT ) -inf f ≤ ϵ.
Proof. Plugging in one step of the algorithm, we have that
where we used that projections are non-expansive. Taking expectation conditioned on w t , and expanding squares we have that
where we used that g t is an unbiased estimator with bounded squared norm. Taking full expectation and re-arranging gives
From now on we use a constant stepsize γ t ≡ γ . Let C = (1 + aγ 2 ). Next we would like to setup a telescopic sum. To make this happen, we multiply through by α t+1 and impose that α t+1 C = α t .
This holds with α t def = 1 C t . Multiplying through by α t+1 and summing from t = 0, . . . , T -1 we have that
Using that α 0 = 1, dropping the negative term -α T E h w T -w * 2 i , dividing through by 2γ P T -1 k=0 α k+1 and using Jensen's inequality we have that
T -1
gives that
10 Solving VI with stochastic optimization
10.1 Convergence of the algorithms Corollary 12 (Prox-SGD for VI). Consider the VI problem where q w is a multivariate Gaussian distribution (Eq. 2) with parameters w = (m, C) ∈ R d × T d , and assume that this problem admits a solution w * . Suppose that log p(•, x)is M -smooth and concave (resp. µ -strongly concave). Generate a sequence w t by using the Prox-SGD algorithm (Def. 6) applied to l and h (Eq. 16), using g energy (5) as an estimator of ∇l . Let the stepsizes γ t be constant and equal to 1/( p a energy T ) (resp. be decaying as in Theorem 7 with a energy = 2(d + 3)M 2 ). Then, for a certain average wT of the iterates, we have for T ≥ 2 that
Proof. Our assumption on the log target imply that l is smooth and convex (resp. µ -strongly convex) according to Lemmas 18.1 and 19.1. Furthermore, since log p(•, x) is smooth, we know from Theorem 30 that the gradient estimator g energy is quadratically bounded at every w t with respect to constant parameters (a, b, w * ), with a = 2(d + 3)M 2 . Then, because we work with triangular scale parameters, we know from Lemma 19. We have now verified the hypotheses of Theorem 8 (resp. Theorem 7), which gives us the desired bounds.
Corollary 13 (Proj-SGD for VI). Consider the VI problem where q w is a multivariate Gaussian distribution (Eq. 2) with parameters w = (m, C) ∈ R d × S d , and assume that this problem admits a solution w * . Suppose that log p(•, x)is M -smooth and concave (resp. µ -strongly concave). Generate a sequence w t by using the Proj-SGD algorithm (Def. 9) applied to the function f = l+h (Eq. 16) and the constraint W M (Eq. 3), using g ent (7) or g STL (9) as an estimator of ∇f . Let the stepsizes γ t be constant and equal to p 2/(aT ) (resp. be decaying as in Theorem 10) with a = a ent = 4(d + 3)M 2
or a = a STL = 24(d + 3)M 2 . Then, for a certain average wT of the iterates, we have for T ≥ 2 that
Proof. Our assumption on the log target imply that l is smooth and convex (resp. µ -strongly convex) according to Lemmas 18.1 and 19.1. Regarding the entropy, since log p(•, x)is M -smooth, we know from Lemma 18.3 that h is M -smooth on W M . We also know that h is closed convex, since we consider symmetric scale parameters (see Lemma 19.2). Furthermore, the smoothness of log p(•, x) implies that the gradient estimator g ent is quadratically bounded at every w t ∈ W M , with respect to constant parameters (a, b, w * ). This is stated in Theorem 30, where we can see that for g ent we can take a ent = 4(d + 3)M 2 , and for g STL we can take a STL = 2(d + 3)(2K 2 + (2M) 2 ). Since we know from Theorem 4 that K ≤ 2M , we can consider that a STL ≤ 24(d + 3)M 2 . Finally, we know from Lemma 20 that the minimizers of f belong to W M , meaning that argmin(f ) = argmin W M (f ) .
We have now verified the hypotheses of Theorem 11 (resp. Theorem 10), which gives us the desired bounds.
Corollary 14 (Proj-SGD for VI -Gaussian target). Consider the setting of Corollary 13, in the scenario that log p(•, x)is µ -strongly concave, that we use the g STL estimator, and that we take a constant stepsize γ t ≡ γ ∈ 0, min{
µ 2a STL , 2 µ } . Assume further that p(•|x) is Gaussian. Then, E ∥w T -w * ∥ 2 2 ≤ 1 -µγ 2 T ∥w 0 -w * ∥ 2 2 .
Proof. We use the same arguments as in the proof of Corollary 13, but this time we use a constant stepsize, which gives us the bound (13) from Theorem 10, where b = 4(d + 3)K 2 ∥w * -ˆw∥ 2 according to Theorem 4. All we need to conclude is to verify that b = 0 . This comes from the assumption that the target is Gaussian, from which we deduce thatK = 0 (see again Theorem 4).
Proof. This is essentially proved in [ 13,Theorem 13]. We provide a proof here for the sake of completeness, also because our entropy here enforces positive definiteness. So we need to compute
Because the objective here is separated in z and X , we can optimizae those variables separately. It is clear that the optimal z will be z = m . For the matrix component, it remains to solve argmin
Now remember that for a triangular matrix, being positive definite is equivalent to have positive diagonal elements (Lemma 19.3). Because X is a lower triangular matrix, we have-log det X = P d i=1 -log X ii and X ij = 0 when i < j . So our problem becomes
This is a gain a separated optimization problem with respect to the variables X ij . So we deduce that when i ≥ j we must take Ĉij = C ij , and when i = j we need to take
This is the proximal operator of -log, which can be classically computed [55, Example 6.9], and gives the desired result.
Proof. Consider g = δW M to be the indicator function of W M , which is equal to 0 on W M and +∞ elsewhere. Then proj W M = prox g , and we can compute it with [55, Theorem 7.18]. For this, we need to write g(C) = F (λ(C)) , where λ(C) is the set of eigenvalues of C , and F is the indicator function of K := {z ∈ R d : z i ≥ 1/ √ M} . In that way, we see that g is a symmetric spectral function in the sense of [55,Definition 7.12], and the conclusion follows.
1. Algorithm 1 is equivalent to the Prox-SGD algorithm (Def. 6) applied to l and h (Eq. 16), using g energy (5) as an estimator of ∇l and using lower triangular covariance factors.
2. Algorithm 2 is equivalent to the Proj-SGD algorithm (Def. 9) applied to the function f = l + h (Eq. 16) and the constraint W M (Eq. 3), and using g ent (7) as an estimator of ∇f .
3. Algorithm 3 is equivalent to the Proj-SGD algorithm (Def. 9) applied to the function f = l + h (Eq. 16) and the constraint W M (Eq. 3), and using g STL (9) as an estimator of ∇f .
Proof. In this proof, we fix an iteration t ∈ N , consider u t ∼ N (0, I) and we note w t = (m t , C t ).
1. From (Def. 6) we know that w t+1 = prox γ t h (w t -γ t g energy (ut )
) . According to Lemma 38, we have w t -γ t g energy (ut ) = (m t , C t ) -γ t (πt , proj T d (πt u ⊤ t )), where π t = -∇ z log p(Ct u t + m t , x). Moreover, we know from Lemma 35 whatprox γ t h is: prox γ t h (m, C) = (m, R γ t (C)), where R γ t (C) ij = 1 2 (Cii + p C 2 ii + 4γ t ) if i = j, C ij if i ̸ = j. First, we see that with respect to the m variable the proximal operator is the identity, meaning that we indeed have m t+1 = m t -γ t π t . Second, we see that with respect to the C variable we have C t+1 = R γ t (Ct -γ t proj T d (πt u ⊤ t )). It is immediate to see by induction that since we initialize with C 0 ∈ T d , then C t ∈ T d as well. Therefore we can use the linearity of proj T d to write that C t+1 = R γ t proj T d (Ct -γ t π t u ⊤ t ) . The conclusion follows after observing that (R γ t • proj T d ) (C)(i, j) = C(i, j) if i > j, 1 2 C(i, i) + p C(i, i) 2 + 4γ t if i = j, 0 if i < j. 2. From (Def. 9) we know that w t+1 = proj W M (w t -γ t g ent (ut )) . According to Lemma 38, we have w t -γ t g ent (ut ) = (m t , C t ) -γ t (πt , proj S d (πt u ⊤ t -C -⊤ t
)). Moreover, we know from Lemma 36 whatproj W M is:
where R(C) = U DU ⊤ , with UDU ⊤ being a SVD decomposition of C , and D(i, i) = max{D(i, i); M -1/2 } . First, we see that with respect to the m variable the projection is the identity, meaning that we indeed have m t+1 = m t -γ t π t .
Second, we see that with respect to the C variable we have
It is immediate to see by induction that since we initialize with C 0 ∈ S d , then C t ∈ S d as well. Therefore we can use the linearity of proj S d to write that C t+1 = R proj S d (Ct -γ t (πt u ⊤ t -C -⊤ t
)) .
The conclusion follows after setting Ĉt+1 = C t -γ t (πt u ⊤ t -C -⊤ t ), and using the fact that proj S d (C) = (C + C ⊤ )/2 .
3. Use the same reasoning as in the previous item, with the appropriate computation for the STL estimator from Lemma 38.
One could also build an estimator without holding w constant in this way. However, in the case of Gaussian variational families, this turns out to be mathematically equivalent to using a closed form entropy because log qw (Tw (u)) = -1
∥u∥2 2 -d 2 log(2π) -log |C| , which has the same gradient (independent of u ) as h(w) .