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

arXiv: https://arxiv.org/abs/2511.06042 · HTML: https://mlsanity.com/paper/2511.06042/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 33 of its 44 references that are in the map.
- First posted on arXiv 8 Nov 2025.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 30 of 33)

- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Glow: Generative Flow with Invertible 1x1 Convolutions](https://mlsanity.com/paper/1807.03039/) (2018)
- [Normalizing Flows for Probabilistic Modeling and Inference](https://mlsanity.com/paper/1912.02762/) (2019)
- [Diffusion Models: A Comprehensive Survey of Methods and Applications](https://mlsanity.com/paper/2209.00796/) (2022)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [DGM: A deep learning algorithm for solving partial differential equations](https://mlsanity.com/paper/1708.07469/) (2017)
- [Diffusion Schr"odinger Bridge with Applications to Score-Based Generative Modeling](https://mlsanity.com/paper/2106.01357/) (2021)
- [The Deep Ritz method: A deep learning-based numerical algorithm for solving variational problems](https://mlsanity.com/paper/1710.00211/) (2017)
- [Mean Flows for One-step Generative Modeling](https://mlsanity.com/paper/2505.13447/) (2025)
- [Solving high-dimensional partial differential equations using deep learning](https://mlsanity.com/paper/1707.02568/) (2017)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Discrete Flow Matching](https://mlsanity.com/paper/2407.15595/) (2024)
- [Rectified Flow: A Marginal Preserving Approach to Optimal Transport](https://mlsanity.com/paper/2209.14577/) (2022)
- [Optimal transport mapping via input convex neural networks](https://mlsanity.com/paper/1908.10962/) (2019)
- [Adversarial Latent Autoencoders](https://mlsanity.com/paper/2004.04467/) (2020)
- [OT-Flow: Fast and Accurate Continuous Normalizing Flows via Optimal Transport](https://mlsanity.com/paper/2006.00104/) (2020)
- [Consistency Flow Matching: Defining Straight Flows with Velocity Consistency](https://mlsanity.com/paper/2407.02398/) (2024)
- [On the relation between optimal transport and Schr"odinger bridges: A stochastic control viewpoint](https://mlsanity.com/paper/1412.4430/) (2014)
- [Wasserstein-2 Generative Networks](https://mlsanity.com/paper/1909.13082/) (2019)
- [Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark](https://mlsanity.com/paper/2106.01954/) (2021)
- [Optimal Flow Matching: Learning Straight Trajectories in Just One Step](https://mlsanity.com/paper/2403.13117/) (2024)
- [Light and Optimal Schr"odinger Bridge Matching](https://mlsanity.com/paper/2402.03207/) (2024)
- [Adaptive Deep Learning for High-Dimensional Hamilton-Jacobi-Bellman Equations](https://mlsanity.com/paper/1907.05317/) (2019)
- [Neural Monge Map estimation and its applications](https://mlsanity.com/paper/2106.03812/) (2021)

## Abstract (arXiv)

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.