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

2025-11-08

AI Summary

Existing flow-matching methods for optimal transport (OT) rely on nested optimization subproblems, compromising the balance between generation consistency and sampling flexibility. Method: We propose a consistency modeling framework based on OT flows, centered on a physics-informed partially input-convex neural network (PICNN) that explicitly encodes the convex potential structure underlying displacement interpolation, thereby directly parameterizing the Brenier map’s velocity field. By jointly optimizing the Hamilton–Jacobi equation residual and flow-matching loss, we eliminate the need for inner-loop optimization inherent in conventional one-step OT matching. Contribution/Results: The method unifies one-step Brenier mapping and multi-step ODE-based sampling, leveraging the straight-line property of OT geodesics to enhance path consistency and modeling flexibility. Extensive evaluation on standard OT benchmarks demonstrates its efficiency, scalability, and capability in probabilistic path modeling.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Computer Vision: Learning & Optimization for CV
• Machine Learning: Learning with Manifolds

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Systems and Infrastructure for Web, Mobile and WoT: Experiences and lessons learnt from Web-based algorithms and system deployments

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.

Problem

• Designing physics-informed neural networks for optimal transport flow
• Coupling Hamilton-Jacobi residual with flow matching loss function
• Enabling both one-step and multi-step sampling from learned potential

Innovation

• Physics-informed design of input convex neural networks
• Coupling Hamilton-Jacobi residual with flow matching loss
• One-step and multi-step ODE sampling from learned potential