Multiscale Supervised Unbalanced Optimal Transport Flow Matching
Abstract
Unbalanced optimal transport (UOT) provides a principled framework for modeling single-cell transitions and birth-death dynamics, but its high computational cost limits scalability to large-scale datasets. Although single-cell data often contain hierarchical annotations and known transition priors, existing UOT approximations rarely exploit this multiscale structure or prior knowledge. We introduce Multiscale Supervised Unbalanced Optimal Transport Flow Matching (MUST-FM), a simulation-free framework that scales UOT by leveraging hierarchical data structure. MUST-FM further supports an optional supervised formulation that incorporates transition priors, such as cell lineages, to guide the learning of displacement fields and mass variations. Experiments show that MUST-FM reduces computational overhead while achieving robust and biologically meaningful trajectory inference, enabling dynamic modeling of atlas-scale single-cell datasets. The code is available at https://github.com/QiangweiPeng/MUST-FM.
Related
- A Network Simplex Framework for Discrete Unbalanced Optimal Transport
- Existence, Equivalence, and Convergence of Unbalanced Optimal Transport
- Gaussian-Restricted Barycenters for KL-Unbalanced Optimal Transport: Variational Theory and Fixed-Point Convergence
- Dynamical Unbalanced Optimal Transport: Duality, Reachability, and Computation
- One-Step Generative Modeling via Unbalanced Optimal Transport
BibTeX
@misc{peng2026multiscale,
title = {{Multiscale Supervised Unbalanced Optimal Transport Flow Matching}},
author = {Qiangwei Peng and Lezhi Chen and Peijie Zhou},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2605.16529},
eprint = {2605.16529},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2605.16529},
note = {CC BY-NC-SA 4.0},
}
RIS
TY - UNPB
TI - Multiscale Supervised Unbalanced Optimal Transport Flow Matching
AU - Qiangwei Peng
AU - Lezhi Chen
AU - Peijie Zhou
PY - 2026
DA - 2026/10/30
PB - arXiv
DO - 10.48550/arxiv.2605.16529
UR - https://arxiv.org/abs/2605.16529
AB - Unbalanced optimal transport (UOT) provides a principled framework for modeling single-cell transitions and birth-death dynamics, but its high computational cost limits scalability to large-scale datasets. Although single-cell data often contain hierarchical annotations and known transition priors, existing UOT approximations rarely exploit this multiscale structure or prior knowledge. We introduce Multiscale Supervised Unbalanced Optimal Transport Flow Matching (MUST-FM), a simulation-free framework that scales UOT by leveraging hierarchical data structure. MUST-FM further supports an optional supervised formulation that incorporates transition priors, such as cell lineages, to guide the learning of displacement fields and mass variations. Experiments show that MUST-FM reduces computational overhead while achieving robust and biologically meaningful trajectory inference, enabling dynamic modeling of atlas-scale single-cell datasets. The code is available at https://github.com/QiangweiPeng/MUST-FM.
ER -