The theory of optimal transport (OT) defines a natural geometry on the space of probability measures (Santambrogio, 2015;Villani, 2009) and has become ubiquitous in modern data-driven tasks. In this area, optimal transport maps are a central object of study: suppose P and Q are two probability distributions with finite second moments, with P having a density with respect to the Lebesegue measure on R d . Then, Brenier's theorem (see Section 2.1) states that there exists a convex function φ 0 whose gradient defines a unique optimal transport map between P and Q. This map is optimal in the sense that it minimizes the following objective function:
where T (P, Q) := {T : R d → R d | X ∼ P, T (X) ∼ Q} is the set of transport maps between P and Q. The optimal value of the objective function in Equation ( 1) is called the (squared) 2-Wasserstein distance, written explicitly as
though a more general formulation is available (see Section 2.1). Computing or approximating S 0 (P, Q) as well as ∇φ 0 has found use in several academic communities, such as economics (Carlier et al., 2016;Chernozhukov et al., 2017;Gunsilius & Xu, 2021;Torous et al., 2021), computational biology (Bunne et al., 2021;2022;Demetc ¸i et al., 2022;Lübeck et al., 2022;Moriel et al., 2021;Schiebinger et al., 2019;Yang et al., 2020), and computer vision (Feydy et al., 2017;Solomon et al., 2015;2016), among many others.
Practitioners seldom have access to P or Q, but instead have access to i.i.d. samples X 1 , . . . , X n ∼ P and Y 1 , . . . , Y n ∼ Q. On the basis of these samples, practitioners face both computational and statistical challenges when estimating ∇φ 0 . From a theoretical perspective, the statistical task of estimating optimal transport maps has attracted much interest in the last few years (Deb et al., 2021;Divol et al., 2022;Ghosal & Sen, 2022;Hütter & Rigollet, 2021;Manole et al., 2021;Muzellec et al., 2021;Pooladian & Niles-Weed, 2021).
The first finite-sample analysis of this problem was performed by Hütter & Rigollet (2021), who proposed an estimator for ∇φ 0 under the assumption that φ 0 is s + 1-times continuously differentiable, for s > 1. They showed that a wavelet-based estimator φW satisfies
and that this rate is minimax optimal up to logarithmic factors. Their analysis requires that P and Q have bounded densities with compact support Ω ⊆ R d , and that φ 0 be both strongly convex and smooth. Implementing the estimator φW is computationally challenging even in moderate dimensions, and is practically infeasible for d > 3. Follow up work has proposed alternative estimators which improve upon φW either in computational efficiency or in the generality in which they apply. Though these subsequent works go significantly beyond the setting considered by Hütter & Rigollet (2021), none has eliminated the crucial assumption that φ 0 is smooth, i.e., that the transport map ∇φ 0 is Lipschitz.
We highlight two estimators proposed in this line of work that are particularly practical. Manole et al. (2021) study the 1-Nearest Neighbor estimator T1NN . This estimator is obtained by solving the empirical optimal transport problem between the samples, which is then extended to a function defined on R d using a projection scheme; see Section 4 for more details. Given n samples from the source and target measures in R d , T1NN has a runtime of O(n 3 ) via the Hungarian Algorithm (see Peyré & Cuturi, 2019, Chapter 3), and, for d ≥ 5, achieves the rate
whenever the optimal Brenier potential φ 0 is smooth and strongly convex, and under mild regularity conditions on P . In another work, Pooladian & Niles-Weed (2021) conducted a statistical analysis of an estimator originally proposed by Seguy et al. (2018) based on entropic optimal transport. The efficiency of Sinkhorn's algorithm for large-scale problems (Cuturi, 2013;Peyré & Cuturi, 2019) makes this estimator attractive from a computational perspective, and Pooladian & Niles-Weed (2021) also give statistical guarantees, though these fall short of being minimax-optimal.
Despite this progress, none of the aforementioned results can be applied in situations where ∇φ 0 is not Lipschitz. And in practice, even requiring the continuity of the transport map can be far too stringent. It is indeed too much to hope for that an underlying data distribution (e.g. over the space of images) has one single connected component; this is supported by recent work that stipulates that the underlying data distribution is the union of disjoint manifolds of varying intrinsic dimension (Brown et al., 2022). In such a setting, the transport map ∇φ 0 will not be continuous, demonstrating the need of considering the problem of the statistical estimation of discontinuous transport maps to get closer to real-world situations.
As a first step, we choose to focus on the case where the target distribution Q = J j=1 q j δ yj is discrete while the source measure P has full support, often called the semidiscrete setting in the optimal transport literature. In this setting, the optimal transport map ∇φ 0 is constant over regions known as Laguerre cells (each cell corresponding to a different atom of the discrete measure), while displaying discontinuities on their boundaries (see Section 2.1.1 for more details). Figure 1 provides such an example. Semidiscrete optimal transport therefore provides a natural class of discontinuous transport maps. We focus on this setting for two reasons. First, it has garnered a lot of attention in recent years, in both computational and theoretical circles (see, e.g., Altschuler et al., 2022;Chen et al., 2022;Mérigot et al., 2021), due in particular to its connection with the quantization problem (Graf & Luschgy, 2007). Second, the semi-discrete setting is intriguing from a statistical perspective: existing results show that statistical estimation problems involving semi-discrete optimal transport can escape the curse of dimensionality (del Barrio & Loubes, 2019;del Barrio et al., 2022a;Forrow et al., 2019;Hundrieser et al., 2022). For example, Hundrieser et al. (2022, Theorem 3.2) show that if P n and Q n are empirical measures consisting of i.i.d. samples from P and Q, then the semi-discrete assumption implies
These results offer the tantalizing possibility that semidiscrete transport maps can be estimated at the rate n -1/2 , in sharp contrast to the dimension-dependent rates obtained in bounds such as (2). However, the optimal rates of estimation for semi-discrete transport maps are not known, and no estimators with finite-sample convergence guarantees exist.
We show that the computationally efficient estimator Tε based on entropically regularized optimal transport, originally studied in (Pooladian & Niles-Weed, 2021;Seguy et al., 2018), provably estimates discontinuous semi-discrete optimal transport maps at the optimal rate. More precisely, our contributions are the following:
1. For Q discrete and P with full support on a compact, convex set, we show that Tε achieves the following dimension-independent convergence rate to the optimal transport map (see Theorem 3.1)
when the regularization parameter ε ≍ n -1/2 . We further show (Proposition 4.1) that this rate is minimax optimal.
2. As a by-product of our analysis, we give new parametric rates of convergence to the entropic Brenier map T ε , a result which improves exponentially on prior work in the dependence on ε (see Theorem 3.7 and Remark 3.8).
3. Our proof technique requires several new results, including a novel stability bound for the entropic Brenier maps (Proposition 3.9), and a new stability result for the entropic dual Brenier potentials in the semi-discrete case (Proposition 3.11).
4. We show that, unlike Tε , the 1-Nearest-Neighbor estimator is provably suboptimal in the semi-discrete setting (see Proposition 4.2) by exhibiting a discrete measure Q such that the risk suffers from the curse of dimensionality:
In Section 4, we verify our theoretical findings on synthetic experiments. We also show by simulation that the entropic estimator appears to perform well even outside the semi-discrete setting, suggesting it as a promising choice for estimating other types of discontinuous maps.
The Euclidean ball centered at a with radius r > 0 is written as B(a; r).
for the variance of f with respect to ρ.
We define P(Ω) to be the space of probability measures whose support lies in a compact subset Ω ⊆ R d . If a probability measure P has a density with respect to the Lebesgue measure on R d with support Ω ⊆ R d , then we write P ∈ P ac (Ω).
For two probability measures P, Q ∈ P(Ω), we define the (squared) 2-Wasserstein distance to be (Kantorovitch, 1942)
where π ∈ Γ(P, Q) ⊆ P(Ω × Ω) such that for any event A,
We call Γ(P, Q) the set of couplings between P and Q. In this work, we focus on the squared-Euclidean cost but Equation (4) is well-defined for convex, lower-semicontinuous costs; see (Santambrogio, 2015;Villani, 2009) for more information on optimal transport under general costs.
Equation ( 4) is a convex optimization problem on the space of joint measures, and a minimizer, denoted π 0 , always exists; we call π 0 an optimal plan from P to Q. Moreover, Equation (4) possesses the following dual formulation,
where M 2 (P ) := ∥x∥ 2 dP (x) (similarly for M 2 (Q)) and the functions (φ, ψ)
As with the primal formulation, the infimum in Equation ( 5) is attained at functions (φ 0 , ψ 0 ). These minimizers are called (optimal) Brenier potentials. In particular, at optimality, we have that these Brenier potentials are convex conjugates of one another, i.e. the Legendre transform of one of the potentials gives the other:
and vice-versa.
Apart from these two formulations of optimal transport under the squared-Euclidean cost, there exists a third, known as the Monge problem:
where T (P, Q) is the set of admissible transport maps, i.e. for X ∼ P , T (X) ∼ Q. This optimization problem is non-convex in T , and a solution is not always guaranteed to exist for arbitrary P and Q.
The following theorem unifies these three formulations of optimal transport under the squared-Euclidean cost:
Theorem 2.1 (Brenier's theorem;Brenier, 1991). Let P ∈ P ac (Ω) and let Q ∈ P(Ω), then 1. the solution to Equation (7) exists and is of the form T 0 = ∇φ 0 , where φ 0 solves Equation (5)
2. π 0 is also uniquely defined as
When we want to place emphasis on the underlying measures, we will write φ 0 = φ P →Q 0 , ψ 0 = ψ P →Q 0 and T 0 = T P →Q 0 .
In optimal transport, the semi-discrete setting refers to the case where P has as density with respect to the Lebesgue measure on R d , and Q is a discrete measure supported on points. The following theorem characterizes the optimal transport map in this situation, which exhibits a particular structure compared to the general results in the previous section. Let [J] = {1, . . . , J}.
Proposition 2.2 (Aurenhammer et al., 1998). If P ∈ P ac (Ω) and Q is a discrete measure supported on the points y 1 , . . . , y J , then the optimal transport map ∇φ 0 is given by
where ψ 0 is the dual to φ 0 in the sense of Equation (6).
Here, the optimal dual Brenier potential ψ 0 can be identified with a vector in R J , defined by the number of atoms, and the optimal Brenier potential is consequently given by
{⟨x, y j ⟩ -ψ 0 (y j )} .
Although φ 0 is not differentiable, only subdifferentiable, we still use the gradient notation as ∇φ 0 is well-defined P -almost everywhere.
The map ∇φ 0 partitions the space into J convex polytopes L j := ∇φ -1 0 ({y j }) called Laguerre cells; recall Figure 1. From this definition, it is clear that for a given x ∈ L j , x → ∇φ 0 (x) = y j is the optimal transport mapping. The difficulty in finding this map lies in determining the cells L j , or equivalently the dual variables ψ 0 (y j ).
Entropic regularization was introduced to both optimal transport and machine learning communities in the seminal paper by Cuturi (2013), allowing approximate optimal transport distances to be computed at unprecedented speeds. Entropic optimal transport (EOT) is defined as the following regularized version of Equation ( 4): for ε > 0 S ε (P, Q) := min
where KL(µ∥ν) = log dµ dν dµ when µ ∈ P(Ω) is absolutely continuous with respect to ν ∈ P(Ω). This speedup is due to the elegant connection of (9) to Sinkhorn's algorithm; we refer the interested reader to (Peyré & Cuturi, 2019, Chapter 4) for more information. The computational tractability of S ε compared to S 0 when dealing with many samples lends itself to being a central object of study in its own right (see, e.g., Chizat et al., 2020;Genevay et al., 2019;Gonzalez-Sanz et al., 2022;Mena & Niles-Weed, 2019;Rigollet & Stromme, 2022). Equation ( 9) admits the following dual formulation, which is now an unconstrained optimization problem (Genevay, 2019;Marino & Gerolin, 2020)
where (φ, ψ) ∈ L 1 (P )×L 1 (Q). When P and Q have finite second moments, Equation ( 9) admits a unique minimizer, π ε and we have the existence of minimizers to Equation ( 10), which we denote as (φ ε , ψ ε ). We call π ε the entropic optimal plan and (φ ε , ψ ε ) are called entropic Brenier potentials.
The following optimality relation further relates these primal and dual solutions (Csiszár, 1975):
As a consequence, the following relationship holds at optimality:
and, moreover, we can define versions of φ ε and ψ ε such that the following relationships hold (see Mena & Niles-Weed, 2019;Nutz & Wiesel, 2022) over all x ∈ R d and y ∈ R d , respectively:
which are smoothed version of the Legendre transform, see Appendix A for details. In what follows, we always assume that we have selected φ ε and ψ ε so that these identities hold.
If (X, Y ) ∼ π ε , we may define the conditional probability
The barycentric projection of the optimal entropic coupling π ε , or entropic Brenier map, is a central object of study in several works e.g. (del Barrio et al., 2022b;Goldfeld et al., 2022;Pooladian & Niles-Weed, 2021;Rigollet & Stromme, 2022), defined as
where π x ε is as in Equation ( 13). Note that this quantity is well defined for all x ∈ R d as long as the source and target measures have compact support; in particular, it applies to both discrete and continuous measures. The second equality follows from Equation ( 11) and the dominated convergence theorem. As in the unregularized case, we will write
when we want to emphasize on the dependency with respect to the underlying measures.
This particular barycentric projection was proposed as a tool for large-scale optimal transport by Seguy et al. (2018), but analyzed statistically for the first time by Pooladian & Niles-Weed (2021) as an estimator for the optimal transport map. We mention some of their results to highlight the differences with our new results for the semi-discrete setting in Section 3. First, they prove the following approximation result for T ε . Proposition 2.3 (Pooladian & Niles-Weed, 2021, Corollary 1). Let P, Q be compactly supported absolutely continuous measures on a compact set Ω ⊆ R d with densities p and q, that are bounded away from 0 and ∞. Assume that φ 0 is smooth and strongly convex, and that φ * 0 is at least C 3 . Then,
Their main statistical result is the following theorem: Proposition 2.4 (Pooladian & Niles-Weed, 2021, Theorem 3). Suppose the same assumptions as Proposition 2.3, and let P n and Q n denote the empirical measures of P and Q constructed from i.i.d. samples. Let Tε = T Pn→Qn ε denote the entropic Brenier map from P n to Q n and let T 0 = ∇φ 0 be the optimal transport map from
where d ′ = 2⌈d/2⌉.
Note that in particular the the rate of convergence of the entropic estimator critically depends on the ambient dimension d in the continuous-to-continuous case.
Characterizing the convergence of entropic objects (e.g. potentials, cost, plans) to their unregularized counterparts in the ε → 0 regime has been a topic of several works in recent years. Convergence of the costs S ε to S 0 with precise rates was investigated in (Chizat et al., 2020;Conforti & Tamanini, 2021;Pal, 2019). The works (Bernton et al., 2022;Carlier et al., 2017;Ghosal et al., 2022;Léonard, 2012) study the convergence of the minimizers π ε to π 0 under varying assumptions. Convergence of the potentials in a very general setting was established in (Nutz & Wiesel, 2022), though without a rate of convergence in ε. In the semidiscrete case, this gap was closed in (Altschuler et al., 2022) followed closely by (Delalande, 2022), which gave nonasymptotic rates. The Sinkhorn Divergence, a non-negative, symmetric version of S ε , was introduced in (Genevay et al., 2018), was statistically analysed in (Goldfeld et al., 2022) and also in (del Barrio et al., 2022b;Gonzalez-Sanz et al., 2022), and was connected to the entropic Brenier map in (Pooladian et al., 2022). The recent pre-print by (Rigollet & Stromme, 2022) proved parametric rates of estimation between the empirical entropic Brenier map and its population counterpart, though with an exponentially poor dependence on the regularization parameter (see Remark 3.8). Using covariance inequalities, the entropic Brenier potentials were used give a new proof of Caffarelli's contraction theorem; see (Chewi & Pooladian, 2022); this approach was recently generalized in (Conforti, 2022). Entropic optimal transport has also come into contact with the area of deep generative modelling through the following works (De Bortoli et al., 2021;Finlay et al., 2020), among others.
Let P n and Q n be the empirical measures associated with two n-samples from P and Q. We make the following regularity assumptions on P , already introduced by Delalande (2022).
(A) The measure P has a compact convex support Ω ⊆ B(0; R), with a density p satisfying 0 < p min ≤ p ≤ p max < ∞ for positive constants p min , p max and R.
For example, P can be the uniform distribution over Ω, or a truncated Gaussian distribution. Furthermore, we will need the following assumption on Q.
(B) The discrete probability measure Q = J j=1 q j δ yj is such that q j ≥ q min > 0 and y j ∈ B(0; R) for all j ∈ [J].
The goal of this section is to prove the following theorem: Theorem 3.1. Let P satisfy (A) and let Q satisfy (B). Let Tε = T Pn→Qn ε . Then, for ε ≍ n -1/2 and n large enough,
Remark 3.2. We remark that the hidden constants in Theorem 3.7 and related results depend on J, p min , p max , q min and R. Remark 3.3 (Fixing the support via rounding). At present, the entropic map need not necessarily map exactly to one of {y 1 , . . . , y J }. In fact, Tε :
where conv(A) is the convex hull for some set A. In turn, the support of the entropic map does not in general match that of Q. However, this can be readily fixed with a rounding scheme. We can replace our estimator by Tε which is obtained by mapping the output of Tε to its nearest neighbor in the support of Q -this projection step is easy to compute, given that we essentially know the support of Q via samples. By viewing this as a projection onto an appropriate set (namely, the set of transport maps with codomain equal to the support of Q), and applying the triangle inequality, it holds that
but Tε matches the support of Q.
Let T ε = T P →Q ε denote the entropic Brenier map associated to P and Q. Our proof relies on the following bias-variance decomposition:
Following the next two results (Theorem 3.4 and Theorem 3.7) and the preceding decomposition, the proof of Theorem 3.1 is merely a balancing act in the regularization parameter ε.
Theorem 3.4. Let P satisfy (A) and let Q satisfy (B). Then, for ε small enough,
The proof of Theorem 3.4 relies on the following qualitative picture: if a point x belongs to some Laguerre cell L j , and is far away from the boundary of L j , then the entropic optimal plan π ε will send almost all of its mass towards the point y j = T 0 (x), sending an exponentially small amount of mass to the other points y j . Such a picture is correct as long as x is at distance at least ε from the boundary of the Laguerre cell L j , incurring a total error of order ε. A rigorous proof of Theorem 3.4 can be found in Appendix B.
Note that this rate is slower than the rate appearing in Proposition 2.3 in the continuous-to-continuous case. The following example shows that the dependency in ε is optimal in Theorem 3.4, indicating that the presence of discontinuities necessarily affects the approximation properties of the entropic Brenier map.
Example 3.5. Let P be a probability measure on R having a symmetric bounded density p continuous at 0, and let Q = 1 2 (δ -1 + δ 1 ). Following (Altschuler et al., 2022, Section 3), one can check that the entropic Brenier map in this setting is the following scaled sigmoidal function
whereas the optimal transport map T 0 (x) = sign(x). Then, performing a computation
where in the last step we invoked the dominated convergence theorem, and computed the limiting integral. Remark 3.6. Assumption (A) can be relaxed for Theorem 3.4 to hold. More precisely, it can be replaced by Assumptions 2.2 and 2.9 of Altschuler et al. (2022), that hold for unbounded measures such as the normal distribution.
Finally, we present the sample-complexity result:
Theorem 3.7. Let P satisfy (A) and let Q satisfy (B). Then, for 0
Remark 3.8. In (Rigollet & Stromme, 2022), the authors show that if P and Q are merely compactly supported with
where c > 0 is some absolute positive constant. Thus, under the additional structural assumptions of the semi-discrete formulation, we are able to significantly improve the rate of convergence between the empirical and population entropic Brenier maps.
The proof of Theorem 3.7 relies on a novel stability result, reminiscent of (Manole et al., 2021, Theorem 6), which is of independent interest. We provide the proof in Appendix C.
Proposition 3.9. Let µ, ν, µ ′ , ν ′ be four probability measures supported in B(0; R). Then the entropic maps T µ→ν ε and
Remark 3.10. The right side of the bound in Proposition 3.9 is equal to
where
. Proposition 3.9 is therefore the entropic analogue of the stability bounds of Manole et al. (2021, Theorem 6) and Ghosal & Sen (2022, Lemma 5.1). Unlike those results, Proposition 3.9 allows both the source and target measure to be modified, and does not require any smoothness assumptions.
To prove Theorem 3.7, we first consider the one-sample setting, where we assume that we only have access to samples Y 1 , . . . , Y n ∼ Q, but we have full access to P . We then consider the one-sample entropic estimator T P →Qn ε . We apply Proposition 3.9 with µ = µ ′ := P , ν := Q n and ν ′ := Q, yielding (see Corollary C.1 for details)
Let χ 2 (P ∥Q) denote the χ 2 -divergence between probability measure. Young's inequality (see Lemma H.1) and the inequality KL(Q n ∥Q) ≤ χ 2 (Q n ∥Q) yield the following bound:
To complete our proof sketch, we use a new stability result on the entropic dual Brenier potentials, catered for the semidiscrete setting.
Proposition 3.11. Let µ be a measure that satisfies (A). Let ν, ν ′ be two discrete probability measures supported on {y 1 , . . . , y J }, with ν ′ ≥ λν for some λ > 0. Then, for 0
where C depends on R, p min and p max .
Moreover, a computation provided in Lemma H.2 shows that E[χ 2 (Q n ∥Q)] = J-1 n , which is enough to conclude the proof of the one-sample case, see Appendix E for details. The two-sample setting is tackled using similar reasoning, where we ultimately prove in Appendix F that the risk E∥ Tε -
Such a quantity can again be related to the estimation of the dual potentials ψ P →Qn ε and ψ Pn→Qn ε . Using the same reasoning as before, we expect a parametric rate of convergence for this term as well. Merging the two results completes the proof of Theorem 3.7. We refer to Appendix F for full details.
The upper bound of Theorem 3.7 shows that our estimator achieves the n -1/2 rate. In fact, the following simple proposition tells us that this rate is optimal in the semi-discrete case.
Proposition 4.1. Let P be the uniform distribution on [-1/2, 1/2] d and for any J ≥ 2, let Q J denote the space of of probability measures with at most J atoms, supported on [-1/2, 1/2] d . Define the minimax rate of estimation
Then, it holds that R n (Q J ) ≥ n -1/2 /64.
Proof. Let e be a vector of the canonical basis of R d , scaled by 1/2. Fix 0 < r < 1/2 and let
Therefore, by Le Cam's lemma (see, e.g., Wainwright, 2019, Chapter 15),
Let d H 2 (Q 0 , Q 1 ) denote the (squared) Hellinger distance between measures. We have
We obtain the conclusion by picking r = n -1/2 /4.
The 1-Nearest-Neighbor estimator, henceforth denoted T1NN , was proposed by (Manole et al., 2021) as a computational surrogate for estimating optimal transport maps in the low smoothness regime. Written succinctly, their estimator is
are Voronoi regions i.e.
and π is the optimal transport plan between the empirical measures P n and Q n , which amounts to a permutation.
Computing the closest X i to a new sample x has runtime O(n log(n)), though the complexity of this estimator is determined by computing the plan π, which takes O(n 3 ) time via, e.g., the Hungarian Algorithm (see Peyré & Cuturi, 2019, Chapter 3).
When φ 0 is smooth and strongly convex, Manole et al. (2021) showed that, for d ≥ 5,
. In contrast to the rate optimality of the entropic Brenier map, we now show that T1NN is proveably suboptimal in the semi-discrete setting. Not only does it fail to recover the minimax rate obtained by the entropic Brenier map, but its performance in fact degrades in comparison to the smooth case. A proof appears in Appendix G. Proposition 4.2. There exist a measure P satisfying (A) and a discrete measure Q satisfying (B) such that for d ≥ 3
We briefly verify our theoretical findings on synthetic experiments. To create the following plots, we draw two sets of n i.i.d. points from P , (X 1 , . . . , X n ) and (X ′ 1 , . . . , X ′ n ), and create target points
, where T 0 is known to us in advance in order to generate the data. Our estimators are computed on the data (X 1 , . . . , X n ) and (Y 1 , . . . , Y n ), and we evaluate the Mean-Squared error criterion
of a given map estimator T using Monte Carlo integration, using 50000 newly sampled points from P . We plot the means across 10 repeated trials, accompanied by their standard deviations.
First consider P = Unif([0, 1] d ) and create atoms {y 1 , . . . , y J } by partitioning the points along the first coordinate for all j ∈ [J]:
We choose uniform q j = 1/J for j ∈ [J]. In this case, it is easy to see that the optimal transport map T 0 (x) is uniquely defined by the first coordinate of x 1 . Figure 2 illustrates the rate-optimal performance of the entropic Brenier map, and the proveably suboptimal performance of the 1-Nearest-Neighbor estimator.
We now consider a synthetic experiment with far less symmetry. Let P = Unif([0, 1] d ), and fix J ∈ N. We randomly generate y 1 , . . . , Y J ∈ [0, 1] d , and also randomly generate ψ 0 ∈ R J , and consider the optimal transport map T 0 (x) = argmin j∈[J] {x ⊤ y j -(ψ 0 ) j }. We define Q = (T 0 ) ♯ P , leading to the same setup as before, but with a less structured optimal transport map. We consider J = 5 and d = 50, and repeat the procedure of the preceding section to generate our data, and the resulting estimator.
We turn our attention to a discontinuous transport map, where for x ∈ R d , all the coordinates are fixed except for the first one
We choose P = Unif([-1, 1] d ) to exhibit a discontinuity in the data. Focusing on d = 10, we see in Figure 4 that the entropic map estimator avoids the curse of dimensionality and enjoys a faster convergence rate, with better constants.
Understanding optimal transport maps in the semi-discrete case is a natural stepping-stone to understanding the case for general discontinuous transport maps. In this work, we propose a tractable, minimax optimal estimator of the Brenier map in the semi-discrete setting, where the rate of estimation is dimension independent. To prove our result, we require several new results and techniques, and, as a by-product of our analysis, give the first parametric rates of estimation the entropic Brenier map, without exponential dependence in the regularization parameter. Our synthetic experiments indicate that the entropic Brenier map might be useful in estimating other variants of discontinuous transport maps, which constitutes an interesting direction for future research.
We recall in this section some known results on entropic optimal transport that will be needed later. Let µ, ν ∈ P(Ω), where Ω ⊂ B(0; R) is a compact set.
Lemma A.1 (Genevay et al., 2019). The entropic potential (φ µ→ν ε , ψ µ→ν ε ) have a bounded amplitude, in the sense that
for some absolute constant c, and similarly for ψ µ→ν ε .
Assume now that ν = J j=1 ν j δ yj is a discrete measure. In this situation, only the values of the dual potential ψ µ→ν ε on the points y 1 , . . . , y J are relevant. We therefore consider ψ µ→ν ε as a vector in R J . The potentials φ µ→ν ε and ψ µ→ν ε are dual of one another, in the sense of the ε-Legendre transform. Given a finite measure ρ, the ε-Legendre transform of a function h with respect to ρ is given by
Relations ( 11) and ( 12) express that
) and vice-versa. In the semi-discrete setting, it is also convenient to introduce the ε-Legendre transform with respect to the counting measure σ on {y 1 , . . . , y J }. For a vector ψ ∈ R J , we have
The Φ ε transform and the Φ ν ε transform are linked through the relation
where we call ψ a shifted potential. With this notation, the optimality condition on the potentials can be rephrased. Let
Then, the function F µ→ν ε is minimized at ψµ→ν ε . For ψ ∈ R J and x ∈ R d , we introduce the probability measure supported on {y 1 , . . . , y J } given by ∀i ∈
e (⟨x,yi⟩-ψ(yi))/ε J j=1 e (⟨x,yj ⟩-ψ(yj ))/ε = e (⟨x,yi⟩-Φε(ψ)(x)-ψ(yi))/ε . (28)
A computation gives ∇F µ→ν ε (ψ) = π x ε [ψ] dµ(x) -ν, so that at optimality, we have
In this case,
is the conditional distribution of the second marginal of π ε given that the first is equal to x, as in Section 2.2.1. More generally, for any potential ψ, the first order condition implies that ψ is equal to ψµ→ν ψ ε , the optimal dual potential between µ an ν ψ = π x ε [ψ] dµ(x).
Proof of Theorem 3.4. Let i, j ∈ [J]. We define the jth slack at x ∈ L i by
As φ 0 is the Legendre transform of ψ 0 , we have ∆ ij (x) ≥ 0. If the cells L i and L j have a nonempty intersection, the set
= t} represents the trace on L i of the hyperplane spanned by the boundary between L i and L j , shifted by t. It is stated in (Altschuler et al., 2022) that for every nonnegative measurable function
where
is the (weighted) surface of the boundary between the i th and j th Laguerre cells (should it exist). Given x ∈ L i , let s(x) = min j̸ =i 1 2 ∆ ij (x). When the point x is sufficiently inside its Laguerre cell, the conditional probability π x ε becomes extremely concentrated around the point y i , as the next lemma shows. Note that π x 0 = δ yi when x ∈ L i . Lemma B.1. Let x ∈ L i . For ε small enough, it holds that for every j ∈ [J], |π x ε (y j ) -π x 0 (y j )| ≤ ce -s(x)/ε , where c depends on J, the distances ∥y i -y j ∥ and on the quantities w ij .
Such a result was already stated in (Delalande, 2022, Corollary 2.2), although while requiring that the source measure P has a Hölder continuous density. Only assumption (A) is needed here.
Proof. According to (Altschuler et al., 2022, Proposition 4.6), for ε small enough,
where ψε is the shifted version of ψ ε (see ( 25)) and C depends on the distances ∥y i -y j ∥ and on the w ij s. Following (Delalande, 2022, Proof of Corollary 2.2) and ( 28), we have for j ̸ = i
We can bound for any x ∈ L i ,
Therefore, letting C ′ denote a constant, which may depend on J, whose value may change from line to line, we obtain
where in the second equality, we used the definition of s(x). Assumption (A) ensures that the functions h ij s are bounded, which implies that the right-hand side in ( 35) is of order ε.
Proof of Proposition 3.9. Note that we may assume without loss of generality that ν ≪ ν ′ and that KL(ν∥ν ′ ) < ∞, for otherwise the bound is vacuous. For notational convenience, we omit the dependence on ε in the subscripts.
Write π µ,ν = γ µ,ν (x, y)dµ(x)dν(y) for the entropic optimal plan between µ and ν, where γ µ,ν = exp 1 ε (⟨x, y⟩ -φ µ→ν (x) -ψ µ→ν (y)) , and analogously define
Consider the measure γ µ ′ ,ν ′ (x, y) dµ(x) dν ′ (y). The first-order optimality condition for (φ µ ′ →ν ′ , ψ µ ′ →ν ′ ) implies that
so that γ µ ′ ,ν ′ (x, y) dν ′ (y) is a probability measure. Let us write dπ x (y) = γ µ,ν (x, y) dν(y) and dρ x (y) = γ µ ′ ,ν ′ (x, y) dν ′ (y).
We make the following observations: first, T µ→ν (x) = y dπ x (y) and T µ ′ →ν ′ (x) = y dρ x (y). Second, the support of ρ x lies inside B(0; R); since any Lipschitz function f on B(0; R) satisfies sup x f (x) -inf x f (x) ≤ 2R, Hoeffding's lemma (see Boucheron et al., 2013, Lemma 2.2) implies that if f is Lipschitz and f dρ x = 0, then
This implies (Bobkov & Götze, 1999, Theorem 3.1) that
Third, Jensen's inequality implies that for any coupling γ between π x and ρ x ,
so that in particular,
Combining these facts, we obtain
Integrating both sides of this equation with respect to µ yields
Expanding the definition of γ µ,ν and γ µ ′ ,ν ′ and using that log dν dν ′ (y) dπ µ,ν (x, y) = log dν dν ′ (y) dν(y) = KL(ν∥ν ′ ) yields the claim.
We now record two corollaries of this bound, which apply when either the source or the target measures of the entropic maps agree. Corollary C.1. For any µ, ν, ν ′ supported in B(0; R),
Proof. We apply Proposition 3.9 with µ = µ ′ , which yields (once again omitting the dependency in ε)
By definition, (φ µ→ν ′ , ψ µ→ν ′ ) minimizes the expression φ dµ + ψ dν ′ + ε e (⟨x,y⟩-φ(x)-ψ(y))/ε dµ(x) dν ′ (y) -ε, so, recalling that e (⟨x,y⟩-φ µ→ν ′ (x)-ψ µ→ν ′ (y))/ε dµ(x) dν ′ (y) = 1, we have in particular
where we have used that the first-order optimality condition for (φ µ→ν , ψ µ→ν ) implies that e (⟨x,y⟩-φ µ→ν (x)-ψ µ→ν (y))/ε dµ(x) dν ′ (y) = 1 as well (see ( 11)). This implies
Applying this inequality to (42) yields
Corollary C.2. For any µ, µ ′ , ν supported in B(0; R),
Proof. We apply Proposition 3.9 with ν = ν ′ , yielding (dropping the dependency on ε)
An argument analogous to the one used in the proof of Corollary C.1 gives the inequality
or, equivalently,
and combining this inequality with (45) proves the claim. ). Let ν = J j=1 ν j δ yj be a measure supported on {y 1 , . . . , y J } ⊆ B(0; R) and let µ supported on a compact convex set Ω ⊆ B(0; R) with a density p satisfying p min ≤ p ≤ p max for some p max ≥ p min > 0. For ψ ∈ R J , define ν ψ = π x ε [ψ] dµ(x) and assume that ν ψ ≥ λν for some 0 < λ ≤ 1. Then, we have for ε ∈ (0, 1)
where C = e 2R 2 pmax pmin + ε -1 pmin pmax .
Proof. As µ and ε are fixed, we will simply write ψ ν instead of ψ µ→ν ε , and write similarly F ν = F µ→ν ε . Recall the definition (25) of the shifted potential ψν (y j ) = ψ ν (y j ) -ε log ν j . According to (Delalande, 2022, Theorem 3.2), the functional F ν is minimized at the vector ψν , with
For t ∈ [0, 1], let ψ t = ψν + t(ψ -ψν ) and let ν t = π x ε [ψ t ] dµ(x). The potential ψ t is the (shifted) entropic Brenier potential between µ and ν t , so that it minimizes the functional F νt (see Appendix A). Also, note that ∇ 2 F ν does not depend on ν, so that
Lemma D.2. Write ν t = J j=1 ν t,j δ yj . Then, for all t ∈ [0, 1] and j ∈ [J], we have ν t,j ≥ pmin pmax ν 1-t 0,j ν t 1,j .
This lemma is enough to conclude the proof. Indeed, ν 1 = ν ψ ≥ λν, so that it implies that Var νt (v) ≥ pmin pmax λVar ν (v).
Proof of Lemma D.2. According to (Delalande, 2022, Proof of Proposition 4.1),
Therefore, if we let h t (x) = e (⟨x,yj ⟩-ψt(yj )-Φε(ψt)(x))/ε , then we have h t (tx + (1 -t)y) ≥ h 0 (x) t h 1 (y) 1-t . By the Prékopa-Leindler inequality,
Proof of Proposition 3.11. As in the previous proof, we drop the ε and µ dependency in our notation. Write ν k = J j=1 ν k,j δ yj for k = 0, 1, and define as before the shifted potentials ψν k (y j ) = ψ ν1 (y j ) -ε log ν k,j . Let θ > 0 be a parameter to fix. According to Proposition D.1, Lemma H.1, and using the inequality F ν1 ( ψν1 ) ≤ F ν1 ( ψν0 ), we have
We pick θ = Cλ to conclude that
Therefore, using the inequality | log(a/b)| ≤ |a -b|/ min{a, b} for a, b > 0,
Lemma E.1 (Sample complexity in the one-sample case). Assume that P satisfy (A) and that Q satisfy (B). Then, it holds that
Proof. To ease notation, we write T ε,n = T P →Qn ε and ψ ε,n = ψ P →Qn ε . As explained in Section 3, the stability result Proposition 3.9 implies that
Write Q = J j=1 q j δ yj and Q n = J j=1 qj δ yj , and introduce the event
If E is not satisfied, we use the fact that the entropic potentials have a bounded amplitude (see Lemma A.1), to obtain that
Lemma E.2. Let E be the event that Q n ≥ Q/2. Then P(E c ) ≤ Je -cqminn for some c > 0.
Proof. By (Vershynin, 2018, Exercise 2.3.2), we have P(E c ) ≤ J j=1 P(q j < q j /2) ≤ Je -cqminn for some c > 0.
We obtain
by Lemma H.2.
The goal of this section is to prove Theorem 3.7. We will actually prove a more general result, and show that for any discrete measure ν = J j=1 ν j δ yj supported on {y 1 , . . . , y J } with ν j ≥ ν min > 0 for all j ∈ [J], we have for log(1/ε) ≲ n/ log(n),
Theorem 3.7 follows from (58) by conditioning on Q n . Let E be the event that
We obtain Theorem 3.7 by combining this bound with Lemma E.1.
To prove (58), we first use Corollary C.2 which yields
where we recall that for a potential ψ, the shifted potential ψ is given by ψj = ψ j -ε log ν j . The remainder of the proof consists in bounding this integral by using localization arguments and standard bounds on suprema of empirical processes. Our first goal is to show that the potential ψ Pn→ν ε is close to to the potential ψ P →ν ε for the ∞-norm. It will be convenient to work with the "L ∞ -variance"
As the measure ν is lower bounded, it holds that
Lemma F.1 (Supremum of ε-Legendre transforms). Let ψ 0 be a fixed potential and let τ > 0. Then, for all j ∈ [J],
for some absolute constant C.
Proof. For a metric space (A, d) and u > 0, we let N (u, A, d) be the covering number of A at scale u, that is the smallest number of balls of radius u needed to cover A. Let B be the L ∞ -ball of radius τ in R J , centered at ψ 0 , and let ∥
We start with the second inequality. Note that ψ → Φ ε (ψ) is 1-Lipschitz continuous, and that the functional
As c d(P -P n ) = 0, we can therefore restrict the supremum to vectors ψ ∈ B. Furthermore, an envelope function of the class {Φ ε (ψ) -Φ ε (ψ 0 ) : ψ ∈ B} is the constant function equal to τ . Therefore, by Lemma H.3, we obtain
We repeat the same argument for the first inequality. The functional π x ε is invariant by translation:
Remarking furthermore that 0 ≤ π x ε (ψ) j ≤ 1 (so that the class of functions {x → π x ε (ψ) j : ψ ∈ B} admits the constant function 1 as an envelope function), we obtain the following control using Lemma H.3:
where c 0 , c 1 and c 2 are absolute constants, and the last line follows from arguing whether c 1 < τ /ε or not.
Proposition F.2. Assume that P satisfies (A) and let ν = J j=1 ν j δ yj be a measure supported on {y 1 , . . . , y J } ⊂ B(0; R), with ν j ≥ q min for all j ∈ [J]. Then, for all 0 < ε ≤ 1 with log(1/ε) ≲ n/ log(n), it holds that
Proof. To alleviate notation, we will write ψ n = ψ Pn→ν
Let us bound P(E c ). As ψn is the minimum of F n , we have ν = π x ε ( ψn ) j dP n (x) (see Appendix A). Therefore, we may write ν n,j = π x ε ( ψn ) j dP n (x) + π x ε ( ψn ) j d(P -P n )(x) = ν j + Z j , where Z j = π x ε ( ψn ) j d(P -P n )(x) = (π x ε ( ψn ) j -π x ε ( ψ0 ) j ) d(P -P n )(x).
Note that Var ∞ ( ψn -ψ0 ) ≲ R 2 (see Lemma A.1), so that by Lemma F.1 and Lemma H.
Proof. Let Z = Var ∞ ( ψn -ψ0 ). Let once again a k = 2 k / √ n for k ≥ 1, with a 0 = 0. Fix some p > 2, with q = p p-1 . For a > 0, let D a = sup Var∞(ψ-ψ0)≤a 2 (Φ ε (ψ) -Φ ε ( ψ0 )) d(P -P n ) . By Hölder inequality and Markov inequality, we obtain,
where we use Proposition F.2, Lemma F.1 and Lemma H.3 at the last line. Equation (59) then gives the conclusion.
In this section, we prove Proposition 4.2. We let P be the Lebesgue measure on Ω = [0, 1] d , and let y 0 = (0, 1/2, . . . , 1/2) and y 1 = (1, 1/2, . . . , 1/2). We denote by P n an empirical measure consisting of i.i.d. samples from P . As in Appendix F, we work in a general setting of a generic discrete target measure ν, which may either be fixed or may be a random measure independent of P n . We let ν = j=0,1 ν j δ yj for ν 0 , ν 1 ≥ 1 4 ; this latter condition will hold with overwhelming probability if ν is an empirical measure Q n corresponding to n i.i.d. samples from Q = 1 2 δ y0 + 1 2 δ y1 . Following Manole et al. (2021), we define the one-nearest neighbor estimator T1NN in this general context by
where π is the empirical optimal coupling between P n and ν.
We first examine the structure of the Brenier map T 0 = ∇φ 0 . The considerations in Section 2.1.1 imply that T 0 (x) = y 0 ⟨e 1 , x⟩ ≤ ν 0 y 1 ⟨e 1 , x⟩ > ν 0 ,
where e 1 is the first elementary basis vector. The potential φ 0 is not differentiable on the separating hyperplane ⟨e 1 , x⟩ = ν 0 , which has measure 0 under P , but we may arbitrarily assign points on this hyperplane to y 0 .
Proof. We can write Q n = J j=1 qj δ yj , where qj is a binomial random variable with parameters n and q j . We obtain
(q j -q j ) 2 q j .
Taking expectations, our bound reads
Var(q j ) q j = J j=1 q j (1 -q j ) nq j = J -1 n .
Lemma H.3 (Control of suprema of empirical processes). Let X 1 , . . . , X n be an i.i.d. sample from some probability measure P on R d , with P n the associated empirical measure. Consider F a class of functions R d → R with ∥f ∥ ∞ ≤ A for all f ∈ F. For u > 0, let N (u) be the u-covering numbers of F, that is the minimal number of balls of radius u for the ∥ • ∥ ∞ -metric required to cover F. Then,
for two positive absolute constants C 0 and C 1 . Furthermore, for all t > 0,
for some positive absolute constant C 2 . Eventually, for all p ≥ 2,
Proof. See (Vaart & Wellner, 1996, Theorem 2.14.2 and Theorem 2.14.5).