Estimation of Stochastic Optimal Transport Maps

TL;DR

This paper introduces a general framework for estimating stochastic optimal transport maps by defining a new Ep error metric that does not require existence or uniqueness of a deterministic Brenier map. It develops efficient estimators (entropic-kernel and rounding) with near-optimal finite-sample guarantees and robustness to adversarial contamination, extending OT map theory to stochastic kernels. Under Hölder continuity of the optimal kernel, it provides improved guarantees via WDRO and discusses plug-in and wavelet-based approaches, achieving minimax-like rates under weaker assumptions than traditional Brenier-based analysis. The results are supported by experiments on irregular OT maps, illustrating both theoretical and practical advantages of the Ep framework for real-world, non-deterministic transport problems. This work thus offers a foundational, broadly applicable theory for stochastic OT map estimation with implications for domains where mass splitting is intrinsic.

Abstract

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.

Figures (6)

• Figure 1: Previous map estimation theory only accounts for the inner circle, despite many OT map applications lying outside it. The proposed estimation framework under $\mathcal{E}_p$ covers all possible $\mathsf{W}_p$ problems (subject to tail bounds for quantitative rates).
• Figure 2: Diagrams of 2 maps and a kernel for $\mathsf{W}_p(\mu,\nu)$, where in each $\mu$ is uniform over the blue line connecting $(0,0)$ and $(0,1)$ and, in orange, $\nu = T^\star_\sharp \mu$ for $T^\star(x) = ((-1)^{\lfloor x_2/\delta \rfloor},x_2)$. We depict $T^\star$ on the left, ${T(x) = (-(-1)^{\lfloor x_2/\delta \rfloor},x_2)}$ in the center, and kernel $\kappa_x = \mathop{\mathrm{Unif}}\nolimits(\{(-1,x_2),(1,x_2)\})$ on the right. While $T$ and $\kappa$ are far from $T^\star$ in an $L^p$ sense, they achieve $\mathcal{E}_p \leq \delta$ (indeed, both $T$ and $\kappa$ achieve zero optimality gap and at most $\delta$ feasibility gap). Taking $\delta \to 0$, $T^\star$ becomes impossible to recover from finite samples, whereas $\kappa$ can be estimated effectively.
• Figure 3: $\mathcal{E}_1$ and $L^1$ performance of nearest-neighbor and rounding estimators in two settings.
• Figure 4: Randomized Rounding for Efficient OT Kernel Estimation
• Figure 5: $\mathcal{E}_1$ (left), optimality gap (middle), and feasibility gap (right) performance of nearest-neighbor and rounding estimators for Setting A.

Theorems & Definitions (39)

• Lemma 1: $\mathsf{W}_p$ empirical convergence, lei2020convergence
• Proposition 1: Relation to $L^p$ loss
• Remark 1: Optimality vs. feasibility
• Remark 2: Reverse $L^2$ comparison
• Lemma 2: Comparison to Monge gap
• Lemma 3: Stability in $\nu$
• Lemma 4: $\mathsf{W}_p$ stability in $\mu$
• Lemma 5: TV stability in $\mu$
• Lemma 6: Kernel composition
• Theorem 1: Entropic kernel estimator