EigenVI: score-based variational inference with orthogonal function expansions

Thanks for your feedback on the paper. We discuss several of the points below.

1. It isn't clear if subsampling results in an unbiased objective

Any unbiased estimate of the gradient to our objective results in an unbiased estimate of the objective. And so EigenVI can scale to large n with subsampling.

Let $p(z,x)$ be the target joint distribution between data $x$ and latent variables $z$.

Let $g(z)$ be an unbiased estimate of the score $\nabla \log p(z,x)$, i.e., $\mathbb{E}[g(z) \mid z] = \nabla \log p(z,x)$.

Expanding the Fisher divergence,

$\mathbb{E}[\|\nabla \log q(z_b) - g(z_b)\|^2 \mid z_b] = \|\nabla \log q(z_b)\|^2 - 2 \langle \nabla \log q(z), \mathbb{E}[g(z_b) \mid z_b] \rangle + \mathbb{E}[\|g(z_b)\|^2 \mid z_b]$.

Using the unbiased property, this becomes

$\|\nabla \log q(z)\|^2 - 2 \langle \nabla \log q(z), \nabla \log p(z,x) \rangle + \mathbb{E}[\|g(z_b)\|^2 \mid z_b]$.

So we get

$\|\nabla \log q(z) - \nabla \log p(z,x)\|^2 + \mathbb{E}[\|g(z_b)\|^2 \mid z_b] - \|\nabla \log p(z,x)\|^2$.

Thus, $\|\nabla \log q(z_b) - g(z_b)\|^2$ is an unbiased estimate of $\|\nabla \log q(z) - \nabla \log p(z,x)\|^2$ up to constants that do not depend on $q$. Because of this, we can use $g(z_b)$ as a drop-in replacement within the Fisher divergence within our method since

$\min_{q \in \mathcal{Q}} \mathbb{E}\_{q(z)} [|\nabla \log q(z) - \nabla \log p(z,x)|^2] = \mathbb{E}_{q(z)}[|\nabla \log q(z) - g(z)|^2] + C.$

To build this unbiased estimator of the score, we can use data subsampling as follows: suppose data is sampled i.i.d such that

$p(z,x) = \prod_{i=1}^n p(z,x_i)$.

Thus,

$\nabla \log p(z,x) = \sum_{i=1}^n \nabla \log p(z,x_i)$.

We can use this to get unbiased estimates of the full score by sampling the data $x_i$. For instance, by sampling $x_i \sim \frac{1}{n}$ we have that

$\mathbb{E}[n \nabla \log p(z_b, x_i) \mid z_b] = \frac{1}{n} \sum_{j=1}^n n \nabla \log p(z_b, x_i) = \nabla \log p(z, x)$.

We could also use any form of mini-batching over data to also build an unbiased estimate of the gradient.

2. Scaling with dimension

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. Another direction is a low-rank approximation of the tensor – this family could be optimized with gradient-based methods. Finally, one may consider refining an existing basis to develop better basis functions.

3. Why not use VI fit as a preconditioner for MCMC?

As noted in above, MCMC struggles with targets with varying curvature (using a preconditioner does not help with multiscale distributions where curvature varies over different regions of the distribution, e.g., Neal's funnel). In the absence of pathological geometries, HMC can work very well in high-dimensions and so dimension itself doesn’t strike us as the relevant criterion to choose between VI or MCMC. Using VI with MCMC is an active research area, and this motivates developing new VI methods to benefit from their synergies.

4. Comparing to non-parametric VI families

We thank the referee for pointing to additional literature on non-parametric BBVI families. We will refer to these in the related works in the revised manuscript. However given the breadth of this literature, individual comparison with each of these is beyond the scope of this work.

There are a few common elements to all these approaches, e.g., they all use gradient based optimization and hence are similarly sensitive as the baselines that we have considered. In addition, there are some individual challenges, e.g., mixtures are prone to mode collapse, and spline VI assumes a factorized posterior distribution.

Given these considerations, we consider our baselines of full rank Gaussian VI and normalizing flows (NFs) to be adequate. Together, these are the most widely used variational families for BBVI, they encompass the two extremes on the spectrum of parametric vs non-parametric families, and especially NFs being universal approximators should typically encompass other variational families suggested by the reviewer.

Finally, note that we do not suggest that EigenVI results in better final approximation than NFs, but instead our focus is introducing a new variational family and approach to inference which does not rely on typical iterative gradient based optimization and associated hyperparameter tuning (which is common to all aforementioned methods).

5. The paper doesn't discuss how to solve the eigenvalue problem.

Without further information on the scores of p, the resulting matrix in our eigenvalue problem is dense and unstructured, thus there is little we can leverage to design a specialized algorithm. In our experiments, we used off-the-shelf eigenvalue solvers. For typical problems, this is not a bottleneck (see the general response on the discussion for computational cost). We will add a discussion of this.

6. References and related work

We will expand our discussion of related nonparametric families, and we will add references to additional literature on orthogonal families and score matching.

7. Given bijectors, constrained supports are rarely a problem. Would there be a way to use the proposed method for problems with discrete support?

Transforming supports to be real-valued can sometimes make the problem more challenging (and adds an additional level of tuning to the problem). If we know that the variables are bounded, it may be more natural to model them in the original space. Regarding discrete support, this is an interesting direction for future work. With an appropriate basis set and replacing the scores with log ratios, one would still get an eigenvalue problem.

Thanks for your feedback; in the revised manuscript, we will add or expand our discussion to address several of the points you bring up, as detailed below and in the main rebuttal comment.

1. A practitioner would want to know how MUCH additional compute is required to obtain any advantages.

While the actual overhead in terms of computational time depends on the efficiency of implementation, we provide qualitative reasons to justify why this is competitive to other approaches:

Standardization can use BBVI approaches (like ADVI, BaM or GSM), and hence the cost of standardization is the same as the cost of doing VI inference using these. Following this, we need to evaluate the scores for importance samples which can be fully parallelized and hence, depending on the implementation, adds minimal overhead in terms of computational time. The final step is to solve the eigenvalue problem which is independent of the dimensions and is fast in the regimes we consider (see general comment and attached pdf for details).

2. EigenVI is clearly sensitive to (a) the number of basis functions and (b) the number of importance samples. [...] This problem is reminiscent of the problem of a poorly chosen learning rate in standard ADVI

We will modify the text to clarify that we do have two hyper-parameters: the number of basis functions and the number of importance samples. But we do find there is an important difference between our two hyperparameters and the learning rate in ADVI and other gradient-based methods. As we use more basis functions and more samples, the resulting fitted $q$ is a better approximation. So, we can increase the number of basis functions and importance samples until a budget is reached, or until the resulting q is a good enough fit. Tuning the learning rate is much more sensitive because it cannot be too large or too small. If it is too large, ADVI may diverge. If it is too small, it may take too long to converge.

Another fundamental difference in setting the # of bases as compared to the learning rate or batch size of gradient based optimization is that once we have evaluated the score of the target distribution for the samples, these same samples can be reused for solving the eigenvalue problem with any choice of the number of basis functions, as these tasks are independent. By contrast, in iterative BBVI, the optimization problem needs to be solved from scratch for every choice of hyperparameters, and the samples from different runs cannot be mixed together. Furthermore, solving the eigenvalue problem is fast, and scores can be computed in parallel.

Finally, for choosing the proposal distribution, in the case of the Hermite family, we have reasonable defaults because we are standardizing and hence can use e.g. an isotropic Gaussian with a larger variance to have stable importance weights.

3. 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? [...] Do they at least have any practical insight to share?

In terms of practical insights, if the target p is in the variational family Q, then we need the number of samples = # of basis functions. If it is very different from Q, we need more samples, and we use a multiple of the number of basis functions (say of order 10). As discussed before, once we have evaluated a set of scores, these can be reused to fit a larger number of basis functions.

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

In the rebuttal pdf Fig. 1, we focus on a particular model considered in the paper, and we show the sensitivity of posterior estimates to the learning rate and the batch size for the normalizing flow and standard ADVI.

5. What is the proposed advantage of EigenVI over Zhang et al.'s energy-based approach?

The energy-based model is not normalized, and thus it cannot be evaluated or sampled from (without using MCMC). Indeed, their paper details an approach that uses HMC with the energy-based VI approach.

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

In this case, the Gaussian fit by BaM is used as the standardizing distribution. Thus, K=1 can beat ADVI if the standardizing distribution converges faster / to a better distribution.

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

We discuss the motivation behind the standardization in Section 2.3; here you can pick your favorite Gaussian approximation to do this, e.g., a Laplace approximation or ADVI or directly estimating the target mean and covariance (we chose the score-based VI methods in the experiments because they require less tuning and often converge faster). In general, we expect a better estimate of the target mean and covariance to help due to the base distribution of the Hermite family being a standard Gaussian (and thus, lower-order expansions will be more accurate).

8. Should "this family is limited to target functions that are close to Gaussian" read "... that are close to mixtures of Gaussians"?

We agree this sentence could be more clear, and we’ll modify it. What we meant was that to get to higher-dimensional spaces, we are limited to target functions that are close to Gaussian. For lower dimensional models (e.g., the 2D targets), we can afford to take higher-order expansions, and are able to model highly non-Gaussian distributions; we generally found this to be true for up to 3 or 4 dimensions.

Thanks for your feedback on the paper. In the revised manuscript, we will work on clarifying the following points.

1. Should Equation 9 be dz?

We could alternatively write $q(z) dz$ here (with the appropriate adjustments of $q(z)$ and $p(z)$).

2. Line 130: the score of the target is not equal to the gradient of the log joint but up to a scaling?

The log target is equal to the log joint + constant, and so the constant disappears when we take the gradient. Thus, $\nabla \log p(z) = \nabla \log \rho(z)$, where $\rho$ represents the unnormalized model (i.e., the joint in the context of Bayesian inference).

3. On Equation 16 and the likelihood:

This equation describes a general set of Bayesian problems from a set of benchmark models, and so we don’t a priori make assumptions on the conditional independence structure of the likelihood. We are happy to write $p(x_{1:N} | z)$ instead.

4. Could the authors explain a bit more why centering the distribution is very important for better approximation?

The standardization is used with the Hermite polynomial basis, which has a centered Gaussian as the base distribution. As we show in Figure 2, higher order expansions are needed to model distributions whose largest mode is shifted away from the origin. While this is fine in low dimensions, in order to handle higher dimensional targets, we apply the standardization technique as a way to reduce the order needed to model the target.

5. Why is forward KL used rather than the commonly used backward KL as in ELBO?

The ELBO is often used when comparing VI algorithms that use the ELBO / reverse KL as the objective for all VI algorithms. Here we are comparing VI algorithms with different objectives, and so we chose an evaluation metric that was agnostic to the algorithms’ objectives.

6. How to choose a good proposal distribution? How were they chosen in experiments?

Please see our general response above for intuition behind choosing the proposal distribution. See Appendix E for the exact proposal distribution used for each experiment.

7. How about the actual running time comparison on synthetic and real-world experiments? [...] Solving large eigen-decomposition problems in practice might be very time-consuming.

We do not need to solve the full eigendecomposition problem, but only need the smallest eigenvalue pair for a KxK symmetric matrix, which can be a much easier problem to solve. As we mention in the general rebuttal comment, for 3000 basis functions, which is around the largest number of basis functions considered in experiments, this takes under a second. In a qualitative context, in the rebuttal pdf, we report the time it takes to run iterative BBVI methods, and we report the time needed for the eigenvalue solve. For instance, our implementation of ADVI takes over 12 seconds to converge in this example.

Thus overall, we expect the overhead in terms of computational time due to solving the eigenvalue problem to be minimal.

There are two other steps in (Hermite) EigenVI - standardization and evaluating scores of importance samples. While the actual computational time for these depends on the efficiency of implementation, expect this to be competitive to iterative BBVI approaches for two reasons -- 1) Standardization is done using BBVI approaches themselves (like ADVI, BaM or GSM), and hence the cost of standardization is the same as the cost of doing VI inference using these. 2) The scores for the importance samples can be evaluated in a fully parallelized manner and hence, depending on the implementation, adds minimal overhead in terms of computational time.

Finally, to present a quantitative comparison, we previously included a running time comparison for the synthetic 2D experiments, see Appendix E (Figure 6 shows both # of gradient evaluations and wallclock time without parallelization). Here we did not need to use standardization, and so we indicate when EigenVI and the iterative method for the Gaussian fit use the same number of gradient evaluations.

Thanks for your feedback on the paper. We discuss several of the comments and questions below.

1. Compared to other BBVI methods, Eigen-VI restricts the variational families to be 𝐾th-order variational family(defined in eq. 2). When generalize to other variational families, gradient-free optimization in this paper may fail.

As the reviewer points out, the EigenVI approach works for any orthogonal basis set. Importantly, this approach does not just work for a single variational family but a very large class of variational families. In particular, there are many orthogonal basis sets beyond the ones listed in the paper in Table 1. For instance, one may apply a procedure such as Gram-Schmidt to obtain such a set. Finally, new variational families can be formed by even mixing different basis sets. Thus, this approach opens the door to applying VI to many new types of families.

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

Classical BBVI methods also have assumptions that restrict the variational families. For instance, it is common that the family supports reparameterization (which is not true here).

3. Different proposal distributions in importance sampling may affect the property of the matrix M, and further affect the numerical efficiency or stability of computing λmin(M). Adding some analysis regarding this matter would enhance the paper.

In principle we agree that a poor proposal distribution might cause such numerical issues; however, in the current experiments, we do not observe this. Standardizing the distribution as we currently do it also helps here, as it provides a good proposal distribution which can easily be adjusted to have heavier tails if necessary. Furthermore, with enough samples, and since solving for minimum eigenvalues is typically so fast even for large matrices, we expect these issues to be less important. Experimenting with different proposal distributions is beyond the scope of this rebuttal, but we acknowledge the good point raised and will add a remark to this end in the revised manuscript on how, e.g., this could affect the spectral gap.

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

The parameters for most experiments considered lie on an unconstrained scale where Hermite polynomials form a natural basis; the Hermite family can also be seen as a natural extension of non-Gaussianity.

The EigenVI method itself applies to any orthogonal basis set; in Figure 1, we additionally show 1D target distributions and fitted distributions arising from Legendre polynomial and Fourier series expansions.

Thanks for your interest in the paper. We will revise the final version of the paper to make the following points more clear.

1. The authors describe a procedure to sample from the variational approximation, so I am not sure why importance sampling is needed to estimate the Fisher divergence? Is it purely for computational reasons or am I missing something?

First, because we are optimizing the empirical Fisher divergence with respect to $q$, we need the samples to be independent of $q$ (we cannot simultaneously sample from and optimize over $q$). In addition, as we show in Appendix D, the importance weights are important and lead to the form of the eigenvalue problem.

The procedure to sample from $q$ is primarily used after the variational distribution has been fit to, for instance, approximate functionals of $p$.

We will revise the paragraph starting from Line 133 to explain these points more thoroughly.

2. How is the proposal distribution chosen that is used for importance sampling?

Please see our general response above for intuition behind choosing the proposal distribution. See Appendix E for the exact proposal distribution used for each experiment.

We will add a paragraph in the Section 2.3 that details this intuition behind choosing the proposal distribution.