Estimation of Stochastic Optimal Transport Maps

The optimal transport (OT) map is a geometry-driven transformation between high-dimensional probability distributions which underpins a wide range of tasks in statistics, applied probability, and machine learning. However, existing statistical theory for OT map estimation is quite restricted, hinging on Brenier’s theorem (quadratic cost, absolutely continuous source) to guarantee existence and uniqueness of a deterministic OT map, on which various additional regularity assumptions are imposed to obtain quantitative error bounds. In many real-world problems these conditions fail or cannot be certified, in which case optimal transportation is possible only via stochastic maps that can split mass. To broaden the scope of map estimation theory to such settings, this work introduces a novel metric for evaluating the transportation quality of stochastic maps. Under this metric, we develop computationally efficient map estimators with near-optimal finite-sample risk bounds, subject to easy-to-verify minimal assumptions. Our analysis further accommodates common forms of adversarial sample contamination, yielding estimators with robust estimation guarantees. Empirical experiments are provided which validate our theory and demonstrate the utility of the proposed framework in settings where existing theory fails. These contributions constitute the first general-purpose theory for map estimation, compatible with a wide spectrum of real-world applications where optimal transport may be intrinsically stochastic.

💡 Research Summary

This paper presents a groundbreaking framework for estimating stochastic optimal transport (OT) maps, significantly broadening the scope of statistical theory beyond the restrictive settings of prior work. Traditional OT map estimation theory relies heavily on Brenier’s theorem (quadratic cost, absolutely continuous source) to guarantee a unique deterministic map, and further imposes strong regularity assumptions on the map and densities to derive quantitative error bounds. These conditions often fail or are unverifiable in real-world applications like domain adaptation across manifolds or modeling branching cellular trajectories, where optimal transport is inherently stochastic.

The core innovation is the introduction of a novel error metric, E_p(κ; μ, ν), for evaluating stochastic maps (Markov kernels) κ. E_p combines an optimality gap (the excess transportation cost of κ compared to the optimum W_p(μ, ν)) and a feasibility gap (the Wasserstein distance between the pushforward measure κ♯μ and the target ν). Crucially, E_p does not require the existence or uniqueness of a deterministic OT map, enabling analysis of a vastly wider range of problems.

The technical foundation is laid in a study of E_p’s stability properties (Lemmas 3-5). These lemmas quantify how E_p responds to perturbations in the source (μ) and target (ν) distributions under Wasserstein and Total Variation (TV) distances. This stability is key to all subsequent finite-sample and robust estimation analyses.

Under minimal assumptions (ν sub-Gaussian, μ with bounded 2p-th moments), the authors propose a computationally efficient rounding-based estimator ˆκ_n. It achieves a near-optimal finite-sample risk bound of E