Learning Elastic Costs to Shape Monge Displacements

Given a source and a target probability measure supported on $\mathbb{R}^d$, the Monge problem asks to find the most efficient way to map one distribution to the other. This efficiency is quantified by defining a \textit{cost} function between source and target data. Such a cost is often set by default in the machine learning literature to the squared-Euclidean distance, $\ell^2_2(\mathbf{x},\mathbf{y})=\tfrac12|\mathbf{x}-\mathbf{y}|_2^2$. Recently, Cuturi et. al '23 highlighted the benefits of using elastic costs, defined through a regularizer $\tau$ as $c(\mathbf{x},\mathbf{y})=\ell^2_2(\mathbf{x},\mathbf{y})+\tau(\mathbf{x}-\mathbf{y})$. Such costs shape the \textit{displacements} of Monge maps $T$, i.e., the difference between a source point and its image $T(\mathbf{x})-\mathbf{x})$, by giving them a structure that matches that of the proximal operator of $\tau$. In this work, we make two important contributions to the study of elastic costs: (i) For any elastic cost, we propose a numerical method to compute Monge maps that are provably optimal. This provides a much-needed routine to create synthetic problems where the ground truth OT map is known, by analogy to the Brenier theorem, which states that the gradient of any convex potential is always a valid Monge map for the $\ell_2^2$ cost; (ii) We propose a loss to \textit{learn} the parameter $\theta$ of a parameterized regularizer $\tau_\theta$, and apply it in the case where $\tau_{A}(\mathbf{z})=|A^\perp \mathbf{z}|^2_2$. This regularizer promotes displacements that lie on a low dimensional subspace of $\mathbb{R}^d$, spanned by the $p$ rows of $A\in\mathbb{R}^{p\times d}$.

Paper

Similar papers

Peer review

Reviewer 3gJN7/10 · confidence 2/52024-06-27

Summary

Efficiently mapping one distribution of points to another is crucial in numerous machine learning applications. Optimal transport (OT) theory has become a favored approach for solving these matching problems across various scientific disciplines. This work zeroes in on a significant OT subtask: solving the Monge problem. The objective is to determine a map $T$ that converts high-dimensional source data $(x_1, \ldots x_n)$ into target data $(y_1,\ldots,y_m)$. The map $T$ should act as a pushforward map, ensuring that the transformed source samples align with the target distribution, and be efficient, meaning the transformed points $T(x_i)$ are, on average, close to the original points $x_i$. Efficiency is defined by a cost function $c$ that evaluates the cost of mapping a point $x$ to $T(x)$. Elastic costs in optimal transport (OT), augmented with a regularizer $\tau$, promise structured OT maps, but understanding is limited beyond early experiments by Cuturi et al. (2023). This contrasts with Brenier's detailed work on squared-Euclidean costs. Regularizers also introduce hyperparameter selection challenges. To address this the paper shows OT maps for any elastic cost $h$ can be generated using proximal gradient descent with the proximal operator of $\tau$. This allows visualization of Monge maps beyond gradient-convex Brenier maps and synthetic generation in high dimensions. Introduces subspace elastic costs, promoting displacements in low-dimensional subspaces defined by matrix $A$. It provides sample-complexity estimates for the MBO estimator of Cuturi et al. 2023, and links to the spiked transport model. Since different regularizers create diverse OT map structures, the paper also proposes parametrized families $\tau_{\theta}$ and a loss function to adaptively select $\theta$. Results demonstrate the effectiveness of the MBO estimator and recovery of $A$ from i.i.d. samples in synthetic and single-cell data, showing improved predictive ability over baseline estimators.

Strengths

The first strength of this paper is the first visualizations of Monge maps beyond the standard gradient-convex Brenier maps, as well as synthetic generation in high dimensions, showcasing the practical applicability of the proposed methods. The results, particularly the improved performance of the MBO estimator and the recovery of $A$ in synthetic and single-cell data tasks, provide empirical validation of the proposed methods.

Weaknesses

It would be good if the authors have an elaboration on the meaning of jargon terms in the literature of optimal transport in the introduction itself before discussing the main contributions and prior work.

Questions

I do not have any questions.

Rating

7

Confidence

2

Soundness

3

Presentation

2

Contribution

3

Limitations

Authors have addressed limitations.

Reviewer zNJm5/10 · confidence 3/52024-07-03

Summary

This paper explores the optimal transport problem with an elastic regularizer, to help deal with an underparametrized model. Specifically, the regularizer acts in a specific subspace, defined by a matrix $A$, which is also learned. A convex optimization problem is proposed and a method used to find this optimal transport map, and the solution is applied to single-cell application, with improved observed performance

Strengths

The math seems sound and interesting, and the application seems to benefit from this framework. It is unclear to me the novelty of the work, but I am not aware of any major competing work.

Weaknesses

I got a bit confused as to all the moving parts in this paper. Each particular component makes sense, but how does it all fit together? I think this paper would benefit from having some kind of "algorithms box", either in the main paper or appendix, that lists from start to finish how an application problem is solved using their method. In particular, the interplay between learning A and learning the transport map, and any decisions required to tune that part of the method. There should be some numerical comparisons of this method compared to others, attempting to solve similar problems. An obvious one is OT without the subspace constraint. Are there also non-OT methods that are competitive in this space?

Questions

Eq (8), is this really an entropy term? Seems more like a kernel function. Is this prox operator really easy to use for all types of tau? Also, are there non-euclidean distances that you look at when looking at this elastic prox? What happens if A is not on the stiefel manifold? Does it have to be? are there cases where you may want some "directional penalization" but not necessarily wiping out all movement in that direction? Is it challenging to constrain A to be on the stiefel manifold? What does MBO stand for? What is the computational complexity for this method? What kind of data sizes is it meant for? Is it challenging to find T? is there some kind of method for choosing the hyperparameters, such as for the soft min? What are the p hats and such in Fig 4 signifying?

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

no negative societal impact

Authorsrebuttal2024-08-10

Thanks for your comments!

Many thanks for acknowledging our rebuttal! >> The proximal operator is known in closed-form for a huge family of regularizers . See for example [1] > **Ok, I think I was caught up on if tau = ||Ax||_2^2, then you at least have to solve a linear system with A. But i guess you consider that a small cost, if A is small, and you can compute its SVD easily?** This is a great question that was on our mind for a while while drafting the idea of the paper. While it might be natural to start from $\tau_A(z) :=\|Az\|^2$, three things happen when making that choice: (1) as you mention, one needs to compute an inverse to evaluate the proximal operator. While this can be managed, the real downside is that one would need to differentiate *through* that inverse (more precisely, through a linear system) when computing the Jacobian of our regularizer, i.e. the term $\partial R(\theta)^*[P^\star(\theta)]$ that appears in Eq. 19; This sounds fairly costly. (2) if one uses loss (Eq.19), then one seeks a projection where, ultimately, $\|Az\|^2$ is small. For very high dimensional data, $p\ll d$, it's usually very easy to find a subspace where "nothing happens", and $\|Az\|^2$ is more or less always zero. Therefore using this approach would likely select small / noisy / uninformative projections, which is somewhat the contrary of what we seek. (3) In the absence of any constraint on $A$, one would certainly need to deal with the *magnitude* of $A$ in the optimization, as one can imagine that $A$ going to $0$ or $\infty$ may easily impact most relevant objectives. Taken together, these reasons explain largely why we converged instead towards the definition $\tau_A(z):=\|A^\perp z\|$^2, the norm of $z$ in the **orthogonal** of $A$, with $A$ a matrix in the Stiefel manifold: that way, we get instead that $A$ is where most of the action in the transport occurs, $A$ is easy to optimize and properly normalized. We have explained this intuition in Appendix A, "More on Subspace Elastic Costs", but given the many questions you have asked, we are now convinced this will benefit from more details. > > The original MBO paper looked at L1 and k-overlap norms. In our revision, we have added a plot with the k-overlap norm for the ground truth OT map. > **Is the k overlap norm the same as the latent group norm by Obozinski et al? If so, actually the prox is a bit complex, but the dual norm is often easier to compute. I wonder if that's something you can leverage.** We believe you refer to Obozinski et al's group overlap norm [1]. The $k$-overlap norm that we use (and which was used in the MBO paper) was proposed in [2] and generalized in [3]. The $k$-overlap norm can be viewed as a L1-norm type regularization that does not suffer from shrinkage. Because of this, while being still sparse (i.e. most go down or left), the arrows shown in Figure A exhibit more displacements than those shown with the $\ell_1$ norm (second plot of Fig. 1 in the draft), and none stays put. [1] Laurent Jacob, Guillaume Obozinski, Jean-Philippe VertAuthors, *Group lasso with overlap and graph lasso*, ICML '09: [2] Andreas Argyriou, Rina Foygel, and Nathan Srebro. *Sparse prediction with the k-support norm*, NIPS, 2012. [3] Andrew M. McDonald, Massimiliano Pontil, Dimitris Stamos, *Spectral k-Support Norm Regularization*, NIPS 2014 >> The approach outlined in S.5 requires computing Sinkhorn solutions, and differentiate them. This has compute/memory footprint in , where is the batch-size and depends on the regularization in the soft-min. These blocks are very well implemented by now. For hyperparameters, we can think of cross-validation, but this was not used. > **I think we still need some hard numbers for the numerical experiment section (dataset sizes). Not to say that the experiments provided aren't valuable, and certainly for an OT paper I do expect smaller datasets than for a typical DNN paper, but still, we need a reference.** This is indeed a very good point, we did not realize we had not given these numbers, we apologize for this oversight. These datasets are part of the Sciplex dataset ([Srivatsan et al. 2020], a science paper), which has gained much popularity in recent years in OT, and is featured prominently in [4] Bunne et. al, *Learning single-cell perturbation responses using neural optimal transport*, Nature Methods 2023 The data we use has the following size: | Cell line | Control | Dac | Bel | Quis | |-----|--------:|------:|------:|------:| | A549 | 3274 | 558 | 669 | 475 | | K562 | 3346 | 388 | 656 | 339 | | MCF7 | 6346 | 1562 | 1684 | 1520 | For each drug, we therefore have 3 sets of source/target datasets. We average results over these cell-lines. We have run experiments for 2 other drugs (similar size) as mentioned in our general response. We will report them in the final draft. Our method is substantially better with Givinostat (better than Dac, Bel and Quis shown in draft) and is slightly worse on Hesperadin.

Reviewer zNJm2024-08-11

Thanks for all the responses. I have no further questions and am happy with all answers.

Authorsrebuttal2024-08-12

Thanks for your questions and comments.

Your comments and questions have helped improved our draft, we have incorporated these points. We remain available until the closing of the discussion period to address any other aspect of our draft. the authors

Reviewer JR354/10 · confidence 3/52024-07-10

Summary

The paper studies the problem of computing Monge maps between two distributions. Such maps are usually found with respect to the $\ell_2^2$ cost function, however the focus of the paper is on elastic cost functions which have an additional penality term (denoted by $\tau$). A numerical method is formulated for estimating the Monge maps w.r.t elastic costs. Moreover, a loss function is proposed for learning the parameter $\theta$ of a parametrized regularizer $\tau_{\theta}$, and is applied to the case where $\theta$ corresponds to the low dimensional subspace. Experiments are performed on synthetic data and and also on single-cell data tasks.

Strengths

The paper builds on the work of Cuturi et al. 2023 which proposed the idea of using elastic costs. However that paper proposes an estimator for finding an approximate Monge map (via the MBO estimator), while the present paper proposes a projected gradient descent scheme whose iterates converge to the Monge map. The paper also provides insights for the case where $\tau$ involves a subspace constraint, and the results are likely to be of interest for researchers working on related problems.

Weaknesses

1. The exposition of the paper can be improved substantially. At present, it is difficult to parse the technical details since the notation is unclear at several places. For e.g., in eq. (7), where are the quantities a,b, X and Y defined? 2. The first main contribution involves showing that a projected gradient scheme recovers the ground truth Monge map as the iterates approach infinity. So for finitely many iterates, we would have an approximation to the ground-truth. But even for the MBO estimator, the estimate is a good approximation to the ground-truth provided the number of samples is sufficiently large. So, while this is a nice observation, I feel that its a bit incremental in terms of contribution. 3. The second contribution involves studying the case where the regularizer corresponds to a subspace constraint penalty, for which estimation rates are derived for the MBO estimator. However these appear to follow readily from existing work, so the theoretical contributions are a bit weak in my view.

Questions

1. In Theorem 1, the rate depends exponentially on d but the subspace dimension p does not appear anywhere. Then what is the benefit of the (low-dimensional) subspace constraint?

Rating

4

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

I do not see the limitations discussed anywhere.

Reviewer EFNS6/10 · confidence 2/52024-07-12

Summary

This paper studies the problem of optimal transport with elastic costs. They demonstrate the performance of the MBO estimator on subspace losses. They then introduce a loss for learning the elastic cost function and provide experimental results on the performance of this cost learning scheme.

Strengths

* Clearly and well-written (except for some notational confusion that I mention in the Questions section) * Optimizing over cost functions to find low-dimensional subspace is a cool idea.

Weaknesses

* The theoretical contributions of this work (sections 3 and 4) seem limited. The main machinery used is the application of the MBO estimator, which is from pre-existing work, to straightforward cost functions. It would be more interesting if some theoretical results for the setting of sections 5 and 6 were demonstrated (such as, i) does the true $\theta$ optimize the elastic costs loss when the data is drawn from, e.g., $T^h_g$ and ii) what is the sample complexity of parameter recovery?)

Questions

* Why is it called elastic costs? * What are $\mathbf{a}$ and $\mathbf{b}$ in equation 7? * What is $\tilde{T}$ in Theorem 1? * Am I missing something or is the setting of section 4 very simple given the pre-existing results for $h = \frac{1}{2} ||\cdot||^2$? In particular, because $h$ is a squared Mahalanobis-norm, this OT problem seems easily solvable by first pre-processing the data s.t. $h$ (on the new space) becomes the usual $\frac{1}{2} ||\cdot||^2$ function and then second, solving the usual OT map with cost $\frac{1}{2} ||\cdot||^2$. It seems like the proof of Theorem 1 acknowledges this with the map $W$, but I was wondering if there is anything more to these results than rescaling the data. * Related to the previous question, but can the results of section 4 be extended directly to general squared Mahalanobis norms?

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors have addressed their limitations.

Authorsrebuttal2024-08-13

Many thanks for acknowledging our rebuttal.

We sincerely appreciate your time reading our rebuttal. We are very grateful for your score increase, and remain available to answer any further questions on the paper you may have. The Authors

Reviewer zNJm2024-08-10

- There is a mild confusion that we would like to clarify: While the task of computing the entropic map given a cost and samples weighted as is convex (Eq.10), as you point out, the main task we tackle is that of learning the parameters of a regularizer (S.5). While we pair learning cost parameters with the MBO estimator to improve predictive performance on a real task (S.6.3), learning the cost is not a convex pipeline. Got it, yes that is an important point - We thank you for this great suggestion, and added these algorithms in the attached pdf. Thanks, this is helpful. - We have such experiments in Fig. 5. The dotted line is our baseline: the “standard” MBO estimator that does not use any subspace assumption, computing OT maps with a cost (L.284). The other lines are "our" MBO estimators computed with learned costs. Remarkably, they perform better, even in sense. Nice, these are great. - The proximal operator is known in closed-form for a huge family of regularizers . See for example [1] Ok, I think I was caught up on if tau = ||Ax||_2^2, then you at least have to solve a linear system with A. But i guess you consider that a small cost, if A is small, and you can compute its SVD easily? - The original MBO paper looked at L1 and k-overlap norms. In our revision, we have added a plot with the k-overlap norm for the ground truth OT map. Is the k overlap norm the same as the latent group norm by Obozinski et al? If so, actually the prox is a bit complex, but the dual norm is often easier to compute. I wonder if that's something you can leverage. - The approach outlined in S.5 requires computing Sinkhorn solutions, and differentiate them. This has compute/memory footprint in , where is the batch-size and depends on the regularization in the soft-min. These blocks are very well implemented by now. For hyperparameters, we can think of cross-validation, but this was not used. I think we still need some hard numbers for the numerical experiment section (dataset sizes). Not to say that the experiments provided aren't valuable, and certainly for an OT paper I do expect smaller datasets than for a typical DNN paper, but still, we need a reference. For all the other comments, I read them and am happy with them.

Reviewer 3gJN2024-08-12

Response to rebuttal

Thanks to the authors for addressing my question, and looking forward to improved writing. My score remains the same.

Authorsrebuttal2024-08-13

Many thanks for reading our rebuttal

We are grateful for your time during this reviewing process, and for supporting the paper. the authors

Reviewer JR352024-08-12

Authors response

Thank you for your response to my queries, I acknowledge that I have read the rebuttal.

Authorsrebuttal2024-08-13

Many thanks for acknowledging our rebuttal

We thank you for the time you have taken to read our rebuttal and have already benefitted from your comments (as well as those of all other reviewers) to improve presentation. Although remaining discussion time is short, we are happy to answer any further concerns you may have on the paper. the authors

Reviewer EFNS2024-08-12

Thanks for your response. The response clarifies the contributions and I will increase my score.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC