Title: Energy-Based Sliced Wasserstein Distance
Venue: NeurIPS 2023
Accepted: yes
Keywords: Sliced Wasserstein

Monte Carlo Methods

Point-Cloud

Optimal Transport
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

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 1 .

1 Introduction

The sliced Wasserstein [2] (SW) distance is a sliced probability metric that is derived from the Wasserstein distance [5, 36] as the base metric. Utilizing the closed-form solution of optimal transport on one-dimension [36], the SW distance can be computed very efficiently at the time complexity of O ( n log n ) and the space complexity of O ( n ) when dealing with two probability measures that have at most n supports. Moreover, the sample complexity of the SW is only O ( n -1 / 2 ) [27, 34] which indicates that it does not suffer from the curse of dimensionality in statistical inference. Therefore, the SW distance has been applied to various domains of applications including point-cloud applications e.g., reconstruction, registration, generation, and upsampling [31, 41], generative models [8], domain adaptation [22], clustering [20], gradient flows [24, 1], approximate Bayesian computation [26], variational inference [47], and many other tasks.

The SW distance can be defined as the expectation of the one-dimensional Wasserstein distance between two projected probability measures. The randomness comes from a random projecting direction which is used to project two original probability measures to one dimension. The probability distribution of the random projecting direction is referred to as the slicing distribution. Therefore, a central task that decides the effectiveness of the SW in downstream applications is designing

slicing distribution. The conventional sliced Wasserstein distance [2] simply takes the uniform distribution over the unit-hypersphere as the slicing distribution. Despite being easy to sample from, the uniform distribution is not able to differentiate between informative and non-informative projecting distributions in terms of discriminating two interested probability measures through projection [7]. To avoid a flat prior distribution like the uniform distribution, a different approach tries to find the best slicing distribution that can maximize the expectation. This distribution is found inside a parametric family of distribution over the unit-hypersphere [30, 32]. However, searching for the best slicing distribution often requires an iterative procedure which is often computationally expensive and unstable. Moreover, choosing the family for the slicing distribution is challenging since the number of distributions over the unit-hypersphere is limited. Widely used and explicit spherical distributions such as von Mises-Fisher distribution [17, 32] might be misspecified while implicit distributions [30] are expensive, hard to adapt to downstream applications and uninterpretable.

In this paper, we aim to develop new choices of slicing distributions that are both discriminative in comparing two given probability measures and do not require optimization. Motivated by energybased models [25, 16], we model the slicing distribution by an unnormalized density function which gives a higher density for a more discriminative projecting direction. To induce that property, the density function at a projecting direction is designed to be proportional to the value of the onedimensional Wasserstein distance between the two corresponding projected probability measures.

Contribution. In summary, our contributions are three-fold:

1. We propose a new class of slicing distribution, named energy-based slicing distribution, which has the density function proportional to the value of the projected one-dimensional Wasserstein distance. We further control the flexibility of the slicing distribution by applying an energy function e.g., the polynomial function, and the exponential function, to the projected Wasserstein value. By using the energy-based slicing distribution, we derive a novel metric on the space of probability measures, named energy-based sliced Wasserstein (EBSW) distance.

2. We derive theoretical properties of the proposed EBSW including topological properties, statistical properties, and computational properties. For topological properties, we first prove that the EBSW is a valid metric on the space of probability measures. After that, we show that the weak convergence of probability measures is equivalent to the convergence of probability measures under the EBSW distance. Moreover, we develop the connection of the EBSW to existing sliced Wasserstein variants and the Wasserstein distance. We show that the EBSW is the first non-optimization variant that is an upper bound of the sliced Wasserstein distance. For the statistical properties, we first derive the sample complexity of the EBSW which indicates that it does not suffer from the curse of dimensionality. For computational properties, we propose importance sampling, sampling importance resampling, and Markov Chain Monte Carlo methods to derive empirical estimations of the EBSW. Moreover, we discuss the time complexities and memory complexities of the corresponding estimations. Finally, we discuss the statistical properties of estimations.

3. We apply the EBSW to various tasks including gradient flows, color transfer, and point-cloud applications. According to the experimental result, the EBSW performs better than existing projectionselection sliced Wasserstein variants including the conventional sliced Wasserstein [2] (SW), max sliced Wasserstein [7] (Max-SW), and distributional sliced Wasserstein (DSW) [30, 32]. More importantly, the importance sampling estimation of the EBSW is as efficient and easy to implement as the conventional SW, i.e., its implementation can be obtained by adding one to two lines of code.

Organization. The remainder of the paper is organized as follows. We first review the background on the sliced Wasserstein distance and its projection-selection variants in Section 2. We then define the energy-based sliced Wasserstein distance, derive their theoretical properties, and discuss its computational methods in Section 3. Section 4 contains experiments on gradient flows, color transfer, and point-cloud applications. We conclude the paper in Section 5. Finally, we defer the proofs of key results, and additional materials in the Appendices.

Notations. For any d ≥ 2 , we denote S d -1 := { θ ∈ R d | || θ || 2 2 = 1 } and U ( S 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 ≥ 1 , P p ( X ) is the set of all probability measures on the set X that have finite p -moments. For any two sequences a n and b n , the notation a n = O ( b n ) means that a n ≤ Cb n for all n ≥ 1 , where C is some universal constant. We denote θ♯µ is the push-forward measures of µ

through the function f : R d → R that is f ( x ) = θ ⊤ x . For a vector X ∈ R dm , X := ( x 1 , . . . , x m ) , P X denotes the empirical measures 1 m ∑ m i =1 δ x i .

2 Background

In this section, we first review the sliced Wasserstein distance and its projection-selection variants including max sliced Wasserstein distance and distributional sliced Wasserstein distance.

Sliced Wasserstein. The definition of sliced Wasserstein (SW) distance [2] between two probability measures µ ∈ P p ( R d ) and ν ∈ P p ( R d ) is:

where the Wasserstein distance has a closed form which is W p p ( θ♯µ, θ♯ν ) = ∫ 1 0 | F -1 θ♯µ ( z ) -F -1 θ♯ν ( z ) | p dz where F θ♯µ and F θ♯ν are the cumulative distribution function (CDF) of θ♯µ and θ♯ν respectively. However, the expectation in the definition of the SW distance is intractable to compute. Therefore, the Monte Carlo scheme is employed to approximate the value:

where θ 1 , . . . , θ L i.i.d ∼ U ( S d -1 ) and are referred to as projecting directions. The pushfoward measures θ 1 ♯µ, . . . , θ L ♯µ are called projections of µ (similar to ν ). 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 ( Ln log n ) and O ( L ( d + n )) respectively. It is worth noting that ̂ SW p p ( µ, ν ; L ) is an unbiased estimation of SW p p ( µ, ν ) , however, ̂ SW p ( µ, ν ; L ) is only asymptotically unbiased estimation of SW p ( µ, ν ) . Namely, we have ̂ SW p ( µ, ν ; L ) → SW p ( µ, ν ) when L →∞ (law of large numbers).

Distributional sliced Wasserstein. As discussed, using the uniform distribution over projecting directions is not suitable for two general probability measures. A natural extension is to replace the uniform distribution with a "better" distribution on the unit-hypersphere. Distributional sliced Wasserstein [30] suggests searching this distribution in a parametric family of distributions by maximizing the expected distance. The definition of distributional sliced Wasserstein (DSW) distance [30] between two probability measures µ ∈ P p ( R d ) and ν ∈ P p ( R d ) is:

where σ ψ ( θ ) ∈ P ( S d -1 ) , e.g., von Mises-Fisher [17] distribution with unknown location parameter σ ψ ( θ ) := vMF ( θ | ϵ, κ ) , ψ = ϵ [32]. After using T ≥ 1 (projected) stochastic (sub)-gradient ascent iterations to obtain an estimation of the parameter ˆ ψ T , Monte Carlo samples θ 1 , . . . , θ L i.i.d ∼ σ ˆ ψ T ( θ ) are used to approximate the value of the DSW. Interestingly, the metricity DSW holds for nonoptimal ˆ ψ T as long as σ ˆ ψ T ( θ ) are continuous on S d -1 e.g., vMF with κ < ∞ [31]. In addition, the unbiasedness property of the DSW is the same as the SW, namely, when L → ∞ , the empirical estimation of the DSW converges to the true value. The time complexity and space complexity of the DSW are O ( LTn log n ) and O ( L ( d + n )) in turn without counting the complexities of sampling from σ ˆ ψ T ( θ ) . We refer to Appendix B.1 for more details, e.g., equations, algorithms, and discussion. PAC-Bayesian generalization bounds for DSW are investigated in [35].

Max sliced Wasserstein. By letting the concentration parameter κ → ∞ , the vMF distribution degenerates to the Dirac distribution vMF ( θ | ϵ, κ ) → δ ϵ , we obtain the max sliced Wasserstein distance [7]. The definition of max sliced Wasserstein (Max-SW) distance between two probability measures µ ∈ P p ( R d ) and ν ∈ P p ( R d ) is:

Similar to the DSW, the Max-SW is often computed by using T ≥ 1 iterations of (projected) (sub)-gradient ascent to obtain an estimation of the "max" projecting direction ˆ θ T . After that, the

estimated value of the Max-SW is W p ( ˆ θ T ♯µ, ˆ θ T ♯ν ) . The time complexity and space complexity of the Max-SW are O ( Tn log n ) and O ( d + n ) . It is worth noting that the Max-SW is only a metric at the global optimum θ ⋆ , hence, we cannot guarantee the metricity of the Max-SW due to the non-convex optimization [34] problem even when T →∞ . Therefore, the performance of Max-SW is often unstable in practice [28]. We refer the reader to Appendix B.1 for more details e.g., equations, algorithms, and discussions about the Max-SW.

3 Energy-Based Sliced Wasserstein Distance

From the background, we observe that using a fixed slicing distribution e.g., in the SW, is computationally efficient but might be not effective. In contrast, using optimization-based slicing distributions is computationally expensive e.g., in the DSW, and is unstable e.g., in the Max-SW. Therefore, we address previous issues by introducing a novel sliced Wasserstein variant that uses optimization-free slicing distribution which can highlight the difference between two comparing probability measures.

3.1 Energy-Based Slicing Distribution

We first start with the key contribution which is the energy-based slicing distribution.

Definition 1. For any p ≥ 1 , dimension d ≥ 1 , an energy function f : [0 , ∞ ) → Θ ⊂ (0 , ∞ ) and two probability measures µ ∈ P p ( R d ) and ν ∈ P p ( R d ) , the energy-based slicing distribution σ µ,ν ( θ ) supported on S d -1 is defined as follow:

where the image of f is in the open interval (0 , ∞ ) is for making σ µ,ν ( θ ) continuous on S d -1

In contrast to the approach of the DSW which creates the dependence between the slicing distribution and two input probability measures via optimization, the energy-based slicing distribution obtains the dependence by exploiting the value of the projected Wasserstein distance between two input probability measures at each support.

Monotonically increasing energy functions. Similar to previous works, we again assume that "A higher value of projected Wasserstein distance, a better projecting direction" . Therefore, it is natural to use a monotonically increasing function for the energy function f . We consider the following two functions: the exponential function: f e ( x ) = e x , and the shifted polynomial function: f q ( x ) = x q + ε with q, ε > 0 . The shifted constant ε helps to avoid the slicing distribution undefined when two input measures are equal. In a greater detail, when µ = ν , we have W p p ( θ♯µ, θ♯ν ) = 0 for all θ ∈ S d -1 due to the identity property of the Wasserstein distance. Hence, σ µ,ν ( θ ; f, p ) ∝ 0 for θ ∈ S d -1 and f ( x ) = x q ( q > 0 ). Therefore, the slicing distribution σ µ,ν ( θ ; f, p ) is undefined due to an invalid density function. In practice, it is able to set ε = 0 since we rarely deal with two coinciding measures.

Other energy functions. We can choose any positive function for energy function f and it will result in a valid slicing distribution. However, it is necessary to come up with an assumption for the choice of the function. Since there is no existing other assumption for the importance of projecting direction, we will leave the investigation of non-increasing energy function f to future works.

Example 1. Let µ = N ( m , v 2 1 I ) and ν = N ( m , v 2 2 I )) are two location-scale Gaussian distributions with the same means, we have their projections are θ♯µ = N ( θ ⊤ m , v 2 1 ) and θ♯ν = N ( θ ⊤ m , v 2 2 ) . Based on the closed form of the Wasserstein distance between two Gaussians [10], we have W 2 2 ( θ♯µ, θ♯ν ) = ( v 1 -v 2 ) 2 for all θ ∈ S d -1 which leads to σ µ,ν ( θ ; f, p ) = U ( S d -1 ) for definitions of energy function f in Definition 1.

Example 1 gives a special case where we can have the closed form of the slicing function.

Applications to other sliced probability metrics and mutual information. In this paper, we focus on comparing choices of slicing distribution in the basic form of the SW distance. The proposed energy-based slicing distribution can be adapted to other variants of the SW that are not about designing slicing distribution e.g., non-linear projecting [19], orthogonal projecting directions [38], and so on. Moreover, the energy-based slicing approach can be applied to other sliced probability metrics e.g., sliced score matching [42], and sliced mutual information [13].

3.2 Definitions, Topological, and Statistical Properties of Energy Based Sliced Wasserstein

With the definition of energy-based slicing distribution in Definition 1, we now are able to define the energy-based sliced Wasserstein (EBSW) distance.

Definition 2. For any p ≥ 1 , dimension d ≥ 1 , two probability measures µ ∈ P p ( R d ) and ν ∈ P p ( R d ) , the energy function f : [0 , ∞ ) → (0 , ∞ ) , and the energy-based slicing distribution σ µ,ν ( θ ; f, p ) , the energy-based sliced Wasserstein (EBSW) distance is defined as follows:

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 ≥ 1 , energy-function f , the energy-based sliced Wasserstein EBSW p ( · , · ; f ) is a semi-metric in the probability space on R d , namely EBSW satisfies non-negativity, symmetry, and identity of indiscernibles.

The proof of Theorem 1 in given in Appendix A.1. Next, we establish the connections among the EBSW, the SW, the Max-SW, and the Wasserstein.

Proposition 1. (a) For any p ≥ 1 and increasing energy function f , we find that

The equality holds when f ( x ) = c for some positive constant c for all x ∈ [0 , ∞ ) .

(b) For any p ≥ 1 and energy function f , we have

Proof of Proposition 1 is in Appendix A.2. The results of Proposition 1 indicate that for increasing energy function f , the EBSW is lower bounded by the SW while it is upper bounded by the Max-SW. It is worth noting that the EBSW is the first variant that changes the slicing distribution while still being an upper bound of the SW.

Theorem 2. For any p ≥ 1 and energy function ¯ f , the convergence of probability measures under the energy-based sliced Wasserstein distance EBSW p ( · , · ; ¯ f ) implies weak convergence of probability measures and vice versa.

Theorem 2 implies that for any sequence of probability measures ( µ k ) k ∈ N and µ in P p ( R d ) , lim k → + ∞ EBSW p ( µ k , µ ; ¯ f ) = 0 if and only if for any continuous and bounded function f : R d → R , lim k → + ∞ ∫ f d µ k = ∫ f d µ . The proof of Theorem 2 is in Appendix A.3.

Statistical Properties. From Proposition 1, we derive the sample complexity of the EBSW.

Proposition 2. Let X 1 , X 2 , . . . , X n be i.i.d. samples from the probability measure µ being supported on compact set of R d . We denote the empirical measure µ n = 1 n ∑ n i =1 δ X i . Then, for any p ≥ 1 and energy function f , there exists a universal constant C > 0 such that

where the outer expectation is taken with respect to the data X 1 , X 2 , . . . , X n .

The proof of Proposition 2 is given in Appendix A.4. From this proposition, we can say that the EBSW does not suffer from the curse of dimensionality. We will discuss other statistical properties of approximating the EBSW in the next section.

3.3 Computational Methods and Computational Properties

Calculating the expectation with respect to the slicing distribution σ µ,ν ( θ ; f, p ) is intractable. Therefore, we propose some Monte Carlo estimation methods to approximate the value of EBSW.

3.3.1 Importance Sampling

The most simple and computationally efficient method that can be used is importance sampling (IS) [18]. The idea is to utilize an efficient-sampling proposal distribution σ 0 ( θ ) ∈ P ( S d -1 ) to provide Monte Carlo samples. After that, we use the density ratio between the original slicing distribution and the proposal distribution to weight samples. We can rewrite the EBSW distance as:

where σ 0 ( θ ) ∈ P ( S d -1 ) is the proposal distribution, and:

is the importance weighted function. The detailed derivation is given in Appendix B.2. Let θ 1 , . . . , θ L be i.i.d samples from σ 0 ( θ ) , the importance sampling estimator of the EBSW (IS-EBSW) is:

where ˆ w µ,ν,σ 0 ,f,p ( θ l ) = w µ,ν,σ 0 ,f,p ( θ l ) ∑ L l ′ =1 w µ,ν,σ 0 ,f,p ( θ l ′ ) is the normalized importance weights. When σ 0 ( θ ) = U ( S d -1 ) = Γ( d/ 2) 2 π d/ 2 (a constant of θ ), we can replace w µ,ν,σ 0 ( θ l ) with f ( W p p ( θ l ♯µ, θ l ♯ν )) . When we choose the energy function f ( x ) = e x , computing the normalized importance weights

is equivalent to computing the Softmax function.

Computational algorithms and complexities. The computational algorithm of IS-EBSW can be derived from the algorithm of the SW distance by adding only one to two lines of code for computing the importance weights. For a better comparison, we give algorithms for computing the SW distance and the EBSW distance in Algorithm 1 and Algorithm 4 in Appendix B.1 and Appendix B.2 respectively. When µ and ν are two discrete measures that have at most n supports, the time complexity and the space complexity of the IS-EBSW distance are O ( Ln log n + Lnd ) and O ( L ( n + d )) which are the same as the SW.

Unbiasedness. The IS approximation is asymptotically unbiased for EBSW p p ( µ, ν ; 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 p p ( µ, ν ; f ) is not very vital. Moreover, we can show that ̂ IS-EBSW p p ( µ, ν ; f, L ) is an 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, it is safe to use the IS-EBSW in practice.

Gradient Estimation. In statistical inference, we might want to estimate the gradient ∇ ϕ EBSW p ( µ ϕ , ν ; f ) for doing minimum distance estimator [45]. Therefore, we derive the gradient 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 σ µ,ν ( θ ; f, p ) . For example, when we have θ 1 , . . . , θ L are approximately distributed following σ µ,ν ( θ ; f, p ) , we can take ( 1 L ∑ L l =1 W p p ( θ l ♯µ, θ l ♯ν ) ) 1 p 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-Hasting algorithm: independent Metropolis-Hasting (IMH) and random walk Metropolis-Hasting (RMH).

Sampling Importance Resampling. Similar to importance sampling in Section 3.3.1, the SIR uses a proposal distribution σ 0 ( θ ) to obtain L samples θ ′ 1 , . . . , θ ′ L and the corresponding normalized

importance weights:

After that, the SIR creates the resampling distribution which is a Categorical distribution ˆ q ( θ ) = ∑ L l =1 ˆ w µ,ν,σ 0 ,f,p ( θ l ) δ θ ′ l . Finally, the SIR draws L samples θ 1 , . . . , θ L i.i.d ∼ ˆ q ( θ ) . We denote the SIR estimation of the EBSW distance as SIR-EBSW.

Markov Chain Monte Carlo. MCMCcreates a 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-Hastings algorithm. Let the starting sample follow a prior distribution θ 1 ∼ σ 0 ( θ ) ∈ P ( S d -1 ) , a transition distribution σ t ( θ t | θ t -1 ) ∈ P ( S d -1 ) for any timestep t > 1 is used to sample a candidate θ ′ t . After that, the new sample θ t is set to θ ′ t with the probability α and is set to θ t -1 with the probability 1 -α with

In theory, T should be large enough to help the Markov chain to mix to the stationary distribution and the first M < T samples are often dropped as burn-in samples. However, to keep the computational complexity the same as the previous computational methods, we set T = L and M = 0 . The first choice of transition distribution is σ t ( θ t | θ t -1 ) = U ( S d -1 ) which leads to independent Metropolis-Hasting (IMH). The second choice is σ t ( θ t | θ t -1 ) = vMF ( θ t | θ t -1 , κ ) (the von MisesFisher distribution [17] with location θ t -1 ) which leads to random walk Metropolis-Hasting (RMH). Since both above transition distributions are symmetric σ t ( θ t | θ t -1 ) = σ t ( θ t -1 | θ t ) , the acceptance probability turns into:

which means that the acceptance probability equals 1 and θ ′ t is always accepted as θ t if it can increase the energy function. We refer to the IMH estimation and the RMH estimation of the EBSW distance as IMH-EBSW and RMH-EBSW in turn.

Computational algorithms and complexities. We refer the reader to Algorithm 5, Algorithm 6, and Algorithm 7 in Appendix B.3 for the detailed algorithms of the SIR-EBSW, the IMH-EBSW, and the RMH-EBSW. 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 O ( Ln log n + Lnd ) and O ( 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 p ( µ, ν ; f ) . However, they are also unbiased estimations of the power p of a 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 µ ϕ . In the SIR estimation and the MCMC estimation, the expectation is with respect to the slicing distribution σ µ ϕ ,ν ( θ, f ) that depends on µ ϕ . 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 is to create the slicing distribution from an independent copy of µ ϕ . In particular, we denote µ ϕ ′ is the independent copy of µ ϕ with ϕ ′ equals ϕ in terms of value. Therefore, we can obtain the slicing distribution σ µ ϕ ′ ,ν ( θ ; f ) that does not depend on µ ϕ . This approach still gives the same value of distance, we refer to Appendix B.3 for a more careful discussion.

4 Experiments

In this section, we first visualize the shape of the energy-based slicing distribution in a 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.

0.2

0.4

0.6

0.8

x

1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

x

2

3

2

1

0

1

2

3

0.000000

0.000025

0.000050

0.000075

0.000100

0.000125

0.000150

0.000175

0.000200

Density

f

1(

x

) =

x

,

(

;

f

1)

SIR

,

(

;

f

1)

IMH

,

(

;

f

1)

RMH

,

(

;

f

1)

vMF

(

,

)

3

2

1

0

1

2

3

0.0000

0.0002

0.0004

0.0006

0.0008

0.0010

Density

f

e

(

x

) =

e x

,

(

;

f

e

)

SIR

,

(

;

f

e

)

IMH

,

(

;

f

e

)

RMH

,

(

;

f

e

)

vMF

(

,

)

In experiments, we denote EBSW-e for the exponential energy function i.e., f ( x ) = e x , and EBSW-1 for the identity energy function i.e., f ( x ) = 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-based slicing distribution in two dimensions in Figure 1. In particular, we consider comparing two empirical distributions in the left-most figure (taken from [11]). We utilize the SIR, the IMH, and the RMH to obtain 10 4 Monte Carlo samples from the energy-based slicing distribution. For the IMH and the RHM, we burn in the 10 4 samples. After that, we use the von Mises kernel density estimation to 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 (tuning κ ∈ { 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 f 1 ( x ) = x and the right-most figure is according to the energy function f e ( x ) = e x . 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 µ ( t ) flowing with time t along the gradient flow of a loss functional µ ( t ) → D ( µ ( t ) , ν ) that drives it towards a target distribution ν [40] where D is our sliced Wasserstein variants. We consider discrete flows, namely. we set ν = 1 n ∑ n i =1 δ Y i as a fixed empirical target distribution and the model distribution µ ( t ) = 1 n ∑ n i =1 δ X i ( t ) . Here, the model distribution is parameterized by a time-varying point cloud X ( t ) = ( X i ( t )) n i =1 ∈ ( R d ) n . Starting from an initial condition at time t = 0 , we integrate the ordinary differential equation ˙ X ( t ) = -n ∇ X ( t ) [ D ( 1 n ∑ n i =1 δ X i ( t ) , ν )] for each iteration. We choose µ (0) and ν 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 µ ( t ) and ν ( t ∈ { 0 , 100 , 200 , 300 , 400 , 500 } ) and the computational time of the SW variants in Table 1. Here, we set L = 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 . 1 } . For the v-DSW, we report the best result for ( L, T ) ∈ { (10 , 10) , (50 , 2) , (2 , 50) } , κ ∈ { 1 , 10 , 50 } , and the learning rate for finding the location in { 0 . 001 , 0 . 01 , 0 . 1 } . We observe that the IS-EBSW-e helps to drive the flow to converge faster than baselines. More importantly, the computational time of the IS-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 3 in Appendix C.1. 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 ). Despite having comparable performance, the SIR-EBSW-e, the IMH-EBSW-e, and the RMH-EBSW-e

( κ = 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 4 for comparing gradient estimators of the EBSW.

Qualitative Results. We show the point-cloud flows of the SW, the Max-SW, the v-DSW, and the IS-EBSW-e, in Figure 2. The flows from the SIR-EBSW-e, the IMG-EBSW-e, and the RMH-EBSW-e are added in Figure 4 in the Appendix C.1. 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 a gradient flow that starts from the empirical distribution over the normalized color palette (RGB) of the source image to the empirical distribution over the normalized color palette (RGB) of the target image. Since the value of the color palette is 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 as in the previous section. We show both transferred images, corresponding computational times, and Wasserstein-2 distances in Figure 3. We observe the same phenomenon as in the previous section, namely, the 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 is only slightly higher. Moreover, the transferred image from the 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-e, the IMH-EBSW-e, the RMH-EBSW-e, the results for the identity energy function, the results for changing hyperparamters ( L, T ) , and the results for comparing gradient estimators. Overall, we observe similar phenomenons as in the previous gradient flow section.

4.4 Deep Point-Cloud Reconstruction

We follow [33] to train point-cloud autoencoders with sliced Wasserstein distances on the ShapeNet Core-55 dataset [4]. In short, we aim to estimate an autoencoder that contains an encoder f ϕ that maps a point cloud X ∈ R nd to a latent code z ∈ R h , and a decoder g ψ that maps the latent

Source

SW 85.04(s),

W

2 =412.76

Max-SW 126.06(s),

W

2 =460.82

v-DSW 92.41(s),

W

2 =340.92

IS-EBSW-e 86.7(s),

W

2 =264.95

Target

code z to a reconstructed point cloud ˜ X ∈ R nd . We want to have the pair f ϕ and g ψ such that ˜ X = g ψ ( f ϕ ( X )) ≈ X for all X ∼ p ( X ) which is our data distribution. To do that, we solve the following optimization problem: min ϕ,γ E X ∼ µ ( X ) [ S ( P X , P g γ ( f ϕ ( X ))) ] , where S is a sliced Wasserstein variant, and P X 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 an SGD optimizer with a learning rate of 1e-3, a batch size of 128, a momentum of 0.9, and a weight decay of 5e-4. We give more detail in Appendix C.3

Quantitative Results. We evaluate the trained autoencoders on a 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 are the average of corresponding distances on all point-clouds. From the table, we observe that the IS-EBSW-e can help to train the autoencoder faster in terms of the Wasserstein distance and the sliced Wasserstein distance. We refer to Table 5 in Appendix C.3 for similar ablation studies as in the previous sections including the results for the SIR-EBSW-e, the IMH-EBSW-e, the RMH-EBSW-e, the results for the identity energy function, the results for changing hyperparameters ( L, T ) , and the results for comparing gradient estimators.

Qualitative Results. We show some ground-truth point-clouds ModelNet40 and their corresponding reconstructed point-clouds from different models ( L = 100 ) at epochs 200 and 20 in Figure 6- 7 respectively. 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 in 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-based 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-based 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 at a 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.

References

C. Bonet, N. Courty, F. Septier, and L. Drumetz. Efficient gradient flows in sliced-Wasserstein space. Transactions on Machine Learning Research , 2022.

N. Bonneel, J. Rabin, G. Peyré, and H. Pfister. Sliced and Radon Wasserstein barycenters of measures. Journal of Mathematical Imaging and Vision , 1(51):22-45, 2015.

N. Bonnotte. Unidimensional and evolution methods for optimal transportation . PhD thesis, Paris 11, 2013.

A. X. Chang, T. Funkhouser, L. Guibas, P. Hanrahan, Q. Huang, Z. Li, S. Savarese, M. Savva, S. Song, H. Su, et al. Shapenet: An information-rich 3d model repository. arXiv preprint arXiv:1512.03012 , 2015.

M. Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems , pages 2292-2300, 2013.

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.

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.

I. Deshpande, Z. 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.

L. Devroye, L. Györfi, and G. Lugosi. A probabilistic theory of pattern recognition , volume 31. Springer Science & Business Media, 2013.

D. Dowson and B. Landau. The fréchet distance between multivariate normal distributions. Journal of multivariate analysis , 12(3):450-455, 1982.

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 Statistics , pages 2681-2690, 2019.

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.

Z. Goldfeld and K. Greenewald. Sliced mutual information: A scalable measure of statistical dependence. Advances in Neural Information Processing Systems , 34, 2021.

N. J. Gordon, D. J. Salmond, and A. F. Smith. Novel approach to nonlinear/non-gaussian bayesian state estimation. In IEE proceedings F (radar and signal processing) , volume 140, pages 107-113. IET, 1993.

M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter. GANs trained by a two time-scale update rule converge to a local Nash equilibrium. In Advances in Neural Information Processing Systems , pages 6626-6637, 2017.

G. E. Hinton. Training products of experts by minimizing contrastive divergence. Neural computation , 14(8):1771-1800, 2002.

P. E. Jupp and K. V. Mardia. Maximum likelihood estimators for the matrix von Mises-Fisher and bingham distributions. The Annals of Statistics , 7(3):599-606, 1979.

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.

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.

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 Pattern Recognition , pages 3427-3436, 2018.

A. Krizhevsky, G. Hinton, et al. Learning multiple layers of features from tiny images. Master's thesis, Department of Computer Science, University of Toronto , 2009.

C.-Y. Lee, T. Batra, M. H. Baig, and D. Ulbricht. Sliced Wasserstein discrepancy for unsupervised domain adaptation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10285-10295, 2019.

Z. Liu, P. Luo, X. Wang, and X. Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV) , December 2015.

A. Liutkus, U. 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.

K. P. Murphy. Probabilistic Machine Learning: Advanced Topics . MIT Press, 2023.

K. Nadjahi, V. De Bortoli, A. Durmus, R. Badeau, and U. ¸ Sim¸ sekli. Approximate Bayesian computation with the sliced-Wasserstein distance. In ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) , pages 5470-5474. IEEE, 2020.

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-20812, 2020.

K. Nguyen and N. Ho. Amortized projection optimization for sliced Wasserstein generative models. Advances in Neural Information Processing Systems , 2022.

K. Nguyen and N. Ho. Control variate sliced wasserstein estimators. arXiv preprint arXiv:2305.00402 , 2023.

K. Nguyen, N. Ho, T. Pham, and H. Bui. Distributional sliced-Wasserstein and applications to generative modeling. In International Conference on Learning Representations , 2021.

K. Nguyen, D. Nguyen, and N. Ho. Self-attention amortized distributional projection optimization for sliced Wasserstein point-cloud reconstruction. Proceedings of the 40th International Conference on Machine Learning , 2023.

K. Nguyen, S. Nguyen, N. Ho, T. Pham, and H. Bui. Improving relational regularized autoencoders with spherical sliced fused Gromov-Wasserstein. In International Conference on Learning Representations , 2021.

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.

S. Nietert, R. Sadhu, Z. Goldfeld, and K. Kato. Statistical, robustness, and computational guarantees for sliced wasserstein distances. Advances in Neural Information Processing Systems , 2022.

R. Ohana, K. Nadjahi, A. Rakotomamonjy, and L. Ralaivola. Shedding a pac-bayesian light on adaptive sliced-wasserstein distances. In International Conference on Machine Learning , pages 26451-26473. PMLR, 2023.

G. Peyré and M. Cuturi. Computational optimal transport, 2020.

C. R. Qi, H. Su, K. Mo, and L. J. Guibas. Pointnet: Deep learning on point sets for 3d classification and segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 652-660, 2017.

M. Rowland, J. 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.

T. Salimans, I. Goodfellow, W. Zaremba, V. Cheung, A. Radford, and X. Chen. Improved techniques for training GANs. Advances in Neural Information Processing Systems , 29, 2016.

F. Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY , 55(58-63):94, 2015.

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.

Y. Song, S. Garg, J. Shi, and S. Ermon. Sliced score matching: A scalable approach to density and score estimation. In Uncertainty in Artificial Intelligence , pages 574-584. PMLR, 2020.

C. Villani. Optimal transport: old and new , volume 338. Springer Science & Business Media, 2008.

M. J. Wainwright. High-dimensional statistics: A non-asymptotic viewpoint . Cambridge University Press, 2019.

J. Wolfowitz. The minimum distance method. The Annals of Mathematical Statistics , pages 75-88, 1957.

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.

M. Yi and S. Liu. Sliced Wasserstein variational inference. In Fourth Symposium on Advances in Approximate Bayesian Inference , 2021.

Supplement to 'Energy-Based 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 1

Non-negativity and Symmetry. the non-negativity and symmetry properties of the EBSW follow directly by the non-negativity and symmetry of the Wasserstein distance since it is an expectation of the one-dimensional Wasserstein distance.

Identity. Weneed to show that EBSW p ( µ, ν ; f ) = 0 if and only if µ = ν . First, from the definition of EBSW, we obtain directly µ = ν implies EBSW p ( µ, ν ; f ) = 0 . For the reverse direction, we use the same proof technique in [3]. If EBSW p ( µ, ν ; f ) = 0 , we have ∫ S d -1 W p ( θ♯µ, θ♯ν ) d σ µ,ν ( θ ; f, p ) = 0 . Hence, we have W p ( θ♯µ, θ♯ν ) = 0 for σ µ,ν ( θ ; f, p ) -almost surely θ ∈ S d -1 . Since σ µ,ν ( θ ; f, p ) is continuous, we have W p ( θ♯µ, θ♯ν ) = 0 for all θ ∈ S d -1 . From the identity property of the Wasserstein distance, we obtain θ♯µ = θ♯ν for σ µ,ν ( θ ; f, p ) -a.e θ ∈ S d -1 . Therefore, for any t ∈ R and θ ∈ S d -1 , we have:

where F [ γ ]( w ) = ∫ R d ′ e -i ⟨ w,x ⟩ dγ ( x ) denotes the Fourier transform of γ ∈ P ( R d ′ ) . By the injectivity of the Fourier transform, we obtain µ = ν which concludes the proof.

A.2 Proof of Proposition 1

(a) We first provide the proof for the inequality SW p ( µ, ν ) ≤ EBSW p ( µ, ν ; f ) . It is equivalent to prove that

From the law of large number, it is sufficient to demonstrate that

for any L ≥ 1 and θ 1 , . . . , θ L i.i.d. ∼ U ( S d -1 ) . To ease the presentation, we denote a i = W p p ( θ i ♯µ, θ i ♯ν ) and b i = f ( W p p ( θ i ♯µ, θ i ♯ν )) for all 1 ≤ i ≤ L . The inequality (9) becomes:

We prove the inequality (10) via an induction argument. It is clear that this inequality holds when L = 1 . We assume that this inequality holds for any L . We now verify that the inequality (10) also holds for L +1 . Without loss of generality, we assume that a 1 ≤ a 2 ≤ . . . ≤ a L ≤ a L +1 . Since the function f is an increasing function, it indicates that b 1 ≤ b 2 ≤ . . . ≤ b L ≤ b L +1 . Applying the induction hypothesis for a 1 , . . . , a L and b 1 , . . . , b L , we find that

This inequality leads to

Therefore, to obtain the conclusion of the hypothesis for L +1 , it is sufficient to demonstrate that

which is equivalent to show that

Since a L +1 ≥ a i and b L +1 ≥ b i for all 1 ≤ i ≤ L , we have ( a L +1 -a i )( b L +1 -b i ) ≥ 0 , which is equivalent to a L +1 b L +1 + a i b i ≥ a L +1 b i + b L +1 a i for all 1 ≤ i ≤ L . By taking the sum of these inequalities over i from 1 to L , we obtain the conclusion of inequality (11). Therefore, we obtain the conclusion of the induction argument for L +1 , which indicates that inequality (10) holds for all L . As a consequence, we obtain the inequality SW p ( µ, ν ) ≤ EBSW p ( µ, ν ; f ) .

(b) We recall the definition of the Max-SW:

Since S d -1 is compact and the function θ → W p ( θ♯µ, θ♯ν ) is continuous, we have θ ⋆ = argmax θ ∈ S d -1 W p ( θ♯µ, θ♯ν ) . From Definition 2, for any p ≥ 1 , dimension d ≥ 1 , energy-function f , and µ, ν ∈ P p ( R d ) we have:

Furthermore, by applying the Cauchy-Schwartz inequality, we have:

after taking the p -rooth, we completes the proof.

A.3 Proof of Theorem 2

We aim to show that for any sequence of probability measures ( µ k ) k ∈ N and µ in P p ( R d ) , lim k → + ∞ EBSW p ( µ k , µ ; ¯ f ) = 0 if and only if for any continuous and bounded function f : R d → R , lim k → + ∞ ∫ f d µ k = ∫ f d µ .

̸

We now show that lim k →∞ EBSW p ( µ k , µ ; ¯ f ) = 0 implies ( µ k ) k ∈ N converges weakly to µ . Let ( µ k ) k ∈ N be a sequence such that lim k →∞ EBSW p ( µ k , µ ; ¯ f ) = 0 , we suppose ( µ k ) k ∈ N does not converge weakly to µ . So, let D P be the Lévy-Prokhorov metric, lim k →∞ D P ( µ k ,µ ) = 0 that implies

there exists ε > 0 and a subsequence ( µ ψ ( k ) ) k ∈ N with an increasing function ψ : N → N such that for any k ∈ N : D P ( µ ψ ( k ) , µ ) ≥ ε . Using the Holder inequality with µ, ν ∈ P p ( R d ) , we have:

̸

Therefore, lim k →∞ SW 1 ( µ ψ ( k ) , µ ) = 0 which implies that there exists a subsequence ( µ ϕ ( ψ ( k )) ) k ∈ N with an increasing function ϕ : N → N such that ( µ ϕ ( ψ ( k )) ) k ∈ N converges weakly to µ by Lemma S1 in [27]. Therefore a contradiction appears from the assumption of lim k →∞ D P ( µ ϕ ( ψ ( k )) , µ ) = 0 . Therefore, lim k →∞ EBSW p ( µ k , µ ; ¯ f ) = 0 , ( µ k ) k ∈ N converges weakly to µ .

When ( µ k ) k ∈ N converges weakly to µ , we have ( θ♯µ k ) k ∈ N converges weakly to θ♯µ for any θ ∈ S d -1 by the continuous mapping theorem. From [43], the weak convergence implies the convergence under the Wasserstein distance. So, we have lim k →∞ W p ( θ♯µ k , µ ) = 0 . Moreover, using the fact that the Wasserstein distance is also bounded, hence, the bounded convergence theorem implies:

Again, using the continuous mapping theorem with function x → x 1 /p , we have lim k →∞ EBSW p ( µ k , µ ; ¯ f ) → 0 . We conclude the proof.

A.4 Proof of Proposition 2

We first show that the following inequality holds

where C > 0 is some universal constant and the outer expectation is taken with respect to the random variables X 1 , . . . , X n . We now follow the proof technique from in [30]. Let F -1 n,θ and F -1 θ be the inverse cumulative distributions of two push-forward measures θ♯µ n and θ♯µ . Using the closed-form expression of the Wasserstein distance in one dimension, we obtain the following equations and inequalities:

where X ⊂ R d is the compact set of the probability measure µ . We can check that

where H is the set of half-spaces { z ∈ R d : θ ⊤ z ≤ x } for all θ ∈ R d such that ∥ θ ∥ ≤ 1 . From VC inequality (Theorem 12.5 in [9]), we have:

with S ( H , n ) is the growth function which is upper bounded by ( n +1) V C ( H ) due to the Sauer Lemma (Proposition 4.18 in "High-Dimensional Statistics: A Non-Asymptotic Viewpoint", Wainwright et al). From Example 4.21 in [44], we get V C ( H ) = d +1 .

Using the Jensen inequality and the tail sum expectation for non-negative random variable, we have:

Since the inequality holds for any u , we optimize for u . Let f ( u ) = u +256 S ( H , n ) e -nu/ 32 n , we have f ′ ( u ) = 1 + 8 S ( H , n ) e -nu/ 32 . Setting f ′ ( u ) = 0 , we obtain the minima u ⋆ = 32 log(8 S ( H ,n )) n . Therefore, we have:

by using Sauer Lemma. As a consequence, we obtain:

which completes the proof.

B Additional Materials

B.1 Additional Background

Sliced Wasserstein. When two probability measures are empirical probability measures on n supports: µ = 1 n ∑ n i =1 δ x i and ν = 1 n ∑ n i =1 δ y i , the SW distance can be computed by sorting projected supports. In particular, we have θ♯µ = 1 n ∑ n i =1 δ θ ⊤ x i , θ♯ν = 1 n ∑ n i =1 δ θ ⊤ y i , and W p p ( θ♯µ, θ♯ν ) = 1 n ∑ n i =1 ( θ ⊤ x ( i ) -θ ⊤ y ( i ) ) p where θ ⊤ x ( i ) is the ordered projected supports. We provide the pseudo-code for computing the SW in Algorithm 1.

Max sliced Wasserstein. The Max-SW is 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 η . Moreover, the empirical estimation of the Max-SW might not converge to the Max-SW when T →∞ .

Distributional sliced Wasserstein. To solve the optimization of the DSW, we need to use the stochastic (sub)-gradient ascent algorithm. In particular, we first need to estimate the gradient:

Algorithm 1 Computational algorithm of the SW distance

Input: Probability measures µ and ν , p ≥ 1 , and the number of projections L .

end for

Return:

̂ SW p ( µ, ν ; L )

Algorithm 2 Computational algorithm of the Max-SW distance

Input: Probability measures µ and ν , p ≥ 1 , the number of iterations T , and the step size η . Sample ˆ θ 0 ∼ U ( S d -1 )

end for

Return:

̂ Max-SW p ( µ, ν ; T )

Algorithm 3 Computational algorithm of the DSW distance

Input: Probability measures µ and ν , p ≥ 1 , the number of projections L , the number of iterations T , and the step size η .

Initialize ˆ ψ 0

for l = 1 to L do

Sample θ l,ψ ∼ σ ˆ ψ t -1 ( θ ) via reparameterization.

end for

end for

Sample θ l ∼ σ ˆ ψ T ( θ ) via reparameterization.

end for

To estimate the gradient ∇ ψ E θ ∼ σ ψ ( θ ) W p p ( θ♯µ, θ♯ν ) , we need to use reparameterization trick for σ ψ ( θ ) e.g., the vMF distribution. After using the reparameterization trick, we can approximate the gradient ∇ ψ E θ ∼ σ ψ ( θ ) W p p ( θ♯µ, θ♯ν ) = 1 L ∑ l l =1 ∇ ψ W p p ( θ l,ψ ♯µ, θ l,ψ ♯ν ) where θ 1 ,ψ , . . . , θ L,ψ are i.i.d reparameterized samples from σ ψ ( θ ) . Similarly, we approximiate E θ ∼ σ ψ ( θ ) W p p ( θ♯µ, θ♯ν ) = 1 L ∑ L l =1 W p p ( θ l ♯µ, θ l ♯ν ) . We refer to the details in the following papers [6, 32]. We review the algorithm for 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 η .

Minimum Distance Estimator and Gradient Estimation. In statistical inference, we are given the empirical samples X 1 , . . . , X n from the interested distribution ν . Since we do not know the form of ν , we might want to find an alternative representation. In particular, we want to find the best member µ ϕ in a family of distribution parameterized by ϕ ∈ Φ . To do that, we want to minimize the

Input: Probability measures µ and ν , p ≥ 1 , the number of projections L , the energy function f . for l = 1 to L do

Sample

θ

Compute

Compute

end for

l

∼ U

v

l

=

w

l

=

(

S

W

d

p

f

(

-

1

(

θ

)

l

♯µ, θ

W

p

(

θ

l

l

)

♯µ, θ

l

))

distance between µ ϕ and the empirical distribution ν n = 1 n ∑ n i =1 δ X i . This framework is named the minimum distance estimator [45]:

min ϕ ∈ Φ D ( µ ϕ , ν n ) , where D is a discrepancy between two distributions. The gradient ascent algorithm is often used to solve the problem. To do so, we need to compute the gradient ∇ ϕ D ( µ ϕ , ν n ) . When using sliced Wasserstein distances, the gradient ∇ ϕ D ( µ ϕ , ν n ) is often approximated by a stochastic gradient since the SW distances involve an intractable expectation. In previous SW variants, the expectation does not depend on ϕ , 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 ∇ ϕ E θ ∼ σ ( θ ) [ W p p ( θ♯µ, θ♯ν )] = E θ ∼ σ ( θ ) [ ∇ ϕ W p p ( θ♯µ, θ♯ν )] ≈ 1 L ∑ L l =1 ∇ ϕ W p p ( θ l ♯µ, θ l ♯ν ) for θ 1 , . . . , θ L i.i.d ∼ σ ( θ ) .

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:

Algorithms. We provide the algorithm for the IS estimation of the EBSW in Algorithm 4. Compared to the algorithm of the SW in Algorithm 1, 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 µ ϕ be parameterized by ϕ , we derive now the gradient estimator ∇ ϕ EBSW p ( µ, ν ; f ) through importance sampling. We have:

We denote A ( ϕ ) = E θ ∼ σ 0 ( θ ) [ W p p ( θ♯µ ϕ , θ♯ν ) w µ ϕ ,ν,σ 0 ,f ( θ ) ] , B ( ϕ ) = E θ ∼ σ 0 ( θ ) [ w µ ϕ ,ν,σ 0 ,f ( θ ) ] , we

Using Monte Carlo samples θ 1 , . . . , θ L ∼ σ 0 ( θ ) after using the Lebnitz rule to exchange the differentiation and the expectation, we obtain:

which yields the gradient estimation. If we construct the slicing distribution by using a copy of µ ϕ i.e., µ ϕ ′ with ϕ ′ = ϕ in terms of value, the gradient estimator can be derived by:

Using Monte Carlo samples θ 1 , . . . , θ L ∼ σ 0 ( θ ) after using the Lebnitz rule to exchange the differentiation and the expectation, we obtain:

It is worth noting that using a copy of µ ϕ does not change the value of the distance. This trick will show its true benefit when dealing with the SIR, and the MCMC methods. However, we still discuss it in the IS case for completeness. We refer to the "copy" trick is the "parameter-copy" gradient estimator while the original one is the conventional estimator.

Importance Weighted sliced Wasserstein distance. Although the IS estimation of the EBSW is not an unbiased estimation for finite L , it is an unbiased estimation of a valid distance on the space of probability measures. We refer to the distance as the importance weighted sliced Wasserstein distance (IWSW) which has the following definition.

Definition 3. For any p ≥ 1 , dimension d ≥ 1 , energy function f , a continuous proposal distribution σ 0 ( θ ) ∼ P ( S d -1 ) and two probability measures µ ∈ P p ( R d ) and ν ∈ R d , the importance weighted sliced Wasserstein (IWSW) distance is defined as follows:

where the expectation is with respect to θ 1 , . . . , θ L i.i.d ∼ σ 0 ( θ ) , and w µ,ν,σ 0 ,f,p ( θ ) = f ( W p p ( θ♯µ,θ♯ν )) σ 0 ( θ ) .

Algorithm 5 Computational algorithm of the SIR-EBSW distance

Input: Probability measures µ and ν , p ≥ 1 , the number of projections L , the energy function f . for l = 1 to L do

Compute

w

l

=

f

(

W

p

(

end for

end for

end for

Algorithm 6 Computational algorithm of the SW distance and the IMH-EBSW distance

The IWSW is semi-metric, it also does not suffer from the curse of dimensionality, and it induces weak convergence. The proofs can be derived by following directly the proofs of the EBSW in Appendix A.1, 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 the RMHin Algorithm 5-7.

Gradient estimators. We derive the reinforce gradient estimator of the EBSW for the SIR, the IMH, and the RHM sampling.

We have:

θ

l

♯µ, θ

l

))

Algorithm 7 Computational algorithm of the SW distance and the RMH-EBSW distance

Input: Probability measures µ and ν , p ≥ 1 , the number of projections L , the energy function f , the concentration parameter κ .

Sample

θ

Compute

1

∼ U

v

1

for

l

= 2

=

to

Sample

θ

Compute

Sample

(

S

W

L

d

p

do

′

l

vMF

(

∼

α

u

if

α

u

Set

-

(

1

)

θ

1

♯µ, θ

θ

= min

1

l

1

-

, κ

(

1

∼ U

([0

,

,

1])

then

θ

else if

Set end if

v

l

=

W

l

=

θ

′

l

α < u

θ

l

=

θ

then

l

p

(

θ

l

-

1

♯µ, θ

l

)

Return:

and

and

Using L Monte Carlo samples from the SIR (or the IMH or the RMH) to approximate the expectation E θ ∼ σ µ ϕ ,ν ( θ ; f ) , and L samples from U ( S d -1 ) to approximate the expectation E θ ∼U ( S d -1 ) , we obtain the gradient estimator of the EBSW. However, the reinforce gradient estimator is unstable in practice, especially with the energy function f e ( x ) = e x . Therefore, we propose a more simple gradient estimator which is:

The key is to use a copy of the parameter ϕ ′ for constructing the slicing distribution σ µ ϕ ′ ,ν ( θ ; f ) , hence, we can exchange directly the differentiation and the expectation. It is worth noting that using the copy also affects the gradient estimation, it does not change the value of the distance. We refer to the "copy" trick is the "parameter-copy" gradient estimator while the original one is the conventional estimator.

̂

RMH-EBSW

p

(

µ, ν

;

L

)

)

)

f

f

(

(

W

W

p

p

p

p

(

θ

(

θ

-

l

′

l

1

♯µ,θ

♯µ,θ

′

l

l

)))

-

1

)))

)

Population distance. The approximated values of p -power EBSW from using the SIR, the IMH, and the RMH can be all written as 1 L ∑ L l =1 W p p ( θ l ♯µ, θ l ♯ν ) . Here, the distributions of θ 1 , . . . , θ L are different. Therefore, they are not an unbiased estimation of the EBSW p p ( µ, ν ; f ) . However, the population distance of the estimation can be defined as in Definition 4.

Definition 4. For any p ≥ 1 , dimension d ≥ 1 , energy function f , and two probability measures µ ∈ P p ( R d ) and ν ∈ R d , the projected sliced Wasserstein (PSW) distance is defined as follows:

where the expectation is with respect to ( θ 1 , . . . , θ L ) ∼ σ ( θ 1 , . . . , θ L ) which is a distribution defined by the SIR (the IMH or the RHM).

The PSW is a valid metric since it satisfies the triangle inequality in addition to the symmetry, the non-negativity, and the identity. In particular, given three probability measures µ 1 , µ 2 , µ 3 ∈ P p ( R d ) we have:

where the first inequality is due to the triangle inequality of Wasserstein distance and the second inequality is due to the Minkowski inequality. The PSW also does not suffer from the curse of dimensionality, and it induces weak convergence. The proofs can be derived by following directly the proofs of the EBSW in Appendix A.1, Apendix A.3, and Appendix A.4. Therefore, using the SIR, the IMH, and the RMH estimation of the EBSWs are as safe as the SW.

C Additional Experiments

In this section, we provide additional results for point-cloud gradient flows in Appendix C.1, color transfer in Appendix C.2, and deep point-cloud reconstruction in Appendix C.3.

C.1 Point-Cloud Gradient Flows

We provide the full experimental results including the IS-EBSW, the SIR-EBSW, the IMH-EBSW, and the RMH-EBSW with both the exponential energy function and the identity energy function in Table 3. In the table, we also include the results for the number of projections L = 10 . In Table 3, we use the conventional gradient estimator for the IS-EBSW while the "parameter-copy" estimator is used for other variants of the EBSW. Therefore, we also provide the ablation studies comparing the gradient estimators in Table 4 by adding the results for the "parameter-copy" estimator for the IS-EBSW and the conventional estimator for other variants. Experimental settings are the same as in the main text.

Quantitative Results. From the two tables, we observe that the IS-EBSW is the best variant of the EBSW in both performance and computational time. Also, we observe that the exponential energy function is better than the identity energy function in this application. It is worth noting that the EBSW variants of all computational methods and energy functions are better than the baselines in

terms of Wasserstein-2 distances at the last epoch. For all sliced Wasserstein variants, we see that reducing the number of projections leads to worsening performance which is consistent with previous studies in previous works [28, 19]. In Table 3, the IS-EBSW uses the conventional gradient estimator while the SIR-EBSW, the IMH-EBSW, and the RMH-EBSW use the "parameter-copy" estimator. Therefore, we report the IS-EBSW with the "parameter-copy" estimator and the SIR-EBSW, the IMH-EBSW, and the RMH-EBSW with the Reinforce estimator (conventional estimator) in Table 4. From the table, we observe the "parameter-copy" estimator is worse than the conventional estimator in the case of IS-EBSW. For the SIR-EBSW, the IMH-EBSW, and the RMH-EBSW, we cannot use the exponential energy function due to the numerically unstable Reinforce estimator. In the case of the identity energy function, the exponential energy function is also worse than the "parameter-copy" estimator. Therefore, we recommend to use the IS-EBSW-e with the conventional gradient estimator.

Qualitative Results. We provide the visualization of the gradient flows from SW (L=100), MaxSW (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 is consistent. Therefore, the energy-based slicing function helps to improve the convergence of the source point-cloud to the target point-cloud.

C.2 Color Transfer

Similar to the point-cloud gradient flow, we follow the same experimental settings of color transfer in the main text. We provide the full experimental results including the IS-EBSW, the SIR-EBSW, the IMH-EBSW, and the RMH-EBSW with both the exponential energy function and the identity energy function, with both L = 10 and L = 100 , and with both gradient estimators in Figure 5.

Results. From the figure, we observe that IMH-EBSW-e gives the best Wasserstein-2 distance among all EBSW variants. Between the exponential energy function and the identity energy function,

n

n

SIR-EBSW-e 328.69(s),

W

2 =277.2

IMH-EBSW-e 287.53(s),

W

2 =256.77

RMH-EBSW-e 328.69(s),

W

2 =277.2

IS-EBSW-1 86.09(s),

W

2 =374.84

IS-EBSW-e (c) 86.3(s),

W

2 =270.59

IS-EBSW-1 (c) 86.19(s),

W

2 =306.21

SIR-EBSW-1 327.74(s),

W

2 =302.79

IMH-EBSW-1 269.62(s),

W

2 =309.08

RMH-EBSW-1 327.74(s),

W

2 =302.79

SIR-EBSW-1 (c) 252.59(s),

W

2 =374.82

IMH-EBSW-1 (c) 355.21(s),

W

2 =381.76

RMH-EBSW-1 (c) 427.33(s),

W

2 =368.04

Source

SW 8.95(s),

W

2 =422.73

Max-SW 13.1(s),

W

2 =463.5

v-DSW 11.61(s),

W

2 =347.38

IS-EBSW-e 9.09(s),

W

2 =266.62

Target

we see that the exponential energy function yields a better result for all EBSW variants. Similar to the gradient flow, reducing the number of projections to 10 also leads to worse results for all sliced Wasserstein variants For the gradient estimators, the conventional estimator is preferred for the IS-EBSW while the "parameter-copy" estimator is preferred for other EBSW variants.

C.3 Deep Point-cloud Reconstruction

We follow the same experimental settings as in the main text. We provide the full experimental results including the IS-EBSW, the SIR-EBSW, the IMH-EBSW, and the RMH-EBSW with both the exponential energy function and the identity energy function, with both L = 10 and L = 100 in Table 5. In Table 5, we use the conventional gradient estimator for the IS-EBSW while other variants of EBSW use the "parameter-copy" gradient estimator. We also compare gradient estimators for the EBSW by adding the results for the "parameter-copy" gradient estimator for the IS-EBSW (denoted as (c)), and the conventional gradient estimator for the SIR-EBSW, the IMH-EBSW, and the RMH-EBSW in Table 6.

Quantitative Results. From the two tables, we observe that the IS-EBSW-e performs the best for both settings of the number of projections L = 10 and L = 100 in terms of the Wasserstein-2 reconstruction errors. For the SW reconstruction error, it is only slightly worse than the SIR-EBSW-e at epoch 200. Comparing the exponential energy function and the identity energy function, we observe that the exponential function is better in both settings of the number of projections. For the same number of projections, the EBSW variants with both types of energy function give lower errors than the baseline including the SW, the Max-SW, and the v-DSW. For all sliced Wasserstein variants, a higher value of the number of projections gives better results. For the gradient estimator of the EBSW, we see that the conventional gradient estimator is preferred for the IS-EBSW while the "parameter-copy" estimator is preferred for other EBSW variants.

Qualitative Results. We show some ground-truth point-clouds ModelNet40 and their corresponding reconstructed point-clouds from different models ( L = 100 ) at epochs 200 and 20 in Figure 6- 7 respectively. From the top to the bottom is the ground truth, the SW, the Max-SW, the v-DSW, the IS-EBSW-e, the SIR-EBSW-e, the IMH-EBSW-e, and the RMH-EBSW-e.

C.4 Deep Generative Modeling

We follow the same setup as sliced Wasserstein deep generative models in [29]. We compare ISEBSW-e with SW, and DSW [30] on CIFAR10 (image size 32x32) [21], and CelebA [23]. We show both FID score [15] (CIFAR10, CelebA) and IS score [39] (CIFAR10) in Figure 8. Overall, IS-EBSW-e performs better than SW and DSW in this task which suggests that EBSW has a wide range of applications.

D Computational Infrastructure

For the point-cloud gradient flows and the color transfer, we use a Macbook Pro M1 for conducting experiments. For deep point-cloud reconstruction, experiments are run on a single NVIDIA V100 GPU.

Peer reviews:
comments: 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.

STRENGTHS:
- The study of the idea is organized and clear. 
- Theoretical quantities of interest are derived.
- 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. 
- I think it would be good to have a theorem studying the following sample complexity, that is more of interest for practical reasons:
$ \mathbb{E}[EBSW_p(\mu_n,\nu_n;f) - EBSW_p(\mu,\nu;f)]$
- 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.

QUESTIONS:
- Can the authors explain why experimental results are not compared to the DSW distance? 
- 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-03

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\nI 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\nWEAKNESSES:\nThe 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\n\nQUESTIONS:\nIn 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.\n\nLIMITATIONS:\nSome limitations have been discussed, but I feel there is scope to discuss more limitations of the proposed approach.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: SUMMARY:
This paper has proposed an energy-based slicing distribution that maps original distributions into one-dimesion space  to compute the Wasserstein distance.

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.
2. Put a title for each row in Figure 2 could help reader easier to compare different SW methods.

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?

2. 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?

3. Could the proposed distance measurement be used in generative models, like GANs and EBMs?

4. Is there any comparison with other probability distance measurement methods, like KL divergence, in the experiments?

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

comments: 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).

STRENGTHS:
The paper has the following strengths

- 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.
- 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:

- 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.
- 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.

QUESTIONS:
Please see comments in Weaknesses

LIMITATIONS:
The paper does not include a Limitation discussion.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-10