Curly Flow Matching for Learning Non-gradient Field Dynamics

TL;DR

Curly-FM tackles the limitation of gradient-only trajectory models by introducing a two-stage, simulation-free Schrödinger- bridge framework with a non-zero drift reference, enabling the learning of non-gradient, especially cyclic, dynamics from population snapshots and velocity cues. It learns a neural path interpolant to align with a constructed reference drift and then forms a marginal process as a mixture of conditional bridges guided by OT couplings, optimizing a combined flow and score objective. Across synthetic cyclic flows, ocean currents, single-cell trajectories, and CFD data, Curly-FM yields velocity fields and trajectories that closely match the reference dynamics and endpoint marginals, while offering substantial computational efficiency over simulation-based methods. This work broadens flow-matching models to non-gradient dynamics and periodic phenomena, with potential impacts on trajectory inference in biology and physics; future work includes improving velocity-field estimation and extending to unbalanced transport and multiomics data.

Abstract

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 -- in stark contrast to typical zero-drift reference processes -- which is constructed using inferred velocities in addition to population snapshot data. We showcase Curly-FM by solving the trajectory inference problems for single cells, computational fluid dynamics, and ocean currents with approximate velocities. We demonstrate that Curly-FM can learn trajectories that better match both the reference process and population marginals. Curly-FM expands flow matching models beyond the modeling of populations and towards the modeling of known periodic behavior in physical systems. Our code repository is accessible at: https://github.com/kpetrovicc/curly-flow-matching.git

Figures (11)

• Figure 1: Particle trajectories generated between samples drawn from asymmetric circles distribution at $t=0$ and $t=1$ and respective to underlying reference velocity field $f_t(x_t)$. Traditional flow-based models such as OT-CFM and CFM cannot capture cyclical patterns in physical systems. Curly-FM is capable of learning non-gradient field dynamics behavior in the underlying data.
• Figure 2: Visualization of ground truth data and vectorfield (left), OT-CFM predicted trajectories (center), and Curly-FM predictions (right). Curly fits the vortex much better than OT-CFM.
• Figure 3: Ground truth data.
• Figure 4: Vectorfields (top) and trajectory traces (bottom) learned using Curly-FM (left) OT-CFM (center) CFM (right). Curly-FM is the only method able to learn the cell cycle.
• Figure 5: Visualization of ground truth data and vectorfield (left), OT-CFM predicted trajectories (center) and Curly-FM predictions (right). Curly fits the ground truth much better than OT-CFM.

Theorems & Definitions (1)

• Remark 1