PaperlayerView exact revision on arXiv

Source-linked AI summary

Entropic estimation of optimal transport maps

Aram-Alexandre Pooladian, Jonathan Niles-Weed

arXiv:2109.12004v3math.STstat.MLstat.TH

TL;DR

With only independent samples from two distributions, estimating an existing optimal transport map remains the central problem. The paper estimates it through the barycentric projection of an empirical entropic optimal plan, with a dual interpretation tied to Sinkhorn computation. The authors report strong empirical performance and first finite-sample guarantees for this proposal, while noting limits in assumptions and rate optimality.

  • Problem

    Given independent samples from P and Q, the paper studies how to estimate an optimal transport map T0 between them.

    Source: §1 · p. 1

  • Method

    The estimator is the barycentric projection of the empirical entropic optimal coupling, interpreted through an entropic analogue of Brenier’s theorem and computed by solving the entropic dual.

    Source: §1 · p. 2

  • Results

    The estimator empirically outperforms other proposed estimators in computational and statistical performance, and the paper gives the first finite-sample convergence guarantees for this proposal.

    Source: §1 · p. 2, §8 · p. 20

  • Takeaways & Limitations

    The dual interpretation gives a canonical extension to out-of-sample points, and computing the estimator is no more costly than solving the entropic problem.

    Source: §1 · p. 2

    The analysis requires compact support and strong a priori bounds without clear noncompact analogues; in the high-dimensional, α→1 regime, its rate is about a factor of two worse in the exponent than the minimax rate.

    Source: §3 · p. 7

Abstract

from arXiv · show

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

1 Introduction

Given independent samples from two distributions, the paper asks whether Sinkhorn’s practical entropic optimal transport can yield an optimal-map estimator with convergence guarantees. It proposes a barycentric-projection estimator and reports strong empirical performance alongside its first finite-sample guarantees for this proposal.

1 Introduction

  • The goal is to estimate an optimal transport map from independent samples, while existing statistically optimal estimators can be impractical in high dimensions or costly and difficult to parallelize.The cited alternatives include a gridding method with dimension-exponential cost and plug-in methods requiring at least O(n^3) time.Source: §1 · pp. 1–2
  • The estimator empirically outperforms other proposed estimators in computational and statistical performance, and the paper provides the first finite-sample convergence guarantees for this proposal.Its rate is stated in the introduction, but the supplied passage does not preserve the full expression.Source: §1 · p. 2
  • The estimator is the barycentric projection of the empirical entropic optimal coupling, which the authors identify as the gradient of the dual solution and a canonical extension to out-of-sample points.Because Sinkhorn’s algorithm solves the entropic dual, computing the estimator is no more costly than solving the entropic problem.Source: §1 · p. 2
  • The analysis compares the optimal entropic coupling with a Gaussian-conditional coupling characterized by the unregularized optimal map, using a modified duality principle and an upper bound on entropic transport.The proof strategy builds on a method of Pal (2019).Source: §1 · p. 3

Paper outline

  • The paper outline assigns background to Section 2, estimator and main results to Section 3, statistical bounds to Sections 4–5, smoothness adaptation to Section 6, and computation and experiments to Section 7.Source: Paper outline · p. 3

Notation

  • The notation section defines support, convex duality, KL divergence, differential entropy, conditional density, smoothness classes, total variation, Euclidean balls, and asymptotic constants.It also states that the empirical-process suprema used in the proofs are measurable because the processes are separable.Source: Notation · p. 3

2 Background on optimal transport theory

The background reviews classical optimal transport under quadratic cost and its entropic regularization. It connects map existence and dual potentials to entropic plans, modified duality, and bounds relating regularized and unregularized transport.

2.1 Optimal transport under the quadratic cost

  • When the source measure has a density, Brenier’s theorem gives an optimal map as the gradient of a convex potential; if the target also has a density, the reverse map is the gradient of the conjugate potential.If the source lacks a density, the Monge problem may have no solution, while the coupling relaxation admits a minimizer for measures with finite second moments.Source: §2.1 · pp. 3–4
  • The classical transport problem has a dual formulation whose optimal potentials link to the optimal map, and entropic transport adds regularization to the coupling problem.The entropic primal and dual solutions satisfy an optimality relation.Source: §2.1 · p. 4, §2.2 · p. 4

2.2 Entropic optimal transport under the quadratic cost

  • The proofs use a modified entropic duality relation that takes its supremum over a larger function class than the standard dual.The authors identify this extension as a consequence of Fenchel’s inequality and note that they had not found it stated explicitly in the literature.Source: §2.2 · p. 5
  • Prior work studied convergence of entropic plans, potentials, and entropic transport values as regularization vanishes; the paper supplies and briefly proves a bound used in its analysis.The background situates that bound alongside first-order and second-order convergence results.Source: §2.2 · p. 5
  • For bounded-density compactly supported measures, the stated bound relates entropic transport to squared Wasserstein distance, measure entropies, and integrated Fisher information along the Wasserstein geodesic.Source: §2.2 · p. 5

3 Estimator and main results

The paper estimates transport maps using entropic optimal transport, defining the map through an entropic potential and computing it from empirical distributions. Under regularity assumptions, it gives finite-sample guarantees and an estimator that can adapt to unknown smoothness.

3 Estimator and main results

  • The entropic map is motivated as replacing the optimal potential in Brenier’s theorem with its entropic counterpart, and it need not push P to Q exactly.The paper also gives an inverse-map estimator, with guarantees obtained by interchanging the roles of P and Q in the assumptions.Source: §3 · pp. 6–7
  • The estimator is the entropic map between empirical measures, obtained from an optimal entropic plan and computable using its potential from Sinkhorn’s algorithm.Once the optimal entropic potential is found, the map can be evaluated in linear time.Source: §3 · p. 6
  • The guarantees are proved under compact-support and regularity assumptions; the authors note that their proofs rely on strong a priori bounds without clear analogues for non-compact measures.The stated assumptions include bounded densities and uniform bounds on the Hessian of the optimal-map potential.Source: §3 · p. 7
  • Under (A1)–(A3), the empirical entropic map has an L2(P) convergence guarantee with regularization and rate depending on dimension and smoothness capped at 3.In the regime d→∞ and α→1, the stated rate falls short of the minimax rate by approximately a factor of 2 in the exponent, though the estimator is reported competitive in numerical experiments.Source: §3 · p. 7
  • The one-sample entropic map is shown to approach the optimal map in L2(P), with its rate depending on dimension and smoothness capped at 3.The theorem gives a rate for Tε,n between P and Qn under (A1)–(A3).Source: §3 · p. 8

4 One-sample estimates

The one-sample analysis relates the entropic map from P to an empirical target Qn to the optimal map. It uses a comparison with conditional densities and derives a finite-sample error bound.

4 One-sample estimates

  • The analysis compares the optimal entropic plan for P and Qn with a measure built from conditional densities, using a duality-based inequality.The proof then bounds terms in this comparison using Laplace expansion, Sinkhorn-divergence convergence, and an empirical-process argument.Source: §4 · pp. 9, 11
  • The analysis bounds E[Δ1 + Δ2 + Δ3] by ε^((ᾱ−1)/2) + (ε^-1 + ε^-d/2) log(n)n^-1/2 + ε^-d/2 n^-1/2.Here ᾱ is α capped at 3.Source: §4 · p. 11
  • Theorem 4 also yields a population-level estimate between Tε and T0, which the authors identify as potentially independently interesting.The supplied passage does not state the estimate’s explicit bound.Source: §4 · p. 12

5 Two-sample estimates

The two-sample analysis compares the entropic map estimated from empirical source and target measures with the one-sample map. Its bound concerns their expected squared L2(P) difference.

5 Two-sample estimates

  • The two-sample analysis compares Tε,(n,n) with Tε,n while keeping Qn fixed, simplifying the duality argument.The proof uses an analogue of an earlier proposition that does not require the full force of assumptions (A1)–(A3).Source: §5 · p. 12
  • For ε≤1, E∥Tε,n − Tε,(n,n)∥²_L2(P) is bounded by a constant multiple of ε^-d/2 log(n)n^-1/2.The passage states the intermediate bound with (ε^-1 + ε^-d/2) log(n)n^-1/2 before simplifying it.Source: §5 · p. 13
  • The two-sample result is established by choosing a sufficiently large constant in the comparison argument.The proof uses Hoeffding’s inequality before selecting the comparison function.Source: §5 · p. 13

6 Adaptive estimation

The adaptive procedure selects the regularization parameter using sample splitting and Lepski’s method. Under the stated conditions, its risk matches Theorem 3 up to logarithmic factors despite unknown smoothness.

6 Adaptive estimation

  • The data-driven choice of regularization uses Lepski’s method to adapt the estimator to the unknown smoothness parameter α.The optimal regularization parameter in Theorems 3 and 4 depends on n, d, and α.Source: §6 · p. 13
  • The procedure calibrates ε by sample splitting, using an independent copy of the data and a discrete set of smoothness candidates.The grid covers s from 2+ι to 4, where s=α+1.Source: §6 · pp. 13–14
  • Under Theorem 6’s conditions and for sufficiently large K, the selected estimator’s L2(P) risk matches Theorem 3 up to logarithmic factors.The choice uses ε=εŝ as prescribed by the stated selection rule.Source: §6 · p. 14

7 Computational aspects

This section compares the entropic map’s computational cost with prior estimators and reports theoretical and empirical performance, including parallel runs on large datasets.

7.1 Estimator complexities from prior work

  • Earlier estimators require runtime exponential in d or cubic in n, whereas the entropic estimator is nearly quadratic in n.The wavelet estimator’s grid-based Legendre transform has exponential dimension dependence; density-estimation approaches also face exponential dependence, while 1NN uses the Hungarian algorithm.Source: §7.1 · pp. 14–15

7.2 Computational complexity of the Entropic Map

  • Sinkhorn updates entropic potentials iteratively, with each iteration taking O(n^2) time on the empirical supports.The finite-iteration theorem shows that the resulting estimator remains acceptable once k is sufficiently large, even though g^(k) is not an entropic potential for the original problem.Source: §7.2 · pp. 15–16
  • An estimator with the same rate as Theorem 3’s can be computed in ˜O(n^(2+7/(d+ᾱ+1))) = n^(2+o_d(1)) time.The result assumes (A1)–(A3), the specified ε, and at least k ≳ n^(7/(d+ᾱ+1)) log n iterations.Source: §7.2 · p. 16

7.3 Empirical performance

  • The empirical setup samples from P = [−1, 1]^d and Q = (T0)♯P, with T0 formed by applying a monotone scalar function coordinate-wise.The experiments use two Sinkhorn implementations, including a GPU-optimized KeOps implementation, with log-domain stabilization.Source: §7.3 · p. 17

7.3.1 Comparison to a tractable minimax estimator

  • At moderate dimensions, the entropic map consistently outperforms the 1NN estimator in both MSE and runtime; in d = 2, their MSEs are comparable.For large n, the entropic estimator’s error deteriorates, reflecting its slightly sub-optimal estimation rate.Source: §7.3.1 · p. 19

7.3.2 On estimating non-smooth transport maps

  • For a non-smooth map, the entropic estimator outperforms 1NN in both runtime and MSE despite lacking rigorous guarantees for this setting.The experiment’s pushforward distribution has disconnected support separated along the first coordinate.Source: §7.3.2 · p. 19

7.3.3 Parallel estimation on massive data sets

  • With n = 10^4 and d = 10, the parallel GPU implementation takes roughly a third of a second to optimize.Sinkhorn is parallelizable, unlike the 1-Nearest Neighbor estimator and other transport map estimators the authors know.Source: §7.3.3 · p. 19

8 Conclusion and future directions

The paper presents a first finite-sample analysis of an entropic estimator for optimal transport maps. Its rates fall short of minimax optimality, while experiments show strong computational and empirical performance; the authors identify possible directions for improving the theory.

8 Conclusion and future directions

  • The paper gives the first finite-sample analysis of an entropic optimal transport map estimator, which is fast, parallelizable, and empirically outperforms a leading tractable statistically optimal proposal.Its theoretical rates do not reach minimax optimality.Source: §8 · p. 20
  • The authors conjecture that a properly tuned entropic estimator may attain the minimax rate when α is small and that the assumptions could be loosened.They specifically suggest that stringent conditions on densities or support might not be necessary.Source: §8 · p. 20

A Second-order error estimate

The proof of Theorem 2 uses dynamic formulations of the ordinary and entropic transport objectives, then evaluates the entropic formulation at the ordinary problem’s minimizers.

A Second-order error estimate

  • The proof strategy hinges on the dynamic formulations of W2(P, Q) and Sε(P, Q), beginning with the former.The ordinary formulation uses a continuity equation with endpoint densities p and q.Source: A Second-order error estimate · p. 20
  • Plugging the minimizers of the ordinary dynamic formulation into the entropic formulation yields the result of Eq. (11) by optimality.The entropic dynamic formulation is stated for measures with bounded densities and compact support.Source: A Second-order error estimate · p. 21

B Laplace’s method proof

Appendix B bounds a strongly convex integral using a local expansion around its minimum and separate estimates inside and outside a neighborhood. The following proof fragment establishes a variational inequality by pointwise maximization.

B Laplace’s method proof

  • The integral approximation relies on expanding f around its global minimum, where twice-differentiability gives quadratic local behavior resembling a Gaussian integral for small ε.For α ≤ 2, the proof bounds the inner and outer contributions separately; for α > 2, it expands to third order.Source: B Laplace’s method proof · pp. 21–23
  • For f ∈ C^(α+1) that is m-strongly convex and M-smooth, Theorem 8 bounds I(ε) for α > 1 using constants depending on regularity, dimension, and derivative norm.The proof centers the global minimum at zero and splits the integral into two parts.Source: B Laplace’s method proof · p. 21

C.1 Proof of Proposition 1

  • The proof of Proposition 1 reduces to an inequality holding for all a ≥ 0 and b ∈ R, obtained by maximizing a concave function of b at log a.The other direction follows by choosing η(x, y) = f(x) + g(y) and using Eq. (8).Source: §C.1 · p. 24

C.2 Proof of Proposition 4

The appendix proves Proposition 4 through a more general result and supplies proofs of technical lemmas from Section 4. The arguments use entropic dual potentials and bounds on their regularity.

C.2 Proof of Proposition 4

  • Proposition 5 considers a source measure supported on Ω whose total variation distance from P_n is at most δ, with Q_n as the other empirical measure.Its comparison involves the optimal entropic plan for P and Q_n and the optimal plan for the perturbed source measure.Source: §C.2 · p. 24
  • The proof bounds the comparison using optimal entropic potentials for the perturbed source and Q_n, invoking their boundedness and dual optimality.It also uses that the perturbed plan has a density with respect to P ⊗ Q_n.Source: §C.2 · p. 25
  • For compactly distributed measures, derivatives of the optimal entropic potential are bounded by C_s,d,K(1 + ε^(1−s)) on compact sets, which is used with Lemma 7 to conclude the proof.The proof applies Lemma 7 with s = d/2 on a ball containing Ω.Source: §C.2 · p. 25

C.3 Proofs from Section 4

  • The Section 4 proofs use concentration around the mean of a random variable with density q_x and apply the resulting inequality with v = h(x), then integrate over P.The supplied proof passages also invoke Young’s and Jensen’s inequalities.Source: §C.3 · p. 26
  • The Lemma 4 proof uses a local expansion and makes the contribution outside a ball around x* negligible by choosing the constant in τ sufficiently large.The proof states that this contribution can be made smaller than ε.Source: §C.3 · pp. 26–27

D Supplementary results

The supplementary results establish empirical-process and entropic-potential bounds used in the paper’s proofs. The proof of Theorem 6 analyzes smoothness selection against an oracle risk rate.

D Supplementary results

  • Lemma 7 considers s-Hölder-smooth functions with Hölder norm bounded by M, for convex compact K ⊆ R^d, where s ≥ d/2, and probability measure ν supported in K with corresponding empirical measure ν_n.Source: D Supplementary results · p. 29
  • For compactly supported P and Q, derivatives of their optimal entropic dual potentials through order s are bounded on compact sets by C_s,d,K(1 + ε^(1−s)), with the constant independent of ε.Source: D Supplementary results · p. 30

E Proof of Theorem 6

  • The smoothness grid runs from s_min = 2 + ι to s_N = 4 with increments of order (log n)^−1, and ψ_n(s) = (n/log n)^−s/(2(d+s)).Source: E Proof of Theorem 6 · p. 31
  • Theorem 6’s proof compares the risk of the selected estimator with the oracle rate ψ_n(s_i), considering cases where the selected smoothness is above or below the underlying smoothness.The proof describes the lower-smoothness case as more technical.Source: E Proof of Theorem 6 · pp. 31–32
Loading 2109.12004v3…