[Paper Review] Entropic estimation of optimal transport maps

Aram-Alexandre Pooladian, Jonathan Niles‐Weed|arXiv (Cornell University)|Sep 24, 2021

Markov Chains and Monte Carlo Methods49 references6 citations

TL;DR

This paper proposes a computationally efficient estimator for optimal transport maps between probability distributions using entropic regularization and Sinkhorn's algorithm. By taking the barycentric projection of the entropic coupling, the method achieves statistical consistency with finite-sample guarantees, offering a scalable alternative to prior methods that are slow in high dimensions or with large datasets.

ABSTRACT

We develop a computationally tractable method for estimating the optimal map between two distributions over $\mathbb{R}^d$ with rigorous finite-sample guarantees. Leveraging an entropic version of Brenier's theorem, we show that our estimator -- the \emph{barycentric projection} of the optimal entropic plan -- is easy to compute using Sinkhorn's algorithm. As a result, unlike current approaches for map estimation, which are slow to evaluate when the dimension or number of samples is large, our approach is parallelizable and extremely efficient even for massive data sets. Under smoothness assumptions on the optimal map, we show that our estimator enjoys comparable statistical performance to other estimators in the literature, but with much lower computational cost. We showcase the efficacy of our proposed estimator through numerical examples, even ones not explicitly covered by our assumptions. By virtue of Lepski's method, we propose a modified version of our estimator that is adaptive to the smoothness of the underlying optimal transport map. Our proofs are based on a modified duality principle for entropic optimal transport and on a method for approximating optimal entropic plans due to Pal (2019).

Motivation & Objective

• Address the computational intractability of existing optimal transport map estimators in high dimensions or with large sample sizes.
• Develop a tractable, scalable method for estimating optimal transport maps using entropic regularization and Sinkhorn's algorithm.
• Provide finite-sample statistical convergence guarantees for the proposed estimator under smoothness assumptions on the true map.
• Bridge the gap between theoretical optimal transport and practical machine learning applications by enabling efficient, parallelizable computation on large-scale data.
• Establish a canonical connection between the barycentric projection of the entropic plan and Brenier’s theorem, enabling out-of-sample extensions.

Proposed method

• Leverage entropic regularization of the optimal transport problem to transform the original Monge problem into a computationally tractable form.
• Use Sinkhorn's algorithm to efficiently compute the entropically regularized coupling between empirical measures derived from i.i.d. samples.
• Apply the barycentric projection to the entropic coupling to obtain a map estimator that maps input points to their conditional expectations under the coupling.
• Establish a theoretical link between the barycentric projection and the gradient of the dual potential in entropic optimal transport, generalizing Brenier’s theorem.
• Employ a modified duality principle and approximation techniques from Pal (2019) to analyze the statistical properties of the estimator.
• Use chaining and covering number arguments to bound the empirical process and derive finite-sample risk bounds.

Experimental results

Research questions

• RQ1 Can we construct a statistically consistent optimal transport map estimator that is computationally efficient for high-dimensional or large-scale datasets?
• RQ2 Does the barycentric projection of the entropic coupling provide a valid and consistent estimator of the true optimal transport map?
• RQ3 What are the finite-sample convergence rates of such an estimator under smoothness assumptions on the true map?
• RQ4 How does the computational efficiency of this method compare to existing state-of-the-art estimators in terms of time complexity and parallelizability?
• RQ5 Can the barycentric projection be interpreted as a natural extension of classical optimal transport theory under entropic regularization?

Key findings

• The proposed estimator achieves a convergence rate of $\mathbb{E}\|\hat{T}-T_0\|_{L^2(P)}^2 \lesssim n^{-\frac{\alpha+1}{2(d'+\alpha+1)}}\log n$ for $d' = 2\lceil d/2 \rceil$, under $\mathcal{C}^\alpha$ smoothness of the inverse map $T_0^{-1}$ with $\alpha \in (1,3]$.
• The estimator is computationally efficient, with time complexity quadratic in the number of samples due to Sinkhorn's algorithm, and is highly parallelizable on GPUs.
• Finite-sample convergence guarantees are established for the first time for the barycentric projection of the entropic coupling, a method previously used only heuristically.
• The method outperforms existing estimators in both statistical accuracy and computational speed on numerical benchmarks, despite a slightly slower rate than minimax-optimal estimators.
• The barycentric projection is shown to correspond to the gradient of the dual potential in entropic optimal transport, providing a theoretical foundation and enabling out-of-sample predictions.
• The analysis relies on a modified duality principle and covering number bounds for Hölder spaces, with constants depending on $\varepsilon^{-d/2}$, reflecting the trade-off between regularization and approximation error.