Learning Elastic Costs to Shape Monge Displacements

Given a source and a target probability measure supported on ℝd, the Monge problem asks to find the most efficient way to map one distribution to the other. This efficiency is quantified by defining a textitcost function between source and target data. Such a cost is often set by default in the machine learning literature to the squared-Euclidean distance, ℓ22(x,y)=tfrac12|x-y|22. Recently, Cuturi et. al '23 highlighted the benefits of using elastic costs, defined through a regularizer τ as c(x,y)=ℓ22(x,y)+τ(x-y). Such costs shape the textitdisplacements of Monge maps T, i.e., 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 ℓ22 cost; (ii) We propose a loss to textitlearn the parameter θ of a parameterized regularizer τθ, and apply it in the case where τA(z)=|A^perp z|22. This regularizer promotes displacements that lie on a low dimensional subspace of ℝd, spanned by the p rows of A∈ℝp× d.