Physics-Informed Design of Input Convex Neural Networks for Consistency Optimal Transport Flow Matching

TL;DR

This paper introduces a physics-informed neural network approach for optimal transport flow matching that improves training efficiency and supports flexible sampling methods, validated on standard benchmarks.

Contribution

It presents a novel partially input-convex neural network design coupled with Hamilton-Jacobi residuals to enhance OT flow modeling without inner optimization.

Findings

Avoids inner optimization subproblems

Supports both one-step and multi-step sampling

Validated on standard OT benchmarks

Abstract

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.