Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Entropic estimation of optimal transport maps

Aram-Alexandre Pooladian, Jonathan Niles-Weed

TL;DR

This work introduces a computationally tractable estimator for OT maps between high-dimensional distributions by leveraging entropic regularization and Brenier-type interpretations. The estimator uses the barycentric projection of the entropic plan and is computed via Sinkhorn's algorithm, enabling parallel GPU implementation and scalability to massive data. The authors provide finite-sample guarantees under standard regularity assumptions, derive rates for one- and two-sample settings, and develop an adaptive Lepski-based procedure to handle unknown smoothness. Although the achieved rates are slightly below the minimax benchmark, the method demonstrates strong practical performance and efficiency, outperforming competing tractable estimators in many experiments. The work highlights a promising direction for scalable, statistically principled transport-map estimation with broad applicability in statistics and ML.

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).

Entropic estimation of optimal transport maps

TL;DR

This work introduces a computationally tractable estimator for OT maps between high-dimensional distributions by leveraging entropic regularization and Brenier-type interpretations. The estimator uses the barycentric projection of the entropic plan and is computed via Sinkhorn's algorithm, enabling parallel GPU implementation and scalability to massive data. The authors provide finite-sample guarantees under standard regularity assumptions, derive rates for one- and two-sample settings, and develop an adaptive Lepski-based procedure to handle unknown smoothness. Although the achieved rates are slightly below the minimax benchmark, the method demonstrates strong practical performance and efficiency, outperforming competing tractable estimators in many experiments. The work highlights a promising direction for scalable, statistically principled transport-map estimation with broad applicability in statistics and ML.

Abstract

We develop a computationally tractable method for estimating the optimal map between two distributions over 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).

Paper Structure

This paper contains 24 sections, 26 theorems, 194 equations, 4 figures.

Table of Contents

  1. Introduction
  2. Background on optimal transport theory
  3. Optimal transport under the quadratic cost
  4. Entropic optimal transport under the quadratic cost
  5. Estimator and main results
  6. One-sample estimates
  7. Two-sample estimates
  8. Adaptive estimation
  9. Computational aspects
  10. Estimator complexities from prior work
  11. Computational complexity of the Entropic Map
  12. Empirical performance
  13. Comparison to a tractable minimax estimator
  14. On estimating non-smooth transport maps
  15. Parallel estimation on massive data sets
  16. ...and 9 more sections

Key Result

Theorem 1

Bre91 Let $P \in \mathcal{P}_{ac}(\Omega)$ and $Q \in \mathcal{P}(\Omega)$. Then

Figures (4)

  • Figure 1: Visualization of $\hat{T}_\varepsilon$ and $T_0(x)$ in 2 dimensions.
  • Figure 2: Dashed lines are our estimator, solid lines are $\hat{T}^{\text{1NN}}$, and $T_0(x) = \exp(x)$
  • Figure 3: Dashed lines are our estimator, solid lines are $\hat{T}^{\text{1NN}}$, and $T_0(x) = 2|x_1| + x$
  • Figure 4: Performance of a parallel implementation of our estimator on large data sets.

Theorems & Definitions (50)

  • Theorem 1: Brenier's Theorem
  • Remark 1
  • Proposition 1
  • Theorem 2
  • Proposition 2
  • proof
  • Remark 2
  • Theorem 3
  • Theorem 4
  • Remark 3
  • ...and 40 more