Variational Inference with Gaussian Score Matching
Authors: Chirag Modi
Robert M. Gower
Charles Margossian
Yuling Yao
David Blei
Lawrence K. Saul
Venue: NeurIPS 2023 poster (2023)
Keywords: Variational Inference
score matching
KL projection
polyak stepsize
TL;DR: We derive a closed form update for Gaussian variational family by projecting onto the exact score matching equations.

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
Chirag Modi
Charles C. Margossian Center for Computational Mathematics,
Center for Computational Astrophysics, Center for Computational Mathematics, Flatiron Institute, New York cmodi www.reeiiioonniittttte.....
Flatiron Institute, New York
cmargessiandelltttoooossstttttt org
Yuling Yao Center for Computational Mathematics, Flatiron Institute, New York
Robert M. Gower
Center for Computational Mathematics,
Flatiron Institute, New York
jyao@flatironinstiioee
#goweroflatirronnstiibue org
David M. Blei
Lawrence K. Saul Center for Computational Mathematics,
Department of Computer Science, Statisties, Columbia University, New York david. blei@columbie edu
Flatiron Institute, New York
1saul Isauleflattrroninstiueee org
Abstract
Variational inference (VT) is a method to approximate the computationally in- tractable posterior distributions that arise in Bayesian statistics. Typically, VI fits a simple parametric distribution to be close to the target posterior, optimizing appropriate objective such as the evidence lower bound (ELBO). In this work, we present new approach to VI. Our method 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 match the scores between the variational approximation and the exact posterior. At each iteration, score-matee matching VI solves an inner optimization, one that minimally adjusts the current variational estimate to match the scores at newly sampled value of the latent variables. We show that when the variational family is Gaussian, this inner optimization enjoys dosed-form solution, which we call Gaussian score matching VI (GSM-VI). GSM-VI "black box" variational algorithm that it only requires differentiable joint distribution, and as such it can be applied to a wide class of models. We compare GSM- to black box variational inference (BBVI), which has similar requirements but instead optimizes the ELBO. We first study how GSM-VI behaves as 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 VI collection of real-world Bayesian inference problems from the PosteriorDB database of datasets and models. We find that GSM-VI VI is faster than BBVI and equally more accurate. Specifically, target posteriors, GSM-VI requires 10-100x fewer gradient evaluations than BBVI obtain comparable quality of approximation.
wide range of
provide Python implementation of GSM- VI algorithm at https///ghhabbccommmmiiiinnngggg
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Introduction
This paper about variational inference for approximate Bayesian computation. Consider a statistical model p(0, x) of parameters 0 E IR and observations æ. Bayesian inference aims to infer the posterior distribution p(0) æ), which is often intractable to compute. Variational inference is an optimization- based approach to approximate the posterior [5, 18].
The idea behind VI is approximate the posterior with a member of a variational family of dis- tributioos Iw (0), parameterized by variational parameters w [5, 18]. Specifically, VI methods establish measure of closeness between qu (w) and the posterior, and then minimize with an optimization algorithm. Researchers have explored many aspects of VI, including different objec- tives [8, 9, 20, 26, 27, 29, 34] and optimization strategies [1. 15, 28].
In its modern form, VI typically minimizes KL (9w (0) [p) x)) with stochastic optimization, and further satisfies the so-called "black-bon -bxx criteria [1. 28, 33]. Black-box VI (BBVI) only requires the practitioner to specify the log joint log p(O, x)) and (often) its gradient V log p(0, 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 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 variational distribution be close to the exact posterior. The result is score-matching VI. Rather than explicitly minimize divergence, score-matching VI iteratively projects the variational distribution onto the exact sorre-mathhing constraint. This strategy enables new black-box VI algorithm.
Score-matehn matching VI relies on the same ingredients as reparameterization BBVI [21]-a a differen- tiable variational family and a differentiable joint-annn 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, requires careful tuning of learning rates. GSM-VI was inspired by different tradition of constraint-baat 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-be 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 several classes of models, and with both synthetic and real-world data. general, we found that GSM requires 10-100x fewer gradient evaluations to converge to equally good approximation. When the exact posterior is Gaussian, we found that GSM-VI VI scales significantly better with respect to dimensionality and is insensitive to condition number of the target covariance. When the exact posterior is non- non-Gaussian, found that GSM-VI 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 (0) to be close to p(0 æ)). Score-matching VI requires no tunable optimization hyperma- rameters, 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 BBVL Across many models and datasets, we found that GSM- VI enjoys faster convergence to an equally good approximation.
We develop score- more-matching VI in Section 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-differenii which also 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.
variational inference (ADVI) [21] and Pathfinder [35],
A novel aspect of GSM-VI is how it fits the variational parameters. Rather than minimize loss function, 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 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 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 Tu (0) is rich enough to perfectly capture the posterior p(0 æ). In other words, there exists a w² such that
log tw- (0) = log p(0)z),
AA E θ.
(I)
we could solve Eq. I for w", the resulting variational distribution would be a perfect fit. The challenge is that the posterior on the right side is intractable compute.
To help, we appeal to the idea of score matching [16]. We define the score a distribution to be the gradient of its log with respect to the variable?, e.g., Vo log 9w (0). 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,
log p(0 Φ) log p(0, x) log p(æ).
(2)
With this expression, the principle of score matching leads to the following Lemma.
Lemma 2.1. Let θ C Rd be a set that contains the support of p(0, x) and tw- (0) for some w". That is p(0, x) = qw- (0) = 0 for every 8 & O Furthermore, suppose that the support of p x)) and qw- (0), and θ is path-conneeted. The parameter w" satisfies
Vo log qw (0) = Vo log p)0,
AO O
(3)
if and only 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 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 equations in Eq. 3. At each iteration i, the algorithm first samples new 0, from the current variational approximation and then minimally adjusts w to satisfy Eq. 3 for that value of 0, see (4). We want this adjustment to be minimal because, eventually qw, will learn good fit fit-one that increasingly matches the scores over the support of the target posterior-ant we want to preserve this good fit as much as possible.
We make this clear because, in some literature, the score is defined as the gradient with respect to the parameter.
3

---

Score matching variational inference
At iteration t:
1. Sample O₂ Yuu, (0).
2. Update the variational parameters:
W1+1 = arg min KL (Sw., (0) tw (0) Vo log quo (0.) = Vo log log a)).
w
such that
This algorithm for score matching VI was inspired by earlier online algorithms for learning a classifier from stream of labeled examples. One particularly elegant algorithm in this setting is known as aussive-aggressive (PA) learning [7]: this algorithm incrementally updates its model by the minimal amount classify each example correctly by large margin. This approach was subsequently extended to a probabilistic setting, known as confidence-weighte (CW) learning [12] in which one minimally updates 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 O: when this is not the case, it may be impossible to satisfy single score matching constraint in Eq. 4. That said, it is standard (as in, c.g., factorized mean-field variational family) 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 not contained in the variational family, or may converge to degenerate solution if the variational approximation collapses to 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 la illustrates the effect of the update in Eq. 4 when both the target and approximating distribution are 1d Gaussian. The target posterior p(0) x)) is shaded blue. The plot shows the initial variational distribution qwa (light grey curve) and its update to 9w, (medium grey) so that the gradient of the updated distribution matches the gradient of the target at the sampled 6, (dotted tangent line). also shows the update from un to UP2 now matching the gradient at 0, With these two updates, Qw, (dark grey) is very close to the target p(0 ar). With this picture in mind, we now develop the details of this algorithm for more widely applicable setting.
Gaussian Score-Matching VI. Suppose the variational distribution belongs to multivariate Gaussian family gw (0) := NM,,, Σ); this 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 closed-form solution. The solution w,-1 = (1t+1. Σ(+1) has the following form:
[441 14 A. (Vo logr((0 x) Vo log qoo, (wr, (0.))
(5)
E41 = E, + (c,- 0.) 0,) (pr+1 0,) 0,)
(6)
where A, E Rdxd 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 Vo log p(O, .:) and the score of the variational distribution Vo log qw (0).
Before deriving Eqs. 5 and 6 in full, we highlight the intuitive form of this closed-form solution. Consider the approximation
the tth iteration qw, and the current sample O,. First suppose the scores already match at this sample, that is Vo log p(0., x) Vo log Tw. (0:). Then the mean does not change M++1 = b4 and, similarly, the two rank- terms in the covariance update in Eq. 6 cancel out So Σ++1 = Σ. This shows that when lw, (0) = p(0 ( xx) for all 6, the method stops. On the other hand, the scores do not match, then the mean updated proportionally to the difference between the scores, and the covariance is updated by rank-two correction. Figure 1b illustrates the vector field of updates for a one dimensional target p(0, .x) = N/ 1). The vector field points the solution (green star), which is also the only fixed point of the updates.

---

Log Density
1.0
81
0.00,00,4
-2-1.00
Owy
1=399, 1197
1=000
-400.00
Owg p0009 o'=0.92
(b) The vector field of Eq. 4, averaged over 5 independent samples, where p(0 x) = (00, 1). The solution (n, = (0, 1) is the green star.
(a) Two iterations of GS.-VI. The log density of the target posterior is shaded blue; the initial distribution gung light grey: the first update Own is medium grey; and second update Qun dark grey.
We now formalize this result and give the complete expression for the matrix A, that appears the update for the mean of the variational approximation in Eq. 5.
Theorem 2.2. (GSM-VI updates) Let p(0, x)) be given for some 0 E Rd, and let Ow (0) and qw (0) be multivariate normal distributions with means μ4 and μ and covariance matrices E, and Σ, respectively. As shorthand, let 9: := Vo log p(6, (0,, x)) and let
M4+1, Ee+1 = argmin KL( a))
such that
Ve log ((0,) = Volog Vologp(O,,,))
(7)
4.E3
The solution to eq. (7) is given by Eqs. and 6 where
(M 0.)g.
A
(8)
1+p I+ (M4 0,) Jt_
and p is the positive root of the quadratic equation
(w:-0,) p(1+p) E.g. 1o:--0) 0,) g
(9)
With the definition of A, in Eq. 8
that the computational complexity of updating μ and E via Eqs. 5 and 6 is O(d)), where E IRd and we assume the cost of computing Note this is the best possible iteration complexity we can hope for, since we store and maintain the full covariance matrix of elements. (The proof is in the appendix.)
the gradients is Od).
Algorithm 1 presents the full GSM-VI algorithm. Here also mini-batee mini-batching, where average over B E 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 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 not.
Algorithmic details and comparisons. We compare GSM-VI with 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 minimizing KL).
5

---

Algorithm 1: Gaussian Score Matching VI
Input :lnitial mean estimate poo initial covariance estimate Eg, target distribution p(0
number of iterations N E N, batch size B E N.
Output: :rtputtMultivaiiatt normal variational distribution Ow (0) == ( Σ)
for i .....,N -1
B iteration loop
do
for j - 0, B B-1
P batch loop
do
Sample ((1) NN Ej)
g < Vg log (
E + Eig- p4 0
Solve p(L+p) g Egg + [p: 0) g for p >
5p - ltp I 14:-018
I+p+ 0)
( pi =gp1(
ui 6Σ() t (H:-0)(14 -0) - (mL) (m 0)T
end
Update p441 < 14 + E, أل(رر/ /B
Update Σ1+1 f- Ej + E, 8Σ6)) /B
end
Om 0 + N(NN, En)
Algorithm 2: Black-box variational inference
Input :Initial mean estimate μoo Initial covariance estimate Eo, target distribution p(0) number of iterations N, batch size B, learning rate E
:
normal variational distribution zw (0) Nu,, E)
te f N Eo) :
for i O, N...N.- -1
P iteration loop
do
(0 (1)) ,..., 0(B) (0)
ELBO = E; log(p(0(1), xz) log aw (06)) Sample batch of B points;
wn
FI.BO
Optimization step, we use ADAM;
end
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 shown in Algorithm 2.
The only free parameter in GSM-VI the batch size B. We find that for Gaussian targets, GSM performs equally well for all B 2 1. There are marginal gains for larger batches for high dimensional targets, but B is a good conservative default. For non- 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 KL divergence around the optimal solution. Notwithstanding these effects, is not generally necessary to fine-tune this hyperparameter. In all studies of this section report results for = and show that it a good default baseline. We use the same batch size for BBVI.
Both GSM-VI and BBVI require initial distribution for the variational approximation. Unless specified otherwise, we initialize the variational approximation as 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

---

10"
CSMM
NGD 5-128
-loppee)
10
BBW
NGD B-332
100
10
10
10P
10
10P
10P
10A
10P
10"
10-
10?
109
gradient evalzations
grad ent evalaations
gradiant ovaluations
gradient evaluatiors
(a) Evolution of forward KL divergence, KL(pl))).
D
D-16
DeB6
Dw220
-loppte)
IS0V1
NGD S=-32
20*
gradient evallattoons
# grad ent evans
gradient evaluntions
gradient evalintrors
(b) Evolution of reverse KI. divergence, Kl(qglp).
Figure 2: Scaling with dimensionality: evolution of forward (a) and reverse KL (b) with the number of gradient evaluations of the target distribution, which here is Gaussian distribution with dense D x D covariance matrix. Each panel shows the results for different dimensionality (indicated in the title). Translucent lines show the scatter of 10 different runs, and the solid line shows the average. GSM-VI and BBVI use batch sizes of B 2, while the batch sizes for natural gradient descent are shown in the legend.
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 full rank Gaussian target with 2048 dimensions using sample size of in Jax after JIT-compilation. The GSM update takes 230 ms/update while BBVI takes 227 ms/update.
In experiments with synthetic models Sections 3.1, and 3.2 we have access to the true distribution; hence for these experiments, we estimate the forward KI divergence (FKL) empirically by Ee.pt(e) log p(O)) log (0,). 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, do not have access to the samples from the target distribution; for these experiments we monitor the negative ELBO. all experiments, we show the results for independent runs.
3.1 GSM-VI for Gaussian approximation
begin by studying GSM-VI where the target distribution 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 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 worse: even for small problems (/) < 64) requires 100 times more iterations while also converging to 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' with methods for variational inference based on natural gradient descent (NGD-VI) [2, 22]. Natural gradient descent is stochastic gradient method that uses the Fisher matrix as preconditioner. We find NGD-VI is slower by factor of
Speeciiically, the package GMMVI with default settings except changing the batch size.
7

---

1-10
1-100
c-1000 1000
10a
MSM
107
10
10
10.
aclieend avalaations
grudient erallcat aoliiaaiiionss
graddent wnalvationss
graadeent evallationn
Figure 3: Impact of condition number: the target distribution. The target is a 10-dimensional Gaussian albeit with dense covariance matrix of different condition numbers e specified in the title of different panels. Translucent lines show the scatter of 10 different runs and the solid line shows the average.
evolution of FKL with the number of gradient evaluations of
10a
10a
GSM-VI
BBVI
togp00 log 904)
10a
109
10
5=0.2
t=0.1
5-0.6
t-0.5
10a
t=0.9
5=1.4
t=1.3
5-1.8
10-1
10a
10
10a
10
gradient evaluations
gradient evaluations
Figure 4: Impact of -aussiasiatt evolution of FKL with the number of gradient evaluations for Sinh-aresinh distributions with 10-dimensional dense Gaussian the base distribution. Gaussian distribution has 0, t = 1. In the left panel, we vary skewness s while fixing t = 1, and in the right panel we vary the tail-weight& with skewness fixed to s Solid lines are the results for GSM, dashed for BBVI.
5-10x than GSM-VI While NGD-VI 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-iteration- complexity of NGD-VI is cubic in the dimensionality (D*) each iteration requires inverting Hessian approximation; this in contrast to the quadratic (D²) en--iieraaion-copleeiiiiie 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 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 x c. Figure shows the results for 10 dimensional Gaussian where we vary the condition number from 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 converge for > 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 solution when the family of variational distributions contains the target distribution (see Lemma 2). Here we investigate the sensitivity of GSM-VI to this assumption by fitting non- Gaussian target distributions with varying degrees non- non-Gaussianity. Specifically, we suppose that

---

diamends
ark-akk
3oall_drive-hre drtve-hmm
hudson lynx hare
8000
-140
9000
GSM--VI BBVI
160
6000
8500
100
8000
llg
2000
7500
10"
-240
7000
-2000
220
6500
-200
10 10? 10 101
10
10
10
10 10" 10" 10" evaluations
10" 10 10" 10 Gradient
evaluations
Figure 5: Models from posteriordb: Convergence of the ELBO for four models with multivariate normal posteriors. We show results for 10 runs.
low, low_dir www__dir_gguussnegg www.ggguusssmmm mix
up_pois_regr
Seoii_schoos__eeerreee cenieered
eight__chholl.nnnneeeeeed
eightt_scoools noncentered
880
10*
10.0
GSM-VI
1683
BaVI
1682
10P
2681
1680
10
W
-2.5
1678
-5.0
1050
-7.5
100
10
10
10'
10'
10'
10* 10'
101 10' 10 10
10'
10
10
10
evatustions
Figure 6: Models from posteriordb: Covergence of ELBO for four models with non-Gaussian posteriors. We show results for 10 runs.
the target has a multivariate aanh-arcsinh normal distribution [17]
sinh (: [sinh (())
(H, E):
(10)
where the scalar parameters and control, respectively, the skewness and the heaviness of the tails. Note that Eq. 10 reduces as a special case to Gaussian distribution for the choices = and t = 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 and f. As the target departs further from Gaussianity, the quality of variational fit worsens for both GSM-VI VI (solid lines) and BBVI (dashed lines), but they converge to fit of similar quality in terms average FKL. GSM converges this solution least 10 times faster than BBVI. For highly non-Gaus Gon-Gausiian targets (s 2 or It 1 2 0.8), we have found that GSM-VI VI does not converge a fixed point, and it can experience oscillations that are larger in amplitude than BB VI; see for instance 1.8 and = 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) di amonds (generalized linear models, D=26), hudson- -lyxx-lyx-aaaaa (differential equation dynamics, D=8), bball-dris (hidden Markov models, D-8) and arK (time-series, D=7), eight-schooll sechholl--eeteeeedd and non- centered (hierarchical anal-sss,ysis, D=10), gp-pois regr (Gaussian processes, D=13), low-dim gauss-mis mix (Gaussian mixture, D=5).
For each model (except hudson- Matson-lynx-haree we initialize the variational parameter f40 at the mode of the distribution, and we set Eg 0.1.a where Ia is the identity matrix of dimension d.

---

For hudson-lyna -yre- 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 pos- teriors that can be fit with multivariate normal distribution: diamonds, hudson-lymi com--llxx-aar haane bball-drive, and ark. Figure 5 shows the result for these models. The other models have non-Gaussian posteriors: hight-schools-centeer
wett--schools-oonnneeeee gp-pois-regr, ,, and www-iim-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 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, does yield stable estimates of the 1-D marginals and their moments.
4 Conclusion and Future Work
In this paper have proposed Gaussian score matching VI (GSM-VI), new approach for VI when the variational family is multivariate Gaussian. GSM-VI VI does not explicitly minimize a divergence between the variational and target distribution; instead, it repeatedly enforces score- matching constraints with closed-form updates for the mean and covariance matrix of the variational distribution. This is in contrast 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 http://gttth...cmm/mmmdccinnnnaf
Unlike approaches that are rooted in stochastic gradient descent, GSM-VI does not require the tuning of step-size pareme-rreetesss.. Rather, it is able adaptively make large jumps in the initial iterations (see Fig. Ia) and make smaller adjustments as the approximation converges to the target. GSM-VI VI has only one free parameter, the batch size, and we found 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, initialized the covariance matrix 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 larger ELBO, which is surprising given that BBVI explicitly maximizes the ELBO.
An important avenue for future work is provide 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 generally challenging problem, with no known rates of convergence even for BBVI [10, 11]. This and others directions are left to future work.
References
[1] A. Agrawal, D. R. Sheldon, and J. Domke. Advances black-box vi: normalizing flows, importance weighting, and optimization. Neural Information Processing Systems, 2020.
[2] O. Arenz, P. Dahlinger, Z. Ye, M. Volpp, and G. Neumann. A unified perspective on natural gra- dient variational inference with gaussian mixture models. Transactions on Machine Learning Re- search, 2023. ISSN 2835-8856. URL https: http:////pppeeeeeeeee wttprriwww.ieeeeeeeeeeeeeeeeeeee nett:/orummidd-tBBjjaa.....
10

---

[3] 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.
[4] E. Bingham, Chen, M. Jankowiak, F. Obermeyer. N. Pradhan, Karaletsos, Singh.
P. Szerlip. P. Horsfall, and N. Goodman. Pyro: Deep universal probabilistic programming.
arXiv: 1810.09538, 2018.
[5] D. M. Blei, A. Kucukelbir, and J. D. McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association, 112(518):859-888 2017.
[6] 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.
[7] K. Crammer, O. Dekel, J. Keshet, S. Shalev-Shwarz, and Y. Singer. Online passive-aggressive algorithms. Journal of Machine Learning Research, 7:551-585, 2006.
[8] A. K. Dhaka, A. Catalina, M. Welandawe, M. Riis Andersen, J. Huggins, and A. Vehtari. Challenges and Opportunities in High-dimensional Variational Inference. art DxXiv:2103.0108, March 2021.
[9] A. Dieng, D. Tran, R. Ranganath, J. Paisley and D. Blei. Variational inference via x upper bound minimization. In Neural Information Processing Systems, 2017.
[10] J. Domke. Provable gradient variance guarantees for black-box variational inference. In Neural Information Processing Systems, volume 32, 2019.
[1]] Domke. Provable smoothness guarantees for black-box variational inference. In International
Conference on Machine Learning, volume 119, pages 2587-2596. PMLR, 2020.
[12] M. Dredze, K. Crammer, and F. Pereira. Confidence-weighreed linear classification. In Interna- lional Conference on Machine Learning, volume 307, pages 264-271, 2008.
[13] R. M. Gower, A. Defazio, and M. Rabbat. Stochastic polyak stepsize with moving target. arXiv:2106. / 1851, 2021.
[14] A. Gunes Baydin, B. A. Pearlmutter, A. Andreyevich Radul, and J. M. Siskind. Automatic differentiation in machine learning: a survey. arXiv e-prints, art. 1xXiv11502.0567 February 2015. doi: 10.443550/axxi..552..55555
[15] M. Hoffman, D. Blei, C. Wang, and J. Paisley. Stochastic variational inference. Journal of Machine Learning Research, 14:1308 14:1303-1347 2013.
[16] A. Hyvärinen. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6:695-709, 2005.
[17] C. Jones and A. Pewsey. The Sinh-Aresinh Normal Distribution. Significance. 16(2):6-7,06 04 2019. ISSN 1740-9705.
[18] 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.
[19] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
[20] J. Knoblauch, J. Jewson, and T. Damoulas. An optimization view on Bayes' rule: Reviewing and generalizing variational inference. Journal of Machine Learning Research, 23:
1-109, 2022.
[21] A. Kucukelbir, D. Tran, R. Ranganath, A. Gelman, and D. M. Blei. Automatic Differentiation Variational Inference. art. (0l::1603000888 March 2016.
[22] W. Lin, M. Emtiyaz Khan, and M. Schmidt. Fast and Simple Natural-Gradient Variational Infer- ence with Mixture of Exponential-family Approximations. arXiv e-prints, art. arXiv:1906.02911 June 2019, doi: 10.48550/arXiv. 11w.aa..i.........000 1906.02914.
[23] 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.
[24] M. Magnusson, P. Bürkner, and A. Vehtari. posteriordb: set of posteriors for Bayesian inference and probabilistic programming, November 2022. URL. https: httt::///tttttt /gggtggb......g com/ stan-derrposteriorce
11

---

[25] Charles C. Margossian. A review of automatic differentiation and its efficient implementation.
Wiley interdisciplinary reviews: data mining and knowledge discovery, 9, 2019. doi: 10.1002/ WIDM.130S
[26] T. Minka.
Expectation propagation for approximate Bayesian inference. In Uncertainty in Artificial Intelligence (UAI), pages 362-369, 2001.
[27] C. Naesseth, F. Lindsten, and D. Blei. Markovian
elimbing: Variational inference with KL ll q). In Neural Information Processing Systems, 2020.
[28] R. Ranganath, S. Gerrish, and D. M. Blei. Black Box Variational Inference. art. arXiv:1401.0118, December 2013.
[29] R. Ranganath, J. Altosaar, D. Tran, and D. Blei. Operator variational inference. In Neural Information Processing Systems. 2016.
[30] 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: thub. com/ roualdes/bridggestan.
[31] Salvatier, Wiecki T., and Fonnesbeck. Probabilistic programming in Python using PyMC3.
PeerJ Computer Science, 2016.
[32] U. Seljak and B. Yu. Posterior inference unchained with 20. art. arXiv:1901 .04454,444 January 2019.
[33] M. W. M. R. Andersen, A. Vehtari, and J. H. Huggins. Robust, automated, and accurate black-box variational inference. arXiv:2203.15945, 2022.
[34] Y. Yang, Martin, and H. Bondell. Variational approximations using Fisher divergence. arXiv e-prints, art. arv: 905.05284, May 2019.
[35] L. Zhang, B. Carpenter, A. Gelman, and A. Vehtari. Pathfinder: Parallel quasi-newton variational inference. Journal of Machine Learning Research, 23(306):1 --9,, 2022.
A Proof of Lemma 2.1
Lemma 2.1. Let θ C Id be a set that contains the support of p(0, x) and qw- (0) for some w". That p(0,z) x) 9w- (0) 0 for every 0 & O. Furthermore, suppose that the support of p x) The parameter w" satisfies
and qw- (0). and θ is path-connected
Vo log qw (0) = Vo log p0,,z), AA E e,
(3)
if and only wn also satisfies Eq. 1.
Proof. (1)
(3): Differentiating both sides of (1) in 1 gives
Ve log aw- (0) Ve log (0 (x) Ve (logp(O, x)) log p())
Vo log p(0 0 x)
e.
(3) - (1): Let Og be any arbitrary point in O. Because is path-connected, every 0 E θ is connected to On via differentiable path r(t) where r(0) O. and r(1) = 0. Integrating both sides (3) along this path r(t), using again that V, log (0(z) = Vo log p(0, w), and using the Fundamental Theorem of Calculus gives
/ ( log aw (ooo) dt
log 9w- (0) -log Jw- (0o)
(V, (o dt
log logg
00 e.
Rearranging and defining C : log 9m- (0,) logr(00)=)) gives
log qw- (0) = logg(00x))+++
v0 E e
12

---

By exponentiating both sides and integrating in 0 over è we have that
e / p(0)x)d00- e., Jece
Fw
(0)
Jace
where used that ④ contains the support of both p(0 ar) and aw- (0). Consequently C 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 (Mo, Eg) to denote the mean and covariance the previous time step of the method, thus dropping the iteration counter t.
Theorem B.I. (GSM updates) Let p(0, x) be given for some 0 E IRd and let (0) and g(0) be the multivariate normal distributions, respectively, with means Mo and μ and covariance matrices Eg and E. We seek the distribution
KL a)) --==
arg min
such that V (()) = Ve poo,
(11)
As shorthand, let g = Va log p(0, æ)) and let p be the positive root of the quadratic equation
p1++p) = g Eog + [ 0) g]
(12)
Then the solution is given by the following closed-form updates:
(po 0)g T
1 n4 1 1+ (40 =
Eg (g V (o
(13)
Σ = Eo + (N0 (140 0) (μ - 0) O)T T)
(14)
Furthermore, Eu is symmetric positive definite then is E.
Proof. The constraint in this optimization is given
g = Vo log (())
(15)
= Vo 1 111)=110 (0 11) log (1)))111)
(16)
= E -- 1).
(17)
The KI. divergence is given by
Σ Σg
(μ T (p- no) a).
KT q) 2 11 log
(18)
Dropping irrelevant terms from the optimization, we obtain the Lagrangian
1 [t [1| log Σ + (μ (o) - ( mo)] + +
C(u, Σ)
T
19---1- +
0))
2
(19)
is easier to optimize the matrix Σ instead Σ. We can enforce the symmetry --1 by writing
+ +
(20)
and performing unconstrained optimization over क. With respect to the latter, the gradients of the Lagrangian are given by
7c
Σ
dc O
Σ
ac
Oooor + bun
1
1
Oc
ac
0
2
(21)
13

---

Next we examine where the gradients of the Lagrangian vanish:
0C
0 Σ F1 mo)
p 140
(22)
0p
0c
0
0 = g - - E- (11
= μ Eg
(23)
0x
0c
0
0 = Eg - E + (n- ( (OO | 0)T +
(24)
En + ( u4) fO) A(1 0)T (μ 0)
(25)
We claim that these equations (though nonlinear) can be solved in closed form. The first step is to climinate À from eq. (25) using eq. (22). In this way we find
Σ Σg (μ- (1) uo) (p (o) 0)T (0-0)( pm))
(26)
Eg up T + nh +0p + Huflo 1OB 0pD
(27)
Eg (No ( 0)T (μ (1 -
(28)
It is worth highlighting the form of this equation:
Σ Eo + (mo () 0) (μ (1 - 0)
This is a simple rank-two update for Σ. Note that Σ = En if = n4g: also, the solution for Σ is determined by the solution for μr.
Ultimately we must solve for μ., but first it is useful to solve for the intermediate quantity g Eg > 0. From eq. (28), obtain
g Eg = g T Eog + [110 - []] - [( 0) "]]"""
(29)
and from eq. (23), obtain
g Eg 9 Eng + [(M0 -0) =)]=}}² (g [g Egg
(30)
As shorthand, let p = gg Eg. Then from eq. (30) we see that p satisfies the quadratic equation
p(1+p) = =g Eog
+ -og
Note that there are no unknowns on the right side of this equation. The correct solution is given by the positive root since p Eg > 0. Also note that (μ 0) g from eq. (23).
useful to define one final intermediate quantity before solving for μ.. Let
Eu
Mo O. Eng
Note that ED simply measures the degree to which the parameters (o(O) violate the desired
constraint V., log (0) V.. log p0, y). Put another way, if En 0, then have the trivial solution ja = 14 and = En.
Now have everything to express the solution for μ in a highly intuitive form; in particular, it will be immediately evident that μ - 140 as Eg + O. Starting from eqs. (23) and (28), we find
μ1 no
0 p4 + Eg.
(31)
=
0 - 14 + Σ. + (mo - 0)(10 T) (μx - 0) 4 0) 9,
(32)
6D + (ND () 0)T (u 0) | 0) g.
(33)
ED + (Nt ()(0 0)1 (μ po 4o 0) 0)
(34)
ED p(u mo) + (no 0)(( 0) (μ 0) 1g.
(35)
ED p(u mO) + (Ho 0)(m 4) 9.
(36)
ED (pl (w 0)g (mo 4).
(37)
14

---

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 p. Since p can determined via eq. (30), we arrive effectively at system of linear equations for μ.. Collecting terms, we obtain
[11 + p)T (no 0)g') " (1 44) ED-
(38)
We thus arrive at the closed-form update
po + [[L + p]1+ (NO -0g 0] EO
(39)
It is evident from this update that μa - oo as + 0. The matrix inverse in this update can also be computed efficiently from the Woodbury matrix identity.
In sum, the joint update for and E can be efficiently computed as follows:
1. Set g = Vw w log p(0, y) and ED = Eog 04 + 0.
2. Solve p p) 9 Cog [140 0) [ 2 for p > 0.
3. Compute μ μ40 + [1 + p)l + (mo 0)g T
4. Compute Σ = Eo + (no () 0) (μ 0) (μ -)T.
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
1 Eog + [no - 0) T){g} -
P =
2
We can solve the linear equation for μ in step of this procedure most efficiently by using the Sherman-Morrison formula
aI + +ugggg 1
ug
for every u, g: a.
u
Applying this formula in step 3, we find that
(mO 0)g
# = no 1+
Eo-
(40)
1+ p (no 0) g.
Next we relate this update to the difference in the scores of the target distribution and the current estimate of the variational approximation. Since Ve log (0) = E,, (0 mo), we have that
Eo En (g Σn (NO + O)) Eg (g Vo log (0)).
follows, by substituting the above expression into Eq. 40, that
(mo - 0)g 1 14 1 1+p p+ (Po 0)T g
I
Eg (Vo loz, 0)
Ve log (0)),
which gives the update in the statement of the theorem.
Finally prove that if Eo 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 p in eq. (30) in more convenient form. As shorthand, let
Mo
En (no - () O)T,
(41)
so that Mo captures the first two terms on the right side of eq. (28). Note that Mo is positive-definit fact that we will exploit repeatedly in what follows. In addition, recall that 0 Eg from eq. (23). Thus with this notation we can rewrite eqs. (28) and (30) as
Σ
Mo (g)(E))
(42)
p(1 +
g Mog.
(43)
15

---

Now let e be any normalized eigenvector of Σ; we want to show that its corresponding eigenvalue Xe positive. From eq. (42), follows that
de = e T Ee
(44)
T Mo (Eg)
(45)
c e T²)
(46)
Note that e g = then it follows trivially that A e. "Mee > So we only need consider the non-trivial case e * O. To proceed, we note that
(e Mog)² (e (e " < (Me) Mog).
Mogg),
(47)
where have used the Cauchy-Schwartz inequality to bound (e Mog) in terms of (e Mae), the latter of which appears in eq. (46). Substituting this inequality into eq. (46), we find that
(e Mog)²
Ae
g)
(48)
g Mog
To prove that de > we need one more intermediate result. Focusing on the rightmost term in this equality, note that
Ae g) e Eg
T Mo (Eg)(Eg) g e emMog Ae gg Ae(eTg) Eg),
(49)
and rearranging the terms in this equation, find
Mog Ae g)(1 (((TT))(1 +g Eg).
(50)
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 Mog in eq. (48), we find:
[ee + g _gg]] g Mog
de 2
22 g)²
2 g)²
(1 + g Eg)²
1
g ''Mgg
[11 [1+ p)?
X²((e g)²
p(1 +
Xe(eTg)
P
0,
where the final inequality follows because the individual terms X² (e ge) positive; note that de cannot be equal to zero because this contradiets the equality in eq. (46). This completes the proof. Perhaps it is useful that this derivation also gives upper bounds on X namely
g) and are all strictly
1 (e g))
Eg
>
= de 1 ge
(51)
d
P
(e 2)
16


Reviews:
[{"review_id": "GgT2xdYkyf", "reviewer": "Reviewer_Ciug", "summary": "The 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.", "questions": "These 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.", "limitations": "The 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.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "The 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.", "weaknesses": "The 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> I will update my rating if the author's answers address my concerns. \n\nWe thank the referee for their thoughtful reviews and kind consideration. Below, we address the weakness and questions raised above and will include this discussion in the revised manuscript.  \n\n> Why do you think GSM-VI is so effective and scales so well?\n\nThere are 3 intuitions that help explain why GSM-VI is more efficient than BBVI- \n\n1) GSM-VI does not rely on Taylor approximations. BBVI is a stochastic gradient method and hence relies on the 1st order Taylor approximation of the objective function. In contrast, GSM-VI computes the exact projection onto the score matching constraint.\n2) BBVI requires tuning a learning rate and its scheduling. In contrast, GSM-VI does not have a learning rate parameter. Instead, it is able to adaptively make large jumps in the initial iterations (see Fig. 1.a in the paper) and make smaller adjustments as the approximation converges to the target. \n 3) BBVI relies on a scalar signal, that is it attempts to increase the ELBO (scalar function) by using the steepest ascent direction (gradient). GSM-VI instead uses $d$ equations at each iteration to determine the update, one for each element of the score. Since the number of constraints increases linearly with dimensions, we believe it partly explains why GSM-VI scales well with dimension.\n\n> Why is a batch size of 2 best? This seems odd.\n\nReferee's intuition is correct. Our understanding of the impact of batch size has now improved with more experiments-\n\n- For Gaussian targets, GSM performs equally well for all $B\\geq1$. There are minor gains for larger batches as the dimensionality increases $d\\gtrapprox100$, but these are marginal enough that B=2 is a good conservative default.\n- For non-Gaussian targets, we now recommend a larger batch size. We find that these converge to a more stable solution i.e. smaller oscillations in KL divergence. We show this in Fig. 3 of the above attached pdf. We found that $B \\geq 8$ suffices for $d\\lessapprox100$\n\n> Performance is always measured in terms of number of gradient evaluations. The paper should also report actual running times.  \n\nThere are 2 reasons to measure gradient evaluations- \n1. Actual running time is very sensitive to actual implementation details. \n2. Often in real-world problems, the computational bottleneck is evaluating the target log density and its gradient (for instance it can involve solving complex ODEs). Hence we focus on minimizing these evaluations.\n\nHowever we acknowledge that the referee raises a good point. We have now timed GSM and BBVI updates for fitting a full rank Gaussian target with 2048 dimensions using batch size of 1 in Jax after JIT-compilation. GSM update takes 230 ms/iteration while BBVI takes 227 ms/iteration. \n\n> Will the quadratic update step will worsen GSM-VI's performance for high-dimensional problems. Does BBVI overtake GSM?\n\nIn this paper, we only consider the Gaussian variational family with full-rank covariance matrix and hence both BBVI and GSM have the same $O(d^2)$ complexity. This complexity is unavoidable since we need to store and update the $d \\times d$ covariance matrix.\nThus based on our timing test for D=2048 problem above, and the fact that GSM requires $\\sim10$x less gradient evaluations for almost all examples considered, we are confident that BBVI will not overtake GSM-VI as the dimension increases.\n\n> What was the batch size for BBVI in the benchmarks?\n\nWe fixed the batch size of BBVI to 2 for all experiments, same as GSM, to keep the two algorithms at equal footing. Since we did not tune the batch size of GSM for individual experiments, we did not vary the batch size of BBVI either (however we did tune learning rate of BBVI).  In Figure 1 in the uploaded response pdf, we now show results for varying both- batch and learning rate of BBVI and find that GSM-VI still outperforms BBVI in all cases. Please see general answer for more details.\n\n> How was GSM-VI implemented? What kind of system was it run on? \n\nOur implementation of GSM-VI was in Tensorflow but now we have a Jax package that will be released after the review process. We will also implement GSM in BlackJax package. Also note that if one has access to the score function of the target, GSM updates can be written in native python without any ML or auto-diff. All the experiments were done on a single core of a personal CPU desktop.\n\n> How were the posteriordb examples selected?  What's their dimensionality? \n\nWe compared GSM & BBVI on all the 50 models in PosteriorDB and found GSM to outperform on all of them. For brevity, we chose 8 of these models for the paper based on 2 considerations- i)  Each model represents a different class of statistical models, with different complexities, and ii) we have an equal representation of both Gaussian and non-Gaussian targets. \nAll the models are low-dimensional which allows us to look at marginal and joint distributions of parameters and thus go beyond only comparing KL divergences in evaluating different algorithms. \nSimilar choice of models was made by a recent paper on Pathfinder algorithm for fitting a Gaussian variational distribution (Lu et.al. JMLR ; 23(306):1−49, 2022)\n\nClasses and dimensionality of the 8 models used are- Generalized linear models (d=26), Differential equation dynamics (8), Hierarchical meta-analysis with centered & non-centered parameterization (10), hidden Markov models (8), time-series model (7),\nGaussian processes (13), Gaussian mixture model (d=5).\n\n> Figure 4: what happens for settings where both $s\\neq1$ and $t\\neq1$. \n\nWe did experiments varying both skewness and tail-weights at the same time for different dimensions and found that in all our experiments, GSM with batch $B=2$ always converged to a similar solution as BBVI, but faster. Hence for the sake of clarity, we separated these two axes and showed them separately. An example of $s\\neq1$ & $t\\neq1$ case is in Fig. 3 of the attached pdf."}, {"review_id": "5qsgsGf3ZB", "reviewer": "Reviewer_h2vB", "summary": "This 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", "questions": "The 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", "limitations": "I haven't seen any specific text on the limitations of the approach and I think this could be included in the revision. \n", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "I 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", "weaknesses": "One 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> The 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? \n\nWe can certainly add some more intuition. We know now from Lemma 1 that if Eq 3 holds for every $\\theta$, then our variational family has found a perfect fit. Any practical method based on this observation can however only sample a batch of $\\theta$’s at each iteration. When updating our variational parameters to satisfy the score matching equations over a given batch, we need to be careful that in the remainder of the domain (outside the batch), our variational family remains a reasonable fit. One way to ensure this is to always make the smallest possible change in our variational family to fit the score matching constraints over the current batch. This ensures that for the rest of the domain that has been explored before, the variational family will still be an approximately good fit. Thus the need to minimally adjust $q$ \n is because we only have partial information (a batch) of the constraints in (3). This approach is called the “least change” approach in quasi-Newton methods [Goldfarb] or the “passive-aggressive” approach in PA methods [Cramer].\n\n[Cramer] Online Passive-Aggressive Algorithms (JMLR 2006) \\\n[Goldfarb]  Goldfarb, D. (1970), \"A Family of Variable Metric Updates Derived by Variational Means\", Mathematics of Computation, 24 (109): 23–26,\n\n> 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\nWe thank the referee for this suggestion and will put more emphasis on this in the revision\n\n> If 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.  I 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\nWe agree with the referee’s intuition that mode collapse is a possibility in our algorithm. However this is not necessarily due to score matching, but rather due to generating samples from the variational distribution itself. This is also an issue with BBVI, but unlike BBVI, sampling in this way is not a requirement for GSM, but only convenience. If we have access to another distribution that covers the domain (for e.g. a prior distribution) and allows efficient sampling, we can always generate samples from it. \n\nThis being said, we did investigate different initializations for both Gaussian and non-Gaussian targets. These included starting from–  a standard normal, random Gaussian with broad and narrow covariance, from the mode with identity or using an LBFGS approximation for the inverse covariance matrix, from the resulting approximation of the Pathfinder algorithm etc. \\\nIn all these cases, we always found the same, global solution for GSM. The only setting where the GSM solution depended on initialization was when the target distribution is mulit-modal with separated modes and in this case GSM generally converged to the nearest mode. However this is a known issue with almost all inference algorithms.  We can elaborate on this in the revised manuscript.\n\n> 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\nOur setting is Gaussian families where there are $d^2+d$ variational parameters and $d$ latent variables. Thus the number of variational parameters is always greater than the number of  parameters/latent variables. In the general setting, for non-Gaussian variational families that have fewer parameters than latent variables, this score matching approach is not possible. But we are unaware of any practical examples where such a situation arises. Thus we now find that discussing this in our paper was a digression, and as such we will either remove this discussion, or expand upon it to make it clear. We also apologize that we are not familiar with this setting of “small hyper-net to generate large sets of parameters of a neural net”, and welcome a reference and are open to consider this during the discussion phase.\n\n> 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.\n\nIn this paper, we have focused only on the full-rank Gaussian variational family which cannot be extended to these highly over-parameterized models like deep networks. In the future, we plan on studying applications for mean-field Gaussian families (updates for which can be trivially implemented), and low-rank approximations. Both of these are better suited in the setting that the referee has pointed out. In that regard, we thank the referee for the reference and pointing to an interesting direction of research.\n\n> I haven't seen any specific text on the limitations of the approach and I think this could be included in the revision.\n\nWe have discussed some limitations throughout the paper, especially in the conclusion and future work section, but will be happy to collect them in one sub-section in the revised paper.\n"}, {"review_id": "447kP2xcqm", "reviewer": "Reviewer_DLrt", "summary": "A 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", "questions": "I 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", "limitations": "Yes, limitations are addressed.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "A 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.", "weaknesses": "The 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> I don't concur with the use of the term \"closed-form updates\" of line 229 as GSM-VI\n\nBy closed-form updates, we refer only to the fact that (4) has a closed form solution for the Gaussian variational family for a generated sample. That is, the explicit update equations are given in equations (5) and (6) (detailed in Theorem 2.2). By closed-form we do not mean that the updates in equations (5) and (6) are a solution to global score matching equations in (3). We will clarify this distinction in the revision. \n\n> I would be ready to raise my score if the rebuttal addresses the following points:\n\nWe thank the reviewer for this consideration. We have now addressed both these concerns regarding additional metrics in the experiments, and the presence of oscillations.\n\n> Nuancing the experiments on simulated data with a metric less disadvantageous of VI, e.g. reverse KL, for at least one data set.\n\nFor real-world problems from posteriodb models, we already show the results for reverse KL divergence.\n\nBut we agree with the referee that even for synthetic models, showing the evolution of reverse KL i.e. the objective which is explicitly optimized by BBVI, in addition to forward KL can be instructive. We have now attached a version of this figure in the uploaded pdf  in the general response for a Gaussian target of 16 dimensions (see Figure 1). In this figure, we also present the results by varying hyperparameters of BBVI i.e the learning rate and batch size. This clearly shows that for Gaussian targets, BBVI can converge to the same quality solution as GSM but requires much larger batch size and small learning rates (only B>4, lr<0.01 converged in the 5000 iterations). Hence the conclusions remain unchanged from the Forward KL metrics presented in the original manuscript. Please see the attached pdf and the discussion in the  \"Author Rebuttal\" at the top for more details.\n\nWe will add this discussion and results showing reverse KL for other examples in the supplementary material of the revised manuscript.  \n\n> Addressing the issue raised under weaknesses regarding oscillatory behavior.\n\nThrough additional experiments in the attached pdf, we believe we now have a better understanding of the oscillatory behavior that the referee points out, and we summarize this in the following points:\n- When the variational family is well specified, i.e the target distribution is Gaussian or can be fit by one, we find no oscillations in GSM for any batch size and the algorithm converges to an optimal solution. \n- When the target is non-Gaussian, the monitored KL divergence has oscillations. It turns out that this was because we were using batch size B=2 for GSM for all experiments. For larger batch sizes, these oscillations are suppressed very easily. We have attached a figure (see Figure 3) in the pdf. In the right panel of that figure, we also show the marginal histograms of a parameter for every 200th iteration in the last 1000 iterations for a GSM run with batch size of 8, and these show that the distribution has indeed converged. A detailed discussion of this figure is presented in the general answers section above.\n\nWe would also like to point out that these oscillations in KL divergence are also present in BBVI. Any stochastic optimization can dampen these by scheduling a learning rate to converge to 0, which is also what is often done in BBVI. While we have chosen not to do this for GSM, this same procedure can be applied to dampen the oscillations for GSM. "}, {"review_id": "ojLnyUJnq1", "reviewer": "Reviewer_CVEd", "summary": "The 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 ", "questions": "Why 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). ", "limitations": "The 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.", "rating": 3, "confidence": 5, "soundness": 2, "presentation": 4, "contribution": 2, "strengths": "Originality\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.", "weaknesses": "Technical 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">It would be important to better analyze, if there are any mild conditions for convergence, and which criteria the learned approximation fulfills. \n\nWe agree with the referee that its important, but challenging to analyze given that GSM does not explicitly minimize any divergence. Empirically, we have so far observed that GSM converges to solutions that have the best forward KL divergence (and similar reverse KL divergence as BBVI for non-Gaussian target). Furthermore, as far as we know, there are no theoretical convergence guarantees for BBVI when the variational family is misspecified either. It is a difficult open question still under active research years after the algorithm has been widely used.\n\n> In the setting where the target distribution can be perfectly approximated, ..., 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\n\nWe thank the referee for making this point. We have now updated Lemma 1 to clarify this. \n\n> Natural gradient based methods are known to significantly outperform methods based on the vanilla gradient (such as the BBVI). For changing my opinion, the paper also needs to discuss the relevant methods and ideally add a comparison to the work by Lin et al. 2019 or Arenz et al. 2020.\n\nWe thank the referee for pointing to natural gradient descent based VI as a baseline. We will certainly discuss these works and add this  baseline for comparison in the revision. As requested, we have added a comparison to NGD-VI based on Lin et.al. for a Gaussian target in Fig. 2 in the uploaded pdf. However, note that for full rank Gaussians, NGD has cubic scaling $d^3$ with dimension while GSM has quadratic $d^2$. We now summarize our implementation & results (also see the general answer above) \n\n*Implementation* \n- We use only the first order information for fair comparison to GSM.\n- The updates for variational parameters $\\mu_i$ and $\\Sigma_i$ follow Eq. 16 of Lin et.al (arxiv:1906.02914) adapted for a single Gaussian instead of GMM \n- Eq 16 requires the Hessian of the ELBO. To approximate this with 1st-order information, we implement VOGN update (Eq. 10 of Khan et.al, arXiv:1806.04854). This uses Gauss-Newton approximation for the Hessian of  the likelihood $\\log p(\\theta | x)$ and $\\Sigma_i^{-1}$ for Hessian of $\\log q$.\n\n*Results*\n\nFig. 2 of attached shows results for NGD-VI for a D=64 Gaussian target by varying the hyperparameters- learning rate and batch size. We find that\n- As the referee pointed out, NGD-VI significantly outperforms BBVI.\n- “Optimally tuned\" NGD-VI can be competitive with GSM. \n- However the performance of NGD-VI is very sensitive to the hyper-parameters. Larger batch sizes allow us to use larger learning rate so they need to be tuned together. However a large step size gets stuck oscillating around the minima and hence we need to additionally tune the scheduling of the learning rate.\n- We couldn't experiment with optimizers like Vprop/Vadam from Khan et.al. for the rebuttal but will investigate those for the updated manuscript.  \n- For $D\\geq256$, NGD-VI was prone to diverging for $lr>0.01$. To ensure convergence, we added an additional regularization parameter for Hessian estimates in VOGN. \n\t\n*Compared to GSM*\n1) For VI with a full-rank Gaussian, GSM has a quadratic complexity of $d^2$ while NGD-VI is cubic $d^3$ as it requires estimating both $\\Sigma$ and $\\Sigma^{-1}$ for each iteration. \n2) An optimally tuned NGD-VI can be competitive with GSM, but this will require tuning the learning rate, it's scheduling and batch size. In contrast, GSM has no learning rate parameters to tune, and is largely insensitive to batch size. \n3) For large problems, we need a regularization hyperparameter for Hessian in VOGN updates. GSM has no such issues in scaling. \n\nThus based on our experiments, we argue that as compared to NGD-VI, GSM is faster in terms of gradient evaluations, has a smaller iteration complexity, and is easy to tune. \n\nHowever we agree with the referee that it was a relevant baseline missing from the draft and would like to thank them for pointing this.\n\n> BBVI should not converge to a suboptimal solution for Gaussians. Can you provide code and details on hyperparameter tuning?\n\nWe agree that BBVI will converge for Gaussian targets, and have now uploaded additional figures to show this. In particular, BBVI with a well chosen learning rate and batch size, will converge given a sufficient number of iterations. For instance, in Fig. 1 of uploaded pdf, we show the performance of BBVI for a D=10 Gaussian target for different values of the hyperparameters- learning rate $lr$ & batch size $B$. In this case, BBVI for $B\\geq4$ & $lr \\leq 0.01$ eventually does converge to the same solution as GSM in terms of reverse KL (right panel, blue & orange lines). We hope that this resolves the misunderstanding- we are not claiming that BBVI *never* converges to the correct solution for the Gaussian target. Instead, it did not converge under our finite budget. With more iterations BBVI eventually converges.\n\n*Hyperparameters*- We fixed $B=2$ and then tuned $lr$ with a grid-search, but we didn't tune it's scheduling. As seen in Fig 1, while larger step-sizes converge faster at the beginning, they only converge in a small radius around the minima. Thus one also needs to tune the scheduling of step size. In Fig.1 attached, we now also tune the batch size $B$ of BBVI and find that GSM still outperforms BBVI. We will include this discussion and figure in revision to give more context to the results of BBVI. \n\nWe have provided our code with the initial submission in the supplementary. We will make a package public with all 3 algorithms after the review process. "}]

Author rebuttal:
We thank all the referee for their thoughtful reviews. Based on comments from different reviewers, we have now compiled three new figures from some old, and one entirely new experiment. Please see the attached pdf. These serve to answer some of the questions raised below and so we begin by discussing these figures here as a preamble. We will add these figures and discussion to the revised manuscript.  

**Figure 1. Convergence of BBVI:**

In response to Reviewer **DLrt**'s question we have included an experiment tracking the reverse KL.
In Fig. 1 of the attached pdf, we fit a full-rank Gaussian target distribution of 16 dimensions with BBVI and GSM. Here we monitor the reverse (backward) KL divergence which is explicitly minimized by BBVI, and show results for doing a hyper-parameter search on both the parameters of BBVI- learning rate and the batch size. The top row shows reverse KL (y-axis) on linear scale and the bottom row on log-scale for clarity. We highlight two key takeaways from the figure:

1) Given enough computational budget, BBVI does converge to the same quality of solution as GSM, in the current example for larger batch sizes ($B\geq 4$) and smaller learning rates $lr \leq 0.01$ (see blue and orange lines in last two panels). However its performance is quite sensitive to the hyperparameters. The figure also suggests that performance of BBVI can be improved by scheduling the learning rate, but this will require a greater hyperparameter search. 

2) Even for optimally tuned batch size and learning rate, GSM significantly outperforms BBVI and requires orders of magnitude less gradient evaluations for convergence even in terms of reverse KL, which is the metric that is explicitly minimized by BBVI.
We will include similar figures for other synthetic experiments in the revised version of the paper..

**Figure 2. Natural Gradient Descent (NGD)-VI**

We have added natural gradient descent VI for maximizing ELBO as a new baseline for comparison with GSM since it has been shown to significantly outperform BBVI. We indeed find this to be the case in our experiments. However as shown in Fig. 2, in our regimes of interest, we find that even compared to NGD-VI, GSM is faster in terms of gradient evaluations, has a smaller iteration complexity, and does not require tuning for variational inference with a full-rank Gaussian distribution.

*Implementation of NGD-VI:* We implemented the updates to variational parameters $\mu$ and $\Sigma$ based on Eq. 16 of Lin et.al (arxiv:1906.02914) (adapted for a single Gaussian instead of mixture model by setting $\delta_c=1$). These updates however require the Hessian of the objective function (ELBO). To approximate this with only first-order gradient information for comparison with GSM, we have used the VOGN update (Eq. 16 of Khan et.al., arxiv:1806.04854). This approximates the Hessian of the likelihood term $\log p(\theta | x)$ with a Gauss-Newton approximation and combines it with the correct Hessian of the entropy term $\log q(\theta)$- $\Sigma^{-1}$.

Figure 2 shows results of fitting a 64 dimensional Gaussian target with GSM and NGD-VI while varying its hyperparameters- the learning rate and the batch size. We monitor reverse KL, which is explicitly optimized by NGD-VI, for both linear (top) and log-scale (bottom row) on y-axis for clarity. We find that-
1) While an optimally tuned NGD-VI can be competitive with GSM, its performance is very sensitive to tuning the hyperparameters. This will require not only tuning the batch size and learning rate, but also the scheduling of learning rate. In comparison, GSM has no learning rate parameter to tune. 
2) For large dimensions $d \geq 256$, Gauss-Newton approximation for Hessian in VOGN algorithm starts to diverge. This can possibly be corrected by using a regularization hyper-parameter for Hessian, or a different Hessian approximation, but both these require more tuning. 
GSM does not face any such issues in scaling to high dimensions
3) Finally, we also note that for a  full-rank Gaussian variational family, every iteration of GSM has a quadratic complexity $d^2$ as compared to cubic $d^3$ of NGD-VI which requires one to estimate both $\Sigma$ and $\Sigma^{-1}$ for each update. 

Based on these results, we conclude that  GSM-VI outperforms NGD-VI in our regimes of interest. We will add this baseline to other experiments in our revisions. 

**Figure 3. Batch size, GSM-VI and non-Gaussian targets**

We have run more experiments with GSM-VI for different non-Gaussian targets and varied the batch size to better understand its convergence. Figure 3 summarizes our findings. In it, we compare the performance of GSM-VI on a synthetic non-Gaussian target (D=10) for different batch sizes (B=2, 8, 32) with BBVI (B=8, 32). We show the same forward KL divergence as the paper (left panel), marginal histogram for one of the parameters for all algorithms/configurations (middle panel) and additionally show the same marginal histogram for every 200th iteration in the last 1000 iterations of GSM-VI with B=8 to demonstrate that the distribution has indeed converged (right panel). 

We find that larger batch sizes of GSM ($B \geq 8$)  lead to a more stable solution than both GSM with $B=2$ and BBVI, i.e. they have smaller (or no) oscillations in the forward KL metrics. GSM also converges to a point with a better forward KL divergence. This is also reflected in the marginal histograms which have larger dispersion for GSM than BBVI. Thus GSM does not lead to mode-seeking behavior. This can be desirable, for instance, as it means GSM posteriors can be corrected with importance sampling more easily.

Meta-review: Thank you for your engaging interaction with the reviewers during the discussion. The reviewers and I are generally in agreement that this paper should be accepted. The paper is written very clearly, the methodology is novel and interesting (not just an iteration on existing variational methods), and the method seems to perform well on the baselines tested. However, there are no theoretical guarantees, and the range of baselines in the original submission was very limited.

In the camera ready, the paper should be clear about the lack of analysis, and provide a more extensive set of comparisons to existing approaches with a fair and unbiased analysis of results (among other edits suggested by reviewers) to be suitable for publication. The authors have already done some work in this direction during the discussion phase, so I am confident in the authors' ability to complete this work for the camera ready.
Decision: Accept (poster)