Sliced-Wasserstein Estimation with Spherical Harmonics as Control Variates

TL;DR

This work tackles the efficient estimation of the Sliced-Wasserstein distance by reducing Monte Carlo variance. It introduces Spherical Harmonics Control Variates (SHCV), which uses an orthonormal basis on the sphere to construct many control variates and achieves faster convergence than standard Monte Carlo. The authors establish elementary properties, an asymptotic error bound, and demonstrate substantial empirical improvements across synthetic Gaussian data, 3D point clouds, and SW-based kernel SVM tasks. The approach yields accurate SW estimates with lower sample complexity, facilitating scalable distribution comparison in ML applications.

Abstract

The Sliced-Wasserstein (SW) distance between probability measures is defined as the average of the Wasserstein distances resulting for the associated one-dimensional projections. As a consequence, the SW distance can be written as an integral with respect to the uniform measure on the sphere and the Monte Carlo framework can be employed for calculating the SW distance. Spherical harmonics are polynomials on the sphere that form an orthonormal basis of the set of square-integrable functions on the sphere. Putting these two facts together, a new Monte Carlo method, hereby referred to as Spherical Harmonics Control Variates (SHCV), is proposed for approximating the SW distance using spherical harmonics as control variates. The resulting approach is shown to have good theoretical properties, e.g., a no-error property for Gaussian measures under a certain form of linear dependency between the variables. Moreover, an improved rate of convergence, compared to Monte Carlo, is established for general measures. The convergence analysis relies on the Lipschitz property associated to the SW integrand. Several numerical experiments demonstrate the superior performance of SHCV against state-of-the-art methods for SW distance computation.

Figures (13)

• Figure 1: Visualization of the $L_2$-projection of the integrand $f$ onto the linear space $\operatorname{Span}\{\varphi_1,\ldots,\varphi_s\}$ of control variates.
• Figure 2: MSE curves for sampled Gaussian distributions in dimension $d \in \{3;6\}$.
• Figure 3: Boxplots of the error distribution $\widehat{\operatorname{SW}}_n(\mu_m,\nu_m)-\operatorname{SW}(\mu_m,\nu_m)$ for different SW estimates based on $n$ random projections with $n \in \{100;250;500;1000\}$ obtained over $100$ independent runs.
• Figure 4: MSE and computing time for Gaussian distributions, dimension $d \in \{3;6\}$, obtained over $100$ replications.
• Figure 5: MSE for sampled Gaussian distributions supported on $m=1000$ points, dimension $d \in \{3;6\}$, obtained over $100$ replications.

Theorems & Definitions (22)

• Definition 1: Sliced-Wasserstein distances
• Remark 1: Gaussian case when $p = 2$
• Remark 2: Complexity
• Remark 3: Regularization
• Definition 2: Polynomial spaces
• Definition 3: Spherical harmonics
• Proposition 1: Hilbert basis
• Remark 4: Examples
• Lemma 1: Number of control variates
• Definition 4: SHCV