43
3

Learning Elastic Costs to Shape Monge Displacements

Abstract

Given a source and a target probability measure supported on Rd\mathbb{R}^d, the Monge problem asks to find the most efficient way to map one distribution to the other. This efficiency is quantified by defining a \textit{cost} function between source and target data. Such a cost is often set by default in the machine learning literature to the squared-Euclidean distance, 22(x,y)=12xy22\ell^2_2(\mathbf{x},\mathbf{y})=\tfrac12|\mathbf{x}-\mathbf{y}|_2^2. Recently, Cuturi et. al '23 highlighted the benefits of using elastic costs, defined through a regularizer τ\tau as c(x,y)=22(x,y)+τ(xy)c(\mathbf{x},\mathbf{y})=\ell^2_2(\mathbf{x},\mathbf{y})+\tau(\mathbf{x}-\mathbf{y}). Such costs shape the \textit{displacements} of Monge maps TT, i.e., the difference between a source point and its image T(x)x)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 22\ell_2^2 cost; (ii) We propose a loss to \textit{learn} the parameter θ\theta of a parameterized regularizer τθ\tau_\theta, and apply it in the case where τA(z)=Az22\tau_{A}(\mathbf{z})=|A^\perp \mathbf{z}|^2_2. This regularizer promotes displacements that lie on a low dimensional subspace of Rd\mathbb{R}^d, spanned by the pp rows of ARp×dA\in\mathbb{R}^{p\times d}.

View on arXiv
Comments on this paper