Linearized Wasserstein dimensionality reduction with approximation guarantees

arXiv:2302.07373 · cs.LG, cs.NA, math.NA, stat.ML · Submitted 2023-02-14 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom, and with me are Jane, Lu, senior AI researcher at Tsinghua, Meng, lead engineer at a mysterious AI startup and Lalam, the in-house Large Language Model.

Jane: Today's paper: "Linearized Wasserstein dimensionality reduction with approximation guarantees".

Tom: The gist: This paper introduces LOT Wassmap,

Jane: First, who's behind it and why it matters.

Paper summary: Jane: Looking at the title, "Linearized Wasserstein dimensionality reduction with approximation guarantees," I think the authors are really emphasizing that this isn't just another speed bump; they are providing a way to make these complex mathematical spaces actually usable for analysis.

Tom: Exactly. They take the challenge of finding low-dimensional structures in Wasserstein space—where data is naturally measures—and show how you can use linearized optimal transport approximations to get answers without drowning in massive computational costs from pairwise distance matrices #pg3.

Lu: The implication here is that we can move manifold learning algorithms into the Wasserstein space for large datasets much more practically than before, because they give us those approximation guarantees on the embedding quality #pg3.

Meng: For me, what this means practically is that if I'm working with complex data representations from my startup's AI models, I don't have to wait days for a matrix computation just to get a basic low-dimensional map; I can use this LOT Wassmap approach instead.

Lalam: The paper suggests that as long as the measures are ε-compatible pushforwards of a template measure µ and you use enough samples, then LOT distance is a computationally efficient approximation for the Wasserstein distance, ensuring that the distortion of the LOT Wassmap embedding is small with high probability #pg21.

Jane: It really boils down to being able to handle data that looks like distributions instead of just points in R n, and this method gives us a tractable path there by managing those approximation errors effectively.

Tom: So, the title really sums it up: we’re linearizing the transport process and giving you a safety net—those guarantees—so you know how good your low-dimensional map actually is.

Jane: That's right, so this work opens up more avenues for using these powerful geometric tools in areas where probability measures are the natural language of the data.

Conclusion: Tom: So, we're wrapping up this look at "Linearized Wasserstein dimensionality reduction with approximation guarantees." Basically, the authors are showing how you can find low-dimensional structures in these complex probability spaces without having to crunch massive pairwise distance matrices.

Jane: Right, they take the idea of finding shapes in data that looks like distributions and they show a method that uses approximations to do it efficiently. It’s about using linearized optimal transport to skip the heavy lifting of calculating all those distances between every pair of data points.

Lu: I think what's really interesting here is how they bridge the gap between the theory—the mathematical guarantees—and the actual computation. They don't just propose a fast algorithm; they prove that even with these approximations, you still get an embedding that’s close to what you’d get if you had perfect data.

Meng: From an engineering standpoint, that proof is crucial because we often have limited computing power and we deal with high-dimensional data where pairwise comparisons are just impossible. If the quality guarantee holds under those approximation errors, then this method becomes a viable tool for real-world applications.

Lalam: I see the potential here for improving how we structure complex knowledge. If you can reliably map high-dimensional probability measures onto a low-dimensional space with known error bounds, that means new ways to organize and search information based on those measure relationships.

Tom: It sounds like the authors are really giving us a way to make these powerful geometric tools practical for large datasets. They've shown that you can get good results without doing the most expensive part of the process.

Jane: Exactly, it’s not about getting perfect distances, it’s about finding an embedding that is close enough based on what we actually have available. It gives us a solid foundation when we don't have explicit descriptions of the data measures themselves.

Lu: The work extends beyond just speed; they show how different approximation schemes, like Sinkhorn distances versus linearized transport, affect those error bounds and what kind of target measures they work best for.

Meng: So if I'm building a system, I need to know which approximation scheme is most stable when my data is messy or when the sample size isn't huge. The paper seems to lay out those conditions clearly.

Tom: It’s definitely setting up the next stage of research where we can fine-tune these approximations for specific types of data problems instead of just using a general setup.

Jane: And that's what they do, they provide the framework, and now other researchers can build on it to make these embeddings even more robust.

Lu: This opens up a whole new area for exploring how Wasserstein geometry interacts with generative models and deep learning architectures.

Tom: Next up, we’re going to look at some of the specific experiments they ran to see how this actually performs across different types of data structures.

University of California, San Diego · Halcioglu Data Science Institute of the University of California, San Diego Department of Mathematics at the University of Texas at Arlington Department of Mathematics at the University of North Carolina at Chapel Hill

cs.LG, cs.NA, math.NA, stat.ML

Submitted: 2023-02-14

Updated: 2026-10-08

Comments: 38 pages, 10 figures. Submitted

Journal ref: Applied and Computational Harmonic Analysis, 74 (2025), 101718

Code: https://github.com/varunkhuran/LOTWassMap

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 76/100

The gist: The gist: This paper introduces LOT Wassmap, a computationally feasible algorithm for dimensionality reduction in Wasserstein space that uses linearized optimal transport approximations to avoid

Key concepts

Wasserstein Space: This is a mathematical space used for data analysis where points are probability measures instead of simple vectors in R^n. It allows researchers to study how probability distributions change and find low-dimensional shapes within this space, which is crucial when data naturally represents distributions.
Optimal Transport (OT): Optimal Transport is a mathematical tool that finds the most efficient way to move one distribution into another. In this context, it helps define the distance between two probability measures in a way that respects their underlying structure and allows for dimensionality reduction.
Linearized Optimal Transport (LOT): LOT is an approximation scheme used to speed up optimal transport computations. Instead of solving the complex original problem exactly, LOT uses linearized versions of the transport map, such as Sinkhorn distances, to compute approximate distances much faster and avoid calculating a full pairwise distance matrix.

Terminology

Embeddings — 1× in this paper · explained in 2 episodes
Optimal Transport — 1× in this paper · explained in 5 episodes
Wasserstein Distance — 1× in this paper · explained in 5 episodes

Summary

The gist: This paper introduces LOT Wassmap, a computationally feasible algorithm for dimensionality reduction in Wasserstein space that uses linearized optimal transport approximations to avoid expensive pairwise distance matrix computations, providing theoretical guarantees on embedding quality under these approximations.

Introduction and Motivation

The paper addresses the challenge of finding low-dimensional structures in datasets naturally interpreted as probability measures rather than points in R n, which requires manifold learning algorithms in the Wasserstein space (Page 2). Most existing algorithms rely on computing pairwise Wasserstein distance matrices, which is computationally challenging for large datasets in high dimensions (Page 3). The authors leverage approximation schemes such as Sinkhorn distances and linearized optimal transport (LOT) to speed up computations and avoid this matrix computation entirely (Page 3). Their results provide both approximation guarantees on embedding quality and an efficient algorithm, LOT Wassmap, which computes the truncated SVD of centered transport maps instead of the pairwise LOT distance matrix, leading to only O(N) optimal transport computations (Page 3).

Approximation Guarantees

The paper provides two main types of approximation guarantees:

Approximation guarantees: They provide bounds on the embedding quality of Multidimensional Scaling (MDS) when pairwise Wasserstein distances are only available up to an error τ, studying the size of τ in schemes like entropic regularization and LOT, especially when explicit descriptions of data points are unavailable (Page 4).
Efficient algorithm (LOT Wassmap): They show that by using linearized Wasserstein distance approximations through LOT, Algorithm 2 produces data points that are CN,W (τ1 +τ2)-close to the isometric manifold W up to rigid transformations, where τ1 and τ2 represent the errors from approximation schemes and finite samples (Page 4).

LOT Wassmap Algorithm

The algorithm LOT Wassmap is inspired by the Wassmap algorithm [21] but uses LOT distances (or Sinkhorn distances in the approximate case) to linearly approximate W2(µi,µj) (Page 3). It computes only O(N) optimal transport maps instead of computing the pairwise LOT distance matrix, which represents a significant computational saving given the high cost of computing a single optimal transport map for densely sampled measures (Page 3). The output is equivalent to MDS using pairwise squared LOT distances, but it uses the truncated SVD trick to avoid actually computing the distance matrix (Page 3).

Main Theorem and Bounds

Theorem 3.3 states that if approximations of pairwise W2(R n) distances are used, specifically if W2(µi,µj)2 −λ2i j ≤ τ1 and W2(µi,µj)2 −λ2i j ≤ τ2 for some τ1 and τ2, then the resulting embedding points satisfy min Q∈O(d) kZ −QY kF ≤ (1+p 2)kY†kN (τ1 +τ2) (Page 3). Corollary 3.4 specializes this to Algorithm 2, showing that if the approximation error τ2 is bounded by a term derived from the empirical linearized Wasserstein-2 distance WcLOT2,σ(µb,νb;γ), then the embedding quality bound is attained (Page 14).

Approximation Schemes and Conditions

The paper details how different approximation schemes affect the bound τ2. For compactly supported target measures, bounds are derived using linear programming estimators (Theorem 4.2) or Sinkhorn estimators (Theorem 4.7) (Page 16). For non-compactly supported targets, it is shown that the error can be bounded by relating the true Wasserstein distance to an approximation involving a compactly supported measure µe, where W1(µ,µe) < η (Page 20). Conditions on H and µ are derived in Sections 6 and 7 to ensure these approximation guarantees hold for both barycentric map and entropic map estimators (Page 31).

Experiments

Experiments demonstrate that Algorithm 2 attains correct embeddings given finite sampling without explicitly computing the pairwise Wasserstein distances (Page 8). They test both LP and Sinkhorn variants, showing that the embedding quality improves with increased sample size m and that LOT Wassmap significantly reduces computational cost compared to Wassmap [21] (Page 8). The experiments cover various manifolds, including circle translation, rotation manifold, grid translation manifold, and dilation manifold (Pages 8–26).

Conclusion

The paper concludes that as long as the measures are ε-compatible pushforwards of a template measure µ and the number of samples used is large enough, LOT distance is a computationally efficient and tractable approximation for the Wasserstein distance, ensuring that the distortion of the LOT Wassmap embedding is small with high probability (Page 21). The results extend to non-compactly supported measures provided their densities decay fast enough (Page 20). The conditions derived in Section 6 and 7 ensure that these theoretical guarantees are satisfied for both barycentric and entropic map estimators (Page 38).

References

[1] Akram Aldroubi, Shiying Li, and Gustavo K Rohde. Partitioning signal classes using transport transforms for data analysis and machine learning. Sampling Theory, Signal Processing, and Data Analysis, 19(1):1–25, 2021 (Page 33).

[2] Jason Altschuler, Jonathan Weed, and Philippe Rigollet. Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration. Advances in Neural Information Processing Systems, 2017-December:1965–1975, 2017 (Page 33).

[3] Luigi Ambrosio and Nicola Gigli. A user’s guide to optimal transport. In Modelling and optimisation of flows on networks, pages 1–155. Springer, 2013 (Page 8).

[4] Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient Flows: in Metric Spaces and in the Space of Probability Measures. Springer Science & Business Media, 2008 (Page 8).

[5] Ery Arias-Castro, Adel Javanmard, and Bruno Pelletier. Perturbation bounds for procrustes, classical scaling, and trilateration, with applications to manifold learning. Journal of machine learning research, 21 (Page 33).

[6] Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In Doina Precup and Yee Whye Teh, editors, Proceedings of Machine Learning Research, volume 70, pages 214–223 PMLR (Page 33).

[7] Saurav Basu, Soheil Kolouri, and Gustavo K. Rohde. Detecting and visualizing cell phenotype differences from microscopy images using transport-based morphometry. Proceedings of the National Academy of Sciences, 111(9):3448–3453, 2014 (Page 28).

[8] Mikhail Belkin and Partha Niyogi. Laplacian eigenmaps for dimensionality reduction and data representation. Neural Computation, 15(6):1373–1396, 2003 (Page 28).

[9] R.J. Berman. Convergence rates for discretized Monge–Ampère equations and quantitative stability of optimal transport. Found Comput Math, 21:1099–1140, 2021 (Page 33).

[10] Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Communications on pure and applied mathematics, 44(4):375–417, 1991 (Page 8).

[11] Luis A. Caffarelli. Boundary regularity of maps with convex potentials. Communications on Pure and Applied Mathematics, 45(9):1141–1151, 2002 (Page 33).

[12] Luis A. Caffarelli. The regularity of mappings with a convex potential. Journal of the American Mathematical Society, 5(1):99–104, 2002 (Page 33).

[13] Luis A. Caffarelli. Boundary regularity of maps with convex potential–II. Annals of Mathematics, 144(3):453–496, 1996 (Page 8).

[14] Yongxin Chen, Filemon Dela Cruz, Romeil Sandhu, Andrew L Kung, Prabhjot Mundi, Joseph O. Deasy, and Allen Tannenbaum. Pediatric sarcoma data forms a unique cluster measured via the earth mover’s distance. Scientific Reports, 7(1):7035 (Page 28).

[15] Ronald R Coifman and Stéphane Lafon. Diffusion maps. Applied and Computational Harmonic Analysis, 21(1):5–30, 2006 (Page 33).

[16] Michael AA Cox and Trevor F Cox. Multidimensional scaling. In Handbook of data visualization, pages 315–347 Springer, 2008 (Page 7).

[17] Marco Cuturi. Sinkhorn distances: lightspeed computation of optimal transport. In NIPS, volume 2, page 4, 2013 (Page 8).

[18] Nabarun Deb, Promit Ghosal, and Bodhisattva Sen. Rates of estimation of optimal transport maps using plug-in estimators via barycentric projections.

Improvements for AI systems

Improved AI systems can perform low-dimensional manifold learning on probability measures by using LOT Wassmap to uncover low-dimensional structures in Wasserstein space, which is computationally feasible compared to pairwise Wasserstein distance matrix computations.
The system can achieve guaranteed embedding quality under approximations, as Theorem 3.3 provides bounds on the embedding quality where The second source of error can always be made arbitrarily small given sufficient computation time or judicious choice of parameters.
The system can efficiently compute embeddings by leveraging the truncated SVD of the centered transport maps, which is computationally more efficient than computing the pairwise Wasserstein distance matrix, as described in Section 1.3.
The system can handle finite samples for probability measures by using empirical linearized Wasserstein-2 distance, denoted as WcLOT2σ(µb,νb), and bounding the error between it and the true Wasserstein-2 distance using Theorem 4.7 or Theorem 5.5 bounds on τ2.
The system can adapt to non-compactly supported target measures by utilizing Lemmas 5.3 and Theorems 5.4/5.5, which provide bounds on the embedding quality when approximating a non-compactly supported measure with a compactly supported one, ensuring the resulting approximation error is bounded by "η."
The system can be optimized for specific data structures by selecting appropriate estimators; for instance, it can use the entropic map that is derived from using Sinkhorn’s algorithm when applicable to achieve faster computation time.
The system can be trained on complex manifolds, such as rotation manifolds or translation manifolds, by using the LOT Wassmap Embedding Algorithm (Algorithm 2) with different regularization parameters, allowing for qualitative assessment and quantitative measurement of relative error versus sample size "m."