Curly Flow Matching for Learning Non-gradient Field Dynamics

Modeling the transport dynamics of natural processes from population-level observations is a ubiquitous problem that arises in the natural sciences. A key challenge in these settings is making important modeling assumptions over the scientific process at hand that enable faithful learning of governing dynamics that mimic actual system behavior. Traditionally, the de-facto assumption present in approaches relies on the principle of least action that result in gradient field dynamics, that lead to trajectories that minimize 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 optimal transport based conditional flow 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 reference process with non-zero drift---in stark contrast from zero-drift reference processes---which is constructed using inferred velocities in addition to population snapshot data. We instantiate Curly-FM by solving the single-cell trajectory inference problem with approximate velocities inferred using RNA velocity. We demonstrate that Curly-FM can learn trajectories that match both RNA velocity and population marginals. Curly-FM expands flow matching models beyond the modeling of populations and towards the modeling known periodic behavior observed in cells.