---
title: Wasserstein Distance-Preserving Embeddings
url: https://www.emergentmind.com/topics/wasserstein-distance-preserving-embedding
type: topic
---

# Wasserstein Distance-Preserving Embeddings

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(\mu, \nu)^p = \inf_{\pi\in \Pi(\mu,\nu)} \int_{X\times X} d(x, y)^p \, d\pi(x, y),
\]
where $\Pi(\mu,\nu)$ denotes couplings with $\mu, \nu$ as marginals. This metric extends the geometry of $X$ to the space $\mathcal{P}_p(X)$ of probability measures, endowing it with rich geometric structure. Preservation of $W_p$ 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$ (notably $W_2$) is exactly sensitive to such changes, providing a metric aligned with semantic and geometric meaning [2204.06645]. Theoretical results confirm that for manifolds of translated or dilated distributions, $W_p$ 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 [2204.06645] 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 $\mu_i$.
- Compute the matrix $D_{ij} = W_p(\mu_i, \mu_j)^2$ for all pairs.
- Apply classical Multidimensional Scaling (MDS): Double-center $D$, diagonalize to extract a $d\ll n$ dimensional embedding where Euclidean distances best match $W_p$.

For families like translation and dilation manifolds, MDS on $W_p^2$ 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 [2302.07373] linearizes the $W_2$ geometry via a reference measure $\sigma$:
- Each data measure $\mu_i$ is mapped to its optimal transport map $T_\sigma^{\mu_i}$ (from $\sigma$ to $\mu_i$).
- The $L^2(\sigma)$ distance $\|T_\sigma^{\mu_i} - T_\sigma^{\mu_j}\|$ approximates $W_2(\mu_i, \mu_j)$.
- 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($N^2$) or faster (in the linearized case) algorithmic scaling.

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

## 3. Fourier Sliced-Wasserstein and Hilbert Space Embeddings

**Fourier Sliced-Wasserstein (FSW) Embedding:** FSW [2405.16519, 2504.02544] constructs finite-dimensional Euclidean embeddings that approximately (and, for finite multisets, bi-Lipschitzly) preserve sliced Wasserstein distance, defined as
\[
SW_p(\mu, \nu) = \left( \int_{S^{d-1}} W_p^p\left( P_{\theta\#} \mu,\, P_{\theta\#} \nu \right) d\theta \right)^{1/p}.
\]
FSW exploits the Fourier transform of projected measures,
producing a vector per direction and frequency via
\[
\Phi(\mu)_{(j,k,1)} = \mathrm{Re}\, \hat{P_{\theta_j\#}\mu}(\omega_k),\quad
\Phi(\mu)_{(j,k,2)} = \mathrm{Im}\, \hat{P_{\theta_j\#}\mu}(\omega_k),
\]
with $m$ directions and $K$ frequencies. For multisets of size $N$ in $\mathbb{R}^d$, the FSW embedding is injective and bi-Lipschitz with optimal output dimension $D\gtrsim 2Nd$.

A fundamental result is that there exists *no* finite-dimensional bi-Lipschitz embedding of the full Wasserstein space of measures into $\mathbb{R}^d$ for $d\geq 2$, 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) [2405.16519, 2504.02544].

## 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) [1710.07457]:** A Siamese network learns $\varphi: x \mapsto z$ (encoder), paired with a decoder $\psi: z \mapsto x$, training so that $\|\varphi(x) - \varphi(y)\|^2 \approx W_2^2(x, y)$ and $\psi(\varphi(x)) \approx x$. This enables fast evaluation of $W_2$-like distances and downstream barycenter/geodesic computations within the embedding space.

- **Wasserstein Wormhole [2404.09411]:** 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) [2604.04343]:** 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_2$, achieving high-fidelity surrogates for $W_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_2$ regression [2405.16519].

## 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 [2604.23912]:** 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 [2110.02115]:** For finite metric spaces $X$, if $X$ stochastically embeds into trees with distortion $D$, the Wasserstein space $P_1(X)$ admits a bi-Lipschitz embedding into $\ell^1$ 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 [2204.06645].
- **Approximation Quality:** LOT Wassmap provides first-order approximations to $W_2$ with tight perturbation bounds, and FSW gives approximate isometries for the sliced case with constants converging to 1 as the embedding dimension increases [2302.07373, 2405.16519].
- **Optimality and No-go Theorems:** It is impossible to obtain finite-dimensional bi-Lipschitz isometries of the full Wasserstein space of general measures ($d\geq 2$), both for full $W_p$ and $SW_p$ [2405.16519, 2504.02544], 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 [2204.06645, 2405.16519].

## 7. Applications, Implementation Considerations, and Limitations

Wasserstein distance-preserving embeddings enable:
- Fast estimation of OT distances for large datasets (e.g., replacing expensive $W_2$ oracles in manifold learning pipelines) [2404.09411, 2604.04343, 1710.07457].
- Improved geometric and semantic structure discovery in imaging, clustering, and word embeddings (e.g., Wasserstein t-SNE, Gaussian word embedding) [2205.07531, 1808.07016].
- Robust aggregation and pooling in point cloud networks (e.g., PointNet with FSW-pooling) [2405.16519].
- Efficient inference of barycenters and geodesics directly in the embedding domain (deep embeddings with decoders) [1710.07457, 2404.09411].

Key practical notes:
- Computation of pairwise $W_p$ or FSW codes is substantial but can be dramatically accelerated via Sinkhorn approximation, randomized SVD (for MDS), and vectorized feature evaluation [2204.06645, 2405.16519].
- 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:** [2204.06645], [2302.07373], [2405.16519], [2504.02544], [1710.07457], [2404.09411], [2604.04343], [2110.02115], [2205.07531], [1808.07016], [2604.23912].

Source: https://www.emergentmind.com/topics/wasserstein-distance-preserving-embedding