Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

WFR-FM: Simulation-Free Dynamic Unbalanced Optimal Transport

About

The Wasserstein-Fisher-Rao (WFR) metric extends dynamic optimal transport (OT) by coupling displacement with change of mass, providing a principled geometry for modeling unbalanced snapshot dynamics. Existing WFR solvers, however, are often unstable, computationally expensive, and difficult to scale. Here we introduce WFR Flow Matching (WFR-FM), a simulation-free training algorithm that unifies flow matching with dynamic unbalanced OT. Unlike classical flow matching which regresses only a transport vector field, WFR-FM simultaneously regresses a vector field for displacement and a scalar growth rate function for birth-death dynamics, yielding continuous flows under the WFR geometry. Theoretically, we show that minimizing the WFR-FM loss exactly recovers WFR geodesics. Empirically, WFR-FM yields more accurate and robust trajectory inference in single-cell biology, reconstructing consistent dynamics with proliferation and apoptosis, estimating time-varying growth fields, and applying to generative dynamics under imbalanced data. It outperforms state-of-the-art baselines in efficiency, stability, and reconstruction accuracy. Overall, WFR-FM establishes a unified and efficient paradigm for learning dynamical systems from unbalanced snapshots, where not only states but also mass evolve over time. The Python code is available at https://github.com/QiangweiPeng/WFR-FM.

Qiangwei Peng, Zihan Wang, Junda Ying, Yuhao Sun, Qing Nie, Lei Zhang, Tiejun Li, Peijie Zhou• 2026

Related benchmarks

TaskDatasetResultRank
Trajectory InferenceEB dataset 5D (test)
W1 (t=1)0.324
23
Population Growth Modeling50D CITE t=2
W1 Score27.478
12
Population Growth Modeling50D CITE t=1
W127.831
12
Population Growth Modeling50D CITE t=3
W134.784
12
Cell population dynamics prediction50D Mouse (t=1)
W1 Score5.486
11
Trajectory Inference100D EB dataset t=1
W19.941
11
Trajectory Inference100D EB dataset t=2
W111.04
11
Trajectory Inference100D EB dataset (t=3)
W1 Error11.516
11
Trajectory InferenceDyngen t=1
W1 Score0.11
11
Trajectory InferenceDyngen (t=2)
W10.098
11
Showing 10 of 28 rows

Other info

Follow for update