Particle-based Variational Inference with Preconditioned Functional Gradient Flow
Authors: Hanze Dong
Xi Wang
LIN Yong
Tong Zhang
Venue: ICLR 2023 poster (2023)
Keywords: Posterior Sampling
Particle-based VI

Abstract: Particle-based variational inference (VI) minimizes the KL divergence between model samples and the target posterior with gradient flow estimates. With the popularity of Stein variational gradient descent (SVGD), the focus of particle-based VI algorithms has been on the properties of functions in Reproducing Kernel Hilbert Space (RKHS) to approximate the gradient flow. However, the requirement of RKHS restricts the function class and algorithmic flexibility. This paper offers a general solution to this problem by introducing a functional regularization term that encompasses the RKHS norm as a special case.  This allows us to propose a new particle-based VI algorithm called  preconditioned functional gradient flow (PFG). Compared to SVGD, PFG has several advantages. It has a larger function class, improved scalability in large particle-size scenarios, better adaptation to ill-conditioned distributions, and provable continuous-time convergence in KL divergence. Additionally, non-linear function classes such as neural networks can be incorporated to estimate the gradient flow. Our theory and experiments demonstrate the effectiveness of the proposed framework.

Published as a conference paper at ICLR 2023
PARTICLE- VARIATIONAL INFERENCE WITH PRECONDITIONED FUNCTIONAL GRADIENT FLOW
Hanze Dong", Xi Wang', Yong Lin', Tong Zhang*!
* Department of Mathematics, HKUST
i Department of Computer Science and Engineering, HKUST
1 College of Information and Computer Science, UMass Amherst
ABSTRACT
Particle-based variational inference (VI) minimizes the KL divergence between model samples and the target posterior with gradient flow estimates. With the popularity of Stein variational gradient descent (SVGD), the focus of particle- based VI algorithms has been on the properties of functions in Reproducing Kernel Hilbert Space (RKHS) approximate the gradient flow. However, the requirement of RKHS restricts the function class and algorithmic flexibility. This paper offers a general solution to this problem by introducing functional regularization term that encompasses the RKHS norm as a special case. This allows us to propose a new particle-based VI algorithm called preconditioned functional gradient flow (PFG). Compared to SVGD, PFG has several advantages. It has a larger function class, improved scalability in large particle-size scenarios, better adaptation ill-conditioned distributions, and provable continuous- convergence in KL divergence. Additionally, non-linear function classes such as neural networks be incorporated estimate the gradient flow. Our theory and experiments demonstrate the effectiveness of the proposed framework.
INTRODUCTION
Sampling from unnormalized density is a fundamental problem in machine learning and statistics, especially for posterior sampling. Markov Chain Monte Carlo (MCMC) (Welling & Teh 2011 Chen et al. 2014) and Variational inference (VI) (Ranganath et 2014 Fordan Hoffman et al. 2he4
jtt al. 1999; Blei et al. 2017) are two mainstream solutions: MCMC is asymptotically unbiased but exhauteexhausted: is computationally efficient but usually biased. Recently, particle- based VI algorithms (Liu & Wang 2016 Detommaso et al., 2018; Liu et al., 2019) tend to minimize the Kullback-Leibler (KL) divergence between particle samples and the posterior, and absorb the advantages of both MCMC and VI: (1) non-parametric flexibility and asymptotic unbiasedness; (2) sample efficiency with the interaction between particles; (3) deterministic updates. Thus, these
algorithms are competitive in sampling tasks, such as Bayesian inference (Liu & Wang 2016 Feng
al. 2017 Detommaso et al. 2018), probabilistic models (Wang & Liu 2016 Pu et al.) 2017).
Given target distribution p. (x), particle-hase VI aims to find g(,, æ). so that starting with X, ~ Po.
the distribution p(t, of the following method: d.X,- = Xt)dt, converges p=(x) as + co. By the continuity equation (Jordan et al. 1998), we can capture the evolution of p(t, x) by
Öp(t,z) x) Ot
V ()(t,))
In order to measure the "closeness" between p(t, -) and p.. we typically adopt the KL divergence, DKL(t) = f p(t, x))
p(t,,)
dr.
(2)
p=(x)
Using chain rule and integration by parts, we have
(1)x((1) =- =- p(t,x))V g(t,z) + g(t,x)) "V, V..nn.
(3)
dt

---

Published as a conference paper at ICLR 2023
which captures the evolution of KL divergence.
To minimize the KL divergence, one needs to define a "gradient" to update the particles as our g(t, z). The most standard approach, Wasserstein gradient (Ambrosio et al. 2005). defines gradient for p(t, æ) in the Wasserstein space, which contains probability measures with bounded second moments. In particular, for any functional L that maps probability density p(t, x)) to non-negative scalar, we say that the particle density p(t, x) follows the Wasserstein gradient flow of C if g(t, x) is the gradient field of (F(IR))ffncttion derivative of L [Villani 2009). For KL divergence, the solution However, the computation of deterministic and imm-innhoooeeneuus Wasserstein
is V In (x(a) gradient is non-trivial. It is necessary to restrict the function class of g(4,xr) to obtain a tractable form.
p(t.2))
Stein variational gradient descent (SVGD) is the most popular particle-based algorithm, which provides tractable form to update particles with the kernelized gradient flow (Chewi et al. 2020 Liu 2017). It updates particles by minimizing the KL divergence with a functional gradient measured in RKHS. By restricting the functional gradient with bounded RKHS norm, it has an explicit formulation: g(t, z) can be obtained by minimizing Eq. (3). Nonetheless, there are still some limitations due to the restriction of RKHS: (1) the expressive power is limited because kernel method is known to suffer from the curse of dimensionality (Geenens) 2011); (2) with n particles, the O(n²) computational overhead of kernel matrix is required. Further, we identify another crucial limitation of SVGD: the kernel design is highly non-trivial. Even in the simple Gaussian case, where particles start with (0, I) and p. (p.,, ).), commonly used kernels such as linear and RBF kernel, have fundamental drawbacks in SVGD algorithm (Example
Our motivation originates from functional gradient boosting (Friedman) 2001: Nitanda & Suzuki 2018 Johnson & Zhang 2019). For each p(t, z), we find a proper function as g(. x) in the function class to minimize Eq. (3). this context, we design regularizer for the functional gradient to approximate variants of "gradient" explicitly. We propose regularization family to penalize
the particle distribution's 's functional gradient output. For well-conditioned - Imp we can adapt
approximate the Wasserstein gradient directly; For ill-conditioned - V In p., regularizer to approximate a preconditioned one. Thus, our functional gradient is an approximation to the preconditioned Wasserstein gradient. Regarding the function space, we do not restrict the function in RKHS. Instead, we can use non-linear function classes such as neural networks to obtain a better approximation capacity. The flexibility of the function space can lead to better sampling algorithm, which is supported by our empirical results.
Contributions. We present novel particle-based VI framework that incorporates functional gradient flow with general regularizers. We leverage a special family of regularizers to approximate the preconditioned Wasserstein gradient flow, which proves to be more effective than SVGD, The functional gradient in our framework explicitly approximates the preconditioned Wasserstein gradient, making well-suited to handle ill-conditioned cases and delivering provable convergence rates. Additionally, our proposed algorithm eliminates the need for the computationally expensive O(n²) kernel matrix, resulting in increased computational efficiency for larger particle sizes. Both theoretical and empirical results demonstrate the superior performance of our framework and proposed algorithm.
2 ANALYSIS
Notations. In this paper, use z to denote particle samples in Rdd The distributions assumed to be absolutely continuous w.r.t. the Lebesgue measure. The probability density function of the posterior is denoted by p.. p(t, x)) (or pt) refers particle distribution at time t. For scalar function x)t, Vzp(t, x)) denotes its gradient w.r.t. x. For vector function g(t, x)). V..(tt, x),, V, gtt, x), Vāg(t, I) denote its Jacobian matrix, divergence, Hessian w.r.t. z. partial derivative w.r.t. t. Without ambiguity, V stands for V. for conciseness. Notation ||0|| 2 stands for x: T H. and ||>|| is denoted by ||=||. Notation || || denotes the RKHS norm on Rdd.
@p(t.x) and @g(t.z) denotes the 0t
Ør
We 1et x) belong to a vector-valued function class F. and find the best functional gradient direction. Inspired by the gradient boosting algorithm for regression and classification problems, we approximate the gradient flow by a function g x)) E F with a regularization term, which solves the
For any dositive-definite matrix, the condition number is the ratio of the maximal eigenvalue to the minimal eigenvalue. low condition number well-conditioned, while a high condition number ill-conditioned.
2

---

Published as a conference paper at ICLR 2023
(3) ]l Rg (lineur)
2 11
7e (RBF)
06
2 Pp I.FI
(8)
2 Epp
111
(e) Optimal transport
Figure 1: Illustration of the functional gradient (t, compared with optimal transport. (a)-(d) denotes the corresponding Q of the regularized functional gradient. (a)-(b) are SVGD algorithms with linear and RBF kemel. Optimal transport denotes the direction of the shortest path towards p..
$ 1l QU) +1 Rd
(RBF)
()Q)))
F 1/1
(0 F E,
Figure 2: Evolution of partiele distribution from (11, 0 to (V 20) diag(100,))) 1)) (first
row: evolution particle mean f4; second row: particle distribution p(5, x) at 5)
following minimization formulation:
[-f P(. f(x) (a) +Q5]
g(t, x²) arg min JEF
where Q(-) is regularization term that limit the output magnitude of f. This regularization term also implicitly determines the underlying "distance metric" used to define the gradient estimates g(t, x)) in our framework. When (z)) = =.2 p(t, ff)
ddx,, Eq. (4) is equivalent to
/ p(t,, (F60) Vn (x (x(x)
g(t, x)) arg min
dx.
(5)
Jer
p=(a) which is the direction of Wasserstein gradient flow. Interestingly, despite the computational intractability of Wasserstein gradient, Eq. provides a tractable variational approximation.
If F is well-specified, i.e. In p-(2)
F, we have (t,, In (=(2)
Pe(r)
For RKHS, we will show that SVGD a special case of Eq. 4 where Q()) = _ 11 (Sec. 3.2.1
We point out that RKHS norm usually fails to regularize the functional gradient properly since fixed for any p(t, x). Our next example shows that SVGD is suboptimal solution for Gaussian case.
Example Consider that p(t. ( E.), pa ( E.). We consider the SVGD algorithm with linear kernel, RBF kernel, and regularized functional gradient formulation with Q()
} Ep, II and Q(f) = Ep, 1/"1"_.. Starting with N (0, I), Fig. 7 plots the g(0, x) with different
Q and the optimal transport direction; the path of μa and p(5, x.) illustrated in Fig. 2 detailed mathematical derivations and analytical results are provided in Appendix
The
Example shows the comparison of different regularizations for the functional gradient. For RKHS norm, we consider the most commonly used kernels: linear and RBF. Fig. shows the g(0, a)) of different regularizers: only Q()) E || approximates the optimal transport direction,
while other g0.. deviates significantly. SVGD with linear kernel underestimates the gradient of large-variance directions; SVGD with RBF kernel even suffers from gradient vanishing in low-density area. Fig. demonstrates the path of 14 with different regularizers. For linear kernel, due to the
3

---

Published as a conference paper at ICLR 2023
curl component (Lati term in linear SVGD is non-symmetric; See Appendix for details), p(5, x) rotated with an angle. For RBF kernel, the functional gradient of SVGD not linear, leading to slow convergence. The La regularizer is suboptimal due to the ill-conditioned Σ.. We can that Q() 2 Ep | E- produces the optimal path for 14 (the line segment between fl0 and p.).
2.1
GENERAL REGULARIZATION
Inspired by the Gaussian case, consider the general form
Q(f(x))= 1
p(t,
(6)
where is symmetric positive definite matrix. Proposition shows that with Eq. (6). the resulting functional gradient is the approximation of preconditioned Wasserstein gradient flow. It also implies the dell-definedness and existence of g x) when FF is closed, since || |l is lower bounded by 0.
Proposition 1. Consider the case that Q( 1 S p(t, x)), (x) where H is a symmetric positive definite matrix. Then the functional gradient defined by Eq. (4) is equivalent as
f p(t, F(()) - V In (x( JEF
g(t,x) = arg min
dx
pt x)
Remark. Our regularizer is similar the Bregman divergence mirror descent. We can further extend Q(-) with a convex function h(-) R - [0, oo), where the regularizer is defined by Q(f(z)) = (p.h(((r) We can adapt more complex geometry our framework with proper h(-).
2.2 CONVERGENCE ANALYSIS
Equilibrium condition. To provide the theoretical guarantees, we show that the stationary distribu- tion our g .)) update is p.. Meanwhile, the evolution of KL-divergence is well-defined (without the explosion of functional gradient) and descending. We list the regularity conditions below.
[A1] (Regular function class) The discrepancy induced by function class F is positive-defintte: for any A p, there exists f E F such that E,IV f(x) + (xx) V In p(x)] ) > 0. For f E J and e E R, of E F, and F is closed. The tail of f is regular: lim f(0 x)p, p. (x) = 0 for f f F.
[A2] (L-Smoothness) For any I E R4, (x) > and p((,2) > 0 are L-smooth densities, i.e. |V p(())) p(y)l L| y/1. with Ep||=||2 xc.
Particularly, [A1] is similar to the positive definite kernel in RKHS, which guarantees the decrease of KL divergence. [A2] is designed to make the gradient well-defined: RHS Eq. (3 is finite.
(7). we have p(t, 2) # (2). i... gt, x) = O if and only p(t, x) = (x).
Proposition 2. Under [A1]. [A2]. when update X, Eq.
dDsL < for all
Proposition 2 shows that the continuous dynamics of our g(t,xx) :)) is well-defined and the KI. divergence along the dynamics is descending. The only stationary distribution for p(t, is p. (z).
Convergence rate. The convergence rate of our framework mainly depends (1) the capacity of function class F: (2) the complexity of p.. In this section, we analyze that when the approximation error is small and the target p. is Sobolov,,,, the KL divergence converges linearly.
[Aa] (e-approximation) For any 0, there exists fe(x) and e < 1, such that
f p(l..) [((e ( - ' x)
de S l ((l(.))) In p(t,
de
(x(
(x(x)
[A4] The target p. satisfies loooolllllle inequality (μ 0): For any differentiable function g, have Ep. [9² In g²] Ep. [9"] In EEp [9]]]
Ep. | 2
Specifically, [Ag] is the error control of gradient approximation. With universal approximation theorem (Horiik et al. 1989), any continuous function can be estimated by neural networks, which indicates the existence of such a function class. [A4] is a common assumption in sampling literature. It is more general than strongly concave assumption (Vempala & Wibisono, 2019).
4

---

Published as a conference paper at ICLR 2023
Theorem 1. Under [A1]-[A for any O, assume that the largest eigenvalue, Xmax m.
Then we have (Dxl(t) 1= dt - 2 pp
Pe(m) 1 P. H-1 and DKL(t) S exp(-(1 mku
Theorem I shows that when p. is log-Sobolev, with a function class powerful enough to approximate the preconditioned Wasserstein gradient. Moreover, considering the discrete algorithm, we have to make that the step size is proper. i.e., Imp. Σ C, for some constant c > 0. Thus, the preconditioned H = -cV2 In p. is have
algorithm achieves linear convergence rate
better than plain one H cLI, since (-V2 In p.) > L-11. For better understanding, provided a Gaussian case discretized proof in Appendix A. to illustrate this phenomenon.
PFG: PRECONDITIONED FUNCTIONAL GRADIENT FLOW
Algorithm PFG: Preconditioned Functional Gradient Flow
Input: Unnormalized target distribution p. (2) -((), fo(x) R + R initial particles (parameters)
{ 0o, iteration parameter T, T", step size vy, regularization function h(-).
for ........ T do
Assign 0% = 01-1;
for = 1, 1. T do
Compute ((0) 1 E.= (h(fo(xl)) + fo(x)) vU(x)) V fo(x!))
Update 0 = -1 L/
end
Assign 8, 8thh and update particles ot =t (fo. for all 1,...,
end
Return: Optimized particles {-T)==
3.1 ALGORITHM
We will realize our algorithm with parametric fe (such as neural networks) and discretize the update.
Parametric Function Class. We can let J = (ff(x) O} and apply g(t, x) fgg (zr). such that
f [D(.::)-V fe(x) (x V In 2 [ / (x) [dd
0, arg min
(8)
BEE
where H is symmetric positive definite matrix estimated at time t. Eq. 8 is a direct result from Eq. 4 and (6). The parametric function class allows So to be optimized by iterative algorithms.
Choice of H. Considering the posterior mean trajectory, it is equivalent to the conventional opti- mization, so that - p. ideal (Newton's method) to implement discrete algorithms. We use diagonal Fisher information estimators for efficient computation as Adagrad (Duchi al. (011). We approximate the preconditioner H for particles at each time step by moving averaging.
Discrete update. We can present our algorithm by discretizing Eq. 4 and 88. Given Xo Po. we update X₃ X X nfā, (Xx), where dx is obtained by Eq. 8 with (stochastie) gradient descent. The integral over p( .)) is estimated by particle samples. Full procedure presented in Alg. I. where the regularizer h is $ 1|2 by default.
3.2
COMPARISON WITH SVGD
3.2. SVGD FROM FUNCTIONAL GRADIENT VIEW
For simplicity, we prove the case under dimeesiimessooal feature map (x)) Rd Rh. We assume that F (Wv(x)) W [dxh], and let kernel k(æ,, y) = vxx) (1)) Q() 1
|[W]] where RKHS norm is the Frobenius norm of W. The solution is defined by
/ pp(t.z) trace W Vxb()) + W ( Inp. ddr + ||
W. arg min
(9)
WeR 4
manfnite-dimennsonne version is provided in Appendix A.7
5

---

Published as a conference paper at ICLR 2023
which gives We ] pt, x²) Vxb('') (lmp((''))(=') This implies that
g 2) Web(aa) = p(t. p(t. + Vlnp- (v))))
(10)
which is equivalent to SVGD. For linear function classes, such as RKHS, Q can be directly applied to the parameters, such as W here. The regularization SVGD is -11-11:- (the norm defined in RKHS).
For non-linear function classes, such as neural networks, the RKHS norm cannot be defined.
3.2.2 LIMITATIONS OF SVGD
Kernel function class. As in Section 3.2.1 map functions, which suffers from curse of dimensionality (Geenens) 2011). the other hand, some non-linear function classes, such as neural network, performs well on high dimensional data (LeCun) ct al. 2015). The extension to non-linear function classes is needed for better performance.
RKHS only contains linear combination of the feature
Gradient preconditioning in SVGD. In Example | when - V2 In is ill-conditioned, PFG al- gorithm follows the shortest path from (ig to jue. Although SVGD can implement preconditioning matrices (Wang et al. 2019), due to the curl component and time-dependent Jacobian of dus /tt, any symmetric matrix cannot provide optimal preconditioning (detailed derivation in Appendix).
Suboptimal convergence rate. For log-Sobolev p.. SVGD with commonly used bounded smoothing kernels (such as RBF kernel) cannot reach the linear convergence rate (Duncan et al., 2019) and the explicit KL convergence rate is unknown yet. Meanwhile, the Wasserstein gradient flow converges linearly. When the function class is sufficiently large, PFG converges provably faster than SVGD.
Computational cost. For SVGD, main computation particles, we need (n²) memory and computation. Our algorithm uses an iterative approximation to optimize gt, 2z), whose memory cost is independent of and computational cost ((n) (Bertsekas) et al. 2011) The repulsive force between particles is achieved by V operator on each particle.
comes from the kernel matrix: with
4 EXPERIMENT
To validate the effectiveness of our algorithm, we have conducted experiments on both synthetic and real datasets. Without special declarations, use parametric two-layer neural networks with Sigmoid activation as our function class. To approximate in real datasets, we use the approximated diagonal Hessian matrix H, and choose H = f1", where a E 0.1, 0.2, 0.5, 1 the inner loop T' of PFG is chosen from {1, 2, 5, 10), the hidden layer size chosen from f32, 64, 128, 256, 512]. The parameters are chosen by validation. More detailed settings are provided in the Appendix.
Gaussian Mixture. To demonstrate the capacity of non-linear function class, we have conducted the Gaus- sian mixture experiments to show the advantage over linear function class (RBF kernel) with SVGD. We con- sider to sample from a 10-cluster Gaussian Mixture distribution. Both SVGD and our algorithm are trained with 1,000 particles. Fig. shows that the estimated "score" by RBF kernel usually unsatisfactory: (1) In low- density area, it suffers from gradient vanishing, which makes samples stuck at these parts (similar to Fig. [ (b)); (2) The score function cannot distin- guish connected clusters. Specifically, some clusters are isolated while oth- might be connected. The choice of bandwidth is hard. The fixed band- width makes the SVGD algorithm
(a) True Density
(b) SVCD Deniiyy
(d) PFG Density
True Samples
SVGD Sumples
PTG Sumples
Stein Score
SVGD Scure
( PEG Score
Figure 3: Particle-based VI for Gaussian mixture sampling.

---

Published as a conference paper at ICLR 2023
SVGD (Linear)
SVGD (RBF)
Ours (H =
Ours (H I:-)
Innture
Islsoo 30
divergence
@ = [],
(b) p- = [L,
(0) [20,20]
( [20, 20] d1ag(100, 1)
2. 1)ag(1,1)
(1000 1)
2. diag(1,1)
Figure 4: Evolution of particle distribution from (0, I) to Ni,, Σ,) (first row: mean squared emror of p4: l1ie 14a 1=, second row: KI. divergence between p(t, x)) and p.()))
able to determine the gradient near connected clusters. It fails to capture the clusters near (-0.75, 0.5). However, the non-linear function class is able to resolve the above difficulties. PFG, we found that the score function is well-estimated and the resulting samples mimic the true density properly.
Illconnditioned Gaussian distribution. We show the effectiveness of our proposed regularizer. For ill-conditioned Gaussian distribution, the condition number of Σ, is large, i.e., the ratio between maximal and minimal eigenvalue is large. We follow Example [ and compare different μ4, and E. When Σ. is well-conditioned (E, I). L2 regularizer (equivalent to Wasserstein gradient) performs well. However, it will be slowed down significantly with ill-conditioned Σ.. For SVGD with linear kernel, the convergence slows down with shifted μs, ill-conditioned Σ. For SVGD with RBF kernel, the convergence is slow due to the improper function class. Interestingly. for ill-conditioned case, 14 of SVGD (linear) converges faster than method with H I but KL. divergence does not always follow the trend. The reason is that N, of SVGD is highly biased, making KL divergence large. Our algorithm provides feasible Wasserstein gradient estimators and makes the particle-based sampling algorithm compatible with ill-condtiooned sampling case.
Logistic Regression. We conduct Bayesian lo- gistic regression for binary classification task on Sonar and Australian dataset (Dua Graff) 2017). In particular, the prior of weight samples are assigned (0, 1); the step-size is fixed as 0.01. We compared our algorithm with SVGD (200 particles), full batch gradient is used in this part. To measure the quality of posterior approximation, we use 3 metrics: (1) the dis- tance between sample mean 14 and posterior mean pa; (2) the Maximum Mean Discrep- ancy (MMD) distance (Gretton al.] 2012)
SVGD
Ours
100
IAc- 11-2
-22
log MMD
-22
-24
from particle samples to posterior samples; (3) the evidence lower bound (ELBO) of cur- rent particles, Ezzp, [log (z) log p)t, where pa (xx) the unnormalized density p(x, D) with training data D, the entropy term log ((,, estimated by kernel density. The ground truth samples of p+ is estimated by NUTS (Hoffman et al. 2014). Fig. 5 shows that method outperforms SVGD consistently.
Negative ELBO
Iterations (K)
Iterations (K)
Hierarchical Logistic Regression. For hier- archical logistic regression, consider 2- level hierarchical regression, where prior of weight samples is (0, -1). The prior of is
Figure 5: Posterior sampling for Bayesian logistic regression. (200 particles; dataset: sonar (first col- umn) and Australian (second column). f4r particle mean; μ.: posterior mean.)
7

---

Published as a conference paper at ICLR 2023
Gamma(0. 0.01). All step-sizes chosen from {10 ;t 2.... by validation. We compare the test accuracy and negative log- log-likelihood as Liu & Wang (2016). Mini-bated was used to estimate the log-likelihood of training samples (batch size is 200). Tab. shows that our algorithm performs the best against SVGD and SGLD on hierarchical logistic regression.
Bayesian Neural Networks. We compare algo- rithm with SGLD and SVGD variants on Bayesian neural networks. We use two-layer networks with 50 hidden units (100 for Year dataset, 128 for MNIST) and ReLU activation function; the batch- size is 100 (1000 for Year dataset), All step-sizes are chosen from (10-* = 1,2,... by valida- tion. For UCT datasets, we choose the precondi- tioned version for comparison: M-SVGD (Wang) et al., 2019) and PSGLD (Li et al., 2016). Data
Table I: Hierarchical Logistic Regression on UCI datasets (particle size 200)
COVTYPE
GERMAN
ACC.
NLL
ACC.
NLL
SVGD SGLD
75.1
0.58
65.1
0.71
74.9
0.57
64.8
0.65
PFG
75.7
0.51
67.6
0.64
samples are randomly partitioned to two parts: 90% (training), 10% (testing) and 100 particles is
used. For MNIST, data split follows the default setting and compare algorithm with SVGD for 5, 10, 50, 100 particles (SVGD with 100 particles exceeded the memory limit). In Tab. 22] PFG outperforms SVGD and SGLD significantly. In Fig- 6 we found that the accuracy and NLL of PFG is better than SVGD with all particle size. With more particle samples, PFG algorithm also improves.
Table 2: Averaged test root-mean -quare enror (RMSE) and test log-likelihood of Bayesian Neural Networks on UCI datasets (100 particles). Results are computed by 10 trials.
AVERAGE TEST RMSE
AvERAGE TEST LOG-LIKELIHOOD M-SVGD
DATASET
M-SVGD
PSGLD
PFG
PSGLD
PFG
BOSTON
2.72
2.70
2.47=0.11 =0.11
-2.66
-2.85
-2.35 20.12
CONCRETE
g..1
40.14
4.69
20.14
+0.06
-3.21
-2.83
ENERGY
0.89
0.99
0.48 20.04
-1.42+0.03
-1.31
-1.22
PROTEIN
4.55
4.
=.51
-3.07
-3.2210.11
-2.89
WINERED
0.63
0.6410.02
20.02
-1.77
-1.8010.05
-1.61 20.03
WINEWHITE
0.65
00.07
0.59 20.02
-1.75
-1.8210.07
-1.58
YEAR
8.62=0.09
8.66
8.56
-3.59
-3.5640.04
-3.51
particles
10 particles
50 particles
100
Accuracy
95-
20
90-
10
10
20
0.0
2.5
5.0
75
30.5
0.5
0.5
3.4
5.2
0.2
10
10
0.0
2.5
5.0
30.0
BoraBona (R)
Rarudona
Ticrut una
SVED
Ours
ppation
Savittes particas
(x) SVGD vs PFG on MNIST classifcaaiio
(b) PFG with differeni purticles
Figure 6: Test accuracy and NLL of Bayesian Neural Network (MNIST classification)
Time Comparison. As indicated, our algorithm is more scalable than SVGD in terms of particle numbers. Tab. 3 shows that SVGD entails more computation cost with the increase of particle numbers. Our functional gradient is obtained with iterative approximation without kemel matrix computation. When n is large, our (mn) algorithm is much faster than O(n²) kernel-based method. Interestingly, the introduction of particle interactions is a key merit of particle-basee VI algorithms, which intuitively needs On²?) computation. The integration by parts technique draws the connection between the particle interaction and V fo which supports more efficient computational realizations.
8

---

Published as a conference paper at ICLR 2023
Table 3: Time Comparison on Bayesian logistic regression (sonar dataset, 1,000 iterations)
# OF PARTICLES
5
10
100
1000
2000
SVGD
5.20 10.10
5.25
10.13
5.71 10.11
19.10
81 81.50 11.8
PFG
7.87 10.16
7.95
10.15
8.151
11.56
13.22 10.26
RELATED WORKS
Stein Discrepancy and SVGD. Stein discrepancy (Liu et al. 2016 Chwialkowski al. 2016 Liu & Wang 2016) is known as the foundation of the SVGD algorithm. which is defined on the bounded function classes, such as functions with bounded RKHS norm. We provide another view to derive SVGD: SVGD is a functional gradient method with RKHS norm regularizer. From this view, we found that there are other variants of particle-bast based VI. We suggest that the our proposed framework has many advantages over SVGD the function class is more general and the kernel computation can be saved. We use neural networks as our F in real practice, which much more efficient than SVGD, Moreover, the general regularizer in Eq. (6) improves the conditioning of functional gradient, which is proven better than SVGD in both theory and experiments. The preconditioning of our algorithm is more explainable than SVGD, without incorporating
properties.
Learning Discrepancy with Neural Networks. Some recent works (Hu et al. 2018) Grathwohl et al. 2020; di Langosco et al. 2021) leverage neural networks to learn the Stein discrepancy. which are related to our algorithm. They only implement (z) = Ep, [=) to estimate the discrepancy. Grathwohl et al. (2020) measure the similarity between distributions and train energy-basee models (EBM) with the discrepancy. However, they did not further validate the sampling performance of their "learned discrepancy", so is new version of KSD (Liu et al., 2016). rather than SVGD. Hu et al. (2018) train neural "generator" for sampling tasks. They discussed the scaling of Le regularization to align with the step size. In contrast, we restrict our investigation particle-based VI and avoid learning a "generator", which may introduce more parameterization errors. di Langosco al. (202T) an empirical study to implement the I.2 version of our algorithm with comparable performance to SVGD. We extend the design of Q to general case, and emphasize the benefits of our proposed regularizers corresponding to the preconditioned algorithm. We include the Example and Fig. 4 to demonstrate the necessity of preconditioning. Our theoretical and experimental results have demonstrated the improvement of PFG algorithm against SVGD.
KL Wasserstein Gradient Flow. Wasserstein gradient flow (Ambrosio al. 2005 is the continuous dynamics to minimize functionals in Wasserstein space, which is crucial in sampling and optimal transport tasks. However, the numerical computation of Wasserstein gradient is non-trivial. Previous works (Peyré) 2015 Benamou et al. 2016; Carlier et al. (2017) attempt to find tractable formulation with spatial discretization, which suffers from curse of dimensionality. More recently, (Mokrov et al.) 202T; Alvarez-Melis et al., 202T) leverage neural networks to model the Wasserstein gradient and aim to find the full transport path between Po and p.. The computation of full path is extremely large. Salim et al. (2020) defines proximal gradient in Wasserstein space by JKO operator. However, the work mainly focus on theoretical properties and the efficient implementation remains open. Wang et al. (2022) solves SDP to approximate Wasserstein gradient, which considers the dual form of the variational problem. When the functional is KL divergence, Wasserstein gradient can also be realized with Langevin dynamics (Welling & Teh 2011 Bernton, 2018) by adding Gaussian noise to particles, which are also variants of MCMC. However, the deterministic algorithm to approximate Wasserstein gradient flow is still challenging. Our framework provides tractable approximation to fix this gap.
6 CONCLUSION
In particle-base based VI, consider that the particles be updated with preconditioned functional gradient flow. Our theoretical results and Gaussian example led to an algorithm: PFG that can approximate preconditioned Wasserstein gradient directly. Our theory indicates that when the function class is large, PFG is provably faster than conventional particle-based VI. SVGD. The empirical result showed that PFG performs better than conventional SVGD and SGLD variants.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
The work was supported by the General Research Fund (GRF 16310222 and GRF 16201320).
REFERENCES
David Alvarez-Melis, Yair Schiff, and Youssef Mrouch. Optimizing functionals on the space of probabilities with input convex neural networks. arXiv preprint 0rXiv:2106.00774 2021.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005.
Jimmy Ba, Murat A Erdogdu, Marzyeh Ghassemi, Shengyang Sun, Taiji Suzuki, Denny Wu, and Tianzong Zhang. Understanding the variance collapse of svgd in high dimensions. In International Conference on Learning Representations, 2021.
Jean-David Benamou, Guillaume Carlier, Quentin Mérigot, and Edouard Oudet. Discretization of functionals involving the monge-annoaa operator. Numerische mathematik, 134(3):611 666, 2016.
Espen Bernton. Langevin monte carlo and jko splitting. In Conference On Learning Theory, Pp. 1777-1798. PMLR, 2018.
Dimitri P Bertsekas et al. Incremental gradient, subgradient, and proximal methods for convex optimization: A survey. Optimization for Machine Learning, 2010(1-38):3, 2011.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the Americun statistical Association, 112(518):859-57) 2017.
Guillaume Carlier, Vincent Duval, Gabriel Peyré, and Bernhard Schmitzer. Convergence of entropic schemes for optimal transport and gradient flows. SIAM Journal on Mathematical Analysis, 49(2): 1385-1418 2017.
Peng Chen, Keyi Wu, Joshua Chen, Tom O'Leary-Rosebest and Omar Ghattas. Projected stein variational newton: A fast and scalable bayesian inference method in high dimensions. Advances in Neural Information Processing Systems, 32, 2019.
Tianqi Chen, Emily Fox, and Carlos Guestrin. Stochastic gradient hamiltonian monte carlo. In International conference on machine learning, pp. 1683-1691. PMLR, 2014.
Sinho Chewi, Thibaut Le Gouic, Chen Lu, Tyler Maunu, and Philippe Rigollet. Svgd as a kernel- ized wasserstein gradient flow of the chi-squared divergence. Advances in Neural Information Processing Systems, 15:2098-2109 2020.
Kacper Chwialkowski, Heiko Strathmann, and Arthur Gretton. kernel test of goodness of fit. In International conference on machine learning, pp. 2606-2613 PMLR, 2016.
Gianluca Detommaso, Tiangang Cui, Youssef Marzouk, Alessio Spantini, and Robert Scheichl. A stein variational newton method. Advances in Neural Information Processing Systems, 31, 2018.
Lauro Langosco di Langosco, Vincent Fortuin, and Heiko Strathmann. Neural variational gradient descent. arXiv preprint arXiv:2 10731, 2021.
Dheeru Dua and Casey Graff. UCI machine learning repository, 2017. URL httt::///ccchhi [cs. www. www.uuiiieuum...
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Andrew Duncan, Nikolas Nüsken, and Lukasz Szpruch. On the geometry of stein variational gradient descent. arXiv preprint arXiv: 1912.00894, 2019.
Jiaojiao Fan, Amirhossein Taghvaci, and Yongxin Chen. Variational wasserstein gradient flow. arXiv preprint arXiv:2 2021.
10

---

Published as a conference paper at ICLR 2023
Yihao Feng, Dilin Wang. and Qiang Liu. Learning draw samples with amortized stein variational gradient descent. arXiv preprint arXiv:1707.06622 2017.
Jerome H Friedman. Greedy function approximation: a gradient boosting machine. Annals of statistics, pp.
2001.
Alfredo Garbuno- Inigo. Franca Hoffmann, Wuchen Li, and Andrew Stuart. Interacting langevin diffusions: Gradient structure and ensemble kalman sampler. SIAM Journal Applied Dynamical Systems, (9(1):412-444) 2020.
Gery Geenens. Curse of dimensionality and related issues in nonparametrie funcciionl regression. Statistics Surveys, 5:30-43, 2011.
Wenbo Gong, Yingzhen Li, and José Miguel HemándezzLobuto. Sliced kernelized stein discrepancy. arXiv preprint arXhv:2006. 2020.
Will Grathwohl, Kuan-Chieh Wang, Henn-eeriik Jacobsen, David Duvenaud, and Richard Zemel. Learning the stein discrepancy for training and evaluating energy-based models without sampling. In International Conference on Machine Learning, pp. 3732-3747. PMLR, 2020.
Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-773, 2012.
Matthew D Hoffman, Andrew Gelman, al. The no-u-turn sampler: adaptively setting path lengths in hamiltonian monte carlo. J. Mach. Learn. Res., 15(1):1593-1623, 2014.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks universal approximators. Neural networks, 255):35 1989.
Tianyang Hu, Zixiang Chen, Hanxi Sun, Jincheng Bai, Mao Ye, and Guang Cheng. Stein neural sampler. arXiv preprint arXiv:1810.03355 2018
Michael F Hutchinson. stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Statistics -iimuliis-ii and Computation, 1833):1059-106, 1989.
Rie Johnson and Tong Zhang. A framework of composite functional gradient methods for generative adversarial models. IEEE transactions on pattern analysis and machine intelligence, 43(1):17-32, 2019.
Michael Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. An introduction to variational methods for graphical models. Machine learning, 37(2):183-233 1999.
Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the fokker-planck equation. SIAM journal mathematical analysis, 29(1):1-17, 1998.
Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. nature, 521((553):436, 2015.
Chunyuan Li, Changyou Chen, David Carlson, and Lawrence Carin. Preconditioned stochastic gradient langevin dynamics for deep neural networks. In Thirtieth AAAI Conference Artificial Intelligence, 2016.
Wuchen Li and Lexing Ying. Hessian transport gradient flows Research in the Mathematical Sciences, 104) 2019.
Wuchen Li, Alex Tong Lin, and Guido Montáfar. Affine natural proximal learning. In International Conference Geometric Science of Information, pp. 705-714. Springer, 2019.
Alex Tong Lin, Wuchen Li. Stanley Osher, and Guido Montúfar. Wasserstein proximal of gans. In International Conference Geometric Science of Information, pp. 524-533. Springer, 2021.
Chang Liu, Jingwei Zhuo, Pengyu Cheng, Ruiyi Zhang, and Jun Zhu. Understanding and accelerating particle-based variational inference. In International Conference on Machine Learning, pp. 4082- 4092. PMLR, 2019.
11

---

Published as a conference paper at ICLR 2023
Qiang Liu. Stein variational gradient descent gradient flow. Advances in neural information processing systems, 30, 2017.
Qiang Liu and Dilin Wang. Stein variational gradient descent: general purpose bayesian inference algorithm. Advances in neural information processing systems, 29, 2016.
Qiang Liu, Jason Lee, and Michael Jordan. A kernelized stein discrepancy for goodness-of tests. In International conference machine learning, pp. 276 284. PMLR, 2016.
Xing Liu, Harrison Zhu, Jean-Francois Ton, George Wynne, and Andrew Duncan. Grassmann stein variational gradient descent. arXiv preprint arXiv:2202.03297, 2022.
Petr Mokrov, Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, and Evgeny Burnaev. Large-scale wasserstein gradient flows. Advances in Neural Information Processing Systems, 34, 2021.
Atsushi Nitanda and Taiji Suzuki. Gradient layer: Enhancing the convergence of adversarial training for generative models. In International Conference on Artificial Intelligence and Statistics, pp- 1008-1016 PMLR, 2018.
Gabriel Peyré. Entropic approximation of wasserstein gradient flows SIAM Journal on Imaging Sciences,
2015.
Yuchen Pu, Zhe Gan, Ricardo Henao, Chunyuan Li, Shaobo Han, and Lawrence Carin. Vae learning via stein variational gradient descent. Advances in Neural Information Processing Systems, 30, 2017.
Rajesh Ranganath, Sean Gerrish, and David Blei, Black box variational inference. In Artificial intelligence and statistics, pp. 814 822. PMLR, 2014.
Adil Salim, Anna Korba, and Giulia Luise. The wasserstein proximal gradient algorithm. Advances in Neural Information Processing Systems, 33: 1-356-13666, 2020.
Santosh Vempala and Andre Wibisono. Rapid convergence of the unadjusted langevin algorithm:
Isoperimetry suffices. Advances in neural information processing systems, 32, 2019.
Cédric Villani. Optimal transport: old and new, volume 338. Springer, 2009.
Dilin Wang and Qiang Liu. Learning to draw samples: With application to amortized mle for generative adversarial learning. arXiv preprint arXiv: 161 1.01722, 2016.
Dilin Wang. Ziyang Tang, Chandrajit Bajaj, and Qiang Liu. Stein variational gradient descent with matrix-valued kernels. Advances in neural information processing systems, 32, 2019.
Yifei Wang and Wuchen Li. Information newton's flow: second-order optimization method in probability space. arXiv preprint arXiv:2001.043:2 2020.
Yifei Wang and Wuchen Li. Accelerated information gradient flow. Journal of Scientific Computing. 90(1))1-47, 2022.
Yifei Wang, Peng Chen, and Wuchen Li. Projected wasserstein gradient descent for high-dimensional bayesian inference. arXiv preprint arXiv:2102.06350, 2021.
Yifei Wang. Peng Chen, Mert Pilanci, and Wuchen Li. Optimal neural network approximation of wasserstein gradient direction via convex optimization. arXiv preprint arXiv:2205. 13098, 2022.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference machine learning (ICML-11), pp. 681-688. Citeseer, 2011.
Andre Wibisono. Sampling as optimization in the space of measures: The langevin dynamics as composite optimization problem. In Conference on Learning Theory, pp. 2093-3027. PMLR, 2018.
Jingwei Zhuo, Chang Liu, Jiaxin Shi, Jun Zhu, Ning Chen, and Bo Zhang, Message passing stein variational gradient descent. In International Conference on Machine Learning, pp. 6018-6027. PMLR, 2018
12

---

Published as a conference paper at ICLR 2023
A PROOFS
A.1 PROOF OF EQ.
From the continuity equation (or Fokker-Planchs equation), have
Öp(t,z)
- [g(t, x x)p(t, x)] Öt
By chain rule and integration by parts, dD
Op(t,x)
1 + In p(t,z)
dx
dt
Ot
Do (xx)
1 1+th In p(t,z)
V lo(t, x)tt,)
dr
px(x)
f (oll,))p11 x) Vn p(t,z) du
(x)
/ (o(l,))p(( v (l),((,)) - In p. (x)]dd
/ 2(1,0) T [Vp(t,>) p(t V p= + V = / p(t,x)) (lt,t))))))) lnp.
A.2 DISCUSSION OF EXAMPLE
(1) SVGD with linear kerne h(2,, y) = x Ky +1
g(t, x) = Er [[I[ + (Ur - u. u.),,, Kr + 14 - [ + Kx.
(11)
(2) SVGD with RBF kernel kig (x. = +/ exp (- //2 lly - x=1²)
g(l, x) (xxeex((1111 exp(-1|x133)
(12)
(3) Linear function class with L2 regularization QU((=))
-1). (4) Linear function class with Mahalanobis regulariztinnnl Q(f(=)) = Y/2 E (+)
g(1,,) = -x- 1
(13)
g(t,x) = (2 - pb)) + Σ 14).
(14)
For optimal transport g(t,z) = -x nu.) Σ, 1/2 (Σ) / 1220, 1/2 1/2(x ma).
Note that falaa is not symmetric, which makes the distribution rotate. Figure 7 shows that we can
split the curl component by Helmholtz decomposition, which means that this part would rotate the distribution and cannot be compromised by preconditioning method. Thus, we cannot find any to obtain a proper preconditioning.
Proof. We consider these 4 cases seperately.
(1) SVGD with linear kernel.
k(x y) = x T Ky + 1, then V. vyk(I,)) v) = Kx,
g(t,x)) = / p(t, (-)) I (.)) + (y-m))) Ka + 1)dy
/ p(t.v)(2. (y p2) Kx + d))
+ / ((t Tyu - uM ddyK
M5 (2' + (pr - pm) ]K + ue - [m] + Kx
13

---

Published as a conference paper at ICLR 2023
43
40-
40
20
20
-20
20
40
-00
-20
(a) Vector field of (a, x
(b) Curl component
(c) Divergence component
Figure 7: Illustration of p+14g term by Helmholtz decomposition. (#t = [0, 10), f4- 20))
(2) SVGD with RBF kernel.
Voz exp (-2 1 Ilw - xx?
xl1,x)
exp ( (y mi) 1)) TE, (y - mi)
v 2 Ex|
(2, (y) ml.)) - 21 (y - ue)) dy
/ v 2T Ed exp (-20) """ Σ 8) Voz exp (-
20 110 + 1e - xx1²)
20
(Σ. (0 1tt 11,) - 1) dj
2T ((o 221 21 + 1--
1
2 21 + 1ー1
1 ( - pn) (a 2 + tj (ū ps) (T
exp
11 (2(me [/1 - x) " (Mr - x)) (2: (ū + hi - U4) - 2x v) dÿ
exp
1
(10 exp 7(me= - x) TTT (1x - x))
02 21 1; )
E3 14 1b)) E dÿ
=C, (5E) E-1) Us + [ 11,))
where
Pz o (o 21 + Σ, 1) -1 (x mi) = (I + -2 11 (x (x),),,
(0-21
1
1
C = VI02I+ Egl
exp
202 (x - me) T (I (1+ 23 ) (x mi)
Note that gt,, = Ox exp((_1<<==)) When + co, ((,,x)) is vanishing.
(3) and (4) are special case of Proposition 2
Considering the evolution of μ4, we have
(1) dug - [[E+ + (Mt n.)NT KH4 + M4 u.] + K KPt, which contains non-symmetric p.
di
(2) The transport is not linear.
(3) dui
E,4( - pc).-
di
= -oa -1,, which means that the evolution of (4) is equivalent with optimal transport. (4) dr
dun
14

---

Published as a conference paper at ICLR 2023
A.3 PROOF OF PROPOSITION 1
Proof. Assume that Q() - [((+)))))]]] when we have Eq.
4
g(t,x) = arg min JEF
plt, f(x) (1)p+())
2 [l(())llaldx
(15)
VI (x
arg min fEF
p((,))1-1(x)) Vh
(16)
p(l,x))
H H-1V1 (x)
Vln (x(x)
arg min fer
Epr 1
[(x)
p(t,,) + 2 (x))
(17)
p(t, x H
f(x) H V (x) 2
arg min JEF
Ep.
(18)
p(t,x)
H
A.4 PROOF OF PROPOSITION 2
Lemma 1. (A Variant of Lemma. 11 in (Vempala & Wibisono 2019)) Suppose that p(æ) > is L-smooth. Then, we have
J [Vp((\] de s Ld;
Proof. By [A1], there exists f such that
/ p((,=)]V f+f V. (np. < O.
(19)
Thus, consider fe(x)= = cf(x),
L(c) - / plt,x)(V of +ff V. VV (np.(t)) =: -Cye
(20)
((c) / plt. & p(t.z))lflidde CHe
(21)
where CH, Cf > 0.
When 2C'
C₁
we have
C
- Cjc= <.
4CH
By choosing co 2CH
C₁
/ p(t,x)/V ((tx))l-g(ttz)) x(t,z)++++)))))) + x(t, (I)..(x) (V)p))))) + 1 5llo(t, dd
(22)
3 / v(t,z)/V-col of oo (m) VVpp() + 2 |cofl ldr
(23)
<0
(24)
Thus,
dDkL(t)
( - / p(t,z)(V g(t,2) + g(t,z) Vn < 0
dt
15

---

Published as a conference paper at ICLR 2023
By Eq. 3 (t)
p((,))(V g(l,z) x(1,,) V In (x)]dr
(25)
dt
/ plt,x)lo(txxxVV) In P. (x)
p(t, x² dd
(26)
2 2 / p(t,z) 2)) Vn (x)
1
/ p(t, Ilg(t, x)1 dr
dr
(27)
p(t, x)
2 -(1 Amax (H) (H) / p((,2))
V Pa (x) |2 dx
(28)
pil, x))
2 -11 Amax ( Idd + f p(l,,x)
V (0) Tarl dr
(29)
-00
(30)
where ʃ x()
Vnn p=(x)) p(t,a))
-|
dr can be controlled by Ld ([A2) and Lemma and Ep. ||=/||?
A.5 PROOF OF THEOREM 1
Proof. By p-log-Soboley inequality in [A4], suppose g² Pt/ Pc/Pp. P..
then we have
(t - 2p / p(t,x)
1
VIn In p(t. 2)
dz.
(31)
(x)
When consider the H -inducedd norm, for any x, have || H-1
Amaz (HT) ||0||. Thus,
(t) S Xmaa((H)
/
V In x(t,x) 2
p(t,z)
dx.
(32)
2p
pe (x() H-1
Using Cauchy-Schwand inequality and [As), we have
dDki.
/ p(t, x)) In
(x)
g(t, z)dx
dt
p(t, x) (x(x)
/ p(c, x(t, In p(l, x:)
gt,,x) - -H In (x(x)
(x())
+ H 11 V In
dx
p(l,xx)
p(l, xx)
f p(l, In p(t, x
2
dr
p. (x)
11-1
+ / p(t, x)
(x)
In px())
In p(t, x)
g(t, - -H-1Vnn
dx
p(t, x))
S f p((,
In p(t,ar) x))
dx
Ps (x:
H-1
+ / p(l, xD)
Vin (x)
Pe (x)
g(1,z)) H- VV
dx
p(t, xr)
plt, x:) 17
H1
1 / 2(),
In p(t, x))
1
dr 2 ʃ p(t,zz) gg(t,z) x) HH VVnn In (x(
2
- 2
dx
(x(x)
p(t,z)
H-1
(A1 (t)
Amax
Thus, by Grönwall's 's inequality,
(-p(1-c) (H) DKL (0)
(1)(t) S exp
16

---

Published as a conference paper at ICLR 2023
Remark: Theorem I shows that PFG algorithm convergence with linear rate. Moreover, the choice of preconditioning matrix affect the convergence rate with the largest eigenvalue. In real practice, due to the discretized algorithm, one needs normalize the scaling of H, which is equivalent the learning rate in practice. For any preconditioning H, we need to compute the corresponding learning rate get H 17" H. One usually set that !| V2 In (z))ll I 1 to make the discretized (first-order) algorithm feasible (1) S || v in p. (z))l2). otherwise the second order error would be too large. The construction is similar to condition number dependency in optimization literature.
Take the fixed largest learning rate for example in the continuous case:
When M T, we have that y) 1/T,, H T.T,
(t) dt - Ep
1
V In (x
2L
P. (x)
When = In p.. we have that 17 1, H
V2 12 p..
(1)
In (x))
Pe (a) p. (x)
dl S - --p..
-
Ep.
Vh
2
p. (x) Inp.)
2L
A.6 ANALYSIS OF DISCRETE GAUSSIAN CASE
As shown Example I, given that Tu ~ ~N I) and the target distribution N(ux, Ex), we have the following update for our PFG algorithm,
T++1 = nh 1 (z, [p= xz) + +E (xt - -me))
By taking the expectation, we have
Mt+1 #t 1/ (Σ, (He- Me))
Σ+11 - E. =1= (2, - -E. Ei + nلاi (2, )T H
+ +1 11 - - - 21 - H11
Since Eg I, the eigenvectors of E, will not change during the update. For notation simplicity, we write the diagonal case in the following discussion.
Assume that M, = diag(o ,.). M that a; > 0i+1 for any otherwise we can still rescale the initialization to obtain the similar constraint.
= (i) .......),.))........))))). H diag(h), ha). such 1, d 1. Without loss of generality, we assume that or 1,
Given the fact that In(1 +x) < Eea when 0, we have,
2
(I) S
[ - pm 12 1.
0
Thus, define
Co = 2
1 / |1 Ha
+
1
Σ
1
oz
(33)
2
o
2
i=1
Considering the i-th dimension (with notation
[-1,),
juttak (110))
17 hoo?
[eels +
1
[poks 1poks
(34)
hoo?
1((+1)==8 = =1)(1== + 21)0(11)
1
x1)
h;
s;())
o7
+
h
(s,(t)2
o
(35)
17

---

Published as a conference paper at ICLR 2023
To guarantee the convergence of the algorithm, need to make sure that both br and Σg converge. By rearranging 34 and (35). we have
|uttle [pal
-1
[pols
[pals
04
(t 2(t+1)-- of (1
2m (t) o;
of
st - oo 0f
+
1 st)
One needs to construct contraction sequences of both mean and variance gap, which are guided by 1 1-n/ n n/;oo and 1 1 22/ (h;oş). Notice that (0) 2 of. Assume that r S min, hço /2. Then
have sf ( 2 o²-
Thus, we obtain
1
1
h4
o
(t)
is natural to obtain the contraction both mean and variance gap,
| 11all -all S mmax (1
11 lea 1 1
hoo?
02(11 + 1) oo
s2(1)= of
o7 S
h:
of
21
(t)(t) < max
1
Co.
hoo?
where Co is defined in Eq. 33
When h = 1, we have o2
(KL(t)= (1 (1-
od
Co-
20
When h; = =0 we have
1 DKL(t) < Co.
40
This result is equivalent to the Remark in A.5
A.7 PROOF OF INFINITE-DIMENNIONALL CASE OF RKHS
Proof. Assume the feature map, ((rr): Rd H and let kernel k(,, v) (V(x), Q( |llf One can perform spectral decomposition obtain that
k(x v)
where vi : Rdd + R are orthonormal basis and A; is the corresponding eigenvalue. For any E Hd, we have the decomposition,
8
g(x) = Σ giv (uli()))
where gi & and 2 ||::|
18

---

Published as a conference paper at ICLR 2023
The solution is defined by
/1((.x))
(() + =11011k4s
g = arg min
(36)
SEM"
20
20
26
arg min SEM²
/p(t.2)
(.E E
9. Av)(()+++++
Σ (mp. )
Σ
l1g I?
(37)
which gives
9: = VA. / ( +
This implies that
2
xt) Σ
Vigezi(()) ((0. + (mp.()))(
(38)
1=1
which is equivalent SVGD.
19

---

Published as a conference paper at ICLR 2023
B MORE DISCUSSIONS
B. SCORE MATCHING
Score matching (SM) is related to our algorithm due to the integration by parts technique, which uses parameterized models estimate the stein score V In p.. We made some modifications to the original techniques in SM: (1) We extend the score matching V In to Wasserstein gradient approximation V p. V In pri (2) We introduce the geometry-awart regularization approximate the preconditioned gradient. Thus, our proposed modifications make it suitable for sampling tasks.
In a word, both SM and PFG are derived from the integration by parts technique. SM is a very promising framework in generative modeling and our PFG is more suitable for sampling tasks.
B.2 OTHER PRECONDITIONED SAMPLING ALGORITHMS
Preconditioning is popular topic in SVGD literature. Stein variational Newton method (SVN) (Detommaso et al., 2018) approximates the Hessian to accelerate SVGD. However, due to the gap between SVGD and Wasserstein gradient flow, the theoretical interpretation is not clear. Matrix SVGD Waan al. 2019 leverages more general matrix- matrix-valued kernels in SVGD, which includes SVN as variant and is better than SVN. (Chen et al. (219 perform parallel update of the parameter samples projected into low-dimensional subspace by an SVN method. It implements dimension reduction to SVN algorithm. Liu et al. (2022) projects SVGD onto arbitrary dimensional subspaces with Grassmann Stein kernel discrepancy, which implement Riemannian optimization technique to find the proper projection space.
Note that all of these algorithms are based on RKHS norm regularization. As we mentioned in Example due to the introduction of RKHS, the preconditioning matrix is often hard to find. For example, the preconditioning matrix of the linear kernel varies in time, while our algorithm only needs constant matrix. For RBF kernel, the previous works have demonstrated the improvement of the Hessian inverse matrix, but the design is still heuristic due to the gap between the Wasserstein gradient and SVGD. Our algorithm directly approximates the preconditioned Wasserstein gradient, and further analysis is clear and conclusive. Also, the SVGD-basce algorithms suffer from the drawbacks of RKHS norm regularization, as discussed in Example !
Besides, there are some other interesting works (Li et al. 2019; Garbuno-Inig et al. 2020: Wang al. 2021 Lin al. 202T Wang Li 2020 2022 & Ying 2019) also take the preconditioning
methods / local geometry subspace properties into consideration, and they have also shown the improvement over the plain versions, which have similar motivation to our algorithm. Although these methods are out of the particle-based VI scope, we believe that these methods have great potentials in our literature, which can be interesting future works.
B.3 OTHER WASSERSTEIN GRADIENT FLOW ALGORIHTM
There is another line of work to approximate Wasserstein gradient flow through JKO-discretization (Mokroy et al. 2021 Alvarez-Melis et al. 2021 Fan et al., 202T). The continuous versions of these Wasserstein gradient flow algorithms are related our algorithm when the functional is chosen as KI. divergence. These are promising alternative algorithms to PFG in transport tasks. From the theoretical view, they solves the JKO operator with special neural networks (ICNN) and aims to estimate the Wasserstein gradient of general functionals, including KL/JS divergence, etc. On the other hand, our algorithm is motivated by the continuous case of KL Wasserstein gradient flow and only consider the Euler discretization (rather than JKO), which is the setting of particle-based variational inference.
When we consider the KL divergence and posterior sampling task, our proposed method is more efficient than Wasserstein gradient flow based ones, due to the flexibility the gradient estimation function. We provide the empirical results on Bayesian logistic regression below (Covtype dataset).
20

---

Published as a conference paper at ICLR 2023
Table 4: Bayesian logistic regression (Covtype dataset)
METHOD
ACCURACY (%)
NLL
OURS
75.7
0.58
JKO-ICNN JKO- -.O-VARIATIONAT.
75.0
0.62
75.3
0.59
According to the table, our algorithm outperforms JKO-type Wasserstein gradient methods. Moreover, it is important to mention that JKO-type Wasserstein gradient is extremely conputation-exhaustee due to the computation of JKO operator. Thus, it would be more challenging apply them to larger tasks such as Bayesian neural networks.
B.4 VARIANCE COLLAPSE AND HIGH-DIMENSIONAL PERFORMANCE
Some recent works [Ba ct al. 2021 Gong et al. 2020 Zhuo et al. 2018 Liu et al.
2022) suggest
that SVGD suffers from curse of dimensionality and the variance SVGD-trained particles tend to diminish in high- -ich-iimensional spaces. We highlight that one of the key issues is that the SVGD algorithm with RBF function space improper. As shown in Tab. when fitting Gaussian distribution, an ideal function class is the linear function class given the Gaussian initialization. However, provided with RBF-based SVGD, the variance collapse phenomenon is extremely severe. On the contrary, when using proper function class (linear), both SVGD and PFG performs well. More importantly, for some powerful non-linear function classes (ncural networks), still performs well with PFG. The effectiveness of neural network function class and PFG algorithm would be particularly important when the target distribution more complex.
Table 5: Illustration of Variance Collapse in RBF function classes (20-dim (N,
METHOD
TARGET
SVGD (RBF)
SVGD (LINEAR)
PFG (LINEAR)
PFG (NEURAL NETWORK)
VARIANCE
1.00
0.35
1.01
1.00
0.99
To further justify the effectiveness of PFG algorithm in high dimensional cases, we evaluate our model under different settings following Liu et al. (2022). Tab. shows that our PFG algorithm is robust to high dimensional fitting problems, which is comparable GSVGD (Liu al. (022) in Gaussian case.
Table 6: Estimating the dimension-averaaged marginal variance of A A(, 7)
DIMENSION
20
40
60
80
100
SVGD
0.35
0.18
0.12
0.09
0.04
GSVGD
0.96
0.97
1.00
1.01
1.00
PFG
1.00
0.99
0.98
1.00
0.97
We also include a further justification in Tab. 7 which compute the energy distance between particle the target distribution and the particle estimation 4-mode Gaussian mixture distribution. In this case, the function class of gradient should be much larger than linear case. We can find that the resulting performance improves the previous work (GSVGD).
21

---

Published as a conference paper at ICLR 2023
Table 7: Energy distance (x10-2) between the target distribution and the particle estimation on Multimodal Gaussian mixture distribution (4 modes)
DIMENSION
20
40
60
80
100
SVGD
2.1
6.7
10.8
24.9
39.6
GSVGD
1.0
2.3
3.4
3.2
3.5
PFG
0.6
0.8
1.2
1.6
2.1
B.5 ESTIMATION OF PRECONDITIONING MATRIX
Ideally, one would leverage the exact Hessian inverse as the preconditioning matrix to make the algo- rithm well-conditioned, which is also the target of other preconditioned SVGD algorithms. However, the computation exact Hessian is challenging. Thus, motivated by adaptive gradient optimization algorithms, leverage the moving average of Fisher information to obtain the preconditioning matrix. We conduct hierarchical logistic regression to discuss the importance of preconditioning. We compare 3 kinds of preconditioning strategies to estimate: (1) Full Hessian: compute the exact full Hessian matrix; (2) K-FAC: Cronecker-factordd Approximate Curvature to estimate Hessian; (3) Diagonal: Diagonal moving average estimate Hessian. According to the table, can find that the preconditioning indeed improves the performance of Bayesian inference, and the choice of estimation strategies does not differ much. Thus, we choose the most efficient one: diagonal estimation.
Table 8: Average Performance on Hierarchical Logistic Regression (German dataset)
FULL HESSIAN
K-FAC
DIAGONAL
w/o PRECONDITIONING
ACCURACY NLL
67.8
67.4
67.6
66.2
0.62
0.65
0.64
0.70
B.6 MORE COMPARISONS
We also include a more complete empirical results on Bayesian neural networks, including the comparison with plain version of SVGD (Liu & Wang 2016), SGLD (Welling & Teh 2011). and Sliced SVGD (S- (S-SVGD) Gong et al. 2020). Our algorithm is still the most competitive one according to the table. Besides, also conduct ablation study to further justify the importance of preconditioning, which implies the value of Q(-) design. In Table 11 algorithm with proper Q outperforms the plain version.
we have shown that the full
Table 9: Comparison with plain version of SVGD and SGLD. Averaged test root-mean-squan (RMSE) and test log- likelihoodd of Bayesian Neural Networks on UCI datasets (100 particles). Results are computed by 10 trials.
AVERAGE TEST RMSE
AveRAGE TEST LOG-LIKELIHOOD
DATASET
SVGD
SGLD
PFG
SVGD
SGLD
PFG
BOSTON
3.04
2.79+0.16
2.47 +0.11
-2.76-u15
-2.6340.14
-2.3540.12
CONCRETE
5.51
4.97 +0.12
4.69 10.14
-3.07
-3.03
-2.83
ENERGY
1.9640.10
0,84+p.07
20.04
-2.15
-1,82
-1.2240.06
PROTEIN
4.89410.11
4.6110.11
4.51 1:0.06
-3.12-0.14
-2.9910.14
-2.8940.07
WINERED
0.68
0.6710.05
0.60.0.00
-1.92-0.05
-1.81 10.08
-1.61
WINEEHITE
0.69:40.04
0.71 1:0.6
0.59 1.000
-1.85
-1.69 10.07
-1.58
YEAR
8.65+0.08
8.6940.06
8.56
-3.54_0.10
-3.6440.09
-3.51
22

---

Published as a conference paper at ICLR 2023
Table 10: Comparison with Sliced SVGD (S-SVGD). Averaged test root-mean-square emror (RMSE) and test log-likelihood of Bayesian Neural Networks on UCI datasets (100 particles). Results are computed by 10 trials.
AVERAGE TEST RMSE S-SVGD
AVERAGE TEST LL
DATASET
PFG
S-SVGD
PFG
BOSTON
2.87
2..7
-2.5020.16
-2.35
CONCRETE
4.8840.08
4.69
-3.00
-2.8340.16
ENERGY
1.13
0.484
-1.5420.05
-1.2221006
PROTEIN
4.58
40.00
-2.99
-2.89 +0.07
WINERED
0.64
0.60
-1.82
-1.61
WINE HITE
0.62
0.59
-1.71
-1.58
YEAR
8.77
8.56
-3.60
-3.01
Table 11: Ablation Study, Averaged test root-mean-square error (RMSE) and test log-likelihood of Bayesian Neural Networks on UCI datasets (100 particles). Results are computed by 10 trials.
AVERAGE TEST RMSE
AVERAGE TEST LL
DATASET
W/O PRECONDITION
PFG (FULL)
w/o PRECONDITION
PFG (FULL)
BOSTON
2.69
2.47
2..55
-2355
CONCRETE
4.90=0.06
4.69
-2.94
20.
ENERGY
10.05
0.48
11.58
.1.22.0.06
PROTHIN
4.60
4.50
-3.04
-2.89 20.07
WINERED
0.66
0.60
-1.75
-1.61 20.00
WINEWHITE
0.64
0.59
-1.80
-1.58
YEAR
20.06
8.56
-3.5920.06
-3.51
B.7 PARTICLE-BASED VI VS LANGEVIN DYNAMICS
One may be interested in the reason why choose particle-based VI rather than Langevin dynamics. In the continuous case, Langevin dynamics solves Bokker-Plance equation. We have several reasons demonstrate the superiority of proposed framework rather than Langevin dynamies.
Motivation of particle-baase variational inference: Deterministic update and repulsive interactions.
One of the key algorithmic differences between particle-based variational inference and Langevin dynamics is the realization of V In Pt. where particle-based variational inference explicitly estimates the deterministic repulsive function and Langevin dynamies uses Brownian motion.
The deterministic version repulsive force introduces interactions between particles while the stochastic version only maintain the variance with the randomness. Thus, for each particle, only the deterministic algorithm leads to convergence, which is more stable and efficient (wrt the number of particles).
Figure has demonstrated the phenomenon: we use both particle-based variational inference and Langevin dynamics to sample from Gaussian distribution. It is clear that although both algorithms return reasonable particles, the Langevin-indun induced particles are fully random and highly unstable due to the empirical randomness, but particle-basse VI is robust against different random seeds. When number particles is small, sample efficiency particle-based VI should much better than Langevin dynamics.
Besides, the deterministic update can induce transport function, which can used to map input to the target distribution directly, which be great potential of propose framework. For example, we can maintain the composite function all the particle updates. When T++1 = fe(x;) = I; (t, 22), f 1..... T 1, then FT -r-1 painlessly. For the Gaussian case, this composite function is just a linear transform without other overheads. For neural networks, may also distillation to obtain transport function. We believe that the distillation of the transport function have great potential in the future.
fi (x1). then can perform resampling
2. Discretization of Fokker-Plaane equation: forward-Flow discretization vs Fuler discretization.
23

---

Published as a conference paper at ICLR 2023
About the discretization of Fokker-Planck equation, Langevin dynamics is performing Forward-FFoe (FFI) discretization, which not the same as conventional gradient desent (Euler discretization). In word, FFI only discretizes the in term, but solve In term with SDE, so that the discretized gradient is biased in general (Wibisono, 2018). The particle-based variational inference tend to perform Fuler discretization, which is unbiased by discretizing the full (Wasserstein) gradient (similar to conventional gradient descent). Thus, from a theoretical perspective, Euler-type discretization is simpler and more direct, which worthwhile to be further explored as our paper.
Random seed
Random seed
Random seed 3
(a)
(b)
(d)
(f)
Figure Particle-based (a-c) vs Langevin dynamics (d-f). (Blue dots refer to particles obtained)
3. Function classes: non-linear function class (neural networks) vs linear function class (RKHS).
(Note that the linearity wrt function bases rather than the plain linear function)
In our framework, both RKHS and neural networks (or other function class) are valid function classes to estimate the Wasserstein gradient.
In particular, neural networks are proven to outperform conventional RKHS in many areas, because it a non-linear function class with learnable features, that can work with uneven subspaces. The RBF kernel uses the same smoothing operator for all gradients. In Figure 3 we have shown that the RKHS is incapable of capturing functional gradient near connected clusters, while neural networks can do so. As result, the sampling quality can be improved.
24

---

Published as a conference paper at ICLR 2023
C IMPLEMENTATION DETAILS
All experiments conducted Python 3.7 with NVIDIA 2080 Ti. Particularly, PyTorch 1.9 to build models. Besides, Numpy, Seipy, Sklearn Matplotlib, Pillow are used in the models.
C.1 SYNTHETIC EXPERIMENTS
C.1.1 ILL-CONDITIONED GAUSSIAN DISTRIBUTION
For ill-conditioned Gaussian distribution, we reproduce the continuous dynamics with Euler discretiza-
tion. We consider cases: (a) μ,, [1, 1]_ Σ. diag(L (1 1): (b) μsv = [1, 1] E, diag(100, 1);
(c) = [20, 20] Σ, = (iag(1, 1); (d) j = [20, 20 Σ. = diag 100, 1) to illustrate the behavior of SVGD and our algorithm clearly.
Notably, for our algorithms and SVGD (linear kernel), we use step size 10-3 to approximate the continuous dynamics and solve the mean and variance exactly (equivalent infinite particles); for SVGD with RBF kemel,, we use step size 10 with 1, 1,000 particles to approximate mean and variance. For the neural network function class, we assume that the function class two-layer neural network Fe = {So 0 }}. In practice, we may incorporate base shift So as gradient boosting to accelerate the convergence, we use F = {fo + fe : fe E Fo). where fo(x) = or fo(x) cV In (x) with some constant In practice, by validation. Empirically, high dimensional models (such as Bayesian neural networks) can be accelerated significantly with the base function.
select from {0, 0.1, 0.2, 0.5, 1}
C.1.2 GAUSSIAN MIXTURE
For Gaussian Mixture, we sample 2-d 10-cluster Gaussian Mixture, where cluster means are sampled from standard Normal distribution, covariance matrix 0.1²1, the marginal probability of each cluster is 1/10. The function class is 2-layer network with 32 hidden neurons and Tanh activation function. For inner loop, we choose 0.9. For particle optimization, choose siep-size 1e-1. For SVGD, choose RBF kernel with median bandwidth, step-size is 1e-2.
and SGD optimizer le-3) with momentum
C.2 APPROXIMATE POSTERIOR INFERENCE.
To compute the trace, we use the randomized trace estimation (Hutchinson's 's trace estimator) son 1989) to accelerate the computation.
(Hutchin-
Sample ( such that E = and Covê
Compute Ī())- Σm (Q((((at) + fo(x)) - fo(xt) VU(x)) & EK & (TVoo(a))s)
Compute L(0) = 1 m Σm + - fol)) vU(x)) - V fo(oii))
Update of' = -'vi0''
where is chosen Rademacher distribution.
C.2.1 LOGISTIC REGRESSION
For logistic regression, we compare our algorithm with SVGD in terms of the "goodness" of particle distribution. The ground truth computed by NUTS (Hoffman al. 2014). The mean is computed by 40,000 samples, the MMD is computed with 4,000 samples. The H in this part is chosen as H H 0.1
In the experiments, we select step-size from {10-1, 1.022.102, 20-3 10-3, 10-3, 10003110-00000100 1004) for each algorithm by vali- dation. And use Adam optimizer without momentum. For SVGD, we choose RBF kernel with median bandwidth. For PFG, we select inner loop from (1, 2, 5, 10) by validation. The hidden neuron is 32.
25

---

Published as a conference paper at ICLR 2023
C.2.2 HIERARCHICAL LOGISTIC REGRESSION
For hierarchical logistic regression, use the test performance measure the different algorithms, including likelihood and accuracy. The H in this part is chosen as -0.5
For all experiments, batch-size is 200 and we select step-size from {10-1, ,0-2,,00- 2.10000-3 1100$.10-4) for each algorithm by validation. And use Adam optimizer without momentum. For SVGD, choose RBF kernel with median bandwidth. For PFG, we select inner loop from 21, 10) by validation. The hidden neuron is 32.
C.2.3 BAYESIAN NEURAL NETWORKS (BNN)
We have two experiments in this part: UCI datasets, MNIST classification. The metric includes MSE/accuracy and likelihood. The hidden layer size is chosen from {32, 64, 128, 256, 512). To approximate H, we use the approximated diagonal Hessian matrix H, and choose H/ = Ho where E {0, 0.1, 0.2, 0.5, 1}. The inner loop T of PFG chosen from {1,2,5, 5, 10} and with SGD optimizer (Ir 1e-3, momentum=0.9).
For UCT datasets, we follow the setting of SVGD (Liu Wang 2016). Data samples are randomly partitioned to two parts: 90% (training), 10% (testing) and we use 100 particles for inference. We conduct the two-layer BNN with 50 hidden units (100 for Year dataset) with ReLU activation function. The batch-size is 100 (1000 for Year dataset). Step-size is 0.01.
For MNIST, we conduct the two-layer BNN with 128 neurons. We compared the experiments with different particle size. Step-sizes are chosen from {10-1, 10 10--2 101 3, 10-41. The batch-size is 100.
26


Reviews:
[{"review_id": "zLJD7NSNIt", "reviewer": "Reviewer_m3yE", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a new variational inference method for sampling from the distribution. The proposed method is a sort of extension of the SVGD method inspired by functional gradient boosting. That is, the drift term is modeled by the neural networks and is optimized to maximize the discrepancy under an appropriate regularization. This work also provides a convergence guarantee under the assumptions regarding the approximation error to the preconditioned Wasserstein gradient. The effectiveness of the method is empirically verified on several tasks: sampling from Gaussian mixture, ill-conditioned Gaussian distribution, and Bayesian and hierarchical logistic regressions.", "summary_of_the_review": "The proposed method is interesting because it is a proper extension of SVGD. But, the motivation for using neural networks as the drift term is rather weak. More discussion to motivate the use of the method will make the paper stronger.", "strength_and_weaknesses": "**Strengths**:\n\nThe proposed method can be seen as a proper extension of Stein variational gradient descent (SVGD). Indeed, SVGD is a special case of the proposed method through the explicit feature representation of the kernel. In this sense, this work makes a certain contribution to the context by proposing a general form of SVGD. \n\n**Weaknesses**:\n\nThe motivation for using neural networks for representing the drift term is somewhat unclear. Indeed, Langevin dynamics can be directly applied for sampling from Gibbs distribution $p_*$ and this dynamics does not require the approximation. (Note that Langevin dynamics also solves Fokker-Planck equation Eq. (1) with the Wasserstein gradient.) Thus, I think the standard Langevin dynamics is sufficiently powerful in terms of optimization. Moreover, there are also preconditioned variants of Langevin dynamics:\n\n1. Max Welling and Yee Whye Teh. Bayesian Learning via Stochastic Gradient Langevin Dynamics.\n2. Gaétan Marceau-Caron and Yann Ollivier. Natural Langevin Dynamics for Neural Networks.\n3. C. Li et al., Preconditioned Stochastic Gradient Langevin Dynamics for Deep Neural Networks.\n4. A. Wibisono. Proximal Langevin Algorithm: Rapid Convergence Under Isoperimetry\n\nThese preconditioned Langevin dynamics can be competitors to the proposed method. Thus, the discussion of these methods would be helpful. In particular,  [WIbisono] (see Section 4.1) considered the dynamics which optimizes Eq. (7) in the submission. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "srHIMRkvTon", "reviewer": "Reviewer_9msK", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper studies the general gradient flow of KL divergences. The authors consider a generalized metric space and formulate the gradient descent as a least square problem. The metric spaces include L^2, RKHS, and covariance-weighted norms. Some linear and nonlinear families of functions, including neural networks, are used in approximating generalized metric spaces and gradients. Numerical examples are shown the effectiveness of the proposed methods. ", "summary_of_the_review": "The paper is written well with clear mathematics. Some analytical examples and important literature are missed.  It will take effect my ratings. \n", "strength_and_weaknesses": "Strength: The authors propose general metric space-induced gradient descent to design sampling algorithms. The generality of metrics (Q) may improve the sampling algorithm.  \n\nWeakness: \n\n1. No convincing examples are provided analytically. I suggest that the authors study a Gaussian target distribution. In this case, the author may illustrate the advantage of general metrics and gradients. \n\n2. The authors miss many crucial works in this field.  Please illustrate them in the context.\n\nA. Garbuno-Inigo, F. Hoffmann, W.C. Li, A.M. Stuart. Interacting Langevin Diffusions: Gradient Structure And Ensemble Kalman Sampler. SIAM Journal on applied dynamical system, 2019. \n\nW.C. Li, L.X. Ying. Hessian transport gradient flows. Research in the Mathematical Sciences, 2019. \n\nY.F. Wang, W.C. Li. Information Newton flow: second-order optimization method in probability space. 2020. \n\nY.F. Wang, W.C. Li. Accelerated Information Gradient flow. \n\nY.F. Wang, P. Chen, W.C. Li. Projected Wasserstein gradient descent for high-dimensional Bayesian inference. \n\nY.F. Wang, P. Chen, M. Pilanci, W.C. Li. Optimal Neural Network Approximation of Wasserstein Gradient Direction via Convex Optimization. \n\nW.C. Li, A.T. Lin, G. Montufar. Affine Natural Proximal Learning. GSI, 2019.\n\nA.T. Lin, W.C. Li, S. Osher, G. Montufar. Wasserstein Proximal of GANs. GSI, 2021. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "phJv_r1dK1K", "reviewer": "Reviewer_vc1L", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposed a new partical-based variational inference method by minimization the functional gradient of KL with a regularization where preconditioning seems to be crucial for the performance. The efficiency compared to SVGD is illustrated on several standard benchmark data sets. Relations to previous work were also discussed.", "summary_of_the_review": "The paper proposed a new precondition functional gradient flow approach for particle-based VI and demonstrated its efficiency on several benchmark problems against SVGD. The method can be viewed as a score matching approach that learns the discrepancy of the target score and current particle scores, and update the particles according to the learned functional gradient flows. In that sense, the paper is highly related to a recent paper called neural variational gradient descent by Lauro Langosco et al., 2021. Despite this, the paper do provide some theoretical analysis regarding convergence given well trained functional gradients and introduce the precondition trick which seems to be crucial for performance. I, therefore, recommend it for a weak acceptance.\n\n\nminor issue: there is a mistake in the upper bound of dKL/dt term in Theorem 1.", "strength_and_weaknesses": "Strength:\n\n1. The precondition functional gradient flow for particle-based VI seems to be new.\n2. The method has been well motivated with clear theoretical justification.\n3. The advantages over SVGD has been clearly stated and illustrated in the experiments.\n\n\nWeaknesses:\n\n1. The preconditioning seems to be crucial for performance, lack of ablation study except for the toy examples.\n2. No theoretical analysis for the precondition matrix H.\n3. Discretization error of ODE is not explicitly considered.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster