SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ziheng Cheng*
alex-czh@stu.pku.edu.cn
Shiyue Zhang*
zhangshiyue@stu.pku.edu.cn
Longlin Yu
llyu@pku.edu.cn
Cheng Zhang†
School of Mathematical Sciences and Center for Statistical Science
chengzhang@math.pku.edu.cn
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.
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
*Contributed equally to this work.
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, MCMC often 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.
Notations. Throughout this paper, we use  $x$  to denote particle samples in  $\mathbb{R}^d$ . Let  $\mathcal{P}(\mathbb{R}^d)$  denote all the probability distributions on  $\mathbb{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 \in \mathbb{R}^d$  and  $p > 1$ ,  $\| x\|_p := (|x_1|^p + \dots + |x_d|^p)^{1/p}$  stands for the  $\ell_p$ -norm. The Hölder conjugate of  $p$  is denoted by  $q := p / (p - 1)$ . Notation  $g^*(\cdot)$  denotes the Legendre transform of a convex function  $g(\cdot)$  on  $\mathbb{R}^d$ .
2.1 Particle-based Variational Inference
Let  $\pi \in \mathcal{P}(\mathbb{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  $\mu^{*}$  that minimizes the KL divergence
$$
\mu^ {*} := \arg \min  _ {\mu \in \mathcal {P} ^ {\prime}} D _ {\mathrm {K L}} (\mu \| \pi), \tag {1}
$$
where  $\mathcal{P}' \subseteq \mathcal{P}(\mathbb{R}^d)$  is the variational family. Particle-based variational inference methods (ParVIs) is a class of VI methods where  $\mathcal{P}'$  is represented as a set of particles. Assume the current particle distribution is  $\mu$ , then it holds that
$$
\left. \frac {d}{d \epsilon} \right| _ {\epsilon = 0} D _ {\mathrm {K L}} ((i d + \epsilon v) _ {\#} \mu \| \pi) = - \mathbb {E} _ {\mu} \langle \nabla \log \frac {\pi}{\mu}, v \rangle . \tag {2}
$$
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(\cdot, \cdot): \mathbb{R}^d \times \mathbb{R}^d \to \mathbb{R}$  is continuous and bounded from below. Define the optimal transportation cost between two probabilistic measure  $\mu, \nu$  as:
$$
W _ {c} (\mu , \nu) := \inf  _ {\rho \in \Pi (\mu , \nu)} \int c (x, y) d \rho . \tag {3}
$$
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}(\mu ,\nu)$ . Jordan et al. (1998) consider a minimizing movement scheme (MMS) under  $W_{2}$  metric. Given the current distribution  $\mu_{kh}$ , the distribution for next step is determined by
$$
\mu_ {(k + 1) h} := \underset {\mu \in \mathcal {P} _ {2} (\mathbb {R} ^ {d})} {\arg \min } D _ {\mathrm {K L}} (\mu \| \pi) + \frac {1}{2 h} W _ {2} ^ {2} (\mu , \mu_ {k h}). \tag {4}
$$
When the step size  $h\to 0$ $\{\mu_{kh}\}_{k\geq 0}$  converges to the solution of the Fokker-Planck equation
$$
\partial_ {t} \mu_ {t} + \operatorname {d i v} \left(\mu_ {t} \nabla \log \pi\right) = \Delta \mu_ {t}. \tag {5}
$$
Therefore, MMS corresponds to the deterministic dynamics
$$
d x _ {t} = v _ {t} d t, v _ {t} = \nabla \log \pi - \nabla \log \mu_ {t}, \tag {6}
$$
where  $\mu_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 = \nabla \log \pi(x_t) dt + \sqrt{2} dB_t$  ( $B_t$  is the Brownian motion) reproduces the same distribution curve  $\{\mu_t\}_{t \geq 0}$  and thus also corresponds to the Wasserstein gradient flow (Jordan et al., 1998).
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  $\mathbb{R}^d$  is called Young function if  $g(x) = g(-x), g(0) = 0$ , and for any fixed  $z \in \mathbb{R}^d \backslash \{0\}$ ,  $hg\left(\frac{z}{h}\right) \to \infty$ , as  $h \to 0$ .
Theorem 1. Given a continuously differentiable Young function  $g$  and step size  $h > 0$ , define cost function  $c_{h}(x,y) = g(\frac{x - y}{h})h$ . Suppose that  $\pi, \mu_{kh} \in \mathcal{P}_{c_h}(\mathbb{R}^d) := \{\mu \in \mathcal{P}(\mathbb{R}^d) : \mathbb{E}_\mu [g(\frac{2x}{h})] < \infty\}$ . Under some mild conditions of  $g$  (see details in Proposition A.1),  $\mathcal{P}_{c_h}(\mathbb{R}^d)$  is a Wasserstein space equipped with Wasserstein distance. Consider MMS under transportation cost  $W_{c_h}$ :
$$
\mu_ {(k + 1) h} := \underset {\mu \in \mathcal {P} _ {c _ {h}} (\mathbb {R} ^ {d})} {\arg \min } D _ {\mathrm {K L}} (\mu \| \pi) + W _ {c _ {h}} (\mu , \mu_ {k h}). \tag {7}
$$
Denote the optimal transportation map under  $W_{ch}$  from  $\mu_{(k + 1)h}$  to  $\mu_{kh}$  by  $T_{k}(\cdot)$ . Then we have
$$
\frac {T _ {k} (x) - x}{h} = - \nabla g ^ {*} (\nabla \log \pi (x) - \nabla \log \mu_ {(k + 1) h} (x)). \tag {8}
$$
Please refer to Appendix A for full statements and proofs. Informally,  $\mu_{(k + 1)h} \approx \mu_{kh}$  for small step size  $h$  (Santambrogio, 2017). Further note that  $\frac{T_k(x) - x}{h}$  is the optimal velocity field associated with the transport from  $\mu_{(k + 1)h}$  to  $\mu_{kh}$  (and not vice versa). If step size  $h \to 0$ , then following Jordan et al. (1998), we can recover the dynamics in continuous time:
$$
d x _ {t} = v _ {t} d t, v _ {t} = \nabla g ^ {*} (\nabla \log \pi - \nabla \log \mu_ {t}). \tag {9}
$$
We call (9) the generalized Wasserstein gradient (GWG) flow. If we set  $g(\cdot) = \frac{1}{2} \| \cdot \|_2^2$  or any positive definite quadratic form  $g(\cdot) = \frac{1}{2} \| \cdot \|_H^2$ , 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(\cdot) = \frac{1}{p} \| \cdot \|_p^p$  for some  $p > 1$  and then  $g^{*}(\cdot) = \frac{1}{q} \| \cdot \|_q^q$ . Note that if the particles move along the vector field  $v_{t} = \nabla g^{*}(\nabla \log \frac{\pi}{\mu_{t}})$ , then the descent rate of  $D_{\mathrm{KL}}(\mu_t\| \pi)$  is
$$
\partial_ {t} D _ {\mathrm {K L}} \left(\mu_ {t} \| \pi\right) = - \mathbb {E} _ {\mu_ {t}} \left| \left| \nabla \log \frac {\pi}{\mu_ {t}} \right| \right| _ {q} ^ {q}. \tag {10}
$$
If we choose  $q$  such that  $\mathbb{E}_{\mu_t}\left\| \nabla \log \frac{\pi}{\mu_t}\right\|_q^q$  is large, then  $D_{\mathrm{KL}}(\mu_t\| \pi)$  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.
Example 1. Let  $\pi = \frac{1}{2}\mathcal{N}(-m,1) + \frac{1}{2}\mathcal{N}(m,1)$  and  $\mu = \frac{3}{4}\mathcal{N}(-m,1) + \frac{1}{4}\mathcal{N}(m,1)$ . Then for any  $m\geq \frac{1}{80}, q\geq 1$ , the following holds:
$$
\frac {0 . 0 8}{q m} \left(\frac {m}{3}\right) ^ {q} \exp \left(- \frac {m ^ {2}}{2}\right) \leq \mathbb {E} _ {\mu} \left| \left| \nabla \log \frac {\pi}{\mu} \right| \right| _ {q} ^ {q} \leq \frac {0 . 2}{q m} (4 m) ^ {q} \exp \left(- \frac {m ^ {2}}{2}\right). \tag {11}
$$
However, the KL divergence between  $\pi$  and  $\mu$  is large:  $D_{\mathrm{KL}}(\mu \| \pi) \geq \frac{1}{10\sqrt{2}}$ .
Suppose the target distribution is  $\pi$  and we run ParVI with current particle distribution  $\mu$ . 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_{\mathrm{KL}}(\mu \| \pi)$  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.
The forward-Euler discretization of the dynamics (9) is
$$
x _ {(k + 1) h} = x _ {k h} + \nabla g ^ {*} \left(\nabla \log \frac {\pi}{\mu_ {k h}} \left(x _ {k h}\right)\right) h. \tag {12}
$$
However, since the score of current particle distribution  $\mu_{kh}$  is generally unknown, we need a method to efficiently estimate the GWG direction  $\nabla g^{*}(\nabla \log \frac{\pi}{\mu_{kh}})$ . Given the distribution of current particles  $\mu$ , by the definition of convex conjugate, we have
$$
\nabla g ^ {*} (\nabla \log \frac {\pi}{\mu}) = \arg \max _ {v} \mathbb {E} _ {\mu} [ \langle \nabla \log \frac {\pi}{\mu}, v \rangle - g (v) ].
$$
If we parameterize  $v$  as a neural network  $f_{w}$  with  $w \in \mathcal{W}$ , then we can maximize the following objective with respect to  $w$ :
$$
\begin{array}{l} \mathcal {L} (w): = \mathbb {E} _ {\mu} \left[ \left\langle \nabla \log \frac {\pi}{\mu}, f _ {w} \right\rangle - g \left(f _ {w}\right) \right] \tag {13} \\ = \mathbb {E} _ {\mu} \left[ (\nabla \log \pi) ^ {T} f _ {w} + \nabla \cdot f _ {w} - g (f _ {w}) \right] \\ \end{array}
$$
Algorithm 1 GWG: Generalized Wasserstein Gradient Flow
Require: Unnormalized target distribution  $\pi$ , initial particles  $\{x_0^i\}_{i=1}^n$ , initial parameter  $w_0$ , iteration number  $N$ ,  $N'$ , particle step size  $h$ , parameter step size  $\eta$
for  $k = 0,\dots ,N - 1$  do
Assign  $w_{k}^{0} = w_{k}$
for  $t = 0,\dots ,N^{\prime} - 1$  do
$$
\widehat {\mathcal {L}} (w) = \frac {1}{n} \sum_ {i = 1} ^ {n} \nabla \log \pi \left(x _ {k} ^ {i}\right) ^ {T} f _ {w} \left(x _ {k} ^ {i}\right) + \nabla \cdot f _ {w} \left(x _ {k} ^ {i}\right) - g \left(f _ {w} \left(x _ {k} ^ {i}\right)\right) \tag {14}
$$
Update  $w_{k}^{t + 1} = w_{k}^{t} + \eta \nabla_{w}\hat{\mathcal{L}} (w_{k}^{t})$
Update  $w_{k + 1} = w_k^{N'}$
Update particles  $x_{k + 1}^{i} = x_{k}^{i} + hf_{w_{k + 1}}(x_{k}^{i})$  for  $i = 1,\dots ,n$
return Particles  $\{x_N^i\}_{i = 1}^n$
Here the second equation is by Stein's identity (we assume  $\mu$  vanishes at infinity). This way, the gradient of  $\mathcal{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  $\nabla_{x}\cdot f_{w}(x)$  needs  $\mathcal{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.,
$$
\frac {1}{n} \sum_ {i = 1} ^ {n} \nabla \cdot f _ {w} \left(x _ {k} ^ {i}\right) \approx \frac {1}{n} \sum_ {i = 1} ^ {n} \xi_ {i} ^ {T} \nabla \left(f _ {w} \left(x _ {k} ^ {i}\right) \cdot \xi_ {i}\right), \tag {15}
$$
where  $\xi_{i}\in \mathbb{R}^{d}$  are independent random vectors satisfying  $\mathbb{E}\xi_i\xi_i^T = I_d$ . This is still an unbiased estimator but only needs  $\mathcal{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:
$$
X _ {(k + 1) h} = X _ {k h} + v _ {k} \left(X _ {k h}\right) h, \tag {16}
$$
where  $v_{k}$  is the neural-net-estimated GWG at time  $kh$ . Define the interpolation process
$$
X _ {t} = X _ {k h} + (t - k h) v _ {k} \left(X _ {k h}\right), \text {f o r} t \in [ k h, (k + 1) h ], \tag {17}
$$
and let  $\mu_t$  denote the law of  $X_t$ . Note that here we do not assume isoperimetry of target distribution  $\pi$  (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 1.  $g^{*}(\cdot) = \frac{1}{q}\| \cdot \|_{q}^{q}$  for some  $q > 1$ . And for any  $k$ ,  $\mathbb{E}_{\mu_{kh}}\left|\left|v_k - \nabla g^*(\nabla \log \frac{\pi}{\mu_{kh}})\right|\right|_p^p \leq \varepsilon_k$ .
Assumption 2.  $g^{*}(\cdot)$  is  $\alpha$ -strongly convex and  $\beta$ -smooth. Define  $\kappa := \frac{\beta}{\alpha}$ . And for any  $k$ ,  $\mathbb{E}_{\mu_{kh}}\Bigg{|}\Bigg{|}v_k - \nabla g^* (\nabla \log \frac{\pi}{\mu_{kh}})\Bigg{|}\Bigg{|}_2^2\leq \varepsilon_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(\cdot) = \frac{1}{p} \| \cdot \|_p^p$  for some  $p > 1$ . Given current particle distribution  $\mu$ , we can define the training loss  $\mathcal{L}_{train}(v) \coloneqq \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}_{train}^* \coloneqq \mathcal{L}_{train}(v^*) < \infty$ . For any arbitrarily small  $\varepsilon_1 > 0$ , there exists  $\varepsilon_2 \coloneqq \varepsilon_2(\varepsilon_1, p) < \infty$ , such that
$$
\mathbb {E} _ {\mu} \bigg | \bigg | v - \nabla g ^ {*} (\nabla \log \frac {\pi}{\mu}) \bigg | \bigg | _ {p} ^ {p} \leq \varepsilon_ {1} \mathcal {L} _ {t r a i n} ^ {*} + \varepsilon_ {2} [ \mathcal {L} _ {t r a i n} ^ {*} - \mathcal {L} _ {t r a i n} (v) ].
$$
Besides, if  $p \geq 2$ ,  $\varepsilon_1$  can be 0 while  $\varepsilon_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}(\cdot)$  is twice differentiable. For any  $p > 1$
$$
G _ {p} := \sup _ {x, y} \frac {\| v _ {k} (x) - v _ {k} (y) \| _ {p}}{\| x - y \| _ {p}} <   \infty , M _ {p} := \sup _ {x, z} \lim _ {\delta \to 0 ^ {+}} \frac {\| \nabla v _ {k} (x + \delta z) - \nabla v _ {k} (x) \| _ {o p}}{\delta \| z \| _ {p}} <   \infty .
$$
Note that here we do not assume the smoothness of potential  $\log \pi$  explicitly. But informally,  $G_{p}$  and  $M_{p}$  correspond to the Lipschitz constant of the gradient and the Hessian of  $\log \pi$ , respectively.
Let  $\bar{\mu}_{Nh} \coloneqq \frac{1}{Nh}\int_{0}^{Nh}\mu_{t}dt$  and  $K_{0} \coloneqq D_{\mathrm{KL}}(\mu_{0}\| \pi)$ . Now we present our main results.
Theorem 3 (Full version see Theorem D.9). Under Assumption 1, 3, the following bound holds with proper step size  $h$ :
$$
\left. \mathbb {E} _ {\bar {\mu} _ {N h}} \right| \left| \nabla \log \frac {\pi}{\bar {\mu} _ {N h}} \right| \Bigg | _ {q} ^ {q} = \tilde {\mathcal {O}} \left(\left(\frac {M _ {p} K _ {0} d}{N}\right) ^ {\frac {q}{q + 1}} + \frac {G _ {2} K _ {0} d}{N} + \frac {\sum_ {k = 0} ^ {N - 1} \varepsilon_ {k}}{N}\right). \tag {18}
$$
Here  $\tilde{\mathcal{O}} (\cdot)$  hides all the constant factors that only depend on  $q$
Theorem 4 (Full version see Theorem D.10). Under Assumption 2, 3 with  $\alpha = 1$ , the following bound holds with proper step size  $h$ :
$$
\left. \mathbb {E} _ {\tilde {\mu} _ {N h}} \right| \left| \nabla \log \frac {\pi}{\bar {\mu} _ {N h}} \right| _ {2} ^ {2} = \mathcal {O} \left(\left(\frac {\kappa M _ {2} K _ {0} d}{N}\right) ^ {\frac {2}{3}} + \frac {G _ {2} K _ {0} (d + \kappa)}{N} + \frac {\sum_ {k = 0} ^ {N - 1} \varepsilon_ {k}}{N}\right). \tag {19}
$$
The proofs in this section are deferred to Appendix D. To interpret our results, suppose  $g^{*}(\cdot) = \frac{1}{q}\| \cdot \|_{q}^{q}$  and  $\epsilon \lesssim (\frac{M_p}{G_2})^q$ . If the neural net  $v_{k}(\cdot)$  can approximate  $\nabla g^{*}(\nabla \log \frac{\pi}{\mu_{kh}})$  accurately (i.e.,  $\varepsilon_k\lesssim \epsilon$ ), then to obtain a probabilistic measure  $\mu$  such that  $\mathbb{E}_{\mu}\left\| \nabla \log \frac{\pi}{\mu}\right\| _q^q\lesssim \epsilon$ , the iteration complexity is  $\tilde{\mathcal{O}} (M_pK_0d\epsilon^{-(1 + \frac{1}{q})})$ . If we further let  $q = p = 2$ , the complexity is  $\mathcal{O}(K_0d\epsilon^{-\frac{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  $\| \cdot \| _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  $\left\{\frac{1}{p} \| \cdot \|_p^p : p > 1\right\}$  and propose a procedure that adaptively chooses  $p$  to accelerate convergence.
Algorithm 2 Ada-GWG: Adaptive Generalized Wasserstein Gradient Flow
Require: unnormalized target distribution  $\pi$ , initial particles  $\{x_0^i\}_{i=1}^n$ , initial parameter  $w_0$ , iteration number  $N$ ,  $N'$ , step size  $h$ ,  $\eta$ ,  $\tilde{\eta}$ , lower and upper bounds on  $p$ :  $lb$ ,  $ub$
$$
\widehat {\mathcal {L}} (w) = \frac {1}{n} \sum_ {i = 1} ^ {n} \nabla \log \pi \left(x _ {k} ^ {i}\right) ^ {T} f _ {w} \left(x _ {k} ^ {i}\right) + \nabla \cdot f _ {w} \left(x _ {k} ^ {i}\right) - \frac {1}{p _ {k}} \| f _ {w} \left(x _ {k} ^ {i}\right) \| _ {p _ {k}} ^ {p _ {k}} \tag {22}
$$
Compute  $\widehat{A}(p_k) = \frac{1}{n} \sum_{i=1}^{n} \frac{1}{p_k} \|f_{w_{k+1}}(x_k^i)\|_{p_k}^{p_k}$ .
Update  $p_{k + 1} = \mathbf{clip}(p_k + \tilde{\eta}\nabla \widehat{A} (p_k),lb,ub)$
Consider the continuous time dynamics  $dx_{t} = f_{t}(x_{t})dt$  and denote the distribution of particles at time  $t$  as  $\mu_t$ , we have the following proposition.
Proposition 5. For  $g(\cdot) = \frac{1}{p} \| \cdot \|_p^p$ , the derivative of KL divergence has an upper bound:
$$
\left. \right. \partial_ {t} D _ {\mathrm {K L}} \left(\mu_ {t} \| \pi\right) \leq - \frac {1}{p} \mathbb {E} _ {\mu_ {t}} \left|\left| \nabla \log \frac {\pi}{\mu_ {t}} \right|\right| _ {q} ^ {q} + \frac {1}{p} \mathbb {E} _ {\mu_ {t}} \left|\left| \nabla g ^ {*} (\nabla \log \frac {\pi}{\mu_ {t}}) - f _ {t} \right|\right| _ {p} ^ {p}. \tag {20}
$$