SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \nabla_ {\phi} \mathrm {E B S W} _ {p} (\mu_ {\phi}, \nu ; f) = \frac {1}{p} \left(\frac {\mathbb {E} _ {\theta \sim \sigma_ {0} (\theta)} \left[ \mathrm {W} _ {p} ^ {p} \left(\theta \sharp \mu_ {\phi} , \theta \sharp \nu\right) w _ {\mu_ {\phi^ {\prime}} , \nu , \sigma_ {0} , f} (\theta) \right]}{\mathbb {E} _ {\theta \sim \sigma_ {0} (\theta)} \left[ w _ {\mu_ {\phi^ {\prime}} , \nu , \sigma_ {0} , f} (\theta) \right]}\right) ^ {\frac {1 - p}{p}} \\ \frac {\nabla_ {\phi} \mathbb {E} _ {\theta \sim \sigma_ {0} (\theta)} \left[ \mathbf {W} _ {p} ^ {p} (\theta \sharp \mu_ {\phi} , \theta \sharp \nu) w _ {\mu_ {\phi^ {\prime}} , \nu , \sigma_ {0} , f} (\theta) \right]}{\mathbb {E} _ {\theta \sim \sigma_ {0} (\theta)} \left[ w _ {\mu_ {\phi^ {\prime}} , \nu , \sigma_ {0} , f} (\theta) \right]}, \\ \end{array}
$$
$$
\left(\frac {\mathbb {E} _ {\theta \sim \sigma_ {0} (\theta)} \left[ \mathbf {W} _ {p} ^ {p} (\theta \sharp \mu_ {\phi} , \theta \sharp \nu) w _ {\mu_ {\phi^ {\prime}} , \nu , \sigma_ {0} , f} (\theta) \right]}{\mathbb {E} _ {\theta \sim \sigma_ {0} (\theta)} \left[ w _ {\mu_ {\phi^ {\prime}} , \nu , \sigma_ {0} , f} (\theta) \right]}\right) ^ {\frac {1 - p}{p}} \approx \left(\frac {\frac {1}{L} \sum_ {l = 1} ^ {L} \left[ \mathbf {W} _ {p} ^ {p} (\theta_ {l} \sharp \mu_ {\phi} , \theta_ {l} \sharp \nu) w _ {\mu_ {\phi^ {\prime}} , \nu , \sigma_ {0} , f} (\theta_ {l}) \right]}{\frac {1}{L} \sum_ {l = 1} ^ {L} \left[ w _ {\mu_ {\phi^ {\prime}} , \nu , \sigma_ {0} , f} (\theta_ {l}) \right]}\right) ^ {\frac {1 - p}{p}},
$$
$$
\nabla_ {\phi} \mathbb {E} _ {\theta \sim \sigma_ {0} (\theta)} \left[ \mathbf {W} _ {p} ^ {p} (\theta \sharp \mu_ {\phi}, \theta \sharp \nu) w _ {\mu_ {\phi^ {\prime}}, \nu , \sigma_ {0}, f} (\theta) \right] \approx \frac {1}{L} \sum_ {l = 1} ^ {L} \left(\nabla_ {\phi} \mathbf {W} _ {p} ^ {p} (\theta_ {l} \sharp \mu_ {\phi}, \theta_ {l} \sharp \nu)\right) w _ {\mu_ {\phi}, \nu^ {\prime}, \sigma_ {0}, f} (\theta),
$$
$$
\mathbb {E} _ {\theta \sim \sigma_ {0} (\theta)} \left[ w _ {\mu_ {\phi^ {\prime}}, \nu , \sigma_ {0}, f} (\theta) \right] \approx \frac {1}{L} \sum_ {l = 1} ^ {L} w _ {\mu_ {\phi^ {\prime}}, \nu , \sigma_ {0}, f} (\theta).
$$
It is worth noting that using a copy of  $\mu_{\phi}$  does not change the value of the distance. This trick will show its true benefit when dealing with the SIR, and the MCMC methods. However, we still discuss it in the IS case for completeness. We refer to the "copy" trick is the "parameter-copy" gradient estimator while the original one is the conventional estimator.
Importance Weighted sliced Wasserstein distance. Although the IS estimation of the EBSW is not an unbiased estimation for finite  $L$ , it is an unbiased estimation of a valid distance on the space of probability measures. We refer to the distance as the importance weighted sliced Wasserstein distance (IWSW) which has the following definition.
Definition 3. For any  $p \geq 1$ , dimension  $d \geq 1$ , energy function  $f$ , a continuous proposal distribution  $\sigma_0(\theta) \sim \mathcal{P}(\mathbb{S}^{d-1})$  and two probability measures  $\mu \in \mathcal{P}_p(\mathbb{R}^d)$  and  $\nu \in \mathbb{R}^d$ , the importance weighted sliced Wasserstein (IWSW) distance is defined as follows:
$$
I W S W _ {p} (\mu , \nu ; f) = \left(\mathbb {E} \left[ \frac {\frac {1}{L} \sum_ {l = 1} ^ {L} \left[ W _ {p} ^ {p} \left(\theta_ {l} \sharp \mu , \theta_ {l} \sharp \nu\right) w _ {\mu , \nu , \sigma_ {0} , f , p} \left(\theta_ {l}\right) \right]}{\frac {1}{L} \sum_ {l = 1} ^ {L} \left[ w _ {\mu , \nu , \sigma_ {0} , f , p} \left(\theta_ {l}\right) \right]}\right)\right) ^ {\frac {1}{p}}, \tag {12}
$$
where the expectation is with respect to  $\theta_{1},\ldots ,\theta_{L}\stackrel {i.i.d}{\sim}\sigma_{0}(\theta)$ , and  $w_{\mu ,\nu ,\sigma_0,f,p}(\theta) = \frac{f(W_p^p(\theta\sharp\mu,\theta\sharp\nu))}{\sigma_0(\theta)}$
The IWSW is semi-metric, it also does not suffer from the curse of dimensionality, and it induces weak convergence. The proofs can be derived by following directly the proofs of the EBSW in Appendix A.1, Appendix A.3, and Appendix A.4. Therefore, using the IS estimation of the EBSW is as safe as the SW.
B.3 Sampling Importance Resampling and Markov Chain Monte Carlo
Algorithms. We first provide the algorithm for computing the EBSW via the SIR, the IMH, and the RMH in Algorithm 5-7.
Gradient estimators. We derive the reinforce gradient estimator of the EBSW for the SIR, the IMH, and the RHM sampling.
$$
\nabla_ {\phi} \mathrm {E B S W} _ {p} (\mu_ {\phi}, \nu ; f) = \frac {1}{p} \left(\mathbb {E} _ {\theta \sim \sigma_ {\mu_ {\phi}, \nu} (\theta ; f)} \left[ \mathrm {W} _ {p} ^ {p} (\theta \sharp \mu_ {\phi}, \theta \sharp \nu) \right]\right) ^ {\frac {1 - p}{p}} \nabla_ {\phi} \mathbb {E} _ {\theta \sim \sigma_ {\mu_ {\phi}, \nu} (\theta ; f)} \left[ \mathrm {W} _ {p} ^ {p} (\theta \sharp \mu_ {\phi}, \theta \sharp \nu) \right].
$$
$$
\nabla_ {\phi} \mathbb {E} _ {\theta \sim \sigma_ {\mu_ {\phi}, \nu} (\theta ; f)} \left[ \mathbf {W} _ {p} ^ {p} (\theta \sharp \mu_ {\phi}, \theta \sharp \nu) \right] = \mathbb {E} _ {\theta \sim \sigma_ {\mu_ {\phi}, \nu ; f} (\theta)} \left[ \mathbf {W} _ {p} ^ {p} (\theta_ {\phi} \sharp \mu , \theta \sharp \nu) \nabla_ {\phi} \log \left(\mathbf {W} _ {p} ^ {p} (\theta \sharp \mu_ {\phi}, \theta \sharp \nu) \sigma_ {\mu_ {\phi}, \nu} (\theta ; f)\right) \right]
$$
$$
\begin{array}{l} \nabla_ {\phi} \log \left(\mathrm {W} _ {p} ^ {p} (\theta \sharp \mu_ {\phi}, \theta \sharp \nu) \sigma_ {\mu_ {\phi}, \nu} (\theta ; f)\right) = \nabla_ {\phi} \log \left(\mathrm {W} _ {p} ^ {p} \theta \sharp \mu_ {\phi}, \theta \sharp \nu\right)) + \nabla_ {\phi} \log (f \left(\mathrm {W} _ {p} ^ {p} (\theta \sharp \mu_ {\phi}, \theta \sharp \nu)\right)) \\ - \nabla_ {\phi} \log \left(\int_ {\mathbb {S} ^ {d - 1}} f \left(\mathbf {W} _ {p} ^ {p} \left(\theta \sharp \mu_ {\phi}, \theta \sharp \nu\right)\right) d \theta\right) \\ = \frac {1}{\mathbf {W} _ {p} ^ {p} \left(\theta \sharp \mu_ {\phi} , \theta \sharp \nu\right))} \nabla_ {\phi} \mathbf {W} _ {p} ^ {p} \left(\theta \sharp \mu_ {\phi}, \theta \sharp \nu\right) \\ + \frac {1}{f \left(\mathbf {W} _ {p} ^ {p} \left(\theta \sharp \mu_ {\phi} , \theta \sharp \nu\right)\right)} \nabla_ {\phi} f \left(\mathbf {W} _ {p} ^ {p} \left(\theta \sharp \mu_ {\phi}, \theta \sharp \nu\right)\right) \\ - \nabla_ {\phi} \log \left(\int_ {\mathbb {S} ^ {d - 1}} f \left(\mathrm {W} _ {p} ^ {p} (\theta \sharp \mu_ {\phi}, \theta \sharp \nu)\right) d \theta\right), \\ \end{array}
$$
$$
\begin{array}{l} \nabla_ {\phi} \log \left(\int_ {\mathbb {S} ^ {d - 1}} f \left(\mathbf {W} _ {p} ^ {p} \left(\theta \sharp \mu_ {\phi}, \theta \sharp \nu\right)\right) d \theta\right) \\ = \nabla_ {\phi} \log \left(\mathbb {E} _ {\theta \sim \mathcal {U} (\mathbb {S} ^ {d - 1})} \left[ f \left(\mathbf {W} _ {p} ^ {p} \left(\theta \sharp \mu_ {\phi}, \theta \sharp \nu\right)\right) \frac {2 \pi^ {d / 2}}{\Gamma (d / 2)} \right]\right) \\ = \nabla_ {\phi} \log \left(\mathbb {E} _ {\theta \sim \mathcal {U} (\mathbb {S} ^ {d - 1})} \left[ f \left(\mathrm {W} _ {p} ^ {p} \left(\theta \sharp \mu_ {\phi}, \theta \sharp \nu\right)\right) \right]\right) \\ = \frac {1}{\mathbb {E} _ {\theta \sim \mathcal {U} \left(\mathbb {S} ^ {d - 1}\right)} \left[ f \left(\mathrm {W} _ {p} ^ {p} \left(\theta \sharp \mu_ {\phi} , \theta \sharp \nu\right)\right) \right]} \nabla_ {\phi} \mathbb {E} _ {\theta \sim \mathcal {U} \left(\mathbb {S} ^ {d - 1}\right)} \left[ f \left(\mathrm {W} _ {p} ^ {p} \left(\theta \sharp \mu_ {\phi}, \theta \sharp \nu\right) \right. \right]. \\ \end{array}
$$
Using  $L$  Monte Carlo samples from the SIR (or the IMH or the RMH) to approximate the expectation  $\mathbb{E}_{\theta \sim \sigma_{\mu_{\phi},\nu}(\theta ;f)}$ , and  $L$  samples from  $\mathcal{U}(\mathbb{S}^{d - 1})$  to approximate the expectation  $\mathbb{E}_{\theta \sim \mathcal{U}(\mathbb{S}^{d - 1})}$ , we obtain the gradient estimator of the EBSW. However, the reinforce gradient estimator is unstable in practice, especially with the energy function  $f_{e}(x) = e^{x}$ . Therefore, we propose a more simple gradient estimator which is:
$$
\nabla_ {\phi} \mathrm {E B S W} _ {p} (\mu_ {\phi}, \nu ; f) \approx \frac {1}{p} \left(\mathbb {E} _ {\theta \sim \sigma_ {\mu_ {\phi^ {\prime}}, \nu} (\theta ; f)} \left[ \mathrm {W} _ {p} ^ {p} (\theta \sharp \mu_ {\phi}, \theta \sharp \nu) \right]\right) ^ {\frac {1 - p}{p}} \mathbb {E} _ {\theta \sim \sigma_ {\mu_ {\phi^ {\prime}}, \nu} (\theta ; f)} \left[ \nabla_ {\phi} \mathrm {W} _ {p} ^ {p} (\theta \sharp \mu_ {\phi}, \theta \sharp \nu) \right].
$$
The key is to use a copy of the parameter  $\phi'$  for constructing the slicing distribution  $\sigma_{\mu_{\phi'}, \nu}(\theta; f)$ , hence, we can exchange directly the differentiation and the expectation. It is worth noting that using the copy also affects the gradient estimation, it does not change the value of the distance. We refer to the "copy" trick is the "parameter-copy" gradient estimator while the original one is the conventional estimator.
Population distance. The approximated values of  $p$ -power EBSW from using the SIR, the IMH, and the RMH can be all written as  $\frac{1}{L}\sum_{l=1}^{L}\mathbf{W}_p^p(\theta_l\sharp \mu,\theta_l\sharp \nu)$ . Here, the distributions of  $\theta_1,\ldots,\theta_L$  are different. Therefore, they are not an unbiased estimation of the EBSW $^p_p(\mu,\nu;f)$ . However, the population distance of the estimation can be defined as in Definition 4.
Definition 4. For any  $p \geq 1$ , dimension  $d \geq 1$ , energy function  $f$ , and two probability measures  $\mu \in \mathcal{P}_p(\mathbb{R}^d)$  and  $\nu \in \mathbb{R}^d$ , the projected sliced Wasserstein (PSW) distance is defined as follows:
$$
P S W _ {p} (\mu , \nu ; f) = \left(\mathbb {E} \left[ \frac {1}{L} \sum_ {l = 1} ^ {L} W _ {p} ^ {p} \left(\theta_ {l} \sharp \mu , \theta \sharp \nu\right) \right]\right) ^ {\frac {1}{p}}, \tag {13}
$$
where the expectation is with respect to  $(\theta_{1},\ldots ,\theta_{L})\sim \sigma (\theta_{1},\ldots ,\theta_{L})$  which is a distribution defined by the SIR (the IMH or the RHM).
The PSW is a valid metric since it satisfies the triangle inequality in addition to the symmetry, the non-negativity, and the identity. In particular, given three probability measures  $\mu_1, \mu_2, \mu_3 \in \mathcal{P}_p(\mathbb{R}^d)$  we have:
$$
\begin{array}{l} \mathrm {P S W} _ {p} (\mu_ {1}, \mu_ {3}) = \left(\mathbb {E} _ {(\theta_ {1: L}) \sim \sigma (\theta_ {1: L})} \left[ \frac {1}{L} \sum_ {l = 1} ^ {L} W _ {p} ^ {p} \left(\theta_ {l} \sharp \mu_ {1}, \theta_ {l} \sharp \mu_ {3}\right) \right]\right) ^ {\frac {1}{p}} \\ \leq \left(\mathbb {E} _ {\left(\theta_ {1: L}\right) \sim \sigma \left(\theta_ {1: L}\right)} \left[ \frac {1}{L} \sum_ {t = 1} ^ {L} \left(W _ {p} \left(\theta_ {l} \sharp \mu_ {1}, \theta_ {l} \sharp \mu_ {2}\right) + W _ {p} \left(\theta_ {l} \sharp \mu_ {2}, \theta_ {l} \sharp \mu_ {3}\right)\right) ^ {p} \right]\right) ^ {\frac {1}{p}} \\ \leq \left(\mathbb {E} _ {\left(\theta_ {1: L}\right) \sim \sigma \left(\theta_ {1: L}\right)} \left[ \frac {1}{L} \sum_ {t = 1} ^ {L} W _ {p} ^ {p} \left(\theta_ {l} \sharp \mu_ {1}, \theta_ {l} \sharp \mu_ {2}\right) \right]\right) ^ {\frac {1}{p}} \\ + \left(\mathbb {E} _ {\left(\theta_ {1: L}\right) \sim \sigma \left(\theta_ {1: L}\right)} \left[ \frac {1}{L} \sum_ {l = 1} ^ {T} W _ {p} ^ {p} \left(\theta_ {l} \sharp \mu_ {2}, \theta_ {l} \sharp \mu_ {3}\right) \right]\right) ^ {\frac {1}{p}} \\ = \mathrm {P S W} _ {p} \left(\mu_ {1}, \mu_ {2}\right) + \mathrm {P S W} _ {p} \left(\mu_ {2}, \mu_ {3}\right), \\ \end{array}
$$
where the first inequality is due to the triangle inequality of Wasserstein distance and the second inequality is due to the Minkowski inequality. The PSW also does not suffer from the curse of dimensionality, and it induces weak convergence. The proofs can be derived by following directly the proofs of the EBSW in Appendix A.1, Appendix A.3, and Appendix A.4. Therefore, using the SIR, the IMH, and the RMH estimation of the EBSWs are as safe as the SW.
In this section, we provide additional results for point-cloud gradient flows in Appendix C.1, color transfer in Appendix C.2, and deep point-cloud reconstruction in Appendix C.3.
C.1 Point-Cloud Gradient Flows
We provide the full experimental results including the IS-EBSW, the SIR-EBSW, the IMH-EBSW, and the RMH-EBSW with both the exponential energy function and the identity energy function in Table 3. In the table, we also include the results for the number of projections  $L = 10$ . In Table 3, we use the conventional gradient estimator for the IS-EBSW while the "parameter-copy" estimator is used for other variants of the EBSW. Therefore, we also provide the ablation studies comparing the gradient estimators in Table 4 by adding the results for the "parameter-copy" estimator for the IS-EBSW and the conventional estimator for other variants. Experimental settings are the same as in the main text.
Quantitative Results. From the two tables, we observe that the IS-EBSW is the best variant of the EBSW in both performance and computational time. Also, we observe that the exponential energy function is better than the identity energy function in this application. It is worth noting that the EBSW variants of all computational methods and energy functions are better than the baselines in
terms of Wasserstein-2 distances at the last epoch. For all sliced Wasserstein variants, we see that reducing the number of projections leads to worsening performance which is consistent with previous studies in previous works [28, 19]. In Table 3, the IS-EBSW uses the conventional gradient estimator while the SIR-EBSW, the IMH-EBSW, and the RMH-EBSW use the "parameter-copy" estimator. Therefore, we report the IS-EBSW with the "parameter-copy" estimator and the SIR-EBSW, the IMH-EBSW, and the RMH-EBSW with the Reinforce estimator (conventional estimator) in Table 4. From the table, we observe the "parameter-copy" estimator is worse than the conventional estimator in the case of IS-EBSW. For the SIR-EBSW, the IMH-EBSW, and the RMH-EBSW, we cannot use the exponential energy function due to the numerically unstable Reinforce estimator. In the case of the identity energy function, the exponential energy function is also worse than the "parameter-copy" estimator. Therefore, we recommend to use the IS-EBSW-e with the conventional gradient estimator.
Qualitative Results. We provide the visualization of the gradient flows from SW  $(\mathrm{L} = 100)$ , Max-SW  $(\mathrm{T} = 100)$ , v-DSW  $(\mathrm{L} = 10, \mathrm{T} = 10)$ , and all the EBSW-e variants in Figure 4. Overall, we see that EBSW-e variants give smoother flows than other baselines. Despite having slightly different quantitative scores due to the approximation methods, the visualization from the EBSW-e variants is consistent. Therefore, the energy-based slicing function helps to improve the convergence of the source point-cloud to the target point-cloud.
C.2 Color Transfer
Similar to the point-cloud gradient flow, we follow the same experimental settings of color transfer in the main text. We provide the full experimental results including the IS-EBSW, the SIR-EBSW, the IMH-EBSW, and the RMH-EBSW with both the exponential energy function and the identity energy function, with both  $L = 10$  and  $L = 100$ , and with both gradient estimators in Figure 5.
Results. From the figure, we observe that IMH-EBSW-e gives the best Wasserstein-2 distance among all EBSW variants. Between the exponential energy function and the identity energy function,
we see that the exponential energy function yields a better result for all EBSW variants. Similar to the gradient flow, reducing the number of projections to 10 also leads to worse results for all sliced Wasserstein variants. For the gradient estimators, the conventional estimator is preferred for the IS-EBSW while the "parameter-copy" estimator is preferred for other EBSW variants.
C.3 Deep Point-cloud Reconstruction
We follow the same experimental settings as in the main text. We provide the full experimental results including the IS-EBSW, the SIR-EBSW, the IMH-EBSW, and the RMH-EBSW with both the exponential energy function and the identity energy function, with both  $L = 10$  and  $L = 100$  in Table 5. In Table 5, we use the conventional gradient estimator for the IS-EBSW while other variants of EBSW use the "parameter-copy" gradient estimator. We also compare gradient estimators for the EBSW by adding the results for the "parameter-copy" gradient estimator for the IS-EBSW (denoted as (c)), and the conventional gradient estimator for the SIR-EBSW, the IMH-EBSW, and the RMH-EBSW in Table 6.
Quantitative Results. From the two tables, we observe that the IS-EBSW-e performs the best for both settings of the number of projections  $L = 10$  and  $L = 100$  in terms of the Wasserstein-2 reconstruction errors. For the SW reconstruction error, it is only slightly worse than the SIR-EBSW-e at epoch 200. Comparing the exponential energy function and the identity energy function, we observe that the exponential function is better in both settings of the number of projections. For the same number of projections, the EBSW variants with both types of energy function give lower errors than the baseline including the SW, the Max-SW, and the v-DSW. For all sliced Wasserstein variants, a higher value of the number of projections gives better results. For the gradient estimator of the EBSW, we see that the conventional gradient estimator is preferred for the IS-EBSW while the "parameter-copy" estimator is preferred for other EBSW variants.
Qualitative Results. We show some ground-truth point-clouds ModelNet40 and their corresponding reconstructed point-clouds from different models ( $L = 100$ ) at epochs 200 and 20 in Figure 6-7 respectively. From the top to the bottom is the ground truth, the SW, the Max-SW, the v-DSW, the IS-EBSW-e, the SIR-EBSW-e, the IMH-EBSW-e, and the RMH-EBSW-e.
C.4 Deep Generative Modeling
We follow the same setup as sliced Wasserstein deep generative models in [29]. We compare IS-EBSW-e with SW, and DSW [30] on CIFAR10 (image size 32x32) [21], and CelebA [23]. We show both FID score [15] (CIFAR10, CelebA) and IS score [39] (CIFAR10) in Figure 8. Overall, IS-EBSW-e performs better than SW and DSW in this task which suggests that EBSW has a wide range of applications.
D Computational Infrastructure
For the point-cloud gradient flows and the color transfer, we use a MacBook Pro M1 for conducting experiments. For deep point-cloud reconstruction, experiments are run on a single NVIDIA V100 GPU.