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

• URL: http://arxiv.org/abs/2601.20606v1
• Date: Wed, 28 Jan 2026 13:41:52 GMT
• Title: WFR-MFM: One-Step Inference for Dynamic Unbalanced Optimal Transport
• Authors: Xinyu Wang, Ruoyu Wang, Qiangwei Peng, Peijie Zhou, Tiejun Li,
• Abstract summary: We propose a mean-flow framework for unbalanced flow matching that summarizes both transport and mass-growth dynamics over arbitrary time intervals.<n>We further build on this framework to develop Wasserstein-Fisher-Rao Mean Flow Matching (WFR-MFM)<n>WFR-MFM achieves orders-of-magnitude faster inference than a range of existing baselines while maintaining high predictive accuracy.
• Score: 9.661016195647212
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• WFR-FM: Simulation-Free Dynamic Unbalanced Optimal Transport [17.9375439076052]
  Existing WFR solvers are often unstable, computationally expensive, and difficult to scale.<n>Here we introduce WFR Flow Matching (WFR-FM), a simulation-free training algorithm that unifies flow matching with dynamic unbalanced OT.<n> Empirically, WFR-FM yields more accurate and robust trajectory inference in single-cell biology.
  arXiv  Detail & Related papers  (2026-01-11T08:48:08Z)
• Drift No More? Context Equilibria in Multi-Turn LLM Interactions [58.69551510148673]
  contexts drift is the gradual divergence of a model's outputs from goal-consistent behavior across turns.<n>Unlike single-turn errors, drift unfolds temporally and is poorly captured by static evaluation metrics.<n>We show that multi-turn drift can be understood as a controllable equilibrium phenomenon rather than as inevitable decay.
  arXiv  Detail & Related papers  (2025-10-09T04:48:49Z)
• Simulation-Free Differential Dynamics through Neural Conservation Laws [22.4113724471297]
  We present a novel simulation-free framework for training continuous-time diffusion processes over very general objective functions.<n>We propose a coupled parameterization which jointly models a time-dependent density function, or probability path, and the dynamics of a diffusion process that generates this probability path.
  arXiv  Detail & Related papers  (2025-06-23T13:04:23Z)
• Efficient Transformed Gaussian Process State-Space Models for Non-Stationary High-Dimensional Dynamical Systems [49.819436680336786]
  We propose an efficient transformed Gaussian process state-space model (ETGPSSM) for scalable and flexible modeling of high-dimensional, non-stationary dynamical systems.<n>Specifically, our ETGPSSM integrates a single shared GP with input-dependent normalizing flows, yielding an expressive implicit process prior that captures complex, non-stationary transition dynamics.<n>Our ETGPSSM outperforms existing GPSSMs and neural network-based SSMs in terms of computational efficiency and accuracy.
  arXiv  Detail & Related papers  (2025-03-24T03:19:45Z)
• AMR-Transformer: Enabling Efficient Long-range Interaction for Complex Neural Fluid Simulation [33.63726923336252]
  We propose AMR-Transformer, an efficient and accurate neural CFD-solving pipeline.<n>It integrates a novel adaptive mesh refinement scheme with a Navier-Stokes constraint-aware fast pruning module.<n>Our approach achieves an order-of-magnitude improvement in accuracy over baseline models.
  arXiv  Detail & Related papers  (2025-03-13T11:16:42Z)
• Force Matching with Relativistic Constraints: A Physics-Inspired Approach to Stable and Efficient Generative Modeling [26.145559807686706]
  We introduce Force Matching (ForM), a framework for generative modeling that incorporates special relativistic mechanics.<n>ForM imposes a velocity constraint, ensuring that sample velocities remain bounded within a constant limit.<n>We show that ForM provides a promising pathway toward achieving stable, efficient, and flexible generative processes.
  arXiv  Detail & Related papers  (2025-02-12T06:30:01Z)
• Physics-guided Active Sample Reweighting for Urban Flow Prediction [75.24539704456791]
  Urban flow prediction is a nuanced-temporal modeling that estimates the throughput of transportation services like buses, taxis and ride-driven models. Some recent prediction solutions bring remedies with the notion of physics-guided machine learning (PGML) We develop a atized physics-guided network (PN), and propose a data-aware framework Physics-guided Active Sample Reweighting (P-GASR)
  arXiv  Detail & Related papers  (2024-07-18T15:44:23Z)
• Improving and generalizing flow-based generative models with minibatch optimal transport [90.01613198337833]
  We introduce the generalized conditional flow matching (CFM) technique for continuous normalizing flows (CNFs) CFM features a stable regression objective like that used to train the flow in diffusion models but enjoys the efficient inference of deterministic flow models. A variant of our objective is optimal transport CFM (OT-CFM), which creates simpler flows that are more stable to train and lead to faster inference.
  arXiv  Detail & Related papers  (2023-02-01T14:47:17Z)
• Manifold Interpolating Optimal-Transport Flows for Trajectory Inference [64.94020639760026]
  We present a method called Manifold Interpolating Optimal-Transport Flow (MIOFlow) MIOFlow learns, continuous population dynamics from static snapshot samples taken at sporadic timepoints. We evaluate our method on simulated data with bifurcations and merges, as well as scRNA-seq data from embryoid body differentiation, and acute myeloid leukemia treatment.
  arXiv  Detail & Related papers  (2022-06-29T22:19:03Z)
• TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics [74.43710101147849]
  We present TrajectoryNet, which controls the continuous paths taken between distributions to produce dynamic optimal transport. We show how this is particularly applicable for studying cellular dynamics in data from single-cell RNA sequencing (scRNA-seq) technologies.
  arXiv  Detail & Related papers  (2020-02-09T21:00:38Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.