Energy-Based Sliced Wasserstein Distance
Authors: Khai Nguyen
Nhat Ho
Venue: NeurIPS 2023 poster (2023)
Keywords: Sliced Wasserstein
Monte Carlo Methods
Point-Cloud
Optimal Transport
TL;DR: We propose new sliced Wasserstein variants that have the density of the slicing distribution proportional to  an energy-function of the one-dimensional projected Wasserstein distance.

Abstract: 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 Waserstein (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.

Energy-Based Sliced Wasserstein Distance
Khai Nguyen
Nhat Ho Department of Statistics and Data Sciences The University of Texas at Austin Austin, TX 78712
Department of Statistics and Data Sciences The University of Texas at Austin Austin, TX 78712
Khainboutexas edu
rinhnhat@utexas edu
Abstract
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 distri- bution which belongs to 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 enrgy-based distribution that is parameter-fi function of the projected one-dimensions Wasserstein distance. We then derive a novel sliced Wasserstein variant, energy-basel sliced Waserstein (EBSW) distance, and investigate its topological, statistical, and computational properties via im- portance 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
and has the density proportional to an energy
Introduction
The sliced Wasserstein 2 (SW) distance is sliced probability metric that is derived from the Wasserstein distance 36 as base metric, Utilizing the closed-form solution of optimal transport
one-dimension (36]. the SW distance be computed very efficiently at the time complexity of On logn) and the space complexity of (n)) when dealing with two probability measures that have at most n supports. Moreover, the sample complexity of the SW is only O(n -112) [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 [1] 41) generative models (8), domain adaptation (22), clustering [20], gradient flows [24] I. approximate Bayesian computation [26)..
variational inference 477. 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 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 the slicing distribution. Therefore, a central task that decides the effectiveness of the SW in downstream applications is designing
Cooee for this paper is published at htttp:://gttaab..co////aaaa/////
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

slicing distribution. The conventional sliced Wasserstein distance [2] simply takes the uniform distribution over the umit-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 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 mntt-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 [6] we model the slicing distribution by unnormalized density function which gives a higher density for more discriminative projecting direction. To induce that property, the density function at 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 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 valid metric 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, 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 77] (Max-SW), and distributional sliced Wasserstein (DSW) 30, 32 More
importantly, the importance sampling estimation of the EBSW is as efficient and easy 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 follows. We first review the background on the sliced Wasserstein distance and its the energy- based sliced Wasserstein distance, derive their theoretical properties, and discuss its computational methods in Section 3 Scction 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.
variants in Section [2] We then define
|011 = 1} and (($d-1) as the unit hyper-sphere and its corresponding uniform distribution. We denote P(x') as the set of all probability measures on the set X. For p 2 1, (X)))) is the set of all probability measures on the set X' that have finite p-moments. For any two sequences a, and bn, the notation a = (b,,) means that a, S Cb2 for all n > 1, where C is some universal constant. We denote 0#u is the worh-forward measures of μ
Notations. For any d' 2 2, we denote Sd-1 == {0 E Rd

---

through the function f: - R that is f(x) = =0 For vector X E R x == (II,...,,Tml Px denotes the empirical measures m 2i=1 dr.
2 Background
In this section, we first review the sliced Wasserstein distance and its projection-selectee variants including max sliced Wasserstein distance and distributional sliced Wasserstein distance.
Sliced Wasserstein, The definition of sliced Wasserstein (SW) distance measures 1a E Ppp (IR.) and E (RR) is:
between two probability
SW sW,(u, v) = (Eo- (a (0 0dw)))
where the Wasserstein distance has closed form which is (e 02v) So (Fog)()) Foro (z) where Fern and Fezz are the cumulative distribution function (CDF) of #ij and 0:v respectively. However, the expectation in the definition of the SW distance is intractable to compute. Therefore, the Monte Carlo scheme is employed approximate the value:
(g(0ituBBitt)))) Bithr)
SW (Wp(uv;L)
where fi..... .0.......... 0r. 1.dd 1 and are referred to as projecting directions. The pushfoward measures Oisp,.... Or.3pu are called projections of μ (similar to v/). The number of Monte Carlo samples L, is often referred to as the number of projections. When and are discrete measures that have at most n supports, the time complexity and memory complexity of the SW are O(In logn) and ((L(d + n)) respectively. It is worth noting that SW (14, vz L) is an unbiased estimation of (w, v). however, (1, v5 L) is only asymptotically unbiased estimation of SW (1,, v). Namely, we have (W,(u, L) > SW,(u, when - oo (law of large numbers).
Distributional sliced Wasserstein. As discussed, using the uniform distribution over projecting directions not suitable for two general probability measures. natural extension is 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 maxi- mizing the expected distance. The definition of distributional sliced Wasserstein (DSW) distance 30)
between two probability measures μ E Pp(()) and E (RRR)) is:
DSW (1, max (Eg (9)) (WB(tt), 0dv))) NEW
where o. (0) E P(S e.g., von Mises-Fisher [1 distribution with unknown location parameter (0) = vMF(0]e, K), # (32] After using T 2 1 (projected) stochastic (sub)-gradient ascent iterations obtain estimation of the parameter ,e. Monte Carlo samples 01.-.., Or. i.i.d o J (0)
used to approximate the value of the DSW. Interestingly, the metricity DSW holds for non- optimal e as long as ب (0) are continuous 1 e.g., vMF with k oo [II In addition, the unbiasedness property of the DSW is the same as the SW. namely, when + oo, the empirical estimation of the DSW converges to the true value. The time complexity and space complexity of the DSW are (LTn log n) and O(L(d + n)) turn without counting the complexities of sampling from σ Se (0). We refer to Appendix B.I for more details, e.g., equations, algorithms, and discussion. PAC-Bayesine generalization bounds for DSW are investigated in [35].
Max sliced Wasserstein. By letting the concentration parameter K. - Do, the vMF distribution degenerates to the Dirac distribution vMF(0]e, k) - &,, we obtain the max sliced Wasserstein distance 77]. The definition of max sliced Wasserstein (Max-SW) distance between two probability measures μ E P,, (RR") and E (RR) is:
Max-SW, (Wwpp.)) v)= max W (p 0tv).
#8
Similar to the DSW, the Max-SW is often computed by using T 2 1 iterations of (projected) (sub)-gradient ascent to obtain an estimation of the "max" projecting direction 0r. After that, the
3

---

estimated value of the Max-SW is W, (or0 (y) The time complexity and space complexity of the Max-SW are OTTn n) and O(d+N) is worth noting that the Max-SW is only metric at the global optimum 0*, hence, we cannot guarantee the metricity of the Max-SW due to the non-convex optimization 34 problem even when T -> co. 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.
Energy-Based Sliced Wasserstein Distance
From the background, we observe that using fixed slicing distribution e.g.. in the SW, is computa- tionally efficient but might be not effective. In contrast, using optimization- -based slicing distributions is computationally expensive c.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 2 1, dimension d 2 1, an energy Sunction J [0, x)) + θ c (0, oo) and two probability measures μ. E (P)( and v E Pp (P(( the energy-based slicing distribution on (0) supported on gd-1 is defined as follow:
f F(W 02v))
op (0: f.p) x ff(fff/ww (OWp 0tv))
(5)
Sga-1 f (WPp(12.. (zv))dd
where the image of f is in the open interval (0, oc) is for making (0) continuous on Sd-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 monotonically increasing function for the energy function f. We consider the following two functions: the exponential function: fa(x) e", and the shifted polynomial function: fz) x + with q, > 0. The shifted constant helps to avoid the slicing distribution undefined when two input measures are equal. In a greater detail, when μ = v, we have (D 0+v) = for all 0 E Sd-1 due to the identity property of the Wasserstein distance. Hence, Ju,, (0; f. p) x O for 0 E Sd-1 and fx) = r (q > 0). Therefore, the slicing distribution (0, density function. In practice, it is able to set e = 0 since we rarely deal with two coinciding measures.
f.p) [0;
is undefined due an invalid
We choose any positive function for energy function f and it will result in a valid slicing distribution. However, is necessary to come up with an assumption for the choice of the function. Since there is existing other assumption for the importance of projecting direction, we will leave the investigation of non-increasing energy function f to future works.
Other energy functions.
Example Let μ (m,, II and ~(m,, v)I) are two location-scale scale Gaussian distributions with the same means, have their projections are f4p = (( m, vf) and 0iv ~N m, m,
Based on the closed form of the Wasserstein distance between two Gaussians (10), we have
(01p, (tv) (vi 2 for all 0 E 3d-1 which leads (0:
u d for defi-
nitions of energy function f in Definition
Example T gives special case where can have the closed form of the slicing function.
Applications 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 energy-based sliced Wasserstein (EBSW) distance.
we now are able to define the |
Definition 2. For any p 2
dimension d two probability measures P Pp(RR)) and Po ), the energy function f [0, x0) + oo), and the energy-based slicing distribution Fav(0: f f,p), the energy-based sliced Wasserstein (EBSW) distance is defined as follows:
EBSW v5 (0:..).. [Wp( 01v)))*
(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 2 1, energy-function f, the energy-based sliced Wasserstein EB a semi- metime the probability space on RR, namely EBSW satisfies non-negativity, symmetry, and identity of indiscernibles.
The proof of Theorem Min given in Appendix A.I Next, we establish the connections among the EBSW, the SW, the Max-SW, and the Wasserstein.
Proposition (a) For any p 2 and increasing energy function f. we find that
SW,(,, S. EBSW, (p, (H,v; v; f).
The equality holds when f(xx) c for some positive constant efor all [0, o0).
(b) For any p 2 1 and energy function f. we have
EBSW, 2; f) Max-SW, S W)(,.v).
Proof of Proposition is in Appendix A.2 The results of Proposition Lindicate 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 2 1 and energy function f. the convergence of probability measures under the energy-based sliced Wasserstein distance EBSW,(., (: implies weak convergence of probability measures and vice versa.
Theorem implies that for any sequence of probability measures (μK))EN and in P (R limg-x+to EBSW (1R, fe; f) = if and only if for any continuous and bounded function f Rdd - R,
limj-++0o f dpp = f dp. The proof of Theorem is in Appendix A.3
Statistical Properties. From Proposition| D we derive the sample complexity of the EBSW.
Proposition 2.
Let X1, X2..., Xn, be i.i.d. samples from the probability measure being supported
on compact set of TRd. We denote the empirical measure Phn n Ex.1 8x,- Then, for any p 2 1 and energy function f. there exists a universal constant C such that
E|EBSW P5 f) < C Cvd 1) logn//n
where the outer expectation is taken with respect to the data X1, X2, Xn.
The proof of Proposition 2 is given in Appendix A.4 From this proposition, 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 (0; f.p) is intractable. There- fore, we propose some Monte Carlo estimation methods to approximate the value of EBSW.
5

---

3.3.1 Importance Sampling
The most simple and computationally efficient method that can be used is importance sampling (IS) [8]. The idea is to utilize an -amplingg proposal distribution oo(0) E P(s 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:
(Bg [Wp( (tv)mm..o.foo0)) ) ( Eg (.p..oo..1(()))
EBSW (1,
where o0) P(()-1)) is the proposal distribution, and:
W wp..oo...pp00 s(Wp(02y,00b))
oo(0)
the importance weighted function. The detailed derivation is given in Appendix B.2 Let 0,, 0,,...,, be i.i samples from (f), the importance sampling estimator of the EBSW (IS-EBSW) is:
.......,. Or.
IS-EBSW (14, v; f,
(g((tth, (.
(8)
Sp where W... (0: Ei f.p..ro.f.pp0pp is the normalized importance weights. When (o
1($ T(d/2)
2nd/2 (a constant of 0), we can replace wp.roo (8t) with F(WB((rtt 0:tr)). When
choose the energy function f(z) = computing the normalized importance weights
W Wp.v.oo.f.p())
Wp..ao (0)) = Ete (p.
is equivalent to computing the Softmax function.
Computational algorithms and complexities. The computational algorithm of IS-EBBW can be derived from the algorithm the SW distance by adding only one to two lines of code for computing the importance weights. For better comparison, we give algorithms for computing the SW distance and the EBSW distance in Algorithm [ and Algorithm 4 in Appendix [.]] and Appendix B.2 respecivvey When μ and y are two discrete measures that have at most n supports, the time complexity and the space complexity of the IS-EBSW distance are O log n + Lnd) and OLL(n +d) which are the same the SW.
The IS approximation is asymptotically unbiased for EBSW vi fi f). However, having a biased estimation is not severe since the unbiasedness cannot be preserved after taking the p-rooth (p > 1) like in the case of the SW distance. Therefore, an unbiased estimation of EBSW" vı; f) is not very vital. Moreover, we show that IS IS-EBBS"", (1,, v; f. L) is unbiased estimation of the power p of a valid distance. We refer the reader to Appendix B.2 for the detailed definition and properties of the distance. From this insight, is safe
Unbiasedness.
the IS-EBSW in practice.
Gradient Estimation. In statistical inference, we might want to estimate the gradient VeEBSW, (M's: v: J) for doing minimum distance estimator (45 Therefore, we derive the gra- dient estimator of the EBSW with importance sampling in Appendix B.2
3.3.2 Sampling Importance Resampling and Markov Chain Monte Carlo
The second approach is to somehow sample from the slicing distribution au f.P). For exam- ple, when we have 0,,...,,; Or are approximately distributed following (0; f, p),, we can take
(+E/,, W2 (/v) as the approximated value of the EBSW. Here, we consider two famous approaches in statistics: Sampling Importance Resampling [14 (SIR) and Markov Chain Monte Carlo (MCMC). For MCMC, we utilize two variants of the Metropolis- Haastpi algorithm: independent Metsonolis-Hasting (IMH) and random walk Metropolis- Hasting (RMH).
Sampling Importance Resampling. Similar to importance sampling in Section 3.3.1 the SIR uses proposal distribution (0) to obtain L samples 0'....,'', and the corresponding normalized

---

importance weights:
p. W I...
Σ (....)....
After that, the SIR creates the resampling distribution which Categorical distribution ((0)
Σl= t (8 dōg Finally, the SIR draws L samples .1..... O2 i.d q(f). We denote the SIR
estimation of the EBSW distance as SIR-EBSW.
Markov Chain Monte Carlo. MCMC creates Markov chain that has the stationary distribution as the target distribution. The most famous way to construct such a Markov chain is through the Metropolis- Hastingss algorithm. Let the starting sample follow a prior distribution 01 oo(() E P(Sd-1), a transition distribution (0, 0z-1) E P(S for any timestep t >> 1 is used to sample a candidate 6,. After that, the new sample 0g set to 83 with the probability a and is set to 0:-1 with the probability 1 - with
1, (1, (uu(ffif or (0(0--10)) (',,,0(0-11;1] or (0%)0--))
min (1, S(Wp ((F2))) Te (0--110)
== min a=min
((W2(0:-1 tot,ttt,,-tt (.)) or (0-(0,11)
In theory. T should be large enough help the Markov chain mix to the stationary distribution and the first M < T samples are often dropped as burn-in samples. However, to keep the compu- tational complexity the same as the previous computational methods, we set 7 = L and M = 0. The first choice of transition distribution is oz (0,(0-))) U($211) which leads to independent Metropolis-Hasting (IMH). The second choice is (0,(0),,))) VMF(0.]0r-1,, (the
Mises- Fisher distribution 17 with location 02-1) which leads to random walk Metropolis-Hasting (RMH). Since both above transition distributions are symmetric or (0(0(0=-1)) (0.- (0.) the acceptance probability turns into:
a min (1. [(W ((tv))
5(W3(O,- 1tu, (--1
which means that the acceptance probability equals and 0 is always accepted 9g if it can increase the energy function. We refer to the IMH estimation and the RMH estimation the EBSW distance as IMH-EBSW and R.MH-BBSW in turn.
Computational algorithms and complexities. We refer the reader to Algorithm 5 Algorithm and Algorithm 7 in Appendix B.3 for the detailed algorithms of the SIR-EBSW. the IMH-EBSW, and the RMH-EBSW W. Without counting the complexity of the sampling algorithm, the time complexity and the space complexity of both the SIR and the MCMC estimation of EBSW are OL log n Lnd) and Ó(L(n + d)) which are the same as the IS-EBSW and the SW distance. However, the practical computational time and memory of the SIR and the MCMC estimation depend on the efficiency of implementing the sampling algorithm e.g., resampling and acceptance-rejection.
Unbiasedness. The SIR and the MCMC sampling do not give an unbiased estimation for EBSW
p (n, v;; f). However, they are also unbiased estimations of the power p of valid distance. We refer to Appendix B.3 for detailed definitions and properties of the distance. Therefore, it is safe to use the approximation from the SIR and the MCMC.
Gradient Estimation: In the IS estimation, the expectation is with respect to the proposal distribution (0) that does not depend on two input measures μgg In the SIR estimation and the MCMC estimation, the expectation is with respect to the slicing distribution (0, f) that depends on uag- Therefore, the log-derivative trick (Reinforce) should be used to derive the gradient estimation. We give the detailed derivation in Appendix B.3 However, the log-derivative trick is often unstable in practice. A simpler solution particular, we denote Mar is the independent copy of μt with ø' equals ø in terms of value. Therefore, we can obtain the slicing distribution op f) that does not depend on plg. This approach still gives the same value of distance, we refer to Appendix B.3 Bo3 for more careful discussion.
create the slicing distribution from an independent copy of μ4a In
Experiments
In this section, we first visualize the shape of the energy-based slicing distribution in simple case. After that, we focus on showing the favorable performance of the EBSW compared to the other sliced Wasserstein variants in point-cloud gradient flows, color transfer, and deep point-cloud reconstruction.
7

---

F -
Felet =
0.001
0,000
E600225
PAC10: 5)
Darecity
0g:
0.004
Visualization of the and the sampled energy-hased slicing distributions, the optimal vMF distribution from the v-DSW, and the max projecting direction from the Max-SW.
Figure 1:
In experiments, we denote EBSW-e for the exponential energy function i.e.. f(x) - e", and EBSW-1
for the identity energy function i.e., fz) x. We use p = 2 for all sliced Wasserstein variants.
4.1 Visualization of energy-based slicing distribution
We visualize the shape of the energy- slicing distribution in two dimensions in Figure In particular, we consider comparing two empirical distributions in the left-most figure (taken from I.I.. We utilize the SIR, the IMH, and the RMH to obtain 10 Monte Carlo samples from the energy- based slicing distribution. For the IMH and the RHM, we burn in the 10 samples. After that, we use the von Mises kernel density estimation obtain the density function. We also present the ground-truth density of the energy-based slicing distribution by uniform discretizing the unit-sphere. Moreover, we also show the optimal vMF distribution from the DSW (wnning A {1, 5, 10, 50, 100}) and the "max" projecting direction from the Max-SW (T 100, step size 0.1). The middle figure is according to the energy function (x) = x and the right-most figure according to the energy function fz(z) = e". From the figures, we observe that all sampling methods can approximate well the true slicing distribution. In contrast, the vMF distribution from v-DSW is misspecified to approximate the energy distribution, and the "max" projecting direction from the Max-SW can capture only one mode. We also observe that the exponential energy function makes the density more concentrated around the modes than the identity energy function (polynomial of degree 1).
4.2 Point-Cloud Gradient Flows
We model a distribution p(t) flowing with time t along the gradient flow of a loss functional jut) - D(u(t), v) that drives it towards target distribution [40] where D is our sliced Wasserstein variants. We consider discrete flows, namely. we set v = Σm= Sy, as a fixed empirical target distribution and the model distribution ju(t) = 1.2 Ei 6x,(1) Here, the model distribution is parameterized by time- virying point cloud X(t)= (xx(t)) E (Rd)". Starting from initial condition time
t = 0, we integrate the ordinary diffcrential equation X (t) = V [D nt E 6x.(t), v)] for each iteration. We choose μ(0) and y are two point-cloud shapes in ShapeNet Core-55 dataset 4. After that, we solve the flows by using the Euler scheme with 500 iterations and step size 0.0001.
Quantitative Results. We show the Wasserstein-2 distance between µ²(() and " (1. (0, 100, 200, 300, 400, 500)) and the computational time of the SW variants in Table we set 100 for SW, and EBSW variants. For the Max-SW we set T 100, and report the best result for the step size for finding the max projecting direction in (0.001, 0.01, 0.13. For the v-DSW, we report the best result for () E {100 10), (50, 2), (2, 50)), rate for finding the location in f0.001, 0.01, 0.1}. We observe that the IS IS-EBSW-e helps to drive the flow to converge faster than baselines. More importantly, the computational time of the IB-EBSW-e is approximately the same as the SW and is faster than both the Max-SW and the v-DSW. We report more detailed experiments with other EBSW variants and different settings of hyperparameters (L,T) in Table in Appendix C.I From the additional experiments, we see that the EBSW-e variants give lower Wasserstein-2 distances than the baseline with the same scaling of complexity (same L/T). De- spite having comparable performance, the -BSSEB the IMH-EBSW-c, and the RMH-EBSW-c
Here,
{1, 10, 50). and the learning
8

---

Figure 2:
Gradient flows from the SW, the Max-SW, the v-DSW, and the IS-EBSW-e in turn.
Table I: Summary of Wasserstein-2 scores 12 (multiplied by 10") from three different runs, computational
time in second (s) to reach step 500 of different sliced Wasserstein variants in gradient flows.
Distances
Step (Wa L)
Step 100 (Wa 1)
Step 200 (W, J)
Step 300 (Ws 1)
Step 400(W, .)
Step 500 (Wa 4)
Time (s J)
SW
2048.29
988.93 9.55
350.06 5.32
93.60
1.85
27.03 0.65
9.41 0.27
17.06 0.45
Max-SW
2048.29
506
9.28
93.54
22
0.79
9.62 0.22
6,83 0.22
28.38 0.05
v.DSW
2048.29 0.0
(49.333- + 8.77
127.4 + 5.06
20.44 1.25
10:95 + 1.0
5.68 0.56
21.2 0.02
(k = 10) are slower than the IS-EBSW-e variant. Also, we see that the EBSW-e variants are better than the EBSW-1 variants. We refer to Table for comparing gradient estimators of the EBSW.
Qualitative Results. We show the point-cloud flows the SW, the Max-SW, the v-DSW, and the IS-EBSW-e, in Figure The flows from the SIR-EBSW-e, the IMG-EBSW-e, and the RMH-EBSW-e are added Figure in the Appendix C.I From the figure, the transitions of the flows from the EBSW-e variants are smoother to the target than other baselines.
4.3 Color Transfer
We build gradient flow that starts from the empirical distribution over the normalized color palette (RGB) of the source image the empirical distribution over the normalized color palette (RGB) of the target image. Since the value of the color palette in the set {0, 255]3, we must do an additional rounding step at the final step of the Euler scheme with 2000 steps and step size 0.0001.
Results. We use the same setting for the SW variants in the previous section. We show both transferred images, corresponding computational times, and Wasserstein-2 distances in Figure We observe the same phenomenon in the previous section, namely, the IS- IS-EBSW-e variants perform the best in terms of changing the color of the source image to the target in the Wasserstein-2 metric while the computational time only slightly higher. Moreover, the transferred image from IS-EBSW-e is visually more similar to the target image in color than other baselines, namely, it has a less orange color. We refer to Figure 5 in Appendix C.2 for additional experiments including the results for the SIR-EBSW-c. the IMH-EBSW-C the RMH-EBSW-c, the results for the identity energy function, the results for changing hyperparamters (L, I',, and the results for comparing gradient estimators. Overall, observe similar phenomenons in the previous gradient flow section.
4.4 Deep Point-Cloud Reconstruction
33 to train point-cloud autoencoders with sliced Wasserstein distances on the ShapeNet Core-55 dataset 4 In short, we aim estimate an autoencoder that contains encoder Jo that maps point cloud X E gnd to latent code g²t and decoder gy that maps the latent
follow

---

Source
SW M455 My IN
Max SW 32% (Resi N. 603.2
7ow No (315 My 36000
Eel Pork
Target
Figure 3:
The ligures show the source image, the target image. the transferred images from sliced Wasserstein variants, the corresponding Wasserstein-2 distances to the target color palette, and the computational time.
Table 2:
Reconstruction emrors from three different runs of autoencoders trained by different distances. The sliced Wasserstein distance and the Wasserstein distance multiplied by 100.
Epoch 20
Epoch 100
Epoch 200
Distance
SW²(O
Wz(1)
SW2 (L)
Wz(U)
SW₂ (1)
Wz(U)
SW
0.14
12.67 0.18
2.29 0.04
10.63 0.05
2.15 + 0.04
9.97 0.08
Max-SW
2.91 0.06
12.33 + 0.05
2.24 0.05
10.40 0.06
2.14 0.10
9.84 0.12
v-DSW
0.02
12.64 0.02
2.21 0.01
1.00
2.07 0.09
9.81 0.05
IS-EBSW-e
2.68 0.03
11.90 0.04
2.18 0.04
1.02
2.04 0.09
9.69 0.14
code to a reconstructed point cloud X E Rnd We want to have the pair fō and gv such that X = (x)( s X for all X p(X) which is our data distribution. To do that, we solve the following optimization problem: ming. E [x= P (fx(()))) where S is a sliced Wasserstein variant, and Px denotes the empirical distribution over the point cloud X. The backbone for the autoencoder is a variant of Point-Net [37] with an embedding size of 256. We train the autoencoder for 200 epochs using SGD optimizer with learning rate of 1e-3, batch size of 128, a momentum of 0.9, and weight decay of 5e-4. We give more detail in Appendix C.3
Quantitative Results. We evaluate the trained autoencoders on different dataset: ModelNet40 dataset [46) using two distances: sliced Wasserstein distance (L = 1000), and the Wasserstein distance. We follow the same hyper-parameters settings as the previous sections and show the reconstruction errors at epochs 20, 100, and 200 from the SW, the Max-SW, the DSW, and the IS-EBSW-e in Table 2] The reconstruction errors all point-clouds. From the table, we observe that the ES-EBSW- can help to train the autoencoder faster in terms of the Wasserstein distance and the sliced Wasserstein distance. We refer to Table []] in Appendix SIR-EBSW-e, the IMH-EBSW-e, the R.HH-EBSW-e, the results for the identity energy function, the results for changing hyperparameters (I, T), and the results for comparing gradient estimators.
the average of corresponding distances
C.3 for similar ablation studies as in the previous sections including the results for the
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 Figure B- respectively. Overall, the qualitative results are visually consistent with the quantitative results.
5 Limitations and Conclusion
Limitations. The first limitation of EBSW is that its MCMC variants are not directly parallelizable, resulting slow computation. Additionally, a universal effective choice of the energy-based function is an open question. In the paper, we use simple and computationally effective energy- functions, such as the exponential function and the polynomial function. Finally, proving the triangle inequality of EBSW is challenging due to the expressiveness of the energy- slicing distribution.
Conclusion. We have presented a new variant of sliced Wasserstein distance, named energy-based sliced Wasserstein (EBSW) distance. The key ingredient of the EBSW is the energy-based slicing distribution which has a density projecting direction proportional to an increasing function of the one-dimensional Wasserstein value of that direction. We provide theoretical properties of the EBSW including the topological properties, and statistical properties. Moreover, we propose to compute the EBSW with three different techniques including importance sampling, sampling importance resampling, and Markov Chain Monte Carlo. Also, we discuss the computational properties of different techniques. Finally, we demonstrate the favorable performance of the EBSW compared to existing projecting directions selection sliced Wasserstein variants by conducting experiments on point-cloud gradient flows, color transfer, and deep point-cloud reconstruction.
10

---

Acknowledgements
We would like to thank Peter Mueller for his insightful discussion during the course of this project. NH acknowledges support from the NSF IFML 2019844 and the NSF AI Institute for Foundations of Machine Learning.
References
[1] C. Bonet, N. Courty, F. Septier, and L. Drumetz. Efficient gradient flows in sliced- W Wasserstein space. Transactions on Machine Learning Research, 2022.
[2] N. Bonneel, J. Rabin, G. Peyré, and Pfister. Sliced and Radon Wasserstein barycenters of measures. Journal of Mathematical Imaging and Vision, 1(51))22-45, 2015.
[3] N. Bonnotte. Unidimensional and evolution methods for optimal transportation. PhD thesis, Paris 11, 2013.
[4] A. X. Chang, Funkhouser, L. Guibas, P. Hanrahan, Q. Huang, Li, S. Savarese, M. Savva, S. Song, H. Su, et al. Shapenet: An information-rich 3d model repository. arXiv preprint arXiv: 1512.03012. 2015.
[5] M. Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems, pages 2292-2300, 2013.
[6] T. R. Davidson, L. Falorsi, N. De Cao, T. Kipf, and J. M. Tomczak. Hyperspherical variational auto-encoders. In 34th Conference on Uncertainty in Artificial Intelligence 2018, UAI 2018, pages 856 865. Association For Uncertainty in Artificial Intelligence (AUAI), 2018.
[7] I. Deshpande, Y.-T. Hu, R. Sun, A. Pyrros, N. Siddiqui, S. Koyejo, Z. Zhao, D. Forsyth, and A. G. Schwing. Max-sliced Wasserstein distance and its use for GANs. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 10648-10656, 2019.
[8] I. Deshpande, 7. Zhang, and A. G. Schwing. Generative modeling using the sliced Wasserstein distance. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 3483-3491, 2018.
[9] L. Devroye, L. Györfi, and G. Lugosi. A probabilistic theory of pattern recognition, volume 31. Springer Science & Business Media, 2013.
[10] D. Dowson and B. Landau. The fréchet distance between multivariate normal distributions. Journal of multivariate analysis, 12(3):450-455 1982.
[11] J. Feydy, T. Séjourné, F.-X. Vialard, S.-i. Amari, A. Trouve, and G. Peyré. Interpolating between optimal transport and MMD using Sinkhorn divergences. In The 22nd International Conference on Artificial Intelligence and Statisties, pages 2681-2690, 2019.
[12] R. Flamary, N. Courty, A. Gramfort, M. Z. Alaya, A. Boisbunon, S. Chambon, L. Chapel, A. Corenflos, K. Fatras, N. Fournier, L. Gautheron. N. T. Gayraud, H. Janati, A. Rakotomamonjy. I. Redko, A. Rolet, A. Schutz, V. Seguy, D. J. Sutherland, R. Tavenard, A. Tong, and T. Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8, 2021.
[13] Z. Goldfeld and K. Greenewald. Sliced mutual information: A scalable measure of statistical dependence. Advances in Neural Information Processing Systems, 34, 2021.
[14] N. J. Gordon, D. J. Salmond, and A. F. Smith. Novel approach to nomlinear/non-gaussiaa bayesian state estimation. In IEE proceedings F (radar and signal processing), volume 140, pages 107-113, 13. IET, 1993.
[15] M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter. GANs trained by a two time-scale update rule converge a local Nash equilibrium. In Advances in Neural Information Processing Systems, pages 6626-6637, 2017.
[16] G. F. Hinton. Training products of experts by minimizing contrastive divergence. Neural computation, 14(8):1771-1800, 2002.

---

[17] P. E. Jupp and K. V. Mardia. Maximum likelihood estimators for the matrix von Mises-Fisher and bingham distributions. The Annals of Statistics, (53)5599-606, 1979.
[18] T. Kloek and H. K. Van Dijk. Bayesian estimates of equation system parameters: an application of integration by monte carlo. Econometrica: Journal of the Econometric Society, pages 1-19, 1978.
[19] S. Kolouri, K. Nadjahi, U. Simsekli, R. Badeau, and G. Rohde. Generalized sliced Wasserstein distances. In Advances in Neural Information Processing Systems, pages 261-272, 2019.
[20] S. Kolouri, G. K. Rohde, and H. Hoffmann. Sliced Wasserstein distance for learning Gaussian mixture models. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition, pages 3427-3436, 2018.
[21] A. Krizhevsky, G. Hinton, al. Learning multiple layers of features from tiny images. Master's thesis, Department of Computer Science, University of Toronto, 2009.
[22] C.-Y. Lee, T. Batra, M. H. Baig, and D. Ulbricht. Sliced Wasserstein discrepancy for unsuper- vised domain adaptation. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10285-10295, 2019.
[23] Z. Liu, P. Luo, X. Wang, and X. Tang. Deep learning face attributes in the wild. In Proceedings of International Conference Computer Vision (ICCV), December 2015.
[24] A. Liutkus, Simsekli, S. Majewski, A. Durmus, and F.-R. Stöter. Sliced- Wasserstein flows: Nonparametric generative modeling via optimal transport and diffusions. In International Conference on Machine Learning, pages 4104 4113. PMLR, 2019.
[25] K. P. Murphy. Probabilistic Machine Learning: Advanced Topics. MIT Press, 2023.
[26] K. Nadjahi, V. De Bortoli, A. Durmus, R. Badcau, and U. Şimşekli. Approximate Bayesian computation with the wasserressn distance. In ICASSP 2020-2020 IEEE International Conference Acoustics, Speech and Signal Processing (ICASSP), pages 5470-5474. IEEE, 2020.
[27] K. Nadjahi, A. Durmus, L. Chizat, S. Kolouri, S. Shahrampour, and U. Simsekli. Statistical and topological properties of sliced probability divergences. Advances in Neural Information Processing Systems, 33:20802-208112 12, 2020.
[28] K. Nguyen and N. Ho. Amortized projection optimization for sliced Wasserstein generative models. Advances in Neural Information Processing Systems, 2022.
[29] K. Nguyen and N. Ho. Control variate sliced wasserstein estimators. arXiv preprint arXiv:2305.004000 2023.
[30] K. Nguyen, N. Ho, T. Pham, and H. Bui. Distributional sliced- Wasserseenn and applications to generative modeling. In International Conference on Learning Representations, 2021.
[31] K. Nguyen, D. Nguyen, and N. Ho. Self-attention amortized distributional projection optimiza- tion for sliced Wasserstein point-cloud reconstruction. Proceedings of the 40th International Conference on Machine Learning. 2023.
[32] K. Nguyen, S. Nguyen, N. Ho, T. Pham, and H. Bui. Improving relational regularized au- toencoders with spherical sliced fused Gromov-Wasse Wasserstein. In International Conference on Learning Representations, 2021.
[33] T. Nguyen, Q.-H. Pham, T. Le, T. Pham, N. Ho, and B.-S. Hua. Point-set distances for learning representations of 3d point clouds. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2021.
[34] S. Nietert, R. Sadhu, Z. Goldfeld, and K. Kato. Statistical, robustness, and computational guarantees for sliced wasserstein distances. Advances Neural Information Processing Systems, 2022.
12

---

[35] R. Ohana, K. Nadjahi, A. Rakotomamonjy, and L. Ralaivola. Shedding pac-bayesial light on adaptive sliced-weet wasserstenn distances. In International Conference on Machine Learning, pages 26451-26473. PMLR, 2023.
[36] G. Peyré and M. Cuturi. Computational optimal transport, 2020.
[37] C. R. Qi, H. Su, K. Mo, and L. J. Guibas. Pointnet: Deep learning point sets for 3d classification and segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 652-660 2017.
[38] M. Rowland, Hron, Y. Tang, K. Choromanski, T. Sarlos, and A. Weller. Orthogonal estimation of Wasserstein distances. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 186-195. PMLR, 2019.
[39] Salimans, I. Goodfellow, W. Zaremba, V. Cheung, A. Radford, and Chen. Improved techniques for training GANs. Advances in Neural Information Processing Systems, 29, 2016.
[40] F. Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58-63):94,
2015.
[41] A. Savkin, Y. Wang, S. Wirkert, N. Navab, and F. Tombari. Lidar upsampling with sliced Wasserstein distance. IEEE Robotics and Automation Letters, 8(1):392-399, 2022.
[42] Y. Song, S. Garg, Shi, and S. Ermon. Sliced score matching: A scalable approach density and score estimation. In Uncertainty in Artificial Intelligence, pages 574-584, PMLR, 2020.
[43] C. Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008.
[44] M. J. Wainwright. High-dimensionul statistics: non-asymptotic viewpoint. Cambridge University Press, 2019.
[45] J. Wolfowitz. The minimum distance method. The Annals of Mathematical Statistics, pages 75-88, 1957.
[46] Z. Wu, S. Song, A. Khosla, F. Yu, L. Zhang, X. Tang, and J. Xiao. 3d shapenets: A deep representation for volumetric shapes. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1912 1920, 2015.
[47] M. Yi and S. Liu. Sliced Wasserstein variational inference. In Fourth Symposium on Advances Approximate Bayesian Inference, 2021.
13

---

Supplement to "Energy-Basel Sliced Wasserstein Distance"
We first provide skipped proofs in the main text in Appendix A We then provide additional materials including additional background, detailed algorithms, and discussion in Appendix B Additional experimental results in point-cloud gradient flows, color transfer, and deep point-cloud reconstruction in Appendix C Finally, we report the computational infrastructure in Appendix D
A Proofs
A.1 Proof of Theorem
Non-negativity and Symmetry. the non-negativity and symmetry properties of the EBSW follow directly by the non-negative non-negativity and symmetry of the Wasserstein distance since it an expectation of the one-dimensional Wasserstein distance.
Identity. We need to show that EBSW ((,, v; f) if and only if μa = u. First, from the definition of EBSW, we obtain directly same proof technique in 33. EBSW, p (1, v: f) = 0, we have for W, (0tu, 00t) u (0; p) = 0. Hence, we have W (0tp, 0iv) 0 for uu [o; p)-mmmost surely 0 E gd-1, Since Tn (0; f,p) continuous, we have (offt, 0fv) = 0 for all 0 E S From the identity property of the Wasserstein distance, we obtain 05μ = fiv for (p: f.pe A E sd-1, Therefore, for any E R and gd-1, we have:
implies EBSW, v; f] = 0. For the reverse direction, we use the
Flul(t0) - fra e- 10000000 (-(0000)) / e c0ttdd0tp(2)) = (t)
= = JR.
it= =0tu())) = e e- =
3dd
where (w)) = SR e-((wr)) dr(()) denotes the Fourier transform of y E P(R By the injectiv- ity of the Fourier transform, obtain μ which concludes the proof.
A.2 Proof of Proposition
(a) We first provide the proof for the inequality SW SW,(,,)) I EBSW (,, ff.. It is equivalent to prove that
(B [Wt(Ot,, [ S Ea-an..(0111))) [Wg(Orn, 0gv)]
From the law of large number, is sufficient to demonstrate that
1
p (o,tt 0,42) S WP(0,tu,
0,tv)
T.
Σ (( 0,tv))
(9)
for any L 2 and .1....,,y, 0r lid u($a-1). To ease the presentation, we denote a, WP 0, and b; S(WP (0.tM, 0.b0)) for all s i < L. The inequality 9 becomes:
1 L
a,b;
(10)
We prove the inequality 100) via an induction argument. It is clear that this inequality holds when L = 1. We assume that this inquualtyy holds for any L. We now verify that the inequality (T)) also holds for + 1. Without loss of generality, we assume that ai
S a2 S ... S az. S a1411 Since
the function f is an increasing function, it indicates that by < by S S br. < br. Applying the induction hypothesis for 01.-.., az and .1y...,he ,b, we find that
L iml
a.b.
14

---

This inequality leads to
L+1
b,) S L
a,b;
+
a.bbs+1 + b.)at+1 + #a++bb++1
i=1
==1
Therefore, to obtain the conclusion of the hypothesis for L + 1, it is sufficient to demonstrate that
L
Lab a,b +
a,bbL+1
b()ar+1 + actIde+1 (L + (){ a,bi),
:=1
which is equivalent to show that
by) aa)a S Σ a,b + Lol+1bL+1
(11)
1=1
Since a*+1 a, and 1 b; for all Sis L, we have (aL+1 a;)(bi+1 by) 2 0, which is equivalent to a1+1bL+1 + a,b; 2 BL+1d; + bL+14; for all $ i L. By taking the sum these inequalities over i from 1 to L, obtain the conclusion of inequality (IT). Therefore, we obtain the conclusion of the induction argument for L +-1 1, which indicates that inequality (10) holds for all L. As consequence, we obtain the inequality SW (W, S EBSW (1,
(b) We recall the definition of the Max-SW:
Max-SW, Max-SW,(u.p)) max WptOtu, 0tv). 0E84-1
Since Sd-1 compact and the function 0 urgmaxaesa-1 Wp (0tu, 0tv). From Definition for any 2 1, dimension d 2 1, mmergy-function f. and μ, E (RR) we have:
W Wptttu, 0tv) is continuous, we have f
EBSW (wW,p(,.)) v) (Eo-o,,. [W7 (0Hu, 0)w) 1
S (Eeroo... [W (0"du, 0x)=] $ WP (0 dis, 00) = Max-SW,(1,,))
Furthermore, by applying the Cauchy-Schwand Schwartz inequality, we have:
REIL(H. [ 10 - 0 0Tll" d
Max-SW?(p.r) (w)
max 0E8
inf
max RE
inf ( RA"
|10||||æ dd
inf rEll(p.)) xBd f
|0101/00 y/l'ddd((e" (z, w)
inf =EIl(uw)) Rdd
|| y1 v)
{
< inf aET((uP) RaRa
x y/ v)
W7 (v, v),
after taking the p-rooth, we completes the proof.
A.3 Proof of Theorem
We aim to show that for any sequence of probability measures (μk)eEN and μ in Pp (R lim EBSW (1k, μt, f) if and only if for any continuous and bounded function f Rd R. lim f du = f dp.
We now show that lim EBSW, (1te μz; f) 0 implies (H+ ken converges weakly to p. Let (uR)eer a sequence such that limg- FBSW (URR μe f) O, we suppose (uR) does not converge weakly to μ. So, let Dp be the ProohPro metric, limp- Dr( # 0 that implies
15

---

there exists > and subsequence (RO()) KEN with an increasing function & : N + N such that for any l E N: D 2 Using the Holder inequality with μa, v E P. (TE"). we have:
EBSW (1,
E B (0:F.r) [W (#tu, 0[vi]
2 Bgwa . (0.f P (01ne 0tv)]
(B (Wp (0tg,, 0tv)]
BB 1) [Wi (0tm, 0iv)]
SWI (W,
Therefore, lim SWI(Hem)) p)) = which implies that there exists a subsequence (Po(e())) KEN with increasing function @ N N such that (Ha(0())) KEN converges weakly by Lemma S1 in (27]. Therefore a contradiction appears from the assumption of lim4-+00 Dp (HOCN(O)::)) 1)) + O. Therefore, lim EBSW, (14. f) = 0, (uR)REN converges weakly to p.
When (uK)ee converges weakly to we have (05pk) )ken converges weakly to 0ip for any E Sd-1 by the continuous mapping theorem. From 43 the weak convergence implies the convergence under the Wasserstein distance. So. we have limg-+oc Wp(0nu μI) = O. Moreover, using the fact that the Wasserstein distance is also bounded, hence, the bounded convergence theorem implies:
lim EBSWP 14; lim EB [W (Otue, Otu)) k0000
10000
Ee [Wp(OtuR, W (0
lim 1+000
Ex [w F.
Again, using the continuous mapping theorem with function
I//P, we have
lim EBSW, (pu, ie f) 0. We conclude the proof.
A.4 Proof of Proposition
2
We first show that the following inequality holds
TE SMMax-SWp((m.//le S CV (d lognn//
where C >>> 0 is some universal constant and the outer expectation is taken with respect to the random variables X1, XI....,X X.. We now follow the proof technique from in [30]. Let Fr.e and F3 be the inverse cumulative distributions of two push-forward measures 0in,, and 0tp. Using the closed-form expression of the Wasserstein distance in one dimension, we obtain the following equations and inequalities:
[ (1 Fg (u
max Max-SW (p..
BE5a
L'' |F (FT)(() Fg Pu
max
DER
(2)
max $E
| (n,e(x))
Fe(x)
where & C IRd is the compact set of the probability measure μ. We can check that
max VER* 101
Fn.((z) Fo(x)) sup (A) μ(A)), AEH
where H is the set of half-spaces {z E IRd : - x) for all E gd such that ||0] S 1. From VC
inequality (Theorem 12.5 in (9), we have:
IP sup (A) u(A))> $8S(H, n)e nn)-//1) /32
AEH
16

---

with S(H. r)) is the growth function which is upper bounded by (1+11 due to the Sauer Lemma (Proposition 4.18 in "High-Dimensional Statistics: A Nom-Asympiotic Viewpoint", Wainwright et al).
From Example 4.21 in [4)) we get VC(H) = d 1.
Let 8,(H,, nt /32 8, we have t? 2
32 log BS(H.n) n
Therefore, obtain
32
(, n)
sup
|m (A) ((A) S
log
2 1 -
5
AEW
Using the Jensen inequality and the tail sum expectation for non non-negane -on-negative random variable, we have:
sup Im (A) - -x((A)) AEH
sup (Mx(A) (A)(4)- ma()] 1
//- P
((SUR) AEH
sup ((n(A) m(a))) 21) dt
E
LAH
f TP
sup (Mn(A) m(()))))))))))))) MEH
> dl +
00 P
sup AEH
Im (A) m(4))) 21 dl
I- Idt +
{ nSe = u+ 256S n))
e enn/33
71
Since the inequality holds for any u, we optimize for u. Let [u) = u + 2565 n) na we have f'(u) 1+ 8S n)ee mn/332 Setting f'(u) we obtain the minima u² 32 (85())()))))) Therefore, we have:
7
32 (ogg(8((H,,)) + +22 S CV
(d (d+ / [d+1)log(() log(n +
E sup (um(A) - (A)))| S vf(w))
LEH
n
n
by using Sauer Lemma, As a consequence, we obtain:
CIEBSW, M4 f)] (dd++1 ()++1)log//m.
which completes the proof.
B Additional Materials
B.1 Additional Background
When two probability measures are empirical probability measures on n 1 "i= 8g.- the SW distance be computed by sort- ing projected supports. In particular, we have 0ip 1 Ei dOT Opv Σか døg ye and W2(005 0dv) = = , Fit) - 0 TT )()) where 0 20 the ordered projected supports. We provide the pseudo-code for computing the SW in Algorithm
Sliced Wasserstein.
supports: μ x2 Ex= 8z, and
(0
Max sliced Wasserstein. The Max-SW often computed by the projected gradient ascent. The sub-gradient is used when the one-dimensional optimal matching is not unique e.g., in discrete cases. We provide the projected (sub)-gradient ascent algorithm for computing the Max-SW in Algorithm 2
Compared to the SW. the Max-SW needs two hyperparameters which are the number of iterations T and the step size n) Moreover, the empirical estimation of the Max-SW might not converge the Max-SW when T + co.
Distributional sliced Wasserstein. To solve the optimization of the DSW, need use the stochastic (sub)-gradiee ascent algorithm. In particular, we first need to estimate the gradient:
V (0) Wg 6tv))=
1 P
(Eg-a WP (0
1- V V. (0) WF 0tv).
17

---

Algorithm Computational algorithm of the SW distance
Input: Probability measures μ and v, p 1, and the number of projections L.
for | to do
Sample 9₁ ~ (8
Compute V/2 (fittt 8:11)
end for
Compute SW v; (1 E+
Return: SW (1, L)
Algorithm 2 Computational algorithm of the Max-SW distance
Input: Probability measures μ and v, p 1, the number of iterations T, and the step size 17.
Sample 0g (((=11)
for 1 to T do
Compute 0r 01-1 +V o... Wp(0r- 1
Compute 0: = 110e|12
6,
end for Compute Max-SW, (T = Wp(Ortu, Return: (Kx,sW,
Algorithm 3 Computational algorithm of the DSW distance
Input: Probability measures and v, p 2 1, the number projections L, the number iterations T. and the step size 17.
Initialize eoo for t 1 to T do
V.
for - 1 to do
Sample fr,v ())) via reparameterizzatin.
Compute O, 110,112
end for Compute 3h = -h-1 + n (+ E/ Wp(00 (.okt)) 12p end for
1 Ve We
for 1 = to T. do Sample ~ o (r) via reparameterizzatio.
end for
Compute DSW (;, - Et WB(0 Return: DSW DSW,(H,v;T.Ll
To estimate the gradient V VilEoa (0) WP (Otu, (jv), need to reparameterization trick for (0) e.g., the vMF distribution. After using the reparameterization trick, we approximate the gradient Ve. (0suu 0tv) = El_ V (or, (.. where fI.p...., 1a....,.... are ii.d reparameterized samples from (f). Similarly, we approximiate Ey (e 8-v)
J, Et= we algorithm computing the DSW in Algorithm 3 Compared to the SW, the DSW needs three hyperparameters i.e., the number of projections L, the number of iterations T. and the step size 1).
(Ordu, 01tv). We refer to the details in the following papers 66 [2]. We review the
Minimum Distance Estimator and Gradient Estimation. In statistical inference, we are given the empirical samples XI,..., X,, from the interested distribution Since we do not know the form of v, we might want to find an alternative representation. particular, we want to find the best member ne a family of distribution parameterized by á 4. To do that, want to minimize the
18

---

Algorithm 4 Computational algorithm of the IS-EBSW distance
Input: Probability measures f and v, p 1, the number of projections L, the energy function f.
for = to L do
Sample 9₁ ~ 4(8--1)
Compute et (0r50 0:11) Compute wi= S(W, (Otiio, 0121))
end for
Compute IS-EBSW, p(1, v5 1,
(Etn 21== E wi
Return: IS-EBSW UE,
distance between Me and the empirical distribution V, = E òx, This framework is named the minimum distance estimator [45]
minge D( Vn)), where D is discrepancy between two distributions. The gradient ascent algo- rithm often used to solve the problem. To do so, we need to compute the gradient D( v.). When using sliced Wasserstein distances, the gradient V,D((u, Kn) is often approximated by stochastic gradient since the SW distances involve an intractable expectation. In previous SW vari- ants, the expectation does not depend on Φi hence, we can use directly the Leibniz rule to exchange the gradient and the expectation, then perform the Monte Carlo approximation. In particular, we have VolBewo(0) (Otn,, 0tv]] = Ee)a(0) W WP (Otp,, Otvi] 2
+ V (ppttrt,, Ortv) for
0 O....., Or. iid (0).
B.2 Importance Sampling
Derivation. We first provide the derivation of the importance sampling estimation of EBSW. From the definition of the EBSW, we have:
EBSW (v, f) = (Bo~o (8:1.r) [Wp )
Sgs WR(0rp, 0tv) (WP (OE,, (
Jgs- f(W/ (4t,. )1v))
Sg Wp(Otu, 0t2) f ((W (2tm.0a)))
on (0)c0
ao(
Jo
((W3(etu.0zv) (o
(0) de
f (wng(tn
Wgr0t 0tv)
(o)
Ee
mo(0)
(Bp0 05V) (dw)waeo
Ea Wu,,oo (
Algorithms. We provide the algorithm for the IS estimation of the EBSW in Algorithm Compared to the algorithm of the SW in Algorithm [ the IS-EBSW can be obtained by only adding one or two lines of code in practice. Therefore, the computation of the IS-EBSW is as fast as the SW while being more meaningful.
Gradient Estimators. Let Ue be parameterized by d, we derive now the gradient estimator VaEBSW p v; f) through importance sampling. We have:
(g [WP (OtHus p (B)]
d EBSW (Ms,, M3 f)
p
Ex W Wpo.
(o (Wp(0tups 04. ( 2oo0 (0)
(a 10 (0)
19

---

We denote A(o) = (0) (Wp (pp))op. (8)]. B(o) BB-ar(()) [uppo we have
A B( - A
V
(())
(²(Ø)
Using Monte Carlo samples 0 .,...., Or (0) after using the Lebnitz rule to exchange the differentiation and the expectation, we obtain:
Ev [Wp(OEpg. 0tv) (wo.
t E
WW ( O pr.
EB [Wss.rro., (0)
1 E [ups)v,oo.{())
L
VoA(o) V
B(b()) ~ T,
E
V 2tl (0
1=1
which yields the gradient estimation. construct the slicing distribution by using copy of μs ie., Hy with o' = á in terms of value, the gradient estimator can be derived by:
1-2
Ee
(p (W Moo.
1
V e EBSW p (M
p
Eenool0) Wug.vooo.
(B)
Eg (Wpsr (0)
Using Monte Carlo samples 0,,...,0z 02 differentiation and the expectation, we obtain:
(0) after using the Lebnitz rule to exchange the
Le
(0) wgrO0ups (au)wwg
H. E
Wp(Ortuo, Oitv) woo mon.)
1
Ea (wn..mo..)())
+ Σ (wo. (0:)
1
V Vo Wp (22) mrooo (0)
~
V PoWpp0:2u0 Octu)) upone
L
1=1
Ea-ar(()) [wpg. 1 L
W (pg.
It is worth noting that using a copy of μ4 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 estimation of the EBSW not unbiased estimation for finite L, it is unbiased estimation of valid distance 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 2 1, dimension d > 1, energy function f. continuous proposal distribution (o) ~ P() and probability measures μ E P, (IR) and sliced Wasserstein (IWSW) distance is defined as follows:
RRd. the importance weighted
W 0.t1) (1t 1)
T
IWSW,
B
(12)
1 E [wp.oo..fpp00 (0i)
where the expectation is with respect to 0,... 0r oo(0), and w...ooo.f (P2,
Go(#)
20

---

Algorithm 5 Computational algorithm of the SIR-EBSW distance
Input: Probability measures μ and v, p 1, the number of projections L, the energy function f.
for | to do
Sample 9₁ ~ U((--1)
Compute wi [WW,(Ottu,, O(dw))
end for
for 1 to L do
/(W (0w,1011u
Compute Vi ELs (WNp(0:Hp.0tt)))
end for for
/ I, do
Sample 0: ~ Cat (1...., w,)
Compute 14 W, w (1t 0110)
end for
Compute SIR-SW (14, v; [, (1 $
Return: JR,,S
Algorithm 6 Computational algorithm of the SW distance and the IMH--BSW distance
Input: Probability measures μ and v, p 2 1, the number of projections L, the energy function f.
Sample 01
Compute 1 WPp01tuu 01tv)
for / to L do
Sample 80 ~ U((
Compute = min
Sample a u( 1))
if 2 then
Set 0: = 8%
else if < then
Set 0; = 01-1 end if
0) = W Worttu (itv)
end for
Compute IMH-EBSW (er,,,,(,,,, (H.
Return: IMH-EBSW LI,
The IWSW is semi-metric, 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.I Apendix 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 RMH in Algorithm 7
Gradient estimators. We derive the reinforce gradient estimator of the EBSW for the SIR, the IMH, and the RHM sampling.
Le
V EBSW (us, vi - (B (So(0)) VoBooa
Wg 0tv)
We have:
V V Vallgnap [Wp(0tmp, 0Hv)] Eg p (0)
[P 05v) V log (WD(0dmas, Ot (0;
21

---

Algorithm 7 Computational algorithm of the SW distance and the RMH-EBSW distance
Input:
Probability measures y and v, p 2 1, the number of projections L, the energy function f. the concentration parameter R.
Sample 01 US
Compute for I = L do
W, Wp(Oitne (0, 01tu)
Sample 0g ~ vMF(01-1,
(W3100Ap. ((t))) J(W 1000
Compute
min
Sample u ~ (0 1])
if 2 then
Set 0 = 0j
else if
then
Set 0₁ = 0
end if'
by W (Orttu 01tv)
end for
Compute RMH-EBSW (M,, 1; (tΣti ox)
Return: RMH-EBSW p LL
and
V log (Wp(0tue, 02u) ((:;))) = V log Jog(WpOtne: 030)) + V (og(s(Wp(Otrn) 0f0))
(f
V log
/(W 62v)
1
( 0tv)
Wpt00ue 0Ht))
+ (OWp OHW))
VWT 0tv))
V log ( S(WB(05H6) 0
and
(J
Vo log
2/2
V& log
Eane(Ss 1)
F(WF 0zve) [(d/2)
V log (E8-2(88-1) [f (bv))]
V Vo Vo [W 0tv]]
BB- [S(W 01v))]
Using L Monte Carlo samples from the SIR (or the IMH or the RMH) to approximate the expectation Eg-a (P:f)) and L samples from ((( to approximate the expectation Eg-u(gt-1)) obtain the gradient estimator of the EBSW. However, the reinforce gradient estimator unstable in practice, especially with the energy function fa(x) e². Therefore, we propose more simple gradient estimator which is:
I-P
V &&BBSW (u ( (Od 0zv)]
(Eawangru
Eg (0:1))
WP (05 0tv)]
The key is to use copy of the parameter o' for constructing the slicing distribution ong (0; f). hence, we can exchange directly the differentiation and the expectation. It 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 the "parameter-copy" gradient estimator while the original one the conventional estimator.
22

---

Population distance. The approximated values of p-power EBSW from using the SIR, the IMH, and the RMH can be all written as 1. Ei= Wp(0rzu, 0ite). Here, the distributions of 01,..., Og are different. Therefore, they are not unbiased estimation of the EBSW (u, (u, v; However, the population distance of the estimation can be defined as in Definition
Definition For any p 1, dimension 1. energy function f. and two probability measures (R((R) and v E RR, the projected sliced Wasserstein (PSW) distance is defined as follows:
w (HEE, 0tv) (110B)H.0900 1)
Ps,
(13)
where the expectation is with respect to (fI...., O2) o ,1) which 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 μt, M2, M4 E P, (R*) we have:
wp (O:2p1, (1tu)))) ))
$
1
PSW (SW,lui.pg
E ((tran
L
(W, (Ortun, 0.tu2) + Wp (0r1pe, 0(tu)))P
T.
$
W (0. (1tr2)
< ((8r)))o((Bn))
L
1
+ (((r))o((eea))
Σ
Wp (Oituz, Oitwa)
T.
PSW (SS,((11, + (P2,
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.]] Apendix [A.3] and Appendix A.4 Therefore, using the SIR, the IMH, and the RMH estimation of the EBSWs are safe as the SW.
C Additional Experiments
In this section, we provide additional results for point-cloud gradient flows in Appendix C.I color transfer in Appendix C and deep point-cloud reconstruction in Appendix
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 use the conventional gradient estimator for the IS-EBSW while the "parameter-copy" estimator used for other variants of the EBSW. Therefore, we also provide the ablation studies comparing the gradient estimators in Table 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 the best variant of the EBSW in both performance and computational time. Also, we observe that the exponential energy function better than the identity energy function this application. worth noting that the EBSW variants of all computational methods and energy functions are better than the basclines in
23

---

Table 3: Summary of Wasserstein-2 scores (multiplied by 10*) from three different runs, computational time in
second (s) to reach step 500 of different sliced Wasserstein variants in gradient flows.
Distances
Step 0 (W2 ()
Step
(W2
Step (Wa
Supp 500 L)
Step 100/W, L)
Step 500 (Wa 4)
(s L)
SW L=100
2018 29 +: 0.0
986.99 96. 2: 9.55
350 66 1: 5.32
99.69 2:1.8
27.03 27:03:2:00 0.65
9.11 #: 0.27
17 06
0.45
MRx-SW 1=100
2848 2018221-1 +: 0.0
5066 +::5::::00 0:228
18.54 +: 3.39
22.2 = 0.79
9.62 0.22
6:83 4: 0:22
28. 38
0:05
L=T=100
127.4 5.06
5.68
21
DSW L-T-100
2048.22 +: 0.0
519.63
2.43
101.90
1.44
24.11 0.49
8.85 0.16
4.62 0.15
24.52 + 0.00
EDS. L=100
2048. 29 4: 0.0
419.09
2.04
71.02 4: 0.46
18.2 0.03
6.9 2: 0.08
3.3 0.05
17:63 1: 0.02
SIR-EBSW-e
2044.22
435.02
85.26
0.11
21.96
7.9 1.9+00.22 0.22
3.79
29.8 0.04
IMH-EBSW-eL=100
2018
+:
4600
3.46
81 01:28
1.10
23.35 0.52
8.26 0.26
3.93 0.14
19.3 0.54
MaH--::5W-=1L-000
20448
+ 0.0
454.922
3.22
87.92
0.660
22. 02.666-0 0.46
8.14 0.31
3.82 0.22
62.5 0.09
IS-EBSW-1 L=100
2018. 29 1:
692.63
7.21
167. 75 1: 3.12
11.8 = 0.95
12.31 0.27
5.35 0.1
17.91 1 0.28
SIR-EBSW-I L-100
2048.
704.08
0:47
41.85
12.58
5.64 0.18
MHH-RRS-TI--100
2048.22
0.0
715.97
4.40
171 422 1.225
42.05 0.42
12.6 0.1
5.63 0.06
50.01 + 0:01
RMH-EBSW-I1 L-000
2048.22 20 + 0.0
712.11 1.64
173.47 1.44
42.94 0.4
12,68 0.15
5.54 +: 0.09
64.01 + 0.08
L=10
351.65
2.63
001.5125 2.45
10.11:0.311
0.04
Max T=10
2048.22 2488.2941
0.0
525.72
7.35
134.8 4.8
34.07
0.34
10.77 0.15
T.38 d: 0.31
0.55 0.06
IS-ERSW.e 1-10
2048.22 + 0.0
519.73 8.63
92.14
1.29
23.94
0.07
9.03 - 0.33
4.59 0.22
5.57 0.03
SIR L-10
2048.29 +:
508.86
8.49
104.47
1.30
28.27 0.68
10.56 0.08
5.61 0.16
6.84 0.06
IMH-ERSW-e 1-10
2814
621.51
22.40
131. 75
7.00
34.42
11.53 0.38
0.56 0.0M
0.04
LB1--BBW-->L100
2018. 29 1: 0.0
612.87 +: 5.25
135. 91
8.39
36.11 2.13
12.57 1 0.75
5.91 0.31
9:69 1: 0.01
ES-DDS L-10
2048.
29 1: 0.0
713.65
5.68
177.
1.19
45.07
13.6 = 0.26
6.16 0.22
5.69 0.0
SIR-EBSW-1
2018.
751.
9.37
5.05
41.99
13.59
6.68
0.03
IMH-EBSW-1 L-10
2048.
+ 0.0
772.86
28.00
199. 29
7.00
48.73
1.69
14.1 = 0.49
6.25 0.35
8.61 0.02
RMH-EBBS-1 L=10
20
230 + 0.0
810.1
10.2
212. 11
9.53
04.62
2.63
15.44 0.53
6.74 +: 0.32
9:80 0.06
Table 4: Summary of Wasscrstein-2 2 scores (multiplicd by 10*) from three different runs, computational time in
second (s) to reach step 500 of different sliced Wasserstein variants in gradient flows.
Distanees
Step (W, L)
Step 100 (W2
Step (Wg J)
Step 500 (W; L)
Step 40002, L)
Step 500 (Wg 4)
Time
ES-EBSW-e 1-1100ccc
2048.29 :: 0.0
435.30 #: 1.82
$5.31 +: 0.44
21.9 :: 0.00
7.81 = 0.06
3.68 :: 0.07
17.51
0.01
18-ERSW-1
1=100
2048.20
+ 0.0
711.33 7.2
170.60 + 2:91
42.2 + 0.79
12.62 0.2
5.7 0.11
17.72
0.02
SIR-EBSW-1 L=100
2018.29
685.87
2:65
12.29 0.32
0.56::0..
11.51
IMH-EBSW-1 L=100
20H8
700.47 = 0.13
173.25
1.26
44.08 0.52
13:03: 0.18
5.93 4: 0.2
63.83 0.02
RMIH-EBSW L=100
2018.29
1 0.0
711.0 =: 10.96
175.70 a: 1.45
44.5 =: 0.56
13.39 1 0.13
6.06 4: 0.05
77.32 0.2
18-ERSW-e
2048.22
524.60 7,38
107.37 +
28.46 0.35
10.13 0.38
4.93 0.37
5.54 0.04
IS-EBSW-1 (c)
2018
+:
729.85 53 6.74
179.35
15.03 0.79
13.32
0.82
6.15 0.46
5.7 0.03
SIR-EBBS--1 L-10
2048
+ 0.0
762.233 +: 9.66
202
+ 6.23
66.48
+: 1.56
19.05 +: 0.83
10:42 4: 0.53
8:45 0.02
IMH- E----BSW-1 L=10
2018.29 1
762.67 11:63
200.5 6.48
51.28 1.17
18.11 0.36
9.29 0.26
10.02 0.02
RMIH-EBSW-I L-10
2048.22
817.92 20.86
220.68
2.55
60.15
1.53
0.7
9.8 0.36
11.35 0.03
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 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.
We provide the visualization of the gradient flows from SW (L=100), Max- SW (T=100), v-DSW (L=10,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 consistent. Therefore, the energy-base based slicing function helps to improve the convergence of the source point-cloud to the target point-cloud.
Qualitative Results.
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 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
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,
24

---

Figure 4: Gradient flows from the SW, the Max-SW, the v-DSW, the IS-EBSW-e, the SIR-EBSW-c, the IMH-EBSW-e, and the RMH-EBSW-c turn.
INEESSO SOWES 2712
49TH WPH-SSPP EDCARSS 2212 15405 EOLL 221881
1S SIONNIN dyn 28-21
Figure 5: The first two rows are with L 100, (c) denotes the "parameter-copy" (the SIR-EBSW-e, the IMIL-EBSW-e, the RMIL-EBSW always the "parameter-copy"" estimator since the conventional estimator not stable for them), and last row with 10.
we see that the exponential energy function yields better result for EBSW variants. Similar to the gradient flow, reducing the number of projections to 10 also leads 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.
25

---

Table 5: Reconstruction errors of different autoencoders measured by the (sliced) Wasserstein distance (x 100).
The results from three different runs.
Epoch 20
Epoch 100
Epoch 200
Distance
SW2(J)
W2(1)
SWa (U)
W2(U)
SW₂ (U)
W2(U)
1000
2.97 0.14
12.67
2.29 0.04
10.63 0.05
0.04
9.97 0.08
Max-SW T=100
2.91 + 0.08
12:33 = 0.05
2.24 0.05
10:40 + 0.06
2.14 = 0.10
9.84 + 0.12
v-DSW 1.
2.84 + 0.02
12.64 + 0.02
2.21 0.01
0.04 10.52
2.07 0.09
9.81 0.05
DSW
2.80
12.59
2.16
10.38
2.07
9.78
IS-EBSW-e 1=100
2.68
11.90 0.04
2.18
10.27 = 0.01
2.01
9.69 0 0.14
SIR-EBSWee==100
2.07
12.16
2.24 0.04
10.40 0.01
2.00
9.72 0.04
IMH-EBSW-e L=100
0.03
1..15
0.09
10.39 0.09
1.99 0.05
+
RMH--FBSWeel.-10
0.02
12.21 + 0.03
2.20 + 0.03
10.38 + 0.07
2.02 - 0.02
9.72 0.03
IS-EBSW-1 L=100
2.01
12.37
0.06
10.59 0.07
2.11=
9.90 0.02
SIR-EBSW- 1=100
2.81 2.81+0.000
12.32 + 0.03
2.26 2.26+0000 0.08
10.56 0.14
2.07 0.01
9.81 0.08
IMH-EBSW-I L=100
2.82
12.32
0.11
10.55 4 0.13
2.03 = 0.02
9
RMH-EBSW-I L=100
0.0
0.02
2.22
10.37
- 0.06
2.01 = 0.02
9.73 0.02
SW L=10
0.12
12.70
2.30 0.01
0.04
0.06
10.01 0.09
Max-SW T=10
3.00
12.68 = 0.05
2.31
10.67 0.06
2.14 = 0.01
9.95 + 0.05
IS-EBSW-e L=10
0.04
12.15 0.06
2.20 + 0.08
10.39 0.10
2.04 0.07
9.77 + 0.10
SIR-EBSW-e L=10
0.03
12.26
2.26 0.08
10.53 0.09
9.87
IMH- 1=110
2.82 + 0.02
12.33 0.02
2.26 0.12
10.53 + 0.20
2.07 - 0.02
9.86 + 0.03
RMH-EBSW-e I.=10
0.04
0.01
10.45 0.03
2.02 0.02
0.00
IS-EBSW-I L=10
0.01
0.00
2.28 0.10
10.63
2.00=
0.05
SIR-EBSW- L=10
2.01
0.001
1.07
10.59 1.19
2.07 0.07
9.88 0.122
IMH-EBSW-I L=10
2.52 0.01
2.28 0.08
10.52 0.05
2.08 0.06
9.80
RMH-EBSW-1 1=
2.83 0.04
12.47 0.03
2.21 0.03
10.45 0.08
2.03 0.03
9.80 0.02
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 = 10 and L = 100 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 T
in Table 5 In Table 5
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 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, 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- trruh point-clouds ModelNet40 and their corresponding reconstructed point-clouds from different models (L 100) at epochs 200 and 20 in Figure 6 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- the IMH-EBSW-c, 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 CIFARIO (image size 32x32) 211, and CelebA 233. We show both FID score [15] (CIFAR10. CelebA) and IS score 39 10) in Figure 8 Overall, IS-EBSW-e performs better than SW and DSW this task which suggests that EBSW has wide range of applications.
26

---

Table 6: Reconstruction errors of different autoencoders measured by the (sliced) Wasserstein distance (x 100).
We use (c) for the "parameter-copy" gradient estimator. The results from three different runs.
Epoch 20
Epoch 100
Epuch 200
Distance
SW2()
W2(()
SWa (U)
W2(()
SWa (U)
Wa())
IS --EEB 1=100 (C)
2.74 1: 0.04
12.14
0.12
2.22 1: 0.07
10.42 1 0.05
2.07 = 0.01
9.77 1: 0.07
IS-EBSW-1 1=100 (c)
2.83
0.011
12.34
2.30 1:
10.60 0.09
2.05
0.07
9.83 0.11
SIR-EBSW-1 L=100
2.80 0.02
12.29
0.01
2.21
10.46 1006::000 0.08
2.04 0.02
9.81
0.07
IMH-EBSW-1 L=100
2.96 1: 0.05
12.67
0.08
2.35 0.005
10.82 0.07
2.20 = 0.11
10.20 0.16
RMII-EBSW-I L=100
1.000.000
12.67
0.10
2.27 0.002
10.66 0.06
2.15 - 0.05
10.11 1 0.11
E 1==10
0.01
12.22 + 0.04
2.28 0.09
0.11
2.07 0.07
9.80 + 0.15
IS-EBSW-1 L= (c)
2.061.0.02
12.42 1 0.02
2.24 0.08
10.52
0.113
2.05 - 0.04
9.84 1 0.10
SIR-EBSW-I L=10
2.87 0.02
12.43 1 0.08
2.36
10.67
10.67
2.08 = 0.10
9.88 1 0.14
IMH-EBSW-1
2.98 + 0.02
12.65 0.04
2.35 0.05
10.84 0.06
2.21 = 0.11
10,22
0.11
RMII-EBSW-1 L=10
3.01
12.82
2.37 0.03
10.87
2.11 0.02
D Computational Infrastructure
For the point-cloud gradient flows and color transfer, we use Macbook Pro MI for conducting experiments. For deep point-cloud reconstruction, experiments are run single NVIDIA V100 GPU.
27

---

Figure 6: From the top to the bottom the ground truth, the reconstructed point-clout at epoch 200 of the SW, the Max-SW, the v-DSW, the IS- EBSW-e, the SIR EBSW-e, the IMH-EBSW. and the RMH-EBST respectively.
28

---

Figure 7: From top to the bottom is the ground truth, the reconstructed point-clouds at epoch 20 of the SW, the Max-SW, the v-DSW, the IS EBSW-e, the SIR EBSW-e, the IMH-EBSW- and the RMH-EB ----EBSW-- respectively.
29

---

SW (FID 14.57, IS 8.14)
DSW (FID 14.18, IS 8.17)
IS-EBSW-e (FID 13.24, IS = 8.21)
SW (FII) 9.61)
DSW (FII) 9.46)
IS-EBSW-e (FID 8.99)
Figure 8: Generative modeling from SW, DSW, and -S-EBSW-e L=100.
30


Reviews:
[{"review_id": "rjngV6Px3D", "reviewer": "Reviewer_QgRs", "summary": "This paper proposed a new variant of sliced Wasserstein distance that is inspired from Energy-Based Model called Energy Based Wasserstein Distance (EBWD). The proposed method models the energy function of the distance, from which the slices can be sampled. Three sampling techniques are proposed. The paper evaluate the proposed distance calculation using 3 tasks, Point Cloud Gradient flow, color transferring, and Point Cloud reconstruction. The results have shown that the new variant helps the training process convert faster and the transition between two distribution is smoother (in term of Wasserstein distance and Sliced Wasserstein distance getting small faster).", "questions": "Please see comments in Weaknesses", "limitations": "The paper does not include a Limitation discussion.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper has the following strengths\n\n- A clear motivation to use the energy function. Similar to SW, the proposed energy-based SW enjoys a non-optimization based computation via sampling using the energy function.\n- A detailed theoretical analysis of the proposed distance, accompanied by a nice experimental setup to evaluate its performance.", "weaknesses": "In general, I like the simple, yet interesting proposal from the paper. However, I also have a few concerns:\n\n- It would be reasonable to evaluate EBSW in a more complex task such as image generation, similar to that where DSW is introduced. It seems like only IS-EBSW enjoys the low-computation benefit but other EBSW variants do not.\n- It is not clear to me why EBSW is better than DSW while DSW finds the most separating directions. Especially the fact that EBSW converges faster.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We wish to extend our appreciation for the valuable feedback and comments provided in the review. In response, we would like to address your questions as follows. We are readily available to address additional inquiries and to engage in further discussion.\n\n**Q1**: It would be reasonable to evaluate EBSW in a more complex task such as image generation, similar to that where DSW is introduced.\n\n**A1**: We acknowledge the importance of image generation as a task that sliced Wasserstein variants can excel at. In our paper, we did not include generative modeling due to the well-known instability associated with training generative models using mini-batch stochastic training algorithms. Instead, we focused on a stable setting for comparison and demonstrated the benefits of EBSW in the gradient flow application, which can be seen as a simplified and stable version of generative modeling. Nonetheless, we appreciate the reviewer's feedback and have taken it into consideration. To address this concern, we have now added the results of generative modeling to our paper. This addition showcases EBSW's applicability to various tasks. For detailed experiments and results, we refer the reviewer to the global rebuttal section, where we provide an overview of the experiments and outcomes. Overall, EBSW still shows a favorable performance compared to SW and DSW in this task. Nevertheless, due to the limited time of the rebuttal period, we suggest using the generative modeling result only as a reference for the versatility of EBSW.\n\n\n**Q2**:  It seems like only IS-EBSW enjoys the low-computation benefit but other EBSW variants do not.\n\n**A2**: Thank you for your insightful comment. The reason IS-EBSW is fast is that its computational algorithm is fully vectorizable, i.e., we can stack the supports of measures into matrices, stack the projecting directions into matrices, and then perform computations on matrices (e.g., matrix multiplication) that are parallelizable.\n\nFor sampling importance resampling (SIR-EBSW), the key bottleneck is the resampling step. In particular, our implemented sampling algorithm is not parallel since the computation library (PyTorch) does not seem to support it to the best of our knowledge.\n\nRegarding Markov chain Monte Carlo (IMH-EBSW and RMH-EBSW), the sampling algorithms are naturally sequential; hence, they are not parallelizable. To the best of our knowledge, we could use multiple parallelizable Markov chains as a solution. However, that kind of speeding-up method needs a careful investigation of the mixing time of the MCMC algorithm to design the number of chains and the number of time steps for keeping a budget-constrained algorithm.\n\nHowever, we would like to recall that all discussed specific computation algorithms (IS, SIR, MCMC) in the paper are just basic variants in their vast literature and can be significantly improved. Therefore, we will leave this investigation for future work since our main focus of the paper is the energy-based slicing distribution and EBSW. At the moment, we recommend importance sampling (IS-EBSW) as the main computation method for EBSW due to its fast computation and simplicity.\n\n\n**Q3**: It is not clear to me why EBSW is better than DSW while DSW finds the most separating directions. Especially the fact that EBSW converges faster.\n\n**A3**: We would like to direct the reviewer to the global rebuttal concerning the discussion on DSW. In our paper, we present a budget-constrained comparison that places a computational time limit on distances, defined by a scaling constraint of $n \\log n$. As a result, DSW might not attain optimality. Furthermore, an optimal design for the slicing distribution in DSW remains unknown. In our study, we employ the von Mises Fisher family for the slicing distribution, chosen for its simplicity and efficiency. Nevertheless, it is still slower than EBSW with importance sampling. In our additional experiment, we also use DSW with an implicit slicing distribution induced by a neural network; however, it still does not perform as effectively as IS-EBSW-e in applications. To clarify, we are not questioning DSW in its population form, but rather its practical usage, which presents challenges such as sub-optimality and high computational complexity. In conclusion, EBSW with importance sampling emerges as the preferred choice due to its simplicity, effectiveness, and efficiency.\n\n**Q4**: The paper does not include a Limitation discussion.\n\n**A4**: Thank you for pointing this out. In the revision, we have gathered all limitations of the papers into a new paragraph. We would like to refer the reviewer to the global rebuttal for the discussion about the limitation of the paper. Here, we would like to summarize the discussion. As discussed in **Q2**, one limitation of EBSW is its non-parallelizable MCMC variations, resulting in sluggish computations.  Furthermore, selecting suitable energy-based functions presents a challenge in applications. Finally, demonstrating the triangle inequality of EBSW is difficult due to the complexity of the energy-based slicing distribution. Overall, as discussed in the global rebuttal, these limitations are specific to the current paper itself, rather than EBSW.\n"}, {"review_id": "ny9XBgSC9f", "reviewer": "Reviewer_dYij", "summary": "This paper has proposed an energy-based slicing distribution that maps original distributions into one-dimesion space  to compute the Wasserstein distance.", "questions": "1. In definitition 1, does $\\theta$ represent parameters in the slicing distribution or just the transport function? If  $\\theta$ are parameters, why could the model be called as parameter-free?\n\n2. In general, the energy function should be defined as $f:(-\\inf, \\inf)^d \\rightarrow [0, \\inf]$. Why the author defines it as in line 144?\n\n3. Could the proposed distance measurement be used in generative models, like GANs and EBMs?\n\n4. Is there any comparison with other probability distance measurement methods, like KL divergence, in the experiments?\n\n\n", "limitations": "This paper has proposed an energy-based slice distribution to compute Wasserstein distance between two distributions. My question mainly lies in the application of energy function and comparison of methods based on other  probability measurement methods.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This paper has proposed an energy-based slicing distribution that maps original distributions into one-dimesion space  to compute the Wasserstein distance. The proposed method shows great advantages over other slicing distribution like sliced Wasserstein, Distributional sliced Wasserstein, and Max sliced Wasserstein. The authors also explore different sampling methods to approximate the value of the EBSW distance.", "weaknesses": "1. An introduction about optim transport should be included in the Background section.\n2. Put a title for each row in Figure 2 could help reader easier to compare different SW methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We begin by expressing our gratitude for the insightful feedback and comments provided in the review. In response, we offer the following explanations and answers to your inquiries. We welcome any further questions and are open to engaging in a more in-depth discussion on the matter.\n\n**Q5**: An introduction about optimal transport should be included in the Background section.\n\n**A5**: Thank you for your suggestion. Due to the space constraint of the main paper, we skip the definition of optimal transport and Wasserstein distance. From your feedback, we realize that it affects the readability of the paper. Therefore, we will add those backgrounds to the paper in the revision.\n\n**Q6**: Put a title for each row in Figure 2 could help reader easier to compare different SW methods.\n\n**A6**: Thank you for your suggestions. We will add the legend to the figure for better readability.\n\n**Q7**: In definition 1, does $\\theta$ represent parameters in the slicing distribution or just the transport function? If $\\theta$  are parameters, why could the model be called as parameter-free?\n\n**A7**: Thank you for your insightful question. In Definition 1, $\\theta$ is the realization of the distribution, not the parameter. For a more detailed explanation, in DSW, $\\theta \\sim vMF(\\epsilon,\\kappa)$, where $\\epsilon$ and $\\kappa$ are parameters, while EBSW does not have such parameters. Therefore, EBSW is parameter-free. To elaborate further, EBSW is defined based on the energy function which non-parametric and is easier and more flexible to choose than choosing a family of distributions as in DSW. In addition, in its computational form, Max-SW and DSW require additional hyperparameters for optimization, such as learning rate, the number of gradient steps, optimizers, and so on, while EBSW has only one hyperparameter, which is the number of projections, like the vanilla SW. Overall, this is the reason we refer to EBSW as parameter-free and optimization-free. We will include this explanation in the revision.\n\n**Q8**: In general, the energy function should be defined as  $f :(−\\infty,\\infty)^d\\to [0,\\infty]$. Why the author defines it as in line 144?\n\nA8: Thank you for your question.  We would like to clarify it as follow. The energy function takes the value of the $p$-power Wasserstein distance between two projected one-dimensional measures and yields a positive number. Since the Wasserstein distance is a scalar and always positive, the pre-image of the energy function is in the range of $[0, \\infty]$. Based on the energy function, we can define the energy-based slicing distribution, which is a distribution on the unit-hypersphere and has a probability density function (pdf): $\\mathbb{S}^{d-1} \\to \\Theta \\in (0, \\infty)$, as described in Definition 1.  The reason we avoid having the point 0 in the image of the energy function is to prevent the slicing distribution from being undefined when two measures coincide. This ensures the distribution's well-defined nature and allows us to use it safely in practical applications. We will add a more detailed explanation to the revision.\n\n\n**Q8**: Could the proposed distance measurement be used in generative models, like GANs and EBMs?\n\n**A8**: The answer is yes. The reason we focus on point-cloud applications since they are disjointed supports that other divergences such as KL divergence, Jensen Shannon divergence, f-divergence cannot handle. EBSW can absolutely be used in generative modeling. We would like to refer the reviewer to the global rebuttal for a detailed additional result on deep generative modeling. Overall, IS-EBSW-e has shown a favorable performance compared to SW and DSW in generative modeling. However, we suggest using this additional result as evidence to show that EBSW is versatile since we only a limited time in the rebuttal period.\n\n**Q9**: Is there any comparison with other probability distance measurement methods, like KL divergence, in the experiments?\n\n**A9**: As stated in **Q8**, the KL divergence (one member of f-divergence) is not suitable for handling disjointed support measures. The reason behind this is that the KL divergence relies on the density ratio, which becomes undefined when the denominator is equal to 0. Consequently, KL divergence cannot be directly applied to compare two-point clouds. If one wishes to use KL divergence between two point clouds, it becomes necessary to convert them into voxels (3D bins). Nevertheless, this transformation is not differentiable, rendering it unsuitable for deep learning applications involving point clouds. Therefore, using KL divergence in such scenarios is not practical.\nMoreover, it is important to note that KL divergence is not symmetric, while Wasserstein variants are all symmetric. This means that the Wasserstein distance and its variants, such as the Earth Mover's Distance, do not depend on the ordering of the points and provide a more symmetrical measure for comparing two probability distributions. \n\nFor a comparison in generative modeling, training conventional SNGAN can be seen as an application of Jensen Shannon divergence (which can be seen as a symmetric version of KL)  which achieves a FID score of   21 [1]. By using the same architecture, we achieve a FID score of 13.24 with IS-EBSW-e in the rebuttal pdf.\n\n[1] Spectral Normalization for Generative Adversarial Networks, Miyato et al\n\nIn summary, due to its limitations with disjointed support measures and its lack of symmetry, KL divergence is not the most appropriate choice for comparing point clouds directly. Wasserstein-based distance metrics offer more suitable alternatives for handling these scenarios, especially in deep learning applications involving point clouds.\n"}, {"review_id": "OWHd8wXwXa", "reviewer": "Reviewer_vaua", "summary": "This paper proposes an extension to Sliced Wasserstein Distance (SW), an approach for measuring distances between distributions by computing the average of the energy of the 1-d Wasserstein distances between 1-d projections. The authors argue that moving towards non-uniform ways of sampling the projections is key, and subsequently show that the benchmark approach currently that does so (Distributional sliced Wasserstein) adds another optimization over a parametric distribution family, which may not yield stable results and is more computationally expensive. Subsequently, they propose a simple non-parametric extension to SW called Energy based SW (EBSW) that achieves the same objective of sampling non-uniformly, i.e. sampling projections with larger Wasserstein distance with a greater probability, by considering energy functions $f$ which are monotonically increasing w.r.t the sliced 1-d Wasserstein distance itself. Theoretical properties are highlighted, including the convergence properties of their proposed metric. Next, the authors propose a range of Monte Carlo estimation methods to approximate EBSW, including deriving the necessary gradient expressions to be used in some of their subsequent applications. Lastly, over multiple experiments ranging from point-cloud gradient flows to deep point-cloud reconstruction, the authors demonstrate that EBSW is faster (Importance Sampling approach), and yields distributions closer to the ground truth in most cases than other SW-based benchmarks.\n", "questions": "In addition to the list of weaknesses above, below are some questions and general suggestions for the authors to address:\n1. I noticed that the authors fix the value of $p$ to 2 for all experiments. I wonder if increasing $p$ can achieve a similar effect to using a monotonically increasing energy function (such as the exponential distribution in this work). Because intuitively, it seems to me that increasing $p$ eventually also gradually assigns more importance to the projections with values near to the maximum. I’m mainly asking this because increasing $p$ still falls under SW-p, and thus could perhaps have similar sample complexity and computational load. Could you please provide a theoretical explanation on how changing $p$ fundamentally differs from EBSW with an increasing energy function?\n2. For Figure 2, the authors state that the gradient flows for EBSW are smoother than other approaches. Perhaps describing some visual cues to elaborate on that point can make it easier to follow.\n3. Although four different ways of estimating the EBSW variants are proposed, I only see the IS-EBSW being reported in the main paper. Furthermore, having looked at the results in the appendix, it seems in almost all cases, the variants have roughly similar performance, and IS-EBSW seems to lead to better results in most cases. Furthermore, it also seems that the time complexity for all algorithms are similar, also seeing Table 3 in C.1, it seems that in most cases IS-EBSW is faster as well. Thus, my question is: in discussing these other approaches, are there benefits to the other proposed MCMC variants?\n4. Are there any cases where the choice of an exponential energy function $f$ doesn't work? It would be intuitive to me that, considering a more general family of exponentials in $e^{(\\lambda x)}$, the optimal \"width\" of the slicing distribution as set by $\\lambda$ may be different for different cases. For instance, can there be any cases where the underlying distribution enforces a much larger or smaller $\\lambda$ than one? Or is $\\lambda=1$ in some sense universal?\n5. As a follow-up to the previous question, I would assume that for the optimization results in Tables 1 and 2, as the generated distribution edges closer to the ground truth distribution, would it be profitable to increase $\\lambda$ higher than one? As then the distances would all get very small, and potentially max-SW would start to look more appealing.   \n6. The theoretical results are interesting and important, however, could the authors put into perspective how the theoretical properties of EBSW compare to other SW based measures? It seems to me that both Theorems 1 and 2 may hold for SW based measures (at least SW, not sure about max-SW and DSW). Similarly, for Proposition 2 it seems from the proof that it would hold for both SW and max-SW (not sure about DSW). Additional theoretical results in this direction, or at the least a rough intuitive discussion on how these results would compare for the other SW counterparts, would be useful. This would put more into perspective how EBSW's theoretical properties compare to the other SW benchmarks.", "limitations": "Some limitations have been discussed, but I feel there is scope to discuss more limitations of the proposed approach.\n", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I really liked reading the paper and going through the various findings in it. The following are what I consider to be some of the main strengths of this work. (i) The paper is really well-written. All sections were intuitive to read and easy to understand for me, and this also applied to the Appendices, including the proofs of the Theorems and Propositions. (ii) The contributions of the paper are very clear and supported by the theoretical and empirical results. All proposed benefits of the approach have been verified by the authors in theory and practice. (iii)  The experiments are well motivated and relatively exhaustive. The proposed approach is tested in a wide range of problems and the results convincingly showcase the benefits of EBSW. (iv) The background section is well organized and is an effective introduction to SW-based metrics in general.\n", "weaknesses": "The paper is overall interesting and well-written. However, there are a few points noted below that can improve the work further.\n1. I feel that the theoretical results can be given a bit more perspective w.r.t. the other SW-based metrics. Overall, I felt the theoretical results in Theorems 1, 2 and Propositions 1 and 2 are relatively intuitive to show. Some more specific insights on EBSW could be interesting (perhaps with constrained energy functions, i.e. Lipschitz constrained $f$).\n2. I understand the overall direction for moving towards measures that look disproportionately at projections which have larger distances (lines 37-38). However, I feel that a few more words to elaborate on intuitively why one must move towards non-uniform distributions in the Sliced Wasserstein setting could be useful.\n3. One of the main questions I have is the role of $p$, which is fixed for all experiments in this work. I feel that the authors can elaborate a bit on whether trying larger values of $p$ can achieve a similar goal to what EBSW-e achieves (More details in questions).\n4. The result tables are convincing and informative w.r.t. the tested methods. However, in some of the visual depictions, I found it hard to see any significant differences between EBSW and some other variants (even in the Appendices). Mainly, the color transfer experiments were a bit hard to discern for me in terms of relative performance", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We'd like to express our gratitude for the constructive feedback. We're here to provide responses as outlined below. We are eager to participate in further discussions.\n\n**Q10**: On Lipschitz constrained $f$\n\n**A10**: Regarding the Lipschitz energy function, the polynomial function is Lipschitz and the exponential function is locally Lipschitz. Lipschitzness could lead to a potential notion of robustness to the distance. Specifically, when two measures are contaminated with additive noise, a significant change in the Wasserstein value from two closed projecting directions could be due to the noise. Therefore, using a Lipschitz energy function could penalize the weight given to such projecting directions and has benefits.\n\nFrom a Lipschitz energy function, we can show that the energy-based slicing distribution has a Lipschitz probability density function (pdf) since the SW distance is Lipschitz with respect to the projecting direction [1]. However, an exact computation of the Lipschitz constant is not trivial in the case of the SW distance. \n\n[1] Statistical, Robustness, and Computational Guarantees for Sliced Wasserstein Distances, Goldfeld et al. \n\n**Q11**: Why one must move towards non-uniform distributions in the  SW setting could be useful?\n\n**A11**: It benefits applications that have a sequence of measures that converge to a target measure e.g., gradient flow, and generative modeling. When using SW as the criteria to drive the sequence, the current measure is derived toward the target measure from all directions equally. This is undesirable since two measures can be closed in some directions and can be significantly different in other directions. Hence, the uniform distribution leads to sequences that converge slower.\n\n**Q12**:whether trying larger values of $p$ can achieve a similar goal to what EBSW-e achieves? How changing  $p$  fundamentally differs from EBSW?\n\n**A12**:  As in Proposition 1, we show that EBSW is an upper bound of SW for any $p$.  Adjusting $p$ equals to change the notation of ground metric on the support sets of two measures. However, changing $p$ in SW does not affect the slicing distribution which is still uniform. Therefore, larger values of $p$ cannot achieve a similar goal as EBSW since EBSW has a non-uniform distribution. To verify further the intuition, we run the experiments on gradient flow again for SW with p=3 and p=10 in the global rebuttal.\n\nIn contrast, changing $p$ does affect the EBSW  since the energy distribution is proportional to the composition between $f$ and $L_p$ norm. Due to the flexibility in choosing $f$, the interaction between $f$ and the norm is expressive and worth a careful future investigation.\n\n**Q13**:  On visual comparison in GF.\n\n**A13**: One way to see the difference from our experience is by looking at the level of “yellow” and “orange\" colors. In more detail, the level of yellow color is more than the level of orange color in the target image while the reverse phenomenon happens in the source image. From that perspective, the image from EBSW has more yellow and less orange than images from other distances. \n\n**Q14**: On some visual cues to elaborate on the performance of EBSW in gradient flow.\n\n**A14**: From the figure, we observe that the legs of the table  (the source point-cloud) are bent in SW, Max-SW, and vDSW in the first step and do not look like a real table. In contrast, EBSW can create a more real table with a sharp connection between the legs and the desktop, and the legs are also not bent. In future steps, we see that the point-clouds from EBSW have fewer jumping points on the surface e.g., the cover of the light bulbs. \n\n**Q15**: benefits of proposed MCMC variants?\n\n**A15**: As discussed in the global rebuttal, MCMC variants are slow since they are not parallelable.  However,  using MCMC variants could benefit the approximation when we have a very peaky slicing density. In that case, important sampling might be not accurate since the proposal distribution could not yield good samples due to the concentrated mass of the slicing density. Moreover,  the importance sampling is also unstable to use due to the huge value of the density ratio between the slicing density and the proposal density. However, MCMC can overcome these issues with a good transition distribution and achieve a better approximation. \n\n**Q16**:  On the energy function $e^{\\lambda x}$... Would it be profitable to increase $\\lambda>1$? \n\n**A16**: With the exponential energy function, the importance sampling reverts to using the Softmax function. For your suggestion, it leads to the annealed exponential function, we get the annealed Softmax function which is better at controlling the energy density. However, it costs one hyperparameter in addition. A potentially useful choice of $\\lambda$ is $1/p$ which helps us to prevent the effect of changing $p$ to the slicing distribution as discussed in **Q12**. \n\nIncreasing $\\lambda$ leads to a peaker slicing density and makes the density closer to the Dirac delta at the max slice. However, a peaky density is hard and unstable to approximate as discussed in **Q15**. Secondly, such kind of peak density is not desirable since they are strongly affected by the noise as discussed in **Q10**. \n\n**Q17**: How EBSW's theoretical properties compare to the other SWs?\n\n**A17**: Theorem 2, and Proposition 2 mean that EBSW retains nice properties as SW and Max-SW i.e., weak convergence, and no curse of dimensionality. However, in Theorem 1, the triangle inequality of EBSW has not been proved yet due to the adaptive slicing distribution. Given a fixed slicing distribution, the triangle inequality can be proved using the metricity of Wasserstein distance and Minkowski inequality. In contrast, it is not trivial how to obtain the triangle inequality with the energy-based slicing distribution in EBSW. \n\n**Q18**: On Limitation\n\n**A18**: We would like to refer the reviewer to the global rebuttal of the limitations of our work.\n"}, {"review_id": "rwj2aZfMPU", "reviewer": "Reviewer_krQV", "summary": "This paper introduces a new distribution of slices of the Sliced-Wasserstein distance based on the energy-based model framework. The authors study the theoretical properties and conduct numerical experiments on EBSW.", "questions": "- Can the authors explain why experimental results are not compared to the DSW distance? \n- It appears to me that EBSW falls into the family of adaptive Sliced-Wasserstein distances ( https://arxiv.org/pdf/2206.03230.pdf , which doesn't seem to be cited in the related work section). Can the authors comment on that and whether this framework can help studying this distance?", "limitations": "None. I am willing to discuss with the authors and accordingly modify my score.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The study of the idea is organized and clear. \n- Theoretical quantities of interest are derived.\n- Various sampling methods are provided.", "weaknesses": "- The method is not compared with the vanilla Distributional Sliced-Wasserstein method which should be the baseline to beat. The fact that only the vMF-DSW method is benchmarked is a weakness, as the learned distribution is unimodal, which prevents the model to be too expressive. \n- I think it would be good to have a theorem studying the following sample complexity, that is more of interest for practical reasons:\n$ \\mathbb{E}[EBSW_p(\\mu_n,\\nu_n;f) - EBSW_p(\\mu,\\nu;f)]$\n- Experiments are a bit light, I think a generative modeling experiment would be nice to have as this is one of the main applications of the SW distances.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the review's feedback and comments and want to express our thanks. Our responses are as follows. We remain open to additional questions and further discussion.\n\n**Q19**: The method is not compared with the vanilla Distributional Sliced-Wasserstein method which should be the baseline to beat. The fact that only the vMF-DSW method is benchmarked is a weakness, as the learned distribution is unimodal, which prevents the model to be too expressive.\n\n**A19**:  We would like to refer the reviewer to the global rebuttal for additional experiments comparing our method with vanilla DSW. In the submission. the primary reason we chose the von Mises-Fisher (vMF) distribution as the family of slicing distributions is its simplicity and parameter efficiency. The vMF distribution has only two parameters: the location parameter (a d-dimensional vector) and the concentration parameter (a scalar). Therefore, performing stochastic optimization for vMF is expected to be more stable. In contrast, using an overparameterized family of distributions, such as the push-forward distribution with neural networks in vanilla DSW, can be computationally expensive. For instance, using an MLP neural network could increase the computational complexity to be proportional to $d^2$ due to matrix multiplications, as opposed to the original scaling of $d$ in SW and EBSW. Additionally, vanilla DSW requires a constraint to control the concentration of the push-forward distribution, making the optimization less interpretable and potentially more challenging. Moreover, solving vanilla DSW necessitates an admissible regularizing constant and involves using duality, which could further complicate the optimization problem.\n\nIt's important to note that we are not challenging DSW in its population form but rather in its computational form. Utilizing a powerful family of slicing distributions with a high number of parameters would incur increased computational time, memory usage, and hyperparameter tuning. In our paper, we present a budget-constrained comparison, specifically the scaling constraint of $n \\log n$. Consequently, we believe that using vanilla DSW may not offer significant improvements over vDSW. We would like to emphasize the main advantage of EBSW, which lies in its simplicity of use (choosing an energy function is easier than selecting a family of distributions) and computational efficiency (parameter-free and optimization-free).\n\n**Q20**: On two-sided sample complexity\n\n**A20**: Thank you for your insightful question. Proposition 1 can be referred to as the one-sided sample complexity, which is useful in certain settings, e.g., quantization of a distribution (clustering). The sample complexity you mentioned is called the two-sided sample complexity. Normally, we can directly derive the two-sided sample complexity from the one-sided sample complexity by using the triangle inequality. Unfortunately, the triangle inequality of EBSW has not been proven in our paper, due to the expressiveness of the energy-based slicing distribution. Specifically, while we can prove the triangle inequality using the triangle inequality of Wasserstein distance and the Minkowski inequality for a fixed slicing distribution, it is not trivial to establish this inequality for the adaptive slicing distribution of EBSW. We believe that proving the triangle inequality for EBSW presents a challenging theoretical problem, and we leave this open problem for future work.\n\n\n**Q21**: Experiments are a bit light, I think a generative modeling experiment would be nice to have as this is one of the main applications of the SW distances.\n\n**A21**: We would like to mention that it takes about one day to complete the training of a deep point-cloud autoencoder on a Tesla V100 GPU. Therefore, we believe our experiments are relatively intensive compared to the current literature on proposing new SW variants [1] [2]. Nevertheless, we agree that generative modeling is one of the main applications of SW distances. For detailed additional experiments on deep generative modeling, we refer the reviewer to the global rebuttal.\n\n\n[1] Generalized sliced Wasserstein distances, Kolouri et al\n\n[2] Spherical Sliced Wasserstein, Bonet et al\n\n\n**Q22**: Can the authors explain why experimental results are not compared to the DSW distance?\n\n\n**A22**: We do use v-DSW in the paper as an instance of DSW, which is more interpretable and simpler. As discussed in **Q21**, vanilla DSW requires more parameters and computation compared to vDSW due to the usage of neural networks for push-forward slicing distribution. However, we have added a comparison with DSW in the rebuttal PDF which shows the favorable performance of EBSW. We would like to refer the reviewer to the global rebuttal for detailed experimental results.\n\n\n**Q23**: It appears to me that EBSW falls into the family of adaptive Sliced-Wasserstein distances ( https://arxiv.org/pdf/2206.03230.pdf , which doesn't seem to be cited in the related work section). Can the authors comment on that and whether this framework can help studying this distance?\n\n\n**A23**; Thank you for providing the reference to the new ICML 2023 paper. We will include this paper in our references and discussions. In the paper, they define adaptive sliced Wasserstein distances (ASW) based on an optimization problem, similar to DSW. In contrast, EBSW is an optimization-free variant of SW. EBSW employs an adaptive slicing distribution, which is explicitly constructed, unlike the implicit construction in ASW and DSW. It would be interesting to explore if we can formulate EBSW as the optimal solution for ASW and DSW, under certain choices in the family of slicing distributions and regularity conditions. With such a connection, we can leverage the theoretical results from the mentioned paper to study EBSW in greater detail. Overall, we believe this direction holds great promise, and we will leave this investigation for future research.\n"}]

Author rebuttal:
First, we would like to thank the reviewers for their time and feedback. We would like to summarize some additional experiments in the rebuttal PDF and address common questions from the reviewers. Other questions are addressed in the corresponding rebuttal of reviewers.

1. **Additional experiments on deep generative modeling.** As suggested by Reviewer **QgRs** and Reviewer **krQV**, we conducted additional experiments on deep generative modeling using the framework proposed in [1]. We compared SW, vanilla DSW (which utilizes an MLP neural network to model the slicing distribution as suggested by Reviewer **krQV**), and IS-EBSW-e on CIFAR10 and CelebA datasets. For SW and IS-EBSW-e, we set L=100. For DSW, we set $L=10, T=10$, and reported the best result for $\lambda  \in \{1, 10, 50\}$ and slice learning rate ($\eta \in \{0.001, 0.01, 0.1\}$). We provided both quantitative results (FID score and IS score) and qualitative results (randomly generated images) in Figure 1 of the rebuttal PDF. Overall, we observed that IS-EBSW-e yielded the lowest FID score and the highest IS score. However, due to time limitations in the rebuttal, we recommend using this result as an example to showcase the versatility of EBSW, acknowledging that we ran experiments only once.

   [1] Amortized Projection Optimization for Sliced Wasserstein Generative Models, Nguyen et al.

2. **Comparison to DSW.** As suggested by Reviewer **krQV**, we also incorporated vanilla DSW into our applications. Firstly, we applied the gradient flow with DSW and presented the results in Figure 2 and Table 2 of the rebuttal PDF. We observed that vanilla DSW performed better than v-DSW but required more computation time. Overall, DSW still falls short compared to IS-EBSW-e in terms of performance and computational efficiency. Similarly, we included DSW in the deep point-cloud reconstruction application, where DSW's performance was comparable to v-DSW but inferior to our proposed IS-EBSW-e.

     Furthermore, we would like to emphasize the advantages of EBSW over DSW. Utilizing DSW involves two main steps: designing a family of distributions over the unit-hypersphere and selecting the best member through an interactive optimization procedure. Both steps present challenges; in the first step, selecting a suitable family remains an open question, and in the second step, optimization often suffers from local minima and is computationally expensive due to its iterative nature. Moreover, when performing statistical inference, using DSW (and other optimization-based SW, e.g., Max-SW) leads to a minimax problem known to be difficult to optimize. In contrast, EBSW is optimization-free, parameter-free, and stable. Although EBSW requires choosing an energy function, this process is easier and more flexible than selecting a family of distributions. Furthermore, EBSW can be computed exactly thanks to the asymptotic mixing property of importance sampling and Markov chains.

    In the paper, our criticism of DSW pertains to its computational form rather than its population form. Specifically, we provide an approximate budget constraint for DSW and EBSW (in terms of scaling constraint of $n \log n$) in our experiments, which means that DSW might not achieve its optimality. Additionally, we use the von Mises Fisher slicing distribution for DSW, which might not be the optimal choice. However, it's worth noting that the optimal choice for DSW is still unknown, and adopting a more expressive family could lead to increased memory and computation costs, which is unfavorable in a budget-constrained comparison setting.

    Overall, EBSW introduces a novel idea of energy-based slicing distribution that directly and explicitly conditions two input measures, resulting in a simple, effective, and efficient computation. Therefore, we believe it is easier to use EBSW in practice.

3.  **Increasing $p$ in SW.** As suggested by Reviewers **vaua**, we increased the value of p in SW, i.e., $p = 3, p = 10$ in gradient flow. However, we found that increasing $p$ does not lead to better performance for SW. The reason is that changing $p$ does not affect the slicing distribution in SW; it remains uniform. In contrast, EBSW adapts its slicing distribution changes adaptively. Therefore, increasing $p$ in SW does not have the same effect as in EBSW.


4.  **Limitations of the paper.** We will add a new paragraph to discuss the limitations of our paper, as suggested by Reviewers **QgRs** and **vaua**. In summary, the first limitation of EBSW is that its MCMC variants are not directly parallelizable, resulting in slow computation. Thus, conducting a more detailed investigation of customized MCMC methods for EBSW is crucial to unlocking its true performance. Additionally, the effective choice of energy-based functions is challenging. In the paper, we use simple and computationally effective energy-based functions, such as the exponential function and the polynomial function. However, there is ample room to explore other types of energy functions, such as the annealed exponential function and Lipschitz function, as suggested by Reviewer **vaua**. Finally, proving the triangle inequality of EBSW is challenging due to the expressiveness of the energy-based slicing distribution, which poses difficulties in deriving the two-sided sample complexity of EBSW, as questioned by Reviewer **krQV**. Overall, we believe these limitations pertain to the current paper itself, not EBSW.


Best regards,

Authors,


Meta-review: There is an overall consensus among reviewers that this is an interesting submission that should be accepted, which is also my evaluation.
Decision: Accept (poster)