Optimal Estimation of Smooth Transport Maps with Kernel SoS

Under smoothness conditions, it was recently shown by Vacher et al. [Proceedings of the 34th Conference on Learning Theory, Proc. Mach. Learn. Res. 134, 2021] that the squared Wasserstein distance between two distributions could be approximately computed in polynomial time with appealing statistical error bounds. In this paper, we propose to extend their result to the problem of estimating in L² distance the transport map between two distributions. Also building upon the kernelized sum-of-squares approach, a way to model smooth positive functions, we derive a computationally tractable estimator of the transport map. Contrary to the aforementioned work, the dual problem that we solve is closer to the so-called semidual formulation of optimal transport that is known to gain convexity with respect to the linear dual formulation. After deriving a new stability result on the semidual and using localization-like techniques through Gagliardo-Nirenberg inequalities, we manage to prove under the same assumptions as in Vacher et al. that our estimator is minimax optimal up to polylog factors. Then we prove that this estimator can be computed in the worst case in O^\~(n⁵) time, where n is the number of samples, and show how to improve its practical computation with a Nystr\"om approximation scheme, a classical tool in kernel methods. Finally, we showcase several numerical simulations in medium dimension, where we compute our estimator on simple examples.