Title: Learning Elastic Costs to Shape Monge Displacements
Venue: NeurIPS 2024
Accepted: yes
Keywords: optimal transport
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

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, ℓ 2 2 ( x , y ) := 1 2 ∥ x -y ∥ 2 2 . The benefits of using elastic costs, defined using a regularizer τ as c ( x , y ) := ℓ 2 2 ( x , y )+ τ ( x -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 ( x ) -x , by giving them a structure that matches that of the proximal operator of τ . 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 groundtruth 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 ℓ 2 2 cost; (ii) We propose a loss to learn the parameter θ of a parameterized regularizer τ θ , and apply it in the case where τ A ( z ) := ∥ A ⊥ z ∥ 2 2 . This regularizer promotes displacements that lie on a low-dimensional subspace of R d , spanned by the p rows of A ∈ R p × 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.

1 Introduction

Finding efficient ways to map a distribution of points onto another is a low-level task that plays a crucial role across many machine learning (ML) problems. Optimal transport (OT) theory [Santambrogio, 2015] has emerged as a 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 1 , . . . , x n ) and ( y 1 , . . . , y m ), to recover a map T : R d → R d that is 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 T ( x i ) is not too far, on average from x i . The notion of efficiency can be made precise by choosing a real-valued cost function c that compares a point x and its mapping via c ( x , T ( x )) ∈ R .

Challenges in the estimation of OT maps. When using standard cost functions such as ℓ 2 2 ( x , y ) := 1 2 ∥ x -y ∥ 2 2 , the estimation of OT maps is hindered, in principle, by the curse of dimensionality [Hütter

+

γ

‖

z

‖

1

h

(

z

) =

1

2

‖

z

‖

2

2

h

(

z

) =

1

2

‖

z

‖

2

+

γ

‖

z

‖

1

h

(

z

) =

1

2

‖

z

‖

2

+

γ

‖

z

‖

ovk

h

(

z

) =

1

2

‖

z

‖

2

+

10

2

‖

b

⊥

z

‖

2

b

Base measure

x

i

Level lines of

g

Transported points

T

h

g

(

x

i

)

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, Niles-Weed and Rigollet, 2022, Lin et 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 a vast majority of applications the cost function is often chosen as ℓ 2 2 and sometimes ℓ 2 .

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 (TI), c ( x , y ) := h ( x -y ) with h : R d → R , can offer practitioners a reasonable middle ground, since many numerical schemes developed for the ℓ 2 2 cost can be extended to TI costs, both for static and dynamic formulations of OT (see e.g., Villani et al. [2009, Chap.7] or Liu [2022]). In particular, we propose to focus on elastic costs of the form h ( z ) = 1 2 ∥ z ∥ 2 + γτ ( z ) , with γ > 0 and τ : R d → R a 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 OTmapestimation results in Monge maps whose displacements satisfy T ( x ) -x = -prox τ ◦∇ f ( 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 is limited to the experimentation provided in [Cuturi et al., 2023]. This stands in stark contrast with the fine-grained characterization provided by Brenier that a map is optimal for the ℓ 2 2 cost if and only if it is the gradient of a convex potential. To this end:

We show in § 3 that OT maps can be generated for any elastic cost h by running a proximal gradient descent scheme, through the proximal operator of τ , on a 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 2 costs (see Figure 1), as well as synthetic generation in high-dimensions;

We introduce subspace elastic costs in § 4, which promote displacements occurring in a lowdimensional subspace spanned by the line vectors of a matrix A , A ∈ R p × d , AA T = I p , setting τ ( z ) := ∥ A ⊥ z ∥ 2 2 . We prove sample-complexity estimates for the Monge-Bregman-Occam (MBO) estimator introduced in [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 γ →∞ .

Since the choice of the regularizer τ in the elastic cost gives rise to a diverse family of OT maps, whose structural properties are dictated by the choice of regularizer, we consider parametrized families τ θ , and propose in § 5 a loss to select adaptively a suitable θ .

We illustrate all above results, showing MBO estimator performance, recovery of A on the basis of i.i.d. samples, in both synthetic (using our first contribution) and single-cell data tasks, where we demonstrate an improved predictive ability compared to baseline estimators that do not learn A .

2 Background: Optimal transport with Elastic Costs

Monge Problem. Let P 2 ( R d ) 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 d → R is strictly convex and, to simplify a few computations, symmetric, i.e., h ( z ) = h ( -z ) . Given two measures µ, ν ∈ P 2 ( R d ) , the Monge problem [1781] seeks a map T : R d → R d minimizing an average transport cost, as quantified by h , of the form:

Because the set of admissible maps T is not convex, solving (1) requires taking a detour that involves relaxing (1) into the so-called Kantorovich dual and semi-dual formulations, involving respectively two functions (or only one in the case of the semi-dual)[Santambrogio, 2015, §1.6]:

A function f is said to be h -concave if there exists a function g such that f is the h -transform of g , i.e., f = ¯ g h , where for any function g : R d → R , we define its h -transform as

We recall a fundamental theorem in optimal transport [Santambrogio, 2015, §1.3]. Assuming the optimal, h -concave, potential for (2), f ⋆ , is differentiable at x 0 (this turns out to be a mild assumption since f ⋆ is a.e. differentiable when h is), we have [Gangbo and McCann, 1996]:

Maps and Elastic Costs. Cuturi et al. [2023] consider TI costs w.r.t. a regularizer τ : for γ > 0 they study elastic costs of the form where the convex conjugate of h reads: h ∗ ( w ) := sup z ⟨ z , w ⟩ -h ( z ) . The classic Brenier theorem [1991], which is by now a staple of OT estimation in machine learning [Korotin et al., 2019, Makkuva et al., 2020, Korotin et al., 2021, Bunne et al., 2023] through input-convex neural networks [Amos et al., 2017], is a particular example, stating for h = 1 2 ∥ · ∥ 2 2 , that T ( x ) = x -∇ f ⋆ ( x 0 ) , since in this case, ∇ h = ( ∇ h ) -1 = Id , see [Santambrogio, 2015, Theorem 1.22].

and show that the resulting Monge map is shaped by the proximal operator of τ :

The MBO Estimator. While the result above is theoretical, as it assumes knowledge of an optimal f ⋆ , the Monge-Bregman-Occam (MBO) estimator proposes to plug into (6) an approximation of f ⋆ , recovered from samples from µ and ν . We write X = [ x 1 , · · · , x n ] and Y = [ y 1 , · · · , y m ] for such samples in R d , possibly weighted by two probability vectors a and b of size n and m respectively. f ⋆ can be estimated using entropy-regularized transport [Cuturi, 2013], with so-called entropic potentials [Pooladian and Niles-Weed, 2021]. This involves choosing a regularization strength ε > 0 , and solving the following dual problem using the Sinkhorn algorithm [Peyré and Cuturi, 2019, § 4.2]:

where K ij = exp( -h ( x i -y j ) /ε ) . The entropy-regularized optimal transport matrix associated with that cost h and on those samples can be derived directly from these dual potentials [Peyré and Cuturi, 2019, Prop. 4.3] as P ⋆ ( X , a , Y , b ; h, ε ) ∈ R n × m with entries at ( i, j ) equal to:

We now introduce the soft-minimum operator, and its gradient, defined for any vector u ∈ R q as

Using vectors ( f ⋆ , g ⋆ ) , we can define estimators ˆ f ε and ˆ g ε for the optimal dual function ( f ⋆ , g ⋆ ):

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 = ℓ 2 2 + γτ and solutions to Eq. (7) , the MBO map estimator [Pooladian and Niles-Weed, 2021, Cuturi et al., 2023] is given by:

where p ( x ) := ∇ min ε ([ h ( x -y j ) -g ⋆ j ] j ) is a probability vector.

Algorithm 1 MBO-ESTIMATOR ( X , Y ; γ, τ, ε )

3 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 is a direct consequence of [Santambrogio, 2015, Theorem 1.17].

Proposition 1. Consider a potential g : R d → R and its h -transform as defined in (3) . Additionally, set T h g := Id -∇ h ∗ ◦ ∇ ¯ g h . Then T h g is the OT Monge map for cost h between µ and ( T h g ) ♯ µ for any measure µ in P ( R d ) .

The ability to compute an OT map for h therefore hinges on the ability to solve numerically the h -transform (3) of a potential function g . This can be done, provably, as long as g is concave and smooth, and prox τ is available, as shown in the following result

Proposition 2. Assume g is concave, L -smooth, and that λ < 2 /L . Setting y = x and iterating

converges to a point y ⋆ ( x ) = arg min y h ( x -y ) -g ( y ) . Furthermore, we have

Proof. Because h is the sum of a quadratic norm with γτ , the proximal operator of λh 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 are given by [Bauschke and Combettes, 2011, Prop. 18.7].

Base function

g

h

(

z

) =

1

2

‖

z

‖

2

2

h

(

z

) =

1

2

‖

z

‖

2

+

γ

‖

z

‖

1

h

(

z

) =

1

2

‖

z

‖

2

+

10

2

‖

b

⊥

z

‖

2

b

As summarized in Algo.2, the proximal operator of τ is the only thing needed to implement iterations (11), and, as a result, the h -transform of a suitable concave potential. We can then plug the solution in (12) to evaluate the pushforward T h g . In practice, we use the JAXOPT [Blondel et al., 2021] library to run proximal gradient descent. We illustrate numerically in 2 D the resulting 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 c on h transforms.

4 Subspace Elastic Costs

Recall that for a rankp matrix A ∈ R p × d , p ≤ d , the projection matrix that maps it to its orthogonal is A ⊥ = I -A T ( AA T ) -1 A . When A lies on the Stiefel manifold (i.e. AA T = I ), we have the simplification A ⊥ = I -A T A . This results in the Pythagorean identity ∥ z ∥ 2 = ∥ A ⊥ z ∥ 2 + ∥ A z ∥ 2 . In order to promote displacements that happen within the span of A , we must set a regularizer that penalizes the presence of z within its orthogonal complement , namely

Since τ A ⊥ is quadratic, its proximal operator can be obtained by solving a linear system [Parikh et al., 2014, §6.1.1]; developing and using the matrix inversion lemma results in two equivalent quantities

To summarize, given an orthogonal sub-basis A of p 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 a proximal operator that can be computed by parameterizing A explicitly , either as a full-rank p × d matrix, or more simply a p × d orthogonal matrix, to recover the suitable proximal operator for τ A ⊥ in (14). Because that operator is simpler when A ∈ S p,d is in the Stiefel manifold,

We propose to restrict the study in this work to elastic costs of the form 14 where A ∈ S p,d . We also present in Appendix A 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 a low-dimensional subspace of R d . 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

Algorithm 2 GROUND-TRUTH OT MAP T h g

1: Inputs : point x , elastic cost h = 1 2 ℓ 2 2 + γτ , potential g .

2: y := x

3: while not converged do

4: y ← x +prox λγ λ +1 τ ( y - x + λ ∇ g ( y ) 1+ λ )

5: end while

6: return: y ▷ T h g ( x ) in (Prop. 1)

studied in a 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 ε , the MBO estimator outlined in Definition 1 is a consistent estimator of T ⋆ as n → ∞ , and prove a rate of convergence in L 2 ( µ ) . We also give preliminary theoretical evidence that, as γ →∞ , maps corresponding to the subspace structured cost 1 2 ℓ 2 2 + γτ A ⊥ can be estimated at a rate that depends only on the subspace dimension p , rather than on the ambient dimension d , thereby avoiding the curse of dimensionality .

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 = 1 2 ℓ 2 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 2 ℓ 2 2 + γτ A ⊥ . As a result of being recast as an estimation task for quadratic cost, the following sample-complexity result for the MBO estimator follows from [Pooladian and Niles-Weed, 2021, Theorem 3], and a computation relating the MBO estimator to a barycentric projection for the costs we consider (see Appendix B for the full statements, proofs, and applicability to general Mahalanobis norms).

Theorem 1. Let A ∈ R p × d be fixed,and suppose ν has an upper- and lower- bounded density, and µ is upper-bounded, both supported over Ω ⊆ R d compact. Consider T ⋆ of the form Equation (23) for some γ ≥ 0 fixed, and suppose we have samples x 1 , . . . , x n ∼ µ and y 1 , . . . , y n ∼ ( T ⋆ ) ♯ µ . Let ˆ T ε 1

be the MBO estimator with ε ≍ n -d +4 . Then it holds that where the underlying constants depend on properties of µ, ν, γ and A .

4.2 Connection to the Spiked Transport Model

The additional structure we impose on the displacements allows us to closely relate our model to the 'spiked transport model" as defined by Niles-Weed and Rigollet [2022]. The authors studied the estimation of the Wasserstein distance in the setting where the Brenier map between µ and ν takes the form,

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 ∈ R p × d on the Stiefel manifold such that ν := ( T spiked ) ♯ µ , with T spiked defined as in Equation (16) . Assume that µ is known explicitly. Given n i.i.d. samples from ν , there exists an estimator ˆ T n satisfying where A ∈ S p,d and S : R p → R p is the gradient of a convex function on R p . Divol et al. [2022] performed a statistical analysis of the map estimation problem under the spiked transport model. They constructed an estimator ˆ T n such that the L 2 ( µ ) risk decays with respect to the intrinsic dimension p ≪ d ; this is summarized in the following theorem.

We now argue that the spiked transport model can be recovered in the large γ limit of subspace structured costs. Indeed, if γ →∞ , 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 a 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 a function of γ for γ 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 a θ 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, is dominated by displacements that have a low regularization value. Since the only moving piece in our pipeline will be θ , we consider all other parameters constant in the computation of the primal solution, to re-write (8) as:

Each entry [ P ⋆ ( θ )] ij quantifies the optimal association strength between a pair ( x i , y j ) when the cost is parameterized by θ , where a given pair can be encoded as a displacement z ij := y j -x i . For the regularizer θ to shape displacements, we expect P ⋆ ( θ ) to have a large entry on displacements z ij that exhibit a low regularizer τ θ ( z ij ) value. In other words, we expect that τ θ ( z ij ) to be as small as possible when P ⋆ ij ( θ ) is high. We can therefore consider the loss

Definition 2 (Elastic Costs Loss) . Given two weighted point clouds a , X , b , Y , and P ⋆ ( θ ) defined implicitly, as an OT solution in Equation (8) , let

Because P ⋆ ( θ ) is itself obtained as the solution to an optimization problem, minimizing L is therefore a bilevel problem. To solve it, we must compute the gradient ∇L ( θ ) , given by the vector-Jacobian operators ∂P ⋆ ( · ) ∗ [ · ] and ∂R ( · ) ∗ [ · ] of P ⋆ and R respectively, borrowing notations from [Blondel and Roulet, 2024, §2.3] (see also § C for a walk-through of this identity)

Learning Subspace Costs. We focus in this section on the challenges arising when optimizing subspace costs, as detailed in Section 5. Learning matrix A in this context is equivalent to learning a subspace in which the displacement between the source and target measures happen mostly in the range of A . As discussed previously, the cost function L ( A ) should be optimized over the Stiefel manifold [Edelman et al., 1998]. We use Riemannian gradient descent [Boumal, 2023] for this task, which iterates, for a step-size η > 0

The first operator ∂P ⋆ ( · ) ∗ [ · ] requires differentiating the solution of an optimization problem P ⋆ ( θ ) . 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 ∂R ( · ) ∗ [ · ] can be trivially evaluated, since it only involves differentiating the regularizer function τ θ ( · ) . These steps are summarized in Algo. 3.

with the Riemannian gradient of L given by ˜ T )

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 ) = ( AA ⊤ ) -1 / 2 A . These updates ensure that one stays on the manifold [Absil and Malick, 2012].

6 Experiments

Thanks to our ability to compute ground-truth h -optimal maps presented in § 3, we generate benchmark 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 a subspace-elastic cost defined by

10

-

1

10

1

10

3

regularizer

γ

0

.

25

0

.

50

0

.

75

1

.

00

MSE ratio

γ

∗

h

(

z

) =

1

2

‖

z

‖

2

2

+

γ

2

‖

A

⊥

z

‖

2

2

10

-

1

10

1

regularizer

γ

0

.

50

0

.

75

1

.

00

γ

∗

h

(

z

) =

1

2

‖

z

‖

2

2

+

γ

‖

z

‖

1

10

0

10

3

regularizer

γ

0

.

8

1

.

0

γ

∗

h

(

z

) =

1

2

‖

z

‖

2

2

+

γ

2

‖

A

⊥

z

‖

2

2

10

-

1

10

1

regularizer

γ

0

.

5

1

.

0

γ

∗

h

(

z

) =

1

2

‖

z

‖

2

2

+

γ

‖

z

‖

1

a matrix A ∗ of size p ∗ × d . The cost is parameterized by a matrix ˆ A of size ˆ p × d , where ˆ p is an estimate of the ground-truth subspace dimension p ∗ (usually not known), equal to or larger than p ∗ . We check with this synthetic task the soundness of the loss L ( θ ) , Definition 2, and of our Riemannian descent approach by evaluating to what extent the ˆ p vectors in ˆ A recovers the subspace spanned by A ∗ . Finally, we consider in § 6.3 a direct application of subspace elastic costs to real data, without any ground-truth knowledge, using perturbations of single-cell data. In this experiment, our pipeline 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 T ε map given in Definition 1.

6.1 MBOon Synthetic Ground-Truth Displacement

In this section, we assume that the regularizer τ is known , using the same τ both for generation and estimation, but that the ground-truth regularization strength γ ∗ used to generate data is not known. We use that cost to evaluate the transport associated with ¯ g h ε on a sample of points, using Proposition 1, and then compare the performance of Sinkhorn based estimators, either with that cost or the standard 1 2 ℓ 2 2 cost (which corresponds to γ = 0 ).

We consider the τ = ℓ 1 and τ A ⊥ = ∥ A ⊥ z ∥ 2 2 regularizers, and their associated proximal operators. While we assume knowledge of τ in the MBO estimator, we do not use the ground-truth regularization γ ∗ which generated the data, and instead consider it a free parameter. The data is generated following § 3 by sampling a concave quadratic function g ( z ) := -1 2 ( z -w ) T M ( z -w ) where M is a Wishart matrix, sampled as M = QQ T , where Q ∈ R d × 2 d is multivariate Gaussian, and w is a random Gaussian vector. We then sample n = 1024 Gaussian points stored in X T and transport each using the map defined in Proposition 1, computed in practice with Proposition 2. This recovers matched train data X T and Y T . We do the same for a test fold X t , Y t of the same size, to report our metric, the mean squared error (MSE), defined as ∥ T ε ( X t ) -Y t ∥ 2 2 , where T ε is obtained from Definition 1 using X T , Y T . We plot this MSE as a function of γ , where γ = 0 corresponds exactly to the MBO using the naked ℓ 2 2 cost. We observe in Figure 3 that the MBO estimator with positive γ outperforms significantly that using ℓ 2 2 only, for any range of the parameter γ .

6.2 Recovery of Ground-Truth Subspace Parameters in Elastic Costs

We propose to test the ability of Algorithm 3 to recover the ground-truth A ∗ parameter of a regularizer τ A ⊥ as defined in (13). To do so, we proceed as follows: For dimension d , we build the ground-truth cost h by selecting A ∗ , sampling a p ∗ × d normal matrix that is then re-projected on the p ∗ × d Stiefel manifold. Next, we sample a random ICNN, and set the base function g to be its negative. We then sample a point cloud X of n = 512 standard Gaussian points, and apply, following Proposition 1, the corresponding ground-truth transport to obtain Y of the same size. We tune the regularization parameter γ for τ , to ensure that the p ∗ first singular values of displacements Y -X captured either

64

32

16

8

4

2

1

ˆ

p

=

p

∗

ground-truth subspace dimension

p

∗

15.04

5.61

7.57

5.10

3.78

4.44

1.88

1.51

1.74

3.85

2.77

0.64

0.56

1.01

1.79

2.95

2.67

0.19

0.19

0.33

0.75

0.57

0.73

0.12

0.09

1.02

50% inertia of disp.

in subspace

6.52

4.36

3.11

4.89

2.58

1.61

1.87

1.42

1.29

0.88

2.77

0.63

0.44

0.38

0.33

2.97

2.79

0.17

0.08

0.08

0.43

0.75

0.60

0.71

0.07

0.02

0.21

0.26

70% inertia of disp.

in subspace

5.25

4.13

2.54

4.85

2.22

1.10

1.90

1.46

1.07

0.36

2.78

0.64

0.42

0.26

0.08

2.97

2.78

0.17

0.06

0.03

0.06

0.75

0.60

0.70

0.07

0.00

0.02

0.03

90% inertia of disp.

in subspace

256

128

64

32

16

8

data dimension

d

64

32

16

8

4

2

1

ˆ

p

=

5

4

p

∗

ground-truth subspace dimension

p

∗

10.17

2.01

3.98

1.12

0.86

2.03

0.79

0.47

0.52

1.59

1.15

0.52

0.44

0.43

1.43

0.95

0.41

0.18

0.16

0.29

0.73

0.15

0.17

0.06

0.09

0.59

256

128

64

32

16

8

data dimension

d

2.33

1.02

0.82

0.89

0.37

0.39

0.74

0.35

0.21

0.30

1.14

0.49

0.36

0.13

0.26

0.95

0.39

0.17

0.04

0.06

0.12

0.74

0.15

0.16

0.03

0.03

0.13

256

128

64

32

16

8

4

data dimension

d

1.21

0.83

0.32

0.85

0.28

0.13

0.74

0.33

0.15

0.06

1.14

0.49

0.34

0.06

0.03

0.95

0.40

0.16

0.02

0.01

0.02

0.73

0.15

0.16

0.02

0.01

0.01

0.15

2.00

4.00

6.00

8.00

10.00

12.00

14.00

50%, 70% or 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 ˆ p preset in advance to either ˆ p = p ∗ (matching ground truth) or ˆ p = 5 4 p ∗ (overbudgeting). We measure recovery of A ∗ by ˆ A through the average (normalized by the basis size) of the residual error, when projecting the vectors in A ∗ in the span of the basis ˆ A , namely ∥ A ∗ -ˆ A ˆ A T A ∗ ∥ 2 2 /p ∗ . For simplicity, we report performance after 1000 iterations of Riemannian gradient descent, with a stepsize η of 0 . 1 / √ i +1 at iteration i . 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.

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 5 drugs (Belinostat, Dacinostat, Givinostat, Hesperadin, and Quisinostat) that have the largest effect. After various standard pre-processings (dropping lowvariability genes, and using a log(1 + · ) scaling), we project the dataset to d = 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 ℓ 2 2 cost, and ours, using an elastic subspace cost, following the learning pipeline outlined in § 5. We plot

the Sinkhorn divergence (cf. Feydy et al. [2019]) for the ℓ 2 2 cost for reference (see the documentation in OTT-JAX [Cuturi et al., 2022]).

10

0

10

2

10

4

γ

19

.

6

19

.

7

19

.

8

19

.

9

Sinkhorn divergence

Belinostat

10

0

10

2

10

4

γ

22

.

2

22

.

4

22

.

6

22

.

8

Dacinostat

10

0

10

2

10

4

γ

23

.

6

23

.

8

24

.

0

24

.

2

24

.

4

Givinostat

10

0

10

2

10

4

γ

19

.

2

19

.

4

19

.

6

19

.

8

Hesperadin

10

0

10

2

10

4

γ

22

.

2

22

.

4

22

.

6

22

.

8

Quisinostat

ˆ

p

= 2

ˆ

p

= 4

ˆ

p

= 8

SqEucl.

Conclusion. In this work, we proposed an algorithmic mechanism to design ground-truth transports for elastic costs. As a first application, we were able to successfully benchmark the MBO estimator of [Cuturi et al., 2023] on two tasks (involving the ℓ 1 and an orthogonal projection norm), showcasing the versatility of the MBO framework. Next, we demonstrated our ability to leverage subspace-penalizing 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 a 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 retractions on matrix manifolds. SIAM Journal on Optimization , 22(1):135-158, 2012.

Luigi Ambrosio and Aldo Pratelli. Existence and stability results in the l1 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 J Zico Kolter. Input Convex Neural Networks. volume 34, 2017.

Heinz H Bauschke and Patrick L Combettes. Convex analysis and monotone operator theory in Hilbert spaces , volume 408. Springer, 2011.

Amir Beck and Marc Teboulle. A fast iterative shrinkage-thresholding 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.14606 , 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:2105.15183 , 2021.

Nicolas Boumal. An introduction to optimization on smooth manifolds . Cambridge University Press, 2023.

Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Communications on Pure and Applied Mathematics , 44(4), 1991. doi: 10.1002/cpa.3160440402.

Charlotte Bunne, Stefan G. Stark, Gabriele Gut, Jacobo Sarabia del Castillo, Mitch Levesque, Kjong-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 International Conference on Machine Learning , pages 2028-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 Meng-Papaxanthos, 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 high-dimensions 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-Sanz, Jean-Michel 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-Alexandre Pooladian. Optimal transport map estimation in general function spaces. arXiv preprint arXiv:2212.03722 , 2022.

Alan Edelman, Tomás A Arias, and Steven T Smith. The geometry of algorithms with orthogonality constraints. SIAM journal on Matrix Analysis and Applications , 20(2):303-353, 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 on Artificial Intelligence and Statistics , pages 2681-2690. PMLR, 2019.

Alessio Figalli and Ludovic Rifford. Mass transportation on sub-riemannian manifolds. Geometric and functional analysis , 20:124-159, 2010.

Alessio Figalli, Ludovic Rifford, and Cédric Villani. On the ma-trudinger-wang 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 Mathematica , 177(2):113-161, 1996.

Ziv Goldfeld, Kengo Kato, Gabriel Rioux, and Ritwik Sadhu. Limit theorems for entropic optimal transport maps and the sinkhorn divergence. arXiv preprint arXiv:2207.08683 , 2022.

Gene H Golub and Charles F Van Loan. Matrix computations . JHU press, 2013.

Daniel Grange, Mohammad Al-Jarrah, Ricardo Baptista, Amirhossein Taghvaei, Tryphon T Georgiou, 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-4455. PMLR, 2021.

Jan-Christian 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 Meng-Papaxanthos, Michael Sterr, Aimée Bastidas-Ponce, Marta Tarquis-Medina, Heiko Lickert, Mostafa Bakhti, Mor Nitzan, Marco Cuturi, and Fabian J. Theis. Mapping cells through time and space with moscot. bioRxiv , 2023. doi: 10.1101/2023.05.11.540374. URL https://www. biorxiv.org/content/early/2023/05/12/2023.05.11.540374 .

Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative 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, Marco 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 I Jordan. On projection robust optimal transport: Sample complexity and model misspecification. In International Conference on Artificial Intelligence and Statistics , pages 262-270. PMLR, 2021.

Qiang Liu. Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint arXiv:2209.14577 , 2022.

Xi-Nan Ma, Neil S 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 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-squares. arXiv preprint arXiv:2112.01907 , 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 , 1 (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.12004 , 2021.

Aram-Alexandre Pooladian, Vincent Divol, and Jonathan Niles-Weed. Minimax estimation of discontinuous optimal transport maps: The semi-discrete case. arXiv preprint arXiv:2301.11302 , 2023a.

Aram-Alexandre Pooladian, Carles Domingo-Enrich, Ricky TQ Chen, and Brandon Amos. Neural optimal transport with Lagrangian costs. In ICML Workshop on New Frontiers in Learning, Control, and Dynamical Systems , 2023b.

Philippe Rigollet and Austin J 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 on 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 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-879, 1964.

Sanjay R Srivatsan, José L McFaline-Figueroa, Vijay Ramani, Lauren Saunders, Junyue Cao, Jonathan Packer, Hannah A Pliner, Dana L Jackson, Riza M Daza, Lena Christiansen, et al. Massively multiplex chemical transcriptomics 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 dynamics. In Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 9526-9536. PMLR, 2020.

Cédric Villani et 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-320, 2005.

A More on Subspace Elastic Costs

Recall that for a rankp matrix A ∈ R p × d , p ≤ d , the projection matrix that maps it to its orthogonal is A ⊥ = I -A T ( AA T ) -1 A . When A lies in the Stiefel manifold (i.e. AA T = I ), we have the simplification A ⊥ = I -A T A . This results in the Pythagorean identity ∥ z ∥ 2 = ∥ A ⊥ z ∥ 2 + ∥ A z ∥ 2 , as intended. In order to promote displacements that happen within the span of A , we must set a regularizer that penalizes the presence of z within its complement :

Since τ A ⊥ is a quadratic form, its proximal operator can be obtained by solving a linear system [Parikh et al., 2014, §6.1.1]; developing and using the matrix inversion lemma results in

To summarize, given an orthogonal sub-basis A of p 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 ⊥ , as a matrix directly specified in the ( d -p ) × p Stiefel manifold. This can alleviate computations to obtain a closed form for its proximal operator:

but requires storing B , a ( d -p ) × d orthogonal matrix, which is cumbersome when p ≪ d .

Parameterizing A explicitly , either as a full-rank p × d matrix, or more simply a p × d orthogonal matrix, to recover the suitable proximal operator for τ A ⊥ , by either

(a) Falling back on the right-most expression in (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) × d and is positive definite, in addition to the linear solve of complexity O ( p 3 ) . This simplifies when A is orthogonal, A ∈ S p,d since in that case,

(b) Alternatively, compute a 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 × d matrix A ⊥ or rank d -p , B := Gram-Schmidt ( A ⊥ ) , 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 a close connection with optimal maps for the standard ℓ 2 2 cost.

Proposition 3. Let T ⋆ be the optimal map between µ and ν for the cost h = 1 2 ℓ 2 2 + γτ A ⊥ . Denote by W the linear map x ↦→ ((1 + γ ) I -γA T A ) 1 / 2 x . Then W ◦ T ⋆ ◦ W -1 is the Brenier map (i.e., ℓ 2 2 optimal map) between W ♯ µ and W ♯ ν . Equivalently, T ⋆ is h -optimal if and only if it can be written

where ˜ T is the gradient of a convex function.

Proof of Proposition 3. The cost h = 1 2 ℓ 2 2 + γτ A ⊥ can be written as

The optimal transport problem we consider is therefore equivalent to minimizing

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 W ♯ µ and W ♯ ν is given by y ′ = ˜ T ( x ′ ) , which implies that the optimal h -coupling between µ and ν is given by T ⋆ = W -1 ◦ ˜ T ◦ W , as desired.

The proof of Theorem 1 requires the following two lemmas.

Lemma 1. For costs of the form h ( z ) = 1 2 z ⊤ Bz where B is positive definite, the MBO estimator between two measures µ and ν can be written as the barycentric projection of the corresponding optimal entropic coupling.

Proof. Note that h ∗ ( w ) = 1 2 w ⊤ B -1 w , and thus ∇ h ∗ ( w ) = B -1 w . Let ( f ε , g ε ) denote the optimal entropic potentials for this cost, with corresponding coupling π ε . Borrowing computations from e.g., Proposition 2 of Pooladian and Niles-Weed [2021], we can compute

where π x ε ( y ) is the conditional entropic coupling (given x ). The proof concludes by taking the expression of the MBO estimator and expanding:

which is the definition of the barycentric projection of π ε for a given x

.

Lemma 2 (Pre-conditioning of MBO) . Let T ε be the MBO estimator between µ and ν for the cost h = 1 2 ℓ 2 + γτ A ⊥ . Let W be denoted as in Proposition 3. Then the MBO estimator is written as

where ˜ T ε is the barycentric projection between W ♯ µ and W ♯ ν .

Proof. The proof here is similar to Proposition 3, which we outline again for completeness. As before, we are interested in solutions to the optimization problem

with optimal coupling π ⋆ ε . Performing a change of variables π ′ = ( W ⊗ W ) ♯ π , we have

where µ ′ := W ♯ µ (and similarly for ν ′ ), where now the optimizer reads ( π ′ ε ) ⋆ . The two optimal plans are related as

It was established in Lemma 1 that the MBO estimator T ⋆ ε is given by the barycentric projection

Y

′

=

WY

and

X

′

=

WX

, we can re-write this as a function of

where we identify ˜ T ε ( · ) := E ( π ′ ε ) ⋆ [ Y ′ | X ′ = · ] ; this completes the proof.

We are now ready to present the main proof.

Proof of Theorem 1. Let T ε,n denote the MBO estimator between samples from µ and ν , and let ˜ T ε,n denote the entropic map estimator from samples µ ′ := W ♯ µ and ν ′ := W ♯ ν , where W has spectrum 0 < λ min ( W ) ≤ λ max ( W ) < + ∞ , where we have access to W since A is known.

Our goal is to establish upper bounds on

Paying for constants that scale like λ max ( W -1 ) , we have the bound

where we can now directly use the rates of convergence from [Pooladian and Niles-Weed, 2021, Theorem 3], as µ ′ 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 a simple aggregation of weighted gradients

Performing the change of variables

π

ε

instead:

′

To write this formula in a more compact way, it is sufficient to notice that, adopting the convention that θ be a parameter in R q , and introducing an arbitrary vector ω ∈ R q ,

The products of all coordinate wise gradients with ω is equivalent to the application of the Jacobians of R and P ⋆ . We write J θ P ⋆ and J θ R for these Jacobians, both being maps taking θ as input, and outputting a linear map J θ R : R q → R n × m , i.e. J θ R ( θ ) is a n × m matrix. As a consequence one has because these maps are linear, one also has

which gives the identification given in the main text.

D Additional Details on Experiments

In § 6.2, and unlike Figure 3, we do not choose a predefined value for γ ⋆ , but instead select it with the following procedure: we start with a small value for γ 0 = 0 . 1 , and increase it gradually, until a certain desirable criterion on these displacements goes above a threshold. To measure this, we first compute the (paired) matrix of displacements on a given sample,

We then consider ratio of singular values on p ∗ subspace (to select γ for ∥ A ⊥ · ∥ 2 2 ), writing σ for the vector of singular values of D , ranked in decreasing order, to compute

D.1 Synthetic experiments

In Algorithm 5, we outline a way to generate ground-truth data used in synthetic experiments mentioned in § 6.1 and Figure 3. For a given regularization strength γ and a regularizer τ θ , we first sample a concave function g , along with the ground-truth parameters for the regularizer θ ∗ . We then sample n source points X from the standard normal distribution and push them through the ground-truth OT map to get n target points Y , as described in Algorithm 4.

D.2 Experimental Procedure

Experiments in Section 6.1, MBOon Synthetic Ground Truth Displacement

For both τ := ℓ 1 and τ A := ∥ A ⊥ · ∥ 2 , run SYNTHETIC-OT-TASK (Algo. 5) to form samples X , Y (store A ∗ for τ A ).

Using these samples, benchmark MBO estimator informed by the cost structure: varying γ , using automatically scaled ε , and ground-truth parameter A ∗ when studying τ A ∗ , for 10 random splits of X , Y into X train , Y train and X test , Y test .

-Run MBO-ESTIMATOR (Alg. 1 using regularization) on X train , Y train , store MSE 1 = ∑ i ∥ T ε [ γ, τ θ ∗ , ε ]( x test i ) -y test i ∥ 2

-Report ratio MSE 1 / MSE 2 , showing MBO seeded with the right regularizer always outperforms original entropic map.

-Run MBO-ESTIMATOR (Alg. 1 without regularization) on X train , Y train , store MSE 2 = ∑ i ∥ T ε [0 , 0 , ε ]( x test i ) -y test i ∥ 2

Experiments in Section 6.2, Recovery of Ground-Truth Subspace Parameters in Elastic Costs

For τ A := ∥ A ⊥ · ∥ 2 , run SYNTHETIC-OT-TASK (Algo. 5) to generate paired samples X , Y in R d . Tune γ to have displacements concentrated (50%, 70%, 90%) in a subspace, as described in Appendix D. Recover A ∗ of dimension p ∗ × d .

Experiments in Section 6.3, Learning Displacement Subspaces for Single-Cell Transport

Run RECOVER-THETA on X , Y using arbitrary γ to output ˆ A . Display Average recovery error ∥ A ∗ -ˆ A ˆ A T A ∗ ∥ 2 2 /p ∗

̸

Here data come from real measurements, X ∈ R n × d , Y ∈ R m × d , n = m . Split into X train , Y train and X test , Y test .

Run RECOVER-THETA on X train , Y train using arbitrary γ and random initialization with varying ˆ p . Output ˆ A .

Run MBO-ESTIMATOR (Alg. 1 using regularization) on X train , Y train . Form predictions ˜ y i = T ε [ γ, τ ˆ A , ε ]( x test i ) .

Compute ℓ 2 2 Sinkhorn divergence between point cloud (˜ y i ) n i =1 and ( y test j ) j m =1 .

To benchmark (dotted-line in Fig. 5), run two steps above, but setting γ = 0 (no subspace regularization).

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]

Justification: The abstract describes the challenges underlying the selection of a cost in optimal transport theory. We advertise in that abstract a new way to benchmark such OT problems for a fairly wide class of exotic costs, and claim statistical consistency results and a new algorithm to learn that cost' parameters. All these claims are substantiated by simple experiments.

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 be 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 as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of 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 no 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 a 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 on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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 fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: While theory is only a part of our contribution (our paper proposes also a 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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs 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 a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or 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 experimental results of the paper to the extent that it 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 notebooks/tutorials.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a 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 a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it 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 one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: Data is open. We provide code to use our 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 NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to 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 a 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 as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 not include experiments.

The experimental setting should be presented in the core of the paper to a level of 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, or as 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 our 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. a mean close to 0 means the std is also close to 0). Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence 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).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It 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 a 96% CI, if the hypothesis of Normality of errors is 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 are out of range (e.g. negative error 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

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Although no claim is made in terms of compute performance, the fairly small scale of the experiments allows to execute these runs on a 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, or 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://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have followed these ethics guidelines.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 a 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.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

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 (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: NA

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a 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 use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 not available online, the authors are encouraged to reach out to the asset'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 of 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 is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an 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 not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to 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 crowdsourcing nor research with human subjects.

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nEfficiently 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nIt 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\nQUESTIONS:\nI do not have any questions.\n\nLIMITATIONS:\nAuthors have addressed limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nI 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?\n\nQUESTIONS:\nEq (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?\n\nLIMITATIONS:\nno negative societal impact", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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\nLIMITATIONS:\nI do not see the limitations discussed anywhere.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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?)\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nThe authors have addressed their limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}