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

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 article, we propose a framework for consistent optimal transport flow matching (COFM). In particular, we parameterize the transport potential using partially input convex neural networks (PICNNs), and incorporate a Hamilton-Jacobi (HJ) 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 multistep 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 $2\times $ reduction in ${\mathcal {L}}^{2}$ -UVP compared with state-of-the-art (SOTA) models while requiring approximately $9\times $ 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.