Computing high-dimensional optimal transport by flow neural networks

TL;DR

This paper introduces a flow neural network approach to compute dynamic optimal transport in high-dimensional spaces, enabling efficient distribution interpolation and downstream tasks like density ratio estimation and domain adaptation.

Contribution

It proposes a novel invertible flow model for dynamic OT that works with finite samples and high-dimensional data, extending beyond static formulations.

Findings

Strong empirical performance on high-dimensional OT benchmarks

Effective for image-to-image translation tasks

Improves density ratio estimation accuracy

Abstract

Computing optimal transport (OT) for general high-dimensional data has been a long-standing challenge. Despite much progress, most of the efforts including neural network methods have been focused on the static formulation of the OT problem. The current work proposes to compute the dynamic OT between two arbitrary distributions $P$ and $Q$ by optimizing a flow model, where both distributions are only accessible via finite samples. Our method learns the dynamic OT by finding an invertible flow that minimizes the transport cost. The trained optimal transport flow subsequently allows for performing many downstream tasks, including infinitesimal density ratio estimation (DRE) and domain adaptation by interpolating distributions in the latent space. The effectiveness of the proposed model on high-dimensional data is demonstrated by strong empirical performance on OT baselines, image-to-image translation, and high-dimensional DRE.