EigenVI: score-based variational inference with orthogonal function expansions

Summary

This work proposes a new variational family based on orthogonal function expansions from the quantum mechanics literature. Furthermore, the paper proposes to find the basis coefficients by solving a score matching problem, where the problem cleverly reduces to an eigenvalue problem. The overall approach, called EigenVI, is compared to recent score-matching-based Gaussian VI algorithms and BBVI.

Strengths

* The idea of using orthogonal function expansions is new, although it reminds me of some other non-parametric/almost parametric variational families. * The fact that the score-matching problems reduce to an eigenvalue problem is beautiful in its own right. * The writing of the article is the best in my batch.

Weaknesses

* While the authors are upright about the major limitation of the work, scaling with respect to dimensionality, this, unfortunately, strongly limits the practicality of the work. To me, for a VI algorithm to be useful, it should at least be scalable with respect to $d$, the dimensionality of the posterior, or $n$, the number of datapoints. Unfortunately, the proposed approach seems to achieve neither. In particular, the eigenvalue problem formulation is very cute but it isn't clear if subsampling results in an unbiased objective to be scalable with respect to $n$. (Please correct me if this is not the case!) * The evaluation is not adequate for a work proposing a new variational family. The two baseline variational families considered here are the full-rank Gaussian variational family and normalizing flows. There is a rather long history of developing less-parametric variational families in VI. For instance, boosting VI [1], copulas [2,3], splines [4], semi-implicit families [5], mixtures [6], and many more. These works should be the main focus of the evaluation, not Gaussian VI. * Unless I am not mistaken, the method suffers in the case the posterior is not *a priori* standardized. It doesn't seem like such standardization would be possible unless some preliminary inference has already been performed. This would defeat the point of using a VI algorithm. If one has already obtained preconditioners, why not just provide them to an MCMC algorithm instead? * The paper doesn't discuss how to solve the eigenvalue problem. Given that we only need to solve a specific eigenvalue problem, there might be clever ways to solve it in a scalable way, maybe by using randomized linear algebra methods, which could technically be interesting. Overall, I hope to see at least some evidence that the proposed methodology will eventually lead to a practical algorithm in the future. For instance, whether subsampling results in an unbiased/consistent algorithm. If this is the case, I would be happy to re-assess the paper. ## Additional Comments * The Fisher divergence and the score-matching problem were introduced to the machine-learning community through the work of Hyvarinen [7], which isn't cited. * Related works: A discussion on the literature of non-parametric/"less-parametric" VI algorithms seems necessary. * Line 30 "One thread of BBVI research focuses on Gaussian variational approximations. (ADVI, BBVI)": I wouldn't quite agree with this statement since ADVI is applicable to any family where the reparameterization gradient is applicable. In fact, when normalizing flows are used with ELBO maximization, this is pretty much the same as ADVI. Furthermore, the original BBVI paper [8] requires conjugacy, meaning that one cannot simply use Gaussian variational families at all in a lot of cases. * A lot of the references are citing the arXiv version. Berg et al. 2018 was published at UAI; Dinh et al. 2016 was published at ICLR 2017; Giordano et al. 2023 was published in JMLR; and etc. Please thoroughly review the reference section. * Line 39, Section 2.1: Citing something when mentioning orthogonal functions would be better so that readers interested in the topic could have a look into more in-depth material. ## References 1. Campbell, T., & Li, X. (2019). Universal boosting variational inference. NeurIPS 2. Smith, M. S., & Loaiza-Maya, R. (2023). Implicit copula variational inference. Journal of Computational and Graphical Statistics, 32(3), 769-781. 3. Han, S., Liao, X., Dunson, D., & Carin, L. (2016, May). Variational Gaussian copula inference. In AISTATS. 4. Shao, Y., Shan, N. Y., & Feng, T. (2024, April). Nonparametric Automatic Differentiation Variational Inference with Spline Approximation. In AISTATS. 5. Yu, L., & Zhang, C. (2023). Semi-Implicit Variational Inference via Score Matching. arXiv preprint arXiv:2308.10014. 6. Lambert, M., Chewi, S., Bach, F., Bonnabel, S., & Rigollet, P. (2022). Variational inference via Wasserstein gradient flows. NeurIPS. 7. Hyvärinen, A., & Dayan, P. (2005). Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4). 8. Ranganath, R., Gerrish, S., & Blei, D. (2014, April). Black box variational inference. In AISTATS.

Questions

* In Line 103 to 105, it is stated that one can circumpass distributions with constrained supports. Given bijectors (ADVI), constrained supports are rarely a problem for VI algorithms these days. But dealing with *discrete support* is still a major challenge. Would there be a way to use the proposed method for problems with discrete support? * Eq (10): Why is importance sampling used here when one can sample directly from $q$?

Rating

5

Confidence

4

Soundness

3

Presentation

4

Contribution

2

Limitations

The paper would benefit from a dedicated limitation section.

Summary

The authors propose EigenVI, a new method for black-box variational inference (BBVI). The method uses the Fisher divergence (rather than the more typical reverse KL divergence) and a variational family based on orthogonal function expansions (which I've never seen used before). Advantages of the approach include the ability to model complex distributions (multimodal, skewed, heavy-tailed, etc.) while also giving closed-form calculations for lower-order moments as well as easy posterior sampling.

Strengths

The paper is exceptionally strong in originality, clarity, and significance. For originality, I have not previously seen a paper in machine learning using orthogonal function expansions to define a variational family. The authors state that the approach comes from quantum mechanics, an atypical source domain for a machine learning paper. The construction is interesting, and normalizes elegantly. Given the importance of orthogonal bases throughout mathematics, it is likely that this approach to variational inference will connect with many fields of academic inquiry. Thus, this approach has the potential to greatly enrich research in variational inference. For clarity, the exposition is exceptionally clear nearly everywhere (although Sec 2.3 could use a couple iterations of rewriting to streamline the narrative). In particular, orthogonal function expansions and score-based divergences were introduced in an exceedingly clear (and interesting) manner. The high-level advantages of their method vs. others in the literature were also exceedingly clearly presented. For significance, the approach is broadly useful (at least for low-dimensional posteriors); this is inherited from the fact that BBVI is constructed so as to be applicable to a wide set of applications without requiring any pen-and-paper calculations. Moreover, by avoiding gradient-based inference in favor of solving an eigenvalue problem, their method would additionally remove the common headache of having to specify or learn a learning rate.

Weaknesses

Some weaknesses: (1) As mentioned by the authors, the number of basis functions scales exponentially in the dimensionality of the posterior, and therefore so does inference time. Given the restriction to low-dimensional posteriors, why use VI at all, rather than MCMC sampling? It would have been nice to see this addressed in the exposition. It would have also been nice to see an experimental comparison showing some advantage over using MCMC samplers. For instance, do the authors expect a computational advantage of their approach over MCMC methods? If so, it's hard to trust this since computation times are not given. (2) Following up on the preceding point, it is completely unclear how long it takes to run this method, since no computation times are given relative to competitors. On pp.9, the authors say, "In these examples, we standardize the target using either GSM or BaM before applying EigenVI; for this reason, we do not compare the costs of these methods to EigenVI." I found this sentence quite puzzling. Presumably the authors intend to argue that since EigenVI must take strictly longer than GSM or BaM, there's not point in reporting the computation times. If so, I don't see how the conclusion follows. A practitioner would want to know how MUCH additional compute is required to obtain any advantages. Moreover, how does the overall (combined) computation time compare to other approaches, such as VI with a normalizing flow family or standard ADVI (Kucukelbir et al. 2017)? (3) The method seems fairly complex to run, since a competitor BBVI method must be run simply as a preprocessing step to standardize the distribution. (They use either Gaussian Score Matching or Batch-and-Match VI, per pp. 8.) (4) A major selling point emphasized by the authors is that unlike gradient-based approaches, EigenVI does not require tunable hyper-parameters, such as learning rates or termination criteria. However, EigenVI is clearly sensitive to (a) the number of basis functions and (b) the number of importance samples. Indeed, Figure 4 (c) demonstrates how inference can go catastrophically bad when the number of importance samples is not sufficiently large relative to the number of basis functions. This problem is reminiscent of the problem of a poorly chosen learning rate in standard ADVI (Kucukelbir et al., 2017). Can the authors provide any guidance on how to choose a good number of basis functions for a given problem, and then, given that, how to choose a good number of importance samples? A theoretical result would be nice, although admittedly it's unreasonable to expect that during the rebuttal period. Do they at least have any practical insight to share? (5) I feel that the authors undersell their own method by neglecting to reinforce the primary selling points in the experiments section. For example, the authors mention that their approach allows sampling, closed-form moment calculations, and modeling of distributions with mixed support. However, none of these things appear explicitly in the experimental section. Similarly, the authors emphasize that the normalizing flow competitor is "often sensitive to the hyperparameters of the flow architecture, optimization algorithm, and parameters of the base distribution" (pp.6). Can the authors demonstrate that the posterior approximation deteriorates along these dimensions more easily than their own approach deteriorates as a function of number of importance samples and basis functions?

Questions

Clarifications: (1) pp.6: What is the proposed advantage of EigenVI over Zhang et al.'s energy-based approach with a closed-form solution? (2) pp.7, Figure 4: If the normalized Hermite polynomials contain Gaussian distributions as a "base" (i.e. when the number of basis functions is K=1), then why does EigenVI still beat ADVI when K=1? I expected them to perform identically in this setting. (3) pp.8: The authors state that they "standardize the target using either GSM or BaM before applying EigenVI". Any reason that they pick those methods relative to others? Would they expect any arbitrary BBVI method (e.g. ADVI of Kucukelbir et al. 2017) to work just as well for this purpose? (4) pp.9: Should "this family is limited to target functions that are close to Gaussian" read "... that are close to mixtures of Gaussians"? For example, the cross-distribution (Figure 3, row 3) does not look close to Gaussian to me, although it does look close to a mixture of Gaussians. Typos: (1) Abstract: "novel class variational approximations" should read "novel class of variational approximations".

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

4

Limitations

Yes, authors describe limitations well.

Summary

This paper proposes EigenVI, which uses orthonormal distribution functions as the basis of the variational distribution family $q$. Then, minimizing the difference between $q$ and the target distribution $p$ via minimizing the Fisher divergence (2-norm score distance) is turned into an eigen-decomposition problem. Authors show that EigenVI is very competitive in distribution approximation compared with alternatives.

Strengths

* The presentation is clear. The math is solid. * The proposed EigenVI is very intuitive and reasonable. * The low-order basis and higher-order bases have their own interpretation of the distribution's feature, which is desirable. * Experiments are extensive, including synthetic and real-world data.

Weaknesses

See Questions.

Questions

* Eq. 9, should be $\mathrm{d} z$? Line 130, the score of the target is not equal to the gradient of the log joint, but up to a scaling? * Could the authors explain a bit more why centering the distribution is very important for better approximation? * Is there any guidance about how to choose a good proposal distribution? How do authors choose their proposal distributions used in the experiments? Is the choice of the proposal distribution very significant or quite robust? * Eq. 16 seems a bit weird, especially the term $L_z(x_{1:N})$. Besides, it is better to clearly state whether $x_{1:N}$ are independent samples or a sequence, and how it relates to a single $z$. * Why is forward KL used in the whole paper rather than the commonly used backward KL as in ELBO? * How about the actual running time comparison on synthetic and real-world experiments? I understand that theoretically comparison is hard because eigenVI's complexity is w.r.t. number of basis (order) but the stochastic algorithm's complexity is mainly w.r.t. # epochs. However, solving large eigen-decomposition problems in practice might be very time-consuming. So it is very helpful to have an empirical understanding of the running time of the algorithm.

Rating

7

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

/

Summary

The paper proposes Eigen-VI, a black-box variational inference method that uses orthogonal function expansions to parameterize the variational distribution and uses Fisher divergence as an objective. The method does not require gradient-based optimization method and behaves well in a set of synthetic and real application experiments.

Strengths

Strengths * The paper is well-written and easy to follow. * The proposed Eigen-VI is simple and easy to implement since the optimization problem can be solved by computing the minimum eigenvalue of a given matrix, and it also achieves comparable or better performances on hierarchical modeling benchmarks compared to other black-box variational inference methods.

Weaknesses

Weaknesses * Compared to other BBVI methods, Eigen-WI restricts the variational families to be $K^{\text {th}}$-order variational family(defined in eq. 2). When generalize to other variational families, gradient-free optimization in this paper may fail. Similarly, if variational distribution is set to be a linear combination of function basis, the optimization problem in other BBVI methods would also be much simpler. Therefore, the contribution of this paper may not be very significant. * Different proposal distributions in importance sampling may affect the property of the matrix $M$, and further affect the numerical efficiency or stability of computing $\lambda_{\text{min}}(M)$. Adding some analysis regarding this matter would enhance the paper.

Questions

* In the experiment part of this paper, why does the author only use the Hermite polynomials as the variational family? Is this because other variational family performs badly or have intractable issues?

Rating

6

Confidence

4

Soundness

3

Presentation

4

Contribution

2

Limitations

See Weaknesses above.

Summary

The authors introduce a novel family of variational distributions based on orthogonal function expansions which is optimized by minimizing a Fisher divergence. They show that an unbiased importance sampling estimate of the divergence results in a quadratic form, which can be minimized by finding its smallest eigenvector.

Strengths

The work is original and of high overall quality. The proposed variational family and training methodology is novel and of interest to the variational inference community. The empirical evaluation is fairly extensive and shows that the method works well for low dimensional problems. At the same time the authors are very forward with discussing the current limitations of their method which opens up interesting questions for future research.

Weaknesses

As discussed by the authors, the work seems to be mostly limited by the fact that higher-order function expansions, which are needed to model non-Gaussianity, become exceedingly expensive in high dimensions as a result of the blow up of the size of the eigenvalue problem.

Questions

- The authors describe a procedure to sample from the variational approximation, so I am not sure why importance sampling needed to estimate the Fisher divergence? Is it purely for computational reasons or am I missing something? - How is the proposal distribution chosen that is used for importance sampling?

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

3

Limitations

The authors have adequately addressed the limitations of their work.

Rebuttal?

Dear authors (and reviewers), There were some positive reviews, but also one negative and one borderline review. At the end of the day, we don't simply average scores, so this paper is not guaranteed acceptance, though it certainly has a good chance. Hence I'd advise the authors to write rebuttals, at least to the two most negative/neutral reviews. Reviewers, please take a look at each others' reviews and see if that changes your mind. As a reminder, the OpenReview platform allows us to have a dialogue, so you can post small questions to the authors and authors can respond. Best regards Area Chair

Response

Thank you for the detailed response. For the comment on unbiasedness, I should have phrased my original comment a little better. It is well known the score-matching objective itself is valid with subsampling. I was more concerned about the fact that it is unclear how to incorporate subsampling within the proposed algorithm in a consistent way. Nevertheless, I am convinced that the proposed methodology is technically interesting and will raise my score. However, I am still concerned that there isn't enough evidence that it would be practically applicable to larger-scale real-world problems. Thus, the paper feels somewhat less complete in its current state. For instance: > There are many directions for future work to help scale to higher dimensions. For instance, in many situations, one might expect non-Gaussianity on a low-dimensional subspace, and that the remaining dimensions can be suitably modeled with a Gaussian. This family is a special case of the orthonormal Hermite family considered in the paper and so is still an eigenvalue problem. Demonstrating this would have been much more convincing and made the work more complete.