Learning Elastic Costs to Shape Monge Displacements

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, ℓ²₂(x, y):= 1/2 ∥x − y∥²₂. The benefits of using elastic costs, defined using a regularizer τ as c(x, y):= ℓ²₂(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 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 ℓ²₂ cost; (ii) We propose a loss to learn the parameter θ of a parameterized regularizer τ_θ, and apply it in the case where τ_A(z):= ∥Az∥²₂. 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.