Optimal Transport Flow Matching (OTFM)

• OTFM is a paradigm that fuses optimal transport theory with flow matching, defining continuous flows between distributions through analytic trajectory regression.
• It directly regresses neural vector fields against optimal transport couplings to achieve fast, stable, and few-step generation across continuous, discrete, and unbalanced scenarios.
• OTFM has broad applications in generative modeling, density control, and trajectory inference, providing robust, geometry-informed regularization and theoretical optimality.

Optimal Transport Flow Matching (OTFM) is a paradigm that unifies optimal transport theory and flow matching for generative modeling, density control, and trajectory inference across continuous, discrete, and unbalanced settings. OTFM enforces straightest-possible or action-minimizing trajectories between distributions by leveraging optimal transport (OT) couplings within simulation-free, regression-based training objectives. This ensures theoretically optimal flows, greatly improves sampling efficiency, enables few-step (even one-step) generation, and provides robust, geometry-informed regularization.

1. Mathematical Foundations and Core Formulation

OTFM formulates generation or measure interpolation as a continuous flow problem underpinned by the geometry of optimal transport. For measures $p_0,p_1$ with quadratic cost $c(x_0,x_1)=\frac{1}{2}\|x_0-x_1\|^2$, OTFM seeks either a deterministic transport map or a stochastic flow whose trajectories minimize kinetic action under the continuity equation constraints: $$\begin{aligned} \min_{\{p_t,v_t\}} &\ \int_0^1 \!\! \int_{\mathbb{R}^d}\|v_t(x)\|^2\,p_t(x)\,dx\,dt\ \text{s.t.} &\ \partial_t p_t + \nabla\cdot(p_t v_t)=0,\quad p_0,p_1 \text{ fixed}. \end{aligned}$$ In classical Benamou–Brenier dynamic OT, the optimal displacement interpolation follows straight lines at constant velocity between coupled $(x_0,x_1)$ pairs drawn from the OT plan $\pi^*$, i.e., $x_t = (1-t)x_0 + t x_1$ with $v_t(x_t) = x_1-x_0$ (Kornilov et al., 2024, Kornilov et al., 31 Oct 2025).

OTFM replaces simulation-based or implicit likelihood objectives with direct regression of neural vector fields against these analytic paths, yielding fast, high-fidelity, and numerically stable training (Tong et al., 2023).

2. Algorithmic Framework and Losses

OTFM leverages optimal couplings $\pi^*$ between source and target measures, computed via Wasserstein OT with relevant cost metrics (e.g., Euclidean, Hamming). In the continuous case, for $(x_0,x_1)\sim\pi^*$, the main flow-matching objective is

$$\mathcal{L}_{\rm OTFM}(\theta) = \mathbb{E}_{(x_0,x_1)\sim\pi^*,\, t\sim U[0,1]} \left\| v_\theta(x_t, t) - (x_1 - x_0) \right\|^2,$$

where $c(x_0,x_1)=\frac{1}{2}\|x_0-x_1\|^2$0 (Kornilov et al., 2024, Tong et al., 2023).

Parameterizations include time-dependent neural vector fields (e.g., ResNetFlow, U-Nets, Graph Transformers) or, for strict quadratic-cost OTFM, an input-convex neural network $c(x_0,x_1)=\frac{1}{2}\|x_0-x_1\|^2$1 whose gradient yields the Brenier map and trajectories (Kornilov et al., 2024).

In discrete spaces (e.g., graphs, sequences), OTFM employs a categorical or convex interpolant construction, with couplings determined by Hamming or edit-distance costs and min-batch OT solvers (Sinkhorn, Hungarian) (Hou et al., 2024, Haxholli et al., 2024).

Loss computation typically follows:

1. Sample paired $c(x_0,x_1)=\frac{1}{2}\|x_0-x_1\|^2$2 via OT.
2. Interpolate $c(x_0,x_1)=\frac{1}{2}\|x_0-x_1\|^2$3 (or $c(x_0,x_1)=\frac{1}{2}\|x_0-x_1\|^2$4 for graphs) for random $c(x_0,x_1)=\frac{1}{2}\|x_0-x_1\|^2$5.
3. Compute regression target $c(x_0,x_1)=\frac{1}{2}\|x_0-x_1\|^2$6 (or appropriate discrete-action direction).
4. Update $c(x_0,x_1)=\frac{1}{2}\|x_0-x_1\|^2$7 to minimize squared error, possibly with weighted mass or path-dependent costs.

Simulation-free variants further offer one-step or few-step mapping: e.g., ODE-free neural flow matching directly learns $c(x_0,x_1)=\frac{1}{2}\|x_0-x_1\|^2$8 in a single pass by enforcing global OT-consistent pairing (Shou, 7 Apr 2026).

3. Extensions: Unbalanced, Discrete, Functional, and Conditional OTFM

• Unbalanced (WFR) OTFM: The Wasserstein–Fisher–Rao (WFR) extension augments OTFM with a scalar growth field $c(x_0,x_1)=\frac{1}{2}\|x_0-x_1\|^2$9 to parameterize local mass change, capturing "birth–death" dynamics in unbalanced measure transport. The loss jointly regresses displacement and mass change to closed-form Dirac geodesics (Peng et al., 11 Jan 2026).
• Discrete OTFM: In categorical and combinatorial domains (e.g., molecular graphs, text), OTFM uses discrete convex interpolants and cost-aware pairings, supporting stable, nearly autoregressive-level modeling and controlled generation (Hou et al., 2024, Haxholli et al., 2024).
• Hilbert-space/Functional OTFM: OTFM has been generalized to infinite-dimensional function spaces, enabling mesh-independent, physically-grounded modeling of high-dimensional fields such as turbulence. Here, the coupling and regression are performed in the Hilbert or RKHS basis (Kunpeng et al., 7 Apr 2026).
• Conditional and multi-marginal OTFM: OTFM supports conditional generation by solving pairwise OT for all condition pairs and training a conditional vector field, yielding all-to-all continuous condition transfer (Ikeda et al., 4 Apr 2025).

4. Theoretical Guarantees, Optimality, and Proximal Perspectives

OTFM is analytically linked to both the Monge–Kantorovich OT dual and Benamou–Brenier dynamic OT formulations (Kornilov et al., 31 Oct 2025). Restricting to OT-induced vector fields (via convex Brenier potentials) recovers the exact quadratic cost OT transport in one minimization, with theoretical equivalence proved via Fenchel duality and potential conjugacy (Kornilov et al., 2024).

Recent work further interprets OT-CFM from the viewpoint of proximal operators: each point along the interpolation path can be expressed as a proximal operator applied to an extended Brenier potential, providing a variational, strongly convex characterization of the flow and rigorous convergence guarantees, even in minibatch or manifold-supported regimes (Fukumizu et al., 13 Feb 2026). Normal hyperbolicity at terminal time ensures contraction in normal directions to the data manifold, with robust generation under data perturbations.

In the unbalanced WFR case, minimizing the OTFM loss provably yields WFR geodesics of evolving density and mass, unifying balanced and unbalanced trajectory inference (Peng et al., 11 Jan 2026).

5. Practical Applications and Empirical Outcomes

OTFM is applied across a wide range of modalities:

• Generative models: OTFM-trained flows exhibit straighter paths and reduced integration costs compared to classical FM or diffusion, with empirically fewer steps required and lower FID/MMD errors in image, graph, and molecular generation (Tong et al., 2023, Hou et al., 2024, Shou, 7 Apr 2026).
• Scene flow, pansharpening, and turbulence: One-to-one OT assignment improves self-supervised scene flow label quality (Li et al., 2021), enables high-resolution image fusion in a single step via unbalanced OTFM (Cao et al., 19 Mar 2025), and generalizes to high-fidelity turbulence synthesis at any mesh resolution (Kunpeng et al., 7 Apr 2026).
• 3D molecular conformations: OTFM aligns simulation-free ODE-based flows with minimization of RMSD cost, outperforming SDE/diffusion-based models in coverage and diversity on multi-conformer QM9 (Tian et al., 2024).
• Density control and collision avoidance: OTFM provides a framework for multiagent trajectory optimization with path-dependent and pairwise-interaction costs, achieving collision-free control in swarm scenarios (Duan et al., 8 Oct 2025).

6. Extensions, Alternatives, and Current Limitations

Several extensions and variants address the following:

• Non-quadratic or task-regularized transport: Unbalanced OT, regularized cost terms, and path-dependent action matching can be flexibly incorporated (Cao et al., 19 Mar 2025, Duan et al., 8 Oct 2025).
• Straightness via acceleration minimization: OAT-FM augments OTFM by explicitly minimizing acceleration in the trajectory space, yielding further decreases in path curvature and enabling two-phase fine-tuning of existing FM models (Yue et al., 29 Sep 2025).
• Model-aligned versus geometry-only OT: "Model-Aligned Coupling" selects OT pairings that are not only geometrically but also functionally learnable given the model's current fit, reducing conflicting flows, especially in few-step settings (Lin et al., 29 May 2025).
• Scalability: Minibatch and approximate Sinkhorn/Hungarian OT solve OT couplings efficiently for large data, with additional stability via embedding averaging (Haxholli et al., 2024).
• Limitations: OTFM with quadratic cost admits strong optimality claims, but extension to arbitrary cost functions is nontrivial and remains an open area. Inner subproblems for convex inversion or continuous optimal assignments limit naive scalability to extreme dimensions (Kornilov et al., 2024, Fukumizu et al., 13 Feb 2026).

7. Summary Table of OTFM Variants

Variant / Setting	Core Objective	OT Coupling	Flow Structure	Reference
Continuous, balanced OT	L2 regression to OT direction	W2 Euclidean	Linear (Benamou-Brenier)	(Kornilov et al., 2024, Tong et al., 2023)
Unbalanced/WFR	Displacement + mass growth	WFR semi-coupling	Displacement + mass ODE	(Peng et al., 11 Jan 2026, Cao et al., 19 Mar 2025)
Discrete (sequences, graphs)	Tok.-sparse jump/graph field	Hamming/Sinkhorn	Convex interpolants	(Hou et al., 2024, Haxholli et al., 2024)
Functional/Hilbert space	RKHS flow-matching	$$\begin{aligned} \min_{\{p_t,v_t\}} &\ \int_0^1 \!\! \int_{\mathbb{R}^d}\|v_t(x)\|^2\,p_t(x)\,dx\,dt\ \text{s.t.} &\ \partial_t p_t + \nabla\cdot(p_t v_t)=0,\quad p_0,p_1 \text{ fixed}. \end{aligned}$$0 in function space	Linear, mesh-free	(Kunpeng et al., 7 Apr 2026)
Model-aligned coupling	Model-prediction-reweighted	Top-k model-aligned	Curved OT, regularized	(Lin et al., 29 May 2025)
OT-Acceleration/OAT-FM	Acceleration action	OT, augmented	Cubic, curvature-minimizing	(Yue et al., 29 Sep 2025)

In total, OTFM establishes a simulation-free, mathematically optimal regime for flow-based generative modeling and dynamic inference, tightly unifying optimal transport couplings with learned dynamical systems, and extensible to unbalanced, discrete, and infinite-dimensional settings. Its empirical and theoretical advantages are validated across a range of domains (Tong et al., 2023, Kornilov et al., 2024, Haxholli et al., 2024, Cao et al., 19 Mar 2025, Lin et al., 29 May 2025, Yue et al., 29 Sep 2025, Duan et al., 8 Oct 2025, Kornilov et al., 31 Oct 2025, Peng et al., 11 Jan 2026, Fukumizu et al., 13 Feb 2026, Kunpeng et al., 7 Apr 2026, Shou, 7 Apr 2026, Hou et al., 2024, Tian et al., 2024, Ikeda et al., 4 Apr 2025, Li et al., 2021).