SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Khai Nguyen
Department of Statistics and Data Sciences
khainb@utexas.edu
Nhat Ho
minhnhat@utexas.edu
The sliced Wasserstein (SW) distance has been widely recognized as a statistically effective and computationally efficient metric between two probability measures. A key component of the SW distance is the slicing distribution. There are two existing approaches for choosing this distribution. The first approach is using a fixed prior distribution. The second approach is optimizing for the best distribution which belongs to a parametric family of distributions and can maximize the expected distance. However, both approaches have their limitations. A fixed prior distribution is non-informative in terms of highlighting projecting directions that can discriminate two general probability measures. Doing optimization for the best distribution is often expensive and unstable. Moreover, designing the parametric family of the candidate distribution could be easily misspecified. To address the issues, we propose to design the slicing distribution as an energy-based distribution that is parameter-free and has the density proportional to an energy function of the projected one-dimensional Wasserstein distance. We then derive a novel sliced Wasserstein variant, energy-based sliced Wasserstein (EBSW) distance, and investigate its topological, statistical, and computational properties via importance sampling, sampling importance resampling, and Markov Chain methods. Finally, we conduct experiments on point-cloud gradient flow, color transfer, and point-cloud reconstruction to show the favorable performance of the EBSW<sup>1</sup>.
The sliced Wasserstein [2] (SW) distance is a sliced probability metric that is derived from the Wasserstein distance [5, 36] as the base metric. Utilizing the closed-form solution of optimal transport on one-dimension [36], the SW distance can be computed very efficiently at the time complexity of  $\mathcal{O}(n\log n)$  and the space complexity of  $\mathcal{O}(n)$  when dealing with two probability measures that have at most  $n$  supports. Moreover, the sample complexity of the SW is only  $\mathcal{O}(n^{-1/2})$  [27, 34] which indicates that it does not suffer from the curse of dimensionality in statistical inference. Therefore, the SW distance has been applied to various domains of applications including point-cloud applications e.g., reconstruction, registration, generation, and upsampling [31, 41], generative models [8], domain adaptation [22], clustering [20], gradient flows [24, 1], approximate Bayesian computation [26], variational inference [47], and many other tasks.
The SW distance can be defined as the expectation of the one-dimensional Wasserstein distance between two projected probability measures. The randomness comes from a random projecting direction which is used to project two original probability measures to one dimension. The probability distribution of the random projecting direction is referred to as the slicing distribution. Therefore, a central task that decides the effectiveness of the SW in downstream applications is designing
$^{1}$ Code for this paper is published at https://github.com/khainb/EBSW.
slicing distribution. The conventional sliced Wasserstein distance [2] simply takes the uniform distribution over the unit-hypersphere as the slicing distribution. Despite being easy to sample from, the uniform distribution is not able to differentiate between informative and non-informative projecting distributions in terms of discriminating two interested probability measures through projection [7]. To avoid a flat prior distribution like the uniform distribution, a different approach tries to find the best slicing distribution that can maximize the expectation. This distribution is found inside a parametric family of distribution over the unit-hypersphere [30, 32]. However, searching for the best slicing distribution often requires an iterative procedure which is often computationally expensive and unstable. Moreover, choosing the family for the slicing distribution is challenging since the number of distributions over the unit-hypersphere is limited. Widely used and explicit spherical distributions such as von Mises-Fisher distribution [17, 32] might be misspecified while implicit distributions [30] are expensive, hard to adapt to downstream applications and uninterpretable.
In this paper, we aim to develop new choices of slicing distributions that are both discriminative in comparing two given probability measures and do not require optimization. Motivated by energy-based models [25, 16], we model the slicing distribution by an unnormalized density function which gives a higher density for a more discriminative projecting direction. To induce that property, the density function at a projecting direction is designed to be proportional to the value of the one-dimensional Wasserstein distance between the two corresponding projected probability measures.
Contribution. In summary, our contributions are three-fold:
1. We propose a new class of slicing distribution, named energy-based slicing distribution, which has the density function proportional to the value of the projected one-dimensional Wasserstein distance. We further control the flexibility of the slicing distribution by applying an energy function e.g., the polynomial function, and the exponential function, to the projected Wasserstein value. By using the energy-based slicing distribution, we derive a novel metric on the space of probability measures, named energy-based sliced Wasserstein (EBSW) distance.
2. We derive theoretical properties of the proposed EBSW including topological properties, statistical properties, and computational properties. For topological properties, we first prove that the EBSW is a valid metric on the space of probability measures. After that, we show that the weak convergence of probability measures is equivalent to the convergence of probability measures under the EBSW distance. Moreover, we develop the connection of the EBSW to existing sliced Wasserstein variants and the Wasserstein distance. We show that the EBSW is the first non-optimization variant that is an upper bound of the sliced Wasserstein distance. For the statistical properties, we first derive the sample complexity of the EBSW which indicates that it does not suffer from the curse of dimensionality. For computational properties, we propose importance sampling, sampling importance resampling, and Markov Chain Monte Carlo methods to derive empirical estimations of the EBSW. Moreover, we discuss the time complexities and memory complexities of the corresponding estimations. Finally, we discuss the statistical properties of estimations.
3. We apply the EBSW to various tasks including gradient flows, color transfer, and point-cloud applications. According to the experimental result, the EBSW performs better than existing projection-selection sliced Wasserstein variants including the conventional sliced Wasserstein [2] (SW), max sliced Wasserstein [7] (Max-SW), and distributional sliced Wasserstein (DSW) [30, 32]. More importantly, the importance sampling estimation of the EBSW is as efficient and easy to implement as the conventional SW, i.e., its implementation can be obtained by adding one to two lines of code.
Organization. The remainder of the paper is organized as follows. We first review the background on the sliced Wasserstein distance and its projection-selection variants in Section 2. We then define the energy-based sliced Wasserstein distance, derive their theoretical properties, and discuss its computational methods in Section 3. Section 4 contains experiments on gradient flows, color transfer, and point-cloud applications. We conclude the paper in Section 5. Finally, we defer the proofs of key results, and additional materials in the Appendices.
Notations. For any  $d \geq 2$ , we denote  $\mathbb{S}^{d-1} := \{\theta \in \mathbb{R}^d \mid ||\theta||_2^2 = 1\}$  and  $\mathcal{U}(\mathbb{S}^{d-1})$  as the unit hyper-sphere and its corresponding uniform distribution. We denote  $\mathcal{P}(\mathcal{X})$  as the set of all probability measures on the set  $\mathcal{X}$ . For  $p \geq 1$ ,  $\mathcal{P}_p(\mathcal{X})$  is the set of all probability measures on the set  $\mathcal{X}$  that have finite  $p$ -moments. For any two sequences  $a_n$  and  $b_n$ , the notation  $a_n = \mathcal{O}(b_n)$  means that  $a_n \leq C b_n$  for all  $n \geq 1$ , where  $C$  is some universal constant. We denote  $\theta \sharp \mu$  is the push-forward measures of  $\mu$
through the function  $f: \mathbb{R}^d \to \mathbb{R}$  that is  $f(x) = \theta^\top x$ . For a vector  $X \in \mathbb{R}^{dm}$ ,  $X := (x_1, \ldots, x_m)$ ,  $P_X$  denotes the empirical measures  $\frac{1}{m} \sum_{i=1}^{m} \delta_{x_i}$ .
In this section, we first review the sliced Wasserstein distance and its projection-selection variants including max sliced Wasserstein distance and distributional sliced Wasserstein distance.
Sliced Wasserstein. The definition of sliced Wasserstein (SW) distance [2] between two probability measures  $\mu \in \mathcal{P}_p(\mathbb{R}^d)$  and  $\nu \in \mathcal{P}_p(\mathbb{R}^d)$  is:
$$
\mathrm {S W} _ {p} (\mu , \nu) = \left(\mathbb {E} _ {\theta \sim \mathcal {U} \left(\mathbb {S} ^ {d - 1}\right)} \left[ \mathrm {W} _ {p} ^ {p} \left(\theta_ {\#} ^ {\#} \mu , \theta_ {\#} ^ {\#} \nu\right) \right]\right) ^ {\frac {1}{p}}, \tag {1}
$$
where the Wasserstein distance has a closed form which is  $\mathbf{W}_p^p (\theta \sharp \mu ,\theta \sharp \nu) = \int_0^1 |F_{\theta \sharp \mu}^{-1}(z) - F_{\theta \sharp \nu}^{-1}(z)|^p dz$  where  $F_{\theta \sharp \mu}$  and  $F_{\theta \sharp \nu}$  are the cumulative distribution function (CDF) of  $\theta \sharp \mu$  and  $\theta \sharp \nu$  respectively. However, the expectation in the definition of the SW distance is intractable to compute. Therefore, the Monte Carlo scheme is employed to approximate the value:
$$
\widehat {\mathbf {S W}} _ {p} (\mu , \nu ; L) = \left(\frac {1}{L} \sum_ {l = 1} ^ {L} \mathbf {W} _ {p} ^ {p} \left(\theta_ {l} \sharp \mu , \theta_ {l} \sharp \nu\right)\right) ^ {\frac {1}{p}}, \tag {2}
$$
where  $\theta_{1},\ldots ,\theta_{L}\stackrel {i,i,d}{\sim}\mathcal{U}(\mathbb{S}^{d - 1})$  and are referred to as projecting directions. The pushforward measures  $\theta_1\sharp \mu ,\dots ,\theta_L\sharp \mu$  are called projections of  $\mu$  (similar to  $\nu$ ). The number of Monte Carlo samples  $L$  is often referred to as the number of projections. When  $\mu$  and  $\nu$  are discrete measures that have at most  $n$  supports, the time complexity and memory complexity of the SW are  $\mathcal{O}(Ln\log n)$  and  $\mathcal{O}(L(d + n))$  respectively. It is worth noting that  $\widehat{\mathrm{SW}}_p^p (\mu ,\nu ;L)$  is an unbiased estimation of  $\mathrm{SW}_p^p (\mu ,\nu)$ , however,  $\widehat{\mathrm{SW}}_p(\mu ,\nu ;L)$  is only asymptotically unbiased estimation of  $\mathrm{SW}_p(\mu ,\nu)$ . Namely, we have  $\widehat{\mathrm{SW}}_p(\mu ,\nu ;L)\to \mathrm{SW}_p(\mu ,\nu)$  when  $L\rightarrow \infty$  (law of large numbers).
Distributional sliced Wasserstein. As discussed, using the uniform distribution over projecting directions is not suitable for two general probability measures. A natural extension is to replace the uniform distribution with a "better" distribution on the unit-hypersphere. Distributional sliced Wasserstein [30] suggests searching this distribution in a parametric family of distributions by maximizing the expected distance. The definition of distributional sliced Wasserstein (DSW) distance [30] between two probability measures  $\mu \in \mathcal{P}_p(\mathbb{R}^d)$  and  $\nu \in \mathcal{P}_p(\mathbb{R}^d)$  is:
$$
\mathrm {D S W} _ {p} (\mu , \nu) = \max  _ {\psi \in \Psi} \left(\mathbb {E} _ {\theta \sim \sigma_ {\psi} (\theta))} \left[ \mathrm {W} _ {p} ^ {p} \left(\theta \sharp \mu , \theta \sharp \nu\right) \right]\right) ^ {\frac {1}{p}}, \tag {3}
$$
where  $\sigma_{\psi}(\theta) \in \mathcal{P}(\mathbb{S}^{d-1})$ , e.g., von Mises-Fisher [17] distribution with unknown location parameter  $\sigma_{\psi}(\theta) := \mathrm{vMF}(\theta|\epsilon, \kappa)$ ,  $\psi = \epsilon$  [32]. After using  $T \geq 1$  (projected) stochastic (sub)-gradient ascent iterations to obtain an estimation of the parameter  $\hat{\psi}_T$ , Monte Carlo samples  $\theta_1, \ldots, \theta_L \stackrel{i.i.d.}{\sim} \sigma_{\hat{\psi}_T}(\theta)$  are used to approximate the value of the DSW. Interestingly, the metricity DSW holds for non-optimal  $\hat{\psi}_T$  as long as  $\sigma_{\hat{\psi}_T}(\theta)$  are continuous on  $\mathbb{S}^{d-1}$  e.g., vMF with  $\kappa < \infty$  [31]. In addition, the unbiasedness property of the DSW is the same as the SW, namely, when  $L \to \infty$ , the empirical estimation of the DSW converges to the true value. The time complexity and space complexity of the DSW are  $\mathcal{O}(LTn\log n)$  and  $\mathcal{O}(L(d + n))$  in turn without counting the complexities of sampling from  $\sigma_{\hat{\psi}_T}(\theta)$ . We refer to Appendix B.1 for more details, e.g., equations, algorithms, and discussion. PAC-Bayesian generalization bounds for DSW are investigated in [35].
Max sliced Wasserstein. By letting the concentration parameter  $\kappa \rightarrow \infty$ , the vMF distribution degenerates to the Dirac distribution  $\mathrm{vMF}(\theta|\epsilon, \kappa) \to \delta_{\epsilon}$ , we obtain the max sliced Wasserstein distance [7]. The definition of max sliced Wasserstein (Max-SW) distance between two probability measures  $\mu \in \mathcal{P}_p(\mathbb{R}^d)$  and  $\nu \in \mathcal{P}_p(\mathbb{R}^d)$  is:
$$
\operatorname {M a x} - \mathrm {S W} _ {p} (\mu , \nu) = \max  _ {\theta \in \mathbb {S} ^ {d - 1}} \mathrm {W} _ {p} \left(\theta \sharp \mu , \theta \sharp \nu\right). \tag {4}
$$
Similar to the DSW, the Max-SW is often computed by using  $T \geq 1$  iterations of (projected) (sub)-gradient ascent to obtain an estimation of the "max" projecting direction  $\hat{\theta}_T$ . After that, the
estimated value of the Max-SW is  $\mathrm{W}_p(\hat{\theta}_T\sharp \mu ,\hat{\theta}_T\sharp \nu)$ . The time complexity and space complexity of the Max-SW are  $\mathcal{O}(Tn\log n)$  and  $\mathcal{O}(d + n)$ . It is worth noting that the Max-SW is only a metric at the global optimum  $\theta^{\star}$ , hence, we cannot guarantee the metricity of the Max-SW due to the non-convex optimization [34] problem even when  $T\to \infty$ . Therefore, the performance of Max-SW is often unstable in practice [28]. We refer the reader to Appendix B.1 for more details e.g., equations, algorithms, and discussions about the Max-SW.
3 Energy-Based Sliced Wasserstein Distance
From the background, we observe that using a fixed slicing distribution e.g., in the SW, is computationally efficient but might be not effective. In contrast, using optimization-based slicing distributions is computationally expensive e.g., in the DSW, and is unstable e.g., in the Max-SW. Therefore, we address previous issues by introducing a novel sliced Wasserstein variant that uses optimization-free slicing distribution which can highlight the difference between two comparing probability measures.
3.1 Energy-Based Slicing Distribution
We first start with the key contribution which is the energy-based slicing distribution.
Definition 1. For any  $p \geq 1$ , dimension  $d \geq 1$ , an energy function  $f: [0, \infty) \to \Theta \subset (0, \infty)$  and two probability measures  $\mu \in \mathcal{P}_p(\mathbb{R}^d)$  and  $\nu \in \mathcal{P}_p(\mathbb{R}^d)$ , the energy-based slicing distribution  $\sigma_{\mu, \nu}(\theta)$  supported on  $\mathbb{S}^{d-1}$  is defined as follows:
$$
\sigma_ {\mu , \nu} (\theta ; f, p) \propto f \left(W _ {p} ^ {p} \left(\theta \sharp \mu , \theta \sharp \nu\right)\right) := \frac {f \left(W _ {p} ^ {p} \left(\theta \sharp \mu , \theta \sharp \nu\right)\right)}{\int_ {\mathbb {S} ^ {d - 1}} f \left(W _ {p} ^ {p} \left(\theta \sharp \mu , \theta \sharp \nu\right)\right) d \theta}, \tag {5}
$$
where the image of  $f$  is in the open interval  $(0, \infty)$  is for making  $\sigma_{\mu, \nu}(\theta)$  continuous on  $\mathbb{S}^{d-1}$ .
In contrast to the approach of the DSW which creates the dependence between the slicing distribution and two input probability measures via optimization, the energy-based slicing distribution obtains the dependence by exploiting the value of the projected Wasserstein distance between two input probability measures at each support.
Monotonically increasing energy functions. Similar to previous works, we again assume that "A higher value of projected Wasserstein distance, a better projecting direction". Therefore, it is natural to use a monotonically increasing function for the energy function  $f$ . We consider the following two functions: the exponential function:  $f_{e}(x) = e^{x}$ , and the shifted polynomial function:  $f_{q}(x) = x^{q} + \varepsilon$  with  $q, \varepsilon > 0$ . The shifted constant  $\varepsilon$  helps to avoid the slicing distribution undefined when two input measures are equal. In a greater detail, when  $\mu = \nu$ , we have  $\mathsf{W}_p^p (\theta \sharp \mu ,\theta \sharp \nu) = 0$  for all  $\theta \in \mathbb{S}^{d - 1}$  due to the identity property of the Wasserstein distance. Hence,  $\sigma_{\mu ,\nu}(\theta ;f,p)\propto 0$  for  $\theta \in \mathbb{S}^{d - 1}$  and  $f(x) = x^{q} (q > 0)$ . Therefore, the slicing distribution  $\sigma_{\mu ,\nu}(\theta ;f,p)$  is undefined due to an invalid density function. In practice, it is able to set  $\varepsilon = 0$  since we rarely deal with two coinciding measures.
Other energy functions. We can choose any positive function for energy function  $f$  and it will result in a valid slicing distribution. However, it is necessary to come up with an assumption for the choice of the function. Since there is no existing other assumption for the importance of projecting direction, we will leave the investigation of non-increasing energy function  $f$  to future works.
Example 1. Let  $\mu = \mathcal{N}(\mathbf{m}, v_1^2\mathbf{I})$  and  $\nu = \mathcal{N}(\mathbf{m}, v_2^2\mathbf{I})$  are two location-scale Gaussian distributions with the same means, we have their projections are  $\theta \sharp \mu = \mathcal{N}(\theta^\top \mathbf{m}, v_1^2)$  and  $\theta \sharp \nu = \mathcal{N}(\theta^\top \mathbf{m}, v_2^2)$ . Based on the closed form of the Wasserstein distance between two Gaussians [10], we have  $W_2^2 (\theta \sharp \mu ,\theta \sharp \nu) = (v_1 - v_2)^2$  for all  $\theta \in \mathbb{S}^{d - 1}$  which leads to  $\sigma_{\mu ,\nu}(\theta ;f,p) = \mathcal{U}(\mathbb{S}^{d - 1})$  for definitions of energy function  $f$  in Definition 1.
Example 1 gives a special case where we can have the closed form of the slicing function.
Applications to other sliced probability metrics and mutual information. In this paper, we focus on comparing choices of slicing distribution in the basic form of the SW distance. The proposed energy-based slicing distribution can be adapted to other variants of the SW that are not about designing slicing distribution e.g., non-linear projecting [19], orthogonal projecting directions [38], and so on. Moreover, the energy-based slicing approach can be applied to other sliced probability metrics e.g., sliced score matching [42], and sliced mutual information [13].
3.2 Definitions, Topological, and Statistical Properties of Energy Based Sliced Wasserstein
With the definition of energy-based slicing distribution in Definition 1, we now are able to define the energy-based sliced Wasserstein (EBSW) distance.
Definition 2. For any  $p \geq 1$ , dimension  $d \geq 1$ , two probability measures  $\mu \in \mathcal{P}_p(\mathbb{R}^d)$  and  $\nu \in \mathcal{P}_p(\mathbb{R}^d)$ , the energy function  $f:[0,\infty) \to (0,\infty)$ , and the energy-based slicing distribution  $\sigma_{\mu,\nu}(\theta; f, p)$ , the energy-based sliced Wasserstein (EBSW) distance is defined as follows:
$$
E B S W _ {p} (\mu , \nu ; f) = \left(\mathbb {E} _ {\theta \sim \sigma_ {\mu , \nu} (\theta ; f, p)} \left[ W _ {p} ^ {p} \left(\theta \sharp \mu , \theta \sharp \nu\right) \right]\right) ^ {\frac {1}{p}}. \tag {6}
$$
We now derive some theoretical properties of the EBSW distance.
Topological Properties. We first investigate the metricity of the EBSW distance.
Theorem 1. For any  $p \geq 1$ , energy-function  $f$ , the energy-based sliced Wasserstein  $EBSW_{p}(\cdot, \cdot; f)$  is a semi-metric in the probability space on  $\mathbb{R}^d$ , namely EBSW satisfies non-negativity, symmetry, and identity of indiscernibles.
The proof of Theorem 1 in given in Appendix A.1. Next, we establish the connections among the EBSW, the SW, the Max-SW, and the Wasserstein.
Proposition 1. (a) For any  $p \geq 1$  and increasing energy function  $f$ , we find that
$$
S W _ {p} (\mu , \nu) \leq E B S W _ {p} (\mu , \nu ; f).
$$
The equality holds when  $f(x) = c$  for some positive constant  $c$  for all  $x \in [0, \infty)$ .
(b) For any  $p \geq 1$  and energy function  $f$ , we have
$$
E B S W _ {p} (\mu , \nu ; f) \leq M a x - S W _ {p} (\mu , \nu) \leq W _ {p} (\mu , \nu).
$$
Proof of Proposition 1 is in Appendix A.2. The results of Proposition 1 indicate that for increasing energy function  $f$ , the EBSW is lower bounded by the SW while it is upper bounded by the Max-SW. It is worth noting that the EBSW is the first variant that changes the slicing distribution while still being an upper bound of the SW.
Theorem 2. For any  $p \geq 1$  and energy function  $\bar{f}$ , the convergence of probability measures under the energy-based sliced Wasserstein distance  $EBSW_{p}(\cdot, \cdot; \bar{f})$  implies weak convergence of probability measures and vice versa.
Theorem 2 implies that for any sequence of probability measures  $(\mu_k)_{k \in \mathbb{N}}$  and  $\mu$  in  $\mathcal{P}_p(\mathbb{R}^d)$ ,  $\lim_{k \to +\infty} \mathrm{EBSW}_p(\mu_k, \mu; \bar{f}) = 0$  if and only if for any continuous and bounded function  $f: \mathbb{R}^d \to \mathbb{R}$ ,  $\lim_{k \to +\infty} \int f \, \mathrm{d}\mu_k = \int f \, \mathrm{d}\mu$ . The proof of Theorem 2 is in Appendix A.3.
Statistical Properties. From Proposition 1, we derive the sample complexity of the EBSW.
Proposition 2. Let  $X_{1}, X_{2}, \ldots, X_{n}$  be i.i.d. samples from the probability measure  $\mu$  being supported on compact set of  $\mathbb{R}^d$ . We denote the empirical measure  $\mu_{n} = \frac{1}{n}\sum_{i=1}^{n}\delta_{X_{i}}$ . Then, for any  $p \geq 1$  and energy function  $f$ , there exists a universal constant  $C > 0$  such that
$$
\mathbb {E} \left[ E B S W _ {p} \left(\mu_ {n}, \mu ; f\right) \right] \leq C \sqrt {(d + 1) \log n / n},
$$
where the outer expectation is taken with respect to the data  $X_{1},X_{2},\ldots ,X_{n}$
The proof of Proposition 2 is given in Appendix A.4. From this proposition, we can say that the EBSW does not suffer from the curse of dimensionality. We will discuss other statistical properties of approximating the EBSW in the next section.
3.3 Computational Methods and Computational Properties
Calculating the expectation with respect to the slicing distribution  $\sigma_{\mu,\nu}(\theta; f, p)$  is intractable. Therefore, we propose some Monte Carlo estimation methods to approximate the value of EBSW.
3.3.1 Importance Sampling
The most simple and computationally efficient method that can be used is importance sampling (IS) [18]. The idea is to utilize an efficient-sampling proposal distribution  $\sigma_0(\theta)\in \mathcal{P}(\mathbb{S}^{d - 1})$  to provide Monte Carlo samples. After that, we use the density ratio between the original slicing distribution and the proposal distribution to weight samples. We can rewrite the EBSW distance as:
$$
\operatorname {E B S W} _ {p} (\mu , \nu ; f) = \left(\frac {\mathbb {E} _ {\theta \sim \sigma_ {0} (\theta)} \left[ \mathbf {W} _ {p} ^ {p} \left(\theta \sharp \mu , \theta \sharp \nu\right) w _ {\mu , \nu , \sigma_ {0} , f , p} (\theta) \right]}{\mathbb {E} _ {\theta \sim \sigma_ {0} (\theta)} \left[ w _ {\mu , \nu , \sigma_ {0} , f , p} (\theta) \right]}\right) ^ {\frac {1}{p}}, \tag {7}
$$