POTENTIAL MATCHING OPTIMAL TRANSPORT: CONTINUOUS NORMALIZING FLOWS FOR EXACT p-WASSERSTEIN DYNAMICS

Standard Continuous Normalizing Flows (CNFs) learn invertible mappings between distributions but do not inherently prioritize any specific transport geometry. While existing methods like OT-Flow align dynamics with quadratic-cost (W2) optimal transport, they struggle to generalize to arbitrary p-cost geometries (cp(x,y) = ||x-y||^p). This paper introduces Potential Matching Optimal Transport (PMOT) to learn continuous transport trajectories that satisfy the Benamou-Brenier optimality conditions for any chosen exponent p.

PMOT parameterizes the velocity field of a CNF using a scalar potential, mirroring the generalized Benamou-Brenier optimality system. Instead of relying on external optimal transport (OT) couplings or complex HJB regularizers, PMOT uses a self-induced matching loss. For every data sample, the model calculates a straight-line bridge between the input and its predicted terminal endpoint; it then forces the potential gradient to match the velocity required for constant-speed movement along that bridge. This approach allows the model to learn the optimal transport map while simultaneously supporting likelihood-based density estimation.

The authors prove that under regularity and uniqueness conditions, any zero-loss solution to the PMOT objective recovers the exact p-optimal Monge map and the corresponding Benamou-Brenier velocity field. Empirical results on synthetic benchmarks (8-Gaussians, Pinwheel) demonstrate that PMOT learns maps that align closely with post-hoc Sinkhorn OT references for specific exponents. Furthermore, the model remains competitive with existing methods on high-dimensional tabular density estimation and demonstrates flexibility in sample-based terminal matching tasks like image color transformation.