WFR-MFM: One-Step Inference for Dynamic Unbalanced Optimal Transport

TL;DR

Dynamic unbalanced OT modeling is essential for capturing mass-changing cellular processes but traditional ODE-based inference is computationally prohibitive. The authors propose a mean-flow framework that summarizes transport and growth over finite intervals via mean velocity $\\mathbf{v}$ and mean growth $\\mathbf{h}$, enabling direct one-step, trajectory-free updates. Specializing to Wasserstein-Fisher-Rao geometry, they introduce WFR-MFM, a simulation-free inference method that delivers orders-of-magnitude speedups while preserving predictive accuracy and supporting large-scale perturbation analyses. Extensive experiments on synthetic and real scRNA-seq data demonstrate fast, scalable inference with controllable speed-accuracy trade-offs, highlighting WFR-MFM’s potential for perturbation response prediction and large combinatorial condition spaces.

Abstract

Reconstructing dynamical evolution from limited observations is a fundamental challenge in single-cell biology, where dynamic unbalanced optimal transport provides a principled framework for modeling coupled transport and mass variation. However, existing approaches rely on trajectory simulation at inference time, making inference a key bottleneck for scalable applications. In this work, we propose a mean-flow framework for unbalanced flow matching that summarizes both transport and mass-growth dynamics over arbitrary time intervals using mean velocity and mass-growth fields, enabling fast one-step generation without trajectory simulation. To solve dynamic unbalanced optimal transport under the Wasserstein-Fisher-Rao geometry, we further build on this framework to develop Wasserstein-Fisher-Rao Mean Flow Matching (WFR-MFM). Across synthetic and real single-cell RNA sequencing datasets, WFR-MFM achieves orders-of-magnitude faster inference than a range of existing baselines while maintaining high predictive accuracy, and enables efficient perturbation response prediction on large synthetic datasets with thousands of conditions.

Figures (7)

• Figure 1: Learned dynamics on the Gene and EMT datasets. Top: Gene dataset; bottom: EMT dataset. (a) Individual cell trajectories with evolving weights. (b) Population-level dynamics obtained by resampling cells according to normalized weights.
• Figure 2: Speed--accuracy trade-off (Q3) on the Simulation dataset. (a) $\mathcal{W}_1$ and RME versus the number of inference steps. (b) Corresponding inference runtime versus the number of steps. For each step number, we perform 1000 independent runs.
• Figure 3: Efficiency on the 100D EB dataset. Methods are compared by $\mathcal{W}_1$ distance, training time, and inference time (log scale), with colors indicating GPU memory usage.
• Figure 4: Predictions on unseen perturbation conditions. Gene expression distributions before and after perturbation for nine unseen conditions, visualized by resampling cells using predicted weights (control denotes the unperturbed baseline).
• Figure 5: Learned dynamics on the Dyngen dataset. (a) Trajectories of individual cells with a fixed cell count, where only cell weights evolve over time. (b) Population-level dynamics obtained by normalizing weights into probabilities and resampling cells, revealing changes in cell abundance.

Theorems & Definitions (5)

• Theorem 4.1
• Remark 4.2
• proof
• Proposition B.1: Failure of equivalence under time-averaged targets
• proof