Normalizing flows as approximations of optimal transport maps via linear-control neural ODEs
This work addresses the challenge of computing optimal transport maps in machine learning, particularly for generative models like normalizing flows, but it is incremental as it builds on existing neural ODE and optimal transport frameworks.
The paper tackles the problem of approximating optimal transport maps between probability measures using linear-control neural ODEs, showing that these flows can recover the optimal map under certain assumptions and providing a numerical scheme for practical computation, tested on a two-dimensional example.
In this paper, we consider the problem of recovering the $W_2$-optimal transport map T between absolutely continuous measures $μ,ν\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 $μ,ν$ 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 $μ_N,ν_N$ of the original measures $μ,ν$ 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.