Multimarginal flow matching with optimal transport potentials

Abstract: Flow matching (FM) has emerged as a powerful framework for learning dynamic transport maps between two empirical distributions. However, less explored is the setting with intermediate observed marginals that can help constrain the flows between the endpoints. This "multimarginal" regime is central to modeling temporal evolution in dynamical systems in many scientific domains that can sample sequential distributions. We tackle this problem with a novel approach that leverages the connection between FM and dynamic optimal transport (OT), softly steering the flow towards the intermediate marginals through potential terms in the dynamic OT action. By extending the conditional FM learning target to incorporate these potentials, we derive an efficient, simulation-free algorithm for multimarginal FM that offers considerable flexibility in the spatiotemporal dynamics of the learned flows. We demonstrate state-of-the-art performance and training efficiency of OT-potential FM (OTP-FM) on diverse single-cell RNA sequencing, oceanographic, and meteorological datasets. Our code is available at https://github.com/Bexorg-Inc/OTP-FM.

• The paper introduces OTP-FM, a framework that relaxes hard multimarginal constraints by incorporating soft potential energy terms for smoother trajectory interpolation.
• It leverages dynamic optimal transport and gradient-based corrections to improve physical plausibility and temporal consistency in trajectory modeling.
• Empirical evaluations show that OTP-FM outperforms traditional CFM and prescriptive methods in accuracy, computational efficiency, and scalability.

Multimarginal Flow Matching with Optimal Transport Potentials

Problem Formulation and Background

Multimarginal trajectory inference is a central question in modeling dynamical systems with observed empirical distributions at multiple time points, prevalent in domains such as developmental biology (scRNA-seq), climate science, and oceanography. Conditional flow matching (CFM) has established itself as a simulation-free, scalable approach to learning continuous-time transport maps, typically by regressing onto conditional optimal transport (OT) solutions between pairs of samples drawn from endpoint distributions. However, the multimarginal setting—with observed intermediate marginals—poses significant challenges for ensuring physically plausible, temporally consistent trajectories.

Standard approaches stitch CFM trajectories piecewise between consecutive marginals, yielding unphysical discontinuities at boundaries. Recent methods like MMFM and 3MSBM introduce prescriptive smoothing, but lack a principled connection to the underlying dynamics, frequently misrepresenting system behavior. The paper "Multimarginal flow matching with optimal transport potentials" (2606.05327) proposes a rigorously justified, flexible generalization: it relaxes hard multimarginal constraints by introducing soft potential energy terms in the dynamic OT action, thereby enabling the specification of spatiotemporal dynamics through tunable parameters and distances.

Figure 1: Comparison of standard/stitched multimarginal CFM, prescriptive approaches (MMFM/3MSBM), and OTP-FM's soft, potential-driven flows; method overview illustrating decomposition into base CFM plus marginal corrections.

Dynamic OT with Soft Potentials

Dynamic OT (DOT) solves for a minimal kinetic energy transport path $\rho_t$, parameterized by velocity field $u_t$, satisfying boundary conditions and (typically) the continuity equation. The classical solution corresponds to straight-line sample trajectories in state space, which, in the conditional case, is the standard regression target for CFM.

The key insight of the paper is recasting piecewise CFM as conditional DOT with singular potential terms enforcing hard constraints on the intermediate marginals; these appear as Dirac-delta penalties weighted by infinite strengths. OTP-FM replaces these singularities with smooth, finite-strength potentials, with temporally localized kernels $\lambda_k(t)$, statistical distances $\mathcal{D}$ (e.g., $\mathcal{W}_2$, MMD, KLD), and strength parameters $w_k$.

The Euler-Lagrange equations induced by this action:

$$u_t(x) = \nabla \varphi_t(x), \quad \partial_t \varphi_t(x) + \frac{1}{2}\|\nabla \varphi_t(x)\|^2 = \sum_{k=1}^K w_k \lambda_k(t) g_k(x, t),$$

enable gradient-based correction towards intermediate marginals, where $g_k(x, t)$ is the functional derivative of $\mathcal{D}$ w.r.t. $\rho_{t_k}$.

Figure 2: Exact solutions to the marginal dynamic OTP problem for 1D Gaussian marginals, visualizing tunable effects of potentials on trajectory smoothness and interpolation.

Conditional OTP-FM, Fixed-Point Structure, and Algorithm

The practical algorithm derives a conditional regression target by integrating the force terms over the temporal kernel and employing predictions at the marginal times. Sample trajectories are decomposed as:

$u_t$0

where $u_t$1 denotes the second integral. This target reduces to standard CFM in the hard-potential limit, and yields a simulation-free, closed-form fixed-point for the $u_t$2 potential, or iterative solution for nonlinear estimators (e.g., MMD, KLD).

OTP-FM is agnostic to the choice of consistency model. The authors employ improved MeanFlow (iMF), yielding state-of-the-art one- and two-step inference, but also validate with LSD self-distillation and others.

Figure 3: OTP-FM trajectories inferred from increasing steps per marginal, demonstrating consistency model efficiency for few-step inference.

Theoretical Guarantees and Design Space

The relaxation to soft potentials admits a broad design space: strength $u_t$3, temporal localization $u_t$4 (width of kernel), marginal influence (via $u_t$5), and statistical coupling choice (e.g., product of pairwise OT couplings vs. independent coupling). The paper establishes rigorous bounds, contravening prior beliefs that only hard constraints enforce interpolation:

• Marginal alignment bound: For $u_t$6-type potentials, the induced trajectory's divergence from the ground-truth marginals scales as $u_t$7.
• End-to-end bound: Combined with consistency-model learning rates, overall error of the learned flow map is upper-bounded by a function of training loss, potential strength, and dataset-dependent constants.

These results demonstrate the trade-off between potential strength, interpolation fidelity, and computational stability—directly informing practical hyperparameter tuning.

Empirical Evaluation: Benchmarks and Ablations

Comprehensive empirical studies span high-dimensional single-cell RNA sequencing (EB, CITE), oceanographic, and meteorological datasets. In all interpolation protocols where held-out marginals are evaluated, OTP-FM achieves best or comparable performance across metrics (average $u_t$8, MMD), substantially outpacing simulation-based methods and prescriptive interpolations. Notably, training times are reduced to minutes per run.

Figure 4: Pareto plots of training time vs. performance on CITE 5D and EB 100D datasets; OTP-FM achieves optimal tradeoffs.

Ablation studies validate optimality and robustness with respect to:

• Potential type (hard vs. soft, $u_t$9 vs. MMD/KLD),
• Strength parameter,
• Temporal kernel shape/width,
• Curriculum schedule (homotopy continuation stabilizes fixed-point convergence).

Figure 5: Loss curves for different loss weightings (MSE, adaptive, learned log-variance) and resulting OTP-FM trajectories.

Furthermore, data-driven design flexibility is key for datasets with nontrivial trajectory structure (e.g., Beijing PM2.5 air quality), where prescriptive approaches (e.g., cubic-spline MMFM) fail qualitatively and quantitatively.

Figure 6: Ground truth and simulated trajectories for Beijing air quality; MMFM's cubic interpolation fails to fit marginal distributions, OTP-FM adapts to observed density.

Practical Recommendations and Limitations

OTP-FM's $\lambda_k(t)$0 potential is preferred for its linear fixed-point structure, low gradient variance, and empirical performance. Default parameters: strength $\lambda_k(t)$1, equally spaced Gaussian kernel, sigmoid curriculum. Precomputing OT couplings ($\lambda_k(t)$2) is worthwhile in most applications, but the independent coupling variant ($\lambda_k(t)$3) provides near-equivalent performance scaled to larger datasets.

Limitations include the tuning requirement imposed by the expanded design space, sensitivity to sparsely-sampled or noisily-labeled marginals, and kernel-based estimator bias for MMD/KLD. Integration with manifold learning (e.g., MFM, pixel MeanFlow), domain-specific inductive biases (e.g., cell growth in VGFM), and direct variational multimarginal OT solvers (e.g., WLF-UOT) would further improve trajectory plausibility and consistency.

Conclusion

OTP-FM introduces a rigorously motivated, computationally efficient framework for multimarginal trajectory inference by relaxing hard constraints to soft potentials in dynamic OT. The method yields exact conditional solutions, guarantees interpolation fidelity via tunable parameters, and outperforms state-of-the-art baselines across diverse scientific datasets. Its flexibility, empirical efficiency, and theoretical guarantees make it a robust foundation for future multimarginal generative modeling research, especially as intermediate marginal observations become pervasive in complex temporal systems.