SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
N = \mathcal {O} \left(\frac {M _ {2} d}{\lambda \sqrt {\epsilon}} \log \frac {\lambda D _ {\mathrm {K L}} (\mu_ {0} \| \pi)}{\epsilon}\right) \quad i t e r a t i o n s.
$$
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 \pi$  explicitly and thus our method can tackle more complex distributions.
Proof. For  $t \in (kh, (k + 1)h)$ , we apply Lemma D.8 and log-Sobolev inequality and thus
$$
\partial_ {t} D _ {\mathrm {K L}} \left(\mu_ {t} \| \pi\right) \leq - \frac {\lambda}{3} D _ {\mathrm {K L}} \left(\mu_ {t} \| \pi\right) + 3 \left(M _ {2} d h\right) ^ {2} + 4 \left(1 - h G _ {2}\right) ^ {- d} \epsilon . \tag {47}
$$
By Gronwall's inequality,
$$
D _ {\mathrm {K L}} \left(\mu_ {(k + 1) h} \| \pi\right) \leq e ^ {- \lambda h / 3} D _ {\mathrm {K L}} \left(\mu_ {k h} \| \pi\right) + 3 \lambda^ {- 1} \left(3 \left(M _ {2} d h\right) ^ {2} + 4 \left(1 - h G _ {2}\right) ^ {- d} \epsilon\right) (1 - e ^ {- \lambda h / 3}). \tag {48}
$$
Iterating the recursive bound,
$$
D _ {\mathrm {K L}} \left(\mu_ {N h} \| \pi\right) \leq \exp \left(- \frac {\lambda N h}{3}\right) D _ {\mathrm {K L}} \left(\mu_ {0} \| \pi\right) + 3 \left(3 \left(M _ {2} d h\right) ^ {2} + 4 \left(1 - h G _ {2}\right) ^ {- d} \epsilon\right) \lambda^ {- 1}. \tag {49}
$$
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  $\pi$  satisfies the modified log-Sobolev inequality  $\mathrm{mLSI}(q, \lambda_q)$  if the following holds:
$$
\operatorname {E n t} _ {\pi} (| f | ^ {q}) \leq \frac {q ^ {q - 1}}{\lambda_ {q}} \mathbb {E} _ {\pi} [ \| \nabla f \| _ {q} ^ {q} ], \text {f o r a l l s m o o t h} f: \mathbb {R} ^ {d} \to \mathbb {R}.
$$
Note that  $\mathrm{mLSI}(2,\lambda_2)$  reduces to the conventional log-Sobolev inequality with constant  $\lambda_{2}$ . As a direct corollary of this inequality, for any distribution  $\mu$ , we take  $f = (\frac{\mu}{\pi})^{1 / q}$  and thus
$$
D _ {\mathrm {K L}} (\mu \| \pi) \leq \frac {1}{q \lambda_ {q}} \mathbb {E} _ {\mu} \left| \left| \nabla \log \frac {\pi}{\mu} \right| \right| _ {q} ^ {q}. \tag {50}
$$
Theorem D.12. If we additionally assume that  $\pi$  satisfies (50), then under the same conditions of Theorem D.9 with  $\varepsilon_{k} \leq \epsilon$ , it holds that
$$
D _ {\mathrm {K L}} \left(\mu_ {N h} \| \pi\right) \leq \exp \left(- \frac {q \lambda_ {q} N h}{1 2}\right) D _ {\mathrm {K L}} \left(\mu_ {0} \| \pi\right) + 1 2 \left(A _ {1} \left(M _ {p} d h\right) ^ {q} + A _ {2} (1 - h G _ {2}) ^ {- d} \epsilon\right) \left(q \lambda_ {q}\right) ^ {- 1}.
$$
In particular, if  $\epsilon \lesssim \min \left\{(\frac{M_p}{G_2})^q, (\frac{dM_p}{qG_2})^q\right\}$ , we take  $h \asymp \frac{\epsilon^{1/q}}{M_p d}$  and then we obtain the guarantee  $D_{\mathrm{KL}}(\mu_{Nh} \| \pi) \lesssim \lambda_q^{-1} \epsilon$  after
$$
N = \tilde {\mathcal {O}} \left(\frac {M _ {p} d}{\lambda_ {q} \epsilon^ {1 / q}} \log \frac {\lambda_ {q} D _ {\mathrm {K L}} (\mu_ {0} \| \pi)}{\epsilon}\right) \quad i t e r a t i o n s.
$$
Proof. For  $t \in (kh, (k + 1)h)$ , we apply Lemma D.7 and (50) and thus
$$
\partial_ {t} D _ {\mathrm {K L}} \left(\mu_ {t} \| \pi\right) \leq - \frac {q \lambda_ {q}}{1 2} D _ {\mathrm {K L}} \left(\mu_ {t} \| \pi\right) + A _ {1} \left(M _ {p} d h\right) ^ {q} + A _ {2} \left(1 - h G _ {2}\right) ^ {- d} \epsilon . \tag {51}
$$
$$
D _ {\mathrm {K L}} \left(\mu_ {(k + 1) h} \| \pi\right) \leq e ^ {- q \lambda_ {q} h / 1 2} D _ {\mathrm {K L}} \left(\mu_ {k h} \| \pi\right) + 1 2 \left(q \lambda_ {q}\right) ^ {- 1} \left(A _ {1} \left(M _ {p} d h\right) ^ {q} + A _ {2} (1 - h G _ {2}) ^ {- d} \epsilon\right) \left(1 - e ^ {- q \lambda_ {q} h / 1 2}\right). \tag {52}
$$
$$
D _ {\mathrm {K L}} \left(\mu_ {N h} \| \pi\right) \leq \exp \left(- \frac {q \lambda_ {q} N h}{1 2}\right) D _ {\mathrm {K L}} \left(\mu_ {0} \| \pi\right) + 1 2 \left(A _ {1} \left(M _ {p} d h\right) ^ {q} + A _ {2} \left(1 - h G _ {2}\right) ^ {- d} \epsilon\right) \left(q \lambda_ {q}\right) ^ {- 1}. \tag {53}
$$
Remark 2.  $mLSI(q, \lambda_q)$  cannot hold for  $q > 2$  as mentioned in Barthe & Roberto (2008). However, we only need (50) to hold for all  $\mu = \mu_t$  to prove Theorem D.12. Plus, Gentil et al. (2005, 2007) replace  $\| \cdot \|_q^q$  with  $\max \{\| \cdot \|_2^2, \| \cdot \|_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  $\epsilon$ . On the other hand, it is generally difficult to predict how  $\lambda_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 \in \mathbb{R}^{d \times d}$  with positive eigenvalues, the following holds:
$$
| \log \det  A - \log \det  B | \leq d \| A - B \| _ {2} \max  \left\{\| A ^ {- 1} \| _ {2}, \| B ^ {- 1} \| _ {2} \right\}
$$
Proof. Suppose that the eigenvalues of real matrix  $(A - B)B^{-1}$  are  $\lambda_1, \overline{\lambda}_1, \dots, \lambda_k, \overline{\lambda}_k \in \mathbb{C}, \lambda_{2k+1}, \dots, \lambda_d \in \mathbb{R}$ . Here  $\overline{\lambda_j}$  is the complex conjugate of  $\lambda_j$ . Then it holds that:
$$
\begin{array}{l} \log \det  A - \log \det  B = \log \det  (I + (A - B) B ^ {- 1}) \\ = \log \prod_ {j = 1} ^ {d} (1 + \lambda_ {j}) \\ \leq \sum_ {j = 1} ^ {k} \log (1 + \lambda_ {j}) (1 + \overline {{\lambda_ {j}}}) + \sum_ {j = 2 k + 1} ^ {d} \log (1 + | \lambda_ {j} |) \\ \leq \sum_ {j = 1} ^ {d} \log (1 + | \lambda_ {j} |) \\ \leq d \| (A - B) B ^ {- 1} \| _ {2} \\ \end{array}
$$
Similarly, we have  $\log \det B - \log \det A \leq d\| (B - A)A^{-1}\|_2$  and thus we finish the proof.
Lemma D.14. Define non-negative constants  $c_{1}, c_{2}$  as:
$$
(c _ {1}, c _ {2}) = \left\{ \begin{array}{l l} (0, 2 ^ {p - q}) & \text {i f} q \leq 2, \\ \left(3 ^ {- p}, \min  \left\{3 ^ {q - p} (\frac {p}{q}) ^ {\frac {q - p}{q - 1}} (1 - \frac {p}{q}) ^ {\frac {q - p}{p}}, (q - 1) ^ {p} \left(\frac {\left(\frac {4}{3}\right) ^ {\frac {1}{q - 1}}}{\left(\frac {4}{3}\right) ^ {\frac {1}{q - 1}} - 1}\right) ^ {q - p} \right\}\right) & \text {o t h e r w i s e}. \end{array} \right.
$$
Then for any  $a, b \in \mathbb{R}$ , the following inequality holds:
$$
\left| s g n (a) | a | ^ {q - 1} - s g n (b) | b | ^ {q - 1} \right| ^ {p} \leq c _ {1} | a | ^ {q} + c _ {2} | a - b | ^ {q}
$$
Proof. We shall prove each of the two cases separately.
(1) Case  $q \leq 2$ .
If  $a, b$  have the same sign, we assume they are positive without loss of generality. Then  $|a^{q - 1} - b^{q - 1}| \leq |a - b|^{q - 1}$ , which implies  $|a^{q - 1} - b^{q - 1}|^p \leq |a - b|^q$ .
If  $a, b$  have different signs, we assume  $a \geq 0, b < 0$  without loss of generality. Then by Hölder inequality  $a^{q - 1} + (-b)^{q - 1} \leq 2^{2 - q}|a - b|^{q - 1}$ , i.e.,  $|a^{q - 1} - (-b)^{q - 1}|^p \leq 2^{p - q}|a - b|^q$ .
(2) Case  $q > 2$
If  $a, b$  have different signs, we assume  $a \geq 0, b < 0$  without loss of generality. Then  $|a^{q - 1} + (-b)^{q - 1}| \leq |a - b|^{q - 1}$ , which implies  $|a^{q - 1} - b^{q - 1}|^p \leq |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$ ,
$$
\left| 1 - b ^ {q - 1} \right| ^ {p} \leq c _ {1} + c _ {2} | 1 - b | ^ {q}. \tag {54}
$$
If  $b \leq 1$ , then by simple calculus,
$$
\begin{array}{l} (1 - b ^ {q - 1}) ^ {p} - c _ {2} (1 - b) ^ {q} \leq [ (q - 1) (1 - b) ] ^ {p} - c _ {2} (1 - b) ^ {q} \\ \leq c _ {2} ^ {- \frac {p}{q - p}} (q - 1) ^ {\frac {q + p}{q - p}} (\frac {p}{q}) ^ {\frac {p}{q - p}} (1 - \frac {p}{q}) \\ \leq c _ {1}. \\ \end{array}
$$
If  $1 < b \leq \left(1 + c_1^{1/p}\right)^{\frac{1}{q-1}} = \left(\frac{4}{3}\right)^{\frac{1}{q-1}}$ , then (54) is trivial.
If  $b > \left(\frac{4}{3}\right)^{\frac{1}{q - 1}}$
$$
\begin{array}{l} (b ^ {q - 1} - 1) ^ {p} - c _ {2} (b - 1) ^ {q} = (b - 1) ^ {p} \left[ \left(\frac {b ^ {q - 1} - 1}{b - 1}\right) ^ {p} - c _ {2} (b - 1) ^ {q - p} \right] \\ \leq (b - 1) ^ {p} \left([ (q - 1) b ^ {q - 2} ] ^ {p} - c _ {2} (b - 1) ^ {q - p}\right) \\ \leq 0. \\ \end{array}
$$
The last inequality is due to  $c_2 \geq (q - 1)^p \left( \frac{\left( \frac{4}{3} \right)^{\frac{1}{q - 1}}}{\left( \frac{4}{3} \right)^{\frac{1}{q - 1}} - 1} \right)^{q - p} \geq (q - 1)^p \left( \frac{b}{b - 1} \right)^{q - p}$ .
E Proof of Proposition 5
Proof. Note that  $g(\cdot) = \frac{1}{p} \| \cdot \|_p^p$  and thus  $\nabla g^*(\cdot) = |\cdot|^{q-1} \odot \operatorname{sgn}(\cdot)$ . By Young's inequality,
$$
\begin{array}{l} \partial_ {t} D _ {\mathrm {K L}} \left(\mu_ {t} \| \pi\right) = - \mathbb {E} _ {\mu_ {t}} \left\langle \nabla \log \frac {\pi}{\mu_ {t}}, \nabla g ^ {*} \left(\nabla \log \frac {\pi}{\mu_ {t}}\right) \right\rangle \\ + \mathbb {E} _ {\mu_ {t}} \left\langle \nabla \log \frac {\pi}{\mu_ {t}}, \nabla g ^ {*} (\nabla \log \frac {\pi}{\mu_ {t}}) - f _ {t} \right\rangle \\ \leq - \left(1 - \frac {1}{q} \lambda_ {1} ^ {q}\right) \mathbb {E} _ {\mu_ {t}} \left| \left| \nabla \log \frac {\pi}{\mu_ {t}} \right| \right| _ {q} ^ {q} \tag {55} \\ \left. + \frac {1}{p} \lambda_ {1} ^ {- p} \mathbb {E} _ {\mu_ {t}} \right| \left| \nabla g ^ {*} (\nabla \log \frac {\pi}{\mu_ {t}}) - f _ {t} \right| \Bigg | _ {p} ^ {p}. \\ \end{array}
$$
where  $\lambda_{1}$  could by any positive scalar. Here we set  $\lambda_{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  $\lambda_{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  $\eta = 1\mathrm{e - }3$ . The particle step size is 0.1.
For PFG, following Dong et al. (2023), we set the preconditioning matrix  $H = \hat{H}^{\alpha}$ , where  $\hat{H}$  is the inverse of diagonal variance of particles and  $\alpha$  is 1.0.
For Ada-GWG, we set the initial exponent  $p_0 = 2$  and learning rate  $\tilde{\eta} = 2.5\mathrm{e - }7$ .
Figure 5 shows some quantitative comparisons between different algorithms. Here Exp-GF represents GWG with  $g(\cdot) = \exp (\| \cdot \|_2^2 /(2\sigma^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.
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  $\eta = 1\mathrm{e} - 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 \in \{1.5, 2.0, 2.2\}$  and learning rate  $\tilde{\eta} = 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  $\eta = 1e - 3$  to train  $f_{w}$ . The particle step size is  $3e - 3$  for Ada-GWG and PFG.
For Ada-GWG, we set the initial exponent  $p_0 = 2.2$  and learning rate  $\tilde{\eta} = 0.001$ . The gradient of  $A(p)$  is also clipped within  $[-0.1, 0.1]$ .
For PFG, we set the preconditioning matrix  $H = \hat{H}^{\alpha}$ , where  $\hat{H}$  is the inverse of diagonal variance of particles and  $\alpha$  is chosen from  $\{0.1, 0.5, 1.0\}$  to obtain the best performance.
For SVGD, we use RBF kernel  $\exp \left(-\frac{\|x - y\|^2}{h}\right)$  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  $\tilde{\eta} = 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.
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  $\left\{\frac{1}{p}\| \cdot \| _p^p:p > 1\right\}$ , 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  $\| \cdot \| _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).