Estimation of Stochastic Optimal Transport Maps

We thank the reviewer for their thoughtful feedback and clarifying questions.

Weakness, Monge gap literature: Thank you for bringing up this work, which is indeed quite relevant. When specialized to $p=1$, the Monge gap paper [A] aims to learn a parameterized map $T$ minimizing the following: $\mathcal{E}’\_1(T;P,Q) := \underbrace{\mathbb{E}\_P[\\|T(X) - X\\|] - W\_1(P,T\_\sharp P)}\_{\text{Monge gap}} + \Delta(T\_\sharp P,Q),$ where $\Delta$ is some statistical divergence. In practice, they substitute $P$ and $Q$ with their empirical versions and use entropic OT (EOT) distances to approximate both $W\_1$ and $\Delta$, achieving impressive empirical results. We emphasize that these approximations preclude meaningful statistical guarantees without accounting for implicit bias or other regularization, since an optimal map for the empirical problem can perform arbitrarily bad on the population problem.

We can show that, when $\Delta = W\_1$, the Monge gap objective $\mathcal{E}’\_1$ and our $\mathcal{E}\_1$ coincide up to a constant multiplicative factor of 2. Qualitatively similar results hold for $p>1$, though there is no longer a precise equivalence. Thus, the approach in [A] is quite close to performing empirical risk minimization with our transportation error (up to the use of EOT that is necessary for large-scale data sets). We will include this derivation in a new remark for Section 2. In some sense, this shows that our approach already lends itself to map fitting in practice. Moreover, if we were to directly plug in our $\mathcal{E}\_p$ objective into the neural map estimation experiments from [A] (substituting OT with EOT for tractability), we have the benefit of only having to perform back propagation through one EOT problem instead of two. We are optimistic that this will lead to a performance boost in practice. In the updated manuscript, we will supplement our synthetic experiments, which validated our theory, with a large-scale experiment that compares our performance to the Monge gap approach when learning maps for single-cell genomic data (mirroring the setup of Section 6.4 in [A]). Zooming out, we feel that our framework provides a solid statistical foundation for the practical approach to map estimation established in [A], and we hope that our theoretical insights will boost further practical developments.

Second, our rounding estimation approach can be applied in a black-box way on top of any existing map estimator, including a neural estimator. In the added experiments with single-cell genomic data, we will take the resulting neural estimators and see how their test performance varies if we add on this rounding step, for varied choices of rounding partition (namely, tuning the size of the partition cells). It will be interesting to see if there is extra performance to be eked out of existing methods via the regularizing effect of rounding.

Third, we call the reviewer’s attention to our Remark 4, which draws parallels between our DRO procedure and existing Lipschitz regularization techniques.

Finally, we remark that our current synthetic experiments were computationally bottlenecked by high-precision OT computations, which were necessary to validate our statistical guarantees. The Monge gap work sidesteps these by using EOT with a fixed precision, trading some accuracy for scalability. Another regularization approach which may be worth future exploration is that of sliced OT, where one works only with low dimensional projections of the data distributions. We will add a remark to the end of Section 2 on extensions of our framework to regularized OT problems, which are key for large-scale applications.

Question, toy experiments: We agree that our initial synthetic benchmarks were somewhat limited. In addition to the proposals above, we are happy to include the recommended toy setting in the revised manuscript.

Question, deep learning extensions: We hope that our response above re: Monge gap is satisfactory. In particular, we note that our proposal with the rounding estimator relies on our extension to stochastic maps and is rooted in deep learning (of course, empirical evaluation awaits). One could hope for a neural estimator that involves a more explicit parameterization of kernels, rather than an indirect approach via rounding. This is an interesting question for future research which we will include as an open direction in the updated conclusion. Certainly, well-crafted regularization would be essential, in order to prevent variance terms from exploding due to the very high-dimensional search space in this formulation.

[A]: Uscidda & Cuturi. “The Monge Gap: A Regularizer to Learn All Transport Maps”. ICML 2023.

We thank the reviewer for their thoughtful feedback and clarifying questions.

Weakness 1: “They solve Entropic OT to define conditional kernel in error metric. As far as known, Entropic OT has some drawbacks such as instability (exponents in KL regularization), sharpening of generated data (barycentric projection is averaging) and tuning of hyper parameters.” The reviewer is fair to point out some known drawbacks of EOT. We want to emphasize that EOT is not part of our error metric, and our guarantees for the entropic kernel estimator in Section 3 are strictly worse than those for the rounding estimator. Nonetheless, EOT remains widely used in practice, so we wanted to comment on its performance when applied to our estimation task. We also would like to clarify that our estimators do not use barycentric projection.

Weakness 2, p=1 for experiments: We chose p=1 for the experiments to emphasize the application of our framework even when uniqueness of the OT map fails dramatically (when p=1, there are many optimal maps). However, this is not a contributing factor to computation - we will include an updated set of experiments for p=2 in the updated appendix. Indeed, OT solvers with high accuracy guarantees are mostly agnostic to the choice of cost. For the cases of p=1,2, there are also efficient neural solvers using input-Lipschitz and input-convex neural networks, respectively. (See further discussion in response to the next weakness.)

Weakness 3, scalability concerns: The reviewer fairly points out that we do not use our methods in truly large scale settings. As a minor point, we note that our framework opens up the potential for analysis of future algorithms under far weaker assumptions than currently standard. But, more importantly, we are confident that our techniques may have immediate applications for neural map estimation with large-scale data, via a connection to the Monge gap discussed below. As noted by Reviewer RJR5, there is a promising approach for neural map estimation using the so-called “Monge gap” [A], which we will discuss in the updated related work. When specialized to $p=1$, this work aims to learn a parameterized map $T$ minimizing the following: $\mathcal{E}’\_1(T;P,Q) := \underbrace{\mathbb{E}\_P[\\|T(X) - X\\|] - W\_1(P,T\_\sharp P)}\_{\text{Monge gap}} + \Delta(T\_\sharp P,Q),$ where $\Delta$ is some statistical divergence. In practice, they substitute P and Q with their empirical versions and use entropic OT (EOT) distances to approximate both $W\_1$ and $\Delta$, achieving impressive empirical results. We emphasize that these approximations preclude meaningful statistical guarantees without accounting for implicit bias or other regularization, since an optimal map for the empirical problem can perform arbitrarily bad on the population problem.

We can show that, when $\Delta = W\_1$, the Monge gap objective $\mathcal{E}’_1$ and our $\mathcal{E}_1$ coincide up to a constant multiplicative factor of 2. Qualitatively similar results hold for $p>1$, though there is no longer a precise equivalence. Thus, the approach in [A] is quite close to performing empirical risk minimization with our transportation error (up to the use of EOT that is necessary for large-scale data sets). We will include this derivation in a new remark for Section 2. In some sense, this shows that our approach already lends itself to map fitting in practice. Moreover, if we were to directly plug in our $\mathcal{E}_p$ objective into the neural map estimation experiments from [A] (substituting OT with EOT for tractability), we have the benefit of only having to perform back propagation through one EOT problem instead of two. We are optimistic that this will lead to a performance boost in practice. In the updated manuscript, we will supplement our synthetic experiments, which validated our theory, with a large-scale experiment that compares our performance to the Monge gap approach when learning maps for single-cell genomic data (mirroring the setup of Section 6.4 in [A]). Zooming out, we feel that our framework provides a solid statistical foundation for the practical approach to map estimation established in [A], and we hope that our theoretical insights will boost further practical developments.

Second, our rounding estimation approach can be applied in a black-box way on top of any existing map estimator, including a neural estimator. In the added experiments with single-cell genomic data, we will take the resulting neural estimators and see how their test performance varies if we add on this rounding step, for varied choices of rounding partition (namely, tuning the size of the partition cells). It will be interesting to see if there is extra performance to be eked out of existing methods via the regularizing effect of rounding.

Third, we remark that our current synthetic experiments were computationally bottlenecked by high-precision OT computations, which were necessary to validate our statistical guarantees. The Monge gap work sidesteps these by using EOT with a fixed precision, trading some accuracy for scalability. Another regularization approach which may be worth future exploration is that of sliced OT, where one works only with low dimensional projections of the data distributions. Finally, one could also employ neural OT solvers, which compute W1 and W2 using neural parameterizations of Lipschitz and convex function classes, respectively. We will add a remark to the end of Section 2 on extensions of our framework to regularized/approximate OT problems, which are key for large-scale applications.

Finally, we will add an appendix section which pushes the dimension and sample count as high as possible for our existing high-precision, synthetic experiments. We expect that they can be pushed a fair bit further if bootstrapping is omitted.

Weakness 4, sub-Gaussianity: As noted in Theorem 2, we can relax the sub-Gaussianity assumption on $\mu$ by using a non-uniform partition for the rounding algorithm. One can also employ weaker tail bounds on $\nu$, but we expect that these will necessarily worsen the achievable risk (as remarked after Theorem 2). We note that the existing statistics literature generally required bounded support, a stronger assumption.

Q1, evaluation vs optimization objective: This is certainly a natural question. Since we were primarily focused on statistical guarantees, we knew that performing empirical risk minimization with $\mathcal{E}\_p$ as our objective was unfortunately out of the question, since the optimal empirical map can perform arbitrarily poorly on the population measures. However, our approach in Section 4, using Wasserstein distributionally robust optimization, does use $\mathcal{E}\_p$ as an objective for robust optimization (and can be interpreted as a regularized version of ERM, see Remark 4). Finally, we hope that our promised large-scale experiments above, which compare to the Monge gap approach, will shed some empirical light on the merit of $\mathcal{E}\_p$ for neural map estimation.

Q2a, barycentric projection: We agree that barycentric projection is problematic. This is a feature of some prior work, but we emphasize that none of our estimators use barycentric projection. We would appreciate if the reviewer could point out which part of our text caused this confusion (namely, the wrong impression that barycentric projection was employed in our error metric or estimators) and we will work to elucidate it for the final version.

Q2b, two neural networks: We kindly ask that the reviewer clarify this suggestion (we are happy to discuss further during the discussion period).

Q3, unbiasedness: It is unclear to us what a suitable notion of unbiasedness would be for this paper's setting (since we are not comparing the estimated map to a fixed optimal map $T^\star$). Does the reviewer have any particular definition in mind?

Q4, EOT & barycentric projection: As mentioned above, we do not perform barycentric projection, and our best estimation guarantees do not use EOT.

Q5, sub-Gaussianity: We point the reviewer to our response for Weakness 4.

Q6, intuition for rounding: At a high-level, our rounding procedure allows us to extend a map learned from a discrete point set to the full source domain, while also serving to reduce variance. Recall that the rounding estimator bins the source points into bins of a tuned size, learns a map from the binned points to the (untouched) target points, and, given a new point, sends it to the destination of its bin under the learned map. The variance reduction intuition is as follows: at high spatial resolution, OT analysis is typically dominated by sampling variance terms, and binning serves to keep our analysis at a moderate spatial resolution. Of course, binning also introduces some bias. Balancing these factors was the key to our proof of Theorem 2. We include this approach because it achieves sharper statistical rates than the estimator using EOT.

Q7, large-scale experiments: As mentioned above, the role of our current experiments was to validate our statistical rates, requiring high-precision OT computations that were a bottleneck to very high-dimensions / sample sizes. Please see our response to the scalability concern above for a new set of neural estimation experiments to be included in the revised manuscript. We hope that the single-cell genomics experiment is satisfactory as a large-scale example. This is a compelling option due to the existence of past results to benchmark against.

Limitations, EOT: Please see previous responses on EOT.

[A]: Uscidda & Cuturi. “The Monge Gap: A Regularizer to Learn All Transport Maps”. ICML 2023.

We thank the reviewer for their thoughtful feedback and clarifying questions.

Weakness: The reviewer fairly points out that we do not use our methods to design more accurate / scalable algorithms for truly large scale settings. As a minor point, we note that our framework opens the door for analysis of future algorithms under far weaker assumptions than currently standard. But, more importantly, we are optimistic that our techniques may have immediate applications for neural map estimation with large-scale data, via a connection to the Monge gap discussed below.

As noted by Reviewer RJR5, there is a promising approach for neural map estimation using the so-called “Monge gap” [A], which we will discuss in the updated related work. When specialized to $p=1$, this work aims to learn a parameterized map $T$ minimizing the following: $\\mathcal{E}’\_1(T;P,Q) := \underbrace{\mathbb{E}\_P[\\|T(X) - X\\|] - W\_1(P,T\_\sharp P)}\_{\text{Monge gap}} + \Delta(T\_\sharp P,Q),$ where $\Delta$ is some statistical divergence. In practice, they substitute P and Q with their empirical versions and use entropic OT (EOT) distances to approximate both $W\_1$ and $\Delta$, achieving impressive empirical results. We emphasize that these approximations preclude meaningful statistical guarantees without accounting for implicit bias or other regularization, since an optimal map for the empirical problem can perform arbitrarily bad on the population problem.

We can show that, when $\Delta = W\_1$, the Monge gap objective $\mathcal{E}’\_1$ and our $\mathcal{E}\_1$ coincide up to a constant multiplicative factor of 2. Qualitatively similar results hold for $p>1$, though there is no longer a precise equivalence. Thus, the approach in [A] is quite close to performing empirical risk minimization with our transportation error (up to the use of EOT that is necessary for large-scale data sets). We will include this derivation in a new remark for Section 2. In some sense, this shows that our approach already lends itself to map fitting in practice. Moreover, if we were to directly plug in our $\mathcal{E}\_p$ objective into the neural map estimation experiments from [A] (substituting OT with EOT for tractability), we have the benefit of only having to perform back propagation through one EOT problem instead of two. We are optimistic that this will lead to a performance boost in practice. In the updated manuscript, we will supplement our synthetic experiments, which validated our theory, with a large-scale experiment that compares our performance to the Monge gap approach when learning maps for single-cell genomic data (mirroring the setup of Section 6.4 in [A]). Zooming out, we feel that our framework provides a solid statistical foundation for the practical approach to map estimation established in [A], and we hope that our theoretical insights will boost further practical developments.

Second, our rounding estimation approach can be applied in a black-box way on top of any existing map estimator, including a neural estimator. In the added experiments with single-cell genomic data, we will take the resulting neural estimators and see how their test performance varies if we add on this rounding step, for varied choices of rounding partition (namely, tuning the size of the partition cells). It will be interesting to see if there is extra performance to be eked out of existing methods via the regularizing effect of rounding.

Question, Hölder continuity: This is indeed a strong assumption, but it is in fact weaker than nearly all of the existing work on statistical rates for $L^2$ map estimation (with the exception of the non-comparable setup in [B] that requires $T^\star$ to be piecewise-constant). We agree that this is of less practical interest, which is why we focused on the unrestricted setting first, but we wanted to include this setting since it covered a large set of related work. There are some theoretical results when $p=2$ that guarantee the existence of a regular Brenier map (Caffarelli’s regularity theory - see, e.g., Theorem 3 of [C]), but these are admittedly not quantitative.

We also note that our approach to robustifying estimation to Wasserstein perturbations (for Section 5) uses a black-box method for turning any kernel into a Hölder continuous kernel via Gaussian convolution (at the cost of introducing some bias). Thus, the continuity condition can be of practical interest even if it does not hold exactly in practice.

[A]: Uscidda & Cuturi. “The Monge Gap: A Regularizer to Learn All Transport Maps”. ICML 2023. [B]: A.-A. Pooladian, V. Divol, and J. Niles-Weed. Minimax estimation of discontinuous optimal transport maps: The semi-discrete case. ICML 2023. [C]: Manole, Balakrishnan, Niles-Weed, & Wasserman. Plugin estimation of smooth optimal transport maps. The Annals of Statistics, 2024.

We thank the reviewer for their thoughtful feedback and clarifying questions.

Summary clarification: We wanted to point out that bounding both gaps is required in the analysis of both the entropic kernel estimator and the rounding estimator. If we were to distinguish the estimators’ analysis, we would say that the entropic kernel estimator relies on stability of $\\mathcal{E}\_p$ under Wasserstein perturbations, while the rounding estimator relies on TV stability. We will mention this distinction at the start of Section 3. Also, the analysis of the rounding estimator does not require Hölder continuous kernels (that case is treated by a separate DRO estimator introduced in Section 4).

Q1a “could the finite-sample rates inform how OT maps are parametrized and fitted?”: We agree that using this work to inform map estimation in practice is a very natural direction. As noted by Reviewer RJR5, there is a promising approach for neural map estimation using the so-called “Monge gap” [A], which we will discuss in the updated related work. When specialized to $p=1$, this work aims to learn a parameterized map $T$ minimizing the following: $\mathcal{E}’\_1(T;P,Q) := \underbrace{\mathbb{E}\_P[\\|T(X) - X\\|] - W\_1(P,T\_\sharp P)}\_{\text{Monge gap}} + \Delta(T\_\sharp P,Q),$ where $\Delta$ is some statistical divergence. In practice, they substitute $P$ and $Q$ with their empirical versions and use entropic OT (EOT) distances to approximate both $W\_1$ and $\Delta$, achieving impressive empirical results. We emphasize that these approximations preclude meaningful statistical guarantees without accounting for implicit bias or other regularization, since an optimal map for the empirical problem can perform arbitrarily bad on the population problem.

We can show that, when $\Delta = W\_1$, the Monge gap objective $\mathcal{E}’\_1$ and our $\mathcal{E}\_1$ coincide up to a constant multiplicative factor of 2. Qualitatively similar results hold for $p>1$, though there is no longer a precise equivalence. Thus, the approach in [A] is quite close to performing empirical risk minimization with our transportation error (up to the use of EOT that is necessary for large-scale data sets). We will include this derivation in a new remark for Section 2. In some sense, this shows that our approach already lends itself to map fitting in practice. Moreover, if we were to directly plug in our $\mathcal{E}\_p$ objective into the neural map estimation experiments from [A] (substituting OT with EOT for tractability), we have the benefit of only having to perform back propagation through one EOT problem instead of two. We are optimistic that this will lead to a performance boost in practice. In the updated manuscript, we will supplement our synthetic experiments, which validated our theory, with a large-scale experiment that compares our performance to the Monge gap approach when learning maps for single-cell genomic data (mirroring the setup of Section 6.4 in [A]). Zooming out, we feel that our framework provides a solid statistical foundation for the practical approach to map estimation established in [A], and we hope that our theoretical insights will boost further practical developments.

Second, our rounding estimation approach can be applied in a black-box way on top of any existing map estimator, including a neural estimator. In the added experiments with single-cell genomic data, we will take the resulting neural estimators and see how their test performance varies if we add on this rounding step, for varied choices of rounding partition (namely, tuning the size of the partition cells). It will be interesting to see if there is extra performance to be eked out of existing methods via the regularizing effect of rounding.

Third, we call the reviewer’s attention to our Remark 4, which draws parallels between our DRO procedure and existing Lipschitz regularization techniques.

Finally, we remark that our current synthetic experiments were computationally bottlenecked by high-precision OT computations, which were necessary to validate our statistical guarantees. The Monge gap work sidesteps these by using EOT with a fixed precision, trading some accuracy for scalability. Another regularization approach which may be worth future exploration is that of sliced OT, where one works only with low dimensional projections of the data distributions. We will add a remark to the end of Section 2 on extensions of our framework to regularized OT problems, which are key for large-scale applications.

Q1b “Could you characterize the error functional (i.e. which error term dominates when) at different sample ranges?”: This is an interesting question to explore. The answer would appear to depend on quite a few factors, including the kernel estimator, the sample size, the source distribution, and the target distribution. We are unaware of any simple primitive conditions that would help distinguish the case. To investigate this, we will update all of our experiments to separate the contributions of the two error terms and include some discussion on the resulting balance. For our existing synthetic experiments, we note that the feasibility gap was typically larger than the optimality gap, but we will also include this fine-grained data for the neural estimation experiments promised above.

Q2, extension to conditional OT: This is a great question, certainly of relevance in important applications. In general, we believe that our framework is modular enough to apply to a wide-range of OT variants, including conditional OT. For this case, one could extend existing conditional OT works to measure performance with an $\mathcal{E}\_2$-type functional instead of $L^2$. Depending on the modeling assumptions, one could statistically investigate questions of finite-sample risk. We expect meaningful bounds to require additional regularity assumptions on the relationship between contexts and maps/kernels. We view this as a promising direction for future research and will add a discussion to the summary of our final version.

[A]: Uscidda & Cuturi. “The Monge Gap: A Regularizer to Learn All Transport Maps”. ICML 2023.