Statistical Convergence Rates of Optimal Transport Map Estimation between General Distributions

TL;DR

This paper advances the understanding of optimal transport map estimation by deriving convergence rates under less restrictive conditions, introducing new estimators, and developing scalable neural network algorithms with empirical validation.

Contribution

It broadens the theoretical framework for OT map estimation by removing strong assumptions and introduces practical algorithms for scalable estimation.

Findings

Established non-asymptotic convergence rates without restrictive assumptions.

Introduced a sieve plug-in estimator applicable to common distributions.

Developed neural network algorithms with demonstrated effectiveness.

Abstract

This paper studies the convergence rates of optimal transport (OT) map estimators, a topic of growing interest in statistics, machine learning, and various scientific fields. Despite recent advancements, existing results rely on regularity assumptions that are very restrictive in practice and much stricter than those in Brenier's Theorem, including the compactness and convexity of the probability support and the bi-Lipschitz property of the OT maps. We aim to broaden the scope of OT map estimation and fill this gap between theory and practice. Given the strong convexity assumption on Brenier's potential, we first establish the non-asymptotic convergence rates for the original plug-in estimator without requiring restrictive assumptions on probability measures. Additionally, we introduce a sieve plug-in estimator and establish its convergence rates without the strong convexity assumption on Brenier's potential, enabling the widely used cases such as the rank functions of normal or t-distributions. We also establish new Poincar\'e-type inequalities, which are proved given sufficient conditions on the local boundedness of the probability density and mild topological conditions of the support, and these new inequalities enable us to achieve faster convergence rates for the Donsker function class. Moreover, we develop scalable algorithms to efficiently solve the OT map estimation using neural networks and present numerical experiments to demonstrate the effectiveness and robustness.