Near-optimal estimation of smooth transport maps with kernel sums-of-squares

TL;DR

This paper introduces a computationally feasible method for estimating smooth optimal transport maps with near-minimax statistical error bounds, advancing the practical application of optimal transport in high-dimensional settings.

Contribution

It presents the first tractable algorithm for estimating smooth transport maps with statistical guarantees close to theoretical lower bounds, using a sum-of-squares reformulation.

Findings

Algorithm achieves near-minimax statistical error bounds.

Method has dimension-free polynomial convergence rates.

Constants may grow exponentially with dimension.

Abstract

It was recently shown that under smoothness conditions, the squared Wasserstein distance between two distributions could be efficiently computed with appealing statistical error upper bounds. However, rather than the distance itself, the object of interest for applications such as generative modeling is the underlying optimal transport map. Hence, computational and statistical guarantees need to be obtained for the estimated maps themselves. In this paper, we propose the first tractable algorithm for which the statistical $L^2$ error on the maps nearly matches the existing minimax lower-bounds for smooth map estimation. Our method is based on solving the semi-dual formulation of optimal transport with an infinite-dimensional sum-of-squares reformulation, and leads to an algorithm which has dimension-free polynomial rates in the number of samples, with potentially exponentially dimension-dependent constants.