Estimation of Stochastic Optimal Transport Maps

AI summary

Overview

Research area: Statistical learning theory for optimal transport (OT), specifically the estimation of transport maps between probability distributions (arXiv category stat.ML).

Technical level: Advanced — the paper is dense with measure-theoretic probability, Wasserstein geometry, and minimax statistical theory.

Scope: The paper proposes a new error metric and a general-purpose estimation framework for optimal transport maps that allows stochastic (mass-splitting) maps, with finite-sample guarantees under minimal assumptions.

What This Paper Is About

Optimal transport maps are geometry-driven transformations between probability distributions, but existing statistical theory for estimating them is narrow: it relies on Brenier's theorem (quadratic cost, absolutely continuous source measure) to guarantee that a unique deterministic optimal map exists, and then imposes further smoothness assumptions on that map to get error bounds. Many real applications fail these conditions — for instance when the source distribution lies on a lower-dimensional manifold than the target, or when developmental trajectories branch — so the optimal transport plan must be stochastic. This paper builds an estimation theory that covers those cases, introducing a new way to measure how good a stochastic transport map is and deriving estimators with near-optimal rates under assumptions that are easy to check.

Key Contributions

1. A new error metric, the transportation error ℰ_p. For a Markov kernel κ from source μ to target ν, ℰ_p(κ; μ, ν) is defined as the sum of an optimality gap (the kernel's transport cost above the optimum W_p(μ, ν), truncated at zero) and a feasibility gap (the p-Wasserstein distance between the kernel's pushforward κ_♯μ and ν). The metric is zero exactly when κ is an optimal kernel, and it does not require uniqueness or even existence of a deterministic optimal map.

2. Stability lemmas for ℰ_p. The paper proves how ℰ_p changes under total-variation and Wasserstein perturbations of the source and target measures, and how it behaves under composition of kernels. These lemmas are the technical backbone for all subsequent estimation guarantees.

3. Computationally efficient estimators with near-optimal rates. Two estimators are analyzed — an entropic (EOT-based) kernel estimator and a rounding-based estimator — the latter achieving expected error Õ_{p,d}(n^{−1/(d+2p)}) under mild tail assumptions.

4. Robust estimation under adversarial contamination and minimax lower bounds. The framework extends to data corrupted by both total-variation and Wasserstein-budgeted adversaries, with matching lower bounds. Lower bounds for the clean setting (Ω(n^{−1/(d∨2p)})) show the proposed rates are near-optimal, and a sharper n^{−1/2} rate is proved for d = 1.

Main Findings

• ℰ_p generalizes the existing L^p benchmark. Proposition 1 shows ℰ_p(T; μ, ν) ≤ 2‖T − T^⋆‖_{L^p(μ)} for any map T and any W_p-optimal map T^⋆, but ℰ_p can be finite when L^p distance is meaningless or enormous — the paper's Figure 2 constructs maps whose L^p distance from T^⋆ is arbitrarily large while ℰ_p ≤ δ.

• ℰ_p is closely related to the Monge gap objective. For the alternative objective ℰ'_p that uses the Monge gap of Uscidda and Cuturi (2023), the paper shows ℰ_p ≤ 4ℰ'_p, and for p = 1, ℰ'_1/2 ≤ ℰ_1 ≤ 2ℰ'_1. The authors argue ℰ_p is better suited to quantitative statistical analysis, while the Monge gap objective seems better suited to neural implementation (its gradients carry a stronger signal far from optimality).

• Entropic kernel estimator. When 𝒳, 𝒴 ⊆ [0,1]^d and τ = d^{p/4} n^{−1/(2d∨4)} log n, the conditional kernel of the optimal entropic coupling for the empirical measures satisfies E[ℰ_p(κ̂_n; μ, ν)] ≲_{p,d} n^{−1/(2pd∨4p)} log²(n). For comparison, the authors note that taking the formal α → 0 limit of the Brenier map estimation rate of Pooladian and Niles-Weed (2021) gives n^{−1/(4d+4)}, which is always worse than the n^{−1/(4d∨8)} rate implied here (and which does not require p = 2 or an existing T^⋆).

• Rounding estimator gives the sharpest claimed rate. With μ, ν 1-sub-Gaussian and a regular partition of ℝ^d into cubes of side length r (with R, r, δ tuned independently of μ and ν), E[ℰ_p(κ̂_n; μ, ν)] = Õ_{p,d}(n^{−1/(d+2p)}). The same guarantee holds with a non-uniform partition if sub-Gaussianity of μ is relaxed to E_μ[‖X‖^{2p}] ≤ 1. Computation is dominated by one step — a single low-accuracy call to an entropic OT solver — with running time O((C_∞ + d) n^{2+o_d(1)}), where C_∞ = max_{i,j} ‖X_i − Y_j‖.

• The rounding rate is near-optimal. The upper bound n^{−1/(d+2p)} sits above the minimax lower bound Ω(n^{−1/(d∨2p)}) of Singh and Póczos (2018) inherited from W_p distribution estimation (obtained when μ is a point mass). For d = 1 the authors improve the rate to n^{−1/2} using stability of ℰ_p under the Kolmogorov–Smirnov metric, showing the lower bound cannot be improved in that case.

• Hölder-continuous optimal kernels reduce to distribution estimation. If there is a Lipschitz optimal kernel — W_p((κ^⋆)x, (κ^⋆){x′}) ≲ ‖x − x′‖ — kernel estimation under ℰ_p has the same statistical complexity as estimating μ and ν under W_p, with rate Õ(n^{−1/(d∨2p)}), obtained via an estimator based on Wasserstein distributionally robust optimization. This condition is much weaker than typical assumptions for Brenier map estimation when p = 2; for comparison, Balakrishnan and Manole (2025) obtain Õ(n^{−1/(d∨4)}) for nearest-neighbor estimation of a Brenier map under their own assumptions.

• Robustness to adversarial contamination. Against an adversary with TV budget ε and W_p budget ρ, a convolutional estimator achieves error √d ε^{1/p} + √d ρ^{1/(p+1)} + O_{p,d}(n^{−1/(d+2p)}). The accompanying minimax lower bound is √d ε^{1/p} + d^{1/4} ρ^{1/2} + n^{−1/(d∨2p)}. The gap in the ρ exponent implies a separation between robust map estimation under ℰ_p and robust distribution estimation under W_p, where linear dependence on ρ is attainable.

• Empirical validation. Section 6 reports numerical simulations in two settings with irregular OT maps that are poorly suited to existing theory, demonstrating the performance of the rounding estimator and the advantages of ℰ_p over L^p.

Methodology in Plain English

The authors start by reframing what it means for a transport map to be "good." Instead of comparing an estimated map pointwise to a unique ground-truth map (which only makes sense when such a map exists and is unique), they score a map — or a stochastic kernel — by two things: how much its transport cost exceeds the best possible cost, and how far its output distribution is from the target. This score, ℰ_p, is zero exactly for optimal transport plans and is well-defined even when no deterministic optimal map exists.

They then prove a series of "stability" lemmas showing that this score does not change much when the source or target distribution is slightly perturbed (in total variation or Wasserstein distance) or when kernels are composed. These lemmas let them analyze estimators by decomposing error into an empirical term and a perturbation term.

For estimation, they take i.i.d. samples from source and target. One estimator solves an entropic optimal transport problem on the empirical measures and reads off the conditional kernel. A second, sharper estimator first rounds the source samples onto a grid or partition, solves a transport problem between the rounded empirical measure and the target empirical measure, and then composes the result with the rounding function. The rounding step makes the kernel piecewise constant, which is what allows total-variation stability arguments to give a better rate under weaker assumptions.

They also prove matching lower bounds to show the rates cannot be substantially improved, and extend the whole analysis to a contamination model where an adversary may corrupt samples both globally (total variation) and locally (Wasserstein).

Why This Matters

Impact on research. This is claimed to be the first general-purpose theory for OT map estimation that is compatible with settings where optimal transport is intrinsically stochastic. It removes the dependence on Brenier's theorem and on hard-to-verify regularity assumptions (two-sided density bounds, Lipschitz or Hölder smoothness of the unique optimal map) that constrained prior work by Hütter and Rigollet (2021), Pooladian and Niles-Weed (2021), Deb et al. (2021), Manole et al. (2024), and others. It also provides the missing statistical rates for the neural stochastic map estimators of Korotin et al. (2023a, 2023b), which previously had only empirical justification.

Real-world applications cited or implied by the paper:

• Domain adaptation where the source distribution lies on a lower-dimensional manifold than the target — the paper names text-to-image and sketch-to-photo translation as examples.
• Single-cell genomics, where developmental trajectories branch over time, so any measure-preserving map from an early snapshot to a later one must be stochastic (Schiebinger et al. 2019; Bunne et al. 2023).
• Style transfer (Kolkin et al. 2019; Mroueh 2020).
• Generative modeling (Zhang et al. 2018; Vesseron et al. 2025).

Industry relevance. The framework's assumptions are described as minimal and easy to check, and the rounding estimator's runtime is dominated by a single low-accuracy entropic OT solver call, making the approach computationally practical. The robust variant is relevant to settings where training or reference data may be contaminated or adversarially perturbed.

Future Directions

• Close the gap between the upper and lower rates for d > 1. The rounding estimator gives n^{−1/(d+2p)} while the lower bound is n^{−1/(d∨2p)}. The authors resolve this only in the one-dimensional case (n^{−1/2}), and explicitly raise whether the lower bound can be improved in general.
• Tighten the robust rate in the Wasserstein-contamination budget. The upper bound scales as ρ^{1/(p+1)} while the lower bound scales as ρ^{1/2}; whether the true exponent lies between these is left open.
• Bridge ℰ_p and the Monge gap for statistical analysis. The paper states it is unaware of any existing statistical rates proven under the Monge gap objective 𝒥_p, while noting that 𝒥_p has practical advantages for neural implementation. Deriving rates under 𝒥_p is a natural next step.
• Extend beyond Euclidean support and simple tail conditions. The guarantees rely on tail bounds on μ and ν (sub-Gaussian, or bounded 2p-th moments); the authors note these can be weakened under their analysis but not without worsening the rate, leaving the question of how far the assumptions can be pushed.

Target Audience

Theoretical statisticians and machine learning researchers working on optimal transport, distribution estimation, and minimax theory will get the most from this paper, since the core contribution is a new error metric together with finite-sample upper and lower bounds. Practitioners applying OT to domain adaptation, single-cell trajectory analysis, style transfer, or generative modeling — especially those whose source and target distributions violate Brenier's conditions — will benefit from the new framework and its estimators. Readers will need comfort with measure-theoretic probability, Markov kernels, Wasserstein distances, entropic optimal transport, and minimax lower-bound techniques.

Authors’ 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.