Learning Monge maps with constrained drifting models

TL;DR

This paper introduces a new evolution equation for estimating optimal transport maps, proves its convergence, and develops neural network-based algorithms that outperform standard methods in stability and accuracy.

Contribution

It proposes a novel constrained gradient flow for OT maps, proves its long-term convergence, and implements neural network schemes that improve training stability and approximation quality.

Findings

The implicit scheme converges to the OT map as time approaches infinity.

Neural network parameterization with the proposed flow yields better approximation results.

The method outperforms standard Euclidean gradient descent and Adam in training stability and accuracy.

Abstract

We study the estimation of optimal transport (OT) maps between an arbitrary source probability measure and a log-concave target probability measure. Our contributions are twofold. First, we propose a new evolution equation in the set of transport maps. It can be seen as the gradient flow of a lift of some user-chosen divergence (e.g., the KL divergence, or relative entropy) to the space of transport maps, constrained to the convex set of optimal transport maps. We prove the existence of long-time solutions to this flow as well as its convergence toward the OT map as time goes to infinity, under standard convexity conditions on the divergence. Second, we study the practical implementation of this constrained gradient flow. We propose two time-discrete computational schemes-one explicit, one implicit-, and we prove the convergence of the latter to the OT map as time goes to infinity. We then parameterize the OT maps with convexity-constrained neural networks and train them with these discretizations of the constrained gradient flow. We show that this is equivalent to performing a natural gradient descent of the lift of the chosen divergence in the neural networks' parameter space, similarly to drifting generative models. Empirically, our scheme outperforms the standard Euclidean gradient descent methods used to train convexity-constrained neural networks in terms of approximation results for the OT map and convergence stability, and it still yields better results than the same approach combined with the widely used Adam optimizer.