Energy-Based Sliced Wasserstein (EBSW)

• Energy-Based Sliced Wasserstein (EBSW) is a robust metric that leverages an energy function to adaptively weight projection directions for comparing probability measures.
• It employs Monte Carlo methods like importance sampling, SIR, and MCMC to focus on the most discriminative slices while maintaining computational efficiency.
• Empirical evaluations demonstrate that EBSW improves convergence, reconstruction quality, and sample efficiency in applications such as gradient flows and deep point-cloud reconstruction.

The Energy-Based Sliced Wasserstein (EBSW) distance is a statistically robust and computationally efficient metric for comparing probability measures, extending the classical Sliced Wasserstein (SW) approach by leveraging an adaptive, parameter-free slicing distribution based on an energy function of the projected Wasserstein distance. By up-weighting projection directions that are most discriminative between measures, EBSW emphasizes informative slices and mitigates the limitations of both uniform and parametric slicing distributions. This formulation preserves the computational advantages of SW while improving statistical signal and sample efficiency across a variety of tasks (Nguyen et al., 2023).

1. Definition and Motivation

Given probability measures $\mu, \nu \in \mathcal{P}_p(\mathbb{R}^d)$, the classical $p$-th Sliced Wasserstein distance is:

$$SW_p(\mu, \nu) = \left[ \mathbb{E}_{\theta \sim \text{Uniform}(\mathbb{S}^{d-1})} W_p^p(\theta_\#\mu, \theta_\#\nu) \right]^{1/p}$$

where $\theta_\#\mu$ denotes the push-forward of $\mu$ by the linear projection $x \mapsto \theta^\top x$, and $W_p$ is the 1D Wasserstein distance.

EBSW generalizes this framework by replacing the uniform distribution over projection directions $\theta$ with an energy-based density

$$\sigma_{\mu, \nu}(\theta; f, p) \propto f(W_p^p(\theta_\#\mu, \theta_\#\nu))$$

with $f\!:\! [0, \infty)\!\to\! (0,\infty)$ typically monotonic (e.g., $p$0 for increasing $p$1 or $p$2).

The EBSW metric is then defined as

$p$3

This reweighting causes sampling to concentrate on projections where the differences between $p$4 and $p$5 are most pronounced, thus accentuating statistically informative dimensions and providing higher signal-to-noise ratios for tasks such as gradient flow or generative modeling.

2. Theoretical Properties

EBSW retains the core mathematical properties expected of a metric-like divergence for probability measures:

• Semi-metricity: For every $p$6 and any strictly positive energy $p$7, EBSW satisfies:
  • Non-negativity: $p$8,
  • Symmetry: $p$9,
  • Identity of indiscernibles: $$SW_p(\mu, \nu) = \left[ \mathbb{E}_{\theta \sim \text{Uniform}(\mathbb{S}^{d-1})} W_p^p(\theta_\#\mu, \theta_\#\nu) \right]^{1/p}$$0 if and only if $$SW_p(\mu, \nu) = \left[ \mathbb{E}_{\theta \sim \text{Uniform}(\mathbb{S}^{d-1})} W_p^p(\theta_\#\mu, \theta_\#\nu) \right]^{1/p}$$1.
• Relations to SW, Max-SW, and $$SW_p(\mu, \nu) = \left[ \mathbb{E}_{\theta \sim \text{Uniform}(\mathbb{S}^{d-1})} W_p^p(\theta_\#\mu, \theta_\#\nu) \right]^{1/p}$$2: If $$SW_p(\mu, \nu) = \left[ \mathbb{E}_{\theta \sim \text{Uniform}(\mathbb{S}^{d-1})} W_p^p(\theta_\#\mu, \theta_\#\nu) \right]^{1/p}$$3 is non-decreasing, $$SW_p(\mu, \nu) = \left[ \mathbb{E}_{\theta \sim \text{Uniform}(\mathbb{S}^{d-1})} W_p^p(\theta_\#\mu, \theta_\#\nu) \right]^{1/p}$$4, with equality for constant $$SW_p(\mu, \nu) = \left[ \mathbb{E}_{\theta \sim \text{Uniform}(\mathbb{S}^{d-1})} W_p^p(\theta_\#\mu, \theta_\#\nu) \right]^{1/p}$$5. For any $$SW_p(\mu, \nu) = \left[ \mathbb{E}_{\theta \sim \text{Uniform}(\mathbb{S}^{d-1})} W_p^p(\theta_\#\mu, \theta_\#\nu) \right]^{1/p}$$6, $$SW_p(\mu, \nu) = \left[ \mathbb{E}_{\theta \sim \text{Uniform}(\mathbb{S}^{d-1})} W_p^p(\theta_\#\mu, \theta_\#\nu) \right]^{1/p}$$7.
• Topology: The topology induced by $$SW_p(\mu, \nu) = \left[ \mathbb{E}_{\theta \sim \text{Uniform}(\mathbb{S}^{d-1})} W_p^p(\theta_\#\mu, \theta_\#\nu) \right]^{1/p}$$8 is equivalent to weak convergence plus convergence of $$SW_p(\mu, \nu) = \left[ \mathbb{E}_{\theta \sim \text{Uniform}(\mathbb{S}^{d-1})} W_p^p(\theta_\#\mu, \theta_\#\nu) \right]^{1/p}$$9-th moments, matching the topology induced by the Wasserstein metric.
• Sample Complexity: For empirical measures $\theta_\#\mu$0 with $\theta_\#\mu$1 i.i.d. in a compact set, there exists $\theta_\#\mu$2 such that

$\theta_\#\mu$3

This mirrors the $\theta_\#\mu$4 sample complexity of SW, avoiding the curse of dimensionality (up to logarithmic factors).

3. Monte Carlo Algorithms for EBSW Computation

For discrete measures supported on at most $\theta_\#\mu$5 atoms, the computation of 1D Wasserstein distances costs $\theta_\#\mu$6. EBSW admits several efficient Monte Carlo estimators:

• Importance Sampling (IS): Draw $\theta_\#\mu$7 i.i.d. $\theta_\#\mu$8 (e.g., uniform on $\theta_\#\mu$9). Compute weights $\mu$0, where $\mu$1. Normalize and form the $\mu$2 estimator:

$\mu$3

where $\mu$4.

• Sampling-Importance-Resampling (SIR): After IS on $\mu$5 projections, resample $\mu$6 according to the normalized weights and output the average of the corresponding $\mu$7 values.
• Markov Chain Monte Carlo (MCMC):
  • Independent MH uses a uniform proposal over $\mu$8.
  • Random-Walk MH employs a von Mises–Fisher proposal $\mu$9 with specified concentration. Both achieve $x \mapsto \theta^\top x$0 per step.

These approaches offer asymptotically unbiased estimators for $x \mapsto \theta^\top x$1 as $x \mapsto \theta^\top x$2. Overall complexity is $x \mapsto \theta^\top x$3 for $x \mapsto \theta^\top x$4 projections.

4. Comparative Analysis of Slicing Strategies

Distinct approaches for selecting the slicing distribution yield different computational–statistical tradeoffs:

Approach	Slicing Distribution	Informative Directions	Computational Cost
Classical SW	Uniform on $x \mapsto \theta^\top x$5	No	$x \mapsto \theta^\top x$6
Parametric optimizer	Optimized parametric family	Yes, but limited by parametric family	Often expensive, unstable
Energy-based (EBSW)	$x \mapsto \theta^\top x$7	Yes (data-driven, nonparametric)	$x \mapsto \theta^\top x$8

EBSW uniquely provides a parameter-free, nonparametric adaptation in slicing, focusing computation on directions with the largest observed projected divergences. This yields robustness to misspecification that can hinder parametric approaches and maintains efficiency comparable to classical SW.

5. Empirical Evaluation

All reported experiments use $x \mapsto \theta^\top x$9.

• Point-Cloud Gradient Flows: For driving $W_p$0 to a fixed $W_p$1 using the Euler discretization of gradient flow, IS-EBSW$W_p$2 with $W_p$3 achieves the fastest convergence in $W_p$4, outperforming SW$W_p$5, Max-SW$W_p$6, and v-DSW$W_p$7 within comparable runtime constraints.
• Color Transfer: Modeling images as empirical RGB distributions, IS-EBSW$W_p$8 generates transfers that most closely approximate target colors in $W_p$9 while incurring computational cost nearly equal to SW$\theta$0.
• Deep Point-Cloud Reconstruction: Training a point-cloud autoencoder with IS-EBSW$\theta$1 as the reconstruction loss yields lower slice-$\theta$2 and true $\theta$3 errors compared to SW$\theta$4, Max-SW$\theta$5, and v-DSW$\theta$6 across epochs on held-out ModelNet40. Resulting reconstructions possess increased sharpness and fidelity.

| Method | Epoch 20 (SW$\theta$7, $\theta$8 ×100) | Epoch 100 (SW$\theta$9, $$\sigma_{\mu, \nu}(\theta; f, p) \propto f(W_p^p(\theta_\#\mu, \theta_\#\nu))$$0 ×100) | Epoch 200 (SW$$\sigma_{\mu, \nu}(\theta; f, p) \propto f(W_p^p(\theta_\#\mu, \theta_\#\nu))$$1, $$\sigma_{\mu, \nu}(\theta; f, p) \propto f(W_p^p(\theta_\#\mu, \theta_\#\nu))$$2 ×100) | |-------------|-------------------------------|-------------------------------|-------------------------------| | SW$$\sigma_{\mu, \nu}(\theta; f, p) \propto f(W_p^p(\theta_\#\mu, \theta_\#\nu))$$3 | 2.97, 12.67 | 2.29, 10.63 | 2.15, 9.97 | | Max-SW$$\sigma_{\mu, \nu}(\theta; f, p) \propto f(W_p^p(\theta_\#\mu, \theta_\#\nu))$$4 | 2.91, 12.33 | 2.24, 10.40 | 2.14, 9.84 | | v-DSW$$\sigma_{\mu, \nu}(\theta; f, p) \propto f(W_p^p(\theta_\#\mu, \theta_\#\nu))$$5 | 2.84, 12.64 | 2.21, 10.52 | 2.07, 9.81 | | IS-EBSW$$\sigma_{\mu, \nu}(\theta; f, p) \propto f(W_p^p(\theta_\#\mu, \theta_\#\nu))$$6 | 2.68, 11.90 | 2.18, 10.27 | 2.04, 9.69 |

Qualitative assessment indicates EBSW consistently produces smoother trajectories and reconstructions relative to comparator methods.

6. Practical Implications and Availability

EBSW generalizes SW by replacing uniform slicing with an energy-based density proportional to a strictly positive function of 1D Wasserstein projections. This maintains $$\sigma_{\mu, \nu}(\theta; f, p) \propto f(W_p^p(\theta_\#\mu, \theta_\#\nu))$$7 per slice computational complexity and introduces greater adaptability by focusing on the most discriminative projections.

Monte Carlo estimators—via importance sampling, SIR, or MCMC—are readily implementable. IS-EBSW$$\sigma_{\mu, \nu}(\theta; f, p) \propto f(W_p^p(\theta_\#\mu, \theta_\#\nu))$$8 empirically achieves lower reconstruction and transport errors compared with baseline SW variants, at only a minor increase in computational overhead. All supporting code and data are available at https://github.com/khainb/EBSW (Nguyen et al., 2023).