Potential Matching Optimal Transport: Continuous Normalizing Flows for Exact $p$-Wasserstein Dynamics

2026-08-06

AI Summary

This work proposes a continuous normalizing flow method based on the generalized Benamou–Brenier formulation to solve general $p$-cost optimal transport ($p$-OT) problems. The approach parameterizes the velocity field via the gradient of a scalar potential function and introduces a self-induced matching loss, training the model along the straight-line bridge defined by its own endpoints while leveraging maximum mean discrepancy (MMD) for flexible terminal distribution matching. Under regularity, exact terminal matching, and uniqueness assumptions, the authors theoretically establish that zero-loss solutions strictly satisfy the generalized Benamou–Brenier optimality system, thereby exactly recovering the $p$-OT map and dynamics. Experiments demonstrate that the method accurately reconstructs theoretical $p$-OT maps on synthetic data, achieves strong performance in high-dimensional tabular density modeling, and validates the flexibility of terminal matching in sample-wise color transfer tasks.

Abstract

We introduce Potential Matching Optimal Transport (PMOT), a potential-flow framework for general $p$-cost optimal transport with $c_p(x,y)=\|x-y\|^p$. PMOT parameterizes the CNF velocity field with a scalar potential in the generalized Benamou--Brenier form for the chosen exponent $p$. It trains the potential gradient with a self-induced matching loss along straight bridges determined by the model's own endpoints, while allowing flexible terminal distribution matching. Our main result establishes zero-loss exactness: under the stated regularity, exact terminal matching, and uniqueness assumptions, any zero-loss solution satisfies the generalized Benamou--Brenier optimality system and recovers the corresponding $p$-optimal transport map and dynamics. On synthetic benchmarks, PMOT learns $p$-specific maps that agree with the corresponding $p$-matched OT references. It also remains competitive as a likelihood-based density model on high-dimensional tabular data, and an MMD-based color transformation experiment demonstrates flexible sample-based terminal matching.

Problem

• optimal transport
• p-Wasserstein
• continuous normalizing flows
• potential flow
• terminal distribution matching

Innovation

• Potential Matching Optimal Transport
• Continuous Normalizing Flows
• p-Wasserstein Dynamics
• Benamou–Brenier Formulation
• Self-induced Matching Loss