Physics-Informed Design of Input Convex Neural Networks for Consistency Optimal Transport Flow Matching
This work addresses computational efficiency and flexibility in optimal transport flow matching for machine learning applications, though it appears incremental as it builds on existing flow matching and neural network methods.
The paper tackles the problem of designing consistency models for optimal transport flow matching by proposing a physics-informed approach using partially input-convex neural networks to construct flow fields that emulate displacement interpolation, avoiding inner optimization subproblems in training and supporting flexible sampling methods, with validation on standard benchmarks.
We propose a consistency model based on the optimal-transport flow. A physics-informed design of partially input-convex neural networks (PICNN) plays a central role in constructing the flow field that emulates the displacement interpolation. During the training stage, we couple the Hamilton-Jacobi (HJ) residual in the OT formulation with the original flow matching loss function. Our approach avoids inner optimization subproblems that are present in previous one-step OFM approaches. During the prediction stage, our approach supports both one-step (Brenier-map) and multi-step ODE sampling from the same learned potential, leveraging the straightness of the OT flow. We validate scalability and performance on standard OT benchmarks.