Wasserstein-Space Visualizations

• Wasserstein-space visualizations are mathematical methods that employ optimal transport and geodesic structures to analyze and preserve the shape of probability distributions.
• They enable applications like geodesic PCA, covariance analysis, and discriminative dimensionality reduction, providing interpretable summaries of complex data.
• Recent advances offer scalable techniques, including sliced and kernel methods, to visualize dependencies and overcome embedding challenges in high-dimensional spaces.

Wasserstein-space visualizations refer to a family of mathematical and algorithmic techniques that harness the geometry of optimal transport to visualize, summarize, and explore variability among probability distributions. Rather than treating distributions as elements of a linear space, these methodologies rigorously exploit the weak-Riemannian or geodesic structure of the Wasserstein space $W_2$ and related metric spaces. This paradigm ensures that both the "shape" and support constraints of probability measures are naturally preserved during analysis and visualization. The field has rapidly evolved, catalyzing methodological advances for exploratory analysis, dimension reduction, clustering, statistical inference, and manifold learning on spaces of distributions.

1. Geodesic Structure and Principal Modes of Variation

A fundamental principle of Wasserstein-space visualizations is the use of geodesics—optimal transport interpolations between distributions—as summary descriptors of variability. Geodesic Principal Component Analysis (GPCA) in Wasserstein space $W_2(\Omega)$, where $\Omega \subseteq \mathbb{R}$, exemplifies this approach (Bigot et al., 2013). For a family of densities, GPCA exploits the isometry with a closed convex subset of a Hilbert space via the logarithmic map

$$\log_\mu(\nu) = F^{-1}_{\nu} \circ F_\mu - \operatorname{id}$$

where $F_\mu$ is the cumulative distribution of the reference measure. This characterization permits convex PCA in the tangent space and ensures that principal modes, when mapped back via the exponential map, are actual probability distributions following geodesics in $W_2$.

Key implications:

• The principal geodesic directions correspond to "natural" distributional variations (e.g., shifts in location, changes in scale) and respect nonnegativity and normalization.
• Variability explained by the first geodesic mode can be significantly higher than in classic functional PCA (e.g., explaining 96% vs. 81% for age pyramid data).
• Visualizations of these geodesics provide interpretable, artifact-free summaries of distributional data.

2. Wasserstein Covariance and Visualization of Dependency

Moving beyond summary variation, Wasserstein covariance extends classical second-moment concepts to measure joint variability for vector-valued random distributions (Petersen et al., 2018). The covariance is defined via optimal transport maps and appropriately aligned using parallel transport between tangent spaces. In the quantile representation,

$$\Sigma_\oplus(j, k) = \mathbb{E}[(F_j^{-1}(t) - F^*_{j,-1}(t))(F_k^{-1}(t) - F^*_{k,-1}(t))],\ t \in [0,1]$$

where $F^*_j$ is the Fréchet mean quantile function. This framework is especially meaningful for visualizing dependency structures—e.g., showing group-level covariance differences in brain imaging or temporal dependency patterns in mortality evolution—via heatmaps, correlation matrices, and hierarchical clustering.

3. Dimensionality Reduction and Discriminative Visualization

Dimensionality reduction in Wasserstein space leverages optimal transport to preserve both global and local structures. Methods like Wasserstein Discriminant Analysis (WDA) maximize class separation under a regularized Wasserstein distance objective, using entropy-regularized OT computed via Sinkhorn scaling (Flamary et al., 2016): $$\underset{P\in\Delta}{\max} \frac{\sum_{c<c'} W_\lambda(PX^c, PX^{c'})}{\sum_c W_\lambda(PX^c, PX^c)}$$ The resulting low-dimensional projections capture global class differences and local neighborhood structure, as evidenced by clustering and separation of digits in projected MNIST data and improved classification error rates. Wasserstein embeddings that represent objects as discrete probability distributions further allow for direct, lossless visualization of complex relationship structures without recourse to dimension reduction (Frogner et al., 2019). Embedding points as clouds enables nuanced depiction of semantic similarity (e.g., polysemous word representations).

4. Sliced and Kernel Wasserstein Visualizations

A major computational advance is the usage of sliced Wasserstein distances, which average 1D Wasserstein distances over many linear projections, leading to scalable and statistically efficient visualization tools in high dimensions (Park et al., 2023). For

$$SW(\mu, \nu) = \left( \int_{\mathbb{S}^{d-1}} W_2^2(\langle \theta, \cdot \rangle_\#\mu, \langle \theta, \cdot \rangle_\#\nu) d\theta \right)^{1/2}$$

the geometry is fundamentally different from the classic Wasserstein space: SW is not a geodesic (length) space and possesses a tangent structure related to negative Sobolev spaces. Visualization using SW-based distances is statistically more favorable, particularly in high dimensions, and features parametric convergence rates for empirical measures. Furthermore, kernelized Wasserstein distances (Oh et al., 2019) extend OT to nonlinear feature spaces, enhancing clustering and artifact detection in biomedical applications such as imaging.

Recent work establishes that the equivalence between Wasserstein and sliced metrics breaks in higher dimensions (SW geodesics may only be Hölder continuous in Wasserstein), underscoring the importance of metric choice in visualization and barycenter interpolation (Hopper, 9 Jul 2024).

5. Embeddability, Limitations, and Coarse Geometry

The question of how faithfully Wasserstein and related spaces can be embedded into Hilbert or Banach spaces has direct implications for visualization (Pritchard et al., 2023). There are fundamental obstructions: for $p>1$, the 2-Wasserstein space over $\mathbb{R}^2$ is snowflake universal, and no coarse embedding into Hilbert space exists for either the space of probability measures or the space of persistence diagrams. Stable visualization, therefore, may require working directly in Wasserstein space or using approximations that retain only part of the geometric information.

6. Fast and Scalable Visualization Techniques

Computational efficiency is a priority due to the increasing scale and complexity of distribution-valued data. Recent advances build regression frameworks to approximate Wasserstein distances using collections of sliced Wasserstein distances as predictors, yielding dramatic speedups in large-scale Wasserstein-space visualizations, such as those involving 3D point clouds and ShapeNet data (Nguyen et al., 24 Sep 2025). Linear regression models (unconstrained or constrained by sliced Wasserstein lower and upper bounds) enable accurate, scalable distance estimation, and their integration into deep embedding models (e.g., Wasserstein Wormhole, RG-Wormhole) further accelerates applications while preserving visualization quality.

7. Algorithmic and Statistical Innovations

The differential structure of Wasserstein space—including gradient flows, tangent cones, and Riemannian calculus—enables advanced visual analytics rooted in the geometry of optimal transport (Lanzetti et al., 2022). This encompasses simulating geodesic interpolations (e.g., morphing between measures for barycenter visualization), visualizing flow lines induced by Wasserstein gradients of risk functionals, and constructing principal or tangent directions for complex datasets. Novel depth functions such as Wasserstein spatial depth generalize order-based visualization tools to Wasserstein space, supporting clustering and inferential exploration in a non-linear, non-Euclidean context (Bachoc et al., 16 Nov 2024).

8. Application Domains and Practical Impact

Wasserstein-space visualization techniques have been empirically validated in a diverse range of applications:

• Analysis of population aging (via geodesic PCA on age distributions),
• Brain imaging connectivity (via Wasserstein covariance analysis of network densities),
• 3D shape analysis, point cloud clustering, and single-cell spatial omics (using scalable regression models for Wasserstein distance matrix estimation and UMAP/embedding visualization),
• Functional data analysis, mortality time series, and forecasting,
• Text analytics and hypernymy detection (using Wasserstein elliptical embeddings to represent word uncertainty in semantic tasks).

Each application exploits the interpretability and geometric fidelity of Wasserstein-space visualization, its respect for constraints on probability measures, and, where relevant, its computational tractability.

---

In summary, Wasserstein-space visualizations capitalize on optimal transport geometry, geodesic structure, covariance formulations, scalable approximations, and differential calculus on the space of distributions. They enable interpretable, faithful, and efficient summaries of variability and dependency in high-complexity, distribution-valued data, while overcoming the limitations inherent in linear and Euclidean tools. Advances continue to be motivated by theoretical insights (e.g., geodesic non-equivalence of metrics, embedding obstructions), practical computation (e.g., regression on sliced distances, kernelization), and diverse real-world challenges in statistics, machine learning, and scientific data analysis.