Multimarginal flow matching with optimal transport potentials

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.

Many scientific domains, such as the study of cellular development, ocean currents, and climate, deal with environments that are easy measure as snapshots but hard to follow: we are rarely able to track any single cell or particle through its full evolution. Inferring continuous trajectories from these unconnected snapshots is a central challenge across the biological and physical sciences. We propose OTP-FM, which generalizes the popular flow matching algorithm to treat each snapshot as a potential that exerts an attractive force on the evolving trajectories, guiding them in a theoretically principled manner towards the measured data. The strength, shape, and region the potential acts over are all tunable to best fit the data. OTP-FM further admits an efficient simulation-free training algorithm (i.e., very fast to train) and achieves state-of-the-art accuracy and efficiency on biological, ocean-current, and air-quality datasets.