Optimal Transport CFM

• OT-CFM is a principled framework for learning flow-based generative models by directly regressing optimal transport-induced constant-velocity fields.
• It minimizes path energy through direct OT coupling, enabling fast training and efficient ODE-based sampling with high sample quality.
• The method extends naturally to conditional and multi-domain settings, supporting tasks such as molecular conformations, speech synthesis, and image style transfer.

Optimal Transport Conditional Flow Matching (OT-CFM) is a principled, simulation-free framework for learning flow-based generative models by regressing time-dependent vector fields to optimal transport-induced conditional flows. OT-CFM replaces indirect likelihood or score-matching objectives with direct regression against a constant-velocity field derived from optimal transport pairings between distributions, producing flows with minimal path energy and straight trajectories—thereby enabling both fast training and efficient, high-fidelity sampling through ODE integration. The method extends naturally to conditional settings, aligning prior and data distributions under side information or conditioning variables, and supports both discrete and continuous conditioning as well as equivariant constraints for structured data. This framework is foundational for state-of-the-art approaches in molecular conformation prediction, speech and gesture synthesis, and multi-domain conditional generative modeling (Tian et al., 2024, Tong et al., 2023, Ikeda et al., 4 Apr 2025, Mehta et al., 2023, Mehta et al., 2023, Generale et al., 2024).

1. Mathematical Foundations and Objective

Let $p_0(x_0)$ denote a tractable base distribution (e.g., isotropic Gaussian in $\mathbb{R}^d$) and $p_1(x_1|c)$ a complex data distribution conditioned on side information $c\in\mathcal{C}$ (such as atom/bond types for molecular data). OT-CFM seeks a time-dependent vector field

$$v_\theta : \mathbb{R}^d \times [0,1] \times \mathcal{C} \to \mathbb{R}^d$$

satisfying

$\frac{dx_t}{dt} = v_\theta(x_t, t \mid c)$

so the initial point $x_{t=0}\sim p_0$ is transported to $x_{t=1} \sim p_1(\cdot|c)$ via ODE integration.

The coupling between $x_0$ and $x_1$ is determined by the optimal transport plan

$$\pi^*(dx_0, dx_1|c) \in \arg\min_{\pi\in\Pi(p_0, p_1)} \iint \|x_0-x_1\|^2\, \pi(dx_0, dx_1|c)$$

which induces a straight-line interpolation

$x_t = (1-t) x_0 + t x_1$

with constant velocity

$v^*(x_t, t|c) = x_1 - x_0$

The core flow-matching loss is then

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

which directly regresses $v_\theta$ onto the ground-truth OT velocity (Tian et al., 2024, Tong et al., 2023, Lipman et al., 2022).

2. Algorithmic Structure and Implementation

The OT-CFM workflow involves alternating between OT plan computation, regression on straight-line velocities, and ODE-based sampling:

Training (per batch)

1. Sample mini-batch $\{(c^{(b)}, x_1^{(b)})\}$.
2. Draw noise samples $\{x_0^{(b)}\sim p_0\}$.
3. (If data are point clouds) Align each $(x_0^{(b)}, x_1^{(b)})$ for translation/rotation equivariance (e.g., center-of-mass subtraction, Kabsch algorithm).
4. Solve discrete OT (e.g., Sinkhorn) for pairs $(x_0^{(b)}, x_1^{(b)})$.
5. For each pair and sampled $t\sim U[0,1]$:
   • Compute $x_t$, reference velocity $u = x_1 - x_0$.
   • Compute $v_\theta(x_t, t|c)$.
   • Accumulate loss $\| v_\theta - u \|^2$.
6. Update $\theta$ via backpropagation (Tian et al., 2024, Tong et al., 2023).

Sampling

1. Given $c$ (conditions), sample $x_0 \sim p_0$ and align as appropriate.
2. Numerically solve the ODE $\frac{dx_t}{dt} = v_\theta(x_t, t|c)$ from $t=0$ to $t=1$ (Dormand–Prince or any accurate ODE solver).
3. Output $x_1\approx x_{t=1}$ as a sample from $p_1(\cdot|c)$.

Practical models use graph-based equivariant transformers for $v_\theta$ in structure prediction tasks, U-Nets or 1D CNN+Transformer hybrids for sequential data, and sinusoidal or rotary time embeddings. Optimizers are typically AdamW with moderate batch sizes ($128$–$256$), and minibatch OT is solved using either exact or entropy-regularized solvers (Tian et al., 2024, Tong et al., 2023, Mehta et al., 2023, Mehta et al., 2023).

3. Conditional and All-to-All Generalizations

OT-CFM extends to multi-conditional and all-to-all transfer by defining maps $T_{c_1 \to c_2}: X \to X$ for each $(c_1, c_2) \in \mathcal{C} \times \mathcal{C}$ such that $T_{c_1 \to c_2}\#P_{c_1} = P_{c_2}$ and optimally minimizes

$\int_X \|x - T_{c_1 \to c_2}(x)\|^2\, dP_{c_1}(x)$

Batchwise, this translates to solving for a permutation or assignment minimizing

$$\sum_{i=1}^N \Big( \|x_1^{(i)} - x_2^{(\pi(i))}\|^2 + \beta (\|c_1^{(i)} - c_1^{(\pi(i))}\|^2 + \|c_2^{(i)} - c_2^{(\pi(i))}\|^2) \Big)$$

across provided condition pairs, enabling learning and evaluation across continuous and regressive condition spaces (Ikeda et al., 4 Apr 2025, Generale et al., 2024).

For generalization to settings with unpaired data or continuous conditioning, OT-CFM incorporates kernel-weighted, entropic OT couplings and amortizes the flow field over all $c$, facilitating both scalability and variance reduction without requiring data paired across all conditions (Generale et al., 2024, Ikeda et al., 4 Apr 2025). Extensions enforce cycle consistency or antisymmetry when needed (Ikeda et al., 4 Apr 2025).

4. Theoretical Guarantees and Properties

OT-CFM, when using the true OT plan, yields vector fields realizing the Benamou–Brenier dynamic optimal transport flow. In the small-noise or exact interpolation limit, the marginal drift induced by OT-CFM solves the dynamic OT minimization problem

$$W_2^2 = \inf_{(p_t, u_t)} \int_0^1\!\int p_t(x)\|u_t(x)\|^2\,dx\,dt,\quad \partial_t p_t + \nabla\!\cdot(p_t u_t) = 0$$

with minimal kinetic energy (Tong et al., 2023, Kornilov et al., 31 Oct 2025, Lipman et al., 2022). Empirically, this produces flows with minimal path curvature and reduced trajectory energy as quantified by normalized path energy (NPE) and empirical $W_2^2$ metrics. Variance of the regression target vanishes as the plan converges to OT, permitting faster and more stable convergence in training (Tong et al., 2023).

Equivalences have been established with action-matching and Benamou–Brenier formulations under optimal vector fields, demonstrating that under restriction to OT fields, action-matching and OT problems coincide up to constants (Kornilov et al., 31 Oct 2025).

5. Comparison to Alternative Methods

Method	Training Regime	Coupling	Target Field	Inference
Score Matching	Regression on score	None	$\nabla_x\log p_t$	SDE/ODE, slow
FM / I-CFM	Regression (indep. pairs)	$p_0\otimes p_1$	$x_1 - x_0$	ODE, geometric
OT-CFM	Regression (OT pairs)	OT plan $\pi^*$	$x_1 - x_0$ (OT)	ODE, faster
Diffusion Models	Score regression	None	Time-dependent	SDE/ODE, slow

OT-CFM achieves straight and short trajectories, minimal path energy, and deterministic ODE-based sampling with drastically fewer function evaluations compared to diffusion models (e.g., 2–10 vs hundreds–thousands) while preserving or surpassing sample quality (MOS in TTS, FID in images) (Tian et al., 2024, Mehta et al., 2023, Mehta et al., 2023). Unlike pure independent coupling flow matching (I-CFM), OT-CFM aligns prior and data via OT, materially reducing target variance and inference cost. Weighted CFM (W-CFM) and semidiscrete FM (SD-FM) offer further computational savings or avoid batchwise OT when scaling, but converge to OT-CFM in infinite-batch or dual-potential limits (Calvo-Ordonez et al., 29 Jul 2025, Mousavi-Hosseini et al., 29 Sep 2025).

6. Applications and Variants

OT-CFM has been successfully implemented in:

• 3D molecular conformation prediction, via EquiFlow using an equivariant transformer as $v_\theta$ and geometrically-aware OT (RMSD/Kabsch alignment), yielding higher accuracy and faster sampling over diffusion-based SDEs for the QM9 dataset (Tian et al., 2024).
• Conditional flow transfer across domains, as in all-to-all molecular property optimization and image style transfer, demonstrating state-of-the-art sample efficiency and performance under continuous conditions (Ikeda et al., 4 Apr 2025).
• Fast text-to-speech (TTS) and multimodal speech/gesture synthesis, where OT-CFM yields compact architectures and enables high-fidelity generation in only a handful of ODE steps, outperforming denoising-score diffusion models in real-time factors and mean opinion scores (Mehta et al., 2023, Mehta et al., 2023).
• Amortized conditional forecasting and domain translation, supporting unpaired $\{x, c\}$ datasets with entropic OT and kernel-weighted losses for accurate, efficient conditional generative modeling (Generale et al., 2024).

Extensions address conditional-prior mismatch, anti-symmetric flows, cycle consistency, and computational bottlenecks. Minibatch and semidiscrete OT, entropic regularization, and weighted losses allow OT-CFM to retain efficiency and theoretical guarantees with large, high-dimensional or multi-modal data (Ikeda et al., 4 Apr 2025, Calvo-Ordonez et al., 29 Jul 2025, Mousavi-Hosseini et al., 29 Sep 2025).

7. Limitations and Practical Considerations

The principal computational cost in OT-CFM is the per-batch OT coupling, which scales $\mathcal{O}(B^3)$ (Hungarian) or $\mathcal{O}(B^2)$ (Sinkhorn) in batch size. For problems with large datasets or high dimensionality, approximate global potentials, large-batch weighted methods, or amortized dual estimators (semidiscrete OT) ameliorate this cost (Calvo-Ordonez et al., 29 Jul 2025, Mousavi-Hosseini et al., 29 Sep 2025, Generale et al., 2024). Care must be taken in selecting conditional couplings to avoid prior skew in the conditional setting; conditional OT with appropriate penalty terms (e.g., C$^2$OT, kernel-reweighted losses) is mandatory for preserving correct marginalization during training and inference (Cheng et al., 13 Mar 2025, Generale et al., 2024). The choice of regularization in OT (entropy, condition penalty) must be tuned for convergence, and cycle consistency is not guaranteed without explicit constraints (Ikeda et al., 4 Apr 2025).

---

OT-CFM constitutes the current state-of-the-art in efficient, regularized conditional continuous normalizing flow construction, providing robust connections to optimal transport theory and yielding consistent empirical advantages in pathway straightness, convergence, and computational efficiency for diverse generative modeling tasks (Tian et al., 2024, Tong et al., 2023, Ikeda et al., 4 Apr 2025, Mehta et al., 2023, Mehta et al., 2023, Calvo-Ordonez et al., 29 Jul 2025, Generale et al., 2024, Kornilov et al., 31 Oct 2025).