Curly Flow Matching for Learning Non-gradient Field Dynamics

Conference: NeurIPS 2025
arXiv: 2510.26645
Code: GitHub
Area: Computational Biology
Keywords: Flow Matching, Schrödinger Bridge, Non-gradient Field Dynamics, Trajectory Inference, Optimal Transport

TL;DR

The authors propose Curly Flow Matching (Curly-FM), which designs a Schrödinger Bridge problem with a non-zero reference drift. This allows flow matching to learn non-gradient field dynamics, such as periodic and rotational behaviors, breaking the limitation of traditional methods that can only model gradient fields.

Background & Motivation

Modeling the transport dynamics of natural processes from population-level observations is a fundamental problem in the natural sciences. Current mainstream methods (such as Conditional Flow Matching, OT-CFM) rely on the principle of least action, assuming that systems follow gradient field dynamics, where trajectories minimize the energy functional between two probability measures.

However, many real-world systems naturally exhibit non-gradient field and periodic behaviors: - Single-cell RNA cell cycle: Cells move along periodic trajectories in the gene expression space. - Computational fluid dynamics: Vortex structures generate rotational flow fields. - Ocean currents: Complex rotational and periodic patterns are observed.

These non-gradient (curl) components are fundamentally uncapturable by existing flow matching and bridge matching methods. Trajectories generated by traditional methods are straight-line interpolations, completely ignoring the true dynamics of the underlying system.

Method

Overall Architecture

Curly-FM is a two-stage, simulation-free framework:

First Stage — Learning the Reference Drift (Algorithm 1): A neural path interpolant is utilized to learn the drift of the reference process, allowing the generated paths to curve and match the known velocity field information.

Second Stage — Solving the Transportation Plan (Algorithm 2): Based on the learned reference drift, the optimal coupling is computed via a marginal flow matching objective to train the final drift model.

Key Designs

1. Non-Zero Drift Reference Process

Unlike traditional Schrödinger Bridge methods that use zero-drift (Brownian motion) reference processes, Curly-FM constructs a reference process \(Q\) with a non-zero drift, which is built from velocity information inferred from the data. Specifically, the reference drift \(f_t\) is a continuous field obtained by smoothing velocities at discrete observation points using a kernel function \(\kappa_t\).

2. Neural Path Interpolant

A neural network \(\phi_{t,\eta}(x_0, x_1)\) is employed to parameterize the mean of the conditional bridge, ensuring that it is no longer a straight line between \(x_0\) and \(x_1\), but rather a curved path. The conditional distribution is formulated as:

\[P_{t|0,1}(x_t) = \mathcal{N}(x_t; t x_1 + (1-t)x_0 + t(1-t)\phi_{t,\eta}(x_0, x_1), \sigma_t^2 t(1-t))\]

The interpolant matches the reference drift by minimizing \(\mathbb{E}[\|\partial_t X_{t,\eta} - f_t(X_{t,\eta})\|^2]\).

3. Optimal Coupling Computation

The learned neural path interpolant is used to estimate the OT cost \(c(x_0, x_1)\), and the coupling \(\pi^*\) is obtained by solving the entropy-regularized optimal transport using a mini-batch OT approximation. The cost is approximated by Monte Carlo estimation of the path length integral.

Loss & Training

Stage 1 Loss (learning reference drift matching): $\(\mathcal{L}(\eta) = \mathbb{E}_{(X_0, X_1, t)}\left[\|\partial_t X_{t,\eta} - f_t(X_{t,\eta})\|^2\right]\)$

Stage 2 Loss (flow matching + score matching):

flow_loss = E[||v_t(x_t, t) - x_dot_t||^2]
score_loss = E[||(lambda_t * s_t(x_t, t) + eps)||^2]
loss = flow_loss + score_loss

Where the score model is used to support cases with stochasticity (\(\sigma > 0\)).

Inference: Forward propagation is performed via SDE/ODE integration using the learned drift model.

Key Experimental Results

Main Results

Table 1: Single-cell RNA Trajectory Inference (Cell Cycle Data)

Method	W1 ↓	W2 ↓	Cosine Dist ↓
CFM	High	High	~1.0 (straight trajectories)
OT-CFM	High	High	~1.0
CurlyFM	Lowest	Lowest	~0.03

Table 2: Ocean Current Trajectory Inference

Method	W1 ↓	Cosine Dist ↓	Compute Time (h)
DM-SB	-	-	15.44
TrajectoryNet	-	-	7.44
SBIRR	-	-	4.67
Vanilla-SB	-	-	0.43
CurlyFM	0.062±0.003	0.034±0.006	0.06

On the ocean current task, CurlyFM requires only about 4 minutes of training, which is ~78x faster than SBIRR and ~124x faster than TrajectoryNet.

Ablation Study

Robustness to Noisy Reference Drift (Table R2)

Noise Ratio β	W1 ↓	W2 ↓	Cosine Dist ↓
0.00	0.062±0.003	0.143±0.010	0.034±0.006
0.25	0.057±0.033	0.021±0.036	0.051±0.030
0.50	0.087±0.047	0.301±0.085	0.091±0.046
0.75	0.261±0.123	0.381±0.120	0.145±0.062
1.00	0.428±0.157	0.445±0.121	0.237±0.079

CurlyFM maintains robust performance under a 25% noise level, demonstrating resilience to approximation errors in the reference velocity field.

Stochasticity Ablation (Impact of σ)

σ	W1 ↓	W2 ↓	Cosine Dist ↓
0.01	0.061±0.003	0.141±0.009	0.028±0.066
0.10	0.062±0.002	0.145±0.011	0.066±0.008
1.00	0.145±0.009	0.474±0.058	0.871±0.048

The low-stochasticity setting achieves the best performance, which is consistent with the characteristic of low noise levels in practical applications.

Key Findings

1. Fundamental limitation of traditional flow matching: OT-CFM yields a cosine distance close to 1.0, indicating that its trajectories are almost purely linear, failing completely to recover rotational dynamics.
2. Significant efficiency advantages: As a simulation-free method, CurlyFM is 1-2 orders of magnitude faster than simulation-based methods such as TrajectoryNet and SBIRR.
3. Multi-marginal expansion: CurlyFM naturally extends to multi-time-point setups, training alternately on adjacent marginal pairs.
4. Comparison with GSBM: CurlyFM outperforms GSBM (with 15 control points) in terms of W1 and cosine distance, though with a slightly higher computational cost.

Highlights & Insights

• Core Insight: Incorporating a non-zero drift reference process into the Schrödinger Bridge framework overcomes the limitation that flow matching can only model gradient fields.
• Elegant Two-Stage Design: Fitting the reference drift first and then solving the transport problem avoids the convergence difficulties typical of iterative algorithms.
• High Practicality: It only requires velocity information (e.g., RNA velocity) at the data sampling points, from which a continuous reference field can be constructed through kernel smoothing.
• Simulation-free: Training efficiency is improved by dozens of times compared to methods requiring backpropagation through time or simulation.

Limitations & Future Work

1. Dependency on Velocity Field Quality: Performance relies on the approximation accuracy of the reference velocity field. The method cannot be directly applied when velocity information is unavailable.
2. Focus on Determinism: Main experiments are conducted in the deterministic limit (\(\sigma = 0\)), and performance drops significantly in highly stochastic scenarios.
3. Limited Theoretical Guarantees: The method relies on practical approximations (modeling conditional bridges as Brownian bridges, mini-batch OT) and lacks strict convergence guarantees.
4. Kernel Selection: The choice of kernel \(\kappa\) impacts the quality of the reference drift, but the paper lacks systematic guidance on kernel selection.

Related Work & Insights

• Flow Matching (Lipman et al. 2023): The foundational framework for Curly-FM.
• OT-CFM (Tong et al. 2024): Mini-batch OT + CFM, but limited to gradient fields.
• DSBM (Shi et al. 2024): Diffusion Schrödinger Bridge Matching, dealing with zero-drift references.
• GSBM (Liu et al. 2024): Generalized SBM using spline interpolation, but Curly-FM is more flexible with neural networks.
• TrajectoryNet (Tong et al. 2020): ODE-based trajectory inference requiring simulations and incurring high computational costs.

Rating

Dimension	Score (1-5)
Novelty	4 — The non-zero drift reference process is a key innovation
Technical Quality	4 — The method is sound, although some theoretical details require improvement
Experimental Thoroughness	4 — Validation across multiple scenarios with comprehensive ablation studies
Writing Quality	3 — The method description (Section 3.1) is somewhat difficult to follow
Impact	4 — Holds direct application value for fields like single-cell biology

Related Papers

• [NeurIPS 2025] Energy Matching: Unifying Flow Matching and Energy-Based Models for Generative Modeling
• [ICML 2025] Improving Flow Matching by Aligning Flow Divergence
• [NeurIPS 2025] Prior-Guided Flow Matching for Target-Aware Molecule Design with Learnable Atom Number
• [ICML 2025] Flexibility-conditioned Protein Structure Design with Flow Matching
• [NeurIPS 2025] JAMUN: Bridging Smoothed Molecular Dynamics and Score-Based Learning for Conformational Ensembles