Differentiable Cost-Parameterized Monge Map Estimators

TL;DR

This paper introduces a differentiable estimator for optimal transport maps that learns both the map and an adapted cost function using neural networks, leveraging prior knowledge for improved real-world applicability.

Contribution

It presents a novel neural-based method to jointly learn OT maps and cost functions, incorporating prior information to enhance transport map estimation.

Findings

Method effectively learns adapted OT maps from prior knowledge.

The approach is flexible with different loss functions.

Results demonstrate improved transport map accuracy.

Abstract

Within the field of optimal transport (OT), the choice of ground cost is crucial to ensuring that the optimality of a transport map corresponds to usefulness in real-world applications. It is therefore desirable to use known information to tailor cost functions and hence learn OT maps which are adapted to the problem at hand. By considering a class of neural ground costs whose Monge maps have a known form, we construct a differentiable Monge map estimator which can be optimized to be consistent with known information about an OT map. In doing so, we simultaneously learn both an OT map estimator and a corresponding adapted cost function. Through suitable choices of loss function, our method provides a general approach for incorporating prior information about the Monge map itself when learning adapted OT maps and cost functions.