Particle-based Variational Inference with Generalized Wasserstein Gradient Flow
Authors: Ziheng Cheng
Shiyue Zhang
Longlin Yu
Cheng Zhang
Venue: NeurIPS 2023 poster (2023)
Keywords: Particle-based VI
generalized Wasserstein gradient flow

Abstract: Particle-based variational inference methods (ParVIs) such as Stein variational gradient descent (SVGD) update the particles based on the kernelized Wasserstein gradient flow for the Kullback-Leibler (KL) divergence. However, the design of kernels is often non-trivial and can be restrictive for the flexibility of the method. Recent works show that functional gradient flow approximations with quadratic form regularization terms can improve performance. In this paper, we propose a ParVI framework, called generalized Wasserstein gradient descent (GWG), based on a generalized Wasserstein gradient flow of the KL divergence, which can be viewed as a functional gradient method with a broader class of regularizers induced by convex functions. We show that GWG exhibits strong convergence guarantees. We also provide an adaptive version that automatically chooses Wasserstein metric to accelerate convergence. In experiments, we demonstrate the effectiveness and efficiency of the proposed framework on both simulated and real data problems.

Particle-based Variational Inference with Generalized Wasserstein Gradient Flow
Ziheng Cheng"
Shiyue Zhang'
School of Mathematical Sciences
School of Mathematical Sciences
Peking University alex-czhèste pku. edu
Peking University
zhangshiyue@stu pku edu, cn
Longlin Yu School of Mathematical Sciences Peking University 11yuêpku. edu. cn
Cheng Zhang!
School of Mathematical Sciences and Center for Statistical Science Peking University
chengzhang@math. pku edu.
Abstract
Particle-based variational inference methods (ParVIs) such Stein variational gradient descent (SVGD) update the particles based on the kernelized Wasserstein gradient flow for the Kullback-Leibhe (KL) divergence. However, the design of kernels is often non-trivial and can be restrictive for the flexibility of the method. Recent works show that functional gradient flow approximations with quadratic form regularization termss can improve performance. In this paper, we propose a ParVI framework, called generalized Wasserstein gradient descent (GWG), based a generalized Wasserstein gradient flow of the KI. divergence, which can be viewed as a functional gradient method with a broader class of regularizers induced by convex functions. We show that GWG exhibits strong convergence guarantees. We also provide an adaptive version that automatically chooses Wasserstein metric to accelerate convergence. In experiments, we demonstrate the effectiveness and efficiency of the proposed framework on both simulated and real data problems.
Introduction
Bayesian inference is an important method in modern machine learning that provides powerful tools for modeling complex data and reasoning under uncertainty. The core Bayesian inference to estimate the posterior distribution given the data. As the posterior distribution is intractable in general, various approximation approaches have been developed, of which variational inference and Markov Chain Monte Carlo are two typical examples. By reformulating the inference problem into an optimization problem, variational inference (VI) seeks an approximation within a certain family of distributions that minimizes the Kullback-Leible (KL) divergence to the posterior (Jordan et al.) 1999. Wainwright & Jordan 20088 Blei et al. 2016). Equipped with efficient optimization algorithms, VI allows fast training and easy scaling to large datasets. However, the construction of approximating
"Contributed equally this work. 'Correponding author.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

distributions can be restrictive which may lead to poor approximation. Markov chain Monte Carlo (MCMC) methods simulate a Markov chain to directly draw samples from the posterior (Duane et al.) Welling & Teh, 2011 Chen et al. 2014). While being asympiotically unbiased, MCMC often takes a long time to converge, and it is also difficult to access the convergence.
1987, Neal 2011
Recently, particle based variational inference methods (ParVIs) have been proposed that tend to combine the best of both worlds (Liu & Wang 2016 Chen et 2018 Liu et al. 2019 di Langosco 2022 Alvarez-Melis et al.] (2022). In Par Is, the approximating distribution is represented as a set of particles, which are iteratively updated by minimizing the KL divergence the posterior. This non- parametric nature significantly improves the flexibility of ParVIs upon classical VIs, and the interaction between particles also makes ParVIs more particle-efficien than MCMCs.
et al. 2021 Fan et al.
The most well-known partiele based VI method is Stein Variational Gradient Descent (SVGD) (Liu Wang 2016). It updates the particles by simulating the gradient flows of the KI. divergence on a certain kernel related distribution space, where the gradient flows have tractable form (Liu 2017] Chewi et al., 2020). However, SVGD relies on the choice of an appropriate kernel function whose design highly non-trivial and hence could limit the flexibility of the method. Moreover, the required computation of the kernel matrix scales quadratically with the number of particles, which makes costly to use a large number of particles.
Instead of using kernel induced functional gradients, many attempts have been made to expand the function class for gradient flow approximation (Hu et al.] 2018; Grathwohl et al. [0200] di Langosco al., 202T; Dong et al. 2023). By leveraging the more general neural networks as the function class together with more general regularizers, these approaches have shown improved performance over vanilla SVGD while not requiring expensive kernel computation. However, these methods only use quadratic form regularizers where either the Wasserstein gradient flow or its preconditioned variant is recovered.
In this work, we propose a VI method based on a general formulation of minimizing movement scheme in Wasserstein space, which corresponds to a generalized Wasserstein gradient flow of KL divergence. Using Legendre- Fenchell transformation, our method can also be viewed as functional gradient method with a more general class of regularizers which include the previously used quadratic forms as special cases. We provide theoretical convergence guarantee of Par with netuaal-et--stimated vector field for generalized Wasserstein gradient flow, which to the best of our knowledge, has not been established yet. Perhaps surprisingly, our results show that assuming reasonably accurate vector field estimates, the iteration complexity of ParVIs matches the traditional Langevin Monte Carlo under weaker assumptions on the target distribution. As an extension, we also propose an algorithm that can adaptively adjust the Wasserstein metric to accelerate convergence. Extensive numerical experiments on both simulated and real data sets are conducted to demonstrate the efficiency of our method over existing ones.
2 Background
Throughout this paper, we use to denote particle samples in Rd. Let P(RR) denote all the probability distributions Rd that are absolute continuous with respect to the Lebesgue measure. We do not distinguish probabilistic measure with its density function. For E Rd and p > 1, | (|=r)p + .... /// stands for the f,-norm. The Hölder conjugate of p is denoted by q : p/(p 1). Notation g" (-) denotes the Legendre transform of a convex function g(-) on Rd.
Notations.
2.1 Particle-based Variational Inference
Let π E (((R) be the target distribution we wish to sample from. We can cast the problem of sampling as an optimization problem: to construct a distribution μ²" that minimizes the KI. divergence
μ := arg min DKL (HI(TT WEP
(1)
where P' C P(RR)) is the variational family. Particle-based variational inference methods (ParVIs) class of VI methods where P' is represented set of particles. Assume the current particle distribution is μr, then it holds that
d
DRI(lid + (1) V Voog v).
(2)
de ==

---

ParVIs aim to find the optimal vector field v that minimizes in certain function class. For example, SVGD (Lu & Wang 2016) restricts u the unit ball of reproducing kernel Hilbert space (RKHS) which has closed-form solution by kernel trick. Meanwhile, Hu et al. (2018); Grathwohl et al. (2020); Langosco et (221): Dong et al. (2023) consider a more general class of functions for v, i.e., neural networks, and minimize with some quadratic form regularizers.
2.2 Minimizing Movement Scheme in Wasserstein Space
Assume the cost function c(.,) R R + is continuous and bounded from below. Define the optimal transportation cost between two probabilistic measure μ. as:
/ c(2, y)dp.
W (),
inf pEll(p.))
Specifically, if c(1r, y) = || yI for some > 1, then we get the p-th power of Wasserstein-p distance W Jordan et al. (1998) consider a minimizing movement scheme (MMS) under W2 metric. Given the current distribution Mac, the distribution for next step is determined by
1
M(++1j6 arg min DEL (pllT) + W (1, pkh)-
(4)
(EP2(Bd)
2h
When the step size h (uenkk2 converges the solution of the Fokker-Planne Planck equation
0.p: + div(4 V log π) = Ape..
(5)
Therefore, MMS corresponds to the deterministic dynamics
dz, v,dt V log v log Pr:
(6)
where μ4 is the law of rτ. (6) is also known as the gradient flow of KL divergence under W₂ metric, which we refer to as L2-GF (Ambrosio et al.) (005). Note that the Langevin dynamics dæ, = v (tt + V2dB, (B, is the Brownian motion) reproduces the same distribution curve f2e and thus also corresponds the Wasserstein gradient flow (Jordan al. 1998).
3 Proposed Methods
3.1 Minimizing Movement Scheme with A General Metric
We start with generalizing the scope of the aforementioned MMS in Section 2.2 which is under W2 metric.
Definition 3.1 (Young function). A strictly convex function g Rd is called Young function if g(z) = g((=)) (0)) = 0, and for any fixed E (00), hg( ht ) Co, as h + 0.
Theorem 1. Given a continuously differentiable Young function g and step size h > O, define cost
y hh. Suppose that æ, Pkkk E Pen (R") == {ur E P(RR) Ep 2x
function Cn (x, v) = g( h
h
oo). Under some mild conditions of g (see details in Proposition A. Pa,, (R.) is a Wasserstein space equipped with Wasserstein distance. Consider MMS under transportation cost Wa.:
M(K+1)A == argmin DKL(H]|T)) + Wa, (K.)..).
(EPo,,(R))
Denote the optimal transportation map under We, from μ(K+1)N to Pkb by Tx(.). Then we have
Te(x) Vg" (V log log
(8)
Please refer Appendix A for full statements and proofs. Informally, M(K+1)A R Pun. for small step
size (Santambrogio 2017). Further note that (x is the optimal velocity field associated
h with the transport from 1)4+1)A p&s (and vice versa). step size h + 0, then following Jordan et al. (1998). we can recover the dynamics in continuous time:
dr; v,dt, (V V (.)...
(9)

---

We call 99 the generalized Wasserstein gradient (GWG) flow. If we set g(-) = 221l 12 or any positive
definite quadratic form g(-) - ll ll then (9 reduces to L2-GF or its preconditioned version
(Dong al. 2023 respectively.
3.2 Faster Descent of KI. Divergence
It turns out that we can leverage the general formulation (9) to explore the underlying structure of different probability spaces and further utilize this geometric structure to accelerate sampling. More
specifically, we consider g(-) = ||
Il for some
and then g" (+) =
l Note that if the
particles move along the vector field v1 = Vgg" (V log
π - ). fe
then the descent rate of (KL(Hell (Hell is
9 jg
π
Og (KLllellTT = -Ep,
V log
(10)
If we choose q such that Ey.
V log
is large, then DKL (Me (|))) decreases faster and the sampling
Pr lg
process can be accelerated. We the following example further illustrate our idea. Please refer to Appendix B for detailed analysis.
Example 1. Let -- 1) 1 2 1) and μ
3 N 1) + 1 N(m, 1). Then for any 4
80 9 1, the following holds:
0.08 m qm
m
0.2
m
exp(
S Ey
log
(4m) expl
(11)
2
qm
2
However, the KL divergence between π and μ is large: (KL(H||T) 2
10
Suppose the target distribution is and run ParVI with current particle distribution μ.. We can expect that, if simply using L2 regularization, i.e., 2, then for very large m, the score divergence small and thus the decay of KL divergence is extremely slow. However, DKt. is still large, indicating that it would take a long time for the dynamics to converge to the target. But if we set q much larger, then the derivative of KL divergence would get larger and the convergence can be accelerated.
3.3 Algorithm
The Forward-Euled discretization of the dynamics 9
T(k+ TRk + Vg"
log
π
(Iaa) h.
(12)
Pah
However, since the score of current particle distribution pun is generally unknown, we need method to efficiently estimate the GWG direction Vg" log
Given the distribution of current particles
μ. by the definition of convex conjugate, we have
(Vlog arg max Ea log μ .b) g(v)).
If parameterize neural network fwr with W, then maximize the following objective with respect to w:
C(w) : = E log fw) - g(fw)]
(13)
= [V log T) fmr + V for - g(fw)

---

Algorithm GWG:
Generalized Wasserstein Gradient Flow
Require: Unnormalized target distribution æ, initial particles {ril==1 initial parameter wo, iteration number N, NN, particle step size h. parameter step size 17
for 0,... N do
Assign wg wk for 0,... N' 1 do Compute
[(w)
Σ
log (og)(a)) sm(h))) + V. (m ((fm(r)))
(14)
11. 1=1
Update wik 1+1 = w + 11 1/ (.n((wn)))
end for Update wk+1 wk N'
Update particles x2+1 = 2 + fefwaas (2t) for 1,... n end for return Particles ("
Here the second equation is by Stein's identity (we assume vanishes at infinity). This way, the gradient of L(w) can be estimated via Monte Carlo methods given the current particles. We summarize the procedures in Algorithm
The exact computation of the divergence term Vz (z) needs (dd) times back propaaaop,g,,ion, where d is the dimension of z. In order to reduce computation cost, we refer to Hutchinson's 's estimator (Hutchinson 1989 i.e.,
(x))
Σ of (fin &.),
(15)
where & E Rd are independent random vectors satisfying Eg = Ia. This is still an unbiased estimator but only needs ((1) times back-propagation.
4 Convergence Analysis without Isoperimetry
this section, we state our main theoretical results of Algorithm Consider the discrete dynamics:
X(K+1)A = -Xa + Us (Xx)),,
(16)
where un is the neural-net-estimate GWG time kh. Define the interpolation process
X₁ Xkhh
(Xan), for E [kh, (k+ 1)h],
(17)
and let μr denote the law of X,. Note that here we do not assume isoperimetry target distribution (e.g., log-Sobolev inequality) and hence establish the convergence of dynamics in terms of score divergence, following the framework of non-log-concave sampling (Balasubramanian al. 2022
We first make some basic assumptions. For simplicity, only two types of Young function g are considered here, which are also the most common choices.
Assumption 1. (-(-) = -1- 11.11g.oo some > 1. And for any k. Eppi
1
Vg
log
π
< Ek.
Assumption 2. g" (-) a-strongly convex and B-smooth. Define k. And for any ki, Eyan
Uk
Vg" (V log
<
Ek-
Pich
The two assumptions above ensure the estimation accuracy of neural nets. Note that the preconditioned quadratic form in Dong et al. (203) is included in Assumption 2 Although the estimation error is not exactly the training objective used in Algorithm the following proposition shows the equivalence between them in some sense.
5

---

Proposition 2. Suppose g(-) = 11 for some > 1. Given current particle distribution we can
define the training loss Liaiaa (u) = Ey log v) g(v)). ) The maximizer v" Vgg (V log and the maximum value is Lirain = Lumin (v") 0o. For any arbitrarily small eι there exists 82 = S2(11))) < Oc, such that
"
Ey Vgg( ") S Si + E2lCimain Lrmin (o)nin(v))
Besides, ifp 2 2, & can be while E2 still finite.
Similar results also hold if satisfies Assumption since it is equivalent to the case when p 2. Additionally, we expect some properties of the estimated vector fields.
Assumption (Smoothness of neural nets). For any k. (-) is twice differentiable. For any 1, (x) Uk (u)lll 0c, Mp sup lim 2,5 8-0-
Vex + 82) (ex(x))lo) < oc.
Gp : sup x.9 x y|l
8||=||
Note that here we do not assume smoothness of potential log π explicitly, But informally, G, and Mp correspond to the Lipschitz constant of the gradient and the Hessian of log respectively.
Nh /
Nh
Let ANN=
prdl and Ko DKI. (uolFT.. Now present our main results.
Theorem 3 (Full version see Theorem D.9). Under Assumption 7 proper step size h;
the following bound holds with
11 = 14
(
( (M futr + G2Kod
NN Ek
Easu
V
log
π PNN
(18)
N
N
Here O-)) hides all the constant factors that only depend on Theorem 4 (Full version see Theorem D..0). Under Assumption
2
with a 1, the following
bound holds with proper step size h:
(
=0 O ( (oMaa N
G2Kold k)
Lk=D EK
EaNa
log
(19)
NA
N
The proofs in this section are deferred to Appendix
To interpret our results, suppose g g²()) =1-
and < Mp yg. If the neural net (-) can approximate Vg"
π
(V log
accurately (i.e., EN S c).
then to obtain probabilistic measure y such that Ep
V log
S c, the iteration complexity is
OM,KKod -(+ If we further let = p 2, the complexity OKrde 2), which matches the complexity of Langevin Monte Carlo (LMC) under the Hessian smoothness and the growth order assumption (Balasubramanian et al., (022). However, noticing that Assumption is similar to Hessian smoothness informally, we obtain this rate without additional assumption target distribution. This suggests the potential benefits of particle-based methods.
In addition, our formulation allows wider range of choices of Young function, including || II and the preconditioned quadratic form (Dong al. (023). This provides wider options of convergence metrics. We refer the readers to Appendix D.4 for more discussions.
Extensions: Adaptive Generalized Wasserstein Gradient Flow
The GWG framework also allows adaption of the Young function g, instead of fixed one. Similar ideas are also presented in Wang et al. (2018) In this section, we consider special Young function
class
and propose a procedure that adaptively chooses to accelerate convergence. 18 :

---

Algorithm 2 Ada-GWG: Adaptive Generalized Wasserstein Gradient Flow
Require: unnormalized target distribution и, initial particles {at {=0]i=1- initial parameter wo, iteration number N, N", step size h, 1), 17, lower and upper bounds on lb, ub for 0,... N do
Assign ugg wk
for = 0, ,N do
Compute
[(w) = n i=1
v log (x1)) fu(z)) + v fa(())
1
(22)
Pk
Update w ++1 = w + v end for
Update wk+11 wk N'
Compute A(pk) 11
1
Ww 1 (1
Pk
Update Pt+1 = clip(pk + vj (pk)). ) 1b, ub)
Update particles 4 x + f (2 for = 1,.., end for return Particles {=
Consider the continuous time dynamics dx, = fe(x))dt and denote the distribution of particles at time as Mr,, we have the following proposition.
Proposition 5. For g(-) -| Ipg the derivative of KL divergence has an upper bound:
p
M.DKL(HelET
B log
1 (gg((loo) ")
Euc
1
(20)
Me
p
The proof is in Appendix If the neural network ft can approximate the objective well, i.e. "). then informally we can omit the second term and thus
Jt Vg" (V log
14
- - Vog
0 (KL(
R Bp, IIfall =: - A(p)-
(21)
14
9 p
In order let KL divergence decrease faster, we can choose p such that A(p) larger. This leads to simple adaptive procedure that updates p by gradient ascent w.r.t. A(p). In practice, the adjustment of p is delicate and would cause numerical instability if p becomes excessively small or large. Therefore it is necessary to clip within reasonable range. We call this adaptive version of GWG, Ada-GWG. The whole training procedure of Ada-GWG is shown Algorithm [ Note that (22) can be replaced with Hutchinson's s estimator
15
to improve computational efficiency as before.
6 Numerical Experiments
In this section, compare GWG and Ada-GWG with other ParVI methods including SVGD LLu & Wang 2016), 1.2GFF (di Langosco et al., 2021) and PFG (Dong et al. 2023) on both synthetic and real data problems. In BNN experiments, we also test stochastic gradient Langevin dynamics (SGLD). For Ada-GWG, the exponent p is clipped between 1.1 and 4.0 unless otherwise specified. Throughout this section, we choose for be neural network with 2 hidden layers and the initial particle distribution is (0, Ia). We refer the readers to Appendix F for more detailed setups of our experiments. The code available at hhttps://githuu: com/ Alexczh1/GWG
6.1 Gaussian Mixture
Our first example is on multi-mode Gaussian mixture distribution. Following Dong et al. (2023). consider the 10-cluster Gaussian mixture where the variances of the mixture components
7

---

Le
LP 220
bar:
Ado GMS Bon 25
Ado GWG 50
Ado GWG tten 75
Ado GWG Iter 100
Ane GIGG iter. 300
Figure 1: Comparison of sampled particles at different numbers iterations. Upper: L2-GF. Middle: PFG. Lower: Ada-GWG with Po 2.
all 0.1. The number of particles is 1000. Figure I shows the scatter plots of the sampled particles at different numbers of iterations. We see that on this simple toy example, PFG performs similarly to the standard L2-GF which does not involve the preconditioner, while Ada-GWG with the initial Po = 2 significantly accelerates the convergence compared to these two baseline methods. Please refer appendix for further quantitative comparisons.
6.2 Monomial Gamma
To illustrate the effectiveness and efficiency of the adaptive method compared the non-adaptive counterparts, we consider the heavy tailed Monomial Gamma distribution where the target oc exp(-0.3((zr) 10.9 + ().)..
We test GWG and Ada-GWG with different choices of the initial values of p. The number of particles is 1000. Figure 2 demonstrates the KL divergence different methods against the number of iterations. The dotted line represents GWG with fixed while the solid line represents the corresponding Ada-GWG that starts from the same p initialization.
We see that the adaptive method outperforms the non-adaptive counterpart consistently. Moreover, Ada-GWG automatically leam the appropriate value of p especially when the initial values of p is set inappropriately. For example, in our case, a relatively small value of p = 1.5 would be inappropriate (the dotted green line) for GWG, while Ada-GWG with the same initial value of 1.5 is able to provide much better approximation by automatically adjusting p during runtime. Consequently, Ada-GWG exhibit greater robustness when determining the initial value of p.
6.3 Conditioned Diffusion
The conditioned diffusion example is a high-ddmmensionl model arising from a Langevin SDE, with state u [0, 1] IR and dynamics given by
10u u²) dt + dx ug 0,
du:
(23)
+
where x = 122 is standard Brownian motion.
This system is commonly used in molecular dynamics to represent the motion of a particle with negligible mass trapped an energy potential with thermal fluctuations represented by the Brownian forcing Detommaso et al. (2018): ; Cui al. (2016). Given the perturbed observations y, the goal

---

150
1.0
850
6V6
810
250
500
Figure 2: KL divergence of different methods. Solid Figure 3:
Comparison among PFG, Ada-GWG SVGD in conditioned diffusion example.
line: Ada-GWG, Dotted line: GWG counterpart.
to infer the posterior of the driving process p(aly)) The forward operator is defined by F(x) (u.,.... Mt2o) R220 where t, 0.05i. This is achieved by discretizing the above SDE (3
using an Marer-Maruyama scheme with step size At 0.01; therefore the dimensionality of the problem 100. The noisy observations are obtained as = F( + R220 where Tire is Brownian motion path and { ((00 with a 0.1. The prior is given by the Brownian motion = =xz))200-
We test three algorithms: PFG, Ada-GWG, and SVGD, with 1000 particles. To obtain the ground truth posterior, we run LMC with 1000 particles in parallel, using a small step size = 10-1 for 10000 iterations. reportss the logarithmic Maximum Mean Discrepancy (MMD) curves against iterations. We observe that Ada-GWG provides best performance compared to the other methods.
6.4 Bayesian Neural Networks
We compare our algorithm with SGLD and SVGD variants on Bayesian neural networks (BNN). Following Liu & Wangl (2016), we conduct the two-layer network with 50 hidden units and ReLU activation function, and we use Gamma 1, 0.1) prior for the inverse covariances. The datasets all randomly partitioned into 90% for training and 10% for testing. The mini-batch size is 100 except for Concrete on which we use 400. The particle size is 100 and the results are averaged over 10 random trials. Table I shows the average test RMSE and NLL and their standard deviation. We see that Ada-GWG can achieve comparable or better results than the other methods. And the adaptive method consistently improves of different methods on the Boston dataset. We can see that for this specific task, setting p 3 produces better results than when p Although L2-GF (i.e., GWG with p 2) is sub-optimal, our adaptive method (i.e., Ada-GWG with Po 2) makes significant improvements and demonstrates comparable performance to the optimal choice of p = This suggests that our adaptive method is robust even if the initial exponent choice is not ideal. More comparisons of convergence results and hyperparameter tuning details can be found in the appendix.
L2-GF. Figure shows the test RMSE against iterations
Table I: Averaged test RMSE and test negative log-likelihood of Bayesian Neural Networks on several UCI datasets. The results are averaged from 10 independent runs.
AvG. TEST RMSE
AvG. TEST NLL
Dataset
SGLD
SVGD
Lg-GF
Ada-GWG
SGLD
SVGD
Lg-GF
Ada-GWG
BOSTON
3.011 2.774+n.08
3.072=0.10
2.721
2.196
2.444=0.02
2.517+0.14
2.43440.02
CONCRETE
5.583+0.25 4.436+a..8
4.343+o.11
3.871-0.0
3.184+0.04
3.035-o.a2
3.053+0.03
2.826
POWER
4.089 3.97244.02
4.014g0.02
3.944_o.01
2.84040.00
2.809
2.82440.01
2.80220.1
WINEWHITE
0.677+0.01 0.66440.01
0.666
0.660
1.033
1.014-..1
1.015+0.01
1.006
WINERED
0.60020.01 0.579
0.581go.01
0.575_o.01
0.91040.0
0.887
0.860 oz
0.83910.0
PROTEIN
4.560
4.779
1.867
1.686
2.934
2.984
3.003
2.964120.00
7 Conclusion
We introduced a new PI method, called GWG, which corresponds to generalized Wasserstein gradient flow of KI. divergence. We show that our method has strong convergence guarantees in discrete time setting. We also propose adaptive version, called Ada-GWa that can automatically

---

6.-
Ada-GWG (p:=3.0)
Ada GWG (p;=2.0)
5.-
GWG (Py=3.0)
GWG (po-2.0)
Tost RASE
4.0-
3.0
2.0
250
503
750
1000
12/00
3500
175a
2000
Itaration
Figure Test RMSE for the Bayesian Neural Networks Boston dataset. The number in parentheses specifies the initial exponent Po- The results are averaged from 10 independent runs.
adjust the Wassertein metric to accelerate convergence. Extensive numerical results showed that Ada-GWG outperforms conventional ParVII methods.
Acknowledgements
This work was supported by National Natural Science Foundation of China (grant no. 12201014 and grant no. 12292983). The research Cheng Zhang was supported part by National Engineering Laboratory for Big Data Analysis and Applications, the Key Laboratory of Mathematics and Its Applications (LMAM) and the Key Laboratory of Mathematical Economics and Quantitative Finance (LMEQF) of Peking University. Ziheng Cheng and Shiyue Zhang partially supported by the elite undergraduate training program of School of Mathematical Sciences Peking University. The authors grateful for the computational resources provided by the High-performance Computing Platform of Peking University. The authors appreciate the anonymous NeurlPS reviewers for their constructive feedback.
References
Radoslaw Adamezak, Witold Bednorz, and Pawel Wolff. Moment estimates implied by modified log-sobolev inequalities. ESAIM: Probability and Statistics, 21:467-499, 2017.
David Alvarez-Melis, Yair Schiff, and Youssef Mroueh. Optimizing functionals the space of probabilities with input convex neural networks. Transactions on Machine Learning Research, 2022. ISSN 2835-8856. URL httptss://ppeerreiiee wttp://wwwww..eetoooodddddd.ddd. op0Y77888
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005.
Krishnakumar Balasubramanian, Sinho Chewi, Murat A. Erdogdu, Adil Salim, and Matthew Shunshi Zhang. Towards a theory of non-log-concav sampling: First-order stationarity guarantees for langevin monte carlo. In Annual Conference Computational Learning Theory, 2022.
Alessandro Barp, Srunciis-Xavie Briol, Andrew Duncan, Mark Girolami, and Lester Mackey. Minimum stein discrepancy estimators. Advances in Neural Information Processing Systems, 32, 2019.
Franck Barthe and Cyril Roberto. Modified logarithmic sobolev inequalities on. Potential Analysis, 29(2):167, 2008
David M. Blei, Alp Kucukelbir, and Jon D. McAuliffe. Variational inference: A review for statisticians. Journal of the American Statistical Association, 1 112:859- 877, 2016.
Sergey G Bobkov and Michel Ledoux. From brunn-ninkowski to brascamp-lich and logarithmic sobolev incqualities. Geometric and Functional Analysis, 10:1028-1052 2000.
10

---

Changyou Chen, Ruiyi Zhang, Wenlin Wang, Bai Li, and Liqun Chen. A unified particle-opimiiation 1659, 2018.
framework for scalable bayesian sampling. ArXiv,
Tianqi Chen, Emily Fox, and Carlos Guestrin. Stochastic gradient hamiltonian monte carlo. In International Conference Machine Learning, pp. 1683-1691, 2014.
Sinho Chewi, Thibaut Le Gouic, Chen Lu, Tyler Maunu, and Philippe Rigollet. Svgd as a kemel- ized wasserstein gradient flow of the chi-squared divergence. Advances Neural Information Processing Systems, 33, 2020.
Tiangang Cui, Kody JH Law, and Youssef M Marzouk. Dimenssion-indeeeddeer likelihood-informe memc. Journal of Computational Physics, 304:109-137, 2016.
Arnak Dalalyan and Avetik Karagulyan. Usser-friendly guarantees for the langevin monte carlo with inaccurate gradient. Stochastic Processes and their Applications, 129(12):5278-5) 5311,, 2019.
Gianluca Detommaso, Tiangang Cui, Youssef Marzouk, Alessio Spantini, and Robert Scheichl. A stein variational newton method. Advances in Neural Information Processing Systems, 31, 2018.
Luc Devroye, Abbas Mehrabian, and Tommy Reddad. The total variation distance between high- dimensional gaussians. arXiv preprint arXiv: 08693, 6, 2018.
Lauro Langosco di Langosco, Vincent Fortuin, and Heiko Strathmann. Neural variational gradient descent. ArXiv, abs/2107.10731, 2021.
Hanze Dong, Xi Wang, LIN Yong, and Tong Zhang. Particle-based variational inference with preconditioned functional gradient flow. The Eleventh International Conference on Learning Representations, 2023. URL https://openrevieen
S. Duane, A. D. Kennedy, B J. Pendleton, and D. Roweth. Hybrid Monte Carlo. Physics Letters B, 195(2):216 222, 1987.
Jiaojiao Fan, Qinsheng Zhang, Amirhossein Taghvaei, and Yongxin Chen. Variational wasserstein gradient flow. In International Conference on Machine Learning, pp. 6185-6215. PMLR, 2022.
Ivan Gentil, Arnaud Guillin, and Laurent Miclo. Modified logarithmic sobolev inequalities and transportation inequalities. Probability theory and related fields, 133:409 2005.
Ivan Gentil, Arnaud Guillin, and Laurent Miclo. Modified logarithmic sobolev inequalities in null curvature. Revista Matematica Iberoamericana, 23(1):235-22) 258-28, 2007.
Will Grathwohl, Kuan-Chich Wang,
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-3732727 3747. PMLR, 2020.
Tianyang Hu, Zixiang Chen, Hanxi Sun, Jincheng Bai, Mao Ye, and Guang Cheng. Stein neural sampler. arXiv preprint arXiv: 11Xx::8810.03555, 2018
Michael F Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications Ctatisiix-Simmlaaii and Computation, 1859)1005 1989.
Michael I. Jordan, Zoubin Ghahramani, T. Jaakkola, and Lawrence K. Saul. An introduction to variational methods for graphical models. Machine Learning. 37:183-233, 1999.
Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the fokker-planck equation. SIAM journal on mathematical analysis, 29(1):1-17, 1998.
Frederic Kochler, Alexander Heckett, and Andrej Risteski. Statistical efficiency of score matching: The view from isoperimetry. In The Eleventh International Conference on Learning Representa- tions, 2023. URL
net/ /
Anna Korba, Adil Salim, Michael Arbel, Giulia Luise, and Arthur Gretton. A non-asymptotic analysis for stein variational gradient descent, Advances Neural Information Processing Systems, 33: 4672-4682, 2020.
11

---

Chang Liu, Jingwei Zhuo, Pengyu Cheng, Ruiyi Zhang, Jun Zhu, and Lawrence Carin. Understanding and accelerating particle- particle-based variational inference. International Conference on Machine Learning. PMER, pp. 4082-4092, 2019.
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: A general purpose bayesian inference algorithm. Advances neural information processing systems, 29, 2016.
Wenlong Mou, Nicolas Flammarion, Martin Wainwright, and Peter L Bartlett. Improved bounds for discretization of langevin diffusions: Near-optimal rates without convexity. Bernoulli, 28(3): 1577-1601, 2022.
HP Mulholland. On generalizations of minkowski's inequality the form of a triangle inequality. Proceedings of the London mathematical society, 2(1):294-300 1949.
Radford Neal. MCMC using hamiltonian dynamics. In S Brooks, Gelman, G Jones, and XL Meng (eds.), Handbook of Markov Chain Monte Carlo, Chapman & Hall/CRC Handbooks of Modern Statistical Methods. Taylor & Francis, 2011. ISBN 9781420079425. URL http://books googlee comete..oommbee http://oooo..dd-gffsAa.....
Whitney Newey and Daniel McFadden. Large sample estimation and hypothesis testing. Handbook of Economeirics, 4:2111-2245, 1986.
Filippo Santambrogio. (Euclidean, metric, and Wasserstein) gradient flows: an overview. Bulletin of Mathematical Sciences, 7:87-154, 2017.
Yang Song. Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach to density and score estimation. Uncertainty in Artificial Intelligence, pp. 574-584. PMLR, 2020.
Santosh Vempala and Andre Wibisono. Rapid convergence of the unadjusted langevin algorithm: Isoperimetry suffices. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-Bus. Buch-BBu,, E. Fox, and R. Gamett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran As- sociates, Inc., 2019.
pdf
Cédric Villani. Topics in optimal transportation, volume 58. American Mathematical Soc., 2021.
Cédric Villani al. Optimal transport: old and new, volume 338. Springer, 2009.
M. J. Wainwright and M. Jordan. Graphical models, exponential families, and variational inference. Foundations and Trends in Maching Learning, 1(1-2):1-305, 2008
Dilin Wang, Hao Liu, and Qiang Liu. Variational inference with tail-adaptive f-divergence. Advances in Neural Information Processing Systems, 31, 2018.
Yifei Wang and Wuchen Li. Information newton's flow: second-order optimization method in probability space. ArXiv, abs/2001.04341. 2020.
Yifei Wang. Peng Chen, Mert Pilanci, and Wuchen Li. Optimal neural network approximation of wasserstein gradient direction via convex optimization. arXiv preprint DrXiv:2205.13064 2022.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In International Conference on Machine Learning, pp. 681-688, 2011.
Andre Wibisono and Kay Yang. Convergence kl divergence the inexact langevin algorithm with application to score- sore-based generative models. ArXiv, abs/2211.01512 2022.
12

---

A Minimizing Movement Scheme
A.1 Geometric Interpretation
In fact, under some mild conditions of g. the transportation cost W (++) can induce Wasserstein metric and thus Pen (R*) == (1r P(R (g
2x
oo) } is indeed Wasserstein space.
Proposition A.I. Let g(-) go (1l 11) where Rt U (0) + IRt t U {0} satisfies (0) = 0 and || |l can be any norm in IR. Then 96 (Wen (:,-) 3) is a metric on Pe,, (R*) if 90 satisfies: (1) is continuous and strictly increasing; (2) go convex; (3) log go()) a convex function of log
Proof. Suppose T, M.V. E Pa.. is obvious that 9n (We, (v. v)) if and only f v. Besides, 90 (Won on (s))) is symmetric.
the rest part of proof we 9o (Wa, en. (V) 2 96 (Wa, (), By Gluing lemma Villani et al. (2009). we can construct random variables X ~ Y ~ μa, Z ~ such that Y),,, (X, 2) are the optimal coupling of ( and (T, 2) for transportation cost o. respectively. Then we have
show that 90 (We, (e, x))
Y - -1- ))
9o (W (H, v)) S 90
Ego
| YI1 +-X [I-- -11
S 90
Ego
h
||X
Y1l )
| Z
S 90
Ego
+ go
Ego
h
h
(Wo, (1, T)) - 90 (Wo, (TT, v)).
Here the last inequality is due to generalized Minkowski's incquality Mulholland (1949).
The conditions in Proposition A.I are mild and the most common choices of Young function satisfy them, and hence can induce Wasserstein space the generalized Wasserstein gradient flow. Some typical examples of go include |æ/|,, exp(ax") 1, I exp(ar"), while the norm in Rd can be I
|s and on.
A.2 Derivation of Generalized Wasserstein Gradient Flow
Theorem A.2 (Restatement of Theorem D. Given continuously disferentiable Young function
and step size O, define cost function en (x(2))) v) = g( - 1h. Suppose that æ, Hkc E Pen (IR")
y
h
2x 2 < oo). If satisfies assumptions Proposition A.1
{μ P(R") : Eulg( Wasserstein space equipped with Wasserstein metric. Consider MMS under transportation cost Wa,
Pan. (R.)
h
M(K+1)6 = aig DKL(H]IT) + We, (u.)..
(24)
(Rd),(RR)
Denote the optimal transportation map under Wo,, from M((++1) Paa. by Tel.). Then we have
Tu(x)) - - (V log log (())(()))()))
(25)
Proof. By Kantorovich duality (Villani 2021), the optimal transportation cost 33) has equivalent definition:
We(u, v) = sup / pdu + ʃ dd where ( == inf c(x, y) - ((x)).
(26)
Take the functional derivative of the optimization problem 24 and define the optimal in 26 as tb.
The following holds:
DKL (4(8+1)4l1)) + tiA const.
(27)
13

---

Here the reason for have constant instead of zero is that we constrain U(K+ in the space of smooth probability density, Note that ea (x, v)
1
9 2x
2y and pun M(++1)4 E Pey (IR"), then
( +
h
by the fundamental theorem of optimal transportation Villani 2021).
(x)) 2h (y) (x. y), for y Tx(x),
(28)
which implies (z) (Z y) Vg( x y Vg (V)))
ie..
(29)
Combine this equation with
27 and thus the optimal map is given by
T (z)
å DKL (((K+1)al(=))
V-g
h
8p π
VV (V log M(k+1)b
B Details of Motivating Example
We use Example 1 to illustrate the benefits of choosing general Young function g", which is also discussed in Balasubramanian et al. (2022); Wibisono & Yang (2022).
We follow the procedures of Wibisono & Yang 2022). For convenience, let To (- 1), T
3 (mm, 1) and rewrite To - 4 To T. The lower bound of KL divergence follows
from Devroye al. 2018) and Pinsker inequality. In addition, Balasubramanian al. (002) shows
To
V Vlog
V Voo μ
(30)
2TH
Also note that
exp( - 2mm).. Therefore for any q 1, the following bound holds:
π
m²
m 19
E
log
dr
2
1
49-1 m
Tg 1
dr
(3T0 1 18--11
(To + 1
2 49 m²
1122
(1 30 2)0-1 dr + /
(1+ (1+e-90 2) (3 + 2-- dx
π
+5-
(31)
9
dx
exp( 7 m)²
22m(g 1)x)dr
Jz 1 22
22
π
1 1 } 1)m?)
Probb (Z > (2q - 1)m
m
(32)
1
exp (-2 ((q 1)m + 1
2q(a 1)m²
(2g- 1)m - I 25
m
1 1
exp(
m² m
2q + )
1
2gm 22x
82
Here the first inequality is by Prob (2 2
t
2
exp(
for any t > 0 and
t
13
V2T
2
1
(2q 1)m + m 2 2. Similarly, we can prove that
9 3 1
1
1
dx
exp(
m
29 +
2
(33)
4 2qm
2
2
14

---

Plug
32 and
33 in
m²²
3
2e
E
log
m
exp(
2
43
V
π
(9(e
e + 1)-8
1+3c-2
18-1
+
9-1
m²
3
2e
4
m
exp(
2
43
V
π
+ 3e -)
(1 + 62) 3e
0.08
m
m
y exp(
qm
3
2
As for the upper bound,
To dr
log
u 01
29
1 190-1m 1
9
4 m3
dr
(34)
(3mh + क ) 1 (To T
1220 (1/0 9
Ti
48- m
dx
8 dx
9-
250 x0
T dx
exp(
1
(x m)² 2m1g 1)
220 n²
2x20 2m
2
-Probw(o.I) 2 (2g - 1) 1m exp(29(4 - 1)m²²
1
1
1 ((2g 1)m) 2 + -22l0 1)m²
(35)
exp
(2q 1)m
2m
1 1
1
exp(-
-m²))
qm 22x
2
1
Here the first inequality is by Probw(o.1) can prove that
(72 1
exp(
for any > 0. Similarly, we
2m
2
1
1
1
dx
exp( m²)
(36)
#00
To
9
qm
2T
2
Therefore
E log
π
< 4
m
2
1
exp(
mm²)
1
p
qm
/2
2
0.2
m
(m)
exp(
2 )
qm
C Asymptotic Normality of Estimator
In practice, there are finite particles and theoretical analysis is based on the population loss. With this concern, we show that the maximum point of estimation (I4) have good statistical properties. To be specific, the estimator converges to true maximum point with asymptotic normality under mild conditions. Similar properties are also studied in Barp et al. (2019;; Song et al. (2020): Koehler al. (2023).
only get Monte Carlo estimation of (33 But our
Define objective function 2(w,x) V log (x:) (ara))TTa fa
V. [ (x)(
Assumption 4. W is compact and L(-) defined in (13) has unique maximum point w" E inr(W).
Assumption 5. fue (x:), (u) are continuous with w. max (1) V.. (z) Mo(x) for some Mg E L (n))
Assumption 6. There exists a neighborhood N wn and ( < M1(x) for all w E N. Additionally, assume M, L. (ji) and == V²L(w") is non-singular.
w" such that E(-, x) is twice differentiable in N
15

---

n Σ 1=1
Theorem C.I. Given 21,
In i.
11, let Wn == arg min (w)) =
f(w,vi). Under
Assumption we have
Vn(wn -w" 3 (O,
(37)
where E = Bp V wlww", & Vwl(w",')
Proof: Note that under Assumption 5 [e(w. 22] - M(x) for some M (n)) By Newey & McFadden (1986, Lemma 2.4, Theorem 2.1), Wn weakly consistent for w Additionally, since w the maximum point of L and by Central Limit Theorem,
VnVin(w")>
Vol(w" ,.) & A0 E).
Vn
By Assumption uniformly, i.e.
6 and Newey & McFadden (9886 Lemma 2.4), the second order derivative converges
(,)
((w(')] 3 0.
sup
WEN
Finally, the result follows
Newey & McFadden 1986 Theorem 3.1).
D Proof in Section
We first justify our Assumption and then present some crucial lemmas and finish the proof of convergence results. Our proof procedure uses interpolation process of discrete dynamics, following Vempala & Wibisono| (2019): (Balasbbramanian et al. (2022). Informally, the difference between discrete dynamics and continuous dynamics consists of parts: discretization error and estimation error (by neural nets). We bound the discretization error in Lemma [.3 D.4 and the estimation error
Lemma D.5 D.6
D.1 Justification for Assumption and 2
Proof of Proposition It suffices to show that there exists E2 < oo such that for any a, E R, the following inequality holds:
la/a b sgn(a)lal"-1p S + + |al
10|P
ab +
(38)
q
p
If a = 0, then E2 2 is sufficient. Without loss of generality, suppose a = (by replacing b with sgn(a)b/la(<-)) We only need to show that
1 b +
16/P
bb 11 Se
+
(39)
9
p
(1) Case p 2
16-11 1P
2/b - 1|8-2
Let E1 0. Since lim
lim b+1
S 2, there exists 8 > 0 such that
b+11
b+
p p-1
b6 1
when b E 6,1 8).
(39
holds if p 1 2. Also that f(b)
b bbp is
p
continuous function on R
5, + 6) and lim f(b) = < +00. Therefore, f(b bounded on b+00
R|[1- 6, 1 8] and thus holds for finite E2. It's obvious that E2 only depends on p in this case. (2) Case p < 2.
Similarly, let 5 = 2 2/p When b E 1, 8). 39 will trivially hold for any 62 > 0. Also.
(bb is bounded on RR(11 1 + 5) and thus there exists finite e2 determined by p. such that (39 holds.
16

---

D.2
Main Lemmas
Lemma D.I. For any E (kh, (k + 1)h), (uellT)) = -Bus
(Vlog log Elux (Xlux((AANall (X [11
v
Pe
Proof. Let Fas denote the law of X, conditioned the filtration Fks time kh. Then by
Fokker-Planck equation, we have
FefiFFa -div (Me)Fas un (Xah))
Then take expectation of the above equation; by Bayesian formula (Vempala & Wibisono 2019).
Defi div Elga Fen (Xxn)]
div (uylElux (Xua)|Xr = .1)
Hence
Elu (Xun) = (X =1))
0, (Hell.))) -Bps V log
Pr
1
Lemma D.2. Suppose that h <
Under Assumption
for any t [kh, (k + 1)hl,
G2
1,
0
1
Byg
v
log
S (M, - -
Pich
Proof. Note that log
Pe
log det (la + (t Ve)) since x (t kh)ow(v) is an
1
piietattonnprresevinng diffeomorphism under
Then the following holds:
G2
V log
(x)
=
sup
log
(I),²
Mkh
||=|||
High
lim || 8->
sup
log det (Ia + (t - kh) Vux( +2)) log (La + (t- kh) Ve(()(())
sup
lim (1 - kh) Ve (# 62) Ven(x))l2(11 - hG2)-
|| 6-0
s (1
Here the first equation is due to Young's s inequality and the second equation follows the definition of gradient. The inequality in the third line is due Lemma D.13 and the last one is by Assumption With this uniform bound we finish the proof.
Lemma D.3. Under Assumption
3 and the same conditions in Lemma D.2
Eng
Vg (V log -)- Vg
log Pkkh
S ca
Ey,
Vlog
+ 02 (Mp(1 11G2)-1d1)
14
where C₁s 02 are defined as:
(0, 2"
2,
(f)
8-p )
(01, 0z)
P_P min
38-P
2) 2 q
(g 1)2
otherwise.
(+)=
(40)
17

---

Proof. Since (x)
11=1l we have (g) = sgn(x) @ |2/1-- Here o means entry-wise
product. Apply Lemma D. ent4 and thus
Vg" (V log
=) VV
log _)
S c
V log
+ 02 1
V log
Ht
M4
S c1 Es.
Vlog π
+ (Mp - hG22)= -dh))
Mt
The second inequality is due to Lemma D.2
Lemma D.4. Under Assumption 2 and the same conditions in Lemmal D.2
Epp
Vg
log " VV log
π
S Ba (M2(1 - hG2) -1h))
Hish
Proof. Note that g is B-smooth and by Lemma D.2
2
V log |
Eyg
Vg
log ")
Vg"
log
S _
Uihh
2
Pkhh 12
32
(M2 hG22 -1))
Lemma D.5. Suppose that h < min
Under Assumption I for any E Ikh, (k+1)h))
4Gp" G2 S
-11
E,s
Vg" V log
-
E Elux (Kan)
pkk
28 (1 1122)
Ek +
(AG,bi) 1 (4G
(1 + ca 1.
Vlog क
+ + (Mp hG2) -l0h))
1 hD
14 l
where 01, C2 are defined in (40)
Proof. By Jensen's inequality,
Exp
Vg
V log
π
)
B eu (Xun) =
pkh
Blux (Xan) ) |x, = =,]""} (41)
< 2--1
E
Vg" V log
}
Uk
Eps
pichh
P
19
< 2-
E
Vg
V
log
)
Uk
+ IE
ux())) - (Xan)
Under Assumption 3 and Jensen's inequality.
uk (Xe) (X(a))))
s (G,,)² E||X,
X 1p
P
(Gph) EE
S (G.h))4P-1
E. uk (Xen) - vx(
+ rx
+ (x,())) Vg" 100((logg
VVog π
(x,) 1"
μhh
+
E
Vg"
v
log
π (x1)
Vg"
क (xx) VVlog
1'0 + Vg"
(Vlog
log pt
(xx)
it
18

---

Rearrange the above inequality and thus
E (X,)
vx((Nn)
$ (1G,h))) (4G,h)
Eyr
Up Vggg VV'(Voog log
(42)
pinh
Vg" V log _)
Vgg (V log
_
+ E
Vgg
log
jicch
14
Fst.
Again note that
Pr_ Psh
det (Ia + (t -hh)
(x) S (II(Ta (t v s hG2)
sup a
Pkh
Then by Assumption T
P
Ex.
Uk
Vg"
log
S (1
-E uua
Us
Vgg
(V log
π
S
- hG22
-ekk-
Hkkk
Pkh
P
p
(43)
Combining Lemma D.3 with
42
43
and plugging them into 41
we finish the proof.
Lemma D.6. Suppose that h
Under Assumption 2
for any E [kh, (k: + 1)h).
4G2
B
Vg" V log
") Elux (Xnt) =
Pkh
2(1 - hG <
2(4G2h)2B2
2 + (Mz(1 (NG(1-
Ek + - (4G2h)² - (4Gzh)²
V log
2
Proof. The procedure is exactly the same with Lemma D.5 applying Lemma D.4 instead of Lemma D.3 in the last step.
The only difference appears when
1-2-
1-
Lemma D.7. Suppose that h & min
Under Assumption l 3 for any t e
36Gp"
G2
(kh, (I + 1)h),
1
12 Ep,
V
log
π
+
(M, + A22(1 hG2))-aEk -dEr
He
l1g
where A,, Ag are constants only depending on
2p+2
2d) ifg
(A1, A2)
(44)
T02
otherwise.
Here Cg is defined in 40
19

---

Proof. By Lemma D.I and Young's 's inequality, for any \1 d2 > 0,
(.Dl(( = -Bp. p
V log
Vg π
T (V
log
Mt
Mr
+ Eyg.
V Vlog
V V
log
Vg (V
log
H4 π
Mkhh
+ pr
Vlog
Vgg (
log
π
(Nen
Xx+ =
1 x
1
π
3 (11
A
llog
14
+ -
1
V V log
Vg"
V
log
π
p
1
+
1
E
Vgg log
) Elux
XX, = -11
p
Huhh
Then we apply Lemma D.3 and Lemma D.5
O, (.DKI,(ullTT) S
1
1
-1117
p
_xg
1 ~\
28 (AG (1 (t))) Epg
V llg
p
9
p
(AG)1)0
14t.
+
C2
17
p + \2
(4G,h))
20
(M₂(1 - -00---AG)--))) -dhh)"
22 20 (1 hG2
8k-
p
1 (1G,1))))
p
h1
1 28 (4G,h))) (1
q > 2 so that 6 =
< 1, take A1 = cpts A2 = 1. Note that for h S
36G p
1 - (1G
4Gpl
4
1
a) 5 1 4Gph
And thus,
3
6'
0 (Hell()) S
2
1
1
E
V log
19
+
7c2
(M (M,dhi + [1 -hhG)) -dek
3
9
p
Pz
P
p
6p
5
3
E
pL
V
log
π
+
7C2
p
(M + = hh2) - dEM
p
$
1
Epp
V
log
π
+
7C2
(Mpdh)²
+ 3 - hG2) -1hh22) EEx
12
Ht 1
p
- ()) Note that for h <
1 2 (AG,)) (1+a)s
q S 2 so that C = 0, take \₁ A2 =
36Gp' 1 - (4Gph)o
(8G,
12
1
S
Gph
S
And thus,
2(1 - 4Gph)
5
15'
(.eKL( S
1
_
4
(-)
3
a
Gph
En,
V log
9
+
C29
3 (M,dh))
(
2" # (1 hG22)) p p
3
3p
P
9
<
1
4
-
Gph
Epp
V llgg
π
19
2
(Mpdhi² + 2P hG22 dex
3
p
Pt
p
2
Ep.
V
log
π
+
(Mpdh)² + 21 hG²2) -dex-
Mt
p
Therefore, define A,, A2 in (44)
and we finish the proof.
where x = 3 2 1. Under Assumption 2 3 for
Lemma D.S. Suppose that h <
4G2 1212 +1
any E (kh, (k + 1)h),
a, (HellT) S - Ep,
V
Vlo
+ a
(BM2dh))
+ (1- (1- 11G2))-ee
Pt
2
20

---

Proof. Similar to Lemma D.7 under g" is a-strongly convex.
(ull - - Byn
V log Vg
π
14 V
log
+ Ey
V Vlog
V (
log
π
Vg" (V
log
14
14t
14hh
+ E
V Vlog
Vgg v
log
Elex ((ex(Xan))X
Jtt
1ihhh
$ -(
2 X1
x E
V log
2
H4t
1 "
+
2
Vg" log
Vg"
log
jkhh
1
1
\/
E
Vg" (V
log
E E[ox
(X [|, =
2
pkhh
Then we apply Lemma D.4 and Lemma D.6
0, (M s
2
1
2\1 X22 1 (4G2h)²
1
(4G21))3
Epp
V log
Mt 12
2(4G21))2
A2
2(1 - hG2) -d
+
2
+
x2²
1 (4G2h)²
(BM2(1 hG2)) ddh)
+
2
(4G2h)² 8k.
(4G2h)²
1
Take \ = A2 =
V 2
Note that for 4G2h S
V12k² +1
we have
1- (4G2h)²
S
12x²
And
thus,
2
4 (1 (1-hG2) -
O. (KL s
E
Vlog
+
(BM2dh)
+
2
a
D.3
Proof of Main Results
1-2
Theorem D.9. Under Assumption
1
3
for any step size h < min
it holds that
36Gp'
G2
1 Nh Nh Jo
Ep, V
log
π
9 dt 12
(KI(()))TT (M
Ao hG2)
-d
=i=0 Ek
Pt
Nh
N
where Ad, A2 defined in Additionally, if (RO(HO))T)) Ko. then for N
44 are constants that only depend q.
Ko Ko (Go V (gG2)) qAI(M,
we can choose h x
Ko a Λ dG2
The following bound holds:
qAi (M,d)NN
= M,Koo oer
(M
G2Kod +
N-11 Ek
Eanh
log
PNN
Here Ö(-) hides all the constant factors that only depend q.
Proof.
Under Lemma D.7 take integral of both sides from kh to (I: 1)h and obtain
1
(k+1)h
19 di A, (M
DKL. (M(())))))) DKL. (pak ( <
E,
V log
12
14
21

---

Rearranging it and summing from 0 to N -
1 Nh Jo
π j
DKL(PolIT)
Ai (M,dh)) A2 hG2)) -d k Ek
E
V
log
di S 12 Nh
N
(45)
Note that for any convex function g" on RR", (a,b) + 9" (q/b)bis also convex on R x R4. Therefore, 1 + Eng (V log "
is convex in the classical sense on the space of probability measures. And thus
1 Nh
V
log
Eyg
log
dt.
(46)
PNN
JNN
Nh
Jo
Pt
9
We finish the proof by plugging the step size h in 445 and hiding all the constants that only depend on q.
B 21
Theorem D.10. Under Assumption
2
3
for any step size h S
where R =
4G2v12k3 +1
holds that
Epp
V
log
dt
S
6
DKO 3 +
BM + -1 -1hh))
Ek==
k=0
Ek
Nh
J0
He
1
Nh
N
Ko
For simplicity, assume = 1. If and DRL (Holl S Ko, then we can choose h x
A
(KM2d))N
The following bound holds:
dGi2
kG2
=0
G2Kold k)
END NN
Ek
ERNa
log
+
NAN
N
N
Proof.
Under Lemma D.8 take integral of both sides from kh to (A: 1)h and obtain
DKL (p(K - (pRAl) m) / 1)h
(1+tt))
Ey
V log
2
dt
3 + - M
- (1-
-
Pt
Rearranging it and summing from 0 to N -
1
Nh
Exp
log
dt
6
(Ko(H TT)
3 3
+ - hG2))
-d =k= Ek
N-1
Nh
Jo
Pt a
Nh
a
N
The remaining part is similar to the proof of Theorem D.9
D.4 Discussions
The convergence of score divergence only guarantees that the particle distribution gets the local structure of correct (Balasubramanian et al. 2022). To obtain a stronger convergence guarantee, we still need isoperimetry condition of target distribution. We start with L2-GF.
Theorem D.11. If we additionally assume that π satisfies log-Sobole inequality with constant X, i.e.
2 F. EEIVV for all smooth RRd IR.
(f²
then under the same conditions of Theorem D.10 with a == 1, S e, it holds that
XNh.
(3 + 4(1 h.2)) -- (KI(FNN||=)) S expp(= D (Kol(=)
M2,, take h
Ve and then we obtain the guarantee DEL (uNulIT) S
In particular, if S
G2
Myd
-te after
Mad
log (pol(T)
N O
iterations.
Ave
22

---

Remark 1. We match the SOTA rate of LMC under log-Sobole inequality, Hessian smoothness and dissipativity assumption (Mou et al. 2022 The assumption on smoothness of target Hessian is known to accelerate convergence rate (Dalalyan & Karagulyan 2019). But here we do not assume the smoothness of log a explicitly and thus our method can tackle more complex distributions.
Proof. For t E (kh, (k 1)h), we apply Lemma D.8 and log- Sobolev inequality and thus
X
0,Dkı (poll) -
3(M2dh)² - 4(1 - 4(1--GG2)
(47)
By Gronwall's inequality,
D & -A/3 DKL (Kmllll 33x (3(M2dh)² - 4 (1 c-
-h/33).
(48)
Iterating the recursive bound,
DRI(UNAl(T) & ANh
+ 3 (3(M2dhh) + + hG2) de) x
(49)
To further interpret results with GWG under Wp metric, assume that the target distribution sat- isfies modified log-Sobolev inequality, which has been considered many classical works (Adamczak) al. 2017, Bobkoo & Ledoux 2000 Barthe & Roberto 2008
Definition D.1 ( modified log-Sobel rog-Sooleev inequality). For > 1, we say a satisfies the modified log-Sobolev inequality mLSI(q. Ag) if the following holds:
Ent, (1519) <
Ex [[VV]) for all smooth f : Rd + R.
Note that mLSI(2, A2) reduces to the conventional log- Sobolev inequality with constant Xz. As direct corollary of this inequality, for any distribution μu, take f = )
1/9 and thus
PKL(UlIT S
E2
V
log
(50)
quy
Theorem D.12. If we additionally assume that a satisfies Theorem D.9 with EN holds that
50)
then under the same conditions of
(KL( S exp( Nh
DKL(HOlIT) + 12 (A1(Madh) + A2 - h - (gda)=1.
12
Mayo
dMp we take h x qG²
2/4
In particular, if S min
y,,
and then we obtain the guarantee M,d
G2
DKL(NNN S Xx after
N O
Mpd Agel/1
Ag DKLL(
log
iterations.
Here Ō()) hides all the constant factors that only depend on
Proof. For t (kh, (k 1)h), we apply Lemma D.7 and (50) and thus
Qda DKL. (HellT) + (Mpdh)* A2(1 -G2))
(.DrI.(pell)) S
(51)
By Gronwall's inequality,
DKL ()All) S e 12, (KL 117 (Ar( + A2(11 - hG2) (1-e" (1-e- -2,b//12
(52)
Iterating the recursive bound,
(KL(HNAlIT) S exp( NdgNh
DKL + (M,, + A2(1 (e(1))G2) -4e) (gd,)
12
(53)
23

---

Remark 2. mLSI(q. Àg) camnot hold for > 2 as mentioned in Barthe & Roberto (2008). However, we only need (50 to hold for all f Ht to prove Theorem D.12 Plus, Gentil al. (2005 2007)
replace || ll with max {1l 11: ll [1/] and show that mLSI can hold for a class of distributions in this way. We leave this for future work.
Remark 3. Theorem 12 illustrates how the choice of q will influence the convergence rate of PI On one hand, larger q would reduce the complexity dependence on On the other hand, it is generally difficult to predict how Dg will change with q Besides, large would also increase the difficulty to train the neural net and obtain well-estimated direction. Overall, it is challenging to determine the optimal q and thus our adaptive method can present significant advantages.
D.5
Technical Lemmas
Lemma D.13. For any two matrices A, B E Rddd with positive eigenvalues, the following holds:
log det - log B v d A - B| max [ 11
Proof. Suppose that the eigenvalues of real matrix (A - B)B are A1 A1.71.... AR. Av
C, A2x+1, - Da E R. Here Xy is the complex conjugate of Ay. Then it holds that:
log det -log det B det + (A
B)B
= log (1 + 22))
j=1
d Σ
log(1 + A1) + X1) +
log(1 + 12,))
/=1
j=2k+1
log(1 +1y1))
j=1
S d]| (A B) B B-
Similarly, we have log det B log det A < d - A)A ||2 and thus we finish the proof.
Lemma D.14. Define non-negative constants C₁2 C2 as:
(0, 2P-9)
8-P
})
(c1,02)
P min 39 (0-p(_ (1- (g 1)2
otherwise.
Then for any a, b E R, the following inequality holds:
sgn(a)la|011 <gn(b))b/0
9-11111111111
S ca la/0 + -22 -22a b1a
Proof. We shall prove each of the two cases separately.
(I) Case y S 2.
If a, have the same sign, we assume they are positive without loss of generality. Then la?-1_
3 3 & la - 8/9-1, which implies !a?-1 b b < la b/9.
If a, have different signs, we assume > 0, < 0 without loss of generality. Then by Hölder
inequality d--1 (-6)8-1 < 22-9/a a 6/9-1 i.e., aa -- (-b)a-1pp < 2" Fla 6/9.
(2) Case > 2.
If b have different signs, assume 2 0,b < without loss of generality. Then Ja (-6)a-1 1 < la 6/0-1 which implies ja 19-1 1p < la 6/9.

---

If a, have the same sign, we assume they are positive without loss of generality. Note that this inequality is homogeneous, we can let a so that we only need show for any
8a-1
S + c2 6/%
(54)
If 6 S L, then by simple calculus,
(1 - 1--b-1p b 02 - 619 < [11 1)(1 - -]]6 - c2 - -)
S S - 1)
2 ) (
S c1
I <b< (1 (1/d11))
ग. =
4 1 )g TI then 54) is trivial.
If () ut
(B9-1 1)2 - c2(b 1) (b 1)8(( 19
)P c2(b 1)9-01
b --
< 1)y (I(g 21p c2(b 1)
0.
(5)
p-p
The last inequality is due to 02 (q 1)P
2 (q - 1)P
(f) -1
b 6-1
E
Proof of Proposition 5
Proof. Note that g(.) = p || 1p and thus V-g = =1-10-1 C sgn(.). By Young's s inequality,
0, (HellT)) = -By,
V log
π Vg V
log
π
14
11t
+ pc
V
log
4 (V
log 144
T
Se
14
(55)
1 x x k.
|19
< -(1
V log
q
Pr l1g
+ E,,.
π Vgg (V log
ft
where d1 could by any positive scalar. Here set λ = 1 and finish the proof.
Remark 4. (27y is not the only way to adaptively choose optimal In fact. (55) provides wide range of methods based on different choices of X1 and thus we can obtain different objectives for p. We leave it for future work.
F
Additional Details of Experiments
F.1
Gaussian Mixture
We follow the same setting as Dong et al. (2023). The marginal probability of each cluster is 1/10. The number of particles is 1000. For L2-GH PFG and Ada-GWG, we parameterize f. as 3-layer neural networks with lanh activation function. Fach hidden layer has 32 neurons. The inner loop iteration is and we use SGD optimizer with Nesterov momentum (momentum 0.9) to train fu with learning rate 7- 1e-3, The particle step size is 0.1.
25

---

For PFG, following Dong et al. (2023), we set the preconditioning matrix H = H where H is the inverse of diagonal variance of particles and a is 1.0.
For Ada-GWG, we set the initial exponent Po = 2 and learning rate ij = 2.5e-7.
Figure 5 shows some quantitative comparisons between different algorithms. Here Exp-GF represents GWG with g(.) = exp(I ||// (2(")) 1. The results are the averaged after random trials. We can observe that Ada-GWG can obtain accurated samples within fewer iterations.
20-
44-0AA burr
1953
1.2
divergeenc
1.0
3.5
3.0
13
335
202
(a) divergence
(b) Evolution of
Figure 5: Quantitative results in Gaussian Mixture experiment
F.2
Monomial Gamma
On heavy tailed distributions, the number of particles is 1000. For GWG and Ada-GWG, the neural network structure is the same as which in the Gaussian mixture experiment. The inner loop iteration also 5, but we use Adam optimizer with learning rate 1-le-3 to train fw for better stability. The particle step size is 1e-3.
For Ada-GWG, we set the initial exponent Po E (1.5, 2.2, and learning rate 1.
We run the experiment on random seeds. The average results and the variances are represented in the figure using lines and shades.
F.3 Conditioned Diffusion
The procedure to generate the true path is exactly the same as in Detommaso et al. (2018) For PFG and Ada-GWG, we parameterize fw as 3-layer neural networks with tanh nonlincarities Each hidden layer has 200 neurons. The inner loop N' is selected from {1, 5, 10, 15) to get the best performance. fw is pre-trained for 100 iterations before particle updates and we use Adam optimizer with learning rate 1-le-3 to train fw. The particle step size is 3e for Ada-GWG and PFG.
For Ada-GWG, we the initial exponent Po 2.2 and learning rate 0.001. The gradient of A(p) is also clipped within [-0.1,0.] 0.1).
For PFG, we set the preconditioning matrix H = n where H is the inverse of diagonal variance of particles and chosen from {0.1, 0.5, 1.0) to obtain the best performance.
For SVGD, we RBF kernel exp( || vl1² where h is the heuristic bandwidth (Liu & Wang 2016). The initial step size is le-3 and is adjusted by AdaGrad.
Additionally, run LMC with step size le-4 for 10000 iterations the ground truth for posterior distribution.
F.4 Bayesian Neural Networks
Our experiment settings are almost similar to SVGD (Liu & Wang 2016). For the UCI datasets, the datascts randomly partitioned into 90% for training and 10% for testing. Then, further split
26

---

the training dataset by 10% to create a validation set for hyperparameter selection as done in Liu Wang 2016). For L2-GF and Ada-GWG, we parameterize f.r as 3-layer neural networks. Each hidden layer has 300 neurons, and we use LeakyRel.U) as the activation function with negative slope 0.1. The inner loop N' is selected from (1,5, 103. We use the Adam optimizer and choose the learning rate from (0.001, 0.0001} to train fw-
For Ada-GWG, choose the initial exponent Po from {3, 4] and set the learning rate rj 0.0001. The gradient of A()) is clipped within [-0.2, 0.2]. We select the step size of particle updates from {0.0001, 0.0002, 0.0005, 0.001 For SVGD, we use the RBF kernel as done in (Liu & Wang 2016). For SVGD, L2-GF, and Ada-GWG, the iteration number is chosen from {2000, 4000 to converge. For SGLD, the iteration number set to 10000 converge.
G Limitations and Future Work
Estimating Wasserstein gradient by neural networks. Our formulation leverages the capability of neural networks to estimate the generalized Wasserstein gradient. This approach successfully resolves the problem of kernel design for conventional ParVI methods. However, in high dimensional regime, the design of neural network structure is still important but subtle. Besides, the computation cost is also expensive. We expect more efficient algorithms training neural works to approximate Wasserstein gradient, e.g. Wang et al. (002)).
Better adaptive method. Our Ada-GWo method is based the idea of maximizing the decent rate of KL divergence and heavily relies on an accurate estimation of generalized Wasserstein gradient. We update exponent p by simply gradient ascent which may cause severe numerical instability. Although this can be alleviated by clipping, it is still delicate when the target distribution is complex.
General Young function class. In this paper,
only consider the function class
{
which is still limited. We expect a more general function class that both have (p 11 13.
numerical stability and can perfectly capture the information from score function. The characteristics of Young function class may be instructive. How to design an adaptive algorithm on a more general class also challenging and important. We leave this future research.
Shortness of theoretical analysis. Although we provide convergence guarantee under weak as- sumptions, our analysis is still preliminary and we believe these results can be strengthened. There are also other important extensions to consider. For example, our analysis is based on the population loss, which is an asymptotic result based on infinite particles limit. We believe this framework can be also generalized to finite-particle system like SVGD (Korba et al.) 2020). Moreover, we only consider Young functions that have the form of || II or are strongly convex and strongly smooth. We believe that better-designed Young functions may have more advantageous theoretical properties, e.g., Wasserstein Newton flow (Wang Li 2020).
27


Reviews:
[{"review_id": "euRzTp7tKb", "reviewer": "Reviewer_t31F", "summary": "This paper proposes a novel particle-based variational inference framework based on generalized Wasserstein gradient flow of KL divergence, named generalized Wasserstein gradient descent (GWG).  The strong convergence guarantee of the proposed algorithm is provided. The authors also provide an adaptive version based on Wasserstein metric to accelerate convergence.\n", "questions": "- How Assumption 3 on the smoothness of neural networks are ensured in practice? How to estimate the constants G_p and M_p given a neural network?\n", "limitations": "Yes.", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is well-written and clearly-organized.\n- It is quite novel to replace the standard Wasserstein-2 metric by using a Wasserstein metric with cost function defined by Young function.\n- Numerical results support effectiveness of proposed method.\n", "weaknesses": "- The description of main results (Theorem 2 and Theorem 3) are unclear. It is not specified how the vector-valued function vk is learned. Is vk the maximizer of equ (14) with finite samples?  \n- The convergence result only holds for \\bar \\mu_{Nh}, which is the average of path instead of the end-iteration solution \\mu_{Nh}. \n- The numerical results do not show a quanitative convergence improvement of the proposed Ada-GWG with other methods like SVGD, L2-GF or PFG. For instance, It would be better to compare these method in the form of Figure 2, while you plot KL divergence or even maximum mean discrepancy between current iteration and posterior as a function of iteration number or cpu time.\n\nTypos:\n- In Algorithm 1, it should be \\hat A(p_k) instead of \\hat A(p).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your careful review and valuable questions! We address your comments and questions as below.\n\n### Weaknesses\n`W1`: The description of main results (Theorem 2 and Theorem 3) are unclear. It is not specified how the vector-valued function $v\\_k$ is learned. Is $v\\_k$ the maximizer of eq (14) with finite samples?\n\nA1: Exactly! $v\\_k$ is the neural net learned by maximizing equation (14) with finite samples. But here for the theoretical part, we only consider the infinite-particle limit, which is very common in the theoretical analysis of particle-based VI (like [1] [2]). To bridge this gap, we discuss the asymptotic normality of the neural net estimator in Appendix C of the supplementary material. But in general, analyzing the convergence of particle-based VI with finite particles is very difficult, and we leave this for future work as mentioned in Appendix G.\n\n\n[1] Liu, Qiang. “Stein Variational Gradient Descent as Gradient Flow.” NIPS (2017).\\\n[2] Salim, Adil et al. “A Convergence Theory for SVGD in the Population Limit under Talagrand's Inequality T1.” International Conference on Machine Learning (2021).\n\n`W2`: The convergence result only holds for $\\bar{\\mu}\\_{Nh}$, which is the average of path instead of the end-iteration solution $\\mu\\_{Nh}$. \n\nA2: (1) In the literature of non-log-concave sampling, it is generally difficult to attain the convergence rate of the end-iteration solution ([1], [2]). If the target distribution satisfies the log-Sobolev inequality, then we are able to get convergence result of $\\\\mu\\_{Nh}$. Please refer to Appendix D.3 of the supplementary material for more discussions. (2) On the other hand, we can still sample from $\\\\bar{\\\\mu}\\_{Nh}$ by sampling $t\\\\sim \\\\mathrm{Unif}[0,Nh]$ first and then getting a sample from $\\\\mu\\_t$ by $X\\_t = X\\_{kh} + (t-kh)v\\_k(X\\_{kh})$ where $k=[\\\\frac{t}{h}]$. Therefore, we don't think this is a serious drawback. \n\n[1] Balasubramanian, Krishnakumar et al. “Towards a Theory of Non-Log-Concave Sampling: First-Order Stationarity Guarantees for Langevin Monte Carlo.” Annual Conference Computational Learning Theory (2022).\n[2] Korba, Anna et al. “A Non-Asymptotic Analysis for Stein Variational Gradient Descent.” ArXiv abs/2006.09797 (2020): n. pag.\n\n`W3`: The numerical results do not show a quantitative convergence improvement of the proposed Ada-GWG with other methods like SVGD, L2-GF or PFG. For instance, It would be better to compare these method in the form of Figure 2, while you plot KL divergence or even maximum mean discrepancy between current iteration and posterior as a function of iteration number or cpu time.\n\nA3: Thanks for your advice! We plotted the Test RMSE of BNN experiments in Appendix F.4 and it shows that even with unproperly selected $p\\_0$, Ada-GWG makes significant improvements and demonstrates comparable performance to the optimal choice. We further plot the JS-divergence of Gaussian Mixture experiment in the PDF file of global response and will add them in our revision. The results also indicate that Ada-GWG can accelerate convergence compared with the baselines.\n\n\n`W4`: In Algorithm 1, it should be \\hat A(p\\_k) instead of \\hat A(p).\n\n\nA4: Thanks for catching the typos, and we have corrected these typos in our revision.\n\n\n### Questions\n\n`Q1`: How Assumption 3 on the smoothness of neural networks are ensured in practice? How to estimate the constants $G\\_p$ and $M\\_p$ given a neural network?\n\nA1: (1) Since for any $p\\_1\\geq p\\_2>1, \\text{and}\\ x\\in \\\\mathbb{R}^d\\\\backslash\\\\{0\\\\}$, $d^{1/p\\_1-1/p\\_2}\\\\leq\\\\frac{\\\\|x\\\\|\\_{p\\_1}}{\\\\|x\\\\|\\_{p\\_2}}\\leq 1$, the constants $G\\_p$ and $M\\_p$ differ from $G\\_2$ and $M\\_2$ by a factor of $d^{|1/2-1/p|}$ at most. Moreover, $G\\_2$ is a standard assumption on the smoothness of neural nets in the literature of score matching ([1]). And as mentioned in Lines 155-156, $M\\_2$ corresponds to the Lipschitz constant of the Hessian of $\\\\log\\\\pi$ informally, which is also a widely-used assumption to analyze Langevin Monte Carlo ([2]). Therefore, we think Assumption 3 is reasonable to assume. In practice, we can utilize a regularizer (e.g. weight decay) to avoid overfitting and control the smoothness if necessary. (2) To estimate $G\\_p$ and $M\\_p$, we can compute the gradients of neural nets at current particles in each iteration and use the one with largest gradient norm for estimation. We did so in the experiment of conditioned diffusion. The results are shown in the PDF file of global response. The magenitudes of $G\\_2$ and $M\\_2$ remain in a reasonable range, further justifying our Assumption 3.\n\n[1] Lee, Holden, et al. ‘Convergence for Score-Based Generative Modeling with Polynomial Complexity’. Advances in Neural Information Processing Systems, vol. 35, 2022, pp. 22870–22882.\n\n[2] Mou, Wenlong, et al. ‘Improved Bounds for Discretization of Langevin Diffusions: Near-Optimal Rates without Convexity’. Bernoulli, vol. 28, no. 3, Bernoulli Society for Mathematical Statistics and Probability, 2022, pp. 1577–1601.\n"}, {"review_id": "hgj9o588E4", "reviewer": "Reviewer_QMGR", "summary": "This paper proposes a general version to solve Wasserstein Gradient Flow for Particle-based VI. The authors show that this approach offers strong convergence guarantees and better performance over SVGD and other methods, as evidenced by extensive experiments on both simulated and real data sets. The paper also introduces a novel theoretical convergence guarantee of ParVIs with neural-net-estimated vector fields for generalized Wasserstein gradient flow.\n", "questions": "What is a proper heuristic strategy to choose $p$?\n\nIs there any geometric interpretation of the proposed $g$?\n\n", "limitations": "The paper lacks some intuitive demonstration of the propose $g$, which would be very helpful for readers.\n\n", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "GWG is an interesting extension of ParVI, which extend the quadratic distance to more general metrics.\n\nBoth theoretical and empirical justification suggest the improvement of the proposed algorithm.\n", "weaknesses": "The title Particle-based Variational Inference with Generalized Wasserstein Gradient Flow is too close to (Dong et al. 2023). I would suggest to use a new title to highlight the improvements. \n\nMore intuitive illustration of the method would be more helpful.\n\nThe implementation details and codes are missing.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your careful review and valuable questions! We address your comments and questions as below.\n\n#### Weaknesses\n`W1`: The title Particle-based Variational Inference with Generalized Wasserstein Gradient Flow is too close to (Dong et al. 2023).\n\nA1: Thanks for the suggestion! We will modify our title in our revision.\n\n`W2`: More intuitive illustration of the method.\n\nA2: As mentioned in Section 3.2, our motivation and intuition for GWG is to accelerate particle-based VI in terms of faster decay of KL divergence, which is determined by equation (10). We showed in Example 1 that conventional $L_2$-GF may suffer from slow decaying rate of KL. In this sense, it is natural to generalize the formulation of minimizing movement scheme and also the corresponding training method. We will make this points clearer in our revision.\n\n`W3`: The implementation details and codes are missing.\n\nA3: We include all the implementation details in Appendix F of the supplementary material. We are still cleaning up the code and it will be released soon.\n\n\n#### Questions\n`Q1`: What is a proper heuristic strategy to choose $p$?\n\nA1: As claimed in Section 3.2, our primary goal is to achieve faster decay of KL divergence. Hence the criterion to choose $p$ is to enlarge the magnitude of derivatives of KL, i.e. $\\\\mathbb{E}\\_{\\\\mu\\_t} \\\\| \\nabla\\\\log\\\\frac{\\\\pi}{\\\\mu_t}\\\\|_q^q$, where $q=p/(p-1)$. Informally, $p$ should be small (so that $q$ is large) if $\\\\|\\nabla\\\\log\\\\pi - \\\\nabla\\\\log\\\\mu_t \\\\|$ is large and vise versa. The intuition is to apply a larger penalty term when there is a significant difference between the target score and the particle score.\nIn practice, this heuristic strategy is difficult to carry out. Therefore, based on this criterion, we propose Ada-GWG to choose $p$ automatically by increasing an lower bound of $|\\\\partial\\_t{\\\\mathrm{KL}}|$.\n\n`Q2`: Is there any geometric interpretation of the proposed $g$?\n\nA2: Different $g$ corresponds to different Wasserstein metric and Wasserstein space. If $g(\\cdot)=g_0(\\\\|\\cdot\\\\|)$, where $\\\\|\\cdot\\\\|$ can be any norm in the Euclidean space and $g_0$ satisfies some mild assumptions, we can easily show by generalized Minkowski's inequality ([1]) that $g_0^{-1}(W\\_{c\\_h}(\\\\mu,\\\\nu))$ is a well-defined metric and hence $\\mathcal{P}_{c_h}$ defined in Theorem 1 is a Wasserstein space. The general class of functions for $g$ allows us to explore the underlying structure of different probability spaces and further utilize this geometric structure to accelerate convergence.\n\n[1] Mulholland, H. P. ‘On Generalizations of Minkowski’s Inequality in the Form of a Triangle Inequality’. Proceedings of The London Mathematical Society, 1949, pp. 294–307.\n"}, {"review_id": "Jpd9e27f76", "reviewer": "Reviewer_GLuk", "summary": "The paper's primary focus is the use of a generalized Wasserstein gradient flow of the KL divergence for solving the sampling problem in a particle-based variational inference framework, which they name Generalized Wasserstein Gradient Descent (GWG). Unlike the usual Wasserstein gradient flow of the KL, which corresponds to the Langevin diffusion, the implementation of this generalized gradient flow is more complex, hence the authors propose learning the gradient vector field via a neural network. They further back their proposal with a theoretical analysis which indicates that if the neural network learns the vector fields sufficiently well, guarantees can be obtained in a $L^q$ variant of the Fisher information.", "questions": "In the context of existing literature, the reference to Bernton (2018) on line 88 seems to be out of place. Would it not be more accurate to reference the Jordan–Kinderlehrer--Otto (JKO) paper?\n\nLines 170-171 contain an inaccurate claim regarding the dissipativity assumption in Balasubramanian et al. (the dissipativity assumption they use is much weaker than the log-Sobolev inequality) – could this be corrected?\n\nThe paper heavily relies on Assumptions 1 and 2 regarding the accuracy of the neural network's learning of the vector fields. Could the authors shed more light on how realistic these assumptions are and discuss any empirical evidence that might support these assumptions?\n", "limitations": "The fact that the analysis hinges on an assumption for which it is unclear if it holds in practice, namely accurate estimation of the vector field via neural networks, should be addressed more clearly.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper presents a novel sampling algorithm using generalized Wasserstein gradient flows, which provides an interesting avenue for designing alternative approaches to SVGD.\n\nThe algorithm is backed by theory. The authors present robust guarantees about the convergence of the proposed method given certain conditions on the learning ability of the neural network.\n", "weaknesses": "While the approach presented in the paper is interesting, the technical novelty appears to be limited, as the underlying analysis follows the framework presented by Balasubramanian et al.\n\nAssumptions 1 and 2, which concern the neural network's capability to accurately learn vector fields, form a significant portion of the work's theoretical underpinning. However, the realism of these assumptions in practical scenarios is unclear.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your careful review and valuable questions! We address your comments and questions as below.\n\n#### Weaknesses\n\n`W1`: The technical novelty appears to be limited, as the underlying analysis follows the framework presented by Balasubramanian et al.\n\nA1: We indeed borrow some ideas from [1]. However, our analysis is much more involved and exhibits great novelty for analyzing particle-based VI. As described in the beginning of Appendix D, we need to bound both discretization error and estimation error. \n  1. For discretization error, the main technical difficulties include how to ensure the smoothness of particle distribution. Note that we even do not assume the smoothness of target distribution. \n  2. For estimation error, since neural networks can only give an estimation of the optimal vector field, the difficulty lies in how to control the evolution trajectory of particle distribution with inaccurate vector field. \n\nMost importantly, the two error terms are actually entangled with each other and the bound of one term will depend on the other. Therefore, the analysis of both terms needs to be handled with great care to achieve the SOTA rate. \n\nTo the best of our knowledge, our result is the first non-asymptotic analysis of particle-based VI with functional gradient. And it's also for the first time, that one can show that particle-based VI is able to outperform traditional Langevin Monte Carlo theoretically. We believe that our techniques can motivate further theoretical analysis of particle-based VI methods.\n\n\n\n#### Questions\n\n`Q1`: In the context of existing literature, the reference to Bernton (2018) on line 88 seems to be out of place.\n\nA1: Thanks for your advice! We will correct the reference in our revision.\n\n`Q2`: Lines 170-171 contain an inaccurate claim regarding the dissipativity assumption in Balasubramanian et al.\n\nA2: Thanks for your advice! We will correct it in our revision. The assumption made by Balasubramanian et al.[1] is to control the order of growth of the potential function, which is more general than the common dissipativity assumption. This assumption indeed does not imply log-Sobolev inequality. However, it is still quite strong and our analysis does not rely on it. Therefore, our claim still holds that \"particle-based methods can outperform the traditional LMC as long as the neural nets can estimate the vector field accurately\".\n\n`Q3`: Could the authors shed more light on how realistic Assumptions 1 and 2 are and discuss any empirical evidence that might support these assumptions?\n\nA3: We discuss the Assumptions 1 and 2 in two aspects.\n1. **Theoretical justification for Assumption 1 and 2**\nGiven current particle distribution $\\mu$, since we estimate the vector field by maximizing eq (14), we can define the training loss $\\\\mathcal{L}\\_{\\\\text{train}}(v):=\\\\mathbb{E}\\_{\\\\mu} [\\\\langle \\\\nabla\\\\log\\\\frac{\\\\pi}{\\\\mu}, v\\\\rangle - g(v)] $. The maximizer is $v^*=\\\\nabla g^*(\\\\nabla\\\\log\\\\frac{\\\\pi}{\\\\mu})$ and the maximum value is $\\\\mathcal{L}\\_{\\\\text{train}}^*:=\\\\mathcal{L}\\_{\\\\text{train}}(v^*)<\\\\infty$. Similar to Lemma D.14, we can show that for any $p>1$ and any arbitrarily small $\\\\varepsilon\\_1>0$, if $g(\\\\cdot)=\\frac{1}{p}\\\\|\\cdot\\\\|_p^p$, there exists $\\\\varepsilon_2:=\\\\varepsilon_2(\\\\varepsilon\\_1, p)>0$, such that \n    $$\n        \\\\mathbb{E}\\_{\\\\mu} [\\\\|v-\\nabla g^*(\\\\nabla\\\\log\\\\frac{\\\\pi}{\\mu})\\\\|_p^p] \\\\leq \\\\varepsilon\\_1 \\\\mathcal{L}\\_{\\\\text{train}}^* + \\\\varepsilon\\_2 [\\\\mathcal{L}\\_{\\\\text{train}}^*-\\\\mathcal{L}\\_{\\\\text{train}}(v)].\n    $$\n\n    This shows that Assumption 1 is realistic as long as we can optimize the training loss function well (Assumption 2 is similar). We also discuss the asymptotic rate of estimating optimal $v^*$ with finite particles in Appendix C to further justify our Assumption 1 and Assumption 2.\n\n2. **Empirical evidence**\n    Since the score of particle distribution is unknown, we cannot compute the objective in Assumption 1 or 2 directly. However, given the great empirical performance of score-based generative model, we believe that the capability of modern neural networks is sufficient to estimate the optimal vector field. \n\n[1] Balasubramanian, Krishnakumar et al. “Towards a Theory of Non-Log-Concave Sampling: First-Order Stationarity Guarantees for Langevin Monte Carlo.” Annual Conference Computational Learning Theory (2022)."}, {"review_id": "f0xgjy5heU", "reviewer": "Reviewer_Yghj", "summary": "This paper proposes to compute gradient flows of the KL divergence in the space of probability measures endowed by a generalization of the Wasserstein distance, which uses more general cost based on Young functions, and called Generalized Wasserstein gradient flows (GWGF).\n\nAuthors provide the forward Euler scheme of such gradient flow and observe that the choice of the Young function has an impact on the convergence rate of the flow. Thus, besides proposing an algorithm to solve the GWGF with neural networks and analyzing its convergence rate, they also propose an algorithm to adapt the Young function in order to have accelerated convergence rate.\n\nFinally, they propose several applications and compare with other gradient flows based methods such as the SVGD or the preconditioned functional gradient flow. They show that they obtain consistently better results and also illustrate the benefits of the adaptative algorithm.\n", "questions": "The existence of the OT map for the Wasserstein cost is not discussed and is assumed. In [1] (Section 1.3), the result is stated when the cost is of the form $c(x,y)=g(x-y)$ with $g$ strictly convex. In the Young function class, $g$ is only assumed to be convex. So I think I am missing something here.\n\nIn equation (7), I believe that a $\\frac{1}{2h}$ is missing under the Wasserstein cost.\n\nIn Figure 1, the SVGD is not reported. Is it because it is hard to find good hyperparameters which converge for multimodal distributions? I think it would have been nice to add a plot of the evolution of p in the adaptative version of GWG.\n\nEquation (13) is referred to as the Euler-Maruyama scheme. In my opinion, this is the forward-euler discretization and not the Euler-Maruyama scheme which refers to the discretization of SDEs.\n\n[1] Santambrogio, F. (2015). Optimal transport for applied mathematicians. Birkäuser, NY, 55(58-63), 94.", "limitations": "The limitations are addressed in appendix.\n", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The paper is overall pretty clear and the idea to perform gradient flows in this class of OT problems with more general costs is very appealing.\n\n- Interesting idea well motivated\n- Good Theoretical analysis\n- Many applications with comparisons with other methods", "weaknesses": "- It is not a big weakness, but the class of Young functions investigated seems not very big (as it only experimentally considers $c'x,y)=\\|x-y\\|_p^p$ with $p\\in [1.1,4]$). A test with a less common cost could have been interesting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your careful review and valuable questions! We address your comments and questions as below.\n\n### Weaknesses\n`W1`: The class of Young functions investigated seems not very big.\n\nA1: The choices of Young functions are quite delicate since it should strick a balance between accelerating convergence and maintaining numerical stability. Apart from the $\\\\|\\cdot\\\\|_p^p$ and the preconditioned quadratic form proposed in ([1]), we further test other functions like MGF $g(\\\\cdot)=\\exp\\left(\\\\|\\cdot\\\\|_2^2/(2\\\\sigma^2)\\right)-1$ on the 2D-Gaussian mixture model. The results show that this class of Young functions may suffer from numerical instability and poor acceleration performance, and thus is not preferred. For even more general $g$ functions, one interesting future direction is that we can parameterize $g$ as an ICNN [2].\n\n### Questions\n`Q1`: The existence of the OT map for the Wasserstein cost is not discussed and is assumed.\n\nA1: Thanks for your advice! The general convexity is indeed not sufficient to ensure the existence of OT map. To be mathematically rigorous, we will add the \"strictly convex\" requirement in our revision.\n\n`Q2`: In equation (7), I believe that a $\\\\frac{1}{2h}$ is missing under the Wasserstein cost.\n\nA2: We beg to differ. Actually, our definition of cost function is $c_h(x,y):= g(\\\\frac{x-y}{h})h$ in Theorem 1. For sanity check, if $g(\\cdot)=\\frac{1}{2}\\\\|\\cdot\\\\|^2$, then in eq (7), $W\\_{c\\_h}(\\\\mu, \\\\mu\\_{kh})=\\\\frac{1}{2h}W\\_2^2(\\\\mu, \\\\mu\\_{kh})$. This corresponds to the $L\\_2-GF$ ([3]).\n\n`Q3`: In Figure 1, the SVGD is not reported. It would have been nice to add a plot of the evolution of p in the adaptative version of GWG.\n\nA3: We reported the results of SVGD and the plot of the evolution of $p$ in the rebuttal PDF. Although SVGD may be fast at first, it fails to converge to target distribution in 300 iterations. On the other hand, after 50 iterations, the neural network in Ada-GWG can estimate the vector field accurately and exhibits significant acceleration effect. \n\n\n`Q4`: Equation (13) should be referred to as the forward-Euler discretization. \n\nA4: Thanks for your advice, we will correct this mistake in our revision.\n\n[1] Dong, Hanze, et al. ‘Particle-Based Variational Inference with Preconditioned Functional Gradient Flow’. arXiv Preprint arXiv:2211. 13954, 2022.\n\n[2] Amos, Brandon, et al. ‘Input Convex Neural Networks’. International Conference on Machine Learning, PMLR, 2017, pp. 146–155.\n\n[3] Ambrosio, Luigi, et al. Gradient Flows: In Metric Spaces and in the Space of Probability Measures. Springer Science & Business Media, 2005."}]

Author rebuttal:
We thank all reviewers for their constructive feedback, and will modify our paper accordingly in our revision. We address some of the common issues raised by the reviewers below. 

**Justification for Assumption 1 and 2**

1. `Theoretical justification`
 Given current particle distribution $\mu$, since we estimate the vector field by maximizing eq (14), we can define the training loss $\\mathcal{L}\_{\\text{train}}(v):=\\mathbb{E}\_{\\mu} [\\langle \\nabla\\log\\frac{\pi}{\mu}, v\\rangle - g(v)] $. The maximizer is $v^*=\\nabla g^*(\\nabla\\log\\frac{\\pi}{\\mu})$ and the maximum value is $\\mathcal{L}_{\\text{train}}^*:=\\mathcal{L}\_{\\text{train}}(v^*)<\\infty$. Similar to Lemma D.14, we can show that for any $p>1$ and any arbitrarily small $\\varepsilon\_1>0$, if $g(\\cdot)=\\frac{1}{p}\\|\\cdot\\|\_p^p$, there exists $\\varepsilon\_2:=\\varepsilon\_2(\\varepsilon\_1, p)>0$, such that 
    $$
        \\mathbb{E}\_{\\mu} [\\|v-\\nabla g^*(\\nabla\\log\\frac{\\pi}{\\mu})\\|\_p^p] \\leq \\varepsilon\_1 \\mathcal{L}\_{\\text{train}}^* + \\varepsilon\_2 [\\mathcal{L}\_{\\text{train}}^*-\\mathcal{L}\_{\\text{train}}(v)].
    $$

    This shows that Assumption 1 is realistic as long as we can optimize the training loss function well (Assumption 2 is similar). We will add this result in our revision. We also discuss the asymptotic rate of estimating optimal $v^*$ with finite particles in Appendix C to further justify our Assumption 1 and Assumption 2.

2. `Empirical evidence`
 Since the score of particle distribution is unknown, we cannot compute the objective in Assumption 1 or 2 directly. However, given the great empirical performance of score-based generative model, we believe that the capability of modern neural networks is sufficient to optimize the training loss and estimate the optimal vector field.


**Theoretical novelty and highlights**

Our analysis is very involved and exhibits great novelty for analyzing particle-based VI. As described in the beginning of Appendix D, we need to bound both discretization error and estimation error. 
  1. For discretization error, the main technical difficulties include how to ensure the smoothness of particle distribution. Note that we even do not assume the smoothness of target distribution. 
  2. For estimation error, since neural networks can only give an estimation of the optimal vector field, the difficulty lies in how to control the evolution trajectory of particle distribution with inaccurate vector field. 

Most importantly, the two error terms are entangled with each other and the bound of one term will depend on the other. Therefore, the analysis of both terms needs to be handled with great care to achieve the SOTA rate. 

To the best of our knowledge, our result is `the first non-asymptotic analysis of particle-based VI with functional gradient. And it's also for the first time, that one can show that particle-based VI is able to outperform traditional Langevin Monte Carlo theoretically`. We believe that our techniques can motivate further theoretical analysis of particle-based VI methods.


**Quantitative results of improvement**

We plotted the Test RMSE of BNN experiments in Appendix F.4 (also included in the PDF file of global response). It shows that even with improperly selected $p\_0$, Ada-GWG makes significant improvements and demonstrates comparable performance to the optimal choice. This suggests that Ada-GWG is able to push $p\_0$ towards the optimal $p$. We further plot the JS-divergence of Gaussian Mixture experiment in the PDF file and will add them in our revision. The results also indicate that Ada-GWG can accelerate convergence significantly compared with other baselines.

We hope our response has adequately addressed the reviewers' questions and concerns, and look forward to reading any other additional comments.


Meta-review: Throughout the reviews, the reviewers acknowledge the novelty to develop particle-based variational inference methods by gradient flows under generalized metrics, and a non-asymptotic analysis that also takes the estimation error of neural network required in the method into account (though roughly following a recent proof framework). There are also concerns on the reality of the assumptions, the significance of empirical performance, and some imprecise statements, which seem to have been addressed to some extent in the rebuttal. After a few discussions with reviewers, the overall evaluation tends towards outweighing the contributions. I hence recommend an accept.

Nevertheless, apart from implementing the mentioned amendments in rebuttal into the final paper, the authors also need to restrain from overclaims unless clearly analyzed, for example, how it is "the first non-asymptotic analysis of particle-based VI with functional gradient" compared to e.g., Korba et al. (2020), and what is additionally required to "show that particle-based VI is able to outperform traditional Langevin Monte Carlo theoretically". The authors are also expected to explain more on the formulation for derivation, e.g., how the extension of the MMS is a solid mathematical construction that can be interpreted as a gradient flow, whether the employed Wasserstein loss is indeed a distance as required by MMS, and what is the difference from the standard MMS using the p-Wasserstein distance.
Decision: Accept (poster)