Nearly Dimension-Independent Convergence of Mean-Field Black-Box Variational Inference

Thank you for your review!

There is a lot of emphasis on the scaling being nearly dimension-free, which -- unless I'm missing something -- obscures that some of the assumptions likely do not scale sub-linearly in the dimension. For example, it would be reasonable to expect the Frobenius norm of the optimal scale matrix, ${\lVert C_* \rVert}_{\mathrm{F}}$, to grow linearly in the dimension. I don't think this is a major issue, but some discussion of how the other parameters in the bound are likely to grow with dimension in practice is warranted.

Thank you for the suggestion. Indeed, the term dimension-free appears to vary across communities and may not be clear to all readers. In the submission, we tried to avoid the ambiguity by stating “explicit dimension dependence” in both the abstract and the introduction, but not everywhere in the introduction. As such, in the next version, we will spell out “explicit” more often in the introduction and add a mention about implicit dimension dependence in both the introduction and results sections.

It would be nice if the authors could discuss how relevant (at least in terms of the intuition) the authors think their results are to practical implementations. For example, not all implementations use proximal SGD, many implementations use parameterizations other than the linear parameterization, the optimal step size and step size schedule depend on generally unknown parameters and so cannot be set like this in practice and so on. Proofs in these other cases are obviously outside the scope of the paper, but it would be helpful if the authors could discuss which of their assumptions are for technical purposes and which they think are likely actually required of algorithms.

Indeed, the design space is large, and the algorithm design we chose in the paper does not exactly match practical implementations. In terms of parametrizations, we briefly mention Lines 124-128 that alternative “non-linear parametrizations” are commonly used. However, as mentioned, these usually rule out obtaining convergence guarantees [1] without setting a compact domain [2]. For optimization algorithms outside of proximal SGD, projected SGD is an obvious candidate. However, as mentioned in Lines 145-148, this requires setting an arbitrary compact domain, and both the analysis and convergence guarantees are analogous up to a factor of x2. (See [3].) However, we agree that a more holistic investigation and discussion over all these design choices would be useful for the community. For this, we plan on writing a paper performing an empirical comparison that would hopefully act as a guide to the various pros and cons of these design elements.

For setting the step size schedule and initial step size, we agree that a discussion about practical considerations would be useful. We will mention the fact that the act of parameter tuning can be seen as identifying the step size (schedules) that depend on the unknown problem constants.

I found the proof of Proposition 4.2 hard to follow. In particular, it was unclear which parts of the proof are just for developing intuition for how they chose $H_{\mathrm{worst}}$ and which parts of the proof are actually rigorous. Relatedly, why is there a $z^{w}$ in the middle of the display after line 911? (And similarly in the second line of the display after line 921)?

Thank you for pointing this out! We agree that the use of $z^w$ here is extremely confusing. The original intent was to convey that the result of the integral $\int^1_0 H_{\mathrm{worst}}(z^w) (z - \bar{z}) \mathrm{d}w$ is the same as the pointwise evaluation of $H_{\mathrm{worst}}(z^w) (z - \bar{z})$ with respect to any $w \in [0, 1]$, but the notation is inappropriate. We will fix this in the next version. In terms of the organization of the proof, we will add more signpost text to clearly demarcate where the requirements for the worst-case example are stated and the actual proof begins.

I found the phrasing around the formal Assumptions in the main text to be somewhat confusing on my first read. For example "Assumptions" 2.1 and 2.2 read more like definitions, and the actual assumption is that various functions have those properties.

Thank you for the suggestion. We agree that “Definition” would have been better here, and have in fact changed this shortly after our initial submission.

For example, I think Assumption 2.7 should always hold for mean-field VI with location-scale families with linear parameterizations. I understand that it is being used to set up a more general result, but some changes to the wording might make that clearer.

Thank you for the suggestion. We agree that some text to explain the relevance to VI would have been better here. We will improve this in the next version.

And thank you for spotting the various typos in the proofs! We will make sure to fix these in the next version.

1. Kim, K., Oh, J., Wu, K., Ma, Y., & Gardner, J. (2023). On the convergence of black-box variational inference. Advances in Neural Information Processing Systems, 36, 44615-44657.
2. Hotti, A. M., Van der Goten, L. A., & Lagergren, J. (2024, April). Benefits of non-linear scale parameterizations in black box variational inference through smoothness results and gradient variance bounds. In International Conference on Artificial Intelligence and Statistics (pp. 3538-3546). PMLR.
3. Domke, J., Gower, R., & Garrigos, G. (2023). Provable convergence guarantees for black-box variational inference. Advances in neural information processing systems, 36, 66289-66327.

Thank you for your review!

The paper is highly theoretical and lacks numerical experiments or synthetic illustrations. I think including illustrations would help the reader follow the concepts more easily.

Thank you for the suggestion. We are open to suggestions for any illustrations or experimental results that would help readers understand the technical content of the paper. Please feel free to suggest anything the reviewer would find useful!

While the paper focuses on the reparameterization gradient and provides a technical justification for its use, it does not adequately explain why this estimator is prioritized over alternatives like the score function gradient.

As mentioned in Lines 154-155, we focus on the reparametrization gradient as it is empirically known to perform much better than the score gradient [1,2]. This fact has also been theoretically established in [3] for Gaussian targets. In addition, as mentioned in Lines 76-78, most probabilistic programming frameworks use the reparametrization gradient by default whenever applicable. Therefore, we believe it is fair to say that the reparametrization gradient is the most popular choice in practice. However, we do agree that the score gradient also warrants some attention in terms of a general theoretical analysis. But this would be worth its own paper.

Providing additional context or justification could strengthen the relevance of the paper's contributions.

Thank you for the suggestion. As mentioned in Section 1, a quantitative convergence analysis of mean-field BBVI is practically relevant regarding the statistical-computational trade-off. That is, it is useful to know how long it would take to obtain a full-rank approximation compared to a mean-field approximation. Our paper implies that, in the high-dimensional regime, one should roughly expect a factor of $d$ difference. (See also the reply to Reviewer P294.)

Typo in lines 144 & 185: It looks like the acronym "PSGD" is a typo. Should it be SPGD instead?

Yes, thank you for spotting this! It should have been SPGD.

Should the matrix $C$ in Definition 2.6 be required to be positive definite and diagonal?

This is somewhat a matter of convention. Mathematically speaking, $C$ only needs to be invertible to result in a well-defined distribution since the resulting covariance $CC^{\top}$ will be positive definite even if $C$ has negative eigenvalues. However, it is indeed common to define $C$ to have strictly positive eigenvalues, since, in the 1-dimensional, allows to interpret $C$ as a standard deviation. For our purpose, this requirement isn't necessary for defining location-scale families, so we didn't mention it Definition 2.6. Though we should clarify that $C$ at least needs to be invertible. We will fix this in the next version and thank you for pointing this out.

1. Titsias, M., & Lázaro-Gredilla, M. (2014, June). Doubly stochastic variational Bayes for non-conjugate inference. In International conference on machine learning (pp. 1971-1979). PMLR.
2. Kucukelbir, A., Tran, D., Ranganath, R., Gelman, A., & Blei, D. M. (2017). Automatic differentiation variational inference. Journal of machine learning research, 18(14), 1-45.
3. Xu, M., Quiroz, M., Kohn, R., & Sisson, S. A. (2019, April). Variance reduction properties of the reparameterization trick. In The 22nd international conference on artificial intelligence and statistics (pp. 2711-2720). PMLR.

Thank you for your review!

line 144: PSGD is never defined, and the sentence "Instead of using PSGD, one can also use projected SGD" is a bit confusing since its easy to think that PSGD refers to projected SGD.

We apologize for the confusion, this is a typo and should have been SPGD for stochastic proximal gradient descent.

Prop 2.11 (and elsewhere): I don't really like the $T \geq O(\ldots)$. Maybe $T \geq \Omega(\ldots)$ is better?

Thank you for the suggestion. Indeed, this is a common point of confusion, but in this case, big-O is correct, whereas big-$\Omega$ is not quite. This is because $T \geq O(\ldots) \Rightarrow {\lVert \lambda_0 - \lambda_* \rVert}_2^2 \leq \epsilon$ implies that “the smallest $T$ one can choose to guarantee the $\epsilon$-accuracy condition grows no faster than $\ldots$.” Therefore, this is an “upper bound” on the sufficient number of iterations one needs to take. On the other hand, if one uses big-$\Omega$, it would read, “the smallest $T$ one can choose grows faster than $\ldots$,” which is a statement of necessity rather than sufficiency.

I wonder if it would be possible to verify the results experimentally?

Thank you for the suggestion. We agree that empirical verification of theoretical results is worth doing. Now, the fact that the mean-field approximation results in better scaling has been known empirically for a long time. Therefore, we found that simply showing that the mean-field approximation scales well is not empirically meaningful. Instead, it would be most interesting to verify that the worst-case $O(\log d)$ dimensional scaling can be observed in practice. However, as noted in Remark 4.3, the theoretical worst-case example that achieves this $O(\log d)$ scaling may not exist. We tried various informed guesses for $f$ to see if we observe an increasing scaling with respect to $d$. However, all examples either resulted in a constant scaling or even decreasing (!) scaling with respect to $d$. Therefore, we couldn’t design an interesting enough experiment to verify the results in the paper.

Are there any concrete useable/practical lessons from this paper, other than one need not worry so much about dimensionality when the data is expected to be not heavy tailed.

Again the main point keeping me from a strong accept is my lack of certainty regarding whether the impact, originality or usefulness is good enough to warrant the top reviewing score. This may just be because I am not super comfortable with the topic. Anything you could do here to elaborate for a non-expert would be helpful.

This is an important point. Our results say a bit more than one needs to worry about dimensionality. As mentioned in Section 1, we believe VI is the most relevant for obtaining a controllable statistical-computational trade-off. That is, we would like to know how much speed up we can obtain by giving away some amount of statistical accuracy. For this, a quantitative convergence analysis is necessary for variational families of various degrees of expressivity. With this in mind, our quantitative convergence analysis implies that, in the high-dimensional regime, our results suggest moving from a full-rank family to a mean-field family will require running SGD roughly $d$ times longer with a smaller step size. For variational families with an expressivity in between mean-field and full-rank, we can expect to see a dimensional scaling in between $O(\log d)$ and $O(d)$, where we can use the $O(\log d)$ result as a baseline for estimating the amount of necessary computation time. (See also the last paragraph of Section 5.2.)

Thank you for your review!

Theorem 3.2 states that there exists a switching time and initial step size for which the result holds. Can you comment on the practical significance of having to choose these values?

Yes. In practice, it is no secret that any application of SGD requires extensive tuning of the step size schedule. Indeed, Proposition 2.11 reflects this fact through the use of the two-step step size schedule in Eq. (4). Thus, in that sense, our theory is qualitatively in agreement with practice.

Furthermore, in previous works [1], the switching time was selected only to depend on the condition number $\kappa$. This suggests that it is sufficient to adapt the step size schedule (or more precisely, the switching time) according to the target problem only. Proposition 2.11 demonstrates that, in order to further optimize the dependence on the initialization ${\lVert\lambda_0 - \lambda_*\rVert}_2$, it is necessary to also adapt the switching time according to the properties of the gradient estimator ($\sigma^2$). That is, our results suggest that whenever the gradient estimator is swapped, the step size schedule should also be adjusted accordingly to obtain faster convergence. We will make sure to mention this fact in the next version.

1. Gower, R. M., Loizou, N., Qian, X., Sailanbayev, A., Shulgin, E., & Richtárik, P. (2019, May). SGD: General analysis and improved rates. In International conference on machine learning (pp. 5200-5209). PMLR.

Thank you for the reply to my question. I maintain that this paper is a clear accept.