On the Convergence of Black-Box Variational Inference

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. The paper notes that the theoretical results show that for the location-scale family, nonlinear scale parameterizations are suboptimal (but widely used in practice). This claim is tested experimentally on a synthetic and realistic problems. The empirical results confirm this claim and suggest that proximal methods converge faster.

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. The 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. There 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. 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. Finally, 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.

Questions

- What is meant by the "infinite sum structure" in line 115? - 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? - 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? - 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. Typos: - line 3: comma should go after the parentheses - line 14: "box" missing - line 60: did you mean "sans-serif"? - line 212: remove "this" - line 230: Theorem 12 is not in the main part of the paper.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

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.

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.

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. * 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$. * 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. * 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.

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

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. * The formatting of lines 67-70 is odd, and perhaps these quantities should just be defined in-line. * In line 102, I believe an additional “diag” is missing: should it be $diag(\phi(s)) = diag(\phi(s_1),\dots, \phi(s_d))$? * 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. * 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$? * Line 130-132 has a typo in phrasing “have been”, “alienating” * 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. * The statement of Theorem 2 differs in the main paper and the supplement, this should be corrected. * Line 174, “becomes” is a typo

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

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.

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.

Strengths

1. The paper is solid and covers a lot of theory behind VI, Figure 1 is quite informative. 2. The figures are great and descriptive and support the theory and text. 3. I am not good with theory, but I found the theorems and inequalities fine 4. Covers relevant and contemporary literature quite well, although a paper with similar ideas and claims recently appeared on arxiv.

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. 2. 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. 3. 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.

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. 2. 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$ ? 3. 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 ? 4. 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 ? 5. Why was the exp conditioner/transformation not used in expts. for Figure 4, when it is the most popular one in my opinion ? 6. 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 approximate MCMC' since MCMC is aymptotically exact while MFVI is not and can be arbitrarily different from target.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Minor 1. Typo in line 115:'inifinite' 2. For the uninitiated, maybe introduce $L_{h}$ as Lipschitz smooth, if it does .. ? 3. 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, I 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.

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.

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

Weaknesses

1) The results and experiments are incremental (but nontrivial) extensions of what is presented in (Domke 2020). 2) The paper missed to cite the recent line of works * https://arxiv.org/abs/2205.15902 * https://proceedings.mlr.press/v202/diao23a.html * https://proceedings.mlr.press/v202/diao23a.html which 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. 3) 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.

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? 2) 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. 3) 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.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

2 fair

Limitations

All limitations are addressed.

Reply to rebuttal

Thanks to the authors for point by point response to my questions. I only wanted to emphasize that treatment of both objective and gradients is important in analysis for convergence, which they agree with. I like the new proposed title much more than the previous one which was my main concern and I must point out that other reviewers also felt the same. Then, I also felt that the authors tried too hard in overselling itself a bit too much in places when they actually didn't need to, which other reviewers also noticed. I am other wise happy to revise my rating and recommend acceptance. I think I am also satsfied with their response in comparison to Domke's old and recent work. Maybe they could have a joint session with the authors if the both papers are accepted. I remember a similar thing happening in last year NeurIPS when two accepted papers had conflicting results and conclusions, which is not necessarily the case here. From practise, I know that BBVI is so much dependent on initialization that it can be hard to say a general statement about parameterisations, and this is talking about exclusive KL, the other commonly used divergence objectives are even more wild.

Rebuttal response

I thank the authors for the detailed response. Most of my points were minor and I am glad to see some will be addressed by the authors for clarity. The title change and the limitations section will help practitioners find this work and understand how the results affect implementations of BBVI.