Computing high-dimensional optimal transport by flow neural networks

Published: 27 Oct 2023, Last Modified: 28 Dec 2023OTML 2023 PosterEveryoneRevisionsBibTeX
Keywords: flow-based neural network, high-dimensional optimal transport
TL;DR: We propose a flow-based neural network to learn the optimal transport between high-dimensional distributions, accessible only via finite samples.
Abstract: Flow-based models are widely used in generative tasks, including normalizing flow, where a neural network transports from a data distribution $P$ to a normal distribution. This work develops a flow-based model that transports from $P$ to an arbitrary $Q$ where both distributions are only accessible via finite samples. We propose to learn the dynamic optimal transport between $P$ and $Q$ by training a flow neural network. The model is trained to find an invertible transport map between $P$ and $Q$ optimally by minimizing the transport cost. The trained optimal transport flow allows for performing many downstream tasks, including infinitesimal density ratio estimation and distribution interpolation in the latent space for generative models. The effectiveness of the proposed model on high-dimensional data is empirically demonstrated in mutual information estimation, energy-based generative models, and image-to-image translation.
Submission Number: 58