Source-linked AI summary
Subspace Robust Wasserstein Distances
François-Pierre Paty, Marco Cuturi
TL;DR
High-dimensional Wasserstein distances can be unstable and difficult to estimate robustly. The paper introduces projection- and subspace-robust variants, derives a tight eigenvalue-based convex relaxation, and provides entropic algorithms whose experiments indicate improved robustness to perturbations and better signal capture. The approach retains several Wasserstein geometric properties but faces computational challenges in directly optimizing the relaxation.
Problem
High-dimensional Wasserstein distances can be unstable, with sample complexity that may grow exponentially with dimension.
Method
The paper defines max-min projection robustness and a min-max subspace-robust formulation, then computes the latter using entropic regularization within a Frank-Wolfe-type algorithm.
Results
The min-max formulation has a tight convex relaxation expressed through the sum of the k largest eigenvalues of the displacement moment matrix.
Takeaways & Limitations
The proposed SRW distances retain important Wasserstein properties while appearing more robust to random perturbations and better able to capture signal.
Takeaways & Limitations
Direct projection onto the transportation polytope would require a costly quadratic network flow problem, and Frank-Wolfe cannot be applied directly to the relevant eigenvalue objective.
Abstract
from arXiv · showhide
Making sense of Wasserstein distances between discrete measures in high-dimensional settings remains a challenge. Recent work has advocated a two-step approach to improve robustness and facilitate the computation of optimal transport, using for instance projections on random real lines, or a preliminary quantization of the measures to reduce the size of their support. We propose in this work a "max-min" robust variant of the Wasserstein distance by considering the maximal possible distance that can be realized between two measures, assuming they can be projected orthogonally on a lower $k$-dimensional subspace. Alternatively, we show that the corresponding "min-max" OT problem has a tight convex relaxation which can be cast as that of finding an optimal transport plan with a low transportation cost, where the cost is alternatively defined as the sum of the $k$ largest eigenvalues of the second order moment matrix of the displacements (or matchings) corresponding to that plan (the usual OT definition only considers the trace of that matrix). We show that both quantities inherit several favorable properties from the OT geometry. We propose two algorithms to compute the latter formulation using entropic regularization, and illustrate the interest of this approach empirically.
1. Introduction
The paper addresses instability in high-dimensional optimal transport by extending projection-based robustness from random lines to k-dimensional subspaces. It formulates max-min and min-max variants, derives an eigenvalue-based convex relaxation, and develops computational tools and experiments.
- High-dimensional Wasserstein distances can require exponentially many samples for faithful estimation, motivating more robust OT metrics.
- Existing robustness strategies trade optimality for regularity through entropic or quadratic regularization, kernel approximations, or preliminary quantization.
- The paper extends projection-based OT analysis from random real lines to orthogonal subspaces of dimension k ≥ 2.
- The optimal projection defines a max-min robust OT problem with desirable properties.
- The corresponding min-max problem has a tight convex relaxation that minimizes the sum of the k largest eigenvalues of a displacement moment matrix.
- The paper studies PRW and SRW geometry, provides entropic-regularization algorithms, and evaluates the approach on simulated and real datasets.
2. Background on Optimal Transport
The paper represents quadratic optimal transport through couplings and their displacement second-order moment matrices. Under this representation, the usual 2-Wasserstein cost is the trace of the matrix associated with a transport plan.
- Monge and Kantorovich Formulations of OT: For probability measures in R^d, optimal transport considers couplings whose marginals are the two input measures.
- Monge and Kantorovich Formulations of OT: The 2-Wasserstein distance is defined by minimizing the expected squared Euclidean transport cost over admissible couplings.
- Monge and Kantorovich Formulations of OT: The paper uses W for the 2-Wasserstein distance because the remainder considers only quadratic costs.
- W as Trace-minimization: For a coupling, the displacement second-order matrix aggregates weighted outer products of displacement vectors.
- W as Trace-minimization: The total transport cost equals the trace of this displacement matrix, so computing W is trace minimization over couplings.
3. Subspace Robust Wasserstein Distances
The paper defines projection- and subspace-robust Wasserstein distances by optimizing transport over low-dimensional projections or subspaces. SRW admits an eigenvalue formulation, a tight convex relaxation, and metric properties.
- Robust formulations: Projection robust Wasserstein distance takes the worst OT cost over k-dimensional orthogonal projections.
- Robust formulations: Subspace robust Wasserstein distance is the corresponding min-max formulation focused on in the paper.
- Robust formulations: PRW and SRW are robust variants of W, with weak duality giving P_k(µ,ν) ≤ S_k(µ,ν).
- Eigenvalue characterization: SRW equals minimizing the sum of the k largest eigenvalues of the displacement second-order moment matrix over transport plans.
- Eigenvalue characterization: The eigenvalue characterization makes SRW a convex relaxation computable as a concave maximization over a convex set.
- Metric properties: Both PRW and SRW are distances on P2(R^d).
4. Geometry of Subspace Robust Distances
SRW preserves several geometric features of Wasserstein space while changing how transport emphasizes displacement directions. Its dependence on subspace dimension is increasing and concave, and it supports geodesic interpolation.
- Geometric properties: SRW shares fundamental geometric properties with the Wasserstein distance, including agreement with the ground metric for Dirac measures.
- Metric equivalence: SRW is equivalent to W, with explicit tight comparison constants.
- Dependence on the dimension: The SRW value increases concavely with the subspace dimension k.
- Geodesics: SRW induces a geodesic length space with explicit constant-speed geodesics for interpolating probability measures.
5. Computation
The paper develops algorithms for computing robust Wasserstein quantities and optimal subspaces for discrete measures, including entropic regularization and iterative initialization strategies.
- The computation targets the robust distance value, an optimal subspace, and an optimal transport plan, with subspaces also usable for dimensionality reduction.An optimal transport plan can additionally support geodesic interpolation between the measures.
- Computational challenges: Direct optimization over the transportation polytope is costly, while Frank-Wolfe cannot be applied directly because the eigenvalue-sum objective is nonsmooth.The direct projection leads to a costly quadratic network flow problem.
- Projected optimization: The maximization formulation can be solved by projecting onto the convex constraint set R with Dykstra’s algorithm.This approach returns the robust value and an optimal subspace, but not necessarily an optimal transport plan because OT plans may be nonunique.
- Entropy regularization: Entropic regularization makes the inner transport plan unique and differentiates the objective, enabling a Frank-Wolfe algorithm without learning-rate tuning.Sinkhorn’s algorithm computes the regularized transport plan at each iteration.
- Initialization and stopping: The algorithms initialize the projection with the top k eigenvectors of the transport displacement moment matrix from an ordinary optimal transport plan.This initialization costs the equivalent of one iteration and is optimal if that plan solves the min-max problem.
- Initialization and stopping: For all dimensions k, the method proceeds iteratively from k=d to k=1, reusing the previous plan’s eigenvectors and the latest Sinkhorn potentials.The regularized algorithm stops when the relative duality gap satisfies Δ(t)/Ŝ_k(t) ≤ ϵ.
6. Experiments
The experiments evaluate SRW on synthetic, noisy Gaussian, computational, and movie-script data, emphasizing dimension recovery, noise robustness, computation time, and semantic structure.
- 6.1. Fragmented Hypercube: For fragmented hypercubes, S2_k(µ,ν) becomes constant at 4k* when k≥k*, while empirical curves increase more slowly after k=k*.The slower increase reflects the remaining d−k* dimensions representing noise.
- 6.1. Fragmented Hypercube: The fragmented-hypercube experiments estimate distance and subspace errors over 500 samples as n ranges from 25 to 1000.Figures 3 and 4 report mean errors with 10%-90% and 25%-75% quantile bands.
- 6.1. Fragmented Hypercube: With n=250 and d=30, SRW geodesics are reported as robust to statistical noise compared with the Wasserstein-space mapping.The comparison uses optimal mappings under Wasserstein and SRW objectives.
- 6.3. Sk is Robust to Noise: Under increasing Gaussian noise, SRW distances seem more robust than Wasserstein distances at moderate to high variance.At small noise levels, SRW computation imprecision adds to the error caused by the noise.
- 6.4. Computation time: SRW computation is quadratic in dimension d in the reported GPU experiment because Algorithm 2 eigendecomposes Vπ.Figure 8 summarizes mean computation times over 100 random samplings with n=100.
- 6.5. Real Data Experiment: For seven movie scripts embedded in R300, films with similar genres or themes tend to have smaller SRW distances, and the optimal subspace recovers semantic dissimilarities.The Kill Bill Vol.1 and Interstellar projection illustrates distinct nearest-neighbor word sets.
7. Conclusion
The paper introduces robust OT distances with favorable theoretical and computational properties, provides tools for computing SRW distances, and identifies faster computation and differentiation as future directions.
- The paper proposes a new family of OT distances with robust properties that retain important properties of the 2-Wasserstein distance.
- SRW distances can be computed with algorithmic tools using regularized OT as the inner loop of a Frank–Wolfe-type algorithm.
- Future work targets faster computation and automatic differentiation schemes analogous to those used for Sinkhorn divergences.
- Projection robust Wasserstein distances are well posed because optimal subspaces exist, and they define a distance over probability measures.
- The projection-based distances satisfy an identity-of-indiscernibles property under projections whose subspaces cover the ambient space.
B. Proofs
The proofs establish existence, convex-relaxation identities, eigenvalue characterizations, projection-coupling correspondences, and geodesic properties for the robust Wasserstein formulations.
- Compactness and semicontinuity yield existence of optimal subspaces for projection robust Wasserstein distances.
- The min-max formulation reduces to minimizing the sum of the k largest eigenvalues of the displacement moment matrix over transport plans.
- The robust Wasserstein geodesic construction uses transported interpolations that achieve constant speed and connect the endpoint measures.
- Fan’s variational characterization provides an optimal projector and matching dual variables for the eigenvalue-sum problem.
- Projected couplings are exactly pushforwards of admissible couplings between the original measures.
C.1. Additional Experiment: Transport from Disk to Annulus
The disk-to-annulus experiments examine subspace recovery, estimation with increasing sample size, transport plans, and robustness to statistical noise.
- The experiment compares a k∗-dimensional disk and annulus embedded in a higher-dimensional space.
- An elbow appears at k = k∗ because the remaining dimensions represent noise recovered by the subspace robust distance.
- For k∗ = 2 and k = 2, the experiment separately evaluates the Wasserstein distance and subspace estimation behavior.
- The study evaluates mean estimation error and subspace estimation error over 500 samples across n ∈{25, 50, 100, 250, 500, 1000}.
- For n = 250 and d = 30, the Wasserstein and SRW optimal mappings are compared, with SRW geodesics reported as robust to statistical noise.
C.2. Details about Experiment of Section 6.5
The movie-script experiment represents documents as discrete measures in R300 using frequencies over an 18000-word vocabulary.
- Each movie script becomes a discrete measure in R300 whose weights are proportional to word frequencies.
- The vocabulary contains the 20000 most common English words except the 2000 most common, leaving 18000 words.
- Digits, capitalized words, and terms outside the vocabulary are removed before constructing the measures.