Curly Flow Matching

Description

Curly Flow Matching (Curly-FM) is a new flow matching framework for learning non-gradient field dynamics by solving non-zero drift Schrödinger Bridge (SB) problem. In contrast, traditional flow-based generative models such as Conditional Flow Matching (CFM) commonly rely on least action objectives modeling gradient field dynamics. As a result, they struggle to represent inherently rotational, cyclical, or periodic behavior in data.

In natural sciences, most processes undergo non-gradient field dynamics. One such example is the cell cycle, where cells follow periodic trajectories in gene expression space, as revealed by RNA velocity. Curly-FM can recover such cyclical dynamics across cell cycle data, developmental trajectories, ocean currents, and turbulent flows, unlike widely-used trajectory inference baselines such as CFM and OT-CFM. We study both single marginal (acess to $p_0$ and $p_1$) and multi-marginal (access to $p_{t, t\in[0,1]}$) settings.

Curly-FM is a simple two-stage framework. First, we learn neural bridges between marginal distributions by training a neural interpolant to match the drift of a chosen reference process. Second, we estimate the transport plan by computing a coupling that minimizes the discrepancy between the induced drifts, and then train the model using marginal flow matching objective.

This repo contains all elements needed to reproduce our results. See this http link for the paper.

The preprocessed data can be downloaded here:

• Preprocessed cell cycle data which contains gene expressions and RNA velocities
• Preprocessed oceans data which contains particle positions and velocities
• Preprocessed mouse erythroid data which contains gene expressions and RNA velocities across multiple marginals
• Pre-processed CFD data

The raw data can be downloaded here: Raw cell cycle data and Raw mouse erythroid data. For usability, we provide the notebook cell_data.ipynb which contain code for the data preprocessing, and code for data visualizations.

When adding the downloaded data to the repo make sure to create a data/ folder to save the downloaded datasets. For CFD the data should be saved under data/CFD/2DTGV. 💡 Remember to also change the data directory names in dataloader configs and notebooks.

If you find this code useful in your research, please cite our work.

K. Petrović, L. Atanackovic, V. Moro, K. Kapuśniak, İ. Ceylan, M. Bronstein, J. Bose*, A. Tong*. "Curly Flow Matching for Learning Non-gradient Field Dynamics." The Thirty-ninth Annual Conference on Neural Information Processing Systems (NeurIPS), 2025.

@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}
}

How to run

Install dependencies

# clone project
git clone https://github.com/kpetrovicc/curly-flow-matching.git
cd curly-flow-matching

# [OPTIONAL] create conda environment
conda create -n curlyfm python=3.9
conda activate curlyfm

# install pytorch according to instructions
# https://pytorch.org/get-started/

# install requirements
pip install -r requirements.txt
pip install -e.

Train cell cycles and CFD

Train model with chosen experiment configuration from configs/experiment/

python src/train.py experiment=experiment_name

You can override any parameter from command line like this

python src/train.py experiment=experiment_name trainer.max_epochs=1234 seed=42

Train ocean currents

To train a model via CurlyFM on the ocean currents, use ocean currents notebook in $\sigma\rightarrow 0$ setting and SDE ocean currents notebook in stochastic setting

Train mouse erythroid

To train a model via CurlyFM on the ocean currents, use 2d mouse erythroid notebook in 2d dimensions and nd mouse erythroid notebook for higher dimensions

Visualizing cell cycle trajectories

To visualize cell cycle trajectories use cell cycle visualization notebook

Synthetic Assymetric Circles experiments

To reproduce trajectories in assymetric circles use this notebook

Contributions

Have a question? Found a bug? Missing a specific feature? Feel free to file a new issue, discussion or PR with respective title and description.

Before making an issue, please verify that:

• The problem still exists on the current main branch.
• Your python dependencies are updated to recent versions.

Suggestions for improvements are always welcome!