Learning Elastic Costs to Shape Monge Displacements
Authors: Michal Klein
Aram-Alexandre Pooladian
Pierre Ablin
Eugene Ndiaye
Jonathan Niles-Weed
marco cuturi
Venue: NeurIPS 2024 poster (2024)
Keywords: optimal transport
Primary area: optimization
TL;DR: a deep dive into how changing cost functions impacts transport, and a new algorithm to learn the parameters of a new family of costs

Abstract: Given a source and a target probability measure, the Monge problem studies efficient ways to map the former onto the latter.
This efficiency is quantified by defining a *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$.
The benefits of using *elastic* costs, defined using a regularizer $\tau$ as $c(\mathbf{x},\mathbf{y}):=\ell^2_2(\mathbf{x},\mathbf{y})+\tau(\mathbf{x}-\mathbf{y})$, was recently highlighted in (Cuturi et al. 2023). Such costs shape the *displacements* of Monge maps $T$, namely 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 *learn* the parameter $\theta$ of a parameterized regularizer $\tau_\theta$, and apply it in the case where $\tau_{A}({\bf z}):=\|A^\perp {\bf 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}$. We illustrate the soundness of our procedure on synthetic data, generated using our first contribution, in which we show near-perfect recovery of $A$'s subspace using only samples. We demonstrate the applicability of this method by showing predictive improvements on single-cell data tasks.

Learning Elastic Costs to Shape Monge Displacements
Michal Klein Apple mi michalk@apple.com com
Aram-Alexandre Pooladian
Pierre Ablin
NYU
Apple p d_ablin@apple com
aram-alexandre. araaaaaleccaddrr.pllalaaaaaa.a..
Eugène Ndiaye
Jonathan Niles-Weed
Marco Cuturi Apple cuturi@apple.com com
Apple e s_ndiaye@@pple..om com
NYU
jnw@cims nyu. edu
Abstract
Given source and target probability measure, the Monge problem studies efficient ways to map the former onto the latter. This efficiency quantified by defining cost function between source and target data. Such cost is often set by default the machine learning literature to the squared- Euclidean distance, (2, y) = =||x yll.. The benefits of using elastic costs, defined using a regularizer as y( = +xxxx) et al., 2023]. Such costs shape the displacements of Monge maps T, namely the difference between source point and its image T(x) X. by giving them a structure that matches that of the proximal operator of r. In this work, we make two important contributions to the study of elastic costs: (i) For any elastic cost, we propose numerical method to compute Monge maps that are provably optimal. This provides 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 always valid Monge map for the & cost; (ii) We propose loss to learn the parameter 0 of parameterized regularizer Te. and apply it in the case where (z) := | 2||2- This regularizer promotes displacements that lie on a low-dimensional subspace of RRd. spanned by the p rows of A E Rexd, We illustrate the soundness of our procedure synthetic data, generated using our first contribution, in which we show near-perfect recovery A's subspace using only samples. We demonstrate the applicability of this method by showing predictive improvements on single-cell data tasks.
was recently highlighted in [Cuturi
Introduction
Finding efficient ways to map distribution of points onto another is low-level task that plays crucial role across many machine learning (ML) problems. Optimal transport (OT) theory [Santam- brogio, 2015] has emerged as tool of choice to solve such challenging matching problems, notably in single-cell genomics [Schiebinger et al., 2019, Tong et al., 2020, Bunne et al., 2023, 2024, Klein et al., 2023]. We focus in this work on the numerical resolution of the Monge problem, which aims, using high-dimensional source and target data samples (x)...., X,) and .. map T 2d - Rd that simultaneously (i) a pushfoward map, in the sense that T applied on source samples recovers the distribution of target samples; (ii) efficient, in the sense that (x;; is not too far, on average from x,. The notion of efficiency can be made precise by choosing a real-valued cost function that compares point and its mapping via c(x, T(x)) IR.
to recover a
Challenges in the estimation OT maps. When using standard cost functions such as (xxx := x
2 2. the estimation of OT maps is hindered, in principle, by the curse of dimensionality [Hütter
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

h(z)= 2 kell
1(()) 1llel1 7 lella
h(z))
Tell
h(x))
Ilb 2/1"
A
0.0 0
8-5 g
Base measure X; Level lines of
Transported points Ty (xx)
Figure I: Illustration of ground-truth optimal transport maps with different costs h, for the same base function g. In this experiment, g is the negative of random ICNN with 2-dimensional inputs, 3 layers and hidden dimensions of sizes [8, 8, 8]. All plots display the level lines of g. The optimal transport map Th are recomputed four times using Prop. with four different costs h, displayed above each plot. (left) When h is the usual Lg cost, we observe a typical OT map that follows from each x,, minus the gradient of g. With the é sparsity-inducing regularizer (middle-left), we obtain sparse displacements: most arrows follow either of the two canonical axes, yet some points do not move all. (middle-right) This is slightly different when using the k-overlap norm, which exhibits less shrinkage. With a cost that penalizes displacements that are orthogonal to a vector b, we obtain displacements that push further to the bottom than in the (left) plot, as in the (right) plot, where displacements are almost parallel to b. When b is not known beforehand, and both source and target samples are given, we present procedure to learn adaptively such parameter in $ 5..
g
and Rigollet, 2021]. A simple workaround is to reduce the dimension of input data, using for instance a variational auto-encoder [Bunne et al., 2023], or learning hyperplane projections jointly with OT estimation [Paty and Cuturi, 2019, Weed and Rigollet, 2022, Lin al., 2020, Huang et al., 2021, Lin et al., 2021]. We consider in this work another approach, which explores alternative choices for ground cost c. While OT theory is rife with rich cost structures [Ambrosio and Pratelli, 2003, Ma et al., 2005, Lee and Li, 2012. Figalli et al., 2010, Figalli and Rifford, 2010], that choice has comparatively received far less attention in machine learning, where for vast majority of applications the cost function is often chosen as E2 and sometimes Ez.
Cost structure impacts map structure. While the usage of Riemannian metrics within OT in ML has been considered [Cohen et al., 2021, Grange et al., 2023, Pooladian et al., 2023b], computational challenges restrict these approaches to low-dimensional manifolds. We argue in this work that costs that are translation invariant (TD, c(x, y) == h(x y) with h Idd - R, can offer practitioners a reasonable middle ground, since many numerical schemes developed for the &g cost can be extended to TI costs, both for static and dynamic formulations of OT (see e.g., Villani et al. [2009, Chap.7] 7] or Liu [2022]) In particular, we propose focus on elastic costs of the form (z) = / ||| + (z(), with Y > and 7 Rdd - IR regularizer, following the name of the elastic net regularization [Zou and Hastie, 2005] proposed in the context of regression. Cuturi et al. [2023] show that using elastic costs in OT map estimation results in Monge maps whose displacements satisfy T(x) - = prox, o oV (x). for some potential f, and are therefore shaped by the proximal operator of
Contributions. While elastic costs offer the promise of obtaining OT maps with prescribed structure inherited from the proximal operator of a regularizer т. our current understanding of how to use and exploit such costs limited the experimentation provided in [Cuturi et al., 2023]. This stands in stark contrast with the fine-graie fine-grained characterization provided by Brenier that a map is optimal for the l2 cost if and only if is the gradient of convex potential. To this end:
We show in $ that OT maps can be generated for any elastic cost h by running proximal gradient descent scheme. through the proximal operator of T, on suitable objective. This results in, to our knowledge, the first visualization of Monge maps that extend beyond the usual grad-convex Brenier maps for &2 costs (see Figure 1), as well as synthetic generation high high-dimensions;
We introduce subspace elastic costs in 44 which promote displacements occurring in a low- dimensional subspace spanned by the line vectors of matrix A, A E Rxxd AA"" Lp. setting
T(2) := | We prove sample-complexity estimates for the Lange-Bregman-Oee -occm (MBO) estimator introduced [Cuturi et al., 2023] with this cost (and more generally Mahalanobis costs), and establish a link with the spiked transport model [Niles- Weed and Rigollet, 2022] when ? co.
2

---

Since the choice of the regularizer in the elastic cost gives rise to diverse family of OT maps. whose structural properties are dictated by the choice of regularizer, we consider parametrized families To, and propose in $ 5 a loss to select adaptively suitable 0.
We illustrate all above results, showing MBO estimator performance, recovery of A the basis of i..d. samples, in both synthetic (using our first contribution) and single-cell data tasks, where we demonstrate improved predictive ability compared to baseline estimators that do not learn A.
2 Background: Optimal transport with Elastic Costs
Monge Problem. Let P2 be the set of probability measures with finite second-order moment. We consider in this work cost functions c of the form c(x.y) == h(x - y), where h R + R is strictly convex and, to simplify a few computations, symmetric, i.c., h(z) h(-2). Given two measures μ4, / E (2R)) the Monge problem [1781] seeks a map T Rd Rd minimizing an average transport cost, as quantified by h, of the form:
T = arg min / 1(x - T(x)) m(dx) Typ=v
Jzd
Because the set of admissible maps T is not convex, solving (1) requires taking a detour that involves relaxing (I) into the so-called Kantorovich dual and semi-dual formulations, involving respectively two functions (or only one the case of the
2015, $1.6]:
[ fdp + [ gdv = arg max 3d
{ fåu+ f:R -B, f is f--comcaee
{ fh
(f",g")
arg max f.RR-R f(x) +o
(2)
JRd
A function f is said to be h-coneave if there exists function g such that f is the h-transform of g. i.., f = g", where for any function g Rd - R, we define its h-transform as
xx) = inf 1(x y) g(y).
We recall fundamental theorem in optimal transport [Santambrogio, 2015, $1.3]. Assuming the optimal, h-concave, potential for (2), f*, is differentiable at Xa (this turns out to be a mild assumption since f is a.e. differentiable when h is), we have [Gangbo and McCann, 1996]:
(x) =
= x Vh fx)
where the convex conjugate of reads: h (w) == sup, (z, - The classic Brenier theorem [1991], which is by now staple of OT estimation in machine learning [Korotin et al., 2019, Makkuva al., 2020, Korotin al., 2021, Bunne et al., 2023] through input-convex neural networks [Amos et al., 2017]. is particular example, stating for h = 1 112. that T(x) =
since in this
case, V = (Vh) = ld, see [Santambrogio, 2015, Theorem 1.22].
Maps and Elastic Costs. Cuturi al. [2023] consider TI costs regularizer for y they study elastic costs of the form
hzz) - 1| + 77(2),
(5)
and show that the resulting Monge map is shaped by the proximal operator of T:
1 2 ||w z||2 + 7T(z))
T (xxx)= = prox,, o v (x). where prox.,, (w) == arg
arg min
(6)
The MBO Estimator. While the result above is theoretical, as it assumes knowledge an optimal (MBO) estimator proposes to plug into (6) approximation of f*,
f², the -eee--rreemmneeee recovered from samples from μ and v. We write X X= [x,, ... xx, and Y = [y1, yym for such samples in Rd, possibly weighted by two probability vectors a and of size v and m respectively. transport [Cuturi, 2013]. with so-called entropic
f* can be estimated using entropy-regularey potentials [Pooladian and Niles- Weed, 2021]. This involves choosing a regularization strength E > 0, and solving the following dual problem using the Sinkhorn algorithm [Peyré and Cuturi, 2019, 4.2]:
(f" g.' = (" a. Y, b;
arg max (f, a) + (g. b) - E(e K).).
(7)
fER" gER"
3

---

where Kú = exp(-h(x; - -y)/e). The entropy- regullrzeg optimal transport matrix associated with that cost and on those samples can be derived directly from these dual potentials [Peyré and Cuturi, 2019, Prop. 4.3] as P P'(X, a, Y, b; h e) E IRux with entries at (i, j) equal to:
gj h(x; Y1)
a Y, b; B)l - exp
(8)
E
We now introduce the soft-minimum operator, and its gradient, defined for any vector R as e e-
mine(u) - log uule and Vmin=(u) E Eimee-u
,")): Using vectors (f* gg), we can define estimators fe and Ge for the optimal dual function (f*,g*)
fe x + min: ([x(x - -y,) - -g])), 9e : y minx(( y) +f)a)
(9)
Plugging these approximations into (6) forms the basis for the MBO estimator outlined in Algo. 1. Definition 1 (MBO Estimator). Given data, an elastic cost function h l2 + VTT and solutions to (7), the MBO map estimator [Pooladian and Niles-Weed Weed, 2021, Cuturi et al., 2023] is given by:
(x) =x x - pr (x Σ
- y,) -
(10)
where p(x)= = Vmin -y)) BززرB ) is probability vector
Algorithm MBO-ESTTMATOR X, Y; 7,,,)
1: Set h L/
if y 0, equivalent to [Pooladian+, b Sinkhorn (Eq. 7).
2: (f' =
= D* a, Y, b,,;
3: p lambda: x + softmax [gj h(x - y;)ll//e
4: M =lambda: x Σx= p(x), (T - - -y,)
5: T, = lambda: X prox, ,7 (x+ + M(x)).
6: return: T,
On Ground-Truth Monge Maps for Elastic Costs
Our strategy to compute examples of ground truth displacements for any elastic cost h rests on the following theorem, which a direct consequence of [Santambrogio, 2015, Theorem 1.17].
Proposition 1. Consider potential g : Rd + R. and its h-transform defined in (3). Additionally, set g = ld -hh o Vigh. Then Th the OT Monge map for cost h between μ and (T yp for any
measure μ in P.
The ability to compute OT map for h therefore hinges the ability solve numerically the h-traasformm (3) of a potential function g. This can be done, provably, as long as is concave and smooth, and prox, is available, as shown the following result
Proposition 2. Assume is concave, L-smooth, and that X < 2/L. Setting y x and iterating
y prox Ar ( - [y-x+AVg(y) 1+
(11)
converges to a point y" (x) = arg miny h(x - y) -gl)). Furthermore, we have
-(x)= xxx-xxxxx"(x) g Vg" = Vh(x yx"(x)) and (x) = (x). (12)
Proof. Because h is the sum a quadratic norm with yT, the proximal operator of Ah can be restated in terms of the proximal operator of [Parikh et al., 2014, $2.1.1]. The convergence of iterates (11)
follows from [Beck and Teboulle, 2009, Thm. 1] or [Rockafellar, 1976, Thm. 1]. The final identities given by [Bauschke and Combettes, 2011, Prop. 18.7].

---

Base function g
h(z) =
A(x) 1 111
A(z)
- Llbba '21"
Figure 2: Illustration of the i--transform computation in 2d. (left): base concave potential here negative quadratic. (other figures) Level lines of the corresponding h-transform g" for different choices of h. The h-transform is computed using the iterations described in Prop. 2.
As summarized in Algo.2, the proximal operator of r is the only thing needed to implement iter- ations (11), and, as result, the h-traansfrmm of suitable concave potential. We can then plug the solution (12) to evaluate the pushforward 9 In practice, we use the JAXOPT [Blondel et al., 2021] library to run proximal gradient descent. We illustrate numerically in 2D the re- sulting transport maps for different choices of regularizer in Fig. 1. In this illustration, we use the same base function g, and see clearly the impact of on transforms.
Algorithm 2 GROUND-TRUTHH OT MAP Th
1: Inputs: point x, elastic cost 3/2 7T. potential g-
2: y :=x
3: while not converged do
y-x+AVo(y)
4: y prox A
1+1
5: end while
6: return:
b Thy (x) in (Prop. 1)
4 Subspace Elastic Costs
Recall that for a rank-p p matrix E IRxd .p p d, the projection matrix that maps it to its orthogonal
) is A 7 A ( -' А. When A lies the Stiefel manifold (i.e. AA T), have the simplification A In order to promote displacements that happen within the span of A, we must set a regularizer that penalizes the presence of within its orthogonal complement, namely
AT A. This results in the Pythagorean identity |z|| 1z|2 ||Azz]]
T (z) = 2
(13)
Since TAL quadratic, its proximal operator can obtained by solving a linear system [Parikh et al., 2014, 16.1.1]; developing and using the matrix inversion lemma results in two equivalent quantities
prox 71 (z) (Ia ((-) == ?
(14)
To summarize, given an orthogonal sub-basis A of vectors (each of size d), promoting that a vector z lies in its orthogonal can be achieved by regularizing its norm in the space orthogonal to the span of A. That norm has proximal operator that can be computed by parameterizing full-rank x d matrix, or more simply p x d orthogonal matrix, to recover the suitable proximal operator for TA- in (14). Because that operator is simpler when Sp.d is in the Stiefel manifold,
explicitly, either as
1
Prox, (z)
1+> (La AA
(15)
We propose restrict the study in this work to elastic costs of the form 14 where A E Sp.d- We also present in Appendix alternative parameterizations left aside for future work.
4.1 Statistical Aspects of Subspace Monge Maps
The family of costs (13) is designed to promote transport maps whose displacements mostly lie in low-dimensional subspace of R... In this section, we consider the statistical complexity of estimating such maps from data, assuming A is known. The question of estimating transport maps was first
5

---

studied in statistical context by Hütter and Rigollet [2021], and subsequent research has proposed alternative estimation procedures, with different statistical and computational properties [Deb et al., 2021, Manole et al., 2021, Muzellec et al., 2021, Pooladian and Niles- Weed, 2021]. We extend this line of work by considering the analogous problem for Monge maps with structured displacements.
We show that with a proper choice of e, the MBO estimator outlined in Definition is a consistent estimator of T as n + co, and prove a rate of convergence in L²)).. We also give preliminary theoretical evidence that, as 7 + 0o, maps corresponding to the subspace structured cost 1 T can be estimated at rate that depends only on the subspace dimension p. rather than on the ambient dimension d, thereby avoiding the curse of dimensionaliry.
Sample Complexity Estimates for the MBO Estimator. The MBO estimator is a generalization of the entropic map estimator, originally defined by Pooladian and Niles-Weed [2021] for the quadratic cost h 4.2. This estimator has been statistically analyzed in several regimes, see e.g., [Pooladian et al., 2023a, Rigollet and Stromme, 2022, del Barrio et al., 2022] and [Goldfeld et al., 2022]. We show that this procedure also succeeds for subspace structured costs of the form h 1 e2 + YTA+. As result of being recast as an estimation task for quadratic cost, the following sample-comllxxig result for the MBO estimator follows from [Pooladian and Niles-Weed computation relating the MBO estimator to barycentric projection for the costs we consider (see Appendix B for the full statements, proofs, and applicability to general Mahalanobis norms).
2
2021, Theorem 31, and
Theorem 1. Let A Rpxd be fixed,and suppose has an upper- and lower- bounded density, and μ is upper-bounded, both supported over Ω C Rd compact. Consider T of the form Equation (23) for some y 2 0 fixed, and suppose we have samples ;Iy..., Nn ~ μ and , Yn ~ (TT) Let To be the MBO estimator with E = n di4. Then it holds that
BIJ. - T (T S n uंड
where the underlying constants depend on properties of ,t u,7 and A.
4.2 Connection to the Spiked Transport Model
The additional structure we impose on the displacements allows us closely relate our model the "spiked transport model" as defined by Niles weed and Rigollet [2022]. The authors studied the estima- tion of the Wasserstein distance in the setting where the Brenier map between f and v takes the form,
T Tspikea = A AAx S(Ax)),
(16)
where A E Sp.d and S Rº + IR is the gradient of convex function on R'. Divol et al. [2022] performed a statistical analysis of the map estimation problem under the spiked transport model. They constructed an estimator fn such that the (1) risk decays with respect to the intrinsic dimension p < d; this is summarized in the following theorem.
Theorem 2 (Divol et al., 2022, Section 4.6). Suppose μ has compact support, with density bounded above and below. Suppose further that there exists a matrix A E ixp on the Stiefel manifold such that v= = (Tepiked)3s,, with T spited defined as in Equation (16). Assume that is known explicitly. Given i.i.d. samples from v, there exists estimator Tn, satisfying
El EI - T T (22(2)) Stog(n) n -
(17)
We now argue that the spiked transport model can be recovered in the large 7 limit of subspace structured costs. Indeed, i + oo, then displacements in the subspace orthogonal to A are heavily disfavored, so that the optimal coupling will concentrate on the subspace given by A, thereby recovering map of the form (16), which by Theorem 2 can be estimated at a rate independent of the ambient dimension. Making this observation quantitative by characterizing the rate of estimation of T'* as function of for y large is an interesting question for future work.
5 A Bilevel Loss to Learn Elastic Costs
Following $ 4, we propose a general loss to learn the parameter of a family of regularizers { given source and target samples only. Our goal is to infer adaptively 9 that promotes

---

regular displacements, apply it within the estimation of Monge maps using MBO, and leverage this knowledge to improve prediction quality. Given input and target measures characterized by point clouds X, Y and probability weights a. b, our loss follows a simple intuition: the ideal parameter should be such that the bulk of the OT cost bore by the optimal Monge map, for that cost, dominated by displacements that have low regularization value. Since the only moving piece in our pipeline will be 0, we consider all other parameters constant in the computation of the primal solution, to re-write (8) as:
P*(0)= = P ax, Y, 1 + +T0,e) E R
(18)
Each entry (0)lss quantifies the optimal association strength between pair (x, when the cost is parameterized by 6, where given pair can be encoded as displacement Zij i= Yj X,. For the regularizer 0 to shape displacements, we expect (0) to have a large entry on displacements Zij that exhibit low regularizer To T0(g))) value. In other words, we expect that To To(2g) be as small as possible when P (0) is high. We can therefore consider the loss
Definition 2 (Flastic Costs Loss). Given two weighted point clouds a, X,b, Y, and (0) defined implicitly, as an OT solution in Equation (8), let
( (P* (0), R(0)), with [R(0)];; To (Zay).
(19)
Because P* (0) is itself obtained as the solution to an optimization problem, minimizing is therefore bilevel problem. To solve it, we must compute the gradient V (())), given by the vector-Jacobian
operators PP [:) and aR(.)" )* of P* and R respectively, borrowing notations from [Blondel and Roulet, 2024, $2.3] (see also $ C for walk-through of this identity)
V ( ap (P) + (R())"(P"(O)
(20)
The first operator aP* (+) requires differentiating the solution of an optimization problem (*))). This can be done [Blondel and Roulet, 2024, $10.3.3] using either unrolling of Sinkhorn iterations or using implicit differentiation. We rely on OTT-JAX [Cuturi et al., 2022] to provide that operator, using unrolling. The second operator [- can be trivially evaluated, since it only involves differentiating the regularizer function TO (-). These steps are summarized in Algo. 3.
Learning Subspace Costs. We focus in this section on the challenges arising when optimiz- ing subspace costs, detailed in Section 5. Learning matrix in this context is equivalent to learning a subspace in which the displace- ment between the source and target measures happen mostly in the range of A. As discussed previously, the cost function C(A) should be optimized over the Stiefel manifold [Edelman et al., 1998]. We use Riemannian gradient de- scent [Boumal, 2023] for this task, which iter- ates, for a step-size 1) > 0
Algorithm 3 RECOVER THETA: (X, Y: 7.
1: for 1= 0,.... ... do do
2: Sample mini-batches Xn. Y, from X, Y
3: Compute coupling:
> -Sinkhomm (Eq. 8)
P(O.) + P p* 1 Y 1. + YTB, e).
4:
Compute loss:
(Eq. 19)
[(O,) = ((())).((0.).
5: gradient g < VL((,) using auto-diff.
6: 0.+1 +- -GRAD-UPDATE(O,,.
A + P(A-7&
7: end for
with the Riemannian gradient of L given by (L(A) == G AGT A where: G = VL(A) the standard Euclidean gradient of A computed with automatic differentiation provided in (20); P is the projection on the Stiefel manifold, with formula P(A) )-122) A. These updates ensure
8: return: 0r
that one stays on the manifold [Absil and Malick, 2012].
6 Experiments
Thanks to our ability to compute ground-truth -ruth h-optimal maps presented in $ 3. we generate bench- mark tasks to measure the performance of Monge map estimators. We propose in & 6.1 to test the MBO estimator [Cuturi et al., 2023] when the ground-truth cost h that has generated those benchmarks is known. In $6.2, we consider the more difficult task of learning simultaneously, and as outlined in 5, an OT map and the ground-truth parameter of subspace-clastic cost defined by
7

---

h(())
112211 2 +
R(z))
12 +
M(2)
||al11 +
((2))
111211 2 +
2/12
1.00
1.0
MSE ratio
0.75
0.75
7"
0.50
0.5
0.50
0.8
0.25
10
10a
10a
10
10"
10a
regularizer 7
regularizer "
regularizer TV
regularizer
ground-truth tasks involving the é and TA- = ||A-21 2112 structured costs, where = 2 in dimension d 5 (two figures to the left) and dimension 10 (Iwo figures the right). We display the MSE ratio between the MSE estimated with a regularizer strength > 0 and that in the absence of regularization (i.e., 7 = 0). The level of regularization used for generating the ground-truth data is 7², whereas performance are shown varying w.r.t. We display curves + s.t.d. estimated over 10 random seeds.
Figure 3: Performance of the MBO estimator
matrix A of size p" x d. The cost is parameterized by matrix A of size x d. where p is an estimate the ground-truth subspace dimension p" (usually known), equal to or larger than p". We check with this synthetic task the soundness of the loss (()), Definition 2, and of our Riemannian descent approach by evaluating to what extent the vectors in Á recovers the subspace spanned by 6.3 a direct application of subspace clastic costs to real data, without
A". Finally, we consider any ground-trun knowledge, using perturbations of single-cell data. In this experiment, learns both an OT map and a subspace. Our code implements a parameterized RegTICost class, added to OTT-JAX [Cuturi et al., 2022]. Such costs can be fed into the Sinkhorn solver, and their output cast as DualPotentials objects that can output the Te map given in Definition
pipeline
6.1 MBO on Synthetic Ground- -rund Displacement
In this section, assume that the regularizer J is known, using the same both for generation and estimation, but that the ground-truth regularization strength y used to generate data not known. We use that cost evaluate the transport associated with Ø. sample of points, using Proposition 1, and then compare the performance of Sinkhorn based estimators, either with that cost or the standard L cost (which corresponds to 0).
We consider the & and TĄ= ||A | regularizers, and their associated proximal operators. While we assume knowledge of T the MBO estimator, we do not use the ground-truth regularization Y which generated the data, and instead consider it free parameter. The data is generated following $ by sampling a concave quadratic function g(z) == ] w M w) where M a Wishart matrix, sampled as M QQT,, where Q E Rd is multivariate Gaussian, and w is random Gaussian vector. We then sample n 1024 Gaussian points stored in XT and transport each using the map defined in Proposition 1. computed in practice with Proposition 2. This recovers matched train data Xr and Yr. We do the same for test fold X,, Y, of the same size, to report our metric, the mean squared error (MSE), defined as [X,) using Xr, Y,. We plot this MSE as a function of y, where y 0 corresponds exactly to the MBO using the naked Ez cost. We observe in Figure that the MBO estimator with positive outperforms significantly that using & only, for any range the parameter 7
Y, V.. where T. is obtained from Definition 1
6.2 Recovery of Ground-Trume Subspace Parameters in Elastic Costs
We propose to test the ability of Algorithm to recover the ground-truth A" parameter of regularizer TA= as defined (13). To do so, we proceed as follows: For dimension d, we build the ground-truth cost h by selecting A", sampling p² x d normal matrix that is then re- e--orjecced on the p" x d Stiefel manifold. Next, we sample a random ICNN, and set the base function to be its negative. We then sample point cloud X of n = 512 standard Gaussian points, and apply, following Proposition 1, the corresponding ground-trufe transport to obtain Y of the same size. We tune the regularization parameter y for to ensure that the p" first singular values of displacements Y X captured either

---

insetia disp. wateguace
50%
sabepoce
5.61
7.57
8.30
3.31
13
giandinsh akgiin dreedi
5.10
3.70
4.64
8.85
2.58
1.62
4.85
2.22
10
74
2:77
0.64
0.56
1.01
2.79
2.77
0.03
0:44
0.30
0:33
2:78
0.04
0:42
0.25
0
2.05
2.67
0.19
0.19
0.38
2:97
2.79
0:17
0.08
0.08
0.43
2.97
2.78
17
0.26
0:03
0.26
10:37
2.30
1.21
14.00
gweehich sateei diveda X
2.01
3.50
3.4Z
2.82
0.32
2:37
0:30
0:13
0:79
0.
0:
1.00
1.15
0.52
0.44
0.43
1.43
2.34
2.49
0.30
0.33
0.20
1.14
0.49
0:34
0.06
0:03
18
0:95
0:39
0.17
0.04
0.06
0.12
0:16
0.02
0.01
0.02
0:73
0:35
0.37
0.00
0.09
0.39
0:94
3.35
0.10
0.28
0.08
0.33
0:73
0.35
0.16
0.02
0
0.01
0:13
54
256
12B
$4
10
64
22
16
(313 dimerdion
43t3 dimerttion of
Figure 4: Error averaged over 5 seeded runs (lower better) in [0, 1] of the P x d orthogonal matrix Á recovered by our algorithm, compared to the ground-truth p* x d cost matrix A". Error bars are not shown for compactness, but are negligible since all quantities are bounded below and close O. Dimensions d. p" vary in each of these 6 plots, whereas ý is fixed to either p" (top row) or 1.25p" (bottom row). Error is quantified the normalized squared- rresiuul error obtained when projecting the p" basis vectors of A onto the span of Á. From left to right, the regularization strength > increases to ensure that 50%, 70% and 90% of the total inertia of all displacements generated by the tround-truth Monge map are borne by the p" highest singular values. As expected, recovery is easier when ý is slightly larger than p" (bottom) compared to being exactly equal (top). It is also easier as the share of inertia captured by p" increases.
50%, 70% 90% of the total inertia. We expect that the larger this percentage, the easier recovery should be. See $ D for details.
We launch our solver fed with these datasets with a subspace dimension preset in advance to either 2 (overbudgeting). We measure recovery of A by Á through the average (normalized by the basis size) of the residual error, when projecting the vectors in A in the span of the basis Å, namely || - ΛÄT A 111//pp /p". For simplicity, we report performance after 1000 iterations of Riemannian gradient descent, with a stepsize 7 of 0.1/ Di.1/vi+II at iteration All results in Figure 4 agree with intuition in the way performance varies with d, p", P.. More importantly, with errors that are often below one percent, we can be confident that our algorithm is sound. We observe that most underperforming experiments could be improved using early stopping.
I p" (matching ground truth)
6.3 Learning Displacement Subspaces for Single-Cell Transport
We borrow the experimental setting in [Cuturi et al., 2023], using single-cell RNA sequencing data from [Srivatsan et al., 2020]. The original dataset shows the responses of cancer cell lines to 188 drug perturbations, downsampled to the drugs (Belinostat, Dacinostat, Givinostat, Hesperadin, and Quisinostat) that have the largest effect. After various standard pre-pree pre-processings (dropping low- variability genes, and using a log(1 +) scaling), we project the dataset to = 256 directions using PCA. In Table 1, we report the total number of cells used for experiments after pre-processing. We then use 80% train/20% test folds to benchmark two MBO estimators: that computed using the &g cost, and ours, using clastic subspace cost, following the learning pipeline outlined in 5. We plot

---

the Sinkhorn divergence (cf. Feydy et al. [2019]) for the Eg cost for reference (see the documentation OTT-JAX [Cuturi al., 2022]).
Belinostat
Dacinostat
Givinostat
Hesperadin
Quisinostat
Sinkharn divergence
19.9
22.5
21.6
91.2
19.8
21.0
19.7-
19:4
22.2
22.2 23.6
11..2 100
7
7
Figure 5: Predictive performance of the MBO estimator on single-cell datasets, d 256, using either the naive baseline lg cost (black dotted line) or elastic subspace cost (13), with varying y and Remarkably, promoting displacements to happen in subspace of much lower dimension improves predictions, even when measured in the squared-Fuce Euclidean distance.
Table I: Number of cells for each cell line and drug/control splits.
Control
Dac.
Giv.
Bel.
Hes.
Quis.
A549
3274
558
703
669
436
475
K562
3346
388
589
656
624
339
MCF7
6346
1562
1805
1684
882
1520
Conclusion. In this work, we proposed an algorithmic mechanism design ground-trus -rutt transports for elastic costs. As first application, we were able to successfully benchmark the MBO estimator of [Cuturi al., 2023] on two tasks (involving the & and an orthogonal projection norm), showcasing the versatility of the MBO framework. Next, we demonstrated our ability to leverage posspace-penalizin costs to learn displacement subspaces by solving an inverse OT problem. We showed successful numerical performance of the MBO estimator when the subspace is known but the regularization strength is not, but also that we were able to learn the ground-truth subspace. We foresee several open directions, the most encouraging being considering other learnable proximal operators beyond subspace approaches, and cementing connections and distinctions between subspace regularized transport (where displacements happen in subspace) vs. the spiked transport model (where all points are projected on a subspace prior to being transported)
References
P-A Absil and Jérôme Malick. Projection-like like retractions on matrix manifolds. SIAM Journal on Optimization, 22(1):135-158 2012.
Luigi Ambrosio and Aldo Pratelli. Existence and stability results in the 11 theory of optimal transportation. Optimal Transportation and Applications: Lectures given at the CIME Summer School, held in Martina Franca, Italy, September 2-8, 2001, 2003.
Brandon Amos, Lei Xu, and Zico Kolter. Input Convex Neural Networks. volume 34, 2017.
Heinz Bauschke and Patrick L Combettes. Convex analysis and monotone operator theory in Hilbert spaces, volume 408. Springer, 2011.
Amir Beck and Mare Teboulle. A fast iterative shrinkage-thresshollnng algorithm for linear inverse problems. SIAM journal on imaging sciences, 2(1):183-202, 2009.
Mathieu Blondel and Vincent Roulet. The elements of differentiable programming. arXiv preprint arXiv:2403. /4606, 2024.
Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares-López. Fabian Pedregosa, and Jean-Philippe Vert. Efficient and modular implicit differentiation. arXiv preprint arXiv: 2021.
10

---

Nicolas Boumal. An introduction to optimization on smooth manifolds. Cambridge University Press, 2023.
Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Commu- nications on Pure and Applied Mathematics, 44(4), 1991. doi: 10. 10002/cpa.31044040000
Charlotte Bunne, Stefan G. Stark, Gabriele Gut, Jacobo Sarabia del Castillo, Mitch Levesque, Kjong-Van Van Lehmann, Lucas Pelkmans, Andreas Krause, and Gunnar Rätsch. Learning single-cell perturbation responses using neural optimal transport. Nature Methods, 20, 2023.
Charlotte Bunne. Geoffrey Schiebinger, Andreas Krause, Aviv Regev, and Marco Cuturi. Optimal transport for single-cell and spatial omics. Nature Reviews Methods Primers, 4(1):58, 2024.
Samuel Cohen, Brandon Amos, and Yaron Lipman. Riemannian convex potential maps. In Interna- tional Conference on Machine Learning, pages 2028-20220 2038. PMLR, 2021.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in neural information processing systems, pages 2292- 2300, 2013.
Marco Cuturi, Laetitia Pang-aaapaxntsoss, Yingtao Tian, Charlotte Bunne, Geoff Davis, and Olivier Teboul. Optimal transport tools (ott): A jax toolbox for all things wasserstein. arXiv preprint arXiv:2201. 12324, 2022.
Marco Cuturi, Michal Klein, and Pierre Ablin. Monge, bregman and occam: Interpretable optimal transport in dighedimessson with feature-sparse maps. In Proceedings of the 40th ICML, 2023.
Nabarun Deb, Promit Ghosal, and Bodhisattva Sen. Rates of estimation of optimal transport maps using plug-in estimators via barycentric projections. arXiv preprint arXiv:2107.01718, 2021.
Eustasio del Barrio, Alberto Gonzalez-San -anzz Jean- Michhl Loubes, and Jonathan Niles- Weed. An improved central limit theorem and fast convergence rates for entropic transportation costs. arXiv preprint arXiv:2204.09105. 2022.
Vincent Divol, Jonathan Niles Weed, and Aram-Alexande Pooladian. Optimal transport map estimation in general function spaces. arXiv preprint arXiv:2212.03322,. 2022.
Alan Edelman, Tomás Arias, and Steven T Smith. The geometry of algorithms with orthogonality constraints. SIAM journal Matrix Analysis and Applications, 20(2):303-333, 1998.
Jean Feydy, Thibault Séjourné, François- Xavier Vialard, Shun-ichi Amari, Alain Trouvé, and Gabriel Peyré. Interpolating between optimal transport and mmd using sinkhorn divergences. In The 22nd International Conference Artificial Intelligence and Statistics, pages 2681-2690 PMLR, 2019.
Alessio Figalli and Ludovic Rifford. Mass transportation sub-riemannian manifolds. Geometric and functional analysis, 20: 119-119, 2010.
Alessio Figalli, Ludovic Rifford, and Cédric Villani. On the ma-trudinger- curvature on surfaces. Calculus of Variations and Partial Differential Equations, 39:307-332, 2010.
Wilfrid Gangbo and Robert J McCann. The geometry of optimal transportation. Acta Mahematica, 177(2):113-16)) 1996.
Ziv Goldfeld, Kengo Kato, Gabriel Rioux, and Ritwik Sadhu. Limit theorems for entropic optimal transport maps and the sinkhor divergence. arXiv preprint arXiv:2207.0862 2022.
Gene H Golub and Charles F Van Loan. Matrix computations. JHU press, 2013.
Daniel Grange, Mohammad Al-Jarrah, Ricardo Baptista, Amirhossein Taghvaei, Tryphon Georgion, Sean Phillips, and Allen Tannenbaum. Computational optimal transport and filtering on riemannian manifolds. IEEE Control Systems Letters, 2023.
Minhui Huang, Shiqian Ma, and Lifeng Lai. A riemannian block coordinate descent method for computing the projection robust wasserstein distance. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 4446 1446-4455, PMLR, 2021.

---

Jan- Chri-Chistiin Hütter and Philippe Rigollet. Minimax estimation of smooth optimal transport maps. The Annals of Statistics, 49(2), 2021.
Dominik Klein, Giovanni Palla, Marius Lange, Michal Klein, Zoe Piran, Manuel Gander, Laetitia Peang-apaxanthos, Michael Stem, Aimée Cestidas-Ponncee Marta Tarquis-Medina, Heiko Lickert, Mostafa Bakhti, Mor Nitzan, Marco Cuturi, and Fabian Theis. Mapping cells through time and space with moscot. bioRxiv, 2023. doi: 10.1 10.1110//023,0000 11122,005.1... 1.540344 URL https://ww. biorxiv.org wttpg//ontent/earryyee http://aaaaaaee2//////////222220 05. 5403744
Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safi,, and Evgeny Burnaev.
Wasserstein-2 generative
2eenertatv networks. 2019.
Alexander Korotin, Lingxiao Li, Aude Genevay, Justin Solomon, Alexander Filippov, and Evgeny Burnaev. Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark. 2021.
Paul WY Lee and Jiayong Li. New examples satisfying ma-trudinger-wang conditions. SIAM Journal on Mathematical Analysis, 44(1):61-73, 2012.
Tianyi Lin, Chenyou Fan, Nhat Ho, Mareo Cuturi, and Michael Jordan.
Projection robust wasserstein distance and riemannian optimization. Advances in neural information processing systems, 33: 9383-9397, 2020.
Tianyi Lin, Zeyu Zheng, Elynn Chen, Marco Cuturi, and Michael Jordan. On projection robust
optimal transport: Sample complexity and model misspecification. In International Conference Artificial Intelligence and Statistics, pages 262-270. PMLR, 2021.
Qiang Liu. Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint arXiv:2209.14:: 14577, 2022.
Xi-Nan Ma, Neil Trudinger, and Xu-Jia Wang. Regularity of potential functions of the optimal transportation problem. Archive for rational mechanics and analysis, 177:151-183, 2005.
Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. volume 37, 2020.
Tudor Manole, Sivaraman Balakrishnan, Jonathan Niles Weed, and Larry Wasserman. Plugin estimation of smooth optimal transport maps. arXiv preprint arXiv:2107. 12364, 2021.
Gaspard Monge. Mémoire sur la théorie des déblais et des remblais. Histoire de I' l'Académie Royale des Sciences, 1781.
Boris Muzellec, Adrien Vacher, Francis Bach, François- Xavier Vialard, and Alessandro Rudi. Near-optimal estimation of smooth transport maps with kernel sums of-uure--quares. arXiv preprint 12.19721 2021.
Jonathan Niles-Weed and Philippe Rigollet. Estimation of wasserstein distances in the spiked transport model. Bernoulli, 28(4):2663-2688 2022.
Neal Parikh, Stephen Boyd, et al. Proximal algorithms. Foundations and trends@ in Optimization, (3):127-239, 2014.
François-Pierre Paty and Marco Cuturi. Subspace robust wasserstein distances. arXiv preprint arXiv:1901.08949, 2019.
Gabriel Peyré and Marco Cuturi. Computational optimal transport. Foundations and Trends in Machine Learning, 11(5-6), 2019. ISSN 1935-8245.
Aram-Alexandre Pooladian and Jonathan Niles-Weed. Entropic estimation of optimal transport maps. arXiv preprint arXiv:2109. 2021.
Area-Alexandre Pooladian, Vincent Divol, and Jonathan Niles Weed. Minimax estimation of discontinuous optimal transport maps: The semi-discrete case. arXiv preprint arXiv:2301.1 / 2023a.
12

---

Aram-A Alexandee Pooladian, Carles Domingo-Enriche Ricky TQ Chen, and Brandon Amos. Neural optimal transport with Lagrangian costs. ICML Workshop New Frontiers Learning, Control, and Dynamical Systems, 2023b.
Philippe Rigollet and Austin Stromme. On the sample complexity of entropic optimal transport. arXiv preprint arXiv:2206.13472. 2022.
R Tyrrell Rockafellar. Monotone operators and the proximal point algorithm. SIAM journal control and optimization, 14(5):877-898, 1976.
Filippo Santambrogio. Optimal transport for applied mathematicians, Springer, 2015.
Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Optimal-transport transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 176(4):928 943, 2019.
Richard Sinkhorn. A relationship between arbitrary positive matrices and doubly stochastic matrices. Ann. Math. Statist,, 35:876-873 1964.
Sanjay R Srivatsan, José L. MicFaline-Figueron Vijay Ramani, Lauren Saunders, Junyue Cao, Jonathan Packer, Hannah A Pliner, Dana L Jackson, Riza M Daza, Lena Christiansen, et al. Massively multiplex chemical transcriptomies at single-cell resolution. Science, 2020.
Alexander Tong, Jessie Huang. Guy Wolf, David Van Dijk, and Smita Krishnaswamy. TrajectoryNet: A dynamic optimal transport network for modeling cellular dynamies. In Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 9526-995 9536. PMLR, 2020.
Cédric Villani al. Optimal transport: old and new, volume 338. Springer, 2009.
Hui Zou and Trevor Hastie. Regularization and variable selection via the elastic net. Journal of the Royal Statistical Society Series B: Statistical Methodology, 67(2):301-33)) 2005.
A
More on Subspace Elastic Costs
Recall that for a rank-p matrix A E Rexd p S d, the projection matrix that maps it to its orthogonal
A- = 1 -00 (AAA")-11A A. When A lies in the Stiefel manifold (i.e. AA I), we have the
I AT A. This results in the Pythagorean identity ||z||² || z + || as intended. In order to promote displacements that happen within the span of A, we must set regularizer that penalizes the presence of within its complement:
simplification
T (2) = + = 2 (A-) 1+z
(Ta- A (AA") A)z.
Since TA4 is a quadratic form, its proximal operator can be obtained by solving linear system [Parikh et al., 2014, 16.1.1]; developing and using the matrix inversion lemma results in
prox (z) (Ia+>(A "== (4
AA
(21)
To summarize, given an orthogonal sub-basis A of vectors (each of size d), promoting that a vector z lies in its orthogonal can be achieved by regularizing its norm in the orthogonal of A. That norm has a proximal operator that can be computed either by
Parameterizing A implicitly, through an explicit parameterization of an orthonormal basis B for A matrix directly specified in the (d - p) p Stiefel manifold. This can alleviate computations to obtain a closed form for its proximal operator:
1 (Ba_11Ba) Bz
-11-11--1
BT TB-
prox (z) = prOx, (z) = z =-B Bz (2- 1+7
1-
z.
but requires storing B, (d - p) x orthogonal matrix, which is cumbersome when p < d.
13

---

2. Parameterizing A explicitly, either full-rank d matrix, or more simply a x d orthogonal matrix, recover the suitable proximal operator for TAL. by either
(a) Falling back on the right-most expression (14) in the linear solve, which can be handled using sparse conjugate gradient solvers, since the application of the right-most linear operator has complexity (p + 1) x and is positive definite, in addition to the linear solve complexity Opp"). This simplifies when A is orthogonal, A Sp.d since in that case,
yT (z) 1+7
(22)
(b) Alternatively, compute matrix in the (d - p) p Stiefel manifold that spans the same linear space as, through the Gram-Schmidt process [Golub and Van Loan, 2013, p.254] of the d x matrix A- or rank p, B Srram-Schmidt// (-+). to fall back on the expression above.
B Proofs from Section 4.1
To perform this analysis, we rely on the following characterization of optimal maps for subspace structured costs, which reveals close connection with optimal maps for the standard E cost.
Proposition 3. Let T be the optimal map between μ and v for the cost h = [ + TTAL Denote by W the linear map - ((1 7)] optimal map) between Wyp and Wiv. Equivalently, T is h-optimal if and only if it can be writlen
/2xx Then W o TT* W is the Brenier map (i.e.. lg
T of o
(23)
where T is the gradient of convex function.
Proof of Proposition 3. The cost 1 + YTAI can be written as
(I- = W
The optimal transport problem we consider is therefore equivalent minimizing
/ 111=' - 3'
min 2 Wy dyllddr(,, xEI(Wpp.W..) (pL
min
(24)
Brenier's theorem implies that the solution to the latter problem is given by the gradient of a convex function, and that this property uniquely characterizes the optimal map. Writing this function as T, we obtain that the optimal coupling between Wyp and W₁y is given by y = =(<x), which implies that the optimal h-coupling between f and y is given by T W-1 o T W, desired.
The proof of Theorem requires the following two lemmas.
Lemma For costs of the form lu(2) 1=2 Bz where B is positive definite, the MBO estimator between two measures μ and can be written the barycentric projection of the corresponding optimal entropic coupling.
Proof. Note that h² (w) f B-ww, and thus Vh'(w) B- Let (fe, denote the optimal entropic potentials for this cost, with corresponding coupling Zg. Borrowing computations from e.g., Proposition 2 of Pooladian and Niles Weed [2021], we compute
fe(x) = / B(x y)d Br BB ydT
where (y) is the conditional entropic coupling (given x). The proof concludes by taking the expression of the MBO estimator and expanding:
BB-1(1 Bx B 10-BByddr yddn((s)) = yydT"(u) (v)
Tx(x) = x (Vh") a (Vfc(z)) r -B-1
which is the definition of the barycentric projection of π. for given
14

---

Lemma 2 (Pre-conditionin of MBO). Let T be the MBO estimator between and v for the cost h yeth +TTAI. Let W be denoted as Proposition Then the MBO estimator written as
T: = W T W.
(25)
where Tg is the barycentric projection between Wap and W.v.
Proof. The proof here is similar Proposition 3, which outline again for completeness. As before, we are interested in solutions to the optimization problem
|l Wyll (d, 9) (
min (ET()
with optimal coupling me. Performing a change of variables
(W @9 W)Fm, we have
ʃ || v|l (a, v) <KL( 8/), min x' T'E
where μí == Wen (and similarly for v/), where now the optimizer reads (x²)*. The two optimal plans are related as
n (W W "+T²'"
was established in Lemma that the MBO estimator T given by the barycentric projection
(x) Ex; [Y|X x||
Performing the change of variables Y WY and X W W,, we can re-write this as function of
Ty instead:
Th = Eny [V = x]
(' æ]
-W -W (Y = Wa]
= W ww'f((Wx))
where we identify T,(:) (oy) X -|; this completes the proof.
We are now ready to present the main proof.
Proof of Theorem I. Let Te.n denote the MBO estimator between samples from and v, and let T fe.. denote the entropic map estimator from samples Jí' == Wgn and v' == Wgy, where W has spectrum 0 < Amia (W) & Amax (w)) < +00, where we have access to W since A is known.
Our goal is to establish upper bounds on
[[Te,n T 1T (122(u)) ||W (Tem W - T W)
Paying for constants that scale like Amaxx (W-1), we have the bound
[Te, T (2: Sw I T112200)
where we can now directly use the rates of convergence from [Pooladian and Weed, 2021, Theorem 3], uı' satisfies our regularity assumptions under the conditions we have imposed on W. this completes the proof.
C Gradient of Elastic Cost Loss
The gradient of the loss L in Definition 2 can be recovered through simple aggregation of weighted gradients
VL(0)-
Σ
[(()),,, V (0),) + (O),,), V. [R(0)),)
(26)
15

---

To write this formula in more compact way, it is sufficient to notice that, adopting the convention that 0 be a parameter in RR, and introducing an arbitrary vector R",
(VL()), = [R(O] (Vo [P (0)]s) : + (P)(0) (Vo [R(O);s) w)
= (R(()) (Vo (0)); )))),)))) (P' (Vo [R(O),s m))
The products of all coordinate wise gradients with is equivalent to the application of the Jacobians of and P We write J and J. for these Jacobians, both being maps taking 0 as input, and matrix. As consequence one
outputting linear map JaR: IR IR" " i.. JgR(0) is has
(VL(0).w) (R(0), JoP" (PP) because these maps are linear, one also has
(VL(0),
PT (0)
= P* (]) R(0))))00))
which gives the identification given in the main text.
D Additional Details on Experiments
In $ 6.2, and unlike Figure 3, do not choose predefined value for for-", but instead select it with the following procedure: we start with a small value for nt = 0.1, and increase it gradually, until a certain desirable criterion these displacements goes above threshold. To measure this, first compute the (paired) matrix of displacements on given sample,
D [TT (x)) x,]: e Rnxd
We then consider ratio of singular values on subspace (to select for ||A+- 112), writing for the vector singular values of D, ranked decreasing order, to compute
sv-ratio(7) = Σ asl Σ os E [0,1].
(27)
D.1 Synthetic experiments
In Algorithm 5, we outline way to generate ground-truth data used in synthetic experiments mentioned $ 6.1 and Figure 3. For a given regularization strength fy and a regularizer Te. we first sample concave function g. along with the ground-truth parameters for the regularizer 0* We then sample n source points X from the standard normal distribution and push them through the ground-truth OT map to get target points Y, described in Algorithm 4.
Algorithm 4 GT-SAMPLES(X:: go)
Algorithm 5 SYNTHETIC- OT- TASK()), on regularizers To-
1: Inputs: points X, elastic cost h, po-
1: Sample concave function (e.g. minus ICNN).
tential Yo Rd lR
2: Sample (* parameter
when 7 admits parameters
2: for i, x in enumerate(X) do
3: Sample n source points X (e.g. Gaussian).
3: y; < Th (x) & Algorithm 2
4: Create n targets GT-SAMPLES(X. _ + yT8*.g)
4: end for
5: return: X, Y. optionally GT parameters 0*.
5: return: Y [yl,... yn]
D.2 Experimental Procedure
Experiments in Section 6.1, MBO Synthetic Ground Truth Displacement
For both = l and FA = ||4+ I12, run SYNTHETIC-OT-TASS (Algo. 5) to form samples
X, Y (store A for T.).
16

---

Using these samples, benchmark MBO estimator informed by the cost structure: varying using automatically scaled E, and ground truth parameter .A" when studying TĄ., for 10 random splits of X, Y into X Yuain and X Yoa
Run MBO-ESTIMATOR (Alg. -y-"
using regularization) on X Yimin, store MSE₁
Σ; [Telh, =x(
Run AIO--SSTIATOR (Alg. without regularization) on X inin Yurain store MSEg 1
>; [0, O, (x! -
Report ratio MSE//MSE, /MSE2, showing MBO seeded with the right regularizer always outperforms original entropic map.
Experiments in Section 6.2, Recovery of Ground-Truth Subspace Parameters in Elastic Costs
For TA ||A+ _1 run SYNTHETIC-OT-TASS (Algo. 5) to generate paired samples X, Y in R Tune y to have displacements concentrated (50%, 70%, 90%) in a subspace, as described in Appendix D. Recover A of dimension p"
Run RECOVER-THETA X, Y using arbitrary y to output A. Display Average recovery error 1/12/"
||A*
Experiments in Section 6.3, Learning Displacement Subspaces for Single-Cell Transport
Here data come from real measurements, X E Raxd Y E ggxd nt + m. Split into Xtrain tiuin and X Yiess.
Run RECOVER- RRCOVER-THETA on X Yuain using arbitrary y and random initialization with varying p.
Output Á.
Run MBO-ESTIMATON (Alg. using regularization) on Xtrain Yirainn Form predictions ý;
Te(7, (x!a))
Compute & Sinkhorn divergence between point cloud (y:) and (y y
To benchmark (dotted-line in Fig. 5), run two steps above, but setting 0 (no subspace regularization).
17

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract describes the challenges underlying the selection of a cost in optimal transport theory. We advertise in that abstract way to benchmark such OT problems for a fairly wide class of exotic costs, and claim statistical consistency results and new algorithm to learn that cost' parameters. All these claims substantiated by simple experiments.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. It is fine to include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
This is discussed in the conclusion section.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
18

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: While theory is only part of our contribution (our paper proposes also new algorithm to identify parameters, experimented on synthetic and real data), we have provided, to the best of our best ability, the proofs of these theoretical results in the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proof' can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: By interfacing our code with the OTT-JAX toolbox, and using public datasets, we believe these resuts and claims can be fully reproduced, with some effort. We will release the entire codebase for experiments in coming weeks, as python notebookssuutoriaas.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often good way to accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
19

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Data is open. We provide code to tools. We do not share all of the more complex config files that were required to run all computations on GPU nodes.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cc.
public/ Poolic//gides/Cooe wwweeeubbmssioooool..o
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / /nips.ccc..cccc wtt.cccccccccccuudcccooooo gwtp///dde///ee// SwwoddeSumissiooooooe
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) is recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Yes, we have provided such details in the description of experiments.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to level detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We have for a part of synthetic results. We also have error stds stored for Figure 4 but decided not to plot them given space constraints and the very low value obtained by positive quantities (i.e. mean close means the std also close to 0).
Guidelines:
20

---

The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi-
dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative етror rates).
If error bars are reported in tables or plots. The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Although claim is made in terms of compute performance, the fairly small scale of the experiments allows to execute these runs on single GPU.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: : htt///euurpp..// cesuurps.ccppultee httppplic/tticcaciieeee..
Answer: [Yes]
Justification:
We have followed these ethics guidelines.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [Yes]
21

---

Justification:
We discuss these impacts in the conclusion.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification: NA
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
All data resources have permissive licenses. We refer to the citations.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
22

---

The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. www.coo////oo//o /ataattss has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors are encouraged reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: NA
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: NA
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
23

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
24


Reviews:
[{"review_id": "nJZuHea4t8", "reviewer": "Reviewer_EFNS", "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.", "questions": "* Why is it called elastic costs?\n* What are $\\mathbf{a}$ and $\\mathbf{b}$ in equation 7?\n* What is $\\tilde{T}$ in Theorem 1?\n* 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.\n* Related to the previous question, but can the results of section 4 be extended directly to general squared Mahalanobis norms?", "limitations": "The authors have addressed their limitations.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* Clearly and well-written (except for some notational confusion that I mention in the Questions section)\n* 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?)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are very grateful to your detailed reviews and many thought-provoking questions.\n\n> **i) does the true $\\theta$ optimize the elastic costs loss when the data is drawn from, e.g., $T_g^h$**\n\nWe cannot prove the ground truth parameter is the global optimum of our loss, but we do observe empirically, in *all* our synthetic experiments, that the loss of the ground truth parameter lower bounds that of our iterates.\n\nTo follow your intuition, we agree that if the regularizer $\\tau\\_\\theta$ admits a parameter $\\theta\\_{\\text{null}}$ such that $\\tau\\_{\\theta\\_{\\text{null}}}$ is identically 0, then our loss would be trivially minimized for $\\theta\\_{\\text{null}}$. This cannot happen with the Stiefel-manifold parameterization presented in the paper. For regularizers where this can occur, our loss might need a contrastive term to avoid such trivial $\\theta\\_{\\text{null}}$.\n\nA possible way to give more intuition for $\\mathcal{L}$ (**Eq.19**) is to cast it in a continuous formulation. With, as you suggest, a ground truth map $T_{g_0}^{h_{\\theta_0}}$, obtained with: a ground truth parameter $\\theta_0$ to parameterize costs and a ground truth potential $g_0$, the loss for parameter $\\theta_0$ is (using Prop. 1 and 2) reads:\n\n$$\\mathcal{L}(\\theta_0) = \\int_{\\mathbb{R^d}} \\tau_{\\theta_0}(x - y^\\star_{\\theta_0}(x)) \\mu(dx),$$\n\nWhere \n\n$$y^\\star(x)\\_{\\theta\\_0} := \\arg\\min\\_{y} \\frac{1}{2} \\|x-y\\|^2 - g(y) + \\tau\\_{\\theta\\_0}(x-y),$$ \n\nFor a different $\\theta$, with the same source/target measures, we would need to recompute first an OT map from $\\mu$ to $\\nu_0:=T_{g}^{h_{\\theta_0}}\\sharp \\mu$ using cost $h\\_\\theta$,\n\n$$T^\\star[\\theta] = \\min\\_{T, T\\sharp\\mu = \\nu\\_0}\\int h\\_{\\theta}(x-T(x))\\mu(dx),$$\n\nand use that map to define the loss,\n\n$$\\mathcal{L}(\\theta)= \\int \\tau\\_{\\theta}\\left(x-T^\\star\\[\\theta\\](x)\\right)\\mu(dx).$$\n\n> **ii) what is the sample complexity of parameter recovery?)**\n\nThis sounds hard. In the much simpler spiked model, $\\gamma\\rightarrow \\infty$ (S4.2), for the projected Wasserstein metrics, existing results [Niles-Weed et al. 22, Lin et al.21] do not study the statistical complexity of parameter recovery, but only that of the modified distances between sample vs. population measures (taking a max, or even a mean as [Lin et al.21] over subspace projections).\n\n> **Why is it called elastic costs?**\n\nThe $\\ell^2_2$ squared norm + regularizer term form in **(Eq.5)** was first introduced as the elastic net (https://en.wikipedia.org/wiki/Elastic_net_regularization) by [Zhou & Tibshirani, 2005], we will cite them.\n\nThis split ($\\ell^2_2$ + regularizer) is natural, and leads to explicit *displacements* in Eq. (6) where $T(x) = x - \\textrm{prox}_{\\tau}(\\cdots)$. Splitting the cost as $\\ell^2_2 +\\tau$ leads to the first term in $x$, and a displacement inheriting the “regularity” of $\\tau$ encoded in its prox, as shown in [Cuturi et al.'23]\n\n> **What are $\\mathbf{a}$ and $\\mathbf{b}$ in equation 7?**\n\nWe went too fast, apologies. In **L. 203**, *“Given input and target measures characterized by point clouds $\\mathbf{X}, \\mathbf{Y}$ and probability weights $\\mathbf{a}, \\mathbf{b}$”* we mean that the data is given as two weighted point clouds, as $$\\mu=\\sum_{i=1}^n a_i \\delta\\_{\\mathbf{x}\\_i},\\quad\\nu=\\sum_{j=1}^m b_j \\delta\\_{\\mathbf{y}\\_j}.$$ \n$(a_i), (b_j)$ are probability vectors summing to 1. In our experiments, $a_i$ and $b_j$ are uniform weights, $1/n$ and $1/m$, with $n$ and $m$ changing.\n\n> **What is $\\tilde{T}$ in Theorem 1?**\n\nWe apologize for this. $\\tilde{T}$ is defined in **L.462**, in Theorem 1's proof. It is the Brenier OT map between the source and target measures pushforwarded by the linear map $W$ in **L.456**, defined from matrix $A$. We will simplify the exposition of this theorem to focus on the rate given in between **L.177 and 178**, since $\\tilde{T}$ only influences constants.\n\n> **Am I missing something or is the setting of section 4 very simple given the pre-existing results for $h=\\tfrac12\\|\\cdot\\|^2$? In particular, because $h$ is a squared Mahalanobis-norm [...] 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.**\n\nThe elastic cost seeded with regularizer $\\|A^{\\perp} z\\|^2$ in **(Eq.13)** is indeed a Mahalanobis norm as written in **L.456**, and the results from Thm1 follow from the proof technique you mention. This is a new result that sheds an important light on the subspace cost function, and that as you point out relies on the results of [Pooladian & Weed, 2021].\n\nWe also want to stress that our parameterization $\\ell_2^2 + \\gamma\\|A^{\\perp}x\\|^2_2$ is instrumental in two aspects:\n- it allows to establish a link with spiked transport models when $\\gamma\\to\\infty$\n- it enables the learning of a subspace $A$ using the loss **Eq.19** described in Sec.5, either to recover a ground-truth parameter $A^*$ used to generate synthetic data (in Sec.6.2, Fig.4), or as a parameterized transport that improves predictive performance on real data (Sec.6.3).\n\n> **can the results of section 4 be extended directly to general squared Mahalanobis norms?**\n\nAs you correctly point out, the results from Sec.4.1 can be applied to any Mahalanobis norm $h(x) = \\frac12 x^T Wx$ with $W$ positive definite. We will clarify this in the text.\n\nThe interest of our parameterization becomes, however, fully apparent when trying to *learn* subspaces **for displacements**, while still estimating an OT map valid in full dimensionality between samples. Learning a \"valid\" full-rank $W$ directly would be intractable computationally for large $d$; learning a low-rank $W$ would be degenerate, and is equivalent to the spiked transport model (Sec.4.2). The spiked model only studies OT maps once the measures have been *entirely* projected in a lower-dim. subspace. We would not know how to extend Sec. 5 to *learn* the parameters of a general Mahalanobis distance.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SG2LVmi6tz", "reviewer": "Reviewer_JR35", "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.", "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?", "limitations": "I do not see the limitations discussed anywhere.", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "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?  \n\n2. 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.\n\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewers for their time and for formulating many interesting questions.\n\n> **A numerical method is formulated for estimating the Monge maps w.r.t elastic costs.**\n\nWe believe there might be a minor confusion here. You are referring to the proximal gradient method in Prop. 2. Please notice this method is only used to construct ground-truth OT maps for elastic costs (as mentioned in the title of Section 3, **L.113**, *On Ground Truth Monge Maps for Elastic Costs*). This method cannot, therefore, be used to *estimate* Monge maps (this was done by [Cuturi+23] with their MBO method), but only to create examples for which one can recover Monge maps for elastic (not $\\ell^2_2$) costs.\n\n> **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.**\n\nPlease see our response above. Prop.2 is only used to generate ground truth OT examples. It starts from a cost and any arbitrary concave functions $g$ to define $T^h_g$ in **L.117**. By contrast, the MBO estimator only takes samples $\\mathbf{X}, \\mathbf{Y}$ weighted by $\\mathbf{a}, \\mathbf{b}$.\n\n\n> **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?**\n\nWe went too fast, apologies, this will be improved, notably with any additional page we might get. \n\nIn **L. 203**, *“Given input and target measures characterized by point clouds $\\mathbf{X}, \\mathbf{Y}$ and probability weights $\\mathbf{a}, \\mathbf{b}$”* : we mean that the data is given as two weighted point clouds, as $$\\mu=\\sum_{i=1}^n a_i \\delta\\_{\\mathbf{x}\\_i},\\quad\\nu=\\sum_{j=1}^m b_j \\delta\\_{\\mathbf{y}\\_j}.$$ \n$(a_i), (b_j)$ are probability vectors summing to 1. In our experiments, $a_i$ and $b_j$ are uniform weights, $1/n$ and $1/m$, with $n$ and $m$ changing.\nWe will also improve other notation, see e.g. comments to Rev. EFNS and zNJm.\n\n> **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.**\n\nLet us mention again that in real-life applications  we do not have access to the ground-truth. The method in Section 3 is therefore usefuless with real data. The method in Section 3 was only used to benchmark the MBO estimator (in Fig. 3) and our \"learning elastic costs\" bit, in Fig. 4. For instance, this method is not used at all on single-cell data, S.6.3, Fig. 5\n\n> **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.**\n\nThe results from Thm1 indeed follow from the work of [Pooladian & Weed, 2021], using a change of variable. While the proof technique is simple, this is a new result that sheds an important light on the subspace cost function.\n\nWe also want to stress that our parameterization $\\ell_2^2 + \\gamma\\|A^{\\perp}x\\|^2_2$ is instrumental in two aspects:\n- it allows to establish a link with spiked transport models when $\\gamma\\to\\infty$\n- it enables the learning of a Stiefel matrix $A$ using the loss mentioned in Sec.5. This approach recovers the subpace ground-truth matrix $A^*$ from samples, as demonstrated in Fig.4, and is useful to predict better transport on real single-cell data in dimension $d=256$ in Fig.5.\n\nWhile the cost presented in Section 4 is simple (it is indeed a quadratic norm), a crucial piece is its parameterization, which fits well with our method in Section 5, to recover a subspace cost. In that sense, Section 4 only takes its full meaning when paired with Section 5.\n\n> **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?**\n\nOur method does not constraint displacements to happen in a subspace of a dimension $p$. We use a regularization mechanism (illustrated in Fig. 1) that promotes displacements taking place *mostly* in that subspace (depending on $\\gamma$). Notice that this is the biggest difference with subspace projection methods mentioned in **L.37**, which can be seen as a degenerate / collapsed case of our method.\n\nBecause our algorithms uses the data in its original space dimension, and Theorem 1 makes no specific assumption on the ground-truth generating mechanism, there is no immediate way to leverage $p$. This result is there to highlight that the MBO estimator, in the case where $A$ is known, is sound.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pQw9fmf3P4", "reviewer": "Reviewer_zNJm", "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", "questions": "Eq (8), is this really an entropy term? Seems more like a kernel function. \n\nIs 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?\n\nWhat 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?\n\nWhat does MBO stand for?\n\nWhat 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?\n\nWhat are the p hats and such in Fig 4 signifying?", "limitations": "no negative societal impact", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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. \n\nThere 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are very thankful for your encouraging grades and for your great comments.\n\n> **A convex optimization problem is proposed […]**\n\nThere is a mild confusion that we would like to clarify: While the task of computing the entropic map given a cost $h$ and samples $\\mathbf{X}, \\mathbf{Y}$ weighted as $\\mathbf{a}, \\mathbf{b}$ is convex (**Eq.10**), as you point out, the main task we tackle is that of *learning* the parameters of a regularizer $\\tau_\\theta$ (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.\n\n> **It is unclear to me the novelty of the work, but I am not aware of any major competing work.**\n\nTo the best of our knowledge, the task of estimating an OT map whose displacements mostly happen in a subspace is indeed new (S.5).\n\nCrucially, without devising a scheme for computing ground-truth maps for elastic costs (S.3), it would be literally impossible to verify our methodology. Our contributions are (i) a method for computing ground-truth OT mappings for arbitrary elastic costs $h$ (**L.10**), (ii) a method to estimate the parameters of such costs (**L.15**), (iii) demonstrated applicability in both synthetic and real settings **L.18**.\n\n> **I think this paper would benefit from having some kind of \"algorithms box\"**\n\nWe thank you for this great suggestion, and added these algorithms in the attached pdf.\n\n> **There should be some numerical comparisons of this method compared to others [...]. An obvious one is OT without the subspace constraint.**\n\nWe 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 $\\ell^2_2$ cost (**L.284**). The other lines are \"our\" MBO estimators computed with learned costs. Remarkably, they perform better, even in $\\ell^2_2$ sense.\n\n> **Are there also non-OT methods that are competitive in this space?**\n\nNot to our knowledge. One might speculate that flow-matching methods could be extended to estimate a map, **knowing** $A^*$ (Sec.6.1), but we do not see how they could *also* learn a subspace (Sec.6.2)\n\n> **Eq (8), is this really an entropy term?.**\n\nThis duality relation links the (primal) entropy-regularized OT solution, with dual potentials in **Eq. 7**, see **[Peyre & Cuturi, 19, Prop. 4.3 & proof]**. We will clarify.\n\n> **Is this prox operator really easy to use for all types of tau?**\n\nA good question! The proximal operator is known in *closed-form* for a huge family of regularizers $\\tau$. See for example [1]\n\n> **Also, are there non-euclidean distances that you look at when looking at this elastic prox?**\n\nThe 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.\n\n> **What happens if $A$ is not on the Stiefel manifold?** \n\nAn inverse pops up in the computation of the prox, **Eq.14**. That inverse disappears when $A$ is in the Stiefel manifold, **Eq.15**. Since we compute a gradient w.r.t. $A$ of the loss $\\mathcal{L}$ in **Eq.19**, if $A$ were not in the Stiefel manifold, one would need to differentiate through this inverse.\n\n> **are there cases where you may want some \"directional penalization\" but not necessarily wiping out all movement in that direction?**\n\nThis is a very good question. Note that we do not \"wipe out\" all movements in the orthogonal of $A$, we penalize this (uniformly on all directions in $A$) with $\\gamma$. This is illustrated in the two rightmost plots in **Fig.1**: they only differ in their penalization $\\gamma, \\tfrac12$ vs $\\tfrac{10}{2}$: The right most displacements agree “more” with $b$, i.e. their norm $\\| \\mathbf{b}^\\perp \\cdot\\|^2$ is *smaller*, because $\\gamma$ is larger.\n\nIf $\\gamma\\rightarrow \\infty$, one would all displacements (and points) would only along $b$, which would be equivalent to projecting all data on $b$ first. This is the spiked model discussed in S.4.2.\n\n> **Is it challenging to constrain A to be on the stiefel manifold?**\n\nAt each Riemannian gradient update, this only requires the SVD of a $p\\times d$ matrix. This is usually not challenging, because $p\\ll d$, we did not encounter challenges. We will detail further **L.224**.\n\n> **What does MBO stand for?**\n\nMBO stands for Monge-Bregman-Occam, the estimator in **[Cuturi et al 2023]** that extends the entropic map **[Pooladian and Niles-Weed, 2021]** to arbitrary translation invariant costs $h$. We went too fast, we will clarify.\n\n> **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?**\n\nThe approach outlined in S.5 requires computing Sinkhorn solutions, and differentiate them. This has compute/memory footprint in $O(Kn^2)$, where $n$ is the batch-size and $K$ depends on the regularization $\\varepsilon$ 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.\n\n> **What are the p hats and such in Fig 4 signifying?**\n\nA ground-truth $A^*$ has dimension $p^* \\times d$. Our goal is to recover the subspace in $A^*$. However, we will likely never know $p^*$ in practice. Yet, we need to choose the size of our estimator, $\\hat{A}$, set as $\\hat{p} \\times d$.\n\nIn synthetic settings, we picked $\\hat{p}$ to be exactly $p^*$, or a bit larger (the latter is more realistic in practice) **L.270**. We show that, regardless, we are able to accurately recover the subspace described in the ground-truth matrix $A^*$, **L.272**.\n\nOn real data (Fig. 5), we only considered $\\hat{p}=2,4,8$.\n\n[1]: G. Chierchia, E. Chouzenoux, P. L. Combettes, and J.-C. Pesquet. \"The Proximity Operator Repository. https://proximity-operator.net/\".", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "japrJPW3gR", "reviewer": "Reviewer_3gJN", "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)$.\n\nElastic 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.", "questions": "I do not have any questions.", "limitations": "Authors have addressed limitations.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are very grateful for your careful review and positive grade, and were happy to see that you mentioned several times the originality of our contribution.\n\n> **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.**\n\nWe appreciate this comment. We believe the ability to create examples for such unusual costs was untapped.\n\n> **The results, particularly the improved performance of the MBO estimator and the recovery of \n in synthetic and single-cell data tasks, provide empirical validation of the proposed methods.**\n\nWe appreciate this comment. We were also surprised to see quasi-perfect recovery of the subspaces, notably when the displacements are not that peaked to belong to a certain subspace, as in the left-most plot.\n\n> **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.**\n\nWe agree that OT has, by now, too much jargon which slows reading for first time readers. As was highlighted by Reviewer **zNJm** for MBO, we will revisit the introduction section to ensure all terms are extensively defined.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We would like to thank all reviewers for the quality of their review, and for taking the time to formulate many insightful questions. It was a pleasure on our end to do our best during this rebuttal week to clarify some of these concerns. As a result, our rebuttal text is fairly long, and we thank again reviewers for their time reading our answers.

Many reviewers have mentioned the *complexity of our paper*. Indeed, the task we consider is fairly complex, and is at the forefront of research in OT: to our knowledge, ours is the first work that considers both structured priors for costs (beyond $\ell^2_2$), studies ground truth OT maps for such costs, and proposes a path to recover these parameters exclusively from samples. This works both in synthetic settings, but also, crucially, reaches demonstrated gains on real datasets when compared to the standard $\ell^2_2$ pipelines.

Because of this, we agree that some parts may not have been clear, notably in our experiments where even the parts that concern synthetic experiments are fairly complex. We agree with Reviewer **zNJm** that taken together, these multiple moving parts may be challenging, and that algorithmic boxes would have helped. 

Following this very good suggestion, **we have drafted an "Algorithms" section** in the attached pdf. With these algorithms, we have tried to strike a balance between rigor (specially when it comes to low-level routines) and intuition (to describe data generation). We have then mapped succinctly our experiments to these algorithms, to highlight what parts matter for each task. For instance, our parameter recovery pipeline **Recover-Theta** can be used with the **MBO** estimator (Section 6.3), but can also be used on its own (Section 6.2) if one only seeks a subspace to gather interpretability.

We have also added two figures in the pdf.
- **Figure A will be added to the current Figure 1**. This shows a ground truth OT map beyond $\ell_1$ and the subspace norm, using the $k$-overlap norm, as was formerly explored in [Cuturi et al. 23]. This illustrates better, in our opinion, the interest of being able to generate ground truth displacements for arbitrary structured norms.
- **Figure B will replace Figure 4**. While we were already positively impressed by the ability of our recovery pipeline to estimate the ground truth subspace $A^*$ exclusively from samples (using **Recover-Theta**, Algorithm 5 in pdf), we noticed that, looking at experiments in more detail, the few failure cases we were observing were simple due to suboptimal seedings for the initialization. Running the algorithm 3 times, with 3 different initializations, and keeping the parameter with lowest loss $\mathcal{L}$ significantly improves performance (as with, e.g., $k$-means).

We have also ran more experiments on drugs, adding 2 more datasets. We do not report these for lack of space, but observe similar trends: our subspace regularizer has a predictive edge.

Many thanks again for your time,

The Authors.

Meta-review: Overall reviewers felt that this paper makes a solid contribution towards extending our understanding of Monge maps under regularized Euclidean loss functions. The paper should serve as a foundation for more follow up work on optimal transport problems with such loss functions. While reviewers felt that the paper was largely well-written, they made several suggestions for improvements that we hope the authors will take into account when preparing the final version.
Decision: Accept (poster)