Curly Flow Matching

Katarina Petrović1, Lazar Atanackovic2,3,4, Viggo Moro1, Kacper Kapuśniak1, İsmail İlkan Ceylan5,6,1, Michael Bronstein1,6, Joey Bose1,7*, Alexander Tong6,7,8*
1University of Oxford, 2Broad Institute of MIT and Harvard, 3University of Toronto,
4Vector Institute, 5TU Wien, 6AITHYRA, 7Mila – Quebec AI Institute, 8Université de Montréal
NeurIPS 2025
*Indicates Equal Advising
Paper Code arXiv
Curly-FM teaser figure

Traditional flow-based models such as Conditional Flow Matching cannot capture periodic patterns. Curly-FM learns non-gradient field trajectories ✅

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.

Introduction

Curly Flow Matching (Curly-FM) is a flow-based model that aims to learn non-gradient field dynamics in complex systems which evolve with respect to some prescribed non-zero drift reference process. In this paper we solve non-zero drift Schrödinger Bridge (SB) problem in order to learn periodic or cyclical patterns in the data.

To train a Curly Flow Matching model, we assume access to velocity information prior. For example, this can be RNA velocity in single-cell data or fluid particle velocity in turbulent flows. We model non-gradient field dynamics by decomposing this problem into a two-staged algorithm where first stage learns neural bridge matching the bridge to designed reference process and second stage where we learn couplings and transport plan between marginal data distributions.

Because of its simulation-free nature, Curly-FM surpasses other simualation-based methods in its computational speed and efficiency. Further, given Curly-FM considers non-zero reference process with real velocity priors, it offers an interpretable framework for analyzing trajectories across real-world domains.

Theory & Framework

Goal

We aim to learn transport dynamics between population marginals while leveraging a reference process over trajectories. Unlike standard gradient-field assumptions, our target systems can exhibit non-gradient and periodic behavior. We model this as a Schrödinger Bridge problem between path measures $\mathbb{P}$ for particle evolution and $\mathbb{Q}$ reference process.

Schrödinger Bridge Problem (SB): Optimal path measure $\mathbb{P}^*$ can be found by minimizing KL divergence
$$\mathbb{P}^{*} = \text{arg min}_{\theta} \Big[ \mathrm{KL}\!\big(\mathbb{P}_{\theta}\,\|\,\mathbb{Q}\big) :\ \mathbb{P}_{0}=\mu_{0},\ \mathbb{P}_{1}=\mu_{1}\Big]$$

Zero-drift Schrödinger Bridge

In classical Schrödinger bridge formulations, the reference process is typically a zero-drift diffusion (e.g., Brownian motion). In such cases particle evolution and reference process are modeled as stochastic process via following Itô SDEs

Schrödinger Bridge (zero drift)
$$\mathrm{d}\mathbf{X}_t \;=\; v_{t,\theta}(\mathbf{X}_t)\,\mathrm{d}t \;+\; g_t\,\mathrm{d}\mathbf{B}_t \quad \text{particle evolution}$$ $$\mathrm{d}\mathbf{X}_t \;=\; g_t\,\mathrm{d}\mathbf{B}_t \quad \text{reference process}$$

The optimal bridge then takes form of Brownian bridge given as

$$\mathbb{Q}_t(\,\cdot \mid x_0,x_1) = \mathcal{N}\!\left(x_t;\; t x_1 + (1-t)x_0,\; t(1-t)\sigma^2\right)$$

Non-zero-drift Schrödinger Bridge

In Curly-FM we consider non-zero drift reference process $\mathbb{Q}$ which unlike zero-drift SB does not have closed form solution.

Schrödinger Bridge (non-zero drift)
$$\mathrm{d}\mathbf{X}_t \;=\; v_{t,\theta}(\mathbf{X}_t)\,\mathrm{d}t \;+\; g_t\,\mathrm{d}\mathbf{B}_t \quad \text{particle evolution}$$ $$\mathrm{d}\mathbf{X}_t \;=\; f_t(\mathbf{X}_t)\,\mathrm{d}t \;+\; g_t\,\mathrm{d}\mathbf{B}_t \quad \text{reference process}$$

To solve non-zero reference SB we construct neural bridge $\mathbb{Q}_{\nu,t}(\,\cdot \mid x_0,x_1)$ such that

$$\mathbb{Q}_t(\,\cdot \mid x_0,x_1) = \mathcal{N}\!\left(x_t;\; \mu_{t,\nu},\; t(1-t)\sigma^2\right)$$

where $\mu_{t,\nu}$ is a neural interpolant with neural network $\varphi_{t,\eta}(x_0,x_1)$ designed to match reference drift $$\mu_{t,\nu} = t x_1 + (1-t)x_0 + t(1-t)\varphi_{t,\eta}(x_0,x_1).$$

Results

Cell Cycles: Curly-FM learns cell sycle trajectories

Result 1 figure

Ocean Currents: Curly-FM learns periodic motion in ocean data

Result 2 figure

Mouse Erythroid: Curly-FM learns developmental pathways in mouse erythroid

Result 3 figure

Poster

BibTeX

@inproceedings{
petrovic2025curly,
title={Curly Flow Matching for Learning Non-gradient Field Dynamics},
author={Katarina Petrovi{\'c} and Lazar Atanackovic and Viggo Moro and Kacper Kapu{\'s}niak and Ismail Ilkan Ceylan and Michael M. Bronstein and Joey Bose and Alexander Tong},
booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
year={2025},
url={https://openreview.net/forum?id=7cqKVDgFZQ}
}