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, such as biology and meteorology. We tackle this problem with a novel approach that leverages the connection between FM and dynamic optimal transport (OT), introducing time-dependent potential terms in the dynamic OT action that can steer the flow towards the intermediate marginals. 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 of OT-potential FM (OTP-FM) on diverse scientific datasets.