Computing high-dimensional optimal transport by flow neural networks

TL;DR

This work addresses the challenge of computing high-dimensional optimal transport between two distributions $P$ and $Q$ from samples by learning a continuous-time invertible flow, the Q-flow, that solves the dynamic OT problem via the Benamou–Brenier formulation. The method trains a neural ODE-based velocity field to minimize transport cost, with bi-directional training and end-to-end refinement, and introduces a flow-ratio net for infinitesimal density-ratio estimation along the OT trajectory. Key contributions include the Q-flow net that directly learns the OT path between arbitrary endpoints, and the flow-ratio net enabling accurate DRE and downstream tasks like MI estimation and energy-based modeling, demonstrated on high-dimensional OT baselines, image-to-image translation, and DRE benchmarks. The approach enables scalable OT in high dimensions and provides practical tools for density-ratio estimation and domain adaptation along OT trajectories, with potential for broader impact in statistics and machine learning applications.

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.

Figures (29)

• Figure 1: The Q-flow model learns the dynamic OT, an invertible transport map $T_0^1$ (parametrized by the velocity field $v(x,t)$) between $P$ and $Q$ over the time interval $[0,1]$ with the least transport cost. The push-forwarded distribution $(T_0^1)_{\#} p$ (as well as $(T_1^0)_{\#} q$, respectively) is close to the target distribution $q$ ($p$, respectively).
• Figure 2: Handbag $\rightarrow$ shoes
• Figure 3: CelebA male $\rightarrow$ female
• Figure 5: Trajectory from $P$ (two-moon) to $Q$ (checkerboard)
• Figure 6: Estimated $\log$-ratio between $P_{t_{k-1}}$ and $P_{t_k}$ by the trained flow-ratio net.

Theorems & Definitions (1)

• Remark 1: Symmetry of dynamic OT and bi-directional training