Wasserstein Distance-Preserving Embeddings

• Wasserstein distance-preserving embeddings are techniques that map probability measures to low-dimensional Euclidean spaces while maintaining the geometric structure defined by the p-Wasserstein metric.
• Methods such as Wassmap, LOT Wassmap, and Fourier Sliced-Wasserstein yield near-isometric embeddings that recover intrinsic manifold structures and facilitate efficient optimal transport computations.
• Recent neural approaches, including Deep Wasserstein Embedding and Wasserstein Wormhole, leverage deep architectures to approximate OT distances for scalable inference in imaging and geometric data analysis.

Wasserstein distance-preserving embedding refers to a broad class of dimensionality reduction, representation learning, and algorithmic embedding techniques in which the geometry of a space of probability measures under an optimal transport metric (typically the $p$-Wasserstein metric, $W_p$) is approximately or (in certain families) exactly preserved in a lower-dimensional, often Euclidean, embedding space. The aim is to produce a representation or mapping $\Phi : \mathcal{P}_p(\Omega) \to \mathbb{R}^d$ of measures (or structured data as measures), such that for any pair $\mu, \nu$, the Euclidean distance $\|\Phi(\mu) - \Phi(\nu)\|$ closely tracks $W_p(\mu, \nu)$. Such embeddings are foundational for manifold learning, representation of point clouds, generative modeling, and analysis of geometric data.

1. Mathematical Foundations of Wasserstein Distance Preservation

The $p$–Wasserstein distance for probability measures $\mu,\nu$ on a metric space $(X,d)$ with finite $p$-th moments is

$W_p$0

where $W_p$1 denotes couplings with $W_p$2 as marginals. This metric extends the geometry of $W_p$3 to the space $W_p$4 of probability measures, endowing it with rich geometric structure. Preservation of $W_p$5 in an embedding context means retaining the true (geodesic or mass transport) relationships between data-distributions in a lower-dimensional code or vector space.

A core motivation is that in imaging and manifold learning, pixelwise Euclidean distance does not faithfully reflect similarity under natural image transformations (translations, dilations), whereas $W_p$6 (notably $W_p$7) is exactly sensitive to such changes, providing a metric aligned with semantic and geometric meaning (Hamm et al., 2022). Theoretical results confirm that for manifolds of translated or dilated distributions, $W_p$8 tracks the underlying parameter space exactly and can thus be used to recover intrinsic manifold structure through embedding.

2. Explicit Wasserstein-Preserving Algorithms: Isometric Mapping and Linearization

Wasserstein Isometric Mapping (Wassmap): Wassmap (Hamm et al., 2022) is a nonlinear dimensionality reduction pipeline that constructs embeddings for collections of images or measures:

• Each object (e.g., image) is represented by a probability measure $W_p$9.
• Compute the matrix $\Phi : \mathcal{P}_p(\Omega) \to \mathbb{R}^d$0 for all pairs.
• Apply classical Multidimensional Scaling (MDS): Double-center $\Phi : \mathcal{P}_p(\Omega) \to \mathbb{R}^d$1, diagonalize to extract a $\Phi : \mathcal{P}_p(\Omega) \to \mathbb{R}^d$2 dimensional embedding where Euclidean distances best match $\Phi : \mathcal{P}_p(\Omega) \to \mathbb{R}^d$3.

For families like translation and dilation manifolds, MDS on $\Phi : \mathcal{P}_p(\Omega) \to \mathbb{R}^d$4 recovers the intrinsic parameterization up to rigid motion. The discrete version applies the same idea to empirical measures/discrete images.

Linearized OT Embedding (LOT Wassmap): For computational scalability, LOT Wassmap (Cloninger et al., 2023) linearizes the $\Phi : \mathcal{P}_p(\Omega) \to \mathbb{R}^d$5 geometry via a reference measure $\Phi : \mathcal{P}_p(\Omega) \to \mathbb{R}^d$6:

• Each data measure $\Phi : \mathcal{P}_p(\Omega) \to \mathbb{R}^d$7 is mapped to its optimal transport map $\Phi : \mathcal{P}_p(\Omega) \to \mathbb{R}^d$8 (from $\Phi : \mathcal{P}_p(\Omega) \to \mathbb{R}^d$9 to $\mu, \nu$0).
• The $\mu, \nu$1 distance $\mu, \nu$2 approximates $\mu, \nu$3.
• Embeddings are produced from these maps using SVD, yielding efficient and theoretically robust isometric embeddings for classes of measures where the linearization is accurate.

Both approaches give exact or near-isometric embeddings under model assumptions (e.g., translation/dilation manifolds, measures close to a reference), and admit O($\mu, \nu$4) or faster (in the linearized case) algorithmic scaling.

Algorithm	Distance Preserved	Method	Recovery Guarantee
Wassmap	$\mu, \nu$5 (exact, manifold-specific)	Pairwise OT + MDS	Exact for translation/dilation
LOT Wassmap	$\mu, \nu$6 (first-order)	OT maps to reference + SVD	Approx., quantifiable error

3. Fourier Sliced-Wasserstein and Hilbert Space Embeddings

Fourier Sliced-Wasserstein (FSW) Embedding: FSW (Amir et al., 2024, Amir et al., 3 Apr 2025) constructs finite-dimensional Euclidean embeddings that approximately (and, for finite multisets, bi-Lipschitzly) preserve sliced Wasserstein distance, defined as

$\mu, \nu$7

FSW exploits the Fourier transform of projected measures, producing a vector per direction and frequency via

$\mu, \nu$8

with $\mu, \nu$9 directions and $\|\Phi(\mu) - \Phi(\nu)\|$0 frequencies. For multisets of size $\|\Phi(\mu) - \Phi(\nu)\|$1 in $\|\Phi(\mu) - \Phi(\nu)\|$2, the FSW embedding is injective and bi-Lipschitz with optimal output dimension $\|\Phi(\mu) - \Phi(\nu)\|$3.

A fundamental result is that there exists no finite-dimensional bi-Lipschitz embedding of the full Wasserstein space of measures into $\|\Phi(\mu) - \Phi(\nu)\|$4 for $\|\Phi(\mu) - \Phi(\nu)\|$5, due to the impossibility of controlling high-frequency structure with a fixed feature set; FSW achieves the best possible guarantee for practically relevant classes of inputs (finite sets, empirical distributions) (Amir et al., 2024, Amir et al., 3 Apr 2025).

4. Trainable Wasserstein-Preserving Neural Embeddings

Several recent approaches involve learning Wasserstein-preserving embeddings via deep neural architectures, leveraging large scale datasets:

• Deep Wasserstein Embedding (DWE) (Courty et al., 2017): A Siamese network learns $\|\Phi(\mu) - \Phi(\nu)\|$6 (encoder), paired with a decoder $\|\Phi(\mu) - \Phi(\nu)\|$7, training so that $\|\Phi(\mu) - \Phi(\nu)\|$8 and $\|\Phi(\mu) - \Phi(\nu)\|$9. This enables fast evaluation of $W_p(\mu, \nu)$0-like distances and downstream barycenter/geodesic computations within the embedding space.
• Wasserstein Wormhole (Haviv et al., 2024): A transformer-based autoencoder maps point clouds to a latent space where Euclidean distances align with (Sinkhorn-regularized) OT distances, with rigorous MDS-inspired error bounds. The architecture allows for linear-time computation of approximate OT distances and explicit decoding for barycenters and geodesics.
• Deep Kuratowski Embedding (KENN/ODE-KENN) (He, 6 Apr 2026): Inspired by the Kuratowski embedding theorem, these neural nets (standard multi-layer or ODE-based) aggregate feature space distances with learned weighting to approximate $W_p(\mu, \nu)$1, achieving high-fidelity surrogates for $W_p(\mu, \nu)$2 oracles on datasets like MNIST.

All these architectures optimize a loss function designed to minimize the discrepancy between Euclidean distances in the embedded space and the Wasserstein distances in the input space, often with auxiliary reconstruction or margin-based terms. Some combine these neural embeddings with additional structures (e.g., MLP on FSW codes) for further accuracy in $W_p(\mu, \nu)$3 regression (Amir et al., 2024).

5. Alternative Approaches: Gromov-Wasserstein and Metric Tree Embeddings

Alternative metric-preserving embedding techniques have been developed for spaces where ground metric structure varies or is relational/multi-view:

• Gromov-Wasserstein Embedding (Eufrazio et al., 26 Apr 2026): Methods such as Bary-GWMDS and Mean-GWMDS-C build Euclidean embeddings by minimizing Gromov-Wasserstein discrepancies across multiple pairwise or view-dependent distance matrices, yielding consensus embeddings and geometry-aware clustering characteristic of non-rigid, relational, or multi-view data.
• Embedding via Metric Trees (Mathey-Prevot et al., 2021): For finite metric spaces $W_p(\mu, \nu)$4, if $W_p(\mu, \nu)$5 stochastically embeds into trees with distortion $W_p(\mu, \nu)$6, the Wasserstein space $W_p(\mu, \nu)$7 admits a bi-Lipschitz embedding into $W_p(\mu, \nu)$8 with exactly the same distortion, via explicit coordinate representations (Evans–Matsen formula for tree metrics).

6. Theoretical Guarantees, Impossibility Results, and Empirical Properties

• Exact Recovery: For classes of manifolds such as translated or dilated measure families, isometric embedding via MDS/Wassmap recovers ground-truth parameters up to rigid motion (Hamm et al., 2022).
• Approximation Quality: LOT Wassmap provides first-order approximations to $W_p(\mu, \nu)$9 with tight perturbation bounds, and FSW gives approximate isometries for the sliced case with constants converging to 1 as the embedding dimension increases (Cloninger et al., 2023, Amir et al., 2024).
• Optimality and No-go Theorems: It is impossible to obtain finite-dimensional bi-Lipschitz isometries of the full Wasserstein space of general measures ($p$0), both for full $p$1 and $p$2 (Amir et al., 2024, Amir et al., 3 Apr 2025), necessitating restriction to empirical or finite-support measures for optimal low-distortion embeddings.
• Empirical Metrics: Distortion, stress, grid/circular arrangement accuracy, and downstream performance (e.g., classification robustness and accuracy under embedding pooling schemes such as PointNet+FSW) are used to assess embedding quality (Hamm et al., 2022, Amir et al., 2024).

7. Applications, Implementation Considerations, and Limitations

Wasserstein distance-preserving embeddings enable:

• Fast estimation of OT distances for large datasets (e.g., replacing expensive $p$3 oracles in manifold learning pipelines) (Haviv et al., 2024, He, 6 Apr 2026, Courty et al., 2017).
• Improved geometric and semantic structure discovery in imaging, clustering, and word embeddings (e.g., Wasserstein t-SNE, Gaussian word embedding) (Bachmann et al., 2022, Sun et al., 2018).
• Robust aggregation and pooling in point cloud networks (e.g., PointNet with FSW-pooling) (Amir et al., 2024).
• Efficient inference of barycenters and geodesics directly in the embedding domain (deep embeddings with decoders) (Courty et al., 2017, Haviv et al., 2024).

Key practical notes:

• Computation of pairwise $p$4 or FSW codes is substantial but can be dramatically accelerated via Sinkhorn approximation, randomized SVD (for MDS), and vectorized feature evaluation (Hamm et al., 2022, Amir et al., 2024).
• There are fundamental trade-offs between embedding dimension, storage, and the bi-Lipschitz (or isometric) constants achievable, determined by the size and structure of the dataset and the class of measures.
• Limitations include impossibility of global isometry for general measures, the necessity of architectural capacity in neural approaches, and computational bottlenecks for high-dimensional or large-support data.

References: (Hamm et al., 2022, Cloninger et al., 2023, Amir et al., 2024, Amir et al., 3 Apr 2025, Courty et al., 2017, Haviv et al., 2024, He, 6 Apr 2026, Mathey-Prevot et al., 2021, Bachmann et al., 2022, Sun et al., 2018, Eufrazio et al., 26 Apr 2026).