SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The proof is in Appendix E. If the neural network  $f_{t}$  can approximate the objective well, i.e.,  $f_{t} \approx \nabla g^{*}(\nabla \log \frac{\pi}{\mu_{t}})$ , then informally we can omit the second term and thus
$$
\partial_ {t} D _ {\mathrm {K L}} \left(\mu_ {t} \| \pi\right) \lesssim - \frac {1}{p} \mathbb {E} _ {\mu_ {t}} \left| \left| \nabla \log \frac {\pi}{\mu_ {t}} \right| \right| _ {q} ^ {q} \approx - \frac {1}{p} \mathbb {E} _ {\mu_ {t}} \| f _ {t} \| _ {p} ^ {p} =: - A (p). \tag {21}
$$
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.
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  $\mathcal{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
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  $\pi \propto \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]\longrightarrow \mathbb{R}$  and dynamics given by
$$
d u _ {t} = \frac {1 0 u \left(1 - u ^ {2}\right)}{1 + u ^ {2}} d t + d x _ {t}, \quad u _ {0} = 0, \tag {23}
$$
where  $x = (x_{t})_{t\geq 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
is to infer the posterior of the driving process  $p(x|y)$ . The forward operator is defined by  $\mathcal{F}(x) = (u_{t_1},\dots ,u_{t_{20}})\in \mathbb{R}^{20}$ , where  $t_i = 0.05i$ . This is achieved by discretizing the above SDE (23) using an Euler-Maruyama scheme with step size  $\Delta t = 0.01$ ; therefore the dimensionality of the problem is 100. The noisy observations are obtained as  $y = \mathcal{F}(x_{\mathrm{true}}) + \xi \in \mathbb{R}^{20}$ , where  $x_{\mathrm{true}}$  is a Brownian motion path and  $\xi \sim \mathcal{N}(0,\sigma^2 I)$  with  $\sigma = 0.1$ . The prior is given by the Brownian motion  $x = (x_{t})_{t\geq 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., GWG with  $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.
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
adjust the Wassertein metric to accelerate convergence. Extensive numerical results showed that Ada-GWG outperforms conventional ParVI methods.
This work was supported by National Natural Science Foundation of China (grant no. 12201014 and grant no. 12292983). The research of Cheng Zhang was supported in 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 are partially supported by the elite undergraduate training program of School of Mathematical Sciences in Peking University. The authors are grateful for the computational resources provided by the High-performance Computing Platform of Peking University. The authors appreciate the anonymous NeurIPS reviewers for their constructive feedback.
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=dp0YN7o8Jm.
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. Minimumstein 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.
A Minimizing Movement Scheme
A.1 Geometric Interpretation
In fact, under some mild conditions of  $g$ , the transportation cost  $W_{c_h}(\cdot ,\cdot)$  can induce a Wasserstein metric and thus  $\mathcal{P}_{c_h}(\mathbb{R}^d)\coloneqq \{\mu \in \mathcal{P}(\mathbb{R}^d):\mathbb{E}_\mu [g(\frac{2x}{h})] < \infty \}$  is indeed a Wasserstein space.
Proposition A.1. Let  $g(\cdot) = g_0(\| \cdot \|)$  where  $g_0: \mathbb{R}^+ \cup \{0\} \to \mathbb{R}^+ \cup \{0\}$  satisfies  $g_0(0) = 0$  and  $\| \cdot \|$  can be any norm in  $\mathbb{R}^d$ . Then  $g_0^{-1}(W_{c_h}(\cdot, \cdot))$  is a metric on  $\mathcal{P}_{c_h}(\mathbb{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  $\pi, \mu, \nu \in \mathcal{P}_{c_h}$ . It is obvious that  $g_0^{-1}(W_{c_h}(\mu, \nu)) = 0$  if and only if  $\mu = \nu$ . Besides,  $g_0^{-1}(W_{c_h}(\cdot, \cdot))$  is symmetric. In the rest part of proof we aim to show that  $g_0^{-1}(W_{c_h}(\mu, \pi)) + g_0^{-1}(W_{c_h}(\nu, \pi)) \geq g_0^{-1}(W_{c_h}(\mu, \nu))$ . By Gluing lemma Villani et al. (2009), we can construct random variables  $X \sim \pi, Y \sim \mu, Z \sim \nu$  such that  $(X, Y)$ ,  $(X, Z)$  are the optimal coupling of  $(\pi, \mu)$  and  $(\pi, \nu)$  for transportation cost  $W_{c_h}$ , respectively. Then we have
$$
\begin{array}{l} g _ {0} ^ {- 1} \left(W _ {c h} (\mu , \nu)\right) \leq g _ {0} ^ {- 1} \left(\mathbb {E} g _ {0} \left(\frac {\| Y - Z \|}{h}\right)\right) \\ \leq g _ {0} ^ {- 1} \left(\mathbb {E} g _ {0} \left(\frac {\| X - Y \| + \| X - Z \|}{h}\right)\right) \\ \leq g _ {0} ^ {- 1} \left(\mathbb {E} g _ {0} \left(\frac {\| X - Y \|}{h}\right)\right) + g _ {0} ^ {- 1} \left(\mathbb {E} g _ {0} \left(\frac {\| X - Z \|}{h}\right)\right) \\ = g _ {0} ^ {- 1} \left(W _ {c _ {h}} (\mu , \pi)\right) + g _ {0} ^ {- 1} \left(W _ {c _ {h}} (\pi , \nu)\right). \\ \end{array}
$$
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  $\| \cdot \|$  in  $\mathbb{R}^d$  can be  $\| \cdot \|_p, \| \cdot \|_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(\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\}$ . If  $g$  satisfies assumptions in Proposition A.1,  $\mathcal{P}_{c_h}(\mathbb{R}^d)$  is a Wasserstein space equipped with Wasserstein metric. 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 {24}
$$
Denote the optimal transportation map under  $W_{c_h}$  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 {25}
$$
Proof. By Kantorovich duality (Villani, 2021), the optimal transportation cost (3) has an equivalent definition:
$$
W _ {c} (\mu , \nu) = \sup  _ {\varphi} \int \varphi d \mu + \int \varphi^ {c} d \nu , \text {w h e r e} \varphi^ {c} (y) := \inf  _ {x \in \mathbb {R} ^ {d}} c (x, y) - \varphi (x). \tag {26}
$$
Take the functional derivative of the optimization problem (24) and define the optimal  $\varphi$  in (26) as  $\psi$ . The following holds:
$$
\frac {\delta}{\delta \mu} D _ {\mathrm {K L}} \left(\mu_ {(k + 1) h} \| \pi\right) + \psi_ {k h} = \text {c o n s t .} \tag {27}
$$
Here the reason for having a constant instead of zero is that we constrain  $\mu_{(k + 1)h}$  in the space of smooth probability density. Note that  $c_{h}(x,y)\leq \frac{1}{2}\left(g(\frac{2x}{h}) + g(\frac{2y}{h})\right)$  and  $\mu_{kh},\mu_{(k + 1)h}\in \mathcal{P}_{c_h}(\mathbb{R}^d)$ , then by the fundamental theorem of optimal transportation (Villani, 2021),
$$
\psi_ {k h} (x) + \psi_ {k h} ^ {c _ {h}} (y) = c _ {h} (x, y), \text {f o r} y = T _ {k} (x), \tag {28}
$$
which implies  $\nabla \psi_{kh}(x) = \nabla_{x}c_{h}(x,y) = \nabla g(\frac{x - y}{h}),i.e.$
$$
x - y = \nabla g ^ {*} (\nabla \psi_ {k h} (x)) h. \tag {29}
$$
Combine this equation with (27) and thus the optimal map is given by
$$
\begin{array}{l} \frac {T _ {k} (x) - x}{h} = - \nabla g ^ {*} (- \nabla \frac {\delta}{\delta \mu} D _ {\mathrm {K L}} (\mu_ {(k + 1) h} \| \pi)) \\ = - \nabla g ^ {*} (\nabla \log \frac {\pi}{\mu_ {(k + 1) h}}). \\ \end{array}
$$
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  $\pi_0 = \mathcal{N}(-m,1),\pi_1 = \mathcal{N}(m,1)$  and rewrite  $\pi = \frac{1}{2}\pi_0 + \frac{1}{2}\pi_1,\mu = \frac{3}{4}\pi_0 + \frac{1}{4}\pi_1$ . The lower bound of KL divergence follows from Devroye et al. (2018) and Pinsker inequality. In addition, Balasubramanian et al. (2022) shows
$$
\nabla \log \pi - \nabla \log \mu = - m \frac {\pi_ {0} \pi_ {1}}{2 \pi \mu}. \tag {30}
$$
Also note that  $\frac{\pi_0}{\pi_1} = \exp(-2mx)$ . Therefore for any  $q \geq 1$ , the following bound holds:
$$
\begin{array}{l} \left. \mathbb {E} _ {\mu} \right| \left| \nabla \log \frac {\pi}{\mu} \right| _ {q} ^ {q} = \frac {m ^ {q}}{2 ^ {q}} \int \frac {\pi_ {0} ^ {q} \pi_ {1} ^ {q}}{\mu^ {q - 1} \pi^ {q}} d x \\ = 4 ^ {q - 1} m ^ {q} \int \frac {\pi_ {0} ^ {q} \pi_ {1} ^ {q}}{\left(3 \pi_ {0} + \pi_ {1}\right) ^ {q - 1} \left(\pi_ {0} + \pi_ {1}\right) ^ {q}} d x \\ \geq 4 ^ {q - 1} m ^ {q} \left(\int_ {x \geq \frac {1}{m}} \frac {\pi_ {0} ^ {q}}{(1 + e ^ {- 2}) ^ {q} \pi_ {1} ^ {q - 1} (1 + 3 e ^ {- 2}) ^ {q - 1}} d x + \int_ {x \leq - \frac {1}{m}} \frac {\pi_ {1} ^ {q}}{(1 + e ^ {- 2}) ^ {q} \pi_ {0} ^ {q - 1} (3 + e ^ {- 2}) ^ {q - 1}} d x\right). \\ \end{array}
$$
$$
\begin{array}{l} \int_ {x \geq \frac {1}{m}} \frac {\pi_ {0} ^ {q}}{\pi_ {1} ^ {q - 1}} d x = \int_ {x \geq \frac {1}{m}} \frac {1}{\sqrt {2 \pi}} \exp (- \frac {1}{2} (x + m) ^ {2} - 2 m (q - 1) x) d x \\ = \operatorname {P r o b} _ {\mathcal {N} (0, 1)} \left\{Z \geq (2 q - 1) m + \frac {1}{m} \right\} \exp (2 q (q - 1) m ^ {2}) \tag {32} \\ \geq \frac {3}{4} \frac {1}{(2 q - 1) m + \frac {1}{m}} \frac {1}{\sqrt {2 \pi}} \exp \left(- \frac {1}{2} ((2 q - 1) m + \frac {1}{m}) ^ {2} + 2 q (q - 1) m ^ {2}\right) \\ \geq \frac {3}{4} \frac {1}{2 q m} \frac {1}{\sqrt {2 \pi}} \exp (- \frac {1}{2} m ^ {2} - 2 q + \frac {1}{2}). \\ \end{array}
$$
Here the first inequality is by  $\mathbf{Prob}_{\mathcal{N}(0,1)}\{Z \geq t\} \geq \left(\frac{1}{t} - \frac{1}{t^3}\right)\frac{1}{\sqrt{2\pi}}\exp\left(-\frac{t^2}{2}\right)$  for any  $t > 0$  and  $(2q - 1)m + \frac{1}{m} \geq 2$ . Similarly, we can prove that
$$
\int_ {x \leq - \frac {1}{m}} \frac {\pi_ {1} ^ {q}}{\pi_ {0} ^ {q - 1}} d x \geq \frac {3}{4} \frac {1}{2 q m} \frac {1}{\sqrt {2 \pi}} \exp \left(- \frac {1}{2} m ^ {2} - 2 q + \frac {1}{2}\right). \tag {33}
$$
Plug (32) and (33) in (31),
$$
\begin{array}{l} \mathbb {E} _ {\mu} \left\| \nabla \log \frac {\pi}{\mu} \right\| _ {q} ^ {q} \geq \frac {1}{q} m ^ {q - 1} \exp (- \frac {m ^ {2}}{2}) \cdot \frac {3}{4 ^ {3}} \sqrt {\frac {2 e}{\pi}} 4 ^ {q} (e ^ {2} + 1) ^ {- q} \left[ (\frac {1}{1 + 3 e ^ {- 2}}) ^ {q - 1} + (\frac {1}{3 + e ^ {- 2}}) ^ {q - 1} \right] \\ \geq \frac {1}{q} m ^ {q - 1} \exp (- \frac {m ^ {2}}{2}) \cdot \frac {3}{4 ^ {3}} \sqrt {\frac {2 e}{\pi}} (1 + 3 e ^ {- 2}) \left[ \frac {4}{(1 + e ^ {2}) (1 + 3 e ^ {- 2})} \right] ^ {q} \\ \geq \frac {0 . 0 8}{q m} (\frac {m}{3}) ^ {q} \exp (- \frac {m ^ {2}}{2}). \\ \end{array}
$$
As for the upper bound,
$$
\begin{array}{l} \mathbb {E} _ {\mu} \left| \left| \nabla \log \frac {\pi}{\mu} \right| \right| _ {q} ^ {q} = \frac {m ^ {q}}{2 ^ {q}} \int \frac {\pi_ {0} ^ {q} \pi_ {1} ^ {q}}{\mu^ {q - 1} \pi^ {q}} d x \\ = 4 ^ {q - 1} m ^ {q} \int \frac {\pi_ {0} ^ {q} \pi_ {1} ^ {q}}{\left(3 \pi_ {0} + \pi_ {1}\right) ^ {q - 1} \left(\pi_ {0} + \pi_ {1}\right) ^ {q}} d x \tag {34} \\ \leq 4 ^ {q - 1} m ^ {q} \left(\int_ {x \geq 0} \frac {\pi_ {0} ^ {q}}{\pi_ {1} ^ {q - 1}} d x + \int_ {x \leq 0} \frac {\pi_ {1} ^ {q}}{\pi_ {0} ^ {q - 1}} d x\right). \\ \end{array}
$$
$$
\begin{array}{l} \int_ {x \geq 0} \frac {\pi_ {0} ^ {q}}{\pi_ {1} ^ {q - 1}} d x = \int_ {x \geq 0} \frac {1}{\sqrt {2 \pi}} \exp \left(- \frac {1}{2} (x + m) ^ {2} - 2 m (q - 1) x\right) d x \\ = \operatorname {P r o b} _ {\mathcal {N} (0, 1)} \{Z \geq (2 q - 1) m \} \exp (2 q (q - 1) m ^ {2}) \\ \leq \frac {1}{(2 q - 1) m} \frac {1}{\sqrt {2 \pi}} \exp \left(- \frac {1}{2} \left(\left(2 q - 1\right) m\right) ^ {2} + 2 q (q - 1) m ^ {2}\right) \tag {35} \\ \leq \frac {1}{q m} \frac {1}{\sqrt {2 \pi}} \exp \left(- \frac {1}{2} m ^ {2}\right). \\ \end{array}
$$
Here the first inequality is by  $\mathbf{Prob}_{\mathcal{N}(0,1)}\{Z \geq t\} \leq \frac{1}{t} \frac{1}{\sqrt{2\pi}} \exp \left(-\frac{t^2}{2}\right)$  for any  $t > 0$ . Similarly, we can prove that
$$
\int_ {x \leq 0} \frac {\pi_ {1} ^ {q}}{\pi_ {0} ^ {q - 1}} d x \leq \frac {1}{q m} \frac {1}{\sqrt {2 \pi}} \exp \left(- \frac {1}{2} m ^ {2}\right). \tag {36}
$$
$$
\begin{array}{l} \left. \mathbb {E} _ {\mu} \right\| \nabla \log \frac {\pi}{\mu} \Big \| _ {q} ^ {q} \leq 4 ^ {q - 1} m ^ {q} \cdot \frac {2}{q m} \frac {1}{\sqrt {2 \pi}} \exp \left(- \frac {1}{2} m ^ {2}\right) \\ \leq \frac {0 . 2}{q m} (4 m) ^ {q} \exp \left(- \frac {m ^ {2}}{2}\right). \\ \end{array}
$$
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  $\ell (w,x)\coloneqq \nabla \log \pi (x)^T f_w(x) + \nabla_x\cdot f_w(x) - g(f_w(x))$
Assumption 4.  $\mathcal{W}$  is compact and  $\mathcal{L}(\cdot)$  defined in (13) has a unique maximum point  $w^{*}\in \operatorname {int}(\mathcal{W})$
Assumption 5.  $f_{w}(x),\nabla_{x}f_{w}(x)$  are continuous with  $w$  . max  $\{\| f_w(x)\| ,|\nabla_x\cdot f_w(x)|,g(f_w(x))\} \leq$ $M_0(x)$  for some  $M_0\in \mathcal{L}^1 (\mu)$
Assumption 6. There exists a neighborhood  $\mathcal{N}$  of  $w^{*}$  such that  $\ell (\cdot ,x)$  is twice differentiable in  $\mathcal{N}$  and  $\| \nabla^2\ell (w,x)\| \leq M_1(x)$  for all  $w\in \mathcal{N}$ . Additionally, assume  $M_{1}\in \mathcal{L}^{1}(\mu)$  and  $H\coloneqq \nabla^{2}\mathcal{L}(w^{*})$  is non-singular.
Theorem C.1. Given  $x_{1},\dots ,x_{n}\stackrel {i.i.d.}{\sim}\mu$  , let  $\widehat{w}_n\coloneqq \underset {w}{\arg \min}\widehat{\mathcal{L}}_n(w)\coloneqq \frac{1}{n}\sum_{i = 1}^{n}\ell (w,x_i)$  . Under
Assumption 4-6, we have
$$
\sqrt {n} \left(\hat {w} _ {n} - w ^ {*}\right) \xrightarrow {p} \mathcal {N} \left(0, H ^ {- 1} \Sigma H ^ {- 1}\right), \tag {37}
$$
where  $\Sigma = \mathbb{E}_{\mu}\nabla_{w}\ell (w^{*},\cdot)\otimes \nabla_{w}\ell (w^{*},\cdot)$
Proof. Note that under Assumption 5,  $|\ell(w,x)| \leq M(x)$  for some  $M \in \mathcal{L}^1(\mu)$ . By Newey & McFadden (1986, Lemma 2.4, Theorem 2.1),  $\widehat{w}_n$  is weakly consistent for  $w^*$ . Additionally, since  $w^*$  is the maximum point of  $\mathcal{L}$  and by Central Limit Theorem,
$$
\sqrt {n} \nabla \widehat {\mathcal {L}} _ {n} (w ^ {*}) = \frac {1}{\sqrt {n}} \sum_ {i = 1} ^ {n} \nabla_ {w} \ell (w ^ {*}, x _ {i}) \xrightarrow {p} \mathcal {N} (0, \Sigma).
$$
By Assumption 6 and Newey & McFadden (1986, Lemma 2.4), the second order derivative converges uniformly, i.e.
$$
\sup _ {w \in \mathcal {N}} \| \nabla^ {2} \widehat {\mathcal {L}} _ {n} (w) - \nabla^ {2} \mathcal {L} (w) \| \stackrel {p} {\to} 0.
$$
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  $\varepsilon_2 < \infty$  such that for any  $a, b \in \mathbb{R}$ , the following inequality holds:
$$
\left| b - \operatorname {s g n} (a) \right| a | ^ {q - 1} | ^ {p} \leq \varepsilon_ {1} \frac {| a | ^ {q}}{q} + \varepsilon_ {2} \left(\frac {| a | ^ {q}}{q} - a b + \frac {| b | ^ {p}}{p}\right). \tag {38}
$$
If  $a = 0$ , then  $\varepsilon_2 \geq p$  is sufficient. Without loss of generality, suppose  $a = 1$  (by replacing  $b$  with  $\operatorname{sgn}(a)b / |a|^{q - 1}$ ). We only need to show that
$$
\left| b - 1 \right| ^ {p} \leq \varepsilon_ {1} + \varepsilon_ {2} \left(\frac {1}{q} - b + \frac {\left| b \right| ^ {p}}{p}\right). \tag {39}
$$
(1) Case  $p \geq 2$ .
Let  $\varepsilon_1 = 0$ . Since  $\lim_{b\to 1}\frac{|b - 1|^p}{\frac{1}{q} - b + \frac{|b|^p}{p}} = \lim_{b\to 1}\frac{2|b - 1|^{p - 2}}{p - 1} \leq 2$ , so there exists  $\delta > 0$  such that when  $b \in [1 - \delta, 1 + \delta]$ , (39) holds if  $\varepsilon_2 \geq p + 1 > 2$ . Also note that  $f(b) = \frac{|b - 1|^p}{\frac{1}{q} - b + \frac{|b|^p}{p}}$  is a continuous function on  $\mathbb{R} \setminus (1 - \delta, 1 + \delta)$  and  $\lim_{b\to \infty} f(b) = p < +\infty$ . Therefore,  $f(b)$  is bounded on  $\mathbb{R} \setminus [1 - \delta, 1 + \delta]$  and thus (39) holds for finite  $\varepsilon_2$ . It's obvious that  $\varepsilon_2$  only depends on  $p$  in this case. (2) Case  $p < 2$ .