Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

WFR-FM: Simulation-Free Dynamic Unbalanced Optimal Transport

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

TL;DR

This work introduces WFR-FM, a simulation-free flow-matching framework for dynamic unbalanced OT under the Wasserstein–Fisher–Rao metric. By jointly regressing a transport velocity field $\mathbf{v}_\theta(t,\mathbf{x})$ and a growth rate $g_\phi(t,\mathbf{x})$, it produces continuous flows that capture both displacement and mass birth–death, with a theoretical guarantee that minimizing the WFR-FM loss recovers WFR geodesics. The method leverages conditional path constructions (CGMP/CGM), a mass-aware CUFM loss, and traveling-Gaussian dynamics to ensure OT-consistent trajectories between snapshots, while enabling fully simulation-free training. Empirically, WFR-FM delivers accurate trajectory inference and robust growth dynamics on multiple single-cell datasets, outperforming ODE-based and other unbalanced-flow baselines in efficiency, stability, and reconstruction fidelity. The framework scales to multi-time-point data and high-dimensional settings, offering a principled, scalable tool for learning dynamic systems where both states and mass evolve over time.

Abstract

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.

WFR-FM: Simulation-Free Dynamic Unbalanced Optimal Transport

TL;DR

This work introduces WFR-FM, a simulation-free flow-matching framework for dynamic unbalanced OT under the Wasserstein–Fisher–Rao metric. By jointly regressing a transport velocity field and a growth rate , it produces continuous flows that capture both displacement and mass birth–death, with a theoretical guarantee that minimizing the WFR-FM loss recovers WFR geodesics. The method leverages conditional path constructions (CGMP/CGM), a mass-aware CUFM loss, and traveling-Gaussian dynamics to ensure OT-consistent trajectories between snapshots, while enabling fully simulation-free training. Empirically, WFR-FM delivers accurate trajectory inference and robust growth dynamics on multiple single-cell datasets, outperforming ODE-based and other unbalanced-flow baselines in efficiency, stability, and reconstruction fidelity. The framework scales to multi-time-point data and high-dimensional settings, offering a principled, scalable tool for learning dynamic systems where both states and mass evolve over time.

Abstract

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.
Paper Structure (37 sections, 8 theorems, 88 equations, 9 figures, 14 tables, 2 algorithms)

This paper contains 37 sections, 8 theorems, 88 equations, 9 figures, 14 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Works
  3. Mathematical Background
  4. Simulation-free Training for WFR Problem
  5. Conditional Path Construction
  6. Conditional Loss Design
  7. Coupling in WFR-FM solves the dynamical WFR problem
  8. WFR-FM Workflow for Trajectory Inference
  9. Experiment Results
  10. Conclusion, Limitation and Discussion
  11. Proofs
  12. Proof of Theorem \ref{['thm:semi-coupling']}
  13. Proof of Theorem \ref{['thm:marginalization']}
  14. Proof of Theorem \ref{['thm:equal gradient']}
  15. Proof of Proposition \ref{['prop:CGMP']}
  16. ...and 22 more sections

Key Result

Theorem 3.1

Let $\gamma$ be the optimal coupling of the OET problem (OET), then the semi-coupling $\gamma_0(\bm{x},\bm{y})=\frac{\gamma(\bm{x},\bm{y})}{\int_\mathcal{X}\gamma(\bm{x},\bm{z})\mathrm{d} \bm{z}}\mu_0(\bm{x}), \gamma_1(\bm{x},\bm{y})=\frac{\gamma(\bm{x},\bm{y})}{\int_\mathcal{X}\gamma(\bm{z},\bm{y})

Figures (9)

  • Figure 1: Overview of WFR-FM. (a) Dynamic unbalanced OT aims to interpolate the intermediate state in a way that minimizes the action. (b) Kantorovich unbalanced OT seeks the semi-coupling that minimizes total cost. (c) Analytical trajectory from a Dirac distribution to another Dirac distribution. (d) Jointly train the velocity field and the growth rate in the simulation-free manner.
  • Figure 2: Learned trajectories on the (a) Simulation Gene dataset (b) EB dataset
  • Figure 3: Efficiency on the 100D EB dataset
  • Figure 4: (a) True growth rate (b) Predicted growth rate by WFR-FM
  • Figure 5: (a) Learned trajectories (b) growth rate on the dygen dataset
  • ...and 4 more figures

Theorems & Definitions (8)

  • Theorem 3.1
  • Theorem 4.1
  • Proposition 4.1
  • Theorem 4.2
  • Proposition 4.2
  • Proposition 4.3
  • Proposition 5.1
  • Lemma A.1