Normalizing flows as approximations of optimal transport maps via linear-control neural ODEs

Abstract: In this paper, we consider the problem of recovering the $W_2$-optimal transport map T between absolutely continuous measures $\mu,\nu\in\mathcal{P}(\mathbb{R}^n)$ as the flow of a linear-control neural ODE, where the control depends only on the time variable and takes values in a finite-dimensional space. We first show that, under suitable assumptions on $\mu,\nu$ and on the controlled vector fields governing the neural ODE, the optimal transport map is contained in the $C^0_c$-closure of the flows generated by the system. Then, we tackle the problem under the assumption that only discrete approximations of $\mu_N,\nu_N$ of the original measures $\mu,\nu$ are available: we formulate approximated optimal control problems, and we show that their solutions give flows that approximate the original optimal transport map $T$. In the framework of generative models, the approximating flow constructed here can be seen as a `Normalizing Flow', which usually refers to the task of providing invertible transport maps between probability measures by means of deep neural networks. We propose an iterative numerical scheme based on the Pontryagin Maximum Principle for the resolution of the optimal control problem, resulting in a method for the practical computation of the approximated optimal transport map, and we test it on a two-dimensional example.