Optimal-Transport Conditional Flow Matching

• OT-CFM is a generative modeling framework that unifies continuous flows and dynamic optimal transport to create minimal-action, efficient sample trajectories.
• It employs ordinary differential equations to regress velocity fields towards optimal couplings, ensuring geodesic paths with minimal transport cost.
• Practical implementations use minibatch OT and condition-specific penalties to address train/test mismatches for high-dimensional and conditional tasks.

Optimal-Transport Conditional Flow Matching (OT-CFM) is a generative modeling framework that unifies methods from continuous normalizing flows, conditional generative modeling, and dynamic optimal transport. OT-CFM leverages the mathematical structure of optimal transport (OT) to train flows with straight (minimal-action, geodesic) sample trajectories via ordinary differential equations (ODEs), yielding both theoretical optimality and practical efficiency. The core innovation is to use optimal couplings—ideally, the true dynamic OT plan—as the target for conditional regression of velocity fields in flow-based models. OT-CFM and its extensions address a range of settings, including continuous or discrete conditions, high-dimensional data, and regressive (sparse) conditional tasks.

1. Mathematical Foundation and Core Principles

OT-CFM is rooted in the Benamou–Brenier dynamic OT formulation, where probability distributions are connected via time-dependent flows minimizing kinetic energy under continuity constraints. The probability flow ODE

$$\frac{d}{dt} x(t) = v_\theta(x(t), t \mid c_1, c_2)$$

is trained to connect samples from a source conditional distribution $P_{c_1}$ to a target $P_{c_2}$, using velocity fields regressed to the straight OT displacement: $u_t(x | x_0, x_1) = x_1 - x_0$ for linear (McCann) interpolants $x_t = (1 - t)x_0 + t x_1$, where pairings $(x_0, x_1)$ are generated according to an OT plan (ideally, the quadratic Wasserstein minimizer). This regression is performed under expectations over interpolation times, couplings, and sample trajectories, yielding the simulation-free training objective

$$L_\mathrm{OT-CFM}(\theta) = \mathbb{E}_{t, (x_0, x_1) \sim \pi, x \sim p_t(x | x_0, x_1)} \| v_\theta(t, x) - (x_1 - x_0) \|^2$$

where $\pi$ is an optimal coupling (see (Tong et al., 2023, Lipman et al., 2022)).

By design, OT-CFM optimizes straightest transport paths, which minimizes the expected squared velocity (kinetic action), yielding geodesics with minimal Wasserstein-2 (quadratic) cost. This path structure greatly improves integration accuracy and inference efficiency for generative ODE flows.

2. Practical Implementation and Computational Aspects

In practice, the true continuous OT plan between arbitrary distributions is rarely tractable. OT-CFM addresses this using minibatch OT: for each batch, the quadratic cost optimal assignment (Hungarian or Sinkhorn algorithm) is computed between batch pairs. In the conditional setting, data are paired according to both their values and auxiliary condition labels:

• For discrete conditions: OT is computed within each class (block-diagonal cost).
• For continuous conditions: A condition-augmented cost (e.g., sum of squared sample and condition differences) is used (see (Ikeda et al., 4 Apr 2025), Eq. 13).

The model is trained on pairs $(x_0, c_0)$ and $(x_1, c_1)$, minimizing a cost that incorporates both data and condition similarity: $$C(\pi) = \sum_{i=1}^N \|x_1^{(i)} - x_2^{\pi(i)}\|^2 + \beta \bigl( \|c_1^{(i)} - c_1^{\pi(i)}\|^2 + \|c_2^{(i)} - c_2^{\pi(i)}\|^2 \bigr)$$ enabling all-to-all transfer maps among conditions (A2A-FM). The vector field $v_\theta$ is trained to match linear (OT) paths between coupled pairs, for all $(c_1, c_2)$ (see (Ikeda et al., 4 Apr 2025), Eq. 14).

Efficient computation is facilitated by:

• Minibatch OT solvers for batchwise assignment.
• Semidiscrete OT approaches that approximate the discrete-to-continuous match (see (Mousavi-Hosseini et al., 29 Sep 2025)) to avoid quadratic scaling in mini-batch size.
• Closed-form pairing strategies for settings with block-triangular structure or repeated conditions (see (Alfonso et al., 2023)).

3. Addressing the Conditional Setting: Pitfalls and Algorithmic Fixes

Directly applying minibatch OT can introduce a train/test mismatch in conditional settings: the OT pairing entwines data and conditions, resulting in skewed priors at training that do not match test-time sampling, where priors and conditions are sampled independently. This issue is analyzed in detail in (Cheng et al., 13 Mar 2025), which shows standard OT coupling can degrade conditional generation performance.

The Conditional Conditional OT (C²OT) method resolves this by adding a condition-dependent penalty to the cost matrix. For discrete labels, only same-condition pairs are allowed; for continuous, the cost penalizes condition mismatch (e.g., cosine distance with adaptive weight). This ensures that the OT assignment respects condition-independence at training, aligning train and generation distributions. C²OT achieves straight, conditionally accurate flows, improving FID and class adherence in image generation over naive OT or vanilla flow matching ((Cheng et al., 13 Mar 2025), Table 2).

4. Theoretical Guarantees

When the coupling cost includes both sample and condition proximity, and the batch size $N \rightarrow \infty$ with conditional penalty $\beta \rightarrow \infty$, the empirical matching converges, in law, to the collection of pairwise quadratically optimal transports between all condition pairs $(P_{c_1}, P_{c_2})$. For almost every pair $(c_1, c_2)$,

$$\int_{C^2}\int_{X^2} \|x_1-x_2\|^2 d\Pi^*(x_1,x_2|c_1,c_2)\,dP_C(c_1)dP_C(c_2) = \int_{C^2} W_2^2(P_{c_1},P_{c_2})\,dP_C(c_1)dP_C(c_2)$$

This guarantees that OT-CFM, with generalized cost and sufficient coupling accuracy, learns the true pairwise OT map between conditionals ((Ikeda et al., 4 Apr 2025), Prop. 1, Thm. 1).

Additionally, restricting the learned velocity field to those determined by OT geodesics (the Benamou–Brenier class) coincides with minimization of the dynamic OT dual (Kornilov et al., 31 Oct 2025). When flow matching or action matching is constrained to these vector fields, the minimizer is the OT solution regardless of the choice of intermediate path; thus, AM, FM, and OT all become equivalent up to additive constants.

5. Empirical Performance and Applications

OT-CFM and its conditional extensions demonstrate broad empirical success:

• Synthetic data: OT-CFM and A2A-FM outperform standard flow matching, minibatch-OT, and multimarginal approaches in matching ground-truth OT coupling and lower path energy ((Ikeda et al., 4 Apr 2025), synthetic benchmarks).
• Chemical property optimization: On datasets with sparse or continuous properties (e.g., QED, logP), OT-CFM and A2A-FM achieve high success rates under structural similarity constraints, outperforming COATI-LDM, MolMIM, partial diffusion, and other baselines in candidate expansion tasks (Ikeda et al., 4 Apr 2025).
• Image generation: On benchmark datasets (CIFAR-10, ImageNet, CelebA), C²OT provides better FID, class adherence, and faster ODE integration than both naive OT and vanilla flow matching (Cheng et al., 13 Mar 2025).
• Speech and gesture synthesis: OT-CFM underlies architectures like Matcha-TTS and joint speech-gesture models, supporting unified multimodal generation with efficient sampling and superior cross-modal coherence (Mehta et al., 2023, Mehta et al., 2023).
• Molecular conformation: EquiFlow applies OT-CFM with SE(3)-equivariant transformers, achieving state-of-the-art performance in 3D molecular prediction (Tian et al., 15 Dec 2024).

6. Extensions, Limitations, and Comparisons

Extensions:

• All-to-all transfer: A2A-FM supports simultaneous learning of $T_{c_1 \to c_2}$ for all conditional pairs, suitable for continuous or high-cardinality $c$ (Ikeda et al., 4 Apr 2025).
• Semidiscrete matching: SD-FM enables scalable pairing for large datasets (Mousavi-Hosseini et al., 29 Sep 2025).
• Weighted CFM (W-CFM): Approximates entropy-regularized OT via Gibbs kernel weighting, avoiding batch-OT computation (Calvo-Ordonez et al., 29 Jul 2025).

Limitations:

• Cyclic consistency is not guaranteed: transport maps between $(c_1, c_2)$ and $(c_2, c_3)$ may not compose to $(c_1, c_3)$, as OT is only pairwise optimal (Ikeda et al., 4 Apr 2025).
• Hyperparameter tuning: Conditional penalty $\beta$ and entropic regularization $\varepsilon$ need empirical selection.
• Discontinuity in $c \to P_c$: Highly discontinuous conditionals reduce estimation quality.

Algorithmic comparisons are summarized below:

Approach	All-Pairs OT?	Train/Test Match	Scalability	Straight Paths	Used In
Naive Minibatch OT	No	No	Poor (quad)	Yes	Early OT-CFM models
C²OT	Yes	Yes	Good	Yes	(Cheng et al., 13 Mar 2025)
A2A-FM	Yes	Yes	Good	Yes	(Ikeda et al., 4 Apr 2025)
SD-FM	Yes	Yes	Excellent	Yes	(Mousavi-Hosseini et al., 29 Sep 2025)
Multimarginal SI	Only finite	Yes	Poor	No	(Ikeda et al., 4 Apr 2025)

7. Broader Significance and Outlook

OT-CFM provides a principled, theoretically grounded framework for training flows in generative modeling, connecting the learning objective directly to dynamic OT. This leads to sample trajectories with minimal transport cost, efficient ODE-based sampling, and improved training stability, especially when extended to complex conditional, multi-attribute, or regressive data. Recent advances address the scalability bottlenecks of minibatch OT, the statistical consistency of estimation with empirical conditionals, and correct alignment between training and deployment distributions.

By unifying regularized OT, flow matching, and conditional regression, OT-CFM and its descendants (C²OT, A2A-FM, W-CFM, SD-FM) define the current methodological state-of-the-art for all-to-all, scalable, and theoretically optimal conditional generative modeling.