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

TL;DR

This paper introduces WFR-MFM, a novel one-step inference method for dynamic unbalanced optimal transport that significantly accelerates predictions in single-cell biology applications without sacrificing accuracy.

Contribution

The paper presents a mean-flow framework for unbalanced flow matching that enables fast, one-step inference, advancing the scalability of dynamic optimal transport models.

Findings

Achieves orders-of-magnitude faster inference than existing methods.

Maintains high predictive accuracy across synthetic and real datasets.

Enables efficient perturbation response prediction on large datasets.

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.