Joint control variate for faster black-box variational inference

Summary

The authors introduce a "dual" control variate for reducing gradient variance in black-box variational inference in the context of models that admit data subsampling (i.e. exhibit the required conditional independence). The dual control variant is joint in that it simultaneously addresses the two sources of Monte Carlo variance in ELBO approximations: that due to latent variable sampling and that due to data subsampling. The basic idea, in essence, is to leverage linear or quadratic ELBO approximations (which admit closed form evaluation for e.g. gaussian mean field variational families) in conjunction with a running average of gradient estimates for each data point using the variational parameters from the last iteration in which each data point was accessed. Experiments demonstrate that the proposed method can substantially reduce gradient variance (in particular that due to data subsampling, which is often dominant), thus yielding better ELBO optimization (both w.r.t. wallclock time and final ELBOs obtained).

Strengths

The main strengths of this submission include the following: - it addresses a general problem of relatively wide interest in the NeurIPS community (namely, how best to do black-box variational inference) - it addresses a particular component of that problem that is often somewhat overlooked compared to other aspects (namely how best to do optimization) - the experiments are pretty convincing in establishing the efficacy of the method - the suggested method is technically sound and would appear to be pretty simple to implement - the discussion in Sec 5 and the variance analysis (appendix B) help the reader conceptually place the proposed control variate alongside other alternatives

Weaknesses

The main weaknesses of this submission, as I see them, are the following: - the notation is a bit confusing in some places - some of the limitations of the method and/or extensions to more general problem settings are either not discussed or are insufficiently discussed Let me expand on these points: While the notation for this kind of paper will necessarily be somewhat clunky given the various sources of sampling variability that have to be carefully tracked, I think some improvements are possible. In particular I find the choice of M for the running gradient particularly suboptimal. Since capital N is a positive integer and little n and little m are used to index integers, one might expect that M is also a positive integer. I suggest that M be replaced by something like G(bold $w$) to avoid this confusion and to emphasize that G depends on the "table" of $w$s. The authors consider a generic but still somewhat limited problem specification, in particular they do not consider local latent variables, model learning, or amortization. More discussion of these points would probably be of considerable interest to the reader (for more discussion see below). In addition one of the weaknesses of this method is the potentially large memory requirements, which are O(Nd) where d is the latent dimension. This needs to be *very clearly emphasized*.

Questions

- Do you expect your method would work more or less unchanged for hierarchical models that feature both global and local latent variables? Or would the algorithm need to be adapted to that setting? - Arguably besides the possibility of data subsampling for appropriate model classes, some of the biggest selling points of variational inference are amortization and model learning (i.e. learning point estimates for model parameters). - I guess you could support model learning directly without any change to your algorithm, although I suspect the performance of your control variate would tend to decrease since now in addition to "$w$ drift" you'd also have "$\theta$ drift" (where theta is the model parameter). Can you please comment? - I guess you could also support amortization (e.g. of local latent variables in hierarchical models) without any changes to your algorithm, although since amortization typically involves a neural network with a largeish number of parameters, this would make the memory requirements needed to cheaply compute the running average M intractably large. Can you please comment? - typo: appproximation - It would be great to see some ablation studies comparing first order to second order taylor approximations. How important is going to second order in practice? - What kind of runtime performance would you expect if used a multivariatiate gaussian distribution for the variational distribution? Would the dual estimator still be ~1.5 to ~2.5 times slower than computing control variant free gradients? Or would the gap increase? - It should be possible to do a more fine-grained theoretical analysis for a model that admits more analytic control, e.g. bayesian linear regression. Such an analysis could be particularly valuable in further delineating the regimes in which the dual variate is expected to perform best. - You state that "the total gradient variance is often dominated by variance from μ" [line 207]. Do you have any intuition for this claim? - Comment: For a large dataset initializing the running mean M at the beginning seems like it could be a waste of compute time. It might make more sense to not use the control variate during the first epoch and use the first epoch to collect M.

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

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

As discussed in the weaknesses section, I believe some of the limitations of the method (e.g. w.r.t.~memory requirements and likely trouble with amortization) should be discussed and/or better emphasized.

Summary

This paper proposes a new control variate for black-box variational inference. In particular, the proposed "dual" control variate attempts to reduce the subsampling noise and Monte Carlo noise at the same time. For this, the paper utilizes an incremental gradient-like scheme. The performance of the new control variate is empirically verified on Bayesian inference tasks with large datasets.

Strengths

* While control variates have been an active area of research for BBVI, reducing the subsampling noise has certainly been a problem that hasn't been addressed. In fact, the paper shows that conventional control variate solutions do not solve this problem *at all*, despite the fact that subsampling is a major source of variance. * The paper motivates the latter point by empirically computing the subsampling noise. Overall, the paper conveys the motivations for the proposed control variate very clearly. * The proposed control variate based on incremental gradients seems fairly simple to implement, but not trivial. Thus, the proposed strategy has clear technical contribution. * Empirical evaluations are thorough and adequate to show the superiority of the proposed control variate.

Weaknesses

* Given that the paper builds on top of incremental gradient methods, which, as typical of the optimization community, comes with a heap of theoretical tools for rigorous analysis. Given this fact, it would have been amazing if the paper also provided some rigorous analysis of the proposed strategy. * There seems to be some room for improvement in terms of paper organization and notation. The notation can be quite confusing at times, for example, lots of things happen behind innocent-looking superscripts. More on this below.

Questions

### Major Comments * What was the motivation behind calling the control variate "dual"? At least to people who have some experience in optimization, I believe the word dual would automatically trigger the notion of Lagrangian duality. How about "doubly" in accordance with the term "doubly stochastic?" * The section organization seems to have room for improvement. For example, a general explanation of "general" control variates appears way back in Section 5.1. Considering readers that are not familiar with control variates, it would be better to start with a high-level explanation of how a control variate is supposed to work. A good place would be after Section 3, where the variance of BBVI is motivated to be a problem, and before Section 4 where a control variate is first introduced. * Similarly, I felt that the whole discussion in Section 5.1 and 5.2 would have been a very good way to inspire the readers before actually unrevealing the dual control variate. While reading Section 5.1 and 5.2 I had to constantly go back to Section 4 to remind how the dual control variate attempts to solve the problem. * Algorithm 1 is not very helpful to precisely understand the implementation. For example, in "Compute the control variate" what "using ... $=M$ " mean? Does this mean we simply plug $M$ in place of $ E_m E_{\eta} \nabla \tilde{f} \left(w^m; m, \eta\right) $ ? Then why denote $ E_m E_{\eta} \nabla \tilde{f} \left(w^m; m, \eta\right) $ ? I think here the authors need to clearly distinguish between values and computed quantities. Similarly, the notation for the running mean $M$ is too uninformative. I think something like $\tilde{g}$ or $\hat{g}$ would be more appropriate. This is also useful to understand the dimensionality/domain of each variable. * Line 207: Doesn't this contradict the conclusions of [2,3]? From my understanding of [2], the $\mathcal{O}\left(d\right)$ dimensional dependence comes from the gradients with respect to the scale. * Figure 4 Tennis: Do you have any insight as why the ELBO for the dual control variate peaks and then decreases? This seems quite weird given that the control variate is unbiased. * Caption of Figure 4: This is perhaps a minor point, but could you elaborate why gradient noise of Tennis is correlated? Is this because reshuffling was used here instead of iid subsampling? * Section 6: Please include a discussion of other non-dual control variates since this works, in the broader scheme of things, is a control variate paper. Also, why not include Section 5.3 in Section 6? ### Minor Comments * Eq. (1) I was a little bit confused at first because I am so used to see the likelihood adjustment ratio to be N/B, where B is the batch size. I think it would be useful to either use N/B overall so that the notion that we're dealing with batches is more clear, or to mention in the text that the batch size is assumed to be 1. * Line 59: The classic paper that first proposed doubly stochstic BBVI was [1]. I recommend adding it here. * Table 1: How about using a bar plot instead of a table? I think it would be more appropriate since the relative magnitude is key here. * Line 118: I think it would be better to explicitly show that the performance of a control variate is entirely driven by the covariance between the estimate and the control variate. ### References * [1] Titsias, Michalis, and Miguel Lázaro-Gredilla. "Doubly stochastic variational Bayes for non-conjugate inference." In International conference on machine learning, pp. 1971-1979. PMLR, 2014. * [2] Domke, Justin. "Provable gradient variance guarantees for black-box variational inference." Advances in Neural Information Processing Systems 32 (2019). * [3] Bhatia, Kush, Nikki Lijing Kuang, Yi-An Ma, and Yixin Wang. "Statistical and computational trade-offs in variational inference: A case study in inferential model selection." arXiv preprint arXiv:2207.11208 (2022).

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

3 good

Presentation

2 fair

Contribution

3 good

Limitations

Yes.

Summary

The paper presents a method for variance reduction in stochastic gradient estimation in doubly stochastic variational inference where there exists two sources of variance: (i) Monte Carlo noise when sampling from the variational distribution and (ii) gradient variance due to the minibatch sampling. The authors consider reparametrizable Gaussian variational inference and introduce a control variate that tries to reduce simultaneously both sources of variance. This control variate combines previous ideas, such as a second order Taylor expansion of the function as in [19], and it seems to reduce the variance in the presented experiments.

Strengths

The paper is very clearly written and all derivations appear to be correct. It contains also intuitive discussions why the proposed "coupled" control variate can be useful as opposed to other control variates that deal separately with each source of variance. Certainly variance reduction is a very important topic for stochastic gradient estimation in variational inference, and the paper proposes a potentially useful method. The experiments provide many details including running times.

Weaknesses

I was not so impressed by the experimental results for two reasons. Firstly the models used are quite small and it would be useful to include e.g. a big neural network. Secondly, I am not sure if the comparison is done in a fair way for methods such the "naive" method. This is because the "dual method is more expensive and requires more gradient evaluations, such as the ones for the numerical approximation of the Hessian-vector products. Given that computations are dominated by the number of gradient evaluations, a fair comparison should try to match this number across different estimators. For example, for the "naive" estimator someone could increase the minibatch size so that the number of gradient evaluations matches the one of the "dual" method.

Questions

Why the control variate in equation (16) controls only the variance for the $\mu$? Is this discussed in [9]?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

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

3 good

Presentation

3 good

Contribution

3 good

Limitations

The limitations are explained above regarding the current experimental comparison.

Summary

Existing stochastic methods for black-box variational inference only attempt to reduce the noise either due to data subsampling or Monte-Carlo sampling of the expectation. This paper proposed a new "dual control variate", which addresses both types of noise at the same time. In experiments, the proposed control variate is shown to perform favorably to existing baselines.

Strengths

1) The paper disentangles the effects of noise through the data and noise from Monte-Carlo estimation of the expectation in the design of the new dual control variate, which addresses both at the same time. 2) The proposed control variate is shown to greatly improve performance on the considered examples, at seemingly minimal overhead. If the x-axis in the plots would have been wall clock time, this would be clearer to see.

Weaknesses

The main weakness of the paper is in my opinion that the experimental evaluation is a bit lacking: a) larger experiments would be desirable, see questions; b) the plots were hard to read (difficult to distinguish the red and green curves), and could be wrt. wall clock time and not iterations. The algorithm looks promising, and I like that it can be used in a "black-box" fashion with any optimizer (e.g. Adam); but I believe it is not ready yet in its current state and could require a more thorough experimental evaluation to warrant a clear acceptance.

Questions

Would the proposed dual control variate approach also work for VI in neural networks? Experiments on modern neural networks (e.g. ResNets, transformers) could be a way to convincingly demonstrate the utility of the proposed method. How much additional overhead (memory & runtime) does it have when compared to, say, Bayes-by-Backprop when actually implemented in practice? From Table 2 it seems that one would require three backpropagation passes rather than a single one for the naive method -- it would be interesting to see whether this additional effort has a real practical benefit (e.g. higher accuracy, faster convergence wrt wall clock time).

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

1: Your assessment is an educated guess. The submission is not in your area or the submission was difficult to understand. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

All limitations are addressed.

Summary

The paper addresses the drawback of the black-box variational inference framework for Bayesian posterior inference by proposing dual control variate that is capable of jointly reducing the variances from both data subsampling and Monte Carlo sampling. The experimental evaluations on various datasets demonstrates reduced variance and improved optimization.

Strengths

The paper is fairly written well. The background of the black-box variational inference (BBVI) is nicely described. The doubly-stochastic optimization problem in BBVI's gradient estimation is clearly explained involving two sources of randomness - Monte Carlo sampling from the variational posterior and data subsampling from the full dataset. The proposed dual control variate that jointly controls Monte Carlo and subsampling noise in BBVI to create approximations of the target for each datum where the Monte Carlo noise can be integrated exactly, addressing both forms of noise and interactions between them. Experimental evaluation and visualization of dual estimator.

Weaknesses

- Although the idea is novel but its usefulness/impact could have explained well. - The approximation function for the gradient estimators g_cv and g_dual could have been clarified. - It is unclear how the noise in Monte Carlo sampling influences the noise in data sampling. Is there a way to measure it? - What is the role of Beta in eqn(14 - 15). Is it experimentally evaluated?

Questions

I am not fully familiar with this subject. However, I did get idea about the problem and how the propose idea could address it. Therefore, I will consider the reviews of other reviewers. The main problem is the writing/explanation of the technical terms to make readers to understand the problem. The paper started well and somehow there is a lost connection between experimental evaluation to demonstrate the impact of the dual control variate real-world applications

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, 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

2 fair

Contribution

3 good

Limitations

There is no issue with potential negative societal impact. However, the article does not address limitations (if any).

Further Response

Thank you for the clarifications. I am happy with the updated pseudocode addressed to Reviewer NBCj. Reorganizing the sections would make the paper much easier to read and free up some space (I see some redundancy from introducing/explaining the concept of control variates multiple times.) Nevertheless, I am now happy to increase my score. As per the naming issue, I have looked up some synonyms to "double" or "two" and found words like "duo" or "bi-". Or, given the paper addresses all of the stochasticity present in BBVI, maybe something like "total" or "complete" could also be considerable. Also, given that Titsias chose "double," I think my original suggestion of "doubly" is also fine as it is not exactly the same. Hope this is helpful!

Thank you for your reply

Thanks for adding the comparison by matching the number of gradient evaluations. This makes the experimental evaluation more complete.

Thank you for the clarifications

Thanks for addressing my concerns. Don't forget to improve the content organization and presentation clarity in the revised paper.