Title: Particle-based Variational Inference with Generalized Wasserstein Gradient Flow
Venue: NeurIPS 2023
Accepted: yes
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

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.

1 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 of Bayesian inference is 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-Leibler (KL) divergence to the posterior (Jordan et al., 1999; Wainwright & Jordan, 2008; Blei et al., 2016). Equipped with efficient optimization algorithms, VI allows fast training and easy scaling to large datasets. However, the construction of approximating

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., 1987; Neal, 2011; Welling & Teh, 2011; Chen et al., 2014). While being asymptotically unbiased, MCMCoften takes a long time to converge, and it is also difficult to access the convergence.

Recently, particle based variational inference methods (ParVIs) have been proposed that tend to combine the best of both worlds (Liu & Wang, 2016; Chen et al., 2018; Liu et al., 2019; di Langosco et al., 2021; Fan et al., 2022; Alvarez-Melis et al., 2022). In ParVIs, the approximating distribution is represented as a set of particles, which are iteratively updated by minimizing the KL divergence to 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-efficient than MCMCs. The most well-known particle based VI method is Stein Variational Gradient Descent (SVGD) (Liu &Wang, 2016). It updates the particles by simulating the gradient flows of the KL divergence on a certain kernel related distribution space, where the gradient flows have a tractable form (Liu, 2017; Chewi et al., 2020). However, SVGD relies on the choice of an appropriate kernel function whose design is 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 it 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., 2020; di Langosco et al., 2021; 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 ParVI 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-Fenchel transformation, our method can also be viewed as a functional gradient method with a more general class of regularizers which include the previously used quadratic forms as special cases. We provide a theoretical convergence guarantee of ParVIs with neural-net-estimated 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

Notations. Throughout this paper, we use x to denote particle samples in R d . Let P ( R d ) denote all the probability distributions on R d that are absolute continuous with respect to the Lebesgue measure. We do not distinguish a probabilistic measure with its density function. For x ∈ R d and p > 1 , ∥ x ∥ p := ( | x 1 | p + · · · + | x d | p ) 1 /p stands for the ℓ p -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 R d .

2.1 Particle-based Variational Inference

Let π ∈ P ( R d ) 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 KL divergence

where P ′ ⊆ P ( R d ) is the variational family. Particle-based variational inference methods (ParVIs) is a class of VI methods where P ′ is represented as a set of particles. Assume the current particle distribution is µ , then it holds that

ParVIs aim to find the optimal vector field v that minimizes (2) in certain function class. For example, SVGD (Liu & Wang, 2016) restricts v in the unit ball of an reproducing kernel Hilbert space (RKHS) which has a closed-form solution by kernel trick. Meanwhile, Hu et al. (2018); Grathwohl et al. (2020); di Langosco et al. (2021); Dong et al. (2023) consider a more general class of functions for v , i.e., neural networks, and minimize (2) with some quadratic form regularizers.

2.2 Minimizing Movement Scheme in Wasserstein Space

Assume the cost function c ( · , · ) : R d × R d → R is continuous and bounded from below. Define the optimal transportation cost between two probabilistic measure µ, ν as:

Specifically, if c ( x, y ) = ∥ x -y ∥ p p for some p > 1 , then we get the p -th power of Wasserstein-p distance W p ( µ, ν ) . Jordan et al. (1998) consider a minimizing movement scheme (MMS) under W 2 metric. Given the current distribution µ kh , the distribution for next step is determined by

When the step size h → 0 , { µ kh } k ≥ 0 converges to the solution of the Fokker-Planck equation

Therefore, MMS corresponds to the deterministic dynamics

where µ t is the law of x t . (6) is also known as the gradient flow of KL divergence under W 2 metric, which we refer to as L 2 -GF (Ambrosio et al., 2005). Note that the Langevin dynamics dx t = ∇ log π ( x t ) dt + √ 2 dB t ( B t is the Brownian motion) reproduces the same distribution curve { µ t } t ≥ 0 and thus also corresponds to the Wasserstein gradient flow (Jordan et 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 W 2 metric.

Definition 3.1 (Young function) . A strictly convex function g on R d is called Young function if g ( x ) = g ( -x ) , g (0) = 0 , and for any fixed z ∈ R d \{ 0 } , hg ( z h ) →∞ , as h → 0 .

Theorem 1. Given a continuously differentiable Young function g and step size h > 0 , define cost function c h ( x, y ) = g ( x -y h ) h . Suppose that π, µ kh ∈ P c h ( R d ) := { µ ∈ P ( R d ) : E µ [ g ( 2 x h )] < ∞} . Under some mild conditions of g (see details in Proposition A.1), P c h ( R d ) is a Wasserstein space equipped with Wasserstein distance. Consider MMS under transportation cost W c h :

Denote the optimal transportation map under W c h from µ ( k +1) h to µ kh by T k ( · ) . Then we have

Please refer to Appendix A for full statements and proofs. Informally, µ ( k +1) h ≈ µ kh for small step size h (Santambrogio, 2017). Further note that T k ( x ) -x h is the optimal velocity field associated with the transport from µ ( k +1) h to µ kh (and not vice versa). If step size h → 0 , then following Jordan et al. (1998), we can recover the dynamics in continuous time:

We call (9) the generalized Wasserstein gradient (GWG) flow. If we set g ( · ) = 1 2 ∥ · ∥ 2 2 or any positive definite quadratic form g ( · ) = 1 2 ∥ · ∥ 2 H , then (9) reduces to L 2 -GF (6) or its preconditioned version (Dong et al., 2023) respectively.

3.2 Faster Descent of KL 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 ( · ) = 1 p ∥ · ∥ p p for some p > 1 and then g ∗ ( · ) = 1 q ∥ · ∥ q q . Note that if the particles move along the vector field v t = ∇ g ∗ ( ∇ log π µ t ) , then the descent rate of D KL ( µ t ∥ π ) is

If we choose q such that E µ t ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∇ log π µ t ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∣ q q is large, then D KL ( µ t ∥ π ) decreases faster and the sampling process can be accelerated. We use the following example to further illustrate our idea. Please refer to Appendix B for detailed analysis.

However, the KL divergence between π and µ is large: D KL ( µ ∥ π ) ≥ 1 10 √ 2 .

Suppose the target distribution is π and we run ParVI with current particle distribution µ . We can expect that, if simply using L 2 regularization, i.e. , q = 2 , then for very large m , the score divergence is small and thus the decay of KL divergence is extremely slow. However, D KL ( µ ∥ π ) 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-Euler discretization of the dynamics (9) is

However, since the score of current particle distribution µ kh is generally unknown, we need a method to efficiently estimate the GWG direction ∇ g ∗ ( ∇ log π µ kh ) . Given the distribution of current particles µ , by the definition of convex conjugate, we have

If we parameterize v as a neural network f w with w ∈ W , then we can maximize the following objective with respect to w :

Require: Unnormalized target distribution π , initial particles { x i 0 } n i =1 , initial parameter w 0 , iteration number N,N ′ , particle step size h , parameter step size η

Update w t +1 k = w t k + η ∇ w ̂ L ( w t k ) end for Update w k +1 = w N ′ k Update particles x i k +1 = x i k + hf w k +1 ( x i k ) for i = 1 , · · · , n end for return Particles { x i N } n i =1

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

The exact computation of the divergence term ∇ x · f w ( x ) needs O ( d ) times back-propagation, where d is the dimension of x . In order to reduce computation cost, we refer to Hutchinson's estimator (Hutchinson, 1989), i.e. ,

where ξ i ∈ R d are independent random vectors satisfying E ξ i ξ T i = I d . This is still an unbiased estimator but only needs O (1) times back-propagation.

4 Convergence Analysis without Isoperimetry

In this section, we state our main theoretical results of Algorithm 1. Consider the discrete dynamics:

where v k is the neural-net-estimated GWG at time kh . Define the interpolation process

and let µ t denote the law of X t . Note that here we do not assume isoperimetry of 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 et 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 2. g ∗ ( · ) is α -strongly convex and β -smooth. Define κ := β α . And for any k , E µ kh ∣ ∣ ∣ ∣ ∣ ∣ v k

The two assumptions above ensure the estimation accuracy of neural nets. Note that the preconditioned quadratic form in Dong et al. (2023) is included in Assumption 2. Although the estimation error is not exactly the training objective used in Algorithm 1,the following proposition shows the equivalence between them in some sense.

Proposition 2. Suppose g ( · ) = 1 p ∥· ∥ p p for some p > 1 . Given current particle distribution µ , we can define the training loss L train ( v ) := E µ [ ⟨∇ log π µ , v ⟩ -g ( v )] . The maximizer is v ∗ = ∇ g ∗ ( ∇ log π µ ) and the maximum value is L ∗ train := L train ( v ∗ ) < ∞ . For any arbitrarily small ε 1 > 0 , there exists ε 2 := ε 2 ( ε 1 , p ) < ∞ , such that

Besides, if p ≥ 2 , ε 1 can be 0 while ε 2 is still finite.

Similar results also hold if g satisfies Assumption 2 since it is equivalent to the case when p = 2 . Additionally, we expect some properties of the estimated vector fields.

Assumption 3 (Smoothness of neural nets) . For any k , v k ( · ) is twice differentiable. For any p > 1 , G p := sup x,y ∥ v k ( x ) -v k ( y ) ∥ p ∥ x -y ∥ p < ∞ , M p := sup x,z lim δ → 0 + ∥∇ v k ( x + δz ) -∇ v k ( x ) ∥ op δ ∥ z ∥ p < ∞ .

Note that here we do not assume the smoothness of potential log π explicitly. But informally, G p and M p correspond to the Lipschitz constant of the gradient and the Hessian of log π , respectively.

Theorem 3 (Full version see Theorem D.9) . Under Assumption 1, 3, the following bound holds with proper step size h :

Here ˜ O ( · ) hides all the constant factors that only depend on q .

Theorem 4 (Full version see Theorem D.10) . Under Assumption 2, 3 with α = 1 , the following bound holds with proper step size h :

The proofs in this section are deferred to Appendix D. To interpret our results, suppose g ∗ ( · ) = 1 q ∥·∥ q q and ϵ ≲ ( M p G 2 ) q . If the neural net v k ( · ) can approximate ∇ g ∗ ( ∇ log π µ kh ) accurately ( i.e. , ε k ≲ ϵ ),

then to obtain a probabilistic measure µ such that E µ ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∇ log π µ ∣ ∣ ∣ ∣ ∣ ∣ ∣ ∣ q q ≲ ϵ , the iteration complexity is

˜ O ( M p K 0 dϵ -(1+ 1 q ) ) . If we further let q = p = 2 , the complexity is O ( K 0 dϵ -3 2 ) , which matches the complexity of Langevin Monte Carlo (LMC) under the Hessian smoothness and the growth order assumption (Balasubramanian et al., 2022). However, noticing that Assumption 3 is similar to Hessian smoothness informally, we can obtain this rate without additional assumption on target distribution. This suggests the potential benefits of particle-based methods.

In addition, our formulation allows a wider range of choices of Young function, including ∥ · ∥ p p and the preconditioned quadratic form (Dong et al., 2023). This provides wider options of convergence metrics. We refer the readers to Appendix D.4 for more discussions.

5 Extensions: Adaptive Generalized Wasserstein Gradient Flow

The GWG framework also allows adaption of the Young function g , instead of a fixed one. Similar ideas are also presented in Wang et al. (2018). In this section, we consider a special Young function class { 1 p ∥ · ∥ p p : p > 1 } and propose a procedure that adaptively chooses p to accelerate convergence.

Algorithm 2 Ada-GWG: Adaptive Generalized Wasserstein Gradient Flow

Require: unnormalized target distribution π , initial particles { x i 0 } n i =1 , initial parameter w 0 , iteration number N,N ′ , step size h, η, ˜ η , lower and upper bounds on p : lb, ub

Update w t +1 k = w t k + η ∇ w ̂ L ( w t k )

Update w k +1 = w N ′ k

end for

Update

p

k

+1

=

clip

(

p

k

Update particles

x

i

k

+1

end for

+ ˜

=

return Particles { x i N } n i =1

Consider the continuous time dynamics dx t = f t ( x t ) dt and denote the distribution of particles at time t as µ t , we have the following proposition.

Proposition 5. For g ( · ) = 1 p ∥ · ∥ p p , the derivative of KL divergence has an upper bound:

The proof is in Appendix E. If the neural network f t can approximate the objective well, i.e. , f t ≈ ∇ g ∗ ( ∇ log π µ t ) , then informally we can omit the second term and thus

In order to let KL divergence decrease faster, we can choose p such that A ( p ) is larger. This leads to a 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 p within a reasonable range. We call this adaptive version of GWG, Ada-GWG. The whole training procedure of Ada-GWG is shown in Algorithm 2. Note that (22) can be replaced with Hutchinson's estimator (15) to improve computational efficiency as before.

6 Numerical Experiments

In this section, we compare GWG and Ada-GWG with other ParVI methods including SVGD (Liu & Wang, 2016), L 2 -GF (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 f w to be a neural network with 2 hidden layers and the initial particle distribution is N (0 , I d ) . We refer the readers to Appendix F for more detailed setups of our experiments. The code is available at https://github.com/Alexczh1/GWG .

6.1 Gaussian Mixture

Our first example is on a multi-mode Gaussian mixture distribution. Following Dong et al. (2023), we consider the 10-cluster Gaussian mixture where the variances of the mixture components are

η

∇

x

i

k

̂

A

+

(

p

hf

k

w

, lb, ub

)

(

x

i

+1

k

)

for

)

k

i

= 1

,

· · ·

, n

L2-GF Iter: 50

L-GF Iter: 75

L2-GF Iter: 100

PFG Iter: 25

PFG Iter: 50

PFG Iter: 75

PFG Iter: 100

PFG Iter: 300

2

Ada-GWG Iter: 25

Ada-GWG Iter: 50

Ada-GWG Iter: 75

Ada-GWG Iter: 100

Ada-GWG Iter: 300

all 0.1. The number of particles is 1000. Figure 1 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 L 2 -GF which does not involve the preconditioner, while Ada-GWG with the initial p 0 = 2 significantly accelerates the convergence compared to these two baseline methods. Please refer to appendix for further quantitative comparisons.

6.2 Monomial Gamma

To illustrate the effectiveness and efficiency of the adaptive method compared to the non-adaptive counterparts, we consider the heavy tailed Monomial Gamma distribution where the target π ∝ exp( -0 . 3( | x 1 | 0 . 9 + | x 2 | 0 . 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 of different methods against the number of iterations. The dotted line represents GWG with fixed p , while the solid line represents the corresponding Ada-GWG that starts from the same p at initialization.

We see that the adaptive method outperforms the non-adaptive counterpart consistently. Moreover, Ada-GWG can automatically learn 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 p = 1 . 5 is able to provide much better approximation by automatically adjusting p during runtime. Consequently, Ada-GWG can exhibit greater robustness when determining the initial value of p .

6.3 Conditioned Diffusion

The conditioned diffusion example is a high-dimensional model arising from a Langevin SDE, with state u : [0 , 1] -→ R and dynamics given by

where x = ( x t ) t ≥ 0 is a standard Brownian motion.

This system is commonly used in molecular dynamics to represent the motion of a particle with negligible mass trapped in an energy potential with thermal fluctuations represented by the Brownian forcing Detommaso et al. (2018); Cui et al. (2016). Given the perturbed observations y , the goal

Ada-GWG(po=1.5)

2.00

Ada-GWG(po=2.0)

1.75

Ada-GWG(po=2.2)

GWG(p=1.5)

1.50

GWG(p=2.0)

GWG(p=2.2)

1.25

1.00

0.75

0.50

0.25

0.00

0

2500

5000

75001000012500150001750020000

0

100

200

300

400

500

Iteration

7

6

5

4

3

2

1

log MMD

SVGD

PFG

Ada-GWG

is to infer the posterior of the driving process p ( x | y ) . The forward operator is defined by F ( x ) = ( u t 1 , · · · , u t 20 ) ∈ R 20 , where t i = 0 . 05 i . This is achieved by discretizing the above SDE (23) using an Euler-Maruyama scheme with step size ∆ t = 0 . 01 ; therefore the dimensionality of the problem is 100 . The noisy observations are obtained as y = F ( x true ) + ξ ∈ R 20 , where x true is a Brownian motion path and ξ ∼ N (0 , σ 2 I ) with σ = 0 . 1 . The prior is given by the Brownian motion x = ( x t ) t ≥ 0 .

We test three algorithms: PFG, Ada-GWG, and SVGD, with n = 1000 particles. To obtain the ground truth posterior, we run LMC with 1000 particles in parallel, using a small step size h = 10 -4 for 10000 iterations. Figure 3 reports 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 & Wang (2016), we conduct the two-layer network with 50 hidden units and ReLU activation function, and we use a Gamma(1 , 0 . 1) prior for the inverse covariances. The datasets are 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 1 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 over L 2 -GF. Figure 4 shows the test RMSE against iterations of different methods on the Boston dataset. We can see that for this specific task, setting p = 3 produces better results than when p = 2 . Although L 2 -GF ( i.e. , GWGwith p = 2 ) is sub-optimal, our adaptive method ( i.e. , Ada-GWG with p 0 = 2 ) makes significant improvements and demonstrates comparable performance to the optimal choice of p = 3 . 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.

7 Conclusion

We introduced a new ParVI method, called GWG, which corresponds to a generalized Wasserstein gradient flow of KL divergence. We show that our method has strong convergence guarantees in discrete time setting. We also propose an adaptive version, called Ada-GWG, that can automatically

0

250

500

750

1000

1250

1500

1750

2000

Iteration

2.0

2.5

3.0

3.5

4.0

4.5

5.0

5.5

6.0

Test RMSE

Ada-GWG (

p

0=3.0)

Ada-GWG (

p

0=2.0)

GWG (

p

0=3.0)

GWG (

p

0=2.0)

adjust the Wassertein metric to accelerate convergence. Extensive numerical results showed that Ada-GWG outperforms conventional ParVI methods.

References

Radosław Adamczak, Witold Bednorz, and Paweł Wolff. Moment estimates implied by modified log-sobolev inequalities. ESAIM: Probability and Statistics , 21:467-494, 2017.

David Alvarez-Melis, Yair Schiff, and Youssef Mroueh. Optimizing functionals on the space of probabilities with input convex neural networks. Transactions on Machine Learning Research , 2022. ISSN 2835-8856. URL https://openreview.net/forum?id=dpOYN7o8Jm .

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-concave sampling: First-order stationarity guarantees for langevin monte carlo. In Annual Conference Computational Learning Theory , 2022.

Alessandro Barp, Francois-Xavier 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 , 112:859 - 877, 2016.

Sergey G Bobkov and Michel Ledoux. From brunn-minkowski to brascamp-lieb and to logarithmic sobolev inequalities. Geometric and Functional Analysis , 10:1028-1052, 2000.

Changyou Chen, Ruiyi Zhang, Wenlin Wang, Bai Li, and Liqun Chen. A unified particle-optimization framework for scalable bayesian sampling. ArXiv , abs/1805.11659, 2018.

Tianqi Chen, Emily Fox, and Carlos Guestrin. Stochastic gradient hamiltonian monte carlo. In International Conference on Machine Learning , pp. 1683-1691, 2014.

Sinho Chewi, Thibaut Le Gouic, Chen Lu, Tyler Maunu, and Philippe Rigollet. Svgd as a kernelized wasserstein gradient flow of the chi-squared divergence. Advances in Neural Information Processing Systems , 33, 2020.

Tiangang Cui, Kody JH Law, and Youssef M Marzouk. Dimension-independent likelihood-informed mcmc. Journal of Computational Physics , 304:109-137, 2016.

Arnak S Dalalyan and Avetik Karagulyan. User-friendly guarantees for the langevin monte carlo with inaccurate gradient. Stochastic Processes and their Applications , 129(12):5278-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 highdimensional gaussians. arXiv preprint arXiv:1810.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. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=6OphWWAE3cS .

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-436, 2005.

Ivan Gentil, Arnaud Guillin, and Laurent Miclo. Modified logarithmic sobolev inequalities in null curvature. Revista Matematica Iberoamericana , 23(1):235-258, 2007.

Will Grathwohl, Kuan-Chieh Wang, Jörn-Henrik Jacobsen, David Duvenaud, and Richard Zemel. Learning the stein discrepancy for training and evaluating energy-based models without sampling. In International Conference on Machine Learning , pp. 3732-3747. PMLR, 2020.

Tianyang Hu, Zixiang Chen, Hanxi Sun, Jincheng Bai, Mao Ye, and Guang Cheng. Stein neural sampler. arXiv preprint arXiv:1810.03545 , 2018.

Michael F Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Statistics-Simulation and Computation , 18(3):1059-1076, 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 Koehler, Alexander Heckett, and Andrej Risteski. Statistical efficiency of score matching: The view from isoperimetry. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=TD7AnQjNzR6 .

Anna Korba, Adil Salim, Michael Arbel, Giulia Luise, and Arthur Gretton. A non-asymptotic analysis for stein variational gradient descent. Advances in Neural Information Processing Systems , 33: 4672-4682, 2020.

Chang Liu, Jingwei Zhuo, Pengyu Cheng, Ruiyi Zhang, Jun Zhu, and Lawrence Carin. Understanding and accelerating particle-based variational inference. International Conference on Machine Learning. PMLR , pp. 4082-4092, 2019.

Qiang Liu. Stein variational gradient descent as 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 in neural information processing systems , 29, 2016.

Wenlong Mou, Nicolas Flammarion, Martin J 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 in the form of a triangle inequality. Proceedings of the London mathematical society , 2(1):294-307, 1949.

Radford Neal. MCMC using hamiltonian dynamics. In S Brooks, A 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. google.com/books?id=qfRsAIKZ4rIC .

Whitney Newey and Daniel McFadden. Large sample estimation and hypothesis testing. Handbook of Econometrics , 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. In 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'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019. URL https://proceedings.neurips.cc/paper_files/paper/2019/ file/65a99bb7a3115fdede20da98b08a370f-Paper.pdf .

Cédric Villani. Topics in optimal transportation , volume 58. American Mathematical Soc., 2021.

Cédric Villani et al. Optimal transport: old and new , volume 338. Springer, 2009.

J. Wainwright and M. I. 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 arXiv:2205.13098 , 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 in kl divergence of the inexact langevin algorithm with application to score-based generative models. ArXiv , abs/2211.01512, 2022.

A Minimizing Movement Scheme

A.1 Geometric Interpretation

In fact, under some mild conditions of g , the transportation cost W c h ( · , · ) can induce a Wasserstein metric and thus P c h ( R d ) := { µ ∈ P ( R d ) : E µ [ g ( 2 x h )] < ∞} is indeed a Wasserstein space.

Proposition A.1. Let g ( · ) = g 0 ( ∥ · ∥ ) where g 0 : R + ∪ { 0 } → R + ∪ { 0 } satisfies g 0 (0) = 0 and ∥ · ∥ can be any norm in R d . Then g -1 0 ( W c h ( · , · )) is a metric on P c h ( R d ) if g 0 satisfies: (1) g 0 is continuous and strictly increasing; (2) g 0 is convex; (3) log g 0 ( x ) is a convex function of log x .

Proof. Suppose π, µ, ν ∈ P c h . It is obvious that g -1 0 ( W c h ( µ, ν )) = 0 if and only if µ = ν . Besides, g -1 0 ( W c h ( · , · )) is symmetric. In the rest part of proof we aim to show that g -1 0 ( W c h ( µ, π )) + g -1 0 ( W c h ( ν, π )) ≥ g -1 0 ( W c h ( µ, ν )) . By Gluing lemma Villani et al. (2009), we can construct random variables X ∼ π, Y ∼ µ, Z ∼ ν such that ( X,Y ) , ( X,Z ) are the optimal coupling of ( π, µ ) and ( π, ν ) for transportation cost W c h , respectively. Then we have

Here the last inequality is due to generalized Minkowski's inequality Mulholland (1949).

The conditions in Proposition A.1 are mild and the most common choices of Young function g satisfy them, and hence can induce a Wasserstein space the generalized Wasserstein gradient flow. Some typical examples of g 0 include | x | p , exp( ax 2 ) -1 , x exp( ax b ) , while the norm ∥ · ∥ in R d can be ∥ · ∥ p , ∥ · ∥ H and so on.

A.2 Derivation of Generalized Wasserstein Gradient Flow

Theorem A.2 (Restatement of Theorem 1) . Given a continuously differentiable Young function g and step size h > 0 , define cost function c h ( x, y ) = g ( x -y h ) h . Suppose that π, µ kh ∈ P c h ( R d ) :=

{ µ ∈ P ( R d ) : E µ [ g ( 2 x h )] < ∞} . If g satisfies assumptions in Proposition A.1, P c h ( R d ) is a Wasserstein space equipped with Wasserstein metric. Consider MMS under transportation cost W c h :

Denote the optimal transportation map under W c h from µ ( k +1) h to µ kh by T k ( · ) . Then we have

Proof. By Kantorovich duality (Villani, 2021), the optimal transportation cost (3) has an equivalent definition:

Take the functional derivative of the optimization problem (24) and define the optimal φ in (26) as ψ . The following holds:

Here the reason for have a constant instead of zero is that we constrain µ ( k +1) h in the space of smooth probability density. Note that c h ( x, y ) ≤ 1 2 ( g ( 2 x h ) + g ( 2 y h ) ) and µ kh , µ ( k +1) h ∈ P c h ( R d ) , then by the fundamental theorem of optimal transportation (Villani, 2021),

Combine this equation with (27) and thus the optimal map is given by

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 π 0 = N ( -m, 1) , π 1 = N ( m, 1) and rewrite π = 1 2 π 0 + 1 2 π 1 , µ = 3 4 π 0 + 1 4 π 1 . The lower bound of KL divergence follows from Devroye et al. (2018) and Pinsker inequality . In addition, Balasubramanian et al. (2022) shows

Here the first inequality is by Prob N (0 , 1) { Z ≥ t } ≥ ( 1 t -1 t 3 ) 1 √ 2 π exp( -t 2 2 ) for any t > 0 and

Plug (32) and (33) in (31),

As for the upper bound,

Here the first inequality is by Prob N (0 , 1) { Z ≥ t } ≤ 1 t 1 √ 2 π exp( -t 2 2 ) for any t > 0 . Similarly, we can prove that

Therefore

C Asymptotic Normality of Estimator

In practice, there are finite particles and we can only get a Monte Carlo estimation of (13). But our theoretical analysis is based on the population loss. With this concern, we show that the maximum point of estimation (14) 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 et al. (2023).

Define objective function ℓ ( w,x ) := ∇ log π ( x ) T f w ( x ) + ∇ x · f w ( x ) -g ( f w ( x )) .

Assumption 4. W is compact and L ( · ) defined in (13) has a unique maximum point w ∗ ∈ int ( W ) . Assumption 5. f w ( x ) , ∇ x f w ( x ) are continuous with w . max {∥ f w ( x ) ∥ , |∇ x · f w ( x ) | , g ( f w ( x )) } ≤ M 0 ( x ) for some M 0 ∈ L 1 ( µ ) .

Assumption 6. There exists a neighborhood N of w ∗ such that ℓ ( · , x ) is twice differentiable in N and ∥∇ 2 ℓ ( w,x ) ∥ ≤ M 1 ( x ) for all w ∈ N . Additionally, assume M 1 ∈ L 1 ( µ ) and H := ∇ 2 L ( w ∗ ) is non-singular.

(1) Case p ≥ 2 .

when b ∈ [1 -δ, 1 + δ ] , (39) holds if ε 2 ≥ p + 1 > 2 . Also note that f ( b ) = | b -1 | 1 q -b + | b | p p is a continuous function on R \ (1 -δ, 1 + δ ) and lim b →∞ f ( b ) = p < + ∞ . Therefore, f ( b ) is bounded on R \ [1 -δ, 1 + δ ] and thus (39) holds for finite ε 2 . It's obvious that ε 2 only depends on p in this case. (2) Case p < 2 .

Similarly, let δ = ε 1 /p 1 . When b ∈ [1 -δ, 1 + δ ] , (39) will trivially hold for any ε 2 > 0 . Also, f ( b ) is bounded on R \ [1 -δ, 1 + δ ] and thus there exists finite ε 2 determined by p, δ such that (39) holds.

Theorem C.1. Given x 1 , · · · , x n i . i . d . ∼ µ , let ̂ w n := arg min w ̂ L n ( w ) := 1 n ∑ i =1 ℓ ( w,x i ) . Under Assumption 4-6, we have √ p

Proof. Note that under Assumption 5, | ℓ ( w,x ) | ≤ M ( x ) for some M ∈ L 1 ( µ ) . By Newey & McFadden (1986, Lemma 2.4, Theorem 2.1), ̂ w n is weakly consistent for w ∗ . Additionally, since w ∗ is the maximum point of L and by Central Limit Theorem,

By Assumption 6 and Newey & McFadden (1986, Lemma 2.4), the second order derivative converges uniformly, i.e.

Finally, the result follows Newey & McFadden (1986, Theorem 3.1).

D Proof in Section 4

We first justify our Assumption 1 and 2, 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); Balasubramanian et al. (2022). Informally, the difference between discrete dynamics and continuous dynamics consists of two parts: discretization error and estimation error (by neural nets). We bound the discretization error in Lemma D.3, D.4 and the estimation error in Lemma D.5, D.6.

D.1 Justification for Assumption 1 and 2

Proof of Proposition 2. It suffices to show that there exists ε 2 < ∞ such that for any a, b ∈ R , the following inequality holds:

If a = 0 , then ε 2 ≥ p is sufficient. Without loss of generality, suppose a = 1 (by replacing b with sgn ( a ) b/ | a | q -1 ). We only need to show that

D.2 Main Lemmas

Proof. Let µ t |F kh denote the law of X t conditioned on the filtration F kh at time kh . Then by Fokker-Planck equation, we have

Then we take expectation of the above equation; by Bayesian formula (Vempala & Wibisono, 2019),

Hence

Lemma D.2. Suppose that h < 1 G 2 . Under Assumption 3, for any t ∈ [ kh, ( k + 1) h ] , q >

,

Proof. Note that log µ t µ kh = -log det ( I d +( t -kh ) ∇ v k ) since x ↦→ x + ( t -kh ) v k ( x ) is an orientation-preserving diffeomorphism under h < G 2

1 . Then the following holds:

Here the first equation is due to Young's inequality and the second equation follows the definition of gradient. The inequality in the third line is due to Lemma D.13 and the last one is by Assumption 3. With this uniform bound we finish the proof.

Lemma D.3. Under Assumption 1, 3, and the same conditions in Lemma D.2,

where c 1 , c 2 are defined as:

Proof. Since g ∗ ( x ) = 1 q ∥ x ∥ q q , we have ∇ g ∗ ( x ) = sgn ( x ) ⊙ | x | q -1 . Here ⊙ means entry-wise product. Apply Lemma D.14 entry-wise and thus

The second inequality is due to Lemma D.2.

Lemma D.4. Under Assumption 2, 3 and the same conditions in Lemma D.2,

Proof. Note that g ∗ is β -smooth and by Lemma D.2,

LemmaD.5. Suppose that h < min { 1 4 G p , 1 G 2 } . Under Assumption 1, 3, for any t ∈ [ kh, ( k +1) h ] ,

where c 1 , c 2 are defined in (40).

Proof. By Jensen's inequality,

Under Assumption 3 and Jensen's inequality,

Rearrange the above inequality and thus

Again note that µ t µ kh = det ( I d +( t -kh ) ∇ v k ) -1 , so

Then by Assumption 1,

Combining Lemma D.3 with (42),(43) and plugging them into (41), we finish the proof.

Lemma D.6. Suppose that h < 1 4 G 2 . Under Assumption 2, 3, for any t ∈ [ kh, ( k +1) h ] ,

Proof. The procedure is exactly the same with Lemma D.5. The only difference appears when applying Lemma D.4 instead of Lemma D.3 in the last step.

Lemma D.7. Suppose that h ≤ min { 1 36 G p , 1 -2 -1 q G 2 } . Under Assumption 1, 3, for any t ∈ ( kh, ( k +1) h ) ,

where A 1 , A 2 are constants only depending on q :

Here c 2 is defined in (40).

Proof. By Lemma D.1 and Young's inequality, for any λ 1 , λ 2 > 0 ,

Then we apply Lemma D.3 and Lemma D.5,

If q > 2 so that c 1 = 3 -p < 1 , take λ 1 = c 1 p + q 1 , λ 2 = 1 . Note that for h ≤ 1 36 G p , 2 p -1 (4 G p h ) p 1 -(4 G p h ) p (1+

c 1 ) ≤ 4 G p h 1 -4 G p h · 4 3 ≤ 1 6 . And thus,

If q ≤ 2 so that c 1 = 0 , take λ 1 = λ 2 = ( q 3 ) 1 q . Note that for h ≤ 1 36 G p , 2 p -1 (4 G p h ) p 1 -(4 G p h ) p (1 + c 1 ) ≤

(8 G p h ) 2 2(1 -4 G p h ) ≤ 12 5 G p h ≤ 1 15 . And thus,

Therefore, define A 1 , A 2 as in (44) and we finish the proof.

Lemma D.8. Suppose that h ≤ 1 4 G 2 √ 12 κ 2 +1 , where κ := β α ≥ 1 . Under Assumption 2, 3, for any t ∈ ( kh, ( k +1) h ) ,

Proof. Similar to Lemma D.7, under g ∗ is α -strongly convex,

Then we apply Lemma D.4 and Lemma D.6,

Take λ 1 = λ 2 = √ α 2 . Note that for 4 G 2 h ≤ 1 √ 12 κ 2 +1 , we have (4 G 2 h ) 2 1 -(4 G 2 h ) 2 ≤ 1 12 κ 2 . And thus,

D.3 Proof of Main Results

Theorem D.9. Under Assumption 1, 3, for any step size h ≤ min { 1 36 G p , 1 -2 -1 q G 2 } , it holds that

where A 1 , A 2 defined in (44) are constants that only depend on q .

Additionally, if D KL ( µ 0 ∥ π ) ≤ K 0 , then for N ≳ K 0 ( G p ∨ ( qG 2 )) q +1 qA 1 ( M p d ) q , we can choose h ≍

( K 0 qA 1 ( M p d ) q N ) 1 q +1 ∧ 1 dG 2

. The following bound holds:

Here ˜ O ( · ) hides all the constant factors that only depend on q .

Proof. Under Lemma D.7, take integral of both sides from kh to ( k +1) h and we obtain

Rearranging it and summing from 0 to N -1 ,

Note that for any convex function g ∗ on R d , ( a, b ) ↦→ g ∗ ( a/b ) b is also convex on R d × R + . Therefore, µ ↦→ E µ g ∗ ( ∇ log π µ ) is convex in the classical sense on the space of probability measures. And thus

We finish the proof by plugging the step size h in (45) and hiding all the constants that only depend on q .

Theorem D.10. Under Assumption 2, 3, for any step size h ≤ 1 4 G 2 √ 12 κ 2 +1 , where κ := β α ≥ 1 , it holds that

For simplicity, assume α = 1 . If and D KL ( µ 0 ∥ π ) ≤ K 0 , then we can choose h ≍ ( K 0 ( κM 2 d ) 2 N ) 1 3 ∧

1 dG 2 ∧ 1 κG 2 . The following bound holds:

Proof. Under Lemma D.8, take integral of both sides from kh to ( k +1) h and we obtain

Rearranging it and summing from 0 to N -1 ,

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

Theorem D.11. If we additionally assume that π satisfies log-Sobolev inequality with constant λ , i.e.

then under the same conditions of Theorem D.10 with α = β = 1 , ε k ≤ ϵ , it holds that

In particular, if ϵ ≲ ( M 2 G 2 ) 2 , we take h ≍ √ ϵ M 2 d and then we obtain the guarantee D KL ( µ Nh ∥ π ) ≲ λ -1 ϵ after

Remark 1. We match the SOTA rate of LMC under log-Sobolev 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 π explicitly and thus our method can tackle more complex distributions.

Proof. For t ∈ ( kh, ( k +1) h ) , we apply Lemma D.8 and log-Sobolev inequality and thus

By Gronwall's inequality,

Iterating the recursive bound,

To further interpret our results with GWG under W p metric, we assume that the target distribution satisfies modified log-Sobolev inequality , which has been considered in many classical works (Adamczak et al., 2017; Bobkov & Ledoux, 2000; Barthe & Roberto, 2008).

Definition D.1 ( modified log-Sobolev inequality ) . For q > 1 , we say π satisfies the modified log-Sobolev inequality mLSI( q, λ q ) if the following holds:

Note that mLSI( 2 , λ 2 ) reduces to the conventional log-Sobolev inequality with constant λ 2 . As a direct corollary of this inequality, for any distribution µ , we take f = ( µ π ) 1 /q and thus

Theorem D.12. If we additionally assume that π satisfies (50), then under the same conditions of Theorem D.9 with ε k ≤ ϵ , it holds that

In particular, if ϵ ≲ min { ( M p G 2 ) q , ( dM p qG 2 ) q } , we take h ≍ ϵ 1 /q M p d and then we obtain the guarantee

D KL ( µ Nh ∥ π ) ≲ λ -1 q ϵ after

Here ˜ O ( · ) hides all the constant factors that only depend on q .

Proof. For t ∈ ( kh, ( k +1) h ) , we apply Lemma D.7 and (50) and thus

By Gronwall's inequality,

Iterating the recursive bound,

Remark 2. mLSI( q, λ q ) cannot hold for q > 2 as mentioned in Barthe & Roberto (2008). However, we only need (50) to hold for all µ = µ t to prove Theorem D.12. Plus, Gentil et al. (2005, 2007) replace ∥ · ∥ q q with max {∥ · ∥ 2 2 , ∥ · ∥ q q } and show that mLSI can hold for a class of distributions in this way. We leave this for future work.

Remark 3. Theorem D.12 illustrates how the choice of q will influence the convergence rate of ParVI. On one hand, larger q would reduce the complexity dependence on ϵ . On the other hand, it is generally difficult to predict how λ q will change with q . Besides, large q would also increase the difficulty to train the neural net and obtain a 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 ∈ R d × d with positive eigenvalues, the following holds:

Proof. Suppose that the eigenvalues of real matrix ( A -B ) B -1 are λ 1 , λ 1 , · · · , λ k , λ k ∈ C , λ 2 k +1 , · · · λ d ∈ R . Here λ j is the complex conjugate of λ j . Then it holds that:

Similarly, we have log det B -log det A ≤ d ∥ ( B -A ) A -1 ∥ 2 and thus we finish the proof.

Lemma D.14. Define non-negative constants c 1 , c 2 as:

Then for any a, b ∈ R , the following inequality holds:

Proof. We shall prove each of the two cases separately.

(1) Case q ≤ 2 .

If a, b have the same sign, we assume they are positive without loss of generality. Then | a q -1 -b q -1 | ≤ | a -b | q -1 , which implies | a q -1 -b q -1 | p ≤ | a -b | q .

If a, b have different signs, we assume a ≥ 0 , b < 0 without loss of generality. Then by Hölder inequality a q -1 +( -b ) q -1 ≤ 2 2 -q | a -b | q -1 , i.e. , | a q -1 -( -b ) q -1 | p ≤ 2 p -q | a -b | q .

(2) Case q > 2 .

If a, b have different signs, we assume a ≥ 0 , b < 0 without loss of generality. Then | a q -1 + ( -b ) q -1 | ≤ | a -b | q -1 , which implies | a q -1 -b q -1 | p ≤ | a -b | q .

If a, b have the same sign, we assume they are positive without loss of generality. Note that this inequality is homogeneous, we can let a = 1 so that we only need to show for any b > 0 ,

If b ≤ 1 , then by simple calculus,

If 1 < b ≤ ( 1 + c 1 /p 1 ) 1 q -1 = ( 4 3 ) 1 q -1 , then (54) is trivial.

E Proof of Proposition 5

Proof. Note that g ( · ) = 1 p ∥ · ∥ p p and thus ∇ g ∗ ( · ) = | · | q -1 ⊙ sgn ( · ) . By Young's inequality,

where λ 1 could by any positive scalar. Here we set λ 1 = 1 and finish the proof.

Remark 4. (21) is not the only way to adaptively choose optimal p . In fact, (55) provides a wide range of methods based on different choices of λ 1 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 L 2 -GF, PFG and Ada-GWG, we parameterize f w as 3-layer neural networks with tanh activation function. Each hidden layer has 32 neurons. The inner loop iteration is 5 and we use SGD optimizer with Nesterov momentum (momentum 0.9) to train f w with learning rate η =1e-3. The particle step size is 0 . 1 .

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 α is 1 . 0 .

For Ada-GWG, we set the initial exponent p 0 = 2 and learning rate ˜ η = 2 . 5 e-7.

Figure 5 shows some quantitative comparisons between different algorithms. Here Exp-GF represents GWGwith g ( · ) = exp( ∥ · ∥ 2 2 / (2 σ 2 )) -1 . The results are the averaged after 10 random trials. We can observe that Ada-GWG can obtain highly-accurated samples within fewer iterations.

0

50

100

150

200

250

300

Iteration

0.0

0.5

1.0

1.5

2.0

JS-divergence

Ada-GWG (

p

0=2.0)

L2-GF

Exp-GF

PFG

SVGD

(a) JS divergence

0

50

100

150

200

250

300

Iteration

1.3

1.4

1.5

1.6

1.7

1.8

1.9

2.0

p

Ada-GWG (

p

0=2.0)

(b) Evolution of

p

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 is also 5, but we use Adam optimizer with learning rate η =1e-3 to train f w for better stability. The particle step size is 1e-3.

For Ada-GWG, we set the initial exponent p 0 ∈ { 1 . 5 , 2 . 0 , 2 . 2 } and learning rate ˜ η = 1 .

We run the experiment on 4 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 f w as 3-layer neural networks with tanh nonlinearities. Each hidden layer has 200 neurons. The inner loop N ′ is selected from { 1 , 5 , 10 , 15 } to get the best performance. f w is pre-trained for 100 iterations before particle updates and we use Adam optimizer with learning rate η =1e-3 to train f w . The particle step size is 3 e -3 for Ada-GWG and PFG.

For Ada-GWG, we set the initial exponent p 0 = 2 . 2 and learning rate ˜ η = 0 . 001 . The gradient of A ( p ) is also clipped within [ -0 . 1 , 0 . 1] .

For PFG, we set the preconditioning matrix H = ˆ H α , where ˆ H is the inverse of diagonal variance of particles and α is chosen from { 0 . 1 , 0 . 5 , 1 . 0 } to obtain the best performance.

For SVGD, we use RBF kernel exp( - ∥ x -y ∥ 2 h ) where h is the heuristic bandwidth (Liu & Wang, 2016). The initial step size is 1e-3 and is adjusted by AdaGrad.

Additionally, we run LMC with step size 1e-4 for 10000 iterations as 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 datasets are randomly partitioned into 90% for training and 10% for testing. Then, we further split

the training dataset by 10% to create a validation set for hyperparameter selection as done in (Liu &Wang, 2016). For L 2 -GF and Ada-GWG, we parameterize f w as 3-layer neural networks. Each hidden layer has 300 neurons, and we use LeakyReLU as the activation function with a negative slope of 0.1. The inner loop N ′ is selected from { 1 , 5 , 10 } . We use the Adam optimizer and choose the learning rate from { 0 . 001 , 0 . 0001 } to train f w .

For Ada-GWG, we choose the initial exponent p 0 from { 3 , 4 } and set the learning rate ˜ η = 0 . 0001 . The gradient of A ( p ) 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, L 2 -GF, and Ada-GWG, the iteration number is chosen from { 2000 , 4000 } to converge. For SGLD, the iteration number is set to 10000 to 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 on training neural works to approximate Wasserstein gradient, e.g. Wang et al. (2022).

Better adaptive method. Our Ada-GWG method is based on 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, we only consider the function class { 1 p ∥ · ∥ p p : p > 1 } , which is still limited. We expect a more general function class that both have 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 is also challenging and important. We leave this to future research.

Shortness of theoretical analysis. Although we provide convergence guarantee under weak assumptions, 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 ∥ · ∥ p p 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).

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\nThe 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\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nThe 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.\n\nLIMITATIONS:\nThe limitations are addressed in appendix.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-18"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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\nWEAKNESSES:\nWhile 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\nQUESTIONS:\nIn 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\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-19"}

comments: 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.

STRENGTHS:
GWG is an interesting extension of ParVI, which extend the quadratic distance to more general metrics.

Both theoretical and empirical justification suggest the improvement of the proposed algorithm.

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. 

More intuitive illustration of the method would be more helpful.

The implementation details and codes are missing.

QUESTIONS:
What is a proper heuristic strategy to choose $p$?

Is there any geometric interpretation of the proposed $g$?

LIMITATIONS:
The paper lacks some intuitive demonstration of the propose $g$, which would be very helpful for readers.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05

comments: 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.

STRENGTHS:
- The paper is well-written and clearly-organized.
- It is quite novel to replace the standard Wasserstein-2 metric by using a Wasserstein metric with cost function defined by Young function.
- Numerical results support effectiveness of proposed method.

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?  
- The convergence result only holds for \bar \mu_{Nh}, which is the average of path instead of the end-iteration solution \mu_{Nh}. 
- 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.

Typos:
- In Algorithm 1, it should be \hat A(p_k) instead of \hat A(p).

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?

LIMITATIONS:
Yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06