Curly Flow Matching for Learning Non-gradient Field Dynamics

Michael Bronstein (USI) · Joey Bose (University of Oxford/Mila) · Alexander Tong (Aithyra) · Katarina Petrović (University of Oxford) · Lazar Atanackovic (Broad Institute of MIT-Harvard) · Viggo Moro (University of Oxford) · Kacper Kapusniak (University of Oxford) · Ismail Ilkan Ceylan (TU Wien, AITHYRA, University of Oxford)

computational fluid dynamicsdrift reference processenergy functionalflow matching modelsgradient field dynamicsleast action principlenon-gradient behaviorocean currentsperiodic behaviorpopulation marginalspopulation-level observationsprobability measuresschrödinger bridge problemtrajectory inferencetransport dynamics

Modeling the transport dynamics of natural processes from population-level observations is a ubiquitous problem in the natural sciences. Such models rely on key assumptions about the underlying process in order to enable faithful learning of governing dynamics that mimic the actual system behavior. The de facto assumption in current approaches relies on the principle of least action that results in gradient field dynamics and leads to trajectories minimizing an energy functional between two probability measures. However, many real-world systems, such as cell cycles in single-cell RNA, are known to exhibit non-gradient, periodic behavior, which fundamentally cannot be captured by current state-of-the-art methods such as flow and bridge matching. In this paper, we introduce Curly Flow Matching (Curly-FM), a novel approach that is capable of learning non-gradient field dynamics by designing and solving a Schrödinger bridge problem with a non-zero drift reference process