Technical Summary: The Observable Wasserstein Distance

Problem Statement

The primary challenge addressed in this paper is the computational intractability of calculating the exact Wasserstein distance between probability measures on large-scale, non-Euclidean datasets. While optimal transport (OT) theory provides a rigorous framework for comparing distributions, exact computation typically scales cubically or quadratically with the number of data points, rendering it impractical for massive datasets.

Existing approximations, such as the Sliced Wasserstein distance, offer computational efficiency by projecting measures onto one-dimensional subspaces. However, these methods rely on the linear structure of Euclidean spaces ($\mathbb{R}^d$), utilizing inner products and norms to define projection directions. Many datasets of interest—such as 3D shapes, graphs, protein structures, and functional surfaces—reside in general metric spaces where linearity and inner products are absent. Consequently, standard sliced Wasserstein approaches cannot be directly applied, and there is a need for a computationally tractable alternative that provides lower bounds on the Wasserstein distance in general Polish metric spaces.

Methodology

The Observable Wasserstein Framework

The authors introduce the Observable Wasserstein distance, a framework that generalizes the slicing concept to arbitrary metric spaces $(X, d)$. Instead of projecting onto linear subspaces, the method projects probability measures onto the real line using metric observables.

1. Metric Observables: An observable is defined as a 1-Lipschitz function $f: X \to \mathbb{R}$, satisfying $|f(x) - f(y)| \leq d(x, y)$. The set of all such functions is denoted $\Lambda(X)$.
2. Pushforward Measures: For a measure $\mu \in \mathcal{P}(X)$ and an observable $f$, the pushforward measure $f_\sharp \mu$ is defined on $\mathbb{R}$.
3. Distance Definition: The observable Wasserstein $p$-distance, $\theta_p(\mu, \nu)$, is defined as the supremum of the Wasserstein distances between the pushforward measures over all possible observables:
   $$\theta_p(\mu, \nu) := \sup_{f \in \Lambda(X)} w_p(f_\sharp \mu, f_\sharp \nu)$$
   Since observables are 1-Lipschitz, it follows that $\theta_p(\mu, \nu) \leq w_p(\mu, \nu)$, establishing it as a lower bound.

Hierarchical Approximation and Identifiability

To manage the computational cost of optimizing over the infinite-dimensional space $\Lambda(X)$, the authors construct a hierarchy of subspaces $\Lambda_n(X) \subseteq \Lambda(X)$.

• Basic Observables: The construction begins with distance-to-a-point functions $f_a(x) = d(x, a)$.
• Wedge Products: Higher-order observables are formed via weighted wedge products (minimums) of these basic functions: $f_{\mathbf{a}}^\alpha = \alpha_0 f_{a_0} \wedge \dots \wedge \alpha_n f_{a_n}$.
• The Hierarchy: $\Lambda_n(X)$ consists of observables formed by wedge products of at most $n+1$ basic functions. This induces a sequence of pseudo-metrics $\theta_{p,n}(\mu, \nu)$.

Theoretical Contributions: Stratified Injectivity

A central theoretical result is a metric-space analogue to the Cramér-Wold Device. The paper establishes that the collection of pushforward measures via observables in $\Lambda_n(X)$ uniquely characterizes a probability measure $\mu$ if the metric covering dimension of its support is $\leq n$.

• Metric Covering Dimension ($\xi_X$): A notion of dimension defined by the refinement of open covers into open balls.
• Stratified Injectivity Theorem: If the support of $\mu$ has metric covering dimension $\leq n$, then $\mu$ is uniquely determined by its projections $f_\sharp \mu$ for all $f \in \Lambda_n(X)$.
• Implication: For measures supported on spaces with low metric covering dimension (e.g., finite sets, graphs, or low-dimensional manifolds), restricting the observable space to $\Lambda_n(X)$ yields a genuine metric (not just a pseudo-metric) that is strictly lower-bounded by the true Wasserstein distance.

Discrete Computational Model

For practical implementation, the authors propose a discrete model where the metric space is approximated by a finite $\delta$-cover (a grid). The observables are restricted to those anchored on points within this finite set. Theoretical bounds are provided showing that the discrete observable distance converges to the continuous version as the grid resolution increases, with error proportional to the grid spacing $\delta$.

Key Results

Theoretical Findings

• Metric Property: $\theta_p$ is a metric on the space of all probability measures. The restricted version $\theta_{p,n}$ is a metric on the subspace of measures with support dimension $\leq n$.
• Equivalence for $p=1$: For $p=1$, the observable Wasserstein distance is exactly equal to the standard Wasserstein distance ($\theta_1 = w_1$), a result derived from the Kantorovich-Rubinstein duality.
• Monotonicity: The hierarchy satisfies $\theta_{p,n} \leq \theta_{p,m} \leq \theta_p \leq w_p$ for $n \leq m$, offering a tunable trade-off between computational efficiency and bound sharpness.

Numerical Experiments

The paper presents several proof-of-concept experiments comparing Observable Wasserstein (OW) distances against Sliced Wasserstein (SW), standard Wasserstein, and Chamfer distances:

1. Gaussian Classification: In high-dimensional Euclidean spaces, OW and SW perform comparably in distinguishing Gaussian classes. However, OW shows a relative performance advantage as dimension increases, potentially due to the concentration of measure phenomenon affecting 1D projections in SW.
2. Graph Distributions: On non-Euclidean graph spaces, OW significantly outperforms standard Wasserstein distance in classification tasks while maintaining superior computational scalability. OW demonstrates robustness to noise that classical Wasserstein lacks.
3. Sphere Distributions: Experiments on spheres illustrate the trade-off between the number of observables/functions and computational cost. Increasing the complexity of observables (higher $n$) reduces relative error but incurs combinatorial costs.
4. Point Cloud Classification (ModelNet10): On real-world 3D object data, OW distances achieve the best classification scores among scalable metrics, particularly in noisy regimes, outperforming both SW and Chamfer distance.
5. Deep Learning Application: The authors integrated OW into a point cloud autoencoder loss function. Using OW (specifically an $L_2$-style variant for differentiability) resulted in higher latent-space classification accuracy and more uniform point density in reconstructions compared to Chamfer distance, which produced sharper but unevenly dense outputs.

Significance and Claims

The paper claims to provide a computationally tractable framework for deriving lower bounds on the Wasserstein distance in general metric spaces, bypassing the limitations of Euclidean-specific methods like Sliced Wasserstein.

• Theoretical Novelty: The work establishes a rigorous link between the metric covering dimension of a measure's support and the order of observables required for unique recovery. This serves as a fundamental theoretical tool for understanding identifiability in non-Euclidean OT.
• Practical Utility: The framework offers a "tunable" approach where practitioners can adjust the order of observables ($n$) to balance computational cost against the sharpness of the distance bound.
• Scalability: The discrete model allows for efficient implementation on large datasets, with empirical evidence suggesting linear scaling in the number of projections/anchors, making it suitable for massive datasets where exact OT is infeasible.
• Robustness: The experiments suggest that OW distances may offer superior robustness to noise compared to classical Wasserstein distances, making them suitable for real-world data applications.

The authors position this work as a foundational step, noting that while the current numerical experiments are proof-of-concept, the framework opens avenues for future research into sampling convergence, bi-Lipschitz equivalence for $p \neq 1$, and the integration of learnable observables in deep learning pipelines.