Optimal Transport Flow Matching

• Optimal Transport Flow Matching is a generative modeling framework that combines OT couplings with flow matching techniques to yield straight, non-crossing transport paths.
• It employs convex potential parameterization to recover exact transport maps in one step, reducing error accumulation and computational overhead.
• Extensions of OTFM cover unbalanced, conditional, and functional cases, achieving state-of-the-art performance in diverse high-dimensional generative applications.

Optimal Transport Flow Matching (OTFM) is a class of generative modeling frameworks that couple the mathematical structure of optimal transport (OT)—the theory of finding cost-minimizing maps between probability distributions—with conditional and marginal flow-matching techniques. OTFM aims to learn transport maps or time-dependent velocity fields that achieve straight, non-crossing probability paths between distributions in a sample-efficient, simulation-free, and theoretically justified manner. The methodology arises in response to the limitations of conventional flow matching, such as path curvature, error accumulation, and computational overhead in high-dimensional generative tasks.

1. Mathematical Formulation: From OT to Flow Matching

OTFM builds upon flow matching (FM), which, for quadratic cost $c(x_0, x_1) = \|x_0 - x_1\|^2/2$, seeks a time-dependent vector field $v_t(x)$ such that the ODE

$\frac{dx_t}{dt} = v_t(x_t)\,,\quad x_{t=0}\sim p_0$

transports a “prior” distribution $p_0$ to a “target” $p_1$ at $t=1$. FM commonly relies on a conditional loss, constructed by regressing $v_t$ to the ground-truth velocity $x_1 - x_0$ along paths $x_t = (1 - t)x_0 + t x_1$ for pairs $(x_0, x_1)$ from a chosen coupling $v_t(x)$0.

The optimal transport plan $v_t(x)$1, minimizing the expected quadratic cost over all couplings with marginals $v_t(x)$2, $v_t(x)$3, imbues the induced displacement field with key geometric structure: straightness and non-crossing. However, practical FM often defaults to simpler but suboptimal couplings (e.g., independent $v_t(x)$4), leading to curved, inefficient trajectories.

OTFM methods directly incorporate OT couplings and associated convex potentials into the definition and training of the flow model, yielding exact or near-exact geodesic (straight-line) flows, thereby regularizing transport and dramatically improving inference efficiency and empirical fidelity (Kornilov et al., 2024, Malnick et al., 2 Jun 2026).

2. Convex Potential Parameterization and One-Step Exactness

A central mathematical innovation in OTFM is the restriction of the learnable vector field to those induced by the gradient of a convex potential, as per Brenier’s theorem. For quadratic cost:

• The unique Monge map $v_t(x)$5 pushes $v_t(x)$6 to $v_t(x)$7, with $v_t(x)$8 convex.
• OTFM restricts $v_t(x)$9, where $\frac{dx_t}{dt} = v_t(x_t)\,,\quad x_{t=0}\sim p_0$0.

The OTFM loss

$\frac{dx_t}{dt} = v_t(x_t)\,,\quad x_{t=0}\sim p_0$1

is shown to coincide (up to constant) with the dual OT objective. A single minimization of this loss over convex $\frac{dx_t}{dt} = v_t(x_t)\,,\quad x_{t=0}\sim p_0$2 recovers the OT potential, enabling the exact Brenier map and straight-line transport in one step—no iterative rectification or flow-straightening required. Empirically, minibatch OT couplings accelerate convergence but do not bias the final solution (Kornilov et al., 2024).

3. Extensions: Unbalanced, Conditional, and Functional OTFM

Unbalanced OTFM

Extensions to unbalanced OT (UOT)—where marginal constraints are relaxed—enable flow matching in settings with varying mass or imperfect alignment (e.g., remote sensing, single-cell dynamics). UOTFM integrates the dual formulation of UOT (with $\frac{dx_t}{dt} = v_t(x_t)\,,\quad x_{t=0}\sim p_0$3-divergence penalties on marginal deviations) into the flow-matching loss. The learned mapping retains straight path structure while accommodating intrinsic discrepancies between marginals and supporting task-specific regularization (e.g., spatial/spectral constraints) (Cao et al., 19 Mar 2025, Peng et al., 15 May 2026).

Conditional and Functional OTFM

In problems with structured targets (e.g., 3D conformation prediction, turbulent field generation), OTFM generalizes to infinite-dimensional Hilbert spaces or conditional frameworks, using mini-batch OT assignments and conditional straight-line interpolations. The resulting models are mesh-independent, support high-resolution super-resolution, and preserve straightness even in function spaces. The theoretical correspondence of the conditional flow-matching loss to the marginal (true) flow persists in Hilbert-space OTFM, guaranteeing training recovers the intended transport (Tian et al., 2024, Kunpeng et al., 7 Apr 2026).

4. Minibatch OT, Plan Consistency, and Algorithmic Aspects

Exact OT between empirical measures is computationally expensive in high dimensions, leading to widespread adoption of minibatch OT as a surrogate. The expected-batch OT plan, defined as the average OT plan over random minibatches, interpolates between the independent coupling ($\frac{dx_t}{dt} = v_t(x_t)\,,\quad x_{t=0}\sim p_0$4) and the true OT plan ($\frac{dx_t}{dt} = v_t(x_t)\,,\quad x_{t=0}\sim p_0$5), with identifiable convergence rates for cost and plan bias (Boïté et al., 12 May 2026). While increasing minibatch size flattens the initial velocity field and sharpens empirical transport, diminishing returns are observed beyond moderate sizes (e.g., $\frac{dx_t}{dt} = v_t(x_t)\,,\quad x_{t=0}\sim p_0$6–$\frac{dx_t}{dt} = v_t(x_t)\,,\quad x_{t=0}\sim p_0$7). Algorithmic implementations generally employ entropic regularization and Sinkhorn solvers to scale to large batches or near-real-time sampling (Kornilov et al., 2024, Boïté et al., 12 May 2026, Cao et al., 19 Mar 2025).

A selection of practical algorithmic features across OTFM methods includes:

• Use of input-convex neural networks (ICNNs) to parametrize convex potentials.
• Recovery of $\frac{dx_t}{dt} = v_t(x_t)\,,\quad x_{t=0}\sim p_0$8 for $\frac{dx_t}{dt} = v_t(x_t)\,,\quad x_{t=0}\sim p_0$9 via convex optimization.
• Simulation-free training: straight-line targets admit closed-form conditional velocities, bypassing the need for expensive simulation or ODE integration at training time (Kornilov et al., 2024, Cao et al., 19 Mar 2025, Tian et al., 2024).

5. Theoretical Guarantees and Geometric Properties

OTFM inherits and extends theoretical properties from optimal transport and flow matching:

• Duality Equivalence: The OTFM loss under quadratic cost is exactly equivalent to the OT dual objective, ensuring that minimizers induce the Brenier map and associated displacement interpolation (Kornilov et al., 2024).
• Straightness and Non-Crossing: Trajectories learned via OTFM are straight lines in the latent space, minimizing acceleration and path curvature. The action-minimizing nature of OTFM has been formalized via connections to optimal acceleration transport (OAT), where straight lines correspond to minimal Lagrangian action (Yue et al., 29 Sep 2025).
• Geometric Attractivity: For manifold-supported targets, OTFM dynamics are terminally normally hyperbolic: normal perturbations decay exponentially while tangential perturbations remain neutral, as proven via Lyapunov analysis of the Moreau-envelope/extended Brenier potential (Fukumizu et al., 13 Feb 2026).
• Convergence of Minibatch OT-FM: As batch size increases, minibatch OTFM converges to the population solution in both cost and plan, and the vector fields achieve uniform convergence on compacts (Fukumizu et al., 13 Feb 2026, Boïté et al., 12 May 2026).
• Rectification Caveats: Merely enforcing a gradient (potential) constraint or zero flow-matching loss in rectified flows is insufficient for OT optimality; additional regularity (support connectivity, smoothness) is required to recover genuine OT transport (Hertrich et al., 26 May 2025).

6. Applications and Empirical Performance

OTFM methods have been applied successfully in diverse domains:

• High-Dimensional Generative Modeling: Significantly improved few-step and single-step image generation (CIFAR-10, ImageNet, FFHQ) with over 2× reduction in path curvature and state-of-the-art FID metrics, especially in low-number-of-function-evaluations regimes (Malnick et al., 2 Jun 2026, Kornilov et al., 2024).
• Remote Sensing and Fusion: One-step pan-sharpening surpasses diffusion-based and regression baselines in quality and speed (Cao et al., 19 Mar 2025).
• Physical and Biomedical Systems: Hilbert-space OTFM models fields arising in turbulent flows, supporting mesh-invariant learning and zero-shot super-resolution (Kunpeng et al., 7 Apr 2026); multiscale unbalanced OTFM enables trajectory inference on atlas-scale single-cell datasets (Peng et al., 15 May 2026).
• Graph and Molecular Generation: OTFM provides the backbone for fast, property-controllable discrete graph generation via permutation-invariant Graph Transformers (Hou et al., 2024); in molecular 3D conformation, equivariant OTFM yields state-of-the-art RMSD and recall/precision metrics (Tian et al., 2024).
• Quantization for Edge Inference: OT-based quantization of FM model parameters preserves fidelity and latent-space structure even at extreme (2–3 bit) compression (Varam et al., 14 Nov 2025).
• All-to-All and Conditional Transport: OTFM paired with minibatch-level OT couplings enables simultaneous learning of (nearly) all pairwise optimal transports between conditional distributions, with provable asymptotic convergence to the family of exact pairwise OT maps (Ikeda et al., 4 Apr 2025).

7. Open Challenges and Methodological Variants

While OTFM achieves exactness and efficiency under quadratic cost and convex potentials, challenges and variants include:

• Computational Overheads of Full OT Solvers: Even with minibatch approaches, scaling exact OT couplings remains difficult for large datasets or high dimensions. Prior design and identity couplings enable OT-optimal flows with tractable sampling (Malnick et al., 2 Jun 2026).
• Unbalanced and Supervised Extensions: Incorporating prior knowledge, hierarchical structure, and transition priors enables OTFM to scale to millions of datapoints and to adapt to biological constraints (Peng et al., 15 May 2026).
• Non-Quadratic or Generalized Costs: Theory and efficient algorithms for OTFM with arbitrary ground costs, or in settings with additional constraints (e.g., path dependence, collision avoidance in mean-field control), remain open areas, with recent advances bridging OTFM and mean-field games (Duan et al., 8 Oct 2025).
• Beyond Geometric OT Coupling: Model-aligned coupling (MAC) augments OTFM by selecting pairs that align with current model trajectories for improved straightness and learnability, though with increased computational cost (Lin et al., 29 May 2025).
• Discrete Data Domains: In categorical flow matching, OT couplings directly minimize state transitions, offering improved control over trajectory properties in settings without deterministic rectification (Haxholli et al., 2024, Hou et al., 2024).

---

References:

• "Optimal Flow Matching: Learning Straight Trajectories in Just One Step" (Kornilov et al., 2024)
• "Optimal Transport Flow Matching by Design" (Malnick et al., 2 Jun 2026)
• "Taming Flow Matching with Unbalanced Optimal Transport into Fast Pansharpening" (Cao et al., 19 Mar 2025)
• "Multiscale Supervised Unbalanced Optimal Transport Flow Matching" (Peng et al., 15 May 2026)
• "Expected Batch Optimal Transport Plans and Consequences for Flow Matching" (Boïté et al., 12 May 2026)
• "EquiFlow: Equivariant Conditional Flow Matching with Optimal Transport for 3D Molecular Conformation Prediction" (Tian et al., 2024)
• "Optimal-Transport-Guided Functional Flow Matching for Turbulent Field Generation in Hilbert Space" (Kunpeng et al., 7 Apr 2026)
• "OAT-FM: Optimal Acceleration Transport for Improved Flow Matching" (Yue et al., 29 Sep 2025)
• "Low-Bit, High-Fidelity: Optimal Transport Quantization for Flow Matching" (Varam et al., 14 Nov 2025)
• "Beyond Optimal Transport: Model-Aligned Coupling for Flow Matching" (Lin et al., 29 May 2025)
• "Improving Molecular Graph Generation with Flow Matching and Optimal Transport" (Hou et al., 2024)
• "On the Relation between Rectified Flows and Optimal Transport" (Hertrich et al., 26 May 2025)
• "Flow Matching from Viewpoint of Proximal Operators" (Fukumizu et al., 13 Feb 2026)
• "Minibatch Optimal Transport and Perplexity Bound Estimation in Discrete Flow Matching" (Haxholli et al., 2024)
• "Pairwise Optimal Transports for Training All-to-All Flow-Based Condition Transfer Model" (Ikeda et al., 4 Apr 2025)
• "Trajectory-Optimized Density Control with Flow Matching" (Duan et al., 8 Oct 2025)