Title: Variational Inference with Gaussian Score Matching
Venue: NeurIPS 2023
Accepted: yes
Keywords: Variational Inference

score matching

KL projection

polyak stepsize
Abstract: Variational inference (VI) is a method to approximate the computationally intractable posterior distributions that arise in
  Bayesian statistics. Typically, VI fits a simple parametric distribution to be close to the target posterior, optimizing an appropriate objective such as the evidence lower bound (ELBO).   In this work, we present a new approach to VI. Our method is based on the principle of score matching---namely, that if two distributions are equal then their score functions (i.e., gradients of the log density) are equal at every point on their support. With this principle, we develop score-matching VI, an iterative algorithm that seeks to match the scores between the variational approximation and the exact posterior. At each iteration, score-matching VI solves an inner optimization, one that minimally adjusts the current variational estimate to match the scores at a newly sampled value of the latent variables. We show that when the variational family is a  Gaussian, this inner optimization enjoys a closed-form solution, which we call Gaussian score matching VI (GSM-VI). GSM-VI is a ``black box'' variational algorithm in that it only requires a differentiable joint distribution, and as such it can be applied to a wide class of models. We compare GSM-VI to black box variational inference (BBVI), which has similar requirements but instead optimizes the ELBO. We first study how GSM-VI behaves as a function of the problem dimensionality, the condition number of the target covariance matrix (when the target is Gaussian), and the degree of mismatch between the approximating and exact posterior distribution. We then study GSM-VI on a collection of real-world Bayesian inference problems from the posteriorDB database of datasets and models. We find that GSM-VI is faster than BBVI and equally or more accurate. Specifically, over a wide range of target posteriors, GSM-VI requires 10-100x fewer gradient evaluations than BBVI to obtain a comparable quality of approximation.

Variational Inference with Gaussian Score Matching

Abstract

Variational inference (VI) is a method to approximate the computationally intractable posterior distributions that arise in Bayesian statistics. Typically, VI fits a simple parametric distribution to be close to the target posterior, optimizing an appropriate objective such as the evidence lower bound (ELBO). In this work, we present a new approach to VI. Our method is based on the principle of score matching-namely, that if two distributions are equal then their score functions (i.e., gradients of the log density) are equal at every point on their support. With this principle, we develop score-matching VI, an iterative algorithm that seeks to match the scores between the variational approximation and the exact posterior. At each iteration, score-matching VI solves an inner optimization, one that minimally adjusts the current variational estimate to match the scores at a newly sampled value of the latent variables. We show that when the variational family is a Gaussian, this inner optimization enjoys a closed-form solution, which we call Gaussian score matching VI (GSM-VI). GSM-VI is a 'black box' variational algorithm in that it only requires a differentiable joint distribution, and as such it can be applied to a wide class of models. We compare GSM-VI to black box variational inference (BBVI), which has similar requirements but instead optimizes the ELBO. We first study how GSM-VI behaves as a function of the problem dimensionality, the condition number of the target covariance matrix (when the target is Gaussian), and the degree of mismatch between the approximating and exact posterior distribution. We then study GSM-VI on a collection of real-world Bayesian inference problems from the posteriorDB database of datasets and models. We find that GSM-VI is faster than BBVI and equally or more accurate. Specifically, over a wide range of target posteriors, GSM-VI requires 10-100x fewer gradient evaluations than BBVI to obtain a comparable quality of approximation. 1

1 Introduction

This paper is about variational inference for approximate Bayesian computation. Consider a statistical model p ( θ , x ) of parameters θ ∈ R d and observations x . Bayesian inference aims to infer the posterior distribution p ( θ | x ) , which is often intractable to compute. Variational inference is an optimization-based approach to approximate the posterior [5, 18].

The idea behind VI is to approximate the posterior with a member of a variational family of distributions q w ( θ ) , parameterized by variational parameters w [5, 18]. Specifically, VI methods establish a measure of closeness between q w ( θ ) and the posterior, and then minimize it with an optimization algorithm. Researchers have explored many aspects of VI, including different objectives [8, 9, 20, 26, 27, 29, 34] and optimization strategies [1, 15, 28].

In its modern form, VI typically minimizes KL ( q w ( θ ) || p ( θ | x )) with stochastic optimization, and further satisfies the so-called 'black-box' criteria [1, 28, 33]. Black-box VI (BBVI) only requires the practitioner to specify the log joint log p ( θ , x ) and (often) its gradient ∇ θ log p ( θ, x ) , which for many models can be obtained by automatic differentiation [14, 25]. BBVI has been widely implemented, and it is available in many probabilistic programming systems [4, 21, 31].

In this paper, we propose a new approach to VI. We begin with the principle of score matching [16], that when two densities are equal then their gradients are equal as well, and we use this principle to derive a new way to fit a variational distribution to be close to the exact posterior. The result is score-matching VI . Rather than explicitly minimize a divergence, score-matching VI iteratively projects the variational distribution onto the exact score-matching constraint. This strategy enables a new black-box VI algorithm.

Score-matching VI relies on the same ingredients as reparameterization BBVI [21]-a differentiable variational family and a differentiable log joint-and so it can be as easily incorporated into probabilistic programming systems as well. Further, when the variational family is a Gaussian, score-matching VI is particularly efficient: each iteration is computable in closed form. We call the resulting algorithm Gaussian score matching VI (GSM-VI).

Unlike BBVI, GSM-VI does not rely on stochastic gradient descent (SGD) for its core optimization. Though SGD has the appeal of simplicity, it requires careful tuning of learning rates. GSM-VI was inspired by a different tradition of constraint-based algorithms for online learning [3, 7, 12, 13, 23]. These algorithms have been extensively developed and analyzed for problems in classification, and under the right conditions, they have been observed to outperform SGD. This paper shows how to extend this constraint-based framework-and the powerful machinery behind it-from the problem of classification to the workhorse of Gaussian VI. The key insight is that score-matching (unlike ELBO maximization) lends itself naturally to a constraint-based formulation.

We empirically compared GSM-VI to reparameterization BBVI on several classes of models, and with both synthetic and real-world data. In general, we found that GSM requires 10-100x fewer gradient evaluations to converge to an equally good approximation. When the exact posterior is Gaussian, we found that GSM-VI scales significantly better with respect to dimensionality and is insensitive to the condition number of the target covariance. When the exact posterior is non-Gaussian, we found that GSM-VI converges more quickly without sacrificing the quality of the final approximation.

This paper makes the following contributions:

We introduce score-matching variational inference , a new black-box approach to fitting q w ( θ ) to be close to p ( θ | x ) . Score-matching VI requires no tunable optimization hypermarameters, to which BBVI can be sensitive.

When the variational family is Gaussian, we develop Gaussian score-matching variational inference (GSM-VI). It establishes efficient closed-form iterates for score-matching VI.

We empirically compare GSM-VI to reparameterization BBVI. Across many models and datasets, we found that GSM-VI enjoys faster convergence to an equally good approximation.

We develop score-matching VI in Section 2 and study its performance in Section 3.

Related work. This paper introduces a new method for black-box variational inference that relies only on having access to the gradients of the variational distribution and the log joint. GSM-VI has similar goals to automatic-differentiation variational inference (ADVI) [21] and Pathfinder [35], which also fit multivariate Gaussian variational families, but do so by maximizing the ELBO using stochastic optimization. As in GSM-VI, the algorithm of ref. [32] also seeks to match the scores of the variational and the target posterior, but it does so by minimizing the L2 loss between them.

A novel aspect of GSM-VI is how it fits the variational parameters. Rather than minimize a loss function, it aims to solve a set of nonlinear equations. Similar ideas have been pursued in the context of fitting a model to data using empirical risk minimization (ERM). For example, passive agressive (PA) methods [7] and the stochastic polyak stepsize (SPS) are also derived via projections onto sampled nonlinear equations [3, 13, 23]. A probabilistic extension of PA methods is known as confidence-weighted (CW) learning [12]. In CW learning, the learner maintains a multivariate Gaussian distribution over the weight vector of a linear classifier. Similarly, the second step of GSM-VI also minimizes a KL divergence between multivariate Gaussians. Unlike CW learning, however, the projection in GSM-VI enforces a score-matching constraint as opposed to a margin-based constraint for linear classification.

2 Score Matching Variational Inference

Suppose for the moment that the variational family q w ( θ ) is rich enough to perfectly capture the posterior p ( θ | x ) . In other words, there exists a w ∗ such that

If we could solve Eq. 1 for w ∗ , the resulting variational distribution would be a perfect fit. The challenge is that the posterior on the right side is intractable to compute.

To help, we appeal to the idea of score matching [16]. We define the score of a distribution to be the gradient of its log with respect to the variable 2 , e.g., ∇ θ log q w ( θ ) . The principle of score matching is that if two distributions are equal at every point in their support then their score functions are also equal.

To use score matching for VI, we first write the log posterior as the log joint minus the normalizing constant, i.e., the marginal distribution of x ,

With this expression, the principle of score matching leads to the following Lemma.

Lemma 2.1. Let Θ ⊂ R d be a set that contains the support of p ( θ, x ) and q w ∗ ( θ ) for some w ∗ . That is p ( θ, x ) = q w ∗ ( θ ) = 0 for every θ ̸∈ Θ . Furthermore, suppose that the support of p ( θ, x ) and q w ∗ ( θ ) , and Θ is path-connected. The parameter w ∗ satisfies

if and only if w ∗ also satisfies Eq. 1.

What is notable about Eq. 3 is that the right side is the gradient of the log joint. Unlike the posterior, the gradient of the log joint is tractable to compute for a large class of probabilistic models. (The proof of the lemma is in the appendix.)

This lemma motivates a new algorithm, score-matching VI . The idea is to iteratively refine the variational parameters w in order to satisfy (as much as possible) the system of equations in Eq. 3. At each iteration t , the algorithm first samples a new θ t from the current variational approximation and then minimally adjusts w to satisfy Eq. 3 for that value of θ t , see (4). We want this adjustment to be minimal because, eventually q w t will learn a good fit-one that increasingly matches the scores over the support of the target posterior-and we want to preserve this good fit as much as possible.

Score matching variational inference

This algorithm for score matching VI was inspired by earlier online algorithms for learning a classifier from a stream of labeled examples. One particularly elegant algorithm in this setting is known as passive-aggressive (PA) learning [7]; this algorithm incrementally updates its model by the minimal amount to classify each example correctly by a large margin. This approach was subsequently extended to a probabilistic setting, known as confidence-weighted (CW) learning [12] in which one minimally updates a distribution over classifiers. Our algorithm is similar in that it minimally updates an approximating distribution for VI, but it is different in that it enforces constraints for score matching instead of large margin classification.

At a high level, what makes this approach to VI likely to succeed or fail? Certainly it is necessary that there are more variational parameters than elements of the latent variable θ ; when this is not the case, it may be impossible to satisfy a single score matching constraint in Eq. 4. That said, it is standard (as in, e.g., a factorized or mean-field variational family) to have at least as many variational parameters as there are elements of the latent variable. It is also apparent that the algorithm may never converge if the target posterior is not contained in the variational family, or it may converge to a degenerate solution if the variational approximation collapses to a point mass, thus terminating the updates altogether. While we cannot dismiss these possibilities out of hand, we did not observe erratic or pathological outcomes in any of the empirical studies of Section 3.

For more intuition, Figure 1a illustrates the effect of the update in Eq. 4 when both the target and approximating distribution are 1d Gaussian. The target posterior p ( θ | x ) is shaded blue. The plot shows the initial variational distribution q w 0 (light grey curve) and its update to q w 1 (medium grey) so that the gradient of the updated distribution matches the gradient of the target at the sampled θ 0 (dotted red tangent line). It also shows the update from w 1 to w 2 , now matching the gradient at θ 1 . With these two updates, q w 2 (dark grey) is very close to the target p ( θ | x ) . With this picture in mind, we now develop the details of this algorithm for a more widely applicable setting.

Gaussian Score-Matching VI. Suppose the variational distribution belongs to a multivariate Gaussian family q w ( θ ) := N ( µ , Σ ) ; this is a common choice especially in systems for automated approximate inference [1, 21]. One of our main contributions is to show in this case that Eq. 4 has a closed-form solution. The solution w t +1 = ( µ t +1 , Σ t +1 ) has the following form:

where A t ∈ R d × d is a matrix defined in the theorem below. Note that these updates, which exactly solve the optimization in Eq. 4, only require the score of the log joint ∇ θ log p ( θ , x ) and the score of the variational distribution ∇ θ log q w ( θ ) .

Before deriving Eqs. 5 and 6 in full, we highlight the intuitive form of this closed-form solution. Consider the approximation at the t th iteration q w t and the current sample θ t . First suppose the scores already match at this sample, that is ∇ θ log p ( θ t , x ) = ∇ θ log q w t ( θ t ) . Then the mean does not change µ t +1 = µ t and, similarly, the two rank-one terms in the covariance update in Eq. 6 cancel out so Σ t +1 = Σ t . This shows that when q w t ( θ ) = p ( θ , x ) for all θ , the method stops. On the other hand, if the scores do not match, then the mean is updated proportionally to the difference between the scores, and the covariance is updated by a rank-two correction. Figure 1b illustrates the vector field of updates for a one dimensional target p ( θ , x ) = N (0 , 1) . The vector field points to the solution (green star), which is also the only fixed point of the updates.

0

Log Density

-5

0o

01

p: μ=0.00, o²=1.00

Qw: μ=1.39, o²=1.97

Qwo: μ=1.00, o²=4.00

Qw: μ=0.09, o²=0.92

1.5

0.5

-1.0

-0.5

0.0

0.5

μ

We now formalize this result and give the complete expression for the matrix A t that appears in the update for the mean of the variational approximation in Eq. 5.

Theorem 2.2. (GSM-VI updates) Let p ( θ , x ) be given for some θ ∈ R d , and let q w t ( θ ) and q w ( θ ) be multivariate normal distributions with means µ t and µ and covariance matrices Σ t and Σ , respectively. As shorthand, let g t := ∇ θ log p ( θ t , x ) and let

The solution to eq. (7) is given by Eqs. 5 and 6 where

and ρ is the positive root of the quadratic equation

With the definition of A t in Eq. 8 we can see that the computational complexity of updating µ and Σ via Eqs. 5 and 6 is O ( d 2 ) , where θ ∈ R d and we assume the cost of computing the gradients is O ( d ) . Note this is the best possible iteration complexity we can hope for, since we store and maintain the full covariance matrix of d 2 elements. (The proof is in the appendix.)

Algorithm 1 presents the full GSM-VI algorithm. Here we also use mini-batching, where we average over B ∈ N independently sampled updates of Eqs. 5 and 6 before updating the mean and covariance.

3 Empirical Studies

We evaluate how well GSM-VI can approximate a variety of target posterior distributions. Recall that GSM-VI uses a multivariate Gaussian distribution as its variational family. We investigate two separate cases-one when the target posterior is in this family, and the other when it is not.

Algorithmic details and comparisons. We compare GSM-VI with a reparameterization variant of BBVI as the baseline, similar to [21]. BBVI uses the same multivariate Gaussian variational family, which we fit by maximizing the ELBO. (Maximizing the ELBO is equivalent to minimizing KL).

Algorithm 1: Gaussian Score Matching VI

Algorithm 2: Black-box variational inference

We use the ADAM optimizer [19] with default settings but vary the learning rate between 10 -1 and 10 -3 . We report results only for the best performing setting. The full BBVI algorithm is shown in Algorithm 2.

The only free parameter in GSM-VI is the batch size B . We find that for Gaussian targets, GSM performs equally well for all B ≥ 1 . There are marginal gains for larger batches for high dimensional targets, but B = 2 is a good conservative default. For non-Gaussian targets, we find that all batch sizes converge to the same solution, but we also note that larger batch sizes help to dampen the oscillations in KL divergence around the optimal solution. Notwithstanding these effects, it is not generally necessary to fine-tune this hyperparameter. In all studies of this section we report results for B = 2 and show that it is a good default baseline. We use the same batch size for BBVI.

Both GSM-VI and BBVI require an initial distribution for the variational approximation. Unless specified otherwise, we initialize the variational approximation as a Gaussian distribution with zero mean and identity covariance matrix.

Evaluation metric. GSM-VI does not explicitly minimize any loss function. Hence to compare its performance against BBVI, we estimate empirical divergences between the variational and the target distribution and show their evolution with the number of gradient evaluations. We measure performance in terms of gradient evaluations as the actual running time can be sensitive

(a) Evolution of forward KL divergence, KL(p||q).

D=4

D=16

D=256

106

GSM-VI

NGD B=128

104

BBVI

NGD B=32

102

100

10-2

10-4

10-6

101

103

105

101

103

105

101

103

105

101

103

105

#gradient evaluations

#gradient evaluations

#gradient evaluations

#gradient evaluations

D=4

D=16

D=256

103

GSM-VI

NGD B=128

BBVI

NGD B=32

101

10-1

10-3

10-5

101

103

105

101

103

105

101

103

105

101

103

105

#gradient evaluations

#gradient evaluations

#gradient evaluations

#gradient evaluations

(b) Evolution of reverse KL divergence, KL(q||p).

to implementation details; moreover, we note that in real-world problems, the computation is often bottlenecked by evaluations of the target log density and its gradient. For completeness, we also measure running times of GSM-VI and BBVI for a full rank Gaussian target with 2048 dimensions using sample size of 1 in Jax after JIT-compilation. The GSM update takes 230 ms/update while BBVI takes 227 ms/update.

In the experiments with synthetic models in Sections 3.1, and 3.2 we have access to the true distribution; hence for these experiments, we estimate the forward KL divergence (FKL) empirically by ∑ θ i ∼ p ( θ ) log p ( θ i ) -log q ( θ i ) . To reduce stochasticity, we always use the same pre-generated set of 1000 samples from the target distribution. In the experiments with real-world models in Section 3.3, we do not have access to the samples from the target distribution; for these experiments we monitor the negative ELBO. In all experiments, we show the results for 10 independent runs.

3.1 GSM-VI for Gaussian approximation

We begin by studying GSM-VI where the target distribution is also a multivariate Gaussian.

Scaling with dimensions. How does GSM-VI scale with respect to the dimensionality, D , of the latent variable? Figure 2 shows the convergence of FKL for GSM-VI and BBVI for different dimensionalities. Empirically, for GSM-VI, we find that the number of iterations required for convergence increases almost linearly with dimensionality. The scaling for BBVI is worse: even for small problems ( D < 64 ) it requires 100 times more iterations while also converging to a sub-optimal solution (as measured by the FKL metric). The figure also shows the evolution of the reverse KL divergence, and it is noteworthy that GSM-VI also outperforms BBVI in this metric.

For completeness, we also experimented on these problems 3 with methods for variational inference based on natural gradient descent (NGD-VI) [2, 22]. Natural gradient descent is a stochastic gradient method that uses the Fisher matrix as a preconditioner. We find NGD-VI is slower by factor of

10 1

10 2

10 3

10 4

# gradient evaluations

10

5

10

3

10

1

10 1

10 3

10 5

x

p

log

p

(

x

)

log

q

(

x

)

c=1

GSM-VI

BBVI

10 1

10 2

10 3

10 4

# gradient evaluations

c=10

10 1

10 2

10 3

10 4

# gradient evaluations

c=100

10 1

10 2

10 3

10 4

# gradient evaluations

c=1000

10 1

10 2

10 3

10 4

# gradient evaluations

10

1

10 0

10 1

10 2

10 3

x

p

log

p

(

x

)

log

q

(

x

)

s=0.2

s=0.6

s=1.0

s=1.4

s=1.8

10 1

10 2

10 3

10 4

# gradient evaluations

10

1

10 0

10 1

GSM-VI

BBVI

t=0.1

t=0.5

t=0.9

t=1.3

t=1.7

5-10x than GSM-VI. While NGD-VI outperforms BBVI, its performance is sensitive to the choice of batch size, with small batch sizes prone to divergence. Furthermore, we note that that per-iterationcomplexity of NGD-VI is cubic in the dimensionality ( D 3 ) as each iteration requires inverting a Hessian approximation; this is in contrast to the quadratic ( D 2 ) per-iteration-complexities of GSM-VI and BBVI. For these reasons, we do not investigate NGD-VI further in this paper.

Impact of condition number. What is the impact of the shape of the target distribution? We again consider a Gaussian target distribution, but vary the condition number of its covariance matrix by fixing its smallest eigenvalue to 0 . 1 and scaling the largest eigenvalue to 0 . 1 × c . Figure 3 shows the results for a 10 dimensional Gaussian where we vary the condition number c from 1 to 1000. Convergence of GSM-VI seems to be largely insensitive to the condition number of the covariance matrix. BBVI on the other hand struggles with poorly conditioned problems, and it does not converge for c > 100 even with 100 times more iterations than GSM.

3.2 GSM-VI for non-Gaussian target distributions

GSM-VI was designed to solve the exact score-matching equations in Eq. 3, but these equations only have a solution when the family of variational distributions contains the target distribution (see Lemma 2.1). Here we investigate the sensitivity of GSM-VI to this assumption by fitting nonGaussian target distributions with varying degrees on non-Gaussianity. Specifically, we suppose that

10 1

10 2

10 3

10 4

# Gradient evaluations

2000

0

2000

4000

6000

8000

x

q

(log

q

log

p

)

diamonds

10 1

10 2

10 3

10 4

# Gradient evaluations

280

260

240

220

200

180

160

140

arK-arK

10 1

10 2

10 3

10 4

# Gradient evaluations

6500

7000

7500

8000

8500

9000

bball_drive-hmm

10 1

10 2

10 3

10 4

# Gradient evaluations

10 2

hudson_lynx_hare

GSM-VI

BBVI

10 1

10 2

10 3

10 4

# Gradient evaluations

1677

1678

1679

1680

1681

1682

1683

1684

x

q

(log

q

log

p

)

low_dim_gauss_mix

10 1

10 2

10 3

10 4

# Gradient evaluations

1050

1000

950

900

850

800

gp_pois_regr

10 1

10 2

10 3

10 4

# Gradient evaluations

10 0

0

10 0

10 1

10 2

10 3

10 4

10 5

eight_schools_centered

10 1

10 2

10 3

10 4

# Gradient evaluations

7.5

5.0

2.5

0.0

2.5

5.0

7.5

10.0

eight_schools_noncentered

GSM-VI

BBVI

the target has a multivariate Sinh-arcsinh normal distribution [17]

where the scalar parameters s and t control, respectively, the skewness and the heaviness of the tails. Note that Eq. 10 reduces as a special case to a Gaussian distribution for the choices s = 0 and t = 1 of these parameters.

Figure 4 shows the result for fitting the variational Gaussian to a 10-dimensional Sinh-arcsinh normal distribution for different values of s and t . As the target departs further from Gaussianity, the quality of variational fit worsens for both GSM-VI (solid lines) and BBVI (dashed lines), but they converge to a fit of similar quality in terms of average FKL. GSM converges to this solution at least 10 times faster than BBVI. For highly non-Gaussian targets ( s ≥ 1 or | t -1 | ≥ 0 . 8 ), we have found that GSM-VI does not converge to a fixed point, and it can experience oscillations that are larger in amplitude than BBVI; see for instance s = 1 . 8 and t = 0 . 1 on the left and right of Figure 4, respectively.

3.3 GSM-VI on real-world data

We evaluate GSM-VI for approximate on real-world data with 8 models from the posteriordb database [24]. The database provides the Stan code, data and reference posterior samples, and we use bridgestan to access the gradients of these models [6, 30]. We study the following statistical models, each with a different complexity and dimensionality (D) : diamonds (generalized linear models, D=26), hudson-lynx-hare (differential equation dynamics, D=8), bball-drive (hidden Markov models, D=8) and arK (time-series, D=7), eight-schools-centered and non-centered (hierarchical meta-analysis, D=10), gp-pois-regr (Gaussian processes, D=13), low-dim-gauss-mix (Gaussian mixture, D=5).

For each model (except hudson-lynx-hare ), we initialize the variational parameter µ 0 at the mode of the distribution, and we set Σ 0 = 0 . 1 I d where I d is the identity matrix of dimension d .

For hudson-lynx-hare , we initialize the variational distribution as standard normal. We also experimented with other initializations. We find that they do not qualitatively change the conclusions, but can have larger variance between different runs.

We show the evolution of the ELBO for 10 runs of these models. Four of the models have posteriors that can be fit with multivariate normal distribution: diamonds, hudson-lynx-hare, bball-drive, and arK . Figure 5 shows the result for these models. The other models have non-Gaussian posteriors: eight-schools-centered , eight-schools-non-centered , gp-pois-regr, , and low-dim-gauss-mix . Figure 6 shows the results.

Overall, GSM-VI outperforms BBVI by a factor of 10-100x. When the target posterior is Gaussian, GSM-VI leads to more stable solutions. When the target is non-Gaussian, it converges to the same quality of variational approximation as BBVI. Finally, we note that while GSM-VI yields noisy estimates of the ELBO, it does yield stable estimates of the 1-D marginals and their moments.

4 Conclusion and Future Work

In this paper we have proposed Gaussian score matching VI (GSM-VI), a new approach for VI when the variational family is multivariate Gaussian. GSM-VI does not explicitly minimize a divergence between the variational and target distribution; instead, it repeatedly enforces scorematching constraints with closed-form updates for the mean and covariance matrix of the variational distribution. This is in contrast to stochastic gradient based methods, such as in BBVI, that rely on first-order Taylor approximations of their objective function.

GSM-VI is implemented in an open-source Python code at https://github.com/modichirag/GSM-VI.

Unlike approaches that are rooted in stochastic gradient descent, GSM-VI does not require the tuning of step-size hyper-parameters. Rather, it is able to adaptively make large jumps in the initial iterations (see Fig. 1a) and make smaller adjustments as the approximation converges to the target. GSM-VI has only one free parameter, the batch size, and we found a batch-size of 2 to perform competitively across all experiments. Another choice is how to initialize the variational distribution. For the experiments in this paper, we initialized the covariance matrix as the identity matrix, but additional gains could potentially be made with more informed choices derived from a Laplace approximation or L-BFGS Hessian approximation [35].

We evaluated the performance of GSM-VI on synthetic targets and real-world models from posteriordb . In general, we found that it requires 10-100x fewer gradient evaluations than BBVI for the target distribution to converge. When the target distribution is itself multivariate Gaussian, we observed that GSM-VI scales nearly linearly with dimensionality, which is significantly better than BBVI, and that GSM-VI is fairly insensitive to the condition number of the target covariance matrix. We also found that GSM-VI converges more quickly than BBVI-and also to a solution with a larger ELBO, which is surprising given that BBVI explicitly maximizes the ELBO.

An important avenue for future work is to provide a proof that GSM-VI converges. We note that good convergence results have been obtained for analogous methods that project onto interpolation equations for empirical risk minimization. For instance the Stochastic Polyak Step achieves the min-max optimal rates of convergence for SGD [23]. Note that convergence of VI is a generally challenging problem, with no known rates of convergence even for BBVI [10, 11]. This and others directions are left to future work.

References

A. Agrawal, D. R. Sheldon, and J. Domke. Advances in black-box vi: normalizing flows, importance weighting, and optimization. Neural Information Processing Systems , 2020.

O. Arenz, P. Dahlinger, Z. Ye, M. Volpp, and G. Neumann. A unified perspective on natural gradient variational inference with gaussian mixture models. Transactions on Machine Learning Research , 2023. ISSN 2835-8856. URL https://openreview.net/forum?id=tLBjsX4tjs .

L. Berrada, A. Zisserman, and M. P. Kumar. Training neural networks for and by interpolation. In Machine Learning , volume 119 of Proceedings of Machine Learning Research , 2020.

E. Bingham, J. Chen, M. Jankowiak, F. Obermeyer, N. Pradhan, T. Karaletsos, R. Singh, P. Szerlip, P. Horsfall, and N. Goodman. Pyro: Deep universal probabilistic programming. arXiv:1810.09538 , 2018.

D. M. Blei, A. Kucukelbir, and J. D. McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association , 112(518):859-877, 2017.

B. Carpenter, A. Gelman, M. D. Hoffman, D. Lee, B. Goodrich, M. Betancourt, M. Brubaker, J. Guo, P. Li, and A. Riddell. Stan: A probabilistic programming language. Journal of Statistical Software , 76(1):1-32, 2017.

K. Crammer, O. Dekel, J. Keshet, S. Shalev-Shwartz, and Y. Singer. Online passive-aggressive algorithms. Journal of Machine Learning Research , 7:551-585, 2006.

A. K. Dhaka, A. Catalina, M. Welandawe, M. Riis Andersen, J. Huggins, and A. Vehtari. Challenges and Opportunities in High-dimensional Variational Inference. art. arXiv:2103.01085, March 2021.

A. Dieng, D. Tran, R. Ranganath, J. Paisley, and D. Blei. Variational inference via χ upper bound minimization. In Neural Information Processing Systems , 2017.

J. Domke. Provable gradient variance guarantees for black-box variational inference. In Neural Information Processing Systems , volume 32, 2019.

J. Domke. Provable smoothness guarantees for black-box variational inference. In International Conference on Machine Learning , volume 119, pages 2587-2596. PMLR, 2020.

M. Dredze, K. Crammer, and F. Pereira. Confidence-weighted linear classification. In International Conference on Machine Learning, , volume 307, pages 264-271, 2008.

R. M. Gower, A. Defazio, and M. Rabbat. Stochastic polyak stepsize with a moving target. arXiv:2106.11851 , 2021.

A. Gunes Baydin, B. A. Pearlmutter, A. Andreyevich Radul, and J. M. Siskind. Automatic differentiation in machine learning: a survey. arXiv e-prints , art. arXiv:1502.05767, February 2015. doi: 10.48550/arXiv.1502.05767.

M. Hoffman, D. Blei, C. Wang, and J. Paisley. Stochastic variational inference. Journal of Machine Learning Research , 14:1303-1347, 2013.

A. Hyvärinen. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research , 6:695-709, 2005.

C. Jones and A. Pewsey. The Sinh-Arcsinh Normal Distribution. Significance , 16(2):6-7, 04 2019. ISSN 1740-9705.

M. I. Jordan, Z. Ghahramani, T. S. Jaakkola, and L. K. Saul. An introduction to variational methods for graphical models. Machine learning , 37:183-233, 1999.

D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In ICLR , 2015.

J. Knoblauch, J. Jewson, and T. Damoulas. An optimization-centric view on Bayes' rule: Reviewing and generalizing variational inference. Journal of Machine Learning Research , 23: 1-109, 2022.

A. Kucukelbir, D. Tran, R. Ranganath, A. Gelman, and D. M. Blei. Automatic Differentiation Variational Inference. art. arXiv:1603.00788, March 2016.

W. Lin, M. Emtiyaz Khan, and M. Schmidt. Fast and Simple Natural-Gradient Variational Inference with Mixture of Exponential-family Approximations. arXiv e-prints , art. arXiv:1906.02914, June 2019. doi: 10.48550/arXiv.1906.02914.

N. Loizou, S. Vaswani, I. Laradji, and S. Lacoste-Julien. Stochastic polyak step-size for sgd: An adaptive learning rate for fast convergence. International Conference on Artificial Intelligence and Statistics , 2021.

M. Magnusson, P. Bürkner, and A. Vehtari. posteriordb: a set of posteriors for Bayesian inference and probabilistic programming, November 2022. URL https://github.com/ stan-dev/posteriordb .

Charles C. Margossian. A review of automatic differentiation and its efficient implementation. Wiley interdisciplinary reviews: data mining and knowledge discovery , 9, 3 2019. doi: 10.1002/ WIDM.1305.

T. Minka. Expectation propagation for approximate Bayesian inference. In Uncertainty in Artificial Intelligence (UAI) , pages 362-369, 2001.

C. Naesseth, F. Lindsten, and D. Blei. Markovian score climbing: Variational inference with KL(p || q). In Neural Information Processing Systems , 2020.

R. Ranganath, S. Gerrish, and D. M. Blei. Black Box Variational Inference. art. arXiv:1401.0118, December 2013.

R. Ranganath, J. Altosaar, D. Tran, and D. Blei. Operator variational inference. In Neural Information Processing Systems . 2016.

E. Roualdes, B. Ward, S. Axen, and B. Carpenter. BridgeStan: Efficient in-memory access to Stan programs through Python, Julia, and R, March 2023. URL https://github.com/ roualdes/bridgestan .

J. Salvatier, Wiecki T., and C. Fonnesbeck. Probabilistic programming in Python using PyMC3. PeerJ Computer Science , 2016.

U. Seljak and B. Yu. Posterior inference unchained with EL_2O. art. arXiv:1901.04454, January 2019.

M. W., M. R. Andersen, A. Vehtari, and J. H. Huggins. Robust, automated, and accurate black-box variational inference. arXiv:2203.15945 , 2022.

Y. Yang, R. Martin, and H. Bondell. Variational approximations using Fisher divergence. arXiv e-prints , art. arXiv:1905.05284, May 2019.

L. Zhang, B. Carpenter, A. Gelman, and A. Vehtari. Pathfinder: Parallel quasi-newton variational inference. Journal of Machine Learning Research , 23(306):1-49, 2022.

A Proof of Lemma 2.1

Lemma 2.1. Let Θ ⊂ R d be a set that contains the support of p ( θ, x ) and q w ∗ ( θ ) for some w ∗ . That is p ( θ, x ) = q w ∗ ( θ ) = 0 for every θ ̸∈ Θ . Furthermore, suppose that the support of p ( θ, x ) and q w ∗ ( θ ) , and Θ is path-connected. The parameter w ∗ satisfies

if and only if w ∗ also satisfies Eq. 1.

Proof. (1) = ⇒ (3): Differentiating both sides of (1) in θ gives

(3) = ⇒ (1): Let θ 0 be any arbitrary point in Θ . Because Θ is path-connected, every θ ∈ Θ is connected to θ 0 via a differentiable path r ( t ) where r (0) = θ 0 and r (1) = θ . Integrating both sides of (3) along this path r ( t ) , using again that ∇ θ log p ( θ | x ) = ∇ θ log p ( θ , x ) , and using the Fundamental Theorem of Calculus gives

Rearranging and defining C := log q w ∗ ( θ 0 ) -log p ( θ 0 | x ) gives

By exponentiating both sides and integrating in θ over Θ we have that

where we used that Θ contains the support of both p ( θ | x ) and q w ∗ ( θ ) . Consequently C = 0 , which gives our result.

B Proof of Theorem 2.2

Here we give the proof for Theorem 2.2. We also re-introduce the theorem with a simplified notation, where we use ( µ 0 , Σ 0 ) to denote the mean and covariance at the previous time step of the method, thus dropping the iteration counter t .

Theorem B.1. (GSM updates) Let p ( θ , x ) be given for some θ ∈ R d , and let q 0 ( θ ) and q ( θ ) be the multivariate normal distributions, respectively, with means µ 0 and µ and covariance matrices Σ 0 and Σ . We seek the distribution

As shorthand, let g := ∇ θ log p ( θ , x ) , and let ρ be the positive root of the quadratic equation

Then the solution is given by the following closed-form updates:

Furthermore, if Σ 0 is symmetric positive definite then so is Σ .

Proof. The constraint in this optimization is given by

The KL divergence is given by

Dropping irrelevant terms from the optimization, we obtain the Lagrangian

It is easier to optimize the matrix Σ -1 instead of Σ . We can enforce the symmetry of Σ -1 by writing

and performing an unconstrained optimization over Φ . With respect to the latter, the gradients of the Lagrangian are given by

Next we examine where the gradients of the Lagrangian vanish:

We claim that these equations (though nonlinear) can be solved in closed form. The first step is to eliminate λ from eq. (25) using eq. (22). In this way we find

It is worth highlighting the form of this equation:

This is a simple rank-two update for Σ . Note that Σ = Σ 0 if µ = µ 0 ; also, the solution for Σ is determined by the solution for µ .

Ultimately we must solve for µ , but first it is useful to solve for the intermediate quantity g ⊤ Σ g > 0 . From eq. (28), we obtain

and from eq. (23), we obtain

As shorthand, let ρ = g ⊤ Σ g . Then from eq. (30) we see that ρ satisfies the quadratic equation

Note that there are no unknowns on the right side of this equation. The correct solution is given by the positive root since ρ = g ⊤ Σ g > 0 . Also note that ρ = ( µ -θ ) ⊤ g from eq. (23).

It is useful to define one final intermediate quantity before solving for µ . Let

Note that ε 0 simply measures the degree to which the parameters of q 0 ( θ ) violate the desired constraint ∇ w log q ( θ ) = ∇ w log p ( θ , y ) . Put another way, if ε 0 = 0 , then we have the trivial solution µ = µ 0 and Σ = Σ 0 .

Now we have everything to express the solution for µ in a highly intuitive form; in particular, it will be immediately evident that µ → µ 0 as ε 0 → 0 . Starting from eqs. (23) and (28), we find

Note what has happened here: eq. (32) is a system of nonlinear equations for µ , but in eq. (37), all the nonlinearity has been expressed in terms of ρ . Since ρ can be determined via eq. (30), we arrive effectively at a system of linear equations for µ . Collecting terms, we obtain

We thus arrive at the closed-form update

It is evident from this update that µ → µ 0 as ε 0 → 0 . The matrix inverse in this update can also be computed efficiently from the Woodbury matrix identity.

In sum, the joint update for µ and Σ can be efficiently computed as follows:

Set g = ∇ w log p ( θ , y ) and ε 0 = Σ 0 g -µ 0 + θ .

Solve ρ (1+ ρ ) = g ⊤ Σ 0 g + [ ( µ 0 -θ ) ⊤ g ] 2 for ρ > 0 .

Compute µ = µ 0 + [ (1 + ρ ) I +( µ 0 -θ ) g ⊤ ] -1 ε 0 .

Compute Σ = Σ 0 +( µ 0 -θ )( µ 0 -θ ) ⊤ -( µ -θ )( µ -θ ) ⊤ .

Note that it straightforward to solve for the value of ρ in step 2 of this procedure. In particular, from the quadratic formula, we obtain the positive root

We can solve the linear equation for µ in step 3 of this procedure most efficiently by using the Sherman-Morrison formula

Applying this formula in step 3, we find that

Next we relate this update to the difference in the scores of the target distribution and the current estimate of the variational approximation. Since ∇ θ log q 0 ( θ ) = -Σ -1 0 ( θ -µ 0 ) , we have that

It follows, by substituting the above expression into Eq. 40, that

which gives the update in the statement of the theorem.

Finally we prove that if Σ 0 is positive definite, then so is Σ . To do so, we show explicitly that all of the eigenvalues of Σ are positive. We begin by rewriting our results for Σ in eq. (28) and ρ in eq. (30) in a more convenient form. As shorthand, let

so that M 0 captures the first two terms on the right side of eq. (28). Note that M 0 is positive-definite, a fact that we will exploit repeatedly in what follows. In addition, recall that µ -θ = Σ g from eq. (23). Thus with this notation we can rewrite eqs. (28) and (30) as

Now let e be any normalized eigenvector of Σ ; we want to show that its corresponding eigenvalue λ e is positive. From eq. (42), it follows that

Note that if e ⊤ g = 0 , then it follows trivially that λ e = e ⊤ M 0 e > 0 . So we only need to consider the non-trivial case e ⊤ g = 0 . To proceed, we note that

̸

where we have used the Cauchy-Schwartz inequality to bound ( e ⊤ M o g ) in terms of ( e ⊤ M 0 e ) , the latter of which appears in eq. (46). Substituting this inequality into eq. (46), we find that

To prove that λ e > 0 we need one more intermediate result. Focusing on the rightmost term in this equality, we note that

and rearranging the terms in this equation, we find

This intermediate result is useful because it relates the two terms on the right side of eq. (48). In particular, using eq. (50) to eliminate the term e ⊤ M 0 g in eq. (48), we find:

where the final inequality follows because the individual terms λ 2 e , ( e ⊤ g ) 2 , and ρ are all strictly positive; note that λ e cannot be equal to zero because this contradicts the equality in eq. (46). This completes the proof. Perhaps it is useful that this derivation also gives upper bounds on λ e , namely

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe submission proposes a method for black-box variational inference based that is based on score matching. The method follows an iterative procedure, where at every iteration the variational approximation is updated by first obtaining a single sample and then updating the approximation such that the gradient of the log-densities of the approximation matches the gradient the gradient of the target distribution on that sample. As in general, several solutions might to this constrained problem might exist, the proposed method aims to also minimize the forward KL divergence to the approximation of the previous iteration.\n\nThe paper shows that the optimal parameters are a stationary point, if the optimal distribution can match the score everywhere. This can be easily seen from the constrained optimization problem, as the last approximation minimizes the KL to the last approximation and also satisfies the constraints. \n\nThe constrained optimization problem (minimizing the forward KL to the previous distribution, while matching the gradients of the target) can in general not be solved in closed form, and the submission does not discuss how the algorithm can be implemented in the general setting, nor does it perform any evaluations. Instead, the work focuses on Gaussian variational approximations, and shows that in this (quite relevant) special case, the constrained optimization problem can be solved in closed form to provide a hyperparameter-free update.\n\nThe method is evaluated on several test problems and compared to a baseline which maximizes the ELBO using the reparameterization trick. The proposed method converges around two orders of magnitude faster than using the reparameterization trick and also often learns better or similar approximations. However, in particular for target distributions that are very different from Gaussian distributions, the methods suffers from high oscillations and does not seem to converge.\n\nSTRENGTHS:\nOriginality\n-------------\nThe proposed method seems to be novel, and performing variational inference without explicitly minimizing a particular divergence could be interesting.\n\nClarity\n--------\nThe method is described perfectly clear. The paper was very to read and follow. I also like that the limitations (the method may not converge when the target distribution is more expressive than the variational approximation, the method does not minimize the KL or any other divergence) is clearly communicated.\n\nRelevance\n--------------\nGaussian variational inference is an important problem setting with various applications.\n\nWEAKNESSES:\nTechnical Soundness\n-----------------------------\nWhile the paper shows that the optimal parameters are a stationary point if the target distribution can be perfectly approximated, it does not provide any analysis for the setting that the approximation can not perfectly approximate the target distribution. Given that the paper focuses on Gaussian variational approximations, this lack of analysis is a major shortcoming. Of course, without specifying a particular divergence it is not even clear which parameter would be \"optimal\" in the setting where the target distribution can be perfectly matched. However, it would be important to better analyze, if there are any mild conditions for convergence, and which criteria the learned approximation fulfills (is there any---potentially obscure---divergence that is minimized?). Further, even in the setting where the target distribution can be perfectly approximated, technically the paper only shows that the optimal approximation is a stationary point, but it does not prove that the method actually converges to that point or that there are no other stationary points. Actually, both can be shown straightforwardly from the fact that a Gaussian distribution has full support, but the current submission does not make this point, and, furthermore, the sample complexity might nevertheless be quite bad.\n\nRelated Work\n------------------\nThe paper focuses on Gaussian variational inference, but misses the most relevant works in that area. Natural gradient based methods are known to significantly outperform methods based on the vanilla gradient (such as the BBVI/Reparameterization-Trick baseline used in the paper).\nFor Gaussian variational approximations, the natural gradient can be approximated extremely efficient (almost as fast as the vanilla gradient) and enjoys faster convergence and empirically also better exploration. There is broad literature of natural gradient descent for Gaussian variational approximations based on zero-order, first-order and second-order information:\n\n- VIPS (Arenz et al.. 2020) uses a modified version of a policy search method (called MORE), which is based on compatible function approximation. This approach use least-squares based on samples from the current approximation to fit a quadratic model to the target distribution (therefore, it only require function evaluations but no differentiable target distributions). It can be shown that the parameters of this quadratic model can be used to compute the natural gradient (also in closed form). Their work focuses on GMM approximation, but it uses independent updates to the Gaussian components.\n\n- Lin et al. (2019) show how first order information can be used for approximating the natural gradient. They also consider the GMM setting, but also use independent update for the Gaussian components. Their method estimates the natural gradient of a Gaussian variational distribution using Stein's Lemma. Their estimate of the natural gradient is a simple linear function of the gradient of the log-ratio (between target and approximation) evaluated on samples from the approximation.\n\n- VOGN (Khan & Nielsen, 2018) estimates the natural gradient using both gradients and Hessians of the joint distribution.\n\nFurthermore, there are also other methods (not based on natural gradients) that can be used for optimizing Gaussian variational approximations. For example, when the Hessian of the target distribution is available, HFSGVI (Fan et al., 2015) or TRUST-VI (Regier et al. 2017) can be applied.\n\n\nEvaluation\n--------------\nThe evaluation is insufficient as it lacks the most important baselines VIPS and the method by Lin et al. (2019). Both methods are applicable to the problem setting of the submission, and as shown by Arenz et al. (2020, see Appendix K) even the zero-order method significantly outperforms the reparameterization-baseline used in the submission.\n\nFurthermore, the performance of the baseline seems to be worse than it should. In the experiments of Section 3.1, where a Gaussian distribution is used as target distribution, the paper states that the reparameterization trick converges to suboptimal solutions. However, the ELBO between two Gaussians has a single stationary point which is the global minimizer. Hence, the reparameterization trick should not converge to a suboptimal solution. I can imagine that the vanilla gradient can lead to crawling, however, when using ADAM this should actually not happen as momentum should build up. I suspect that either the hyperparameters are badly chosen (too large stepsizes may lead to oscillations that prevent ELBO improvements) or problems with the implementation (numerical errors when computing the log-density of the Gaussians).\n\nMinor Comment\n---------------------\n\"In the experiments with synthetic models in Sections 3, 3.1, and 3.1 [sic]\"\n\nReferences\n---------------\nArenz, Oleg, Mingjun Zhong, and Gerhard Neumann. \"Trust-region variational inference with gaussian mixture models.\" The Journal of Machine Learning Research 21.1 (2020): 6534-6593.\n\nLin, W., Khan, M. E., & Schmidt, M. (2019). Stein's Lemma for the Reparameterization Trick with Exponential Family Mixtures. arXiv preprint arXiv:1910.13398.\n\nMohammad Emtiyaz Khan and Didrik Nielsen. Fast yet simple natural-gradient descent for variational inference in complex models. In 2018 International Symposium on Information Theory and Its Applications (ISITA), pp. 31–35. IEEE, 2018.\n\nK. Fan, Z. Wang, J. Beck, J. T. Kwok, and K. Heller. Fast second-order stochastic backpropagation for variational inference. In Advances in Neural Information Processing\nSystems, NIPS’15, pages 1387–1395, 2015.\n\nJ. Regier, M. I. Jordan, and J. McAuliffe. Fast black-box variational inference through stochastic trust-region optimization. In Advances in Neural Information Processing\nSystems 30, pages 2399–2408, 2017.\n\nQUESTIONS:\nWhy does the reparameterization trick converge to suboptimal solutions when the target distribution is Gaussian? Can you provide code and additional details on hyperparameter tuning? Can you provide an explanation for this surprising result?\n\nFor changing my opinion, the paper also needs to discuss the relevant methods in this problem setting (above references), and ideally add a comparison to the work by Lin et al. (2019) or Arenz et al. (2020).\n\nLIMITATIONS:\nThe approach does have serious limitations, mainly by not proving convergence. However, I already stated the limitations under weaknesses and I also think that the limitations are adequately discussed in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nA new variational inference (VI) framework is presented by matching the score of the variational distribution, q, with the target posterior p. Specifically, the closed form score matching equations for the Gaussian variational family are derived and the resulting method is named Gaussian Score Matching VI (GSM-VI). GSM-VI shows favorable results in multiple instances when comparing the method performance to Black Box VI (BBVI) on a multivariate Normal target, a sinh-arcsinh Normal target and real data from posteriordb database.\n\nSTRENGTHS:\nA simple, yet demonstrably fruitful, and novel VI framework is presented. The paper is easy to digest; it follows a natural progression of ideas and theoretical results are completed with intuitive explanations. GSMVI shows faster convergence w.r.t. number gradient estimates for several models and datasets w.r.t. BBVI. Furthermore, the paper presents clear scenarios where GSMVI is favorable to BBVI, e.g., when the covariance matrix is ill-conditioned.\n\nWEAKNESSES:\nThe paper presents a novel VI approach, but the only instantiation of the framework in the paper is restricted to the Gaussian variational family setting, without any non-trivial theoretical results beyond this setting. Where VI provides techniques to find the optimal q within a variational family in terms KL-divergence to p, it is not clear how well score-matching VI performs in this setting (except for when q is in the same family as p).\n\nThe paper lacks explanation or deeper analysis of the large amplitude oscillations presented in the experiments; The cause of variance in BBVI is well-known and therefore research can focus on taming it. I am not well-acquainted with the score matching literature, to me these oscillations require an explanation.\n \nOnly reporting the Forward KL (FKL) could be misleading when comparing GSM-VI to other VI methods which optimization is based on the reversed KL. Reporting only the FKL could favor diffuse variational posteriors over a mode-seeking peaked variational posteriors (which are notoriously produced by VI methods); however, it is not clear that the diffuse variational posterior is to be preferred over the peaked variational posterior.\n\nQUESTIONS:\nI don't concur with the use of the term \"closed-form updates\" of line 229 as GSM-VI relies on sampling to be able to evaluate those update equations. This introduces variance in the updates, as opposed to CAVI updates for e.g. the exponential family where closed-form variational parameter updates can be derived without the need of sampling.\n\nI would be ready to raise my score if the rebuttal addresses the following points:\n1. Nuancing the experiments on simulated data with a metric less disadvantageous of VI, e.g. reverse KL, for at least one data set.\n2. Addressing the issue raised under weaknesses regarding oscillatory behavior.\n\nLIMITATIONS:\nYes, limitations are addressed.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a novel alternative optimization strategy for approximate Bayesian inference in statistical modeling. \nThe starting point of the proposed score-based VI approach is the realization that two distributions are the same if their derivative is the same almost everywhere.\nThis principled is used to derive an objective to match the gradient of the log-density of the posterior and the gradient of the log-density of the approximate variational distribution.\nThe objective is set up in a way so as to promote a minimal change in the approximate variational distribution such that the scores are matched on a set of samples drawn from it.\nInterestingly, in the Gaussian case this minimization has clased form.\nThis work is inspired by previous literature on Passive-Aggressive learning.\n\nSTRENGTHS:\nI believe that alternative optimization strategies for VI are an important area of research and, being very familiar with VI and unfamiliar with the literature on PA learning, I find the proposed idea quite original and well realized.\nSomeone may disagree that the idea is novel due to its derivation from PA learning, but in my opinion convincingly showing that this yields an effective strategy to obtain an approximation to the posterior distribution makes a good contribution to the literature.\nThe results are impressive, showing good performance compared to standard optimization (black-box VI) at a fraction of the number of gradient evaluations.\n\nWEAKNESSES:\nOne weakness is that the work focuses heavily on the case when the approximation is Gaussian due to the close form solution of the proposed optimization strategy.\nIt would have been nice to get more insights into other approximations, e.g., Normalizing Flows.\nHaving said that, I believe that the Gaussian case is important and I understand that it was necessary to extensively study it in this version of the paper.\n\nQUESTIONS:\nThe paper is very well written. However, I think that a bit more intuition on the algorithm with Eq 4 as objective would be ideal. At the moment, it seems a bit disconnected from Eq 3 - why would you want to minimally adjust q under that score-mathing constraint? A more comprehensive explanation could help here.\nAlso, even though it is written in some places, I think it is important to constantly remind the reader that there is no longer the traditional ELBO in the formulation, and that the optimization problem is really a different thing altogether.\n\nI wonder about possible situations in which the optimization of Eq 4 could reach a bad local optimum. For example, it the initial q has variance which is too low, is it possible that samples have not enough diversity to allow for a proper coverage of the score-matching objective - and as a result this could result in the algorithm stopping with scores matched in a narrow region of the space with low posterior density.\nI think that some insights into the behavior of the optimization wrt some of the choices that can be made in the initialization would add a lot to the revised version of the paper. \n\nI was intrigued by the comments on the cases when the number of variational parameters is less than the number of parameters/latent variables. A bit more intuition on what is going on in these cases could also be useful for the revised version. And perhaps ways in which the proposed approach could be extended to handle these cases, which may not be that uncommon (e.g., a small hyper-net to generate large sets of parameters of a neural net). \n\nI think that the experimental campaign is solid. However, it would have been really nice to try the proposed score-matching VI on largely parameterized models where the ELBO is known to be problematic. I wonder whether this could mitigate these effects due to the overparameterization (e.g., [1]).\n\n[1] S. Rossi, P. Michiardi, and M. Filippone. Good Initializations of Variational Bayes for Deep Models. In Proceedings of the 36th International Conference on Machine Learning, ICML 2019, Long Beach, USA, 2019.\n\nLIMITATIONS:\nI haven't seen any specific text on the limitations of the approach and I think this could be included in the revision.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes score matching as a new approach to (black box) variational inference (BBVI) where the variational family is Gaussian.\nThe usual way is to minimize the KL divergence (or equivalently to maximize the ELBO) using stochastic gradient descent (SGD) to update the variational parameters.\nInstead, score matching imposes the constraint that the gradient of the log joint equal the gradient of the logarithm of the variational distribution for all parameter values.\nThe update step picks the new variational parameters to minimize the KL divergence between the new variational distribution and the previous variational distribution, under the score matching constraint at parameter values that are sampled from the previous variational distribution. The paper proves that this update step has a closed form solution if the variational family is Gaussian and can be computed in $O(d^2)$ time assuming the gradients can be computed in $O(d)$ time where $d$ is the dimension of the distribution.\n\nThe experimental evaluation compares the new Gaussian score matching VI (GSM-VI) with standard BBVI. If the target distribution is Gaussian, it finds that the number of iterations needed for convergence scales linearly with the dimension for GSM-VI, but worse for BBVI. It also investigates the effect of the condition number of the covariance matrix of the Gaussian target distribution and finds that it does not affect GSM-VI, but BBVI does not perform well if the condition number is large. Using a sinh-arcsinh normal distribution, the paper looks into what happens as the target distribution departs from being Gaussian. If the target is close to a Gaussian, then GSM-VI continues to require fewer gradient evaluations than BBVI, converging to a similar solution. If the target is far from being Gaussian, GSM-VI does not converge and can experience larger oscillations than BBVI. Finally, GSM-VI and BBVI are compared on real-world data from the posteriordb repository. In most of them, GSM-VI outperforms BBVI by a factor of 10 to 100 in terms of number of gradient evaluations.\n\nSTRENGTHS:\nThe score matching idea is a novel approach to BBVI. It is a simple idea that is communicated clearly in the paper and shown to be very effective. It is clearly interesting and relevant to the variational inference community. The code for the experiments is available, which helps with the reproducibility of the results.\n\nWEAKNESSES:\nThe main weakness of the paper is that it is almost purely empirical and provides no explanations for its (strong) empirical result: why is GSM-VI so effective? Why does it scale so well with the dimension? Why is a batch size of 2 best? (This seems particularly odd. I would expect batching to either provide no benefits or the benefit to increase with larger batch size.) Without even an attempt at an explanation, the empirical results almost seem too good to be true.\n\nA potential weakness in the experimental evaluation is that the performance is always measured in terms of number of gradient evaluations, not actual running time. Given that the update step in GSM-VI takes $O(d^2)$ time, the paper should report the actual measured running times as well. It is also unclear if the quadratic update step will worsen GSM-VI's performance for very high-dimensional problems.\n\nThe experimental evaluation should include more details: how was GSM-VI implemented? Is it built on existing VI implementations? What kind of system was it run on? Not all (hyper-)parameters for the experiments were reported: what is the dimensionality of the real-world benchmarks? What was the batch size for BBVI? How were the posteriordb examples selected? And so on.\n\nQUESTIONS:\nThese questions overlap with what I wrote under \"Weaknesses\". I copy them here.\n- Why do you think GSM-VI is so effective and scales so well?\n- What are the measured running times for the experiments (since the update step takes $O(d^2)$ time)?\n- In a similar vein, does the $O(d^2)$ update step have a negative effect for very high-dimensional problems? Does BBVI overtake GSM-VI again for high dimensions in terms of running time?\n- What was the batch size for BBVI in the benchmarks?\n- What's the dimensionality of the posteriordb examples? How were they picked?\n- Regarding figure 4: it would be interesting to see what happens for settings where both $s \\ne 0$ and $t \\ne 1$. Did you try this?\n\nI will update my rating if the author's answers address my concerns.\n\nLIMITATIONS:\nThe paper mentions some limitations in the text, but they are not collected in one place. It would be helpful if the paper elaborated on the limitations in a separate section/paragraph.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}