Linearized Wasserstein dimensionality reduction with approximation guarantees

Alexander Cloninger; Keaton Hamm; Varun Khurana; Caroline Moosmüller

cs.LGFeb 14, 2023

Linearized Wasserstein dimensionality reduction with approximation guarantees

Authors: Alexander Cloninger, Keaton Hamm, Varun Khurana, Caroline Moosmüller

Organizations: DEPARTMENT OF MATHEMATICS, UNIVERSITY OF CALIFORNIA, SAN DIEGO, CA · HALICIOĦGLU DATA SCIENCE INSTITUTE, UNIVERSITY OF CALIFORNIA, SAN DIEGO, CA · DEPARTMENT OF MATHEMATICS, UNIVERSITY OF TEXAS AT ARLINGTON, ARLINGTON, TX · DEPARTMENT OF MATHEMATICS, UNIVERSITY OF NORTH CAROLINA AT CHAPEL HILL, NC

Abstract

We introduce LOT Wassmap, a computationally feasible algorithm to uncover low-dimensional structures in the Wasserstein space. The algorithm is motivated by the observation that many datasets are naturally interpreted as probability measures rather than points in $\mathbb{R}^n$ , and that finding low-dimensional descriptions of such datasets requires manifold learning algorithms in the Wasserstein space. Most available algorithms are based on computing the pairwise Wasserstein distance matrix, which can be computationally challenging for large datasets in high dimensions. Our algorithm leverages approximation schemes such as Sinkhorn distances and linearized optimal transport to speed-up computations, and in particular, avoids computing a pairwise distance matrix. We provide guarantees on the embedding quality under such approximations, including when explicit descriptions of the probability measures are not available and one must deal with finite samples instead. Experiments demonstrate that LOT Wassmap attains correct embeddings and that the quality improves with increased sample size. We also show how LOT Wassmap significantly reduces the computational cost when compared to algorithms that depend on pairwise distance computations.

Figures & tables

Notation	Definition	Reference
$\Delta$	Square Euclidean distance matrix	Algorithm 1
$\Lambda$	Perturbed distance matrix	Corollary 3.2
$X^{\dagger}$	Moore–Penrose pseudoinverse of matrix $X$	Section 2.1
$\mu$	Template measure	Section 2.4
$\widehat{\mu}$	Empirical measure approximating $\mu$	( 7 )
$\sigma$	Reference measure for LOT	Section 2.4

Table 1. Overview of notation used in the paper.

Figure 1. 1-D Manifold of translations: (Left) reference measure $\sigma\sim\mathcal{N}(0,I)$ in blue and data measures $\mu_{i}$ which are Gaussians with the same covariance matrix and means $x_{i}$ uniformly sampled from the circle of radius $8$ . (Left Middle) Means $x_{i}$ of $\mu_{i}$ which are the true embedding points. (Right Middle) Embedding attained with Algorithm 2 using the linear program. (Right) Embedding attained with Algorithm 2 using the Sinkhorn distance with $\lambda=1$ .

Figure 2. Embedding error vs. $m$ (number of sample points from data and reference distributions for the 1-D translation manifold. Optimal transport maps are computed via the Linear Program (Left) and Sinkhorn with $\lambda=1$ (Right) .

Figure 3. 1-D Manifold of rotations: (Left) reference measure $\sigma\sim\mathcal{N}(0,I)$ in blue and data measures $\mu_{i}$ which are Gaussians with means lying on the circle of radius $8$ and covariance matrices that are rotations of each other. (Left Middle) Means $x_{i}$ of $\mu_{i}$ which are the true embedding points. (Right Middle) Embedding attained with Algorithm 2 using the linear program. (Right) Embedding attained with Algorithm 2 using the Sinkhorn distance with $\lambda=1$ .

Figure 4. Embedding error vs. $m$ (number of sample points from data and reference distributions for the 1-D rotation manifold. Optimal transport maps are computed via the Linear Program (Left) and Sinkhorn with $\lambda=1$ (Right) .

Figure 5. 2-D Manifold of translations: (Left) data measures $\mu_{i}$ which are Gaussians with the same covariance matrix and means $x_{i}$ taken from a $5\times 5$ uniform grid on $[-10,10]^{2}$ . (Left Middle) Means $x_{i}$ of $\mu_{i}$ which are the true embedding points. (Right Middle) Embedding attained with Algorithm 2 using the linear program. (Right) Embedding attained with Algorithm 2 using the Sinkhorn distance with $\lambda=10$ .

Figure 6. Embedding error vs. $m$ (number of sample points from data and reference distributions for the 2-D translation manifold. Optimal transport maps are computed via the Linear Program (Left) and Sinkhorn with $\lambda=10$ (Right) .

Figure 7. 2-D Manifold of Anisotropic Dilations: (Left) data measures $\mu_{i}$ which are Gaussians with mean $0$ and anisotropically dilated covariance matrices where dilations are taken from a $3\times 3$ uniform grid on $[1,4]^{2}$ . (Left Middle) Dilation factors $(x_{i},y_{i})$ of $\mu_{i}$ which are the true embedding points. (Right Middle) Embedding attained with Algorithm 2 using the linear program. (Right) Embedding attained with Algorithm 2 using the Sinkhorn distance with $\lambda=100$ .

Figure 8. Embedding error vs. $m$ (number of sample points from data and reference distributions for the 2-D translation manifold. Optimal transport maps are computed via the Linear Program (Left) and Sinkhorn with $\lambda=10$ (Right) .

Figure 9. Timing vs. sample size $m$ of the reference distribution and data measures. The data set consists of $N=25$ measures translated on a $5\times 5$ uniform grid on $[-10,10]^{2}$ as in Experiment 3. Shown are the computation times to compute the Wassmap embedding and the embeddings of Algorithm 2 using the Linear Program (LP) and Sinkhorn with regularization parameter $\lambda=1$ .

Figure 10. Timing vs. sample size $m$ of the reference distribution and data measures. The data set consists of $N=25$ measures translated on a $5\times 5$ uniform grid on $[-10,10]^{2}$ as in Experiment 3. Shown are the computation times to compute the embeddings of Algorithm 2 using the Linear Program (LP) and Sinkhorn with regularization parameters $\lambda=1$ and $\lambda=10$ .

Dimensionality Reduction

Dimensionality reduction maps high-dimensional data to compact representations for visualization, prediction, compression, and scientific modeling.

See all 145 papers

May 11, 2026

The Observable Wasserstein Distance

We introduce the observable Wasserstein distance, a framework for deriving lower bounds on the Wasserstein distance between probability measures on Polish metric spaces, designed to bypass the computational intractability of exact optimal transport in large-scale, non-Euclidean datasets. Analogous to the sliced Wasserstein distance in

\mathbb{R}^d

, our approach projects measures onto the real line via 1-Lipschitz observables and computes the Wasserstein distances between the resulting pushforward distributions. We define a hierarchy of pseudo-metrics by restricting observables to a nested chain of subspaces. A central theoretical contribution is an injectivity result linking the metric covering dimension of the support of a measure to the specific order in the hierarchy that guarantees unique recovery. This serves as a metric-space analogue to the Cramér-Wold Device for Euclidean distributions. We demonstrate that this hierarchy offers a tunable trade-off between sharpness as a lower bound on the Wasserstein distance and computational efficiency. We also present a discrete computational model for finite grids and numerical experiments validating the efficacy and utility of these approximations.

Edivaldo Lopes dos Santos, Leandro Vicente Mauri, Washington Mio +1

1Departamento de Matem´atica, Universidade Federal de S˜ao Carlos, Brazil · Department of Mathematics, Florida State University, USA

Sliced Wasserstein Distance Optimal Transport Wasserstein Distance

Ask PDF

May 19, 2026

Optimizing Computational-Statistical Runtime for Wasserstein Distance Estimation

Squared Wasserstein distance is a frequently used tool to measure discrepancy between probability distributions. This distance is typically computed between empirical measures of size

n

from two underlying random samples. Unfortunately, even in lower dimensional Euclidean space problems

\left( d \in \{2,3\} \right)

, algorithms for Wasserstein distance computation with approximate or exact precision guarantees scale poorly in the runtime as a function of

n

and the desired precision. In response, we consider the computational-statistical runtime, where the goal is to estimate from samples the Wasserstein distance between potentially smooth measures up to

ε

-additive error in expectation with respect to the sampling; we allow

O(1)

computational cost for collecting a sample. Towards this, we develop a Sample-Sketch-Solve paradigm where we introduce a regular cartesian grid sketch of the samples. We show that (especially under

α

-Hölder smooth distributions) this can compress the data without increasing asymptotic error, and also regularizes the structure which enables faster exact algorithms. Ultimately, we approximate

W_2^2(P,Q)

within

ε

error in

ε^{-\max(2,\frac{d+1+o(1)}{1+α})}

time for

0 < α< 1

Hölder smooth distributions

P,Q

(0,1)^{d}

; an optimal

Θ(ε^{-2})

for

α> 1/2

when

d=2

and nearly optimal as

α\to 1

when

d = 3

Peter Matthew Jacobs, Jeff M. Phillips

Department of Statistics University of Wisconsin-Madison Madison, WI 53706 · Kahlert School of Computing University of Utah Salt Lake City, UT 84101

Randomized Sketching Approximation Algorithms Optimal Transport Wasserstein Distance

Ask PDF

May 12, 2026

FastUMAP: Scalable Dimensionality Reduction via Bipartite Landmark Sampling

Exploratory analysis of high-dimensional data rarely stops at a single embedding. In practice, analysts rerun dimensionality reduction after changing preprocessing, subsets, or hyperparameters, and standard nonlinear methods can quickly become the bottleneck. We introduce FastUMAP (Bipartite Manifold Approximation and Projection), a landmark-based method designed for this repeated-use setting. FastUMAP builds a sparse point-landmark fuzzy graph, computes a Nystrom spectral warm start from the induced landmark affinity, and then refines all sample coordinates with a UMAP-style objective on the bipartite graph. The landmark ratio r = m/n provides a direct way to trade runtime against fidelity. On 9 benchmark datasets spanning 178 to 70,000 samples, FastUMAP has the lowest runtime on 7 datasets in our reported default-implementation comparison on one workstation. On MNIST and Fashion-MNIST (n=70000), it runs in about 4.6 seconds, compared with about 73--75 seconds for Barnes--Hut t-SNE, while reaching 91.4% mean kNN accuracy versus 94.6% for the strongest accuracy baseline. FastUMAP is therefore best viewed as a fast option for repeated exploratory embedding, rather than as a replacement for accuracy-first methods.

Hongmin Li

School of Life Science and Technology, Institute of Science Tokyo · 2-12-1 Ookayama, Meguro-ku, Tokyo 152-8550, Japan 2Department of Computational Biology and Medical Sciences Graduate School of Frontier Sciences, The University of Tokyo

Unsupervised Learning Nyström Method Manifold Learning Dimensionality Reduction

Ask PDF

Abstract

Dimensionality Reduction

Explore similar work

The Observable Wasserstein Distance

Optimizing Computational-Statistical Runtime for Wasserstein Distance Estimation

FastUMAP: Scalable Dimensionality Reduction via Bipartite Landmark Sampling