COFM: Consistent Optimal Transport Flow Matching via Partially Input Convex Neural Networks

Abstract

Optimal transport (OT) provides a principled framework for learning mappings between probability distributions, and has found broad applications in generative modeling, inverse problems and scientific computing. Recently, flow matching methods have emerged as an efficient paradigm for learning continuous-time transport dynamics. However, existing OT-based flow matching methods often suffer from either high computational cost due to inner optimization or limited consistency. Moreover, it remains challenging to design neural architectures that can simultaneously guarantee convexity, stability, and efficient transport learning. In this paper, we propose a framework for consistent optimal transport flow matching. Specifically, we parameterize the transport potential using partially input convex neural networks (PICNN), and incorporate a Hamilton-Jacobi residual into the training objective to enforce dynamical consistency of the learned flow. This design enables a unified formulation that supports both one-step transport and multi-step ODE-based sampling, without requiring costly inner optimization. Extensive experiments on benchmark datasets demonstrate that the proposed method achieves competitive performance compared with existing OT-based and flow matching approaches, while maintaining favorable computational efficiency. In particular, under the D=256 benchmark, COFM achieves more than a 2x reduction in L^2-UVP compared with state-of-the-art (SOTA) models, while requiring approximately 9x less computational time. These results suggest that combining convex potential structures with HJ-based dynamical regularization provides an effective framework for scalable and geometrically consistent transport learning.