Minimax estimation of discontinuous optimal transport maps: The semi-discrete case
Proposes an entropic regularized estimator for semi-discrete OT maps, achieving the minimax rate of n^{-1/2} in high-dimensional settings.
Key Findings
Methodology
Utilizing an entropic OT framework with Sinkhorn algorithm, the estimator Tε balances bias and variance through regularization parameter ε. Theoretical analysis employs stability bounds for entropic Brenier maps and dual potentials, establishing that with ε proportional to n^{-1/2}, the estimator converges at the parametric rate n^{-1/2} independent of dimension. The approach involves bias-variance decomposition, stability analysis, and a comparison with 1NN, which suffers from curse of dimensionality. The method's core innovation lies in the stability bounds for entropic maps and the explicit control of regularization effects, ensuring optimal convergence in the semi-discrete setting.
Key Results
- Under assumptions on P (support full space) and Q (discrete with positive weights), setting ε ≈ n^{-1/2} yields an error E∥Tε - ∇φ0∥²L2(P) that converges at n^{-1/2}, matching the minimax lower bound. Numerical experiments confirm the theoretical rate, showing superior performance over 1NN especially as dimension increases. The estimator remains robust across various sample sizes and supports, with errors diminishing proportionally to n^{-1/2}.
- Stability bounds for entropic Brenier maps and dual potentials are derived, demonstrating that the regularization parameter ε controls the bias, while sample variability influences the variance. These bounds enable precise tuning of ε, leading to optimal convergence. The analysis also reveals that 1NN's convergence rate deteriorates to at least n^{-1/d}, confirming its suboptimality in high dimensions.
- Experimental results validate the theoretical findings, with the estimator achieving near-parametric convergence in simulated semi-discrete scenarios. The method's robustness is evidenced by consistent error reduction across different sample sizes, dimensions, and number of target atoms, highlighting its practical viability for complex, high-dimensional data.
Significance
This work advances the theoretical understanding of non-smooth, semi-discrete optimal transport map estimation, providing the first provably optimal estimator with finite-sample guarantees. It addresses the critical challenge of non-continuous mappings in high-dimensional data, offering a scalable, robust solution that surpasses traditional methods like 1NN. The results have broad implications for applications requiring accurate modeling of non-smooth distributions, such as image registration, quantization, and biological data analysis. By establishing the minimax rate and demonstrating practical effectiveness, this research bridges a vital gap between theory and real-world problems involving complex, non-smooth data structures.
Technical Contribution
The paper introduces a novel estimator Tε based on entropic OT, supported by a rigorous stability analysis of entropic Brenier maps and dual potentials. It establishes that with ε proportional to n^{-1/2}, the estimator attains the minimax rate of convergence, independent of dimension. The work also develops new bounds for the stability of entropic maps, improving previous exponential dependencies on ε, and provides a detailed bias-variance decomposition framework. These contributions significantly enhance the theoretical foundation for high-dimensional, non-smooth OT estimation, enabling practical algorithms with provable guarantees.
Novelty
This is the first work to demonstrate that a regularized entropic OT estimator can achieve the parametric n^{-1/2} rate in the semi-discrete setting with non-smooth, discontinuous maps. It extends the theory of entropic OT beyond smooth, continuous maps, addressing the challenge of non-continuity at boundaries of Laguerre cells. The stability bounds for entropic Brenier maps and the explicit control of regularization effects represent key innovations, setting new standards for statistical guarantees in high-dimensional, non-smooth OT estimation.
Limitations
- The current approach assumes the target distribution Q is discrete with positive weights, limiting direct applicability to continuous or more complex target distributions. Extending the framework to non-discrete targets remains an open challenge.
- Choice of regularization parameter ε, while theoretically guided, still requires empirical tuning in practice, which could affect robustness and accuracy.
- Computational costs grow with sample size and dimension, especially in very high-dimensional scenarios, necessitating further algorithmic optimization or approximation techniques.
Future Work
Future research will focus on generalizing the estimator to continuous target distributions, possibly via discretization or kernel methods. Developing adaptive schemes for ε selection and scalable algorithms for ultra-high-dimensional data are also promising directions. Additionally, exploring the estimator's robustness under model misspecification and extending stability bounds to broader classes of non-smooth maps will further enhance its practical utility.
AI Executive Summary
Optimal transport (OT) provides a powerful geometric framework for comparing probability distributions, with applications spanning economics, biology, and computer vision. Traditional methods assume smooth, continuous maps, which limit their effectiveness in real-world scenarios featuring non-smooth, discontinuous, or discrete structures. This paper addresses this gap by proposing a novel estimator based on entropic regularization, specifically tailored for the semi-discrete setting where the target distribution is discrete and the source has full support. The core innovation lies in establishing that, by carefully tuning the regularization parameter ε proportional to n^{-1/2}, the estimator achieves the minimax optimal convergence rate of n^{-1/2}, independent of the ambient dimension. This marks a significant theoretical breakthrough, as previous methods suffered from curse of dimensionality or lacked finite-sample guarantees in such settings.
The methodology hinges on a detailed bias-variance decomposition, combined with new stability bounds for entropic Brenier maps and dual potentials. These bounds ensure that the regularized estimator remains robust against non-smoothness and boundary discontinuities, which are typical in semi-discrete maps. Numerical experiments on synthetic data validate the theoretical predictions, showing that the estimator consistently outperforms traditional approaches like 1NN, especially in high-dimensional regimes where the latter's error decays at a rate no faster than n^{-1/d}. The results demonstrate that the proposed approach not only achieves optimal statistical efficiency but also exhibits practical robustness, making it suitable for complex applications involving non-smooth data.
Furthermore, the paper's theoretical contributions include new stability results for entropic OT maps, improved bounds on the dependence of the regularized solutions on ε, and a comprehensive analysis of the bias-variance trade-off. These advances significantly deepen the understanding of entropic OT's statistical properties, especially in non-smooth, high-dimensional contexts. The work opens avenues for extending the framework to more general target distributions, adaptive regularization schemes, and scalable algorithms, promising broad impact across data science and machine learning fields. Overall, this research marks a pivotal step toward reliable, high-dimensional estimation of discontinuous optimal transport maps, bridging a crucial gap between theory and real-world data complexities.
Deep Dive
Abstract
We consider the problem of estimating the optimal transport map between two probability distributions, $P$ and $Q$ in $\mathbb R^d$, on the basis of i.i.d. samples. All existing statistical analyses of this problem require the assumption that the transport map is Lipschitz, a strong requirement that, in particular, excludes any examples where the transport map is discontinuous. As a first step towards developing estimation procedures for discontinuous maps, we consider the important special case where the data distribution $Q$ is a discrete measure supported on a finite number of points in $\mathbb R^d$. We study a computationally efficient estimator initially proposed by Pooladian and Niles-Weed (2021), based on entropic optimal transport, and show in the semi-discrete setting that it converges at the minimax-optimal rate $n^{-1/2}$, independent of dimension. Other standard map estimation techniques both lack finite-sample guarantees in this setting and provably suffer from the curse of dimensionality. We confirm these results in numerical experiments, and provide experiments for other settings, not covered by our theory, which indicate that the entropic estimator is a promising methodology for other discontinuous transport map estimation problems.