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.

Machine Learning (cs.LG)

arXiv:2511.06042 [cs.LG]

(or arXiv:2511.06042v2 [cs.LG] for this version)

https://doi.org/10.48550/arXiv.2511.06042 Focus to learn more arXiv-issued DOI via DataCite

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.